wip

Cargo.lock

@@ -2026,7 +2026,7 @@ checksum = "a1d728cc89cf3aee9ff92b05e62b19ee65a02b5702cff7d5a377e32c6ae29d8d"
 [[package]]
 name = "clmul"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "bytemuck",
  "cfg-if",
@@ -4257,7 +4257,7 @@ checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
 [[package]]
 name = "matrix-transpose"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "thiserror 1.0.69",
 ]
@@ -4314,7 +4314,7 @@ dependencies = [
 [[package]]
 name = "mpz-circuits"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "mpz-circuits-core",
  "mpz-circuits-data",
@@ -4323,7 +4323,7 @@ dependencies = [
 [[package]]
 name = "mpz-circuits-core"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "bincode 1.3.3",
  "itybity 0.3.1",
@@ -4338,7 +4338,7 @@ dependencies = [
 [[package]]
 name = "mpz-circuits-data"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "bincode 1.3.3",
  "mpz-circuits-core",
@@ -4348,7 +4348,7 @@ dependencies = [
 [[package]]
 name = "mpz-cointoss"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "futures",
  "mpz-cointoss-core",
@@ -4361,7 +4361,7 @@ dependencies = [
 [[package]]
 name = "mpz-cointoss-core"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "mpz-core",
  "opaque-debug",
@@ -4372,7 +4372,7 @@ dependencies = [
 [[package]]
 name = "mpz-common"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "async-trait",
  "bytes",
@@ -4386,13 +4386,12 @@ dependencies = [
  "thiserror 1.0.69",
  "tokio",
  "tokio-util",
- "uid-mux",
 ]
 
 [[package]]
 name = "mpz-core"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "aes 0.9.0-rc.2",
  "bcs",
@@ -4418,7 +4417,7 @@ dependencies = [
 [[package]]
 name = "mpz-fields"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "ark-ff 0.4.2",
  "ark-secp256r1",
@@ -4438,9 +4437,10 @@ dependencies = [
 [[package]]
 name = "mpz-garble"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "async-trait",
+ "blake3",
  "derive_builder 0.11.2",
  "futures",
  "hashbrown 0.14.5",
@@ -4453,7 +4453,7 @@ dependencies = [
  "mpz-vm-core",
  "opaque-debug",
  "rand 0.9.2",
- "rangeset 0.2.0",
+ "rangeset",
  "serde",
  "serio",
  "thiserror 1.0.69",
@@ -4464,7 +4464,7 @@ dependencies = [
 [[package]]
 name = "mpz-garble-core"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "aes 0.9.0-rc.2",
  "bitvec",
@@ -4484,7 +4484,7 @@ dependencies = [
  "rand 0.9.2",
  "rand_chacha 0.9.0",
  "rand_core 0.9.3",
- "rangeset 0.2.0",
+ "rangeset",
  "rayon",
  "serde",
  "serde_arrays",
@@ -4495,7 +4495,7 @@ dependencies = [
 [[package]]
 name = "mpz-hash"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "blake3",
  "itybity 0.3.1",
@@ -4508,7 +4508,7 @@ dependencies = [
 [[package]]
 name = "mpz-ideal-vm"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "async-trait",
  "futures",
@@ -4516,7 +4516,7 @@ dependencies = [
  "mpz-core",
  "mpz-memory-core",
  "mpz-vm-core",
- "rangeset 0.2.0",
+ "rangeset",
  "serde",
  "serio",
  "thiserror 1.0.69",
@@ -4525,14 +4525,14 @@ dependencies = [
 [[package]]
 name = "mpz-memory-core"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "blake3",
  "futures",
  "itybity 0.3.1",
  "mpz-core",
  "rand 0.9.2",
- "rangeset 0.2.0",
+ "rangeset",
  "serde",
  "thiserror 1.0.69",
 ]
@@ -4540,7 +4540,7 @@ dependencies = [
 [[package]]
 name = "mpz-ole"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "async-trait",
  "futures",
@@ -4558,7 +4558,7 @@ dependencies = [
 [[package]]
 name = "mpz-ole-core"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "hybrid-array",
  "itybity 0.3.1",
@@ -4574,7 +4574,7 @@ dependencies = [
 [[package]]
 name = "mpz-ot"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "async-trait",
  "cfg-if",
@@ -4589,6 +4589,7 @@ dependencies = [
  "rand 0.9.2",
  "rand_chacha 0.9.0",
  "rand_core 0.9.3",
+ "serde",
  "serio",
  "thiserror 1.0.69",
  "tokio",
@@ -4597,7 +4598,7 @@ dependencies = [
 [[package]]
 name = "mpz-ot-core"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "aes 0.9.0-rc.2",
  "blake3",
@@ -4628,7 +4629,7 @@ dependencies = [
 [[package]]
 name = "mpz-share-conversion"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "async-trait",
  "mpz-common",
@@ -4644,7 +4645,7 @@ dependencies = [
 [[package]]
 name = "mpz-share-conversion-core"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "mpz-common",
  "mpz-core",
@@ -4658,7 +4659,7 @@ dependencies = [
 [[package]]
 name = "mpz-vm-core"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "async-trait",
  "futures",
@@ -4671,7 +4672,7 @@ dependencies = [
 [[package]]
 name = "mpz-zk"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "async-trait",
  "blake3",
@@ -4689,7 +4690,7 @@ dependencies = [
 [[package]]
 name = "mpz-zk-core"
 version = "0.1.0-alpha.4"
-source = "git+https://github.com/privacy-ethereum/mpz?rev=9c343f8#9c343f86d386bc1360d6ac0a37eb1af65f48216a"
+source = "git+https://github.com/privacy-ethereum/mpz?rev=d9baf0f#d9baf0fb57e9db118caef92aaf9b693f4343aef1"
 dependencies = [
  "blake3",
  "cfg-if",
@@ -4698,7 +4699,8 @@ dependencies = [
  "mpz-core",
  "mpz-memory-core",
  "mpz-vm-core",
- "rangeset 0.2.0",
+ "rand_chacha 0.9.0",
+ "rangeset",
  "rayon",
  "serde",
  "thiserror 1.0.69",
@@ -5803,15 +5805,6 @@ version = "1.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "acbbbbea733ec66275512d0b9694f34102e7d5406fdbe2ad8d21b28dce92887c"
 
-[[package]]
-name = "rangeset"
-version = "0.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1fc7af00a06ad692080d87495a904677592c662610edb82b4fc8782f4ed2f01f"
-dependencies = [
- "serde",
-]
-
 [[package]]
 name = "rangeset"
 version = "0.4.0"
@@ -6879,7 +6872,7 @@ dependencies = [
  "httparse",
  "pest",
  "pest_derive",
- "rangeset 0.4.0",
+ "rangeset",
  "serde",
  "thiserror 1.0.69",
 ]
@@ -7273,7 +7266,7 @@ dependencies = [
  "once_cell",
  "opaque-debug",
  "rand 0.9.2",
- "rangeset 0.4.0",
+ "rangeset",
  "rstest",
  "rustls-pki-types",
  "rustls-webpki 0.103.8",
@@ -7286,6 +7279,7 @@ dependencies = [
  "tlsn-core",
  "tlsn-deap",
  "tlsn-mpc-tls",
+ "tlsn-mux",
  "tlsn-server-fixture",
  "tlsn-server-fixture-certs",
  "tlsn-tls-client",
@@ -7295,7 +7289,6 @@ dependencies = [
  "tokio-util",
  "tracing",
  "tracing-subscriber",
- "uid-mux",
  "web-spawn",
  "webpki-roots 1.0.4",
 ]
@@ -7314,7 +7307,7 @@ dependencies = [
  "p256",
  "rand 0.9.2",
  "rand06-compat",
- "rangeset 0.4.0",
+ "rangeset",
  "rstest",
  "serde",
  "thiserror 1.0.69",
@@ -7357,7 +7350,7 @@ dependencies = [
  "rand 0.9.2",
  "rand_chacha 0.9.0",
  "rand_core 0.9.3",
- "rangeset 0.4.0",
+ "rangeset",
  "rs_merkle",
  "rstest",
  "rustls-pki-types",
@@ -7392,7 +7385,7 @@ dependencies = [
  "mpz-core",
  "mpz-ideal-vm",
  "mpz-vm-core",
- "rangeset 0.4.0",
+ "rangeset",
  "serde",
  "serio",
  "thiserror 1.0.69",
@@ -7607,6 +7600,7 @@ dependencies = [
  "tlsn-core",
  "tlsn-hmac-sha256",
  "tlsn-key-exchange",
+ "tlsn-mux",
  "tlsn-tls-backend",
  "tlsn-tls-client",
  "tlsn-tls-client-async",
@@ -7615,10 +7609,25 @@ dependencies = [
  "tokio-util",
  "tracing",
  "tracing-subscriber",
- "uid-mux",
  "web-time 0.2.4",
 ]
 
+[[package]]
+name = "tlsn-mux"
+version = "0.1.0"
+source = "git+https://github.com/tlsnotary/tlsn-utils?rev=d9facb6#d9facb6f11daa182301b13c33688795e1affdb4a"
+dependencies = [
+ "blake3",
+ "futures",
+ "log",
+ "nohash-hasher",
+ "parking_lot",
+ "pin-project",
+ "rand 0.9.2",
+ "static_assertions",
+ "web-time 1.1.0",
+]
+
 [[package]]
 name = "tlsn-server-fixture"
 version = "0.0.0"
@@ -8141,23 +8150,6 @@ version = "0.1.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971"
 
-[[package]]
-name = "uid-mux"
-version = "0.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "34769f8c719d9e76d8f29186f1935b5c1fff91d27d839dd8d84dc4527eca8fdd"
-dependencies = [
- "async-trait",
- "blake3",
- "futures",
- "hex",
- "serio",
- "tokio",
- "tokio-util",
- "tracing",
- "yamux",
-]
-
 [[package]]
 name = "uint"
 version = "0.9.5"
@@ -8997,22 +8989,6 @@ dependencies = [
  "tap",
 ]
 
-[[package]]
-name = "yamux"
-version = "0.13.8"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "deab71f2e20691b4728b349c6cee8fc7223880fa67b6b4f92225ec32225447e5"
-dependencies = [
- "futures",
- "log",
- "nohash-hasher",
- "parking_lot",
- "pin-project",
- "rand 0.9.2",
- "static_assertions",
- "web-time 1.1.0",
-]
-
 [[package]]
 name = "yansi"
 version = "1.0.1"

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 = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-circuits-data = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-circuits-data = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-memory-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-memory-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-common = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-common = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-vm-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-vm-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-garble = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-garble = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-garble-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-garble-core = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-ole = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-ole = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-ot = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-ot = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-share-conversion = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-share-conversion = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-fields = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-fields = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-zk = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-zk = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-hash = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-hash = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
-mpz-ideal-vm = { git = "https://github.com/privacy-ethereum/mpz", rev = "9c343f8" }
+mpz-ideal-vm = { git = "https://github.com/privacy-ethereum/mpz", rev = "d9baf0f" }
 
 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/mpc-tls/tests/test.rs

@@ -123,8 +123,8 @@ fn build_pair(config: Config) -> (MpcTlsLeader, MpcTlsFollower) {
 
     let (mut mt_a, mut mt_b) = test_mt_context(8);
 
-    let ctx_a = futures::executor::block_on(mt_a.new_context()).unwrap();
+    let ctx_a = mt_a.new_context().unwrap();
-    let ctx_b = futures::executor::block_on(mt_b.new_context()).unwrap();
+    let ctx_b = mt_b.new_context().unwrap();
 
     let delta_a = Delta::new(Block::random(&mut rng));
     let delta_b = Delta::new(Block::random(&mut rng));

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/context.rs

@@ -1,21 +0,0 @@
-//! Execution context.
-
-use mpz_common::context::Multithread;
-
-use crate::mux::MuxControl;
-
-/// Maximum concurrency for multi-threaded context.
-pub(crate) const MAX_CONCURRENCY: usize = 8;
-
-/// Builds a multi-threaded context with the given muxer.
-pub(crate) fn build_mt_context(mux: MuxControl) -> Multithread {
-    let builder = Multithread::builder().mux(mux).concurrency(MAX_CONCURRENCY);
-
-    #[cfg(all(feature = "web", target_arch = "wasm32"))]
-    let builder = builder.spawn_handler(|f| {
-        let _ = web_spawn::spawn(f);
-        Ok(())
-    });
-
-    builder.build().unwrap()
-}

crates/tlsn/src/error.rs

@@ -0,0 +1,87 @@
+use std::fmt::Display;
+
+/// Crate-level error.
+#[derive(Debug, thiserror::Error)]
+pub struct Error {
+    kind: ErrorKind,
+    msg: Option<String>,
+    source: Option<Box<dyn std::error::Error + Send + Sync>>,
+}
+
+impl Error {
+    pub(crate) fn io() -> Self {
+        Self {
+            kind: ErrorKind::Internal,
+            msg: None,
+            source: None,
+        }
+    }
+
+    pub(crate) fn internal() -> Self {
+        Self {
+            kind: ErrorKind::Internal,
+            msg: None,
+            source: None,
+        }
+    }
+
+    pub(crate) fn with_msg(mut self, msg: impl Into<String>) -> Self {
+        self.msg = Some(msg.into());
+        self
+    }
+
+    pub(crate) fn with_source<T>(mut self, source: T) -> Self
+    where
+        T: Into<Box<dyn std::error::Error + Send + Sync>>,
+    {
+        self.source = Some(source.into());
+        self
+    }
+
+    /// Returns `true` if the error was user created.
+    pub fn is_user(&self) -> bool {
+        todo!()
+    }
+
+    /// Returns `true` if the error originated from an IO error.
+    pub fn is_io(&self) -> bool {
+        self.kind.is_io()
+    }
+
+    /// Returns `true` if the error originated from an internal bug.
+    pub fn is_internal(&self) -> bool {
+        self.kind.is_internal()
+    }
+
+    /// Returns the error message if available.
+    pub fn msg(&self) -> Option<&str> {
+        todo!()
+    }
+}
+
+impl Display for Error {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        todo!()
+    }
+}
+
+#[derive(Debug)]
+enum ErrorKind {
+    User,
+    Io,
+    Internal,
+}
+
+impl ErrorKind {
+    fn is_user(&self) -> bool {
+        matches!(self, ErrorKind::User)
+    }
+
+    fn is_io(&self) -> bool {
+        matches!(self, ErrorKind::Io)
+    }
+
+    fn is_internal(&self) -> bool {
+        matches!(self, ErrorKind::Internal)
+    }
+}

crates/tlsn/src/lib.rs

@@ -4,20 +4,25 @@
 #![deny(clippy::all)]
 #![forbid(unsafe_code)]
 
-pub(crate) mod context;
+mod error;
 pub(crate) mod ghash;
 pub(crate) mod map;
 pub(crate) mod mpz;
 pub(crate) mod msg;
-pub(crate) mod mux;
 pub mod prover;
+mod session;
 pub(crate) mod tag;
 pub(crate) mod transcript_internal;
 pub mod verifier;
 
+pub use error::Error;
+pub use session::Session;
 pub use tlsn_attestation as attestation;
 pub use tlsn_core::{config, connection, hash, transcript, webpki};
 
+/// Result type.
+pub type Result<T, E = Error> = core::result::Result<T, E>;
+
 use std::sync::LazyLock;
 
 use semver::Version;

crates/tlsn/src/mux.rs

@@ -1,90 +0,0 @@
-//! Multiplexer used in the TLSNotary protocol.
-
-use std::future::IntoFuture;
-
-use futures::{
-    AsyncRead, AsyncWrite, Future,
-    future::{FusedFuture, FutureExt},
-};
-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;
-
-/// 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>,
-);
-
-impl MuxFuture {
-    /// Returns true if the muxer is complete.
-    pub(crate) fn is_complete(&self) -> bool {
-        self.0.is_terminated()
-    }
-
-    /// 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());
-        // 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 {
-                    error!("mux error: {:?}", e);
-                },
-            }
-        }
-    }
-}
-
-impl Future for MuxFuture {
-    type Output = Result<(), yamux::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)
-    }
-}
-
-/// 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

@@ -7,14 +7,13 @@ pub mod state;
 
 pub use error::ProverError;
 pub use future::ProverFuture;
+use mpz_common::Context;
 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,
+    prover::error::ErrorKind,
     tag::verify_tags,
 };
 
@@ -45,6 +44,7 @@ use tracing::{Instrument, Span, debug, info, info_span, instrument};
 pub struct Prover<T: state::ProverState = state::Initialized> {
     config: ProverConfig,
     span: Span,
+    ctx: Option<Context>,
     state: T,
 }
 
@@ -53,12 +53,14 @@ impl Prover<state::Initialized> {
     ///
     /// # Arguments
     ///
+    /// * `ctx` - A thread context.
     /// * `config` - The configuration for the prover.
-    pub fn new(config: ProverConfig) -> Self {
+    pub(crate) fn new(ctx: Context, config: ProverConfig) -> Self {
         let span = info_span!("prover");
         Self {
             config,
             span,
+            ctx: Some(ctx),
             state: state::Initialized,
         }
     }
@@ -71,35 +73,30 @@ impl Prover<state::Initialized> {
     /// # 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<S: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
+    pub async fn commit(
-        self,
+        mut self,
         config: TlsCommitConfig,
-        socket: S,
     ) -> Result<Prover<state::CommitAccepted>, ProverError> {
-        let (mut mux_fut, mux_ctrl) = attach_mux(socket, Role::Prover);
+        let mut ctx = self
-        let mut mt = build_mt_context(mux_ctrl.clone());
+            .ctx
-        let mut ctx = mux_fut.poll_with(mt.new_context()).await?;
+            .take()
+            .ok_or_else(|| ProverError::new(ErrorKind::Io, "context was dropped"))?;
 
         // Sends protocol configuration to verifier for compatibility check.
-        mux_fut
+        ctx.io_mut()
-            .poll_with(async {
+            .send(TlsCommitRequestMsg {
-                ctx.io_mut()
+                request: config.to_request(),
-                    .send(TlsCommitRequestMsg {
+                version: crate::VERSION.clone(),
-                        request: config.to_request(),
-                        version: crate::VERSION.clone(),
-                    })
-                    .await?;
-
-                ctx.io_mut()
-                    .expect_next::<Response>()
-                    .await?
-                    .result
-                    .map_err(ProverError::from)
             })
             .await?;
 
+        ctx.io_mut()
+            .expect_next::<Response>()
+            .await?
+            .result
+            .map_err(ProverError::from)?;
+
         let TlsCommitProtocolConfig::Mpc(mpc_tls_config) = config.protocol().clone() else {
             unreachable!("only MPC TLS is supported");
         };
@@ -114,20 +111,15 @@ impl Prover<state::Initialized> {
 
         debug!("setting up mpc-tls");
 
-        mux_fut.poll_with(mpc_tls.preprocess()).await?;
+        mpc_tls.preprocess().await?;
 
         debug!("mpc-tls setup complete");
 
         Ok(Prover {
             config: self.config,
             span: self.span,
-            state: state::CommitAccepted {
+            ctx: None,
-                mux_ctrl,
+            state: state::CommitAccepted { mpc_tls, keys, vm },
-                mux_fut,
-                mpc_tls,
-                keys,
-                vm,
-            },
         })
     }
 }
@@ -150,12 +142,7 @@ impl Prover<state::CommitAccepted> {
         socket: S,
     ) -> Result<(TlsConnection, ProverFuture), ProverError> {
         let state::CommitAccepted {
-            mux_ctrl,
+            mpc_tls, keys, vm, ..
-            mut mux_fut,
-            mpc_tls,
-            keys,
-            vm,
-            ..
         } = self.state;
 
         let (mpc_ctrl, mpc_fut) = mpc_tls.run();
@@ -209,10 +196,7 @@ impl Prover<state::CommitAccepted> {
             let mpc_ctrl = mpc_ctrl.clone();
             async move {
                 let conn_fut = async {
-                    mux_fut
+                    conn_fut.await.map_err(ProverError::from)?;
-                        .poll_with(conn_fut.map_err(ProverError::from))
-                        .await?;
-
                     mpc_ctrl.stop().await?;
 
                     Ok::<_, ProverError>(())
@@ -233,10 +217,7 @@ impl Prover<state::CommitAccepted> {
                     debug!("finalizing mpc");
 
                     // Finalize DEAP.
-                    mux_fut
+                    vm.finalize(&mut ctx).await.map_err(ProverError::mpc)?;
-                        .poll_with(vm.finalize(&mut ctx))
-                        .await
-                        .map_err(ProverError::mpc)?;
 
                     debug!("mpc finalized");
                 }
@@ -258,9 +239,7 @@ impl Prover<state::CommitAccepted> {
                 )
                 .map_err(ProverError::zk)?;
 
-                mux_fut
+                vm.execute_all(&mut ctx).await.map_err(ProverError::zk)?;
-                    .poll_with(vm.execute_all(&mut ctx).map_err(ProverError::zk))
-                    .await?;
 
                 let transcript = tls_transcript
                     .to_transcript()
@@ -269,10 +248,8 @@ impl Prover<state::CommitAccepted> {
                 Ok(Prover {
                     config: self.config,
                     span: self.span,
+                    ctx: Some(ctx),
                     state: state::Committed {
-                        mux_ctrl,
-                        mux_fut,
-                        ctx,
                         vm,
                         server_name: config.server_name().clone(),
                         keys,
@@ -312,9 +289,11 @@ impl Prover<state::Committed> {
     /// * `config` - The disclosure configuration.
     #[instrument(parent = &self.span, level = "info", skip_all, err)]
     pub async fn prove(&mut self, config: &ProveConfig) -> Result<ProverOutput, ProverError> {
+        let ctx = self
+            .ctx
+            .as_mut()
+            .ok_or_else(|| ProverError::new(ErrorKind::Io, "context was dropped"))?;
         let state::Committed {
-            mux_fut,
-            ctx,
             vm,
             keys,
             server_name,
@@ -350,15 +329,10 @@ impl Prover<state::Committed> {
             transcript: partial_transcript,
         };
 
-        let output = mux_fut
+        ctx.io_mut().send(msg).await.map_err(ProverError::from)?;
-            .poll_with(async {
+        ctx.io_mut().expect_next::<Response>().await?.result?;
-                ctx.io_mut().send(msg).await.map_err(ProverError::from)?;
 
-                ctx.io_mut().expect_next::<Response>().await?.result?;
+        let output = prove::prove(ctx, vm, keys, transcript, tls_transcript, config).await?;
-
-                prove::prove(ctx, vm, keys, transcript, tls_transcript, config).await
-            })
-            .await?;
 
         Ok(output)
     }
@@ -366,16 +340,6 @@ 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;
-
-        // Wait for the verifier to correctly close the connection.
-        if !mux_fut.is_complete() {
-            mux_ctrl.close();
-            mux_fut.await?;
-        }
-
         Ok(())
     }
 }

crates/tlsn/src/prover/error.rs

@@ -10,7 +10,7 @@ pub struct ProverError {
 }
 
 impl ProverError {
-    fn new<E>(kind: ErrorKind, source: E) -> Self
+    pub(crate) fn new<E>(kind: ErrorKind, source: E) -> Self
     where
         E: Into<Box<dyn Error + Send + Sync + 'static>>,
     {
@@ -50,7 +50,7 @@ impl ProverError {
 }
 
 #[derive(Debug)]
-enum ErrorKind {
+pub(crate) enum ErrorKind {
     Io,
     Mpc,
     Zk,
@@ -90,8 +90,8 @@ impl From<tls_client_async::ConnectionError> for ProverError {
     }
 }
 
-impl From<uid_mux::yamux::ConnectionError> for ProverError {
+impl From<tlsn_mux::ConnectionError> for ProverError {
-    fn from(e: uid_mux::yamux::ConnectionError) -> Self {
+    fn from(e: tlsn_mux::ConnectionError) -> Self {
         Self::new(ErrorKind::Io, e)
     }
 }

crates/tlsn/src/prover/state.rs

@@ -3,7 +3,6 @@
 use std::sync::Arc;
 
 use mpc_tls::{MpcTlsLeader, SessionKeys};
-use mpz_common::Context;
 use tlsn_core::{
     connection::ServerName,
     transcript::{TlsTranscript, Transcript},
@@ -11,10 +10,7 @@ use tlsn_core::{
 use tlsn_deap::Deap;
 use tokio::sync::Mutex;
 
-use crate::{
+use crate::mpz::{ProverMpc, ProverZk};
-    mpz::{ProverMpc, ProverZk},
-    mux::{MuxControl, MuxFuture},
-};
 
 /// Entry state
 pub struct Initialized;
@@ -24,8 +20,6 @@ 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(crate) mpc_tls: MpcTlsLeader,
     pub(crate) keys: SessionKeys,
     pub(crate) vm: Arc<Mutex<Deap<ProverMpc, ProverZk>>>,
@@ -35,9 +29,6 @@ opaque_debug::implement!(CommitAccepted);
 
 /// State after the TLS transcript has been committed.
 pub struct Committed {
-    pub(crate) mux_ctrl: MuxControl,
-    pub(crate) mux_fut: MuxFuture,
-    pub(crate) ctx: Context,
     pub(crate) vm: ProverZk,
     pub(crate) server_name: ServerName,
     pub(crate) keys: SessionKeys,

crates/tlsn/src/session.rs

@@ -0,0 +1,293 @@
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+    task::{Context, Poll},
+};
+
+use futures::{AsyncRead, AsyncWrite};
+use mpz_common::{ThreadId, context::Multithread, io::Io, mux::Mux};
+use tlsn_core::config::{prover::ProverConfig, verifier::VerifierConfig};
+use tlsn_mux::{Connection, Handle};
+
+use crate::{
+    Error, Result,
+    prover::{Prover, state as prover_state},
+    verifier::{Verifier, state as verifier_state},
+};
+
+/// Maximum concurrency for multi-threaded context.
+const MAX_CONCURRENCY: usize = 8;
+
+/// Session state.
+#[must_use = "session must be polled continuously to make progress, including during closing."]
+pub struct Session<Io> {
+    conn: Option<Connection<Io>>,
+    mt: Multithread,
+}
+
+impl<Io> Session<Io>
+where
+    Io: AsyncRead + AsyncWrite + Unpin,
+{
+    /// Creates a new session.
+    pub fn new(io: Io) -> 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(io, mux_config);
+        let handle = conn.handle().expect("handle should be available");
+        let mt = build_mt_context(MuxHandle { handle: handle });
+
+        Self {
+            conn: Some(conn),
+            mt,
+        }
+    }
+
+    /// Creates a new prover.
+    pub fn new_prover(
+        &mut self,
+        config: ProverConfig,
+    ) -> Result<Prover<prover_state::Initialized>> {
+        let ctx = self.mt.new_context().map_err(|e| {
+            Error::internal()
+                .with_msg("failed to created new prover")
+                .with_source(e)
+        })?;
+
+        Ok(Prover::new(ctx, config))
+    }
+
+    /// Creates a new verifier.
+    pub fn new_verifier(
+        &mut self,
+        config: VerifierConfig,
+    ) -> Result<Verifier<verifier_state::Initialized>> {
+        let ctx = self.mt.new_context().map_err(|e| {
+            Error::internal()
+                .with_msg("failed to created new verifier")
+                .with_source(e)
+        })?;
+
+        Ok(Verifier::new(ctx, config))
+    }
+
+    /// Returns `true` if the session is closed.
+    pub fn is_closed(&self) -> bool {
+        self.conn
+            .as_ref()
+            .map(|mux| mux.is_complete())
+            .unwrap_or_default()
+    }
+
+    /// Closes the session.
+    ///
+    /// This will cause the session to begin closing. Session must continue to be polled until completion.
+    pub fn close(&mut self) {
+        self.conn.as_mut().map(|conn| conn.close());
+    }
+
+    /// Attempts to take the IO, returning an error if it is not available.
+    pub fn try_take(&mut self) -> Result<Io> {
+        let conn = self.conn.take().ok_or_else(|| {
+            Error::io().with_msg("failed to take the session io, it was already taken")
+        })?;
+
+        match conn.try_into_io() {
+            Err(conn) => {
+                self.conn = Some(conn);
+                Err(Error::io()
+                    .with_msg("failed to take the session io, session was not completed yet"))
+            }
+            Ok(conn) => Ok(conn),
+        }
+    }
+
+    /// Polls the session.
+    pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
+        self.conn
+            .as_mut()
+            .ok_or_else(|| {
+                Error::io()
+                    .with_msg("failed to poll the session connection because it has been taken")
+            })?
+            .poll(cx)
+            .map_err(|e| {
+                Error::io()
+                    .with_msg("error occurred while polling the session connection")
+                    .with_source(e)
+            })
+    }
+
+    /// Splits the session into a driver and handle.
+    ///
+    /// The driver must be polled to make progress. The handle is used
+    /// for creating provers/verifiers and closing the session.
+    pub fn split(self) -> (SessionDriver<Io>, SessionHandle) {
+        let should_close = Arc::new(AtomicBool::new(false));
+
+        (
+            SessionDriver {
+                conn: self.conn,
+                should_close: should_close.clone(),
+            },
+            SessionHandle {
+                mt: self.mt,
+                should_close,
+            },
+        )
+    }
+}
+
+impl<Io> Future for Session<Io>
+where
+    Io: AsyncRead + AsyncWrite + Unpin,
+{
+    type Output = Result<()>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        Session::poll(&mut (*self), cx)
+    }
+}
+
+/// The polling half of a split session.
+///
+/// Must be polled continuously to drive the session. Returns the underlying
+/// IO when the session closes.
+#[must_use = "driver must be polled to make progress"]
+pub struct SessionDriver<Io> {
+    conn: Option<Connection<Io>>,
+    should_close: Arc<AtomicBool>,
+}
+
+impl<Io> SessionDriver<Io>
+where
+    Io: AsyncRead + AsyncWrite + Unpin,
+{
+    /// Polls the driver.
+    pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Result<Io>> {
+        let conn = self.conn.as_mut().ok_or_else(|| {
+            Error::io().with_msg("session driver already completed")
+        })?;
+
+        if self.should_close.load(Ordering::Acquire) {
+            conn.close();
+        }
+
+        match conn.poll(cx) {
+            Poll::Ready(Ok(())) => {}
+            Poll::Ready(Err(e)) => {
+                return Poll::Ready(Err(Error::io()
+                    .with_msg("error polling session connection")
+                    .with_source(e)));
+            }
+            Poll::Pending => return Poll::Pending,
+        }
+
+        let conn = self.conn.take().unwrap();
+        Poll::Ready(conn.try_into_io().map_err(|_| {
+            Error::io().with_msg("failed to take session io")
+        }))
+    }
+}
+
+impl<Io> Future for SessionDriver<Io>
+where
+    Io: AsyncRead + AsyncWrite + Unpin,
+{
+    type Output = Result<Io>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        SessionDriver::poll(&mut *self, cx)
+    }
+}
+
+/// The control half of a split session.
+///
+/// Used to create provers/verifiers and control the session lifecycle.
+pub struct SessionHandle {
+    mt: Multithread,
+    should_close: Arc<AtomicBool>,
+}
+
+impl SessionHandle {
+    /// Creates a new prover.
+    pub fn new_prover(
+        &mut self,
+        config: ProverConfig,
+    ) -> Result<Prover<prover_state::Initialized>> {
+        let ctx = self.mt.new_context().map_err(|e| {
+            Error::internal()
+                .with_msg("failed to create new prover")
+                .with_source(e)
+        })?;
+
+        Ok(Prover::new(ctx, config))
+    }
+
+    /// Creates a new verifier.
+    pub fn new_verifier(
+        &mut self,
+        config: VerifierConfig,
+    ) -> Result<Verifier<verifier_state::Initialized>> {
+        let ctx = self.mt.new_context().map_err(|e| {
+            Error::internal()
+                .with_msg("failed to create new verifier")
+                .with_source(e)
+        })?;
+
+        Ok(Verifier::new(ctx, config))
+    }
+
+    /// Signals the session to close.
+    ///
+    /// The driver must continue to be polled until it completes.
+    pub fn close(&self) {
+        self.should_close.store(true, Ordering::Release);
+    }
+}
+
+/// Multiplexer controller providing streams.
+struct MuxHandle {
+    handle: Handle,
+}
+
+impl std::fmt::Debug for MuxHandle {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("MuxHandle").finish_non_exhaustive()
+    }
+}
+
+impl Mux for MuxHandle {
+    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)
+    }
+}
+
+/// Builds a multi-threaded context with the given muxer.
+fn build_mt_context(mux: MuxHandle) -> Multithread {
+    let builder = Multithread::builder()
+        .mux(Box::new(mux) as Box<_>)
+        .concurrency(MAX_CONCURRENCY);
+
+    #[cfg(all(feature = "web", target_arch = "wasm32"))]
+    let builder = builder.spawn_handler(|f| {
+        let _ = web_spawn::spawn(f);
+        Ok(())
+    });
+
+    builder.build().unwrap()
+}

crates/tlsn/src/verifier.rs

@@ -7,17 +7,16 @@ mod verify;
 use std::sync::Arc;
 
 pub use error::VerifierError;
+use mpz_common::Context;
 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,
     tag::verify_tags,
+    verifier::error::ErrorKind,
 };
-use futures::{AsyncRead, AsyncWrite, TryFutureExt};
+use futures::TryFutureExt;
 use mpz_vm_core::prelude::*;
 use serio::{SinkExt, stream::IoStreamExt};
 use tlsn_core::{
@@ -45,16 +44,18 @@ pub struct SessionInfo {
 pub struct Verifier<T: state::VerifierState = state::Initialized> {
     config: VerifierConfig,
     span: Span,
+    ctx: Option<Context>,
     state: T,
 }
 
 impl Verifier<state::Initialized> {
     /// Creates a new verifier.
-    pub fn new(config: VerifierConfig) -> Self {
+    pub(crate) fn new(ctx: Context, config: VerifierConfig) -> Self {
         let span = info_span!("verifier");
         Self {
             config,
             span,
+            ctx: Some(ctx),
             state: state::Initialized,
         }
     }
@@ -63,37 +64,22 @@ impl Verifier<state::Initialized> {
     ///
     /// This initiates the TLS commitment protocol, receiving the prover's
     /// configuration and providing the opportunity to accept or reject it.
-    ///
-    /// # Arguments
-    ///
-    /// * `socket` - The socket to the prover.
     #[instrument(parent = &self.span, level = "info", skip_all, err)]
-    pub async fn commit<S: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
+    pub async fn commit(mut self) -> Result<Verifier<state::CommitStart>, VerifierError> {
-        self,
+        let mut ctx = self
-        socket: S,
+            .ctx
-    ) -> Result<Verifier<state::CommitStart>, VerifierError> {
+            .take()
-        let (mut mux_fut, mux_ctrl) = attach_mux(socket, Role::Verifier);
+            .ok_or_else(|| VerifierError::new(ErrorKind::Io, "context was dropped"))?;
-        let mut mt = build_mt_context(mux_ctrl.clone());
-        let mut ctx = mux_fut.poll_with(mt.new_context()).await?;
 
         // Receives protocol configuration from prover to perform compatibility check.
-        let TlsCommitRequestMsg { request, version } =
+        let TlsCommitRequestMsg { request, version } = ctx.io_mut().expect_next().await?;
-            mux_fut.poll_with(ctx.io_mut().expect_next()).await?;
 
         if version != *crate::VERSION {
             let msg = format!(
                 "prover version does not match with verifier: {version} != {}",
                 *crate::VERSION
             );
-            mux_fut
+            ctx.io_mut().send(Response::err(Some(msg.clone()))).await?;
-                .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?;
-            }
 
             return Err(VerifierError::config(msg));
         }
@@ -101,12 +87,8 @@ impl Verifier<state::Initialized> {
         Ok(Verifier {
             config: self.config,
             span: self.span,
-            state: state::CommitStart {
+            ctx: Some(ctx),
-                mux_ctrl,
+            state: state::CommitStart { request },
-                mux_fut,
-                ctx,
-                request,
-            },
         })
     }
 }
@@ -119,15 +101,14 @@ 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(mut self) -> Result<Verifier<state::CommitAccepted>, VerifierError> {
-        let state::CommitStart {
+        let mut ctx = self
-            mux_ctrl,
+            .ctx
-            mut mux_fut,
+            .take()
-            mut ctx,
+            .ok_or_else(|| VerifierError::new(ErrorKind::Io, "context was dropped"))?;
-            request,
+        let state::CommitStart { request } = self.state;
-        } = self.state;
 
-        mux_fut.poll_with(ctx.io_mut().send(Response::ok())).await?;
+        ctx.io_mut().send(Response::ok()).await?;
 
         let TlsCommitProtocolConfig::Mpc(mpc_tls_config) = request.protocol().clone() else {
             unreachable!("only MPC TLS is supported");
@@ -143,42 +124,27 @@ impl Verifier<state::CommitStart> {
 
         debug!("setting up mpc-tls");
 
-        mux_fut.poll_with(mpc_tls.preprocess()).await?;
+        mpc_tls.preprocess().await?;
 
         debug!("mpc-tls setup complete");
 
         Ok(Verifier {
             config: self.config,
             span: self.span,
-            state: state::CommitAccepted {
+            ctx: None,
-                mux_ctrl,
+            state: state::CommitAccepted { mpc_tls, keys, vm },
-                mux_fut,
-                mpc_tls,
-                keys,
-                vm,
-            },
         })
     }
 
     /// Rejects the proposed protocol configuration.
     #[instrument(parent = &self.span, level = "info", skip_all, err)]
-    pub async fn reject(self, msg: Option<&str>) -> Result<(), VerifierError> {
+    pub async fn reject(mut self, msg: Option<&str>) -> Result<(), VerifierError> {
-        let state::CommitStart {
+        let mut ctx = self
-            mux_ctrl,
+            .ctx
-            mut mux_fut,
+            .take()
-            mut ctx,
+            .ok_or_else(|| VerifierError::new(ErrorKind::Io, "context was dropped"))?;
-            ..
-        } = self.state;
 
-        mux_fut
+        ctx.io_mut().send(Response::err(msg)).await?;
-            .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?;
-        }
 
         Ok(())
     }
@@ -188,17 +154,11 @@ impl Verifier<state::CommitAccepted> {
     /// 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> {
-        let state::CommitAccepted {
+        let state::CommitAccepted { mpc_tls, vm, keys } = self.state;
-            mux_ctrl,
-            mut mux_fut,
-            mpc_tls,
-            vm,
-            keys,
-        } = self.state;
 
         info!("starting MPC-TLS");
 
-        let (mut ctx, tls_transcript) = mux_fut.poll_with(mpc_tls.run()).await?;
+        let (mut ctx, tls_transcript) = mpc_tls.run().await?;
 
         info!("finished MPC-TLS");
 
@@ -207,10 +167,7 @@ impl Verifier<state::CommitAccepted> {
 
             debug!("finalizing mpc");
 
-            mux_fut
+            vm.finalize(&mut ctx).await.map_err(VerifierError::mpc)?;
-                .poll_with(vm.finalize(&mut ctx))
-                .await
-                .map_err(VerifierError::mpc)?;
 
             debug!("mpc finalized");
         }
@@ -232,9 +189,7 @@ impl Verifier<state::CommitAccepted> {
         )
         .map_err(VerifierError::zk)?;
 
-        mux_fut
+        vm.execute_all(&mut ctx).map_err(VerifierError::zk).await?;
-            .poll_with(vm.execute_all(&mut ctx).map_err(VerifierError::zk))
-            .await?;
 
         // Verify the tags.
         // After the verification, the entire TLS trancript becomes
@@ -244,10 +199,8 @@ impl Verifier<state::CommitAccepted> {
         Ok(Verifier {
             config: self.config,
             span: self.span,
+            ctx: Some(ctx),
             state: state::Committed {
-                mux_ctrl,
-                mux_fut,
-                ctx,
                 vm,
                 keys,
                 tls_transcript,
@@ -264,11 +217,12 @@ 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(mut self) -> Result<Verifier<state::Verify>, VerifierError> {
+        let mut ctx = self
+            .ctx
+            .take()
+            .ok_or_else(|| VerifierError::new(ErrorKind::Io, "context was dropped"))?;
         let state::Committed {
-            mux_ctrl,
-            mut mux_fut,
-            mut ctx,
             vm,
             keys,
             tls_transcript,
@@ -278,17 +232,17 @@ impl Verifier<state::Committed> {
             request,
             handshake,
             transcript,
-        } = mux_fut
+        } = ctx
-            .poll_with(ctx.io_mut().expect_next().map_err(VerifierError::from))
+            .io_mut()
+            .expect_next()
+            .map_err(VerifierError::from)
             .await?;
 
         Ok(Verifier {
             config: self.config,
             span: self.span,
+            ctx: Some(ctx),
             state: state::Verify {
-                mux_ctrl,
-                mux_fut,
-                ctx,
                 vm,
                 keys,
                 tls_transcript,
@@ -302,16 +256,6 @@ 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;
-
-        // Wait for the prover to correctly close the connection.
-        if !mux_fut.is_complete() {
-            mux_ctrl.close();
-            mux_fut.await?;
-        }
-
         Ok(())
     }
 }
@@ -324,12 +268,13 @@ impl Verifier<state::Verify> {
 
     /// Accepts the proving request.
     pub async fn accept(
-        self,
+        mut self,
     ) -> Result<(VerifierOutput, Verifier<state::Committed>), VerifierError> {
+        let mut ctx = self
+            .ctx
+            .take()
+            .ok_or_else(|| VerifierError::new(ErrorKind::Io, "context was dropped"))?;
         let state::Verify {
-            mux_ctrl,
-            mut mux_fut,
-            mut ctx,
             mut vm,
             keys,
             tls_transcript,
@@ -338,33 +283,30 @@ impl Verifier<state::Verify> {
             transcript,
         } = self.state;
 
-        mux_fut.poll_with(ctx.io_mut().send(Response::ok())).await?;
+        ctx.io_mut().send(Response::ok()).await?;
 
         let cert_verifier =
             ServerCertVerifier::new(self.config.root_store()).map_err(VerifierError::config)?;
 
-        let output = mux_fut
+        let output = verify::verify(
-            .poll_with(verify::verify(
+            &mut ctx,
-                &mut ctx,
+            &mut vm,
-                &mut vm,
+            &keys,
-                &keys,
+            &cert_verifier,
-                &cert_verifier,
+            &tls_transcript,
-                &tls_transcript,
+            request,
-                request,
+            handshake,
-                handshake,
+            transcript,
-                transcript,
+        )
-            ))
+        .await?;
-            .await?;
 
         Ok((
             output,
             Verifier {
                 config: self.config,
                 span: self.span,
+                ctx: Some(ctx),
                 state: state::Committed {
-                    mux_ctrl,
-                    mux_fut,
-                    ctx,
                     vm,
                     keys,
                     tls_transcript,
@@ -375,30 +317,27 @@ impl Verifier<state::Verify> {
 
     /// Rejects the proving request.
     pub async fn reject(
-        self,
+        mut self,
         msg: Option<&str>,
     ) -> Result<Verifier<state::Committed>, VerifierError> {
+        let mut ctx = self
+            .ctx
+            .take()
+            .ok_or_else(|| VerifierError::new(ErrorKind::Io, "context was dropped"))?;
         let state::Verify {
-            mux_ctrl,
-            mut mux_fut,
-            mut ctx,
             vm,
             keys,
             tls_transcript,
             ..
         } = self.state;
 
-        mux_fut
+        ctx.io_mut().send(Response::err(msg)).await?;
-            .poll_with(ctx.io_mut().send(Response::err(msg)))
-            .await?;
 
         Ok(Verifier {
             config: self.config,
             span: self.span,
+            ctx: Some(ctx),
             state: state::Committed {
-                mux_ctrl,
-                mux_fut,
-                ctx,
                 vm,
                 keys,
                 tls_transcript,

crates/tlsn/src/verifier/error.rs

@@ -10,7 +10,7 @@ pub struct VerifierError {
 }
 
 impl VerifierError {
-    fn new<E>(kind: ErrorKind, source: E) -> Self
+    pub(crate) fn new<E>(kind: ErrorKind, source: E) -> Self
     where
         E: Into<Box<dyn Error + Send + Sync + 'static>>,
     {
@@ -50,7 +50,7 @@ impl VerifierError {
 }
 
 #[derive(Debug)]
-enum ErrorKind {
+pub(crate) enum ErrorKind {
     Io,
     Config,
     Mpc,
@@ -84,8 +84,8 @@ impl From<std::io::Error> for VerifierError {
     }
 }
 
-impl From<uid_mux::yamux::ConnectionError> for VerifierError {
+impl From<tlsn_mux::ConnectionError> for VerifierError {
-    fn from(e: uid_mux::yamux::ConnectionError) -> Self {
+    fn from(e: tlsn_mux::ConnectionError) -> Self {
         Self::new(ErrorKind::Io, e)
     }
 }

crates/tlsn/src/verifier/state.rs

@@ -2,9 +2,7 @@
 
 use std::sync::Arc;
 
-use crate::mux::{MuxControl, MuxFuture};
 use mpc_tls::{MpcTlsFollower, SessionKeys};
-use mpz_common::Context;
 use tlsn_core::{
     config::{prove::ProveRequest, tls_commit::TlsCommitRequest},
     connection::{HandshakeData, ServerName},
@@ -25,9 +23,6 @@ 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(crate) ctx: Context,
     pub(crate) request: TlsCommitRequest,
 }
 
@@ -36,8 +31,6 @@ opaque_debug::implement!(CommitStart);
 /// 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(crate) mpc_tls: MpcTlsFollower,
     pub(crate) keys: SessionKeys,
     pub(crate) vm: Arc<Mutex<Deap<VerifierMpc, VerifierZk>>>,
@@ -47,9 +40,6 @@ opaque_debug::implement!(CommitAccepted);
 
 /// State after the TLS transcript has been committed.
 pub struct Committed {
-    pub(crate) mux_ctrl: MuxControl,
-    pub(crate) mux_fut: MuxFuture,
-    pub(crate) ctx: Context,
     pub(crate) vm: VerifierZk,
     pub(crate) keys: SessionKeys,
     pub(crate) tls_transcript: TlsTranscript,
@@ -59,9 +49,6 @@ opaque_debug::implement!(Committed);
 
 /// State after receiving a proving request.
 pub struct Verify {
-    pub(crate) mux_ctrl: MuxControl,
-    pub(crate) mux_fut: MuxFuture,
-    pub(crate) ctx: Context,
     pub(crate) vm: VerifierZk,
     pub(crate) keys: SessionKeys,
     pub(crate) tls_transcript: TlsTranscript,

crates/tlsn/tests/test.rs

@@ -1,5 +1,6 @@
 use futures::{AsyncReadExt, AsyncWriteExt};
 use tlsn::{
+    Session,
     config::{
         prove::ProveConfig,
         prover::ProverConfig,
@@ -18,9 +19,7 @@ use tlsn_core::ProverOutput;
 use tlsn_server_fixture::bind;
 use tlsn_server_fixture_certs::{CA_CERT_DER, SERVER_DOMAIN};
 
-use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_util::compat::TokioAsyncReadCompatExt;
-use tracing::instrument;
 
 // Maximum number of bytes that can be sent from prover to server
 const MAX_SENT_DATA: usize = 1 << 12;
@@ -37,9 +36,34 @@ async fn test() {
     tracing_subscriber::fmt::init();
 
     let (socket_0, socket_1) = tokio::io::duplex(2 << 23);
+    let mut session_p = Session::new(socket_0.compat());
+    let mut session_v = Session::new(socket_1.compat());
+
+    let prover = session_p
+        .new_prover(ProverConfig::builder().build().unwrap())
+        .unwrap();
+    let verifier = session_v
+        .new_verifier(
+            VerifierConfig::builder()
+                .root_store(RootCertStore {
+                    roots: vec![CertificateDer(CA_CERT_DER.to_vec())],
+                })
+                .build()
+                .unwrap(),
+        )
+        .unwrap();
+
+    let (session_p_driver, session_p_handle) = session_p.split();
+    let (session_v_driver, session_v_handle) = session_v.split();
+
+    tokio::spawn(session_p_driver);
+    tokio::spawn(session_v_driver);
 
     let ((_full_transcript, _prover_output), verifier_output) =
-        tokio::join!(prover(socket_0), verifier(socket_1));
+        tokio::join!(run_prover(prover), run_verifier(verifier));
+
+    session_p_handle.close();
+    session_v_handle.close();
 
     let partial_transcript = verifier_output.transcript.unwrap();
     let ServerName::Dns(server_name) = verifier_output.server_name.unwrap();
@@ -56,15 +80,12 @@ async fn test() {
     );
 }
 
-#[instrument(skip(verifier_socket))]
+async fn run_prover(prover: Prover) -> (Transcript, ProverOutput) {
-async fn prover<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
-    verifier_socket: T,
-) -> (Transcript, ProverOutput) {
     let (client_socket, server_socket) = tokio::io::duplex(2 << 16);
 
     let server_task = tokio::spawn(bind(server_socket.compat()));
 
-    let prover = Prover::new(ProverConfig::builder().build().unwrap())
+    let prover = prover
         .commit(
             TlsCommitConfig::builder()
                 .protocol(
@@ -78,7 +99,6 @@ async fn prover<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
                 )
                 .build()
                 .unwrap(),
-            verifier_socket.compat(),
         )
         .await
         .unwrap();
@@ -150,21 +170,9 @@ async fn prover<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
     (transcript, output)
 }
 
-#[instrument(skip(socket))]
+async fn run_verifier(verifier: Verifier) -> VerifierOutput {
-async fn verifier<T: AsyncWrite + AsyncRead + Send + Sync + Unpin + 'static>(
-    socket: T,
-) -> VerifierOutput {
-    let verifier = Verifier::new(
-        VerifierConfig::builder()
-            .root_store(RootCertStore {
-                roots: vec![CertificateDer(CA_CERT_DER.to_vec())],
-            })
-            .build()
-            .unwrap(),
-    );
-
     let verifier = verifier
-        .commit(socket.compat())
+        .commit()
         .await
         .unwrap()
         .accept()