added MpcSetup and MpcConnection

crates/tlsn/src/prover.rs

@@ -7,7 +7,7 @@ mod error;
 mod prove;
 pub mod state;
 
-pub use conn::{ConnectionFuture, TlsConnection, mpc::MpcConnection};
+pub use conn::{ConnectionFuture, TlsConnection, mpc::MpcSetup};
 pub use control::ProverControl;
 pub use error::ProverError;
 pub use tlsn_core::ProverOutput;
@@ -76,14 +76,40 @@ impl Prover<state::Initialized> {
     ///
     /// * `config` - The TLS commitment configuration.
     #[instrument(parent = &self.span, level = "debug", skip_all, err)]
-    pub async fn commit(
-        self,
-        config: TlsCommitConfig,
-    ) -> Result<(MpcConnection, Prover<state::CommitAccepted>), ProverError> {
+    pub async fn commit(self, config: TlsCommitConfig) -> Result<MpcSetup, ProverError> {
         let (duplex_a, duplex_b) = futures_plex::duplex(BUF_CAP);
 
-        let (mut mux_fut, mux_ctrl) = attach_mux(duplex_a, Role::Prover);
+        let setup = self.commit_inner(config, duplex_b);
+        let mpc_conn = MpcSetup::new(duplex_a, Box::new(setup));
+
+        Ok(mpc_conn)
+    }
+
+    /// Starts the TLS commitment protocol and attaches a socket.
+    ///
+    /// This is a convenience function...
+    ///
+    /// # Arguments
+    ///
+    /// * `config` - The TLS commitment configuration.
+    /// * `socket` - The socket to the TLS verifier.
+    #[instrument(parent = &self.span, level = "debug", skip_all, err)]
+    pub async fn commit_with<S: AsyncWrite + AsyncRead + Send>(
+        self,
+        config: TlsCommitConfig,
+        socket: S,
+    ) -> Result<Prover<state::CommitAccepted>, ProverError> {
+        self.commit_inner(config, socket).await
+    }
+
+    async fn commit_inner<S: AsyncWrite + AsyncRead + Send>(
+        self,
+        config: TlsCommitConfig,
+        transport: S,
+    ) -> Result<Prover<state::CommitAccepted>, ProverError> {
+        let (mut mux_fut, mux_ctrl) = attach_mux(transport, Role::Prover);
         let mut mt = build_mt_context(mux_ctrl.clone());
+
         let mut ctx = mux_fut.poll_with(mt.new_context()).await?;
 
         // Sends protocol configuration to verifier for compatibility check.
@@ -122,7 +148,6 @@ impl Prover<state::Initialized> {
 
         debug!("mpc-tls setup complete");
 
-        let mpc_conn = MpcConnection::new(duplex_b);
         let prover = Prover {
             config: self.config,
             span: self.span,
@@ -135,24 +160,7 @@ impl Prover<state::Initialized> {
             },
         };
 
-        Ok((mpc_conn, prover))
-    }
-
-    /// Starts the TLS commitment protocol and attaches a socket.
-    ///
-    /// This is a convenience function...
-    ///
-    /// # Arguments
-    ///
-    /// * `config` - The TLS commitment configuration.
-    /// * `socket` - The socket to the TLS verifier.
-    #[instrument(parent = &self.span, level = "debug", skip_all, err)]
-    pub async fn commit_with<S: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
-        self,
-        config: TlsCommitConfig,
-        socket: S,
-    ) -> Result<Prover<state::CommitAccepted>, ProverError> {
-        todo!()
+        Ok(prover)
     }
 }
 

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

@@ -1,12 +1,105 @@
 use futures_plex::DuplexStream;
+use std::{
+    io::{Read, Write},
+    pin::Pin,
+    task::Poll,
+};
 
-/// A connection to the verifier.
-pub struct MpcConnection {
-    duplex: DuplexStream,
+use crate::prover::{Prover, ProverError, state};
+
+pub(crate) type MpcSetupFuture =
+    Box<dyn Future<Output = Result<Prover<state::CommitAccepted>, ProverError>> + Send>;
+
+/// MPC setup for preparing a connection to the verifier.
+///
+/// Implements [`Read`] and [`Write`] for doing IO with the verifier.
+pub struct MpcSetup {
+    duplex: Option<DuplexStream>,
+    setup: Pin<MpcSetupFuture>,
+    prover: Option<Prover<state::CommitAccepted>>,
 }
 
-impl MpcConnection {
-    pub(crate) fn new(duplex: DuplexStream) -> Self {
-        Self { duplex }
+impl MpcSetup {
+    pub(crate) fn new(duplex: DuplexStream, setup: MpcSetupFuture) -> Self {
+        Self {
+            duplex: Some(duplex),
+            setup: Box::into_pin(setup),
+            prover: None,
+        }
+    }
+
+    /// Finishes the setup and returns a connection to the verifier and the prover.
+    pub fn finish(&mut self) -> Option<(MpcConnection, Prover<state::CommitAccepted>)> {
+        match self.prover.take() {
+            Some(prover) => Some((
+                MpcConnection {
+                    duplex: self.duplex.take().expect("duplex should be available"),
+                },
+                prover,
+            )),
+            None => None,
+        }
+    }
+}
+
+impl Read for MpcSetup {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        self.duplex
+            .as_mut()
+            .expect("duplex should be available")
+            .read(buf)
+    }
+}
+
+impl Write for MpcSetup {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.duplex
+            .as_mut()
+            .expect("duplex should be available")
+            .write(buf)
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(())
+    }
+}
+
+impl Future for MpcSetup {
+    type Output = Result<(), ProverError>;
+
+    fn poll(
+        mut self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Self::Output> {
+        let setup = Pin::new(&mut self.setup);
+
+        if let Poll::Ready(prover) = setup.poll(cx)? {
+            self.prover = Some(prover);
+            return Poll::Ready(Ok(()));
+        }
+        Poll::Pending
+    }
+}
+
+/// MPC Connection to the verifier.
+///
+/// Implements [`Read`] and [`Write`] for doing IO with the verifier.
+pub struct MpcConnection {
+    pub(crate) duplex: DuplexStream,
+}
+
+impl Read for MpcConnection {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        self.duplex.read(buf)
+    }
+}
+
+impl Write for MpcConnection {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.duplex.write(buf)
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(())
     }
 }