diff --git a/crates/tlsn/src/prover.rs b/crates/tlsn/src/prover.rs index 0642698eb..221a6838b 100644 --- a/crates/tlsn/src/prover.rs +++ b/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 { /// /// * `config` - The TLS commitment configuration. #[instrument(parent = &self.span, level = "debug", skip_all, err)] - pub async fn commit( - self, - config: TlsCommitConfig, - ) -> Result<(MpcConnection, Prover), ProverError> { + pub async fn commit(self, config: TlsCommitConfig) -> Result { 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( + self, + config: TlsCommitConfig, + socket: S, + ) -> Result, ProverError> { + self.commit_inner(config, socket).await + } + + async fn commit_inner( + self, + config: TlsCommitConfig, + transport: S, + ) -> Result, 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 { 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 { }, }; - 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( - self, - config: TlsCommitConfig, - socket: S, - ) -> Result, ProverError> { - todo!() + Ok(prover) } } diff --git a/crates/tlsn/src/prover/conn/mpc.rs b/crates/tlsn/src/prover/conn/mpc.rs index d93b65b53..43f255470 100644 --- a/crates/tlsn/src/prover/conn/mpc.rs +++ b/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, 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, + setup: Pin, + prover: Option>, } -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)> { + 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 { + self.duplex + .as_mut() + .expect("duplex should be available") + .read(buf) + } +} + +impl Write for MpcSetup { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + 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 { + 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 { + self.duplex.read(buf) + } +} + +impl Write for MpcConnection { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.duplex.write(buf) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) } }