Merge branch 'feat/socket-api-debug' into feat/socket-api

Author: th4s

crates/tls/client/src/conn.rs

@@ -549,6 +549,7 @@ impl ConnectionCommon {
         if let Ok(0) = res {
             self.common_state.has_seen_eof = true;
         }
+        let res = res.inspect(|v| println!("read tls:  {v} bytes"));
         res
     }
 

crates/tlsn/src/prover.rs

@@ -10,7 +10,6 @@ pub mod state;
 pub use conn::TlsConnection;
 pub use control::ProverControl;
 pub use error::ProverError;
-use futures_plex::DuplexStream;
 pub use tlsn_core::ProverOutput;
 
 use crate::{
@@ -25,7 +24,7 @@ use crate::{
     utils::{await_with_copy_io, build_mt_context},
 };
 
-use futures::{AsyncRead, AsyncWrite, FutureExt, TryFutureExt};
+use futures::{AsyncRead, AsyncWrite, FutureExt, TryFutureExt, ready};
 use rustls_pki_types::CertificateDer;
 use serio::{SinkExt, stream::IoStreamExt};
 use std::{
@@ -270,6 +269,8 @@ impl Prover<state::Setup> {
         S: AsyncRead + AsyncWrite + Send + Unpin,
         T: AsyncRead + AsyncWrite + Send + Unpin,
     {
+        let (client_to_server, server_to_client) = futures_plex::duplex(BUF_CAP);
+
         Prover {
             config: self.config,
             span: self.span,
@@ -281,6 +282,8 @@ impl Prover<state::Setup> {
                 output: None,
                 client_io: self.state.client_io,
                 verifier_io: self.state.verifier_io,
+                tls_client_duplex: client_to_server,
+                server_duplex: server_to_client,
                 server_socket: server_io,
                 verifier_socket: verifier_io,
                 client_closed: false,
@@ -323,13 +326,13 @@ where
     fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let mut state = Pin::new(&mut self.state).project();
 
-        let (mut duplex_1, mut duplex_2) = futures_plex::duplex(BUF_CAP);
         loop {
             let mut progress = false;
 
             progress |= Self::io_client_conn(&mut state, cx)?;
-            progress |= Self::io_client_server(&mut state, cx, &mut duplex_1, &mut duplex_2)?;
+            progress |= Self::io_client_server(&mut state, cx)?;
             progress |= Self::io_client_verifier(&mut state, cx)?;
+
             _ = state.mux_fut.poll_unpin(cx)?;
 
             if state.output.is_none()
@@ -339,8 +342,14 @@ where
             }
 
             if !progress {
+                cx.waker().wake_by_ref();
                 return Poll::Pending;
-            } else if *state.server_closed && *state.client_closed && state.output.is_some() {
+            }
+
+            if *state.server_closed && state.output.is_some() {
+                ready!(state.client_io.poll_close(cx))?;
+                ready!(state.server_socket.poll_close(cx))?;
+
                 return Poll::Ready(Ok(()));
             }
         }
@@ -359,10 +368,11 @@ where
         let mut progress = false;
 
         // tls_conn -> tls_client
-        if let Poll::Ready(mut simplex) = state.client_io.as_mut().poll_lock_read(cx)
+        if state.tls_client.wants_write()
+            && let Poll::Ready(mut simplex) = state.client_io.as_mut().poll_lock_read(cx)
             && let Poll::Ready(buf) = simplex.poll_get(cx)?
         {
-            if buf.len() > 0 {
+            if !buf.is_empty() {
                 let write = state.tls_client.write(buf)?;
                 if write > 0 {
                     progress = true;
@@ -376,72 +386,64 @@ where
         }
 
         // tls_client -> tls_conn
-        if let Poll::Ready(mut simplex) = state.client_io.as_mut().poll_lock_write(cx)
+        if state.tls_client.wants_read()
+            && let Poll::Ready(mut simplex) = state.client_io.as_mut().poll_lock_write(cx)
             && let Poll::Ready(buf) = simplex.poll_mut(cx)?
+            && let read = state.tls_client.read(buf)?
+            && read > 0
         {
-            if buf.len() > 0
-                && let read = state.tls_client.read(buf)?
-                && read > 0
-            {
-                progress = true;
-                simplex.advance_mut(read);
-            }
+            progress = true;
+            simplex.advance_mut(read);
         }
-
         Ok(progress)
     }
 
     fn io_client_server(
         state: &mut ConnectedProj<S, T>,
         cx: &mut Context,
-        duplex_1: &mut DuplexStream,
-        duplex_2: &mut DuplexStream,
     ) -> Result<bool, ProverError> {
         let mut progress = false;
-        let mut duplex_1 = Pin::new(duplex_1);
-        let mut duplex_2 = Pin::new(duplex_2);
 
         // server_socket -> duplex
-        if let Poll::Ready(write) = duplex_1.poll_write_from(cx, state.server_socket.as_mut())? {
+        if let Poll::Ready(write) = state
+            .server_duplex
+            .poll_write_from(cx, state.server_socket.as_mut())?
+        {
             if write > 0 {
                 progress = true;
-            } else if let Poll::Ready(()) = duplex_1.as_mut().poll_close(cx)? {
+            } else if !*state.server_closed {
                 progress = true;
+                *state.server_closed = true;
+                state.tls_client.server_close()?;
             }
         }
 
         // duplex -> tls_client
-        if let Poll::Ready(mut simplex) = duplex_2.as_mut().poll_lock_read(cx)
+        if state.tls_client.wants_read_tls()
+            && let Poll::Ready(mut simplex) = state.tls_client_duplex.as_mut().poll_lock_read(cx)
             && let Poll::Ready(buf) = simplex.poll_get(cx)?
+            && let read = state.tls_client.read_tls(buf)?
+            && read > 0
         {
-            if buf.len() > 0
-                && let read = state.tls_client.read_tls(buf)?
-                && read > 0
-            {
-                progress = true;
-                simplex.advance(read);
-            } else if !*state.server_closed {
-                progress = true;
-                *state.server_closed = true;
-                state.tls_client.server_close()?;
-            }
+            progress = true;
+            simplex.advance(read);
         }
 
         // tls_client -> duplex
-        if let Poll::Ready(mut simplex) = duplex_2.as_mut().poll_lock_write(cx)
+        if state.tls_client.wants_write_tls()
+            && let Poll::Ready(mut simplex) = state.tls_client_duplex.as_mut().poll_lock_write(cx)
             && let Poll::Ready(buf) = simplex.poll_mut(cx)?
+            && let write = state.tls_client.write_tls(buf)?
+            && write > 0
         {
-            if buf.len() > 0
-                && let write = state.tls_client.write_tls(buf)?
-                && write > 0
-            {
-                progress = true;
-                simplex.advance_mut(write);
-            }
+            progress = true;
+            simplex.advance_mut(write);
         }
 
         // duplex -> server_socket
-        if let Poll::Ready(read) = duplex_1.poll_read_to(cx, state.server_socket.as_mut())?
+        if let Poll::Ready(read) = state
+            .server_duplex
+            .poll_read_to(cx, state.server_socket.as_mut())?
             && read > 0
         {
             progress = true;

crates/tlsn/src/prover/client/mpc.rs

@@ -237,6 +237,7 @@ impl TlsClient for MpcTlsClient {
                 receiver,
             } => {
                 trace!("inner client is active");
+                //println!("client is handshaking: {}", inner.tls.is_handshaking());
 
                 if !inner.tls.is_handshaking()
                     && let Ok(cmd) = receiver.try_recv()
@@ -442,8 +443,7 @@ impl InnerState {
             }
             self.client_closed = true;
         }
-        self.tls.process_new_packets().await?;
-        Ok(self)
+        self.run().await
     }
 
     #[instrument(parent = &self.span, level = "debug", skip_all, err)]

crates/tlsn/src/prover/conn.rs

@@ -18,25 +18,18 @@ use std::{
 /// connection.
 pub struct TlsConnection {
     duplex: DuplexStream,
-    closed: bool,
 }
 
 impl TlsConnection {
     pub(crate) fn new(duplex: DuplexStream) -> Self {
-        Self {
-            duplex,
-            closed: false,
-        }
+        Self { duplex }
     }
 }
 
 impl Drop for TlsConnection {
     fn drop(&mut self) {
-        if !self.closed {
-            if let Err(err) = futures::executor::block_on(self.duplex.close()) {
-                tracing::error!("error closing connection: {}", err);
-            }
-            self.closed = true;
+        if let Err(err) = futures::executor::block_on(self.duplex.close()) {
+            tracing::error!("error closing connection: {}", err);
         }
     }
 }
@@ -47,12 +40,8 @@ impl AsyncRead for TlsConnection {
         cx: &mut Context<'_>,
         buf: &mut [u8],
     ) -> Poll<std::io::Result<usize>> {
-        if !self.closed {
-            let duplex = Pin::new(&mut self.duplex);
-            duplex.poll_read(cx, buf)
-        } else {
-            Poll::Ready(Ok(0))
-        }
+        let duplex = Pin::new(&mut self.duplex);
+        duplex.poll_read(cx, buf)
     }
 }
 
@@ -71,14 +60,7 @@ impl AsyncWrite for TlsConnection {
     }
 
     fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        if !self.closed {
-            let duplex = Pin::new(&mut self.duplex);
-            if let Poll::Ready(()) = duplex.poll_close(cx)? {
-                self.closed = true;
-            } else {
-                return Poll::Pending;
-            };
-        }
-        Poll::Ready(Ok(()))
+        let duplex = Pin::new(&mut self.duplex);
+        duplex.poll_close(cx)
     }
 }

crates/tlsn/src/prover/state.rs

@@ -66,6 +66,10 @@ pin_project_lite::pin_project! {
         pub(crate) server_socket: S,
         #[pin]
         pub(crate) verifier_socket: T,
+        #[pin]
+        pub(crate) tls_client_duplex: DuplexStream,
+        #[pin]
+        pub(crate) server_duplex: DuplexStream,
         pub(crate) client_closed: bool,
         pub(crate) server_closed: bool
     }

crates/tlsn/src/utils.rs

@@ -9,7 +9,7 @@ use futures::{AsyncRead, AsyncWrite, future::FusedFuture, ready};
 use futures_plex::DuplexStream;
 use mpz_common::context::Multithread;
 
-use crate::{BUF_CAP, mux::MuxControl};
+use crate::mux::MuxControl;
 
 /// Maximum concurrency for multi-threaded context.
 pub(crate) const MAX_CONCURRENCY: usize = 8;
@@ -53,7 +53,6 @@ pin_project_lite::pin_project! {
         io: S,
         #[pin]
         duplex: &'a mut DuplexStream,
-        buf: Vec<u8>,
         closed: bool
     }
 }
@@ -63,7 +62,6 @@ impl<'a, S> CopyFlush<'a, S> {
         Self {
             io,
             duplex,
-            buf: Vec::with_capacity(BUF_CAP),
             closed: false,
         }
     }
@@ -85,19 +83,23 @@ where
             if let Poll::Ready(read) = read
                 && read == 0
             {
+                println!("copy flush: read 0 from duplex: closing connection");
                 ready!(this.io.as_mut().poll_close(cx))?;
                 *this.closed = true;
             }
 
             if let Poll::Ready(write) = write
                 && write == 0
             {
+                println!("copy flush: read 0 from io: closing connection");
                 ready!(this.duplex.as_mut().poll_close(cx))?;
                 *this.closed = true;
             }
 
             if matches!(read, Poll::Pending) && matches!(write, Poll::Pending) {
                 return Poll::Pending;
+            } else if *this.closed {
+                return Poll::Ready(Ok(()));
             }
         }
     }

crates/tlsn/tests/test.rs

@@ -153,11 +153,11 @@ async fn prover<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
         .write_all(b"GET / HTTP/1.1\r\nConnection: close\r\n\r\n")
         .await
         .unwrap();
-    tls_connection.close().await.unwrap();
 
     let mut response = vec![0u8; 1024];
     tls_connection.read_to_end(&mut response).await.unwrap();
 
+    tls_connection.close().await.unwrap();
     let _ = server_task.await.unwrap();
 
     let (mut prover, _, mut verifier_socket) = prover_task.await.unwrap().unwrap();