refactor: tlsn-common mux module

tlsn/Cargo.toml

@@ -1,9 +1,10 @@
 [workspace]
 members = [
     "tlsn-core",
+    "tlsn-common",
     "tlsn-verifier",
     "tlsn-prover",
-#    "tlsn-formats",
+    #    "tlsn-formats",
     "tlsn-server-fixture",
     "tests-integration",
     "examples",
@@ -12,6 +13,7 @@ resolver = "2"
 
 [workspace.dependencies]
 tlsn-core = { path = "tlsn-core" }
+tlsn-common = { path = "tlsn-common" }
 tlsn-prover = { path = "tlsn-prover" }
 tlsn-verifier = { path = "tlsn-verifier" }
 tlsn-server-fixture = { path = "tlsn-server-fixture" }

tlsn/tlsn-common/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "tlsn-common"
+description = "Common code shared between tlsn-prover and tlsn-verifier"
+version = "0.1.0-alpha.3"
+edition = "2021"
+
+[features]
+default = ["tracing"]
+tracing = ["uid-mux/tracing"]
+
+[dependencies]
+tlsn-utils-aio.workspace = true
+
+futures.workspace = true
+uid-mux.workspace = true

tlsn/tlsn-common/src/lib.rs

@@ -0,0 +1,17 @@
+//! Common code shared between `tlsn-prover` and `tlsn-verifier`.
+
+#![deny(missing_docs, unreachable_pub, unused_must_use)]
+#![deny(clippy::all)]
+#![forbid(unsafe_code)]
+
+pub mod mux;
+
+/// The parties role in the TLSN protocol.
+///
+/// A Notary is classified as a Verifier.
+pub enum Role {
+    /// The prover.
+    Prover,
+    /// The verifier.
+    Verifier,
+}

tlsn/tlsn-common/src/mux.rs

@@ -0,0 +1,45 @@
+//! Multiplexer used in the TLSNotary protocol.
+
+use utils_aio::codec::BincodeMux;
+
+use futures::{AsyncRead, AsyncWrite};
+use uid_mux::{yamux, UidYamux, UidYamuxControl};
+
+use crate::Role;
+
+/// Multiplexer supporting unique deterministic stream IDs.
+pub type Mux<T> = UidYamux<T>;
+/// Multiplexer controller providing streams with a codec attached.
+pub type MuxControl = BincodeMux<UidYamuxControl>;
+
+const KB: usize = 1024;
+const MB: usize = 1024 * KB;
+
+/// Attach 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 fn attach_mux<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
+    socket: T,
+    role: Role,
+) -> (Mux<T>, MuxControl) {
+    let mut mux_config = yamux::Config::default();
+    // See PR #418
+    mux_config.set_max_num_streams(40);
+    mux_config.set_max_buffer_size(16 * MB);
+    mux_config.set_receive_window(16 * MB as u32);
+
+    let mux_role = match role {
+        Role::Prover => yamux::Mode::Client,
+        Role::Verifier => yamux::Mode::Server,
+    };
+
+    let mux = UidYamux::new(mux_config, socket, mux_role);
+    let ctrl = BincodeMux::new(mux.control());
+
+    (mux, ctrl)
+}

tlsn/tlsn-prover/Cargo.toml

@@ -15,17 +15,17 @@ tracing = [
     "dep:tracing",
     "tlsn-tls-client-async/tracing",
     "tlsn-tls-mpc/tracing",
-    "uid-mux/tracing",
+    "tlsn-common/tracing",
 ]
 
 [dependencies]
 tlsn-tls-core.workspace = true
 tlsn-tls-client.workspace = true
 tlsn-tls-client-async.workspace = true
 tlsn-core.workspace = true
+tlsn-common.workspace = true
 #tlsn-formats = { workspace = true, optional = true }
 tlsn-tls-mpc.workspace = true
-uid-mux.workspace = true
 
 tlsn-utils.workspace = true
 tlsn-utils-aio.workspace = true

tlsn/tlsn-prover/src/lib.rs

@@ -11,9 +11,3 @@
 #[cfg(feature = "formats")]
 pub mod http;
 pub mod tls;
-
-use uid_mux::UidYamuxControl;
-use utils_aio::codec::BincodeMux;
-
-/// A muxer which uses Bincode for serialization, Yamux for multiplexing.
-type Mux = BincodeMux<UidYamuxControl>;

tlsn/tlsn-prover/src/tls/mod.rs

@@ -15,8 +15,11 @@ pub mod state;
 pub use config::{ProverConfig, ProverConfigBuilder, ProverConfigBuilderError};
 pub use error::ProverError;
 pub use future::ProverFuture;
+use tlsn_common::{
+    mux::{attach_mux, MuxControl},
+    Role,
+};
 
-use crate::Mux;
 use error::OTShutdownError;
 use future::{MuxFuture, OTFuture};
 use futures::{AsyncRead, AsyncWrite, FutureExt, StreamExt, TryFutureExt};
@@ -33,8 +36,7 @@ use tls_client::{ClientConnection, ServerName as TlsServerName};
 use tls_client_async::{bind_client, ClosedConnection, TlsConnection};
 use tls_mpc::{setup_components, MpcTlsLeader, TlsRole};
 use tlsn_core::transcript::Transcript;
-use uid_mux::{yamux, UidYamux};
-use utils_aio::{codec::BincodeMux, mux::MuxChannel};
+use utils_aio::mux::MuxChannel;
 
 #[cfg(feature = "formats")]
 use http::{state as http_state, HttpProver, HttpProverError};
@@ -74,19 +76,13 @@ impl Prover<state::Initialized> {
         self,
         socket: S,
     ) -> Result<Prover<state::Setup>, ProverError> {
-        let mut mux_config = yamux::Config::default();
-        // See PR #418
-        mux_config.set_max_num_streams(40);
-        mux_config.set_max_buffer_size(16 * 1024 * 1024);
-        mux_config.set_receive_window(16 * 1024 * 1024);
-        let mut mux = UidYamux::new(mux_config, socket, yamux::Mode::Client);
-        let notary_mux = BincodeMux::new(mux.control());
+        let (mut mux, mux_ctrl) = attach_mux(socket, Role::Prover);
 
         let mut mux_fut = MuxFuture {
             fut: Box::pin(async move { mux.run().await.map_err(ProverError::from) }.fuse()),
         };
 
-        let mpc_setup_fut = setup_mpc_backend(&self.config, notary_mux.clone());
+        let mpc_setup_fut = setup_mpc_backend(&self.config, mux_ctrl.clone());
         let (mpc_tls, vm, _, gf2, ot_fut) = futures::select! {
             res = mpc_setup_fut.fuse() => res?,
             _ = (&mut mux_fut).fuse() => return Err(std::io::Error::from(std::io::ErrorKind::UnexpectedEof))?,
@@ -95,7 +91,7 @@ impl Prover<state::Initialized> {
         Ok(Prover {
             config: self.config,
             state: state::Setup {
-                notary_mux,
+                mux_ctrl,
                 mux_fut,
                 mpc_tls,
                 vm,
@@ -124,7 +120,7 @@ impl Prover<state::Setup> {
         socket: S,
     ) -> Result<(TlsConnection, ProverFuture), ProverError> {
         let state::Setup {
-            notary_mux,
+            mux_ctrl,
             mut mux_fut,
             mpc_tls,
             vm,
@@ -168,7 +164,7 @@ impl Prover<state::Setup> {
                 Ok(Prover {
                     config: self.config,
                     state: state::Closed {
-                        notary_mux,
+                        mux_ctrl,
                         mux_fut,
                         vm,
                         ot_fut,
@@ -237,7 +233,7 @@ impl Prover<state::Closed> {
 #[allow(clippy::type_complexity)]
 async fn setup_mpc_backend(
     config: &ProverConfig,
-    mut mux: Mux,
+    mut mux: MuxControl,
 ) -> Result<
     (
         MpcTlsLeader,

tlsn/tlsn-prover/src/tls/notarize.rs

@@ -36,7 +36,7 @@ impl Prover<Notarize> {
     #[cfg_attr(feature = "tracing", instrument(level = "info", skip(self), err))]
     pub async fn finalize(self) -> Result<NotarizedSession, ProverError> {
         let Notarize {
-            mut notary_mux,
+            mut mux_ctrl,
             mut mux_fut,
             mut vm,
             mut ot_fut,
@@ -62,7 +62,7 @@ impl Prover<Notarize> {
         let merkle_root = session_data.commitments().merkle_root();
 
         let mut notarize_fut = Box::pin(async move {
-            let mut channel = notary_mux.get_channel("notarize").await?;
+            let mut channel = mux_ctrl.get_channel("notarize").await?;
 
             channel
                 .send(TlsnMessage::TranscriptCommitmentRoot(merkle_root))

tlsn/tlsn-prover/src/tls/prove.rs

@@ -74,7 +74,7 @@ impl Prover<ProveState> {
             let channel = if let Some(ref mut channel) = self.state.channel {
                 channel
             } else {
-                self.state.channel = Some(self.state.verify_mux.get_channel("prove-verify").await?);
+                self.state.channel = Some(self.state.mux_ctrl.get_channel("prove-verify").await?);
                 self.state.channel.as_mut().unwrap()
             };
 
@@ -154,7 +154,7 @@ impl Prover<ProveState> {
     /// Finalize the proving
     pub async fn finalize(self) -> Result<(), ProverError> {
         let ProveState {
-            mut verify_mux,
+            mut mux_ctrl,
             mut mux_fut,
             mut vm,
             mut ot_fut,
@@ -170,7 +170,7 @@ impl Prover<ProveState> {
         };
 
         let mut finalize_fut = Box::pin(async move {
-            let mut channel = verify_mux.get_channel("finalize").await?;
+            let mut channel = mux_ctrl.get_channel("finalize").await?;
 
             _ = vm
                 .finalize()

tlsn/tlsn-prover/src/tls/state.rs

@@ -1,9 +1,6 @@
 //! TLS prover states.
 
-use crate::{
-    tls::{MuxFuture, OTFuture},
-    Mux,
-};
+use crate::tls::{MuxFuture, OTFuture};
 use mpz_core::commit::Decommitment;
 use mpz_garble::protocol::deap::{DEAPThread, DEAPVm, PeerEncodings};
 use mpz_garble_core::{encoding_state, EncodedValue};
@@ -12,6 +9,7 @@ use mpz_share_conversion::{ConverterSender, Gf2_128};
 use std::collections::HashMap;
 use tls_core::{handshake::HandshakeData, key::PublicKey};
 use tls_mpc::MpcTlsLeader;
+use tlsn_common::mux::MuxControl;
 use tlsn_core::{
     commitment::TranscriptCommitmentBuilder,
     msg::{ProvingInfo, TlsnMessage},
@@ -27,7 +25,7 @@ opaque_debug::implement!(Initialized);
 /// State after MPC setup has completed.
 pub struct Setup {
     /// A muxer for communication with the Notary
-    pub(crate) notary_mux: Mux,
+    pub(crate) mux_ctrl: MuxControl,
     pub(crate) mux_fut: MuxFuture,
 
     pub(crate) mpc_tls: MpcTlsLeader,
@@ -40,7 +38,7 @@ opaque_debug::implement!(Setup);
 
 /// State after the TLS connection has been closed.
 pub struct Closed {
-    pub(crate) notary_mux: Mux,
+    pub(crate) mux_ctrl: MuxControl,
     pub(crate) mux_fut: MuxFuture,
 
     pub(crate) vm: DEAPVm<SharedSender, SharedReceiver>,
@@ -60,7 +58,7 @@ opaque_debug::implement!(Closed);
 /// Notarizing state.
 pub struct Notarize {
     /// A muxer for communication with the Notary
-    pub(crate) notary_mux: Mux,
+    pub(crate) mux_ctrl: MuxControl,
     pub(crate) mux_fut: MuxFuture,
 
     pub(crate) vm: DEAPVm<SharedSender, SharedReceiver>,
@@ -94,7 +92,7 @@ impl From<Closed> for Notarize {
         );
 
         Self {
-            notary_mux: state.notary_mux,
+            mux_ctrl: state.mux_ctrl,
             mux_fut: state.mux_fut,
             vm: state.vm,
             ot_fut: state.ot_fut,
@@ -111,7 +109,7 @@ impl From<Closed> for Notarize {
 
 /// Proving state.
 pub struct Prove {
-    pub(crate) verify_mux: Mux,
+    pub(crate) mux_ctrl: MuxControl,
     pub(crate) mux_fut: MuxFuture,
 
     pub(crate) vm: DEAPVm<SharedSender, SharedReceiver>,
@@ -131,7 +129,7 @@ pub struct Prove {
 impl From<Closed> for Prove {
     fn from(state: Closed) -> Self {
         Self {
-            verify_mux: state.notary_mux,
+            mux_ctrl: state.mux_ctrl,
             mux_fut: state.mux_fut,
             vm: state.vm,
             ot_fut: state.ot_fut,

tlsn/tlsn-verifier/Cargo.toml

@@ -9,10 +9,11 @@ version = "0.1.0-alpha.3"
 edition = "2021"
 
 [features]
-tracing = ["dep:tracing", "tlsn-tls-mpc/tracing"]
+tracing = ["dep:tracing", "tlsn-tls-mpc/tracing", "tlsn-common/tracing"]
 
 [dependencies]
 tlsn-core.workspace = true
+tlsn-common.workspace = true
 tlsn-tls-core.workspace = true
 tlsn-tls-mpc.workspace = true
 uid-mux.workspace = true

tlsn/tlsn-verifier/src/lib.rs

@@ -5,9 +5,3 @@
 #![forbid(unsafe_code)]
 
 pub mod tls;
-
-use uid_mux::UidYamuxControl;
-use utils_aio::codec::BincodeMux;
-
-/// A muxer which uses Bincode for serialization, Yamux for multiplexing.
-pub(crate) type Mux = BincodeMux<UidYamuxControl>;