From 0f7f22f2907543ad9a724b6de55d951b6381fa61 Mon Sep 17 00:00:00 2001
From: Avril <flanchan@cumallover.me>
Date: Sun, 15 Aug 2021 17:11:07 +0100
Subject: [PATCH] ESock: Added `exchange_unpin()`. Start `Exchange` future for
 `exchange()` method that will work on non-Unpin types too. (Requires specific
 implemented future.)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fortune for rsh's current commit: Middle blessing − 中吉
---
 src/sock/enc.rs | 82 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 82 insertions(+)

diff --git a/src/sock/enc.rs b/src/sock/enc.rs
index 064cff3..5633e14 100644
--- a/src/sock/enc.rs
+++ b/src/sock/enc.rs
@@ -29,6 +29,9 @@ use std::{
     marker::Unpin,
 };
 
+/// Max size to read when exchanging keys
+const TRANS_KEY_MAX_SIZE: usize = 4096;
+
 /// Encrypted socket information.
 #[derive(Debug)]
 struct ESockInfo {
@@ -48,6 +51,85 @@ pub struct ESock<W, R> {
     tx: AsyncSink<W>,
 }
 
+impl<W: AsyncWrite, R: AsyncRead> ESock<W, R>
+{
+    pub fn inner(&self) -> (&W, &R)
+    {
+	(self.tx.inner(), self.rx.inner())
+    }
+    
+    fn inner_mut(&mut self) -> (&mut W, &mut R)
+    {
+	(self.tx.inner_mut(), self.rx.inner_mut())
+    }
+
+    /// Create a future that exchanges keys
+    pub fn exchange(&mut self) -> Exchange<'_, W, R>
+    {
+	Exchange{sock: self}
+    }
+}
+
+impl<W: AsyncWrite+ Unpin, R: AsyncRead + Unpin> ESock<W, R>
+{
+    /// Exchange keys.
+    pub async fn exchange_unpin(&mut self) -> eyre::Result<()>
+    {
+	use tokio::prelude::*;
+	let our_key = self.info.read().await.us.get_public_parts();
+	let (tx, rx) = self.inner_mut();
+	let read_fut = {
+	    
+	    async move {
+		// Read the public key from `rx`.
+		//TODO: Find pubkey max size.
+		let mut sz_buf = [0u8; std::mem::size_of::<u64>()];
+		rx.read_exact(&mut sz_buf[..]).await?;
+		let sz= match usize::try_from(u64::from_be_bytes(sz_buf))? {
+		    x if x > TRANS_KEY_MAX_SIZE => return Err(eyre!("Recv'd key size exceeded max")),
+		    x => x
+		};
+		let mut key_bytes = Vec::with_capacity(sz);
+		tokio::io::copy(&mut rx.take(sz as u64), &mut key_bytes).await?;
+		if key_bytes.len() != sz {
+		    return Err(eyre!("Could not read required bytes"));
+		}
+		let k = RsaPublicKey::from_bytes(key_bytes)?;
+
+		Result::<RsaPublicKey, eyre::Report>::Ok(k)
+	    }
+	};
+	let write_fut = {
+	    let key_bytes = our_key.to_bytes();
+	    assert!(key_bytes.len() <= TRANS_KEY_MAX_SIZE);
+	    let sz_buf = u64::try_from(key_bytes.len())?.to_be_bytes();
+	    async move {
+		tx.write_all(&sz_buf[..]).await?;
+		tx.write_all(&key_bytes[..]).await?;
+		Result::<(), eyre::Report>::Ok(())
+	    }
+	};
+	let (send, recv) = tokio::join! [write_fut, read_fut];
+	send?;
+	let recv = recv?;
+	self.info.write().await.them = Some(recv);
+	Ok(())
+    }
+}
+
+pub struct Exchange<'a, W, R>
+{
+    sock: &'a mut ESock<W, R>,
+}
+
+impl<'a, W: AsyncWrite, R: AsyncRead> Future for Exchange<'a, W, R>
+{
+    type Output = eyre::Result<()>;
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+	todo!("This is going to be dificult to implement... We don't have access to write_all and read_exact")
+    }
+}
+
 /// Write half for `ESock`.
 #[pin_project]
 #[derive(Debug)]