From 077565b9087383a8092ae89f225716ea153841ff Mon Sep 17 00:00:00 2001
From: Brian Picciano <mediocregopher@gmail.com>
Date: Sat, 22 Jul 2023 13:40:19 +0200
Subject: [PATCH] Implemented AsyncReadCapture and AsyncReadWrite

These are utility types which will help capture the ClientHello as it is
being read and parsed, as well as any data following it which might get
buffered in the acceptore. This way we can get that data back out in
case we want to switch into transparent TCP mode again.
---
 src/common/mod.rs | 101 ++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs        |  40 +++++++++++++++++-
 2 files changed, 140 insertions(+), 1 deletion(-)

diff --git a/src/common/mod.rs b/src/common/mod.rs
index fde34c0..16e735c 100644
--- a/src/common/mod.rs
+++ b/src/common/mod.rs
@@ -359,5 +359,106 @@ impl<'a, 'b, T: AsyncRead + Unpin> Read for SyncReadAdapter<'a, 'b, T> {
     }
 }
 
+/// Wraps an AsyncRead and AsyncWrite instance together to produce a single type which implements
+/// AsyncRead + AsyncWrite.
+pub struct AsyncReadWrite<R, W> {
+    r: Pin<Box<R>>,
+    w: Pin<Box<W>>,
+}
+
+impl<R, W> AsyncReadWrite<R, W>
+where
+    R: Unpin,
+    W: Unpin,
+{
+    pub fn new(r: R, w: W) -> Self {
+        Self {
+            r: Box::pin(r),
+            w: Box::pin(w),
+        }
+    }
+
+    pub fn into_inner(self) -> (R, W) {
+        (*Pin::into_inner(self.r), *Pin::into_inner(self.w))
+    }
+}
+
+impl<R, W> AsyncRead for AsyncReadWrite<R, W>
+where
+    R: AsyncRead + Unpin,
+{
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        self.r.as_mut().poll_read(cx, buf)
+    }
+}
+
+impl<R, W> AsyncWrite for AsyncReadWrite<R, W>
+where
+    W: AsyncWrite + Unpin,
+{
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        self.w.as_mut().poll_write(cx, buf)
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        self.w.as_mut().poll_flush(cx)
+    }
+
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        self.w.as_mut().poll_shutdown(cx)
+    }
+}
+
+/// Wraps an AsyncRead in order to capture all bytes which have been read by it into an internal
+/// buffer.
+pub struct AsyncReadCapture<R> {
+    r: Pin<Box<R>>,
+    buf: Vec<u8>,
+}
+
+impl<R> AsyncReadCapture<R>
+where
+    R: AsyncRead + Unpin,
+{
+    /// Initializes an AsyncReadCapture with an empty internal buffer of the given size.
+    pub fn with_capacity(r: R, cap: usize) -> Self {
+        Self {
+            r: Box::pin(r),
+            buf: Vec::with_capacity(cap),
+        }
+    }
+
+    pub fn into_inner(self) -> (R, Vec<u8>) {
+        (*Pin::into_inner(self.r), self.buf)
+    }
+}
+
+impl<R> AsyncRead for AsyncReadCapture<R>
+where
+    R: AsyncRead,
+{
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        let res = self.r.as_mut().poll_read(cx, buf);
+
+        if let Poll::Ready(Ok(())) = res {
+            self.buf.extend_from_slice(buf.filled());
+        }
+
+        res
+    }
+}
+
 #[cfg(test)]
 mod test_stream;
diff --git a/src/lib.rs b/src/lib.rs
index 000245c..576d247 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -49,7 +49,7 @@ pub mod client;
 mod common;
 pub mod server;
 
-use common::{MidHandshake, Stream, TlsState};
+use common::{AsyncReadCapture, AsyncReadWrite, MidHandshake, Stream, TlsState};
 use rustls::{ClientConfig, ClientConnection, CommonState, ServerConfig, ServerConnection};
 use std::future::Future;
 use std::io;
@@ -333,6 +333,44 @@ where
     }
 }
 
+type IOWithCapture<IO> = AsyncReadWrite<AsyncReadCapture<IO>, IO>;
+
+pub struct TransparentConfigAcceptor<IO> {
+    acceptor: LazyConfigAcceptor<IOWithCapture<IO>>,
+}
+
+impl<IO> TransparentConfigAcceptor<IO>
+where
+    IO: AsyncRead + AsyncWrite + Unpin,
+{
+    pub fn new(acceptor: rustls::server::Acceptor, io: IO) -> Self {
+        let r = AsyncReadCapture::with_capacity(io, 512);
+        let rw = AsyncReadWrite::new(r, io);
+        Self {
+            acceptor: LazyConfigAcceptor::new(acceptor, rw),
+        }
+    }
+
+    //pub fn into_lazy_config_acceptor(self) -> LazyConfigAcceptor<IO> {
+    //    self.acceptor
+    //}
+}
+
+//impl<IO> Future for TransparentConfigAcceptor<IO>
+//where
+//    IO: AsyncRead + AsyncWrite + Unpin,
+//{
+//    type Output = io::Result<TransparentStartHandshake<IO>>;
+//
+//    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+//        let this = self.get_mut();
+//    }
+//}
+
+pub struct TransparentStartHandshake<IO> {
+    h: StartHandshake<IO>,
+}
+
 /// Future returned from `TlsConnector::connect` which will resolve
 /// once the connection handshake has finished.
 pub struct Connect<IO>(MidHandshake<client::TlsStream<IO>>);