New socket api

crates/examples/attestation/prove.rs

@@ -124,7 +124,7 @@ async fn prover<S: AsyncWrite + AsyncRead + Send + Sync + Unpin + 'static>(
 
     // Bind the prover to the server connection.
     let (tls_connection, prover_fut) = prover
-        .connect(
+        .connect_with(
             TlsClientConfig::builder()
                 .server_name(ServerName::Dns(SERVER_DOMAIN.try_into()?))
                 // Create a root certificate store with the server-fixture's self-signed

crates/examples/interactive/interactive.rs

@@ -102,7 +102,7 @@ async fn prover<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
 
     // Bind the prover to the server connection.
     let (tls_connection, prover_fut) = prover
-        .connect(
+        .connect_with(
             TlsClientConfig::builder()
                 .server_name(ServerName::Dns(SERVER_DOMAIN.try_into()?))
                 // Create a root certificate store with the server-fixture's self-signed

crates/examples/interactive_zk/prover.rs

@@ -89,7 +89,7 @@ pub async fn prover<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
 
     // Bind the prover to the server connection.
     let (tls_connection, prover_fut) = prover
-        .connect(
+        .connect_with(
             TlsClientConfig::builder()
                 .server_name(ServerName::Dns(SERVER_DOMAIN.try_into()?))
                 // Create a root certificate store with the server-fixture's self-signed

crates/harness/executor/src/bench/prover.rs

@@ -59,7 +59,7 @@ pub async fn bench_prover(provider: &IoProvider, config: &Bench) -> Result<Prove
     let downloaded_preprocess = recv.load(Ordering::Relaxed);
 
     let (mut conn, prover_fut) = prover
-        .connect(
+        .connect_with(
             TlsClientConfig::builder()
                 .server_name(ServerName::Dns(SERVER_DOMAIN.try_into()?))
                 .root_store(RootCertStore {

crates/harness/executor/test_plugins/basic.rs

@@ -47,7 +47,7 @@ async fn prover(provider: &IoProvider) {
         .unwrap();
 
     let (tls_connection, prover_fut) = prover
-        .connect(
+        .connect_with(
             TlsClientConfig::builder()
                 .server_name(ServerName::Dns(SERVER_DOMAIN.try_into().unwrap()))
                 .root_store(RootCertStore {

crates/tlsn/Cargo.toml

@@ -44,6 +44,7 @@ mpz-zk = { workspace = true }
 mpz-ideal-vm = { workspace = true }
 
 aes = { workspace = true }
+bytes = { workspace = true }
 ctr = { workspace = true }
 futures = { workspace = true }
 opaque-debug = { workspace = true }
@@ -57,6 +58,7 @@ serde = { workspace = true, features = ["derive"] }
 ghash = { workspace = true }
 semver = { workspace = true, features = ["serde"] }
 once_cell = { workspace = true }
+pin-project-lite = { workspace = true }
 rangeset = { workspace = true }
 webpki-roots = { workspace = true }
 

crates/tlsn/src/prover.rs

@@ -1,10 +1,14 @@
 //! Prover.
 
 mod client;
+mod conn;
+mod control;
 mod error;
 mod prove;
 pub mod state;
 
+pub use conn::{ConnectionFuture, TlsConnection};
+pub use control::ProverControl;
 pub use error::ProverError;
 pub use tlsn_core::ProverOutput;
 
@@ -21,7 +25,7 @@ use futures::{AsyncRead, AsyncWrite, FutureExt, TryFutureExt};
 use rustls_pki_types::CertificateDer;
 use serio::{SinkExt, stream::IoStreamExt};
 use std::{
-    sync::Arc,
+    sync::{Arc, Mutex},
     task::{Context, Poll},
 };
 use tls_client::{ClientConnection, ServerName as TlsServerName};
@@ -223,6 +227,42 @@ impl Prover<state::CommitAccepted> {
         };
         Ok(prover)
     }
+
+    /// Connects the prover and attaches a socket.
+    ///
+    /// This is a convenience function which returns
+    ///   - [`TlsConnection`] for reading and writing traffic as well as other
+    ///     connection-specific settings.
+    ///   - [`ConnectionFuture`] which has to be polled for driving the
+    ///     connection forward.
+    ///
+    /// # Arguments
+    ///
+    /// * `config` - The TLS client configuration.
+    #[instrument(parent = &self.span, level = "debug", skip_all, err)]
+    pub async fn connect_with<S>(
+        self,
+        config: TlsClientConfig,
+        socket: S,
+    ) -> Result<(TlsConnection, ConnectionFuture<S>), ProverError>
+    where
+        S: AsyncRead + AsyncWrite + Send,
+    {
+        let prover = self.connect(config).await?;
+
+        let prover = Arc::new(Mutex::new(prover));
+        let conn_waker = Arc::new(Mutex::new(None));
+        let fut_waker = Arc::new(Mutex::new(None));
+
+        let conn = TlsConnection::new(
+            Arc::downgrade(&prover),
+            conn_waker.clone(),
+            fut_waker.clone(),
+        );
+        let fut = ConnectionFuture::new(socket, prover, conn_waker, fut_waker);
+
+        Ok((conn, fut))
+    }
 }
 
 impl Prover<state::Connected> {
@@ -317,6 +357,7 @@ impl Prover<state::Connected> {
 
         match self.state.tls_client.poll(cx)? {
             Poll::Ready(output) => {
+                let _ = self.state.mux_fut.poll_unpin(cx)?;
                 self.state.output = Some(output);
                 Poll::Ready(Ok(()))
             }