wip

Cargo.toml

@@ -53,6 +53,7 @@ tlsn-formats = { path = "crates/formats" }
 tlsn-hmac-sha256 = { path = "crates/components/hmac-sha256" }
 tlsn-key-exchange = { path = "crates/components/key-exchange" }
 tlsn-mpc-tls = { path = "crates/mpc-tls" }
+tlsn-mux = { git = "https://github.com/tlsnotary/tlsn-utils", rev = "d9facb6" }
 tlsn-server-fixture = { path = "crates/server-fixture/server" }
 tlsn-server-fixture-certs = { path = "crates/server-fixture/certs" }
 tlsn-tls-backend = { path = "crates/tls/backend" }
@@ -66,26 +67,25 @@ tlsn-harness-runner = { path = "crates/harness/runner" }
 tlsn-wasm = { path = "crates/wasm" }
 tlsn = { path = "crates/tlsn" }
 
-mpz-circuits = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-circuits-data = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-memory-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-common = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-vm-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-garble = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-garble-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-ole = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-ot = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-share-conversion = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-fields = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-zk = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-hash = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
-mpz-ideal-vm = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-circuits = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-circuits-data = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-memory-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-common = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-vm-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-garble = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-garble-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-ole = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-ot = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-share-conversion = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-fields = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-zk = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-hash = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
+mpz-ideal-vm = { git = "https://github.com/privacy-ethereum/mpz", rev = "bc1d4ad" }
 
 rangeset = { version = "0.4" }
 serio = { version = "0.2" }
 spansy = { git = "https://github.com/tlsnotary/tlsn-utils", rev = "6f1a934" }
-uid-mux = { version = "0.2" }
 websocket-relay = { git = "https://github.com/tlsnotary/tlsn-utils", rev = "6f1a934" }
 
 aead = { version = "0.4" }

crates/mpc-tls/Cargo.toml

@@ -60,6 +60,8 @@ mpz-common = { workspace = true, features = ["test-utils"] }
 mpz-ot = { workspace = true, features = ["ideal"] }
 mpz-ideal-vm = { workspace = true }
 
+tlsn-mux = { workspace = true }
+
 cipher-crate = { package = "cipher", version = "0.4" }
 generic-array = { workspace = true }
 rand_chacha = { workspace = true }
@@ -70,6 +72,5 @@ tlsn-tls-client-async = { workspace = true }
 tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread"] }
 tokio-util = { workspace = true, features = ["compat"] }
 tracing-subscriber = { workspace = true }
-uid-mux = { workspace = true, features = ["serio", "test-utils"] }
 rustls-pki-types = { workspace = true }
 rustls-webpki = { workspace = true }

crates/tlsn/Cargo.toml

@@ -20,6 +20,7 @@ web = ["dep:web-spawn"]
 tlsn-attestation = { workspace = true }
 tlsn-core = { workspace = true }
 tlsn-deap = { workspace = true }
+tlsn-mux = { workspace = true }
 tlsn-tls-client = { workspace = true }
 tlsn-tls-client-async = { workspace = true }
 tlsn-tls-core = { workspace = true }
@@ -27,7 +28,6 @@ tlsn-mpc-tls = { workspace = true }
 tlsn-cipher = { workspace = true }
 
 serio = { workspace = true, features = ["compat"] }
-uid-mux = { workspace = true, features = ["serio"] }
 web-spawn = { workspace = true, optional = true }
 
 mpz-circuits = { workspace = true, features = ["aes"] }

crates/tlsn/src/mux.rs

@@ -1,45 +1,93 @@
 //! Multiplexer used in the TLSNotary protocol.
 
-use std::future::IntoFuture;
-
 use futures::{
     AsyncRead, AsyncWrite, Future,
     future::{FusedFuture, FutureExt},
 };
+use mpz_common::{ThreadId, io::Io, mux::Mux};
+use tlsn_mux::{Connection, Handle};
 use tracing::error;
-use uid_mux::yamux;
 
-use crate::Role;
-
-/// Multiplexer supporting unique deterministic stream IDs.
-pub(crate) type Mux<Io> = yamux::Yamux<Io>;
 /// Multiplexer controller providing streams.
-pub(crate) type MuxControl = yamux::YamuxCtrl;
+pub(crate) struct MuxControl {
+    handle: Handle,
+}
+
+impl Mux for MuxControl {
+    fn open(&self, id: ThreadId) -> Result<Io, std::io::Error> {
+        let stream = self
+            .handle
+            .new_stream(id.as_ref())
+            .map_err(std::io::Error::other)?;
+        let io = Io::from_io(stream);
+
+        Ok(io)
+    }
+}
+
+impl From<MuxControl> for Box<dyn Mux + Send> {
+    fn from(val: MuxControl) -> Self {
+        Box::new(val)
+    }
+}
 
 /// Multiplexer future which must be polled for the muxer to make progress.
-pub(crate) struct MuxFuture(
-    Box<dyn FusedFuture<Output = Result<(), yamux::ConnectionError>> + Send + Unpin>,
-);
+#[derive(Debug)]
+pub(crate) struct MuxFuture<T> {
+    conn: Connection<T>,
+}
 
-impl MuxFuture {
-    /// Returns true if the muxer is complete.
-    pub(crate) fn is_complete(&self) -> bool {
-        self.0.is_terminated()
+impl<T: AsyncRead + AsyncWrite + Unpin> MuxFuture<T> {
+    pub(crate) fn new(socket: T) -> Self {
+        let mut mux_config = tlsn_mux::Config::default();
+
+        mux_config.set_max_num_streams(36);
+        mux_config.set_keep_alive(true);
+        mux_config.set_close_sync(true);
+
+        let conn = tlsn_mux::Connection::new(socket, mux_config);
+
+        Self { conn }
     }
 
+    pub(crate) fn handle(&self) -> Result<MuxControl, std::io::Error> {
+        let handle = self.conn.handle().map_err(std::io::Error::other)?;
+
+        Ok(MuxControl { handle })
+    }
+
+    pub(crate) fn close(&mut self) {
+        self.conn.close();
+    }
+
+    pub(crate) fn into_io(self) -> Result<T, std::io::Error> {
+        self.conn
+            .try_into_io()
+            .map_err(|_| std::io::Error::other("unable to return IO, connection is not closed"))
+    }
+}
+
+impl<T: AsyncRead + AsyncWrite + Unpin> FusedFuture for MuxFuture<T> {
+    fn is_terminated(&self) -> bool {
+        self.conn.is_complete()
+    }
+}
+
+impl<T: AsyncRead + AsyncWrite + Unpin> MuxFuture<T> {
     /// Awaits a future, polling the muxer future concurrently.
     pub(crate) async fn poll_with<F, R>(&mut self, fut: F) -> R
     where
         F: Future<Output = R>,
     {
         let mut fut = Box::pin(fut.fuse());
+        let mut mux = self;
         // Poll the future concurrently with the muxer future.
         // If the muxer returns an error, continue polling the future
         // until it completes.
         loop {
             futures::select! {
                 res = fut => return res,
-                res = &mut self.0 => if let Err(e) = res {
+                res = mux => if let Err(e) = res {
                     error!("mux error: {:?}", e);
                 },
             }
@@ -47,44 +95,13 @@ impl MuxFuture {
     }
 }
 
-impl Future for MuxFuture {
-    type Output = Result<(), yamux::ConnectionError>;
+impl<T: AsyncRead + AsyncWrite + Unpin> Future for MuxFuture<T> {
+    type Output = Result<(), tlsn_mux::ConnectionError>;
 
     fn poll(
         mut self: std::pin::Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
     ) -> std::task::Poll<Self::Output> {
-        self.0.as_mut().poll_unpin(cx)
+        self.conn.poll(cx)
     }
 }
-
-/// Attaches a multiplexer to the provided socket.
-///
-/// Returns the multiplexer and a controller for creating streams with a codec
-/// attached.
-///
-/// # Arguments
-///
-/// * `socket` - The socket to attach the multiplexer to.
-/// * `role` - The role of the party using the multiplexer.
-pub(crate) fn attach_mux<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
-    socket: T,
-    role: Role,
-) -> (MuxFuture, MuxControl) {
-    let mut mux_config = yamux::Config::default();
-    mux_config.set_max_num_streams(36);
-
-    let mux_role = match role {
-        Role::Prover => yamux::Mode::Client,
-        Role::Verifier => yamux::Mode::Server,
-    };
-
-    let mux = Mux::new(socket, mux_config, mux_role);
-    let ctrl = mux.control();
-
-    if let Role::Prover = role {
-        ctrl.alloc(32);
-    }
-
-    (MuxFuture(Box::new(mux.into_future().fuse())), ctrl)
-}

crates/tlsn/src/prover.rs

@@ -10,11 +10,10 @@ pub use future::ProverFuture;
 pub use tlsn_core::ProverOutput;
 
 use crate::{
-    Role,
     context::build_mt_context,
     mpz::{ProverDeps, build_prover_deps, translate_keys},
     msg::{ProveRequestMsg, Response, TlsCommitRequestMsg},
-    mux::attach_mux,
+    mux::MuxFuture,
     tag::verify_tags,
 };
 
@@ -77,9 +76,11 @@ impl Prover<state::Initialized> {
         self,
         config: TlsCommitConfig,
         socket: S,
-    ) -> Result<Prover<state::CommitAccepted>, ProverError> {
-        let (mut mux_fut, mux_ctrl) = attach_mux(socket, Role::Prover);
-        let mut mt = build_mt_context(mux_ctrl.clone());
+    ) -> Result<Prover<state::CommitAccepted<S>>, ProverError> {
+        let mut mux_fut = MuxFuture::new(socket);
+        let mux_ctrl = mux_fut.handle()?;
+
+        let mut mt = build_mt_context(mux_ctrl);
         let mut ctx = mux_fut.poll_with(mt.new_context()).await?;
 
         // Sends protocol configuration to verifier for compatibility check.
@@ -122,7 +123,6 @@ impl Prover<state::Initialized> {
             config: self.config,
             span: self.span,
             state: state::CommitAccepted {
-                mux_ctrl,
                 mux_fut,
                 mpc_tls,
                 keys,
@@ -132,7 +132,10 @@ impl Prover<state::Initialized> {
     }
 }
 
-impl Prover<state::CommitAccepted> {
+impl<Io> Prover<state::CommitAccepted<Io>>
+where
+    Io: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+{
     /// Connects to the server using the provided socket.
     ///
     /// Returns a handle to the TLS connection, a future which returns the
@@ -148,9 +151,8 @@ impl Prover<state::CommitAccepted> {
         self,
         config: TlsClientConfig,
         socket: S,
-    ) -> Result<(TlsConnection, ProverFuture), ProverError> {
+    ) -> Result<(TlsConnection, ProverFuture<Io>), ProverError> {
         let state::CommitAccepted {
-            mux_ctrl,
             mut mux_fut,
             mpc_tls,
             keys,
@@ -270,7 +272,6 @@ impl Prover<state::CommitAccepted> {
                     config: self.config,
                     span: self.span,
                     state: state::Committed {
-                        mux_ctrl,
                         mux_fut,
                         ctx,
                         vm,
@@ -294,7 +295,10 @@ impl Prover<state::CommitAccepted> {
     }
 }
 
-impl Prover<state::Committed> {
+impl<Io> Prover<state::Committed<Io>>
+where
+    Io: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+{
     /// Returns the TLS transcript.
     pub fn tls_transcript(&self) -> &TlsTranscript {
         &self.state.tls_transcript
@@ -365,18 +369,12 @@ impl Prover<state::Committed> {
 
     /// Closes the connection with the verifier.
     #[instrument(parent = &self.span, level = "info", skip_all, err)]
-    pub async fn close(self) -> Result<(), ProverError> {
-        let state::Committed {
-            mux_ctrl, mux_fut, ..
-        } = self.state;
+    pub async fn close(mut self) -> Result<Io, ProverError> {
+        let mux_fut = &mut self.state.mux_fut;
+        mux_fut.close();
+        mux_fut.await?;
 
-        // Wait for the verifier to correctly close the connection.
-        if !mux_fut.is_complete() {
-            mux_ctrl.close();
-            mux_fut.await?;
-        }
-
-        Ok(())
+        self.state.mux_fut.into_io().map_err(ProverError::from)
     }
 }
 

crates/tlsn/src/prover/error.rs

@@ -90,8 +90,8 @@ impl From<tls_client_async::ConnectionError> for ProverError {
     }
 }
 
-impl From<uid_mux::yamux::ConnectionError> for ProverError {
-    fn from(e: uid_mux::yamux::ConnectionError) -> Self {
+impl From<tlsn_mux::ConnectionError> for ProverError {
+    fn from(e: tlsn_mux::ConnectionError) -> Self {
         Self::new(ErrorKind::Io, e)
     }
 }

crates/tlsn/src/prover/future.rs

@@ -5,23 +5,25 @@ use futures::Future;
 use std::pin::Pin;
 
 /// Prover future which must be polled for the TLS connection to make progress.
-pub struct ProverFuture {
+pub struct ProverFuture<Io> {
     #[allow(clippy::type_complexity)]
     pub(crate) fut: Pin<
-        Box<dyn Future<Output = Result<Prover<state::Committed>, ProverError>> + Send + 'static>,
+        Box<
+            dyn Future<Output = Result<Prover<state::Committed<Io>>, ProverError>> + Send + 'static,
+        >,
     >,
     pub(crate) ctrl: ProverControl,
 }
 
-impl ProverFuture {
+impl<Io> ProverFuture<Io> {
     /// Returns a controller for the prover for advanced functionality.
     pub fn control(&self) -> ProverControl {
         self.ctrl.clone()
     }
 }
 
-impl Future for ProverFuture {
-    type Output = Result<Prover<state::Committed>, ProverError>;
+impl<Io> Future for ProverFuture<Io> {
+    type Output = Result<Prover<state::Committed<Io>>, ProverError>;
 
     fn poll(
         mut self: Pin<&mut Self>,

crates/tlsn/src/prover/state.rs

@@ -13,7 +13,7 @@ use tokio::sync::Mutex;
 
 use crate::{
     mpz::{ProverMpc, ProverZk},
-    mux::{MuxControl, MuxFuture},
+    mux::MuxFuture,
 };
 
 /// Entry state
@@ -23,20 +23,18 @@ opaque_debug::implement!(Initialized);
 
 /// State after the verifier has accepted the proposed TLS commitment protocol
 /// configuration and preprocessing has completed.
-pub struct CommitAccepted {
-    pub(crate) mux_ctrl: MuxControl,
-    pub(crate) mux_fut: MuxFuture,
+pub struct CommitAccepted<Io> {
+    pub(crate) mux_fut: MuxFuture<Io>,
     pub(crate) mpc_tls: MpcTlsLeader,
     pub(crate) keys: SessionKeys,
     pub(crate) vm: Arc<Mutex<Deap<ProverMpc, ProverZk>>>,
 }
 
-opaque_debug::implement!(CommitAccepted);
+opaque_debug::implement!(CommitAccepted<Io>);
 
 /// State after the TLS transcript has been committed.
-pub struct Committed {
-    pub(crate) mux_ctrl: MuxControl,
-    pub(crate) mux_fut: MuxFuture,
+pub struct Committed<Io> {
+    pub(crate) mux_fut: MuxFuture<Io>,
     pub(crate) ctx: Context,
     pub(crate) vm: ProverZk,
     pub(crate) server_name: ServerName,
@@ -45,18 +43,18 @@ pub struct Committed {
     pub(crate) transcript: Transcript,
 }
 
-opaque_debug::implement!(Committed);
+opaque_debug::implement!(Committed<Io>);
 
 #[allow(missing_docs)]
 pub trait ProverState: sealed::Sealed {}
 
 impl ProverState for Initialized {}
-impl ProverState for CommitAccepted {}
-impl ProverState for Committed {}
+impl<Io> ProverState for CommitAccepted<Io> {}
+impl<Io> ProverState for Committed<Io> {}
 
 mod sealed {
     pub trait Sealed {}
     impl Sealed for super::Initialized {}
-    impl Sealed for super::CommitAccepted {}
-    impl Sealed for super::Committed {}
+    impl<Io> Sealed for super::CommitAccepted<Io> {}
+    impl<Io> Sealed for super::Committed<Io> {}
 }

crates/tlsn/src/verifier.rs

@@ -10,11 +10,10 @@ pub use error::VerifierError;
 pub use tlsn_core::{VerifierOutput, webpki::ServerCertVerifier};
 
 use crate::{
-    Role,
     context::build_mt_context,
     mpz::{VerifierDeps, build_verifier_deps, translate_keys},
     msg::{ProveRequestMsg, Response, TlsCommitRequestMsg},
-    mux::attach_mux,
+    mux::MuxFuture,
     tag::verify_tags,
 };
 use futures::{AsyncRead, AsyncWrite, TryFutureExt};
@@ -71,9 +70,11 @@ impl Verifier<state::Initialized> {
     pub async fn commit<S: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
         self,
         socket: S,
-    ) -> Result<Verifier<state::CommitStart>, VerifierError> {
-        let (mut mux_fut, mux_ctrl) = attach_mux(socket, Role::Verifier);
-        let mut mt = build_mt_context(mux_ctrl.clone());
+    ) -> Result<Verifier<state::CommitStart<S>>, VerifierError> {
+        let mut mux_fut = MuxFuture::new(socket);
+        let mux_ctrl = mux_fut.handle()?;
+
+        let mut mt = build_mt_context(mux_ctrl);
         let mut ctx = mux_fut.poll_with(mt.new_context()).await?;
 
         // Receives protocol configuration from prover to perform compatibility check.
@@ -89,11 +90,8 @@ impl Verifier<state::Initialized> {
                 .poll_with(ctx.io_mut().send(Response::err(Some(msg.clone()))))
                 .await?;
 
-            // Wait for the prover to correctly close the connection.
-            if !mux_fut.is_complete() {
-                mux_ctrl.close();
-                mux_fut.await?;
-            }
+            mux_fut.close();
+            mux_fut.await?;
 
             return Err(VerifierError::config(msg));
         }
@@ -102,7 +100,6 @@ impl Verifier<state::Initialized> {
             config: self.config,
             span: self.span,
             state: state::CommitStart {
-                mux_ctrl,
                 mux_fut,
                 ctx,
                 request,
@@ -111,7 +108,10 @@ impl Verifier<state::Initialized> {
     }
 }
 
-impl Verifier<state::CommitStart> {
+impl<Io> Verifier<state::CommitStart<Io>>
+where
+    Io: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+{
     /// Returns the TLS commitment request.
     pub fn request(&self) -> &TlsCommitRequest {
         &self.state.request
@@ -119,9 +119,8 @@ impl Verifier<state::CommitStart> {
 
     /// Accepts the proposed protocol configuration.
     #[instrument(parent = &self.span, level = "info", skip_all, err)]
-    pub async fn accept(self) -> Result<Verifier<state::CommitAccepted>, VerifierError> {
+    pub async fn accept(self) -> Result<Verifier<state::CommitAccepted<Io>>, VerifierError> {
         let state::CommitStart {
-            mux_ctrl,
             mut mux_fut,
             mut ctx,
             request,
@@ -151,7 +150,6 @@ impl Verifier<state::CommitStart> {
             config: self.config,
             span: self.span,
             state: state::CommitAccepted {
-                mux_ctrl,
                 mux_fut,
                 mpc_tls,
                 keys,
@@ -164,7 +162,6 @@ impl Verifier<state::CommitStart> {
     #[instrument(parent = &self.span, level = "info", skip_all, err)]
     pub async fn reject(self, msg: Option<&str>) -> Result<(), VerifierError> {
         let state::CommitStart {
-            mux_ctrl,
             mut mux_fut,
             mut ctx,
             ..
@@ -174,22 +171,21 @@ impl Verifier<state::CommitStart> {
             .poll_with(ctx.io_mut().send(Response::err(msg)))
             .await?;
 
-        // Wait for the prover to correctly close the connection.
-        if !mux_fut.is_complete() {
-            mux_ctrl.close();
-            mux_fut.await?;
-        }
+        mux_fut.close();
+        mux_fut.await?;
 
         Ok(())
     }
 }
 
-impl Verifier<state::CommitAccepted> {
+impl<Io> Verifier<state::CommitAccepted<Io>>
+where
+    Io: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+{
     /// Runs the verifier until the TLS connection is closed.
     #[instrument(parent = &self.span, level = "info", skip_all, err)]
-    pub async fn run(self) -> Result<Verifier<state::Committed>, VerifierError> {
+    pub async fn run(self) -> Result<Verifier<state::Committed<Io>>, VerifierError> {
         let state::CommitAccepted {
-            mux_ctrl,
             mut mux_fut,
             mpc_tls,
             vm,
@@ -245,7 +241,6 @@ impl Verifier<state::CommitAccepted> {
             config: self.config,
             span: self.span,
             state: state::Committed {
-                mux_ctrl,
                 mux_fut,
                 ctx,
                 vm,
@@ -256,7 +251,10 @@ impl Verifier<state::CommitAccepted> {
     }
 }
 
-impl Verifier<state::Committed> {
+impl<Io> Verifier<state::Committed<Io>>
+where
+    Io: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+{
     /// Returns the TLS transcript.
     pub fn tls_transcript(&self) -> &TlsTranscript {
         &self.state.tls_transcript
@@ -264,9 +262,8 @@ impl Verifier<state::Committed> {
 
     /// Begins verification of statements from the prover.
     #[instrument(parent = &self.span, level = "info", skip_all, err)]
-    pub async fn verify(self) -> Result<Verifier<state::Verify>, VerifierError> {
+    pub async fn verify(self) -> Result<Verifier<state::Verify<Io>>, VerifierError> {
         let state::Committed {
-            mux_ctrl,
             mut mux_fut,
             mut ctx,
             vm,
@@ -286,7 +283,6 @@ impl Verifier<state::Committed> {
             config: self.config,
             span: self.span,
             state: state::Verify {
-                mux_ctrl,
                 mux_fut,
                 ctx,
                 vm,
@@ -301,22 +297,19 @@ impl Verifier<state::Committed> {
 
     /// Closes the connection with the prover.
     #[instrument(parent = &self.span, level = "info", skip_all, err)]
-    pub async fn close(self) -> Result<(), VerifierError> {
-        let state::Committed {
-            mux_ctrl, mux_fut, ..
-        } = self.state;
+    pub async fn close(mut self) -> Result<Io, VerifierError> {
+        let mux_fut = &mut self.state.mux_fut;
+        mux_fut.close();
+        mux_fut.await?;
 
-        // Wait for the prover to correctly close the connection.
-        if !mux_fut.is_complete() {
-            mux_ctrl.close();
-            mux_fut.await?;
-        }
-
-        Ok(())
+        self.state.mux_fut.into_io().map_err(VerifierError::from)
     }
 }
 
-impl Verifier<state::Verify> {
+impl<Io> Verifier<state::Verify<Io>>
+where
+    Io: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+{
     /// Returns the proving request.
     pub fn request(&self) -> &ProveRequest {
         &self.state.request
@@ -325,9 +318,8 @@ impl Verifier<state::Verify> {
     /// Accepts the proving request.
     pub async fn accept(
         self,
-    ) -> Result<(VerifierOutput, Verifier<state::Committed>), VerifierError> {
+    ) -> Result<(VerifierOutput, Verifier<state::Committed<Io>>), VerifierError> {
         let state::Verify {
-            mux_ctrl,
             mut mux_fut,
             mut ctx,
             mut vm,
@@ -362,7 +354,6 @@ impl Verifier<state::Verify> {
                 config: self.config,
                 span: self.span,
                 state: state::Committed {
-                    mux_ctrl,
                     mux_fut,
                     ctx,
                     vm,
@@ -377,9 +368,8 @@ impl Verifier<state::Verify> {
     pub async fn reject(
         self,
         msg: Option<&str>,
-    ) -> Result<Verifier<state::Committed>, VerifierError> {
+    ) -> Result<Verifier<state::Committed<Io>>, VerifierError> {
         let state::Verify {
-            mux_ctrl,
             mut mux_fut,
             mut ctx,
             vm,
@@ -396,7 +386,6 @@ impl Verifier<state::Verify> {
             config: self.config,
             span: self.span,
             state: state::Committed {
-                mux_ctrl,
                 mux_fut,
                 ctx,
                 vm,

crates/tlsn/src/verifier/error.rs

@@ -84,8 +84,8 @@ impl From<std::io::Error> for VerifierError {
     }
 }
 
-impl From<uid_mux::yamux::ConnectionError> for VerifierError {
-    fn from(e: uid_mux::yamux::ConnectionError) -> Self {
+impl From<tlsn_mux::ConnectionError> for VerifierError {
+    fn from(e: tlsn_mux::ConnectionError) -> Self {
         Self::new(ErrorKind::Io, e)
     }
 }

crates/tlsn/src/verifier/state.rs

@@ -2,7 +2,7 @@
 
 use std::sync::Arc;
 
-use crate::mux::{MuxControl, MuxFuture};
+use crate::mux::MuxFuture;
 use mpc_tls::{MpcTlsFollower, SessionKeys};
 use mpz_common::Context;
 use tlsn_core::{
@@ -24,43 +24,39 @@ pub struct Initialized;
 opaque_debug::implement!(Initialized);
 
 /// State after receiving protocol configuration from the prover.
-pub struct CommitStart {
-    pub(crate) mux_ctrl: MuxControl,
-    pub(crate) mux_fut: MuxFuture,
+pub struct CommitStart<Io> {
+    pub(crate) mux_fut: MuxFuture<Io>,
     pub(crate) ctx: Context,
     pub(crate) request: TlsCommitRequest,
 }
 
-opaque_debug::implement!(CommitStart);
+opaque_debug::implement!(CommitStart<Io>);
 
 /// State after accepting the proposed TLS commitment protocol configuration and
 /// performing preprocessing.
-pub struct CommitAccepted {
-    pub(crate) mux_ctrl: MuxControl,
-    pub(crate) mux_fut: MuxFuture,
+pub struct CommitAccepted<Io> {
+    pub(crate) mux_fut: MuxFuture<Io>,
     pub(crate) mpc_tls: MpcTlsFollower,
     pub(crate) keys: SessionKeys,
     pub(crate) vm: Arc<Mutex<Deap<VerifierMpc, VerifierZk>>>,
 }
 
-opaque_debug::implement!(CommitAccepted);
+opaque_debug::implement!(CommitAccepted<Io>);
 
 /// State after the TLS transcript has been committed.
-pub struct Committed {
-    pub(crate) mux_ctrl: MuxControl,
-    pub(crate) mux_fut: MuxFuture,
+pub struct Committed<Io> {
+    pub(crate) mux_fut: MuxFuture<Io>,
     pub(crate) ctx: Context,
     pub(crate) vm: VerifierZk,
     pub(crate) keys: SessionKeys,
     pub(crate) tls_transcript: TlsTranscript,
 }
 
-opaque_debug::implement!(Committed);
+opaque_debug::implement!(Committed<Io>);
 
 /// State after receiving a proving request.
-pub struct Verify {
-    pub(crate) mux_ctrl: MuxControl,
-    pub(crate) mux_fut: MuxFuture,
+pub struct Verify<Io> {
+    pub(crate) mux_fut: MuxFuture<Io>,
     pub(crate) ctx: Context,
     pub(crate) vm: VerifierZk,
     pub(crate) keys: SessionKeys,
@@ -70,19 +66,19 @@ pub struct Verify {
     pub(crate) transcript: Option<PartialTranscript>,
 }
 
-opaque_debug::implement!(Verify);
+opaque_debug::implement!(Verify<Io>);
 
 impl VerifierState for Initialized {}
-impl VerifierState for CommitStart {}
-impl VerifierState for CommitAccepted {}
-impl VerifierState for Committed {}
-impl VerifierState for Verify {}
+impl<Io> VerifierState for CommitStart<Io> {}
+impl<Io> VerifierState for CommitAccepted<Io> {}
+impl<Io> VerifierState for Committed<Io> {}
+impl<Io> VerifierState for Verify<Io> {}
 
 mod sealed {
     pub trait Sealed {}
     impl Sealed for super::Initialized {}
-    impl Sealed for super::CommitStart {}
-    impl Sealed for super::CommitAccepted {}
-    impl Sealed for super::Committed {}
-    impl Sealed for super::Verify {}
+    impl<Io> Sealed for super::CommitStart<Io> {}
+    impl<Io> Sealed for super::CommitAccepted<Io> {}
+    impl<Io> Sealed for super::Committed<Io> {}
+    impl<Io> Sealed for super::Verify<Io> {}
 }