From edeb2ffee7b40d4d4a658ae8ddf467da7d273961 Mon Sep 17 00:00:00 2001
From: Avril <flanchan@cumallover.me>
Date: Mon, 16 Aug 2021 15:13:18 +0100
Subject: [PATCH] Added AsyncWrite and AsyncRead impls to `ESock<W,R>`
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fortune for rsh's current commit: Half blessing − 半吉
---
 src/sock/enc.rs | 40 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/sock/enc.rs b/src/sock/enc.rs
index 97c4425..1cccd2e 100644
--- a/src/sock/enc.rs
+++ b/src/sock/enc.rs
@@ -93,7 +93,6 @@ impl Default for ESockState
     }
 }
 
-
 /// Contains a cc20 Key and IV that can be serialized and then encrypted
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 struct ESockSessionKey
@@ -110,7 +109,6 @@ impl fmt::Display for ESockSessionKey
     }
 }
 
-
 impl ESockSessionKey
 {
     /// Generate a new cc20 key + iv,
@@ -375,6 +373,44 @@ impl<W: AsyncWrite+ Unpin, R: AsyncRead + Unpin> ESock<W, R>
     }
 }
 
+impl<W, R> AsyncWrite for ESock<W, R>
+where W: AsyncWrite
+{
+    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, io::Error>> {
+	//XXX: If the encryption state of the socket is changed between polls, this breaks. Idk if we can do anything about that tho. 
+	if self.state.encw {
+	    self.project().tx.poll_write(cx, buf)
+	} else {
+	    // SAFETY: Uhh... well I think this is fine? Because we can project the container.
+	    // TODO: Can we project the `tx`? Or maybe add a method in `AsyncSink` to map a pinned sink to a `Pin<&mut W>`?
+	    unsafe { self.map_unchecked_mut(|this| this.tx.inner_mut()).poll_write(cx, buf)}
+	}
+    }
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+	// Should we do anything else here?
+	// Should we clear foreign key/current session key?
+	self.project().tx.poll_shutdown(cx)
+    }
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+	self.project().tx.poll_flush(cx)
+    }
+}
+
+impl<W, R> AsyncRead for ESock<W, R>
+where R: AsyncRead
+{
+    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
+	//XXX: If the encryption state of the socket is changed between polls, this breaks. Idk if we can do anything about that tho. 
+	if self.state.encw {
+	    self.project().rx.poll_read(cx, buf)
+	} else {
+	    // SAFETY: Uhh... well I think this is fine? Because we can project the container.
+	    // TODO: Can we project the `tx`? Or maybe add a method in `AsyncSink` to map a pinned sink to a `Pin<&mut W>`?
+	    unsafe { self.map_unchecked_mut(|this| this.rx.inner_mut()).poll_read(cx, buf)}
+	}
+    }
+}
+
 /// Write half for `ESock`.
 #[pin_project]
 #[derive(Debug)]