diff --git a/src/sock/enc.rs b/src/sock/enc.rs index fa57a40..7e79659 100644 --- a/src/sock/enc.rs +++ b/src/sock/enc.rs @@ -16,10 +16,17 @@ use tokio::{ RwLock, }, io::{ - self, DuplexStream, }, }; +use std::{ + io, + task::{ + Context, Poll, + }, + pin::Pin, + marker::Unpin, +}; /// Encrypted socket information. #[derive(Debug)] @@ -28,6 +35,50 @@ struct ESockInfo { them: Option, } +/// A shitty way of reversing `AsyncSink` +//TODO: Implement a decrypting reader in `chacha20stream`. +#[pin_project] +#[derive(Debug)] +struct ReverseSink { + #[pin] temp_sink: AsyncSink, + #[pin] rx: DuplexStream, + #[pin] raw: R, +} + +impl AsyncRead for ReverseSink +{ + fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll> { + let proj = self.project(); + use futures::ready; + macro_rules! try_ready { + ($poll_res:expr) => { + match ready!($poll_res) { + Ok(v) => v, + Err(e) => return Poll::Ready(Err(e)), + } + } + } + alloc_local_bytes(buf.len(), |temp_buf| -> Poll> { + //XXX: Does this work? I dunno... Maybe just pin and poll a local `async{}` block. + + // Read encrypted bytes into `temp_buf`. + let filled_buffer: &[u8] = try_ready!(proj.raw.poll_read(cx, temp_buf).map_ok(|read| { + &temp_buf[..read] + })); + + // Sink the `filled_buffer`. + let sunk: usize = try_ready!(proj.temp_sink.poll_write(cx, filled_buffer)); + debug_assert_eq!(sunk, filled_buffer.len()); + + // Ready from the `rx`. + let bytes_read: usize = try_ready!(proj.rx.poll_read(cx, &mut buf[..sunk])); + debug_assert_eq!(bytes_read, sunk); + + Poll::Ready(Ok(bytes_read)) + }) + } +} + /// A tx+rx socket. #[pin_project] #[derive(Debug)] @@ -51,8 +102,6 @@ pub struct ESockWriteHalf(Arc, #[pin] AsyncSink); #[derive(Debug)] pub struct ESockReadHalf( Arc, - - #[pin] R, // read from this (raw.) - #[pin] AsyncSink, // sink raw from `R` here, outputs decrypted bytes into next. - #[pin] DuplexStream, // read decrypted bytes from here. + + #[pin] ReverseSink, );