diff --git a/Cargo.lock b/Cargo.lock index b99ed5085a..41c6fd9a87 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1741,7 +1741,7 @@ checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675" [[package]] name = "clmul" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "bytemuck", "cfg-if", @@ -3770,7 +3770,7 @@ checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" [[package]] name = "matrix-transpose" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "thiserror 1.0.69", ] @@ -3914,7 +3914,7 @@ dependencies = [ [[package]] name = "mpz-circuits" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "bincode", "itybity 0.3.1", @@ -3930,7 +3930,7 @@ dependencies = [ [[package]] name = "mpz-cointoss" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "futures", "mpz-cointoss-core", @@ -3943,7 +3943,7 @@ dependencies = [ [[package]] name = "mpz-cointoss-core" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "mpz-core", "opaque-debug", @@ -3954,7 +3954,7 @@ dependencies = [ [[package]] name = "mpz-common" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "async-trait", "bytes", @@ -3976,7 +3976,7 @@ dependencies = [ [[package]] name = "mpz-core" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "aes 0.9.0-rc.0", "bcs", @@ -4002,7 +4002,7 @@ dependencies = [ [[package]] name = "mpz-fields" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "ark-ff 0.4.2", "ark-secp256r1", @@ -4022,7 +4022,7 @@ dependencies = [ [[package]] name = "mpz-garble" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "async-trait", "derive_builder 0.11.2", @@ -4049,7 +4049,7 @@ dependencies = [ [[package]] name = "mpz-garble-core" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "aes 0.9.0-rc.0", "bitvec", @@ -4080,7 +4080,7 @@ dependencies = [ [[package]] name = "mpz-hash" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "itybity 0.3.1", "mpz-circuits", @@ -4092,7 +4092,7 @@ dependencies = [ [[package]] name = "mpz-memory-core" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "blake3", "futures", @@ -4107,7 +4107,7 @@ dependencies = [ [[package]] name = "mpz-ole" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "async-trait", "futures", @@ -4125,7 +4125,7 @@ dependencies = [ [[package]] name = "mpz-ole-core" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "hybrid-array", "itybity 0.3.1", @@ -4141,7 +4141,7 @@ dependencies = [ [[package]] name = "mpz-ot" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "async-trait", "cfg-if", @@ -4165,7 +4165,7 @@ dependencies = [ [[package]] name = "mpz-ot-core" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "aes 0.9.0-rc.0", "blake3", @@ -4196,7 +4196,7 @@ dependencies = [ [[package]] name = "mpz-share-conversion" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "async-trait", "mpz-common", @@ -4212,7 +4212,7 @@ dependencies = [ [[package]] name = "mpz-share-conversion-core" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "mpz-common", "mpz-core", @@ -4226,7 +4226,7 @@ dependencies = [ [[package]] name = "mpz-vm-core" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "async-trait", "futures", @@ -4239,7 +4239,7 @@ dependencies = [ [[package]] name = "mpz-zk" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "async-trait", "blake3", @@ -4251,12 +4251,13 @@ dependencies = [ "mpz-zk-core", "serio", "thiserror 1.0.69", + "tokio", ] [[package]] name = "mpz-zk-core" version = "0.1.0-alpha.3" -source = "git+https://github.com/privacy-scaling-explorations/mpz?rev=ccc0057#ccc00572f35c589706ca6be0e2e8749eb9b4414a" +source = "git+https://github.com/privacy-scaling-explorations/mpz?branch=refactor%2Fzk_expose_ot_handle#01262276483268259887508d87234911dbf0c336" dependencies = [ "blake3", "cfg-if", diff --git a/Cargo.toml b/Cargo.toml index 221befbe1e..0eae547f01 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -70,19 +70,19 @@ tlsn-harness-runner = { path = "crates/harness/runner" } tlsn-wasm = { path = "crates/wasm" } tlsn = { path = "crates/tlsn" } -mpz-circuits = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "ccc0057" } -mpz-memory-core = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "ccc0057" } -mpz-common = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "ccc0057" } -mpz-core = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "ccc0057" } -mpz-vm-core = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "ccc0057" } -mpz-garble = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "ccc0057" } -mpz-garble-core = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "ccc0057" } -mpz-ole = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "ccc0057" } -mpz-ot = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "ccc0057" } -mpz-share-conversion = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "ccc0057" } -mpz-fields = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "ccc0057" } -mpz-zk = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "ccc0057" } -mpz-hash = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "ccc0057" } +mpz-circuits = { git = "https://github.com/privacy-scaling-explorations/mpz", branch = "refactor/zk_expose_ot_handle" } +mpz-memory-core = { git = "https://github.com/privacy-scaling-explorations/mpz", branch = "refactor/zk_expose_ot_handle" } +mpz-common = { git = "https://github.com/privacy-scaling-explorations/mpz", branch = "refactor/zk_expose_ot_handle" } +mpz-core = { git = "https://github.com/privacy-scaling-explorations/mpz", branch = "refactor/zk_expose_ot_handle" } +mpz-vm-core = { git = "https://github.com/privacy-scaling-explorations/mpz", branch = "refactor/zk_expose_ot_handle" } +mpz-garble = { git = "https://github.com/privacy-scaling-explorations/mpz", branch = "refactor/zk_expose_ot_handle" } +mpz-garble-core = { git = "https://github.com/privacy-scaling-explorations/mpz", branch = "refactor/zk_expose_ot_handle" } +mpz-ole = { git = "https://github.com/privacy-scaling-explorations/mpz", branch = "refactor/zk_expose_ot_handle" } +mpz-ot = { git = "https://github.com/privacy-scaling-explorations/mpz", branch = "refactor/zk_expose_ot_handle" } +mpz-share-conversion = { git = "https://github.com/privacy-scaling-explorations/mpz", branch = "refactor/zk_expose_ot_handle" } +mpz-fields = { git = "https://github.com/privacy-scaling-explorations/mpz", branch = "refactor/zk_expose_ot_handle" } +mpz-zk = { git = "https://github.com/privacy-scaling-explorations/mpz", branch = "refactor/zk_expose_ot_handle" } +mpz-hash = { git = "https://github.com/privacy-scaling-explorations/mpz", branch = "refactor/zk_expose_ot_handle" } rangeset = { version = "0.2" } serio = { version = "0.2" } diff --git a/crates/components/deap/src/lib.rs b/crates/components/deap/src/lib.rs index 03598cd244..11c8f96c70 100644 --- a/crates/components/deap/src/lib.rs +++ b/crates/components/deap/src/lib.rs @@ -43,6 +43,17 @@ pub struct Deap { /// Outputs of the follower from the ZK VM. The references /// correspond to the MPC VM. outputs: Vec<(Slice, DecodeFuture)>, + /// Whether the DEAP VM should operate in a limited mode. + /// + /// The following limitations will apply: + /// + /// - `wants_flush`, `flush`, `wants_preprocess`, `preprocess` will **NOT** + /// have an effect on the ZK VM, only on the MPC VM. + /// - `call_raw` will **not** accept calls which have AND gates. + /// + /// This mode facilitates the ZK VM preprocessing in an + /// external context concurrently with the execution of the DEAP VM. + limited: bool, } impl Deap { @@ -56,6 +67,7 @@ impl Deap { follower_input_ranges: RangeSet::default(), follower_inputs: Vec::default(), outputs: Vec::default(), + limited: false, } } @@ -91,6 +103,11 @@ impl Deap { self.memory_map.try_get(value.to_raw()).map(T::from_raw) } + /// Sets the limited mode of operation. + pub fn limited(&mut self) { + self.limited = true; + } + #[cfg(test)] fn mpc(&self) -> MutexGuard<'_, Mpc> { self.mpc.try_lock().unwrap() @@ -301,6 +318,12 @@ where Zk: Vm, { fn call_raw(&mut self, call: Call) -> Result { + if self.limited && call.circ().and_count() > 0 { + return Err(VmError::call( + "calls with AND gates not allowed in limited mode", + )); + } + let (circ, inputs) = call.clone().into_parts(); let mut builder = Call::builder(circ); @@ -326,14 +349,29 @@ where Zk: Execute + Send + 'static, { fn wants_flush(&self) -> bool { - self.mpc.try_lock().unwrap().wants_flush() || self.zk.try_lock().unwrap().wants_flush() + if self.limited { + self.mpc.try_lock().unwrap().wants_flush() + } else { + self.mpc.try_lock().unwrap().wants_flush() || self.zk.try_lock().unwrap().wants_flush() + } } async fn flush(&mut self, ctx: &mut Context) -> Result<(), VmError> { - let mut zk = self.zk.clone().try_lock_owned().unwrap(); let mut mpc = self.mpc.clone().try_lock_owned().unwrap(); + + let zk = if self.limited { + None + } else { + Some(self.zk.clone().try_lock_owned().unwrap()) + }; + ctx.try_join( - async move |ctx| zk.flush(ctx).await, + async move |ctx| { + if let Some(mut zk) = zk { + zk.flush(ctx).await.unwrap(); + } + Ok(()) + }, async move |ctx| mpc.flush(ctx).await, ) .await @@ -343,15 +381,30 @@ where } fn wants_preprocess(&self) -> bool { - self.mpc.try_lock().unwrap().wants_preprocess() - || self.zk.try_lock().unwrap().wants_preprocess() + if self.limited { + self.mpc.try_lock().unwrap().wants_preprocess() + } else { + self.mpc.try_lock().unwrap().wants_preprocess() + || self.zk.try_lock().unwrap().wants_preprocess() + } } async fn preprocess(&mut self, ctx: &mut Context) -> Result<(), VmError> { - let mut zk = self.zk.clone().try_lock_owned().unwrap(); let mut mpc = self.mpc.clone().try_lock_owned().unwrap(); + + let zk = if self.limited { + None + } else { + Some(self.zk.clone().try_lock_owned().unwrap()) + }; + ctx.try_join( - async move |ctx| zk.preprocess(ctx).await, + async move |ctx| { + if let Some(mut zk) = zk { + zk.preprocess(ctx).await.unwrap(); + } + Ok(()) + }, async move |ctx| mpc.preprocess(ctx).await, ) .await diff --git a/crates/tlsn/src/mux.rs b/crates/tlsn/src/mux.rs index 1df078a4e7..982350e547 100644 --- a/crates/tlsn/src/mux.rs +++ b/crates/tlsn/src/mux.rs @@ -1,30 +1,51 @@ //! Multiplexer used in the TLSNotary protocol. -use std::future::IntoFuture; +use std::{error::Error, future::IntoFuture, mem}; use futures::{ - AsyncRead, AsyncWrite, Future, - future::{FusedFuture, FutureExt}, + AsyncRead, AsyncWrite, Future, TryFutureExt, + future::{FusedFuture, FutureExt, ready}, }; use tracing::error; use uid_mux::yamux; use crate::Role; +type BoxError = Box; + /// Multiplexer supporting unique deterministic stream IDs. pub(crate) type Mux = yamux::Yamux; /// 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> + Send + Unpin>, -); +pub(crate) struct MuxFuture { + /// The multiplexer future. + main: Box> + Send + Unpin>, + /// An auxiliary future which gets polled implicitly alongside the futures + /// polled with [Self::poll_with]. + /// + /// This future has no effect on the muxer progress. + aux: Box> + Send + Unpin>, + /// The error returned by the `aux` future. + aux_error: Option, +} impl MuxFuture { + /// Creates a new multiplexer future. + pub(crate) fn new( + fut: Box> + Send + Unpin>, + ) -> Self { + Self { + main: fut, + aux: Box::new(ready(Ok(())).fuse()), + aux_error: None, + } + } + /// Returns true if the muxer is complete. pub(crate) fn is_complete(&self) -> bool { - self.0.is_terminated() + self.main.is_terminated() } /// Awaits a future, polling the muxer future concurrently. @@ -33,15 +54,57 @@ impl MuxFuture { F: Future, { let mut fut = Box::pin(fut.fuse()); + // Poll the future concurrently with the muxer future. + // If the muxer returns an error or if the auxiliary future completes, + // continue polling the future until it completes. + loop { + futures::select! { + res = fut => return res, + res = &mut self.main => if let Err(e) = res { + error!("mux error: {:?}", e); + }, + res = &mut self.aux => if let Err(e) = res { + error!("aux future error: {:?}", e); + self.aux_error = Some(e); + }, + } + } + } + + /// Sets an auxiliary future. + pub(crate) fn aux(&mut self, fut: F) + where + F: Future> + Send + Unpin + 'static, + E: Error + Send + Sync + 'static, + { + self.aux = Box::new(fut.map_err(|e| -> BoxError { Box::new(e) }).fuse()); + self.aux_error = None; + } + + /// Awaits the auxiliary future, polling the muxer future concurrently. + pub(crate) async fn await_aux(&mut self) -> Result<(), BoxError> { + if self.aux.is_terminated() { + let err = mem::take(&mut self.aux_error); + return err.map_or(Ok(()), Err); + } + // 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 = &mut self.main => if let Err(e) = res { error!("mux error: {:?}", e); }, + res = &mut self.aux => { + match res { + Err(e) => { + error!("aux future error: {:?}", e); + return Err(e); + }, + _ => return Ok(()) + } + }, } } } @@ -54,7 +117,7 @@ impl Future for MuxFuture { mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> std::task::Poll { - self.0.as_mut().poll_unpin(cx) + self.main.as_mut().poll_unpin(cx) } } @@ -83,8 +146,8 @@ pub(crate) fn attach_mux( let ctrl = mux.control(); if let Role::Prover = role { - ctrl.alloc(32); + ctrl.alloc(36); } - (MuxFuture(Box::new(mux.into_future().fuse())), ctrl) + (MuxFuture::new(Box::new(mux.into_future().fuse())), ctrl) } diff --git a/crates/tlsn/src/prover.rs b/crates/tlsn/src/prover.rs index b6764d1d84..044af32ed2 100644 --- a/crates/tlsn/src/prover.rs +++ b/crates/tlsn/src/prover.rs @@ -10,7 +10,7 @@ pub use error::ProverError; pub use future::ProverFuture; pub use tlsn_core::{ProveConfig, ProveConfigBuilder, ProveConfigBuilderError, ProverOutput}; -use mpz_common::Context; +use mpz_common::{Context, Flush}; use mpz_core::Block; use mpz_garble_core::Delta; use mpz_vm_core::prelude::*; @@ -33,7 +33,7 @@ use futures::{AsyncRead, AsyncWrite, TryFutureExt}; use mpc_tls::{LeaderCtrl, MpcTlsLeader, SessionKeys}; use rand::Rng; use serio::{SinkExt, stream::IoStreamExt}; -use std::sync::Arc; +use std::{pin::Pin, sync::Arc}; use tls_client::{ClientConnection, ServerName as TlsServerName}; use tls_client_async::{TlsConnection, bind_client}; use tls_core::msgs::enums::ContentType; @@ -56,10 +56,12 @@ pub(crate) type RCOTSender = mpz_ot::rcot::shared::SharedRCOTSender< mpz_ot::kos::Sender, mpz_core::Block, >; -pub(crate) type RCOTReceiver = mpz_ot::rcot::shared::SharedRCOTReceiver< - mpz_ot::ferret::Receiver>, - bool, - mpz_core::Block, +pub(crate) type RCOTReceiver = mpz_ot::ferret::Receiver< + mpz_ot::rcot::shared::SharedRCOTReceiver< + mpz_ot::kos::Receiver, + bool, + mpz_core::Block, + >, >; pub(crate) type Mpc = mpz_garble::protocol::semihonest::Garbler>; @@ -104,17 +106,19 @@ impl Prover { let (mut mux_fut, mux_ctrl) = attach_mux(socket, Role::Prover); let mut mt = build_mt_context(mux_ctrl.clone()); let mut ctx = mux_fut.poll_with(mt.new_context()).await?; + // A context for preprocessing. + let prepr_ctx = mux_fut.poll_with(mt.new_context()).await?; // Sends protocol configuration to verifier for compatibility check. mux_fut .poll_with(ctx.io_mut().send(self.config.protocol_config().clone())) .await?; - let (vm, mut mpc_tls) = build_mpc_tls(&self.config, ctx); + let (vm, mut mpc_tls, prepr_fut) = build_mpc_tls(&self.config, ctx, prepr_ctx); // Allocate resources for MPC-TLS in the VM. let mut keys = mpc_tls.alloc()?; - let vm_lock = vm.try_lock().expect("VM is not locked"); + let mut vm_lock = vm.try_lock().expect("VM is not locked"); translate_keys(&mut keys, &vm_lock)?; // Allocate for committing to plaintext. @@ -132,9 +136,12 @@ impl Prover { self.config.protocol_config().max_recv_data(), )?; - drop(vm_lock); - debug!("setting up mpc-tls"); + // Changing the VM mode and setting a preprocessing future to have + // concurrent ZK VM preprocessing and MPC-TLS execution. + vm_lock.limited(); + mux_fut.aux(prepr_fut); + drop(vm_lock); mux_fut.poll_with(mpc_tls.preprocess()).await?; @@ -232,6 +239,10 @@ impl Prover { info!("finished MPC-TLS"); + // Only finalize once the ZK VM preprocessing future is + // complete. + mux_fut.await_aux().await.map_err(ProverError::zk)?; + { let mut vm = vm.try_lock().expect("VM should not be locked"); @@ -587,7 +598,15 @@ impl Prover { } } -fn build_mpc_tls(config: &ProverConfig, ctx: Context) -> (Arc>>, MpcTlsLeader) { +fn build_mpc_tls( + config: &ProverConfig, + ctx: Context, + mut prepr_ctx: Context, +) -> ( + Arc>>, + MpcTlsLeader, + Pin> + Send>>, +) { let mut rng = rand::rng(); let delta = Delta::new(Block::random(&mut rng)); @@ -598,19 +617,21 @@ fn build_mpc_tls(config: &ProverConfig, ctx: Context) -> (Arc (Arc (Arc>, - mpz_core::Block, +pub(crate) type RCOTSender = mpz_ot::ferret::Sender< + mpz_ot::rcot::shared::SharedRCOTSender< + mpz_ot::kos::Sender, + mpz_core::Block, + >, >; pub(crate) type RCOTReceiver = mpz_ot::rcot::shared::SharedRCOTReceiver< mpz_ot::kos::Receiver, @@ -98,6 +100,8 @@ impl Verifier { let (mut mux_fut, mux_ctrl) = attach_mux(socket, Role::Verifier); let mut mt = build_mt_context(mux_ctrl.clone()); let mut ctx = mux_fut.poll_with(mt.new_context()).await?; + // A context for preprocessing. + let prepr_ctx = mux_fut.poll_with(mt.new_context()).await?; // Receives protocol configuration from prover to perform compatibility check. let protocol_config = mux_fut @@ -112,11 +116,12 @@ impl Verifier { .await?; let delta = Delta::random(&mut rand::rng()); - let (vm, mut mpc_tls) = build_mpc_tls(&self.config, &protocol_config, delta, ctx); + let (vm, mut mpc_tls, prepr_fut) = + build_mpc_tls(&self.config, &protocol_config, delta, ctx, prepr_ctx); // Allocate resources for MPC-TLS in the VM. let mut keys = mpc_tls.alloc()?; - let vm_lock = vm.try_lock().expect("VM is not locked"); + let mut vm_lock = vm.try_lock().expect("VM is not locked"); translate_keys(&mut keys, &vm_lock)?; // Allocate for committing to plaintext. @@ -128,9 +133,12 @@ impl Verifier { zk_aes_ctr_recv.set_key(keys.server_write_key, keys.server_write_iv); zk_aes_ctr_recv.alloc(&mut *vm_lock.zk(), protocol_config.max_recv_data())?; - drop(vm_lock); - debug!("setting up mpc-tls"); + // Changing the VM mode and setting a preprocessing future to have + // concurrent ZK VM preprocessing and MPC-TLS execution. + vm_lock.limited(); + mux_fut.aux(prepr_fut); + drop(vm_lock); mux_fut.poll_with(mpc_tls.preprocess()).await?; @@ -250,6 +258,10 @@ impl Verifier { info!("finished MPC-TLS"); + // Only finalize once the ZK VM preprocessing future is + // complete. + mux_fut.await_aux().await.map_err(VerifierError::zk)?; + { let mut vm = vm.try_lock().expect("VM should not be locked"); @@ -611,45 +623,67 @@ fn build_mpc_tls( protocol_config: &ProtocolConfig, delta: Delta, ctx: Context, -) -> (Arc>>, MpcTlsFollower) { + mut preprocess_ctx: Context, +) -> ( + Arc>>, + MpcTlsFollower, + Pin> + Send>>, +) { let mut rng = rand::rng(); let base_ot_send = mpz_ot::chou_orlandi::Sender::default(); let base_ot_recv = mpz_ot::chou_orlandi::Receiver::default(); - let rcot_send = mpz_ot::kos::Sender::new( + let rcot_send_kos = mpz_ot::kos::Sender::new( mpz_ot::kos::SenderConfig::default(), delta.into_inner(), base_ot_recv, ); - let rcot_send = mpz_ot::ferret::Sender::new( + + let rcot_send_kos_shared = mpz_ot::rcot::shared::SharedRCOTSender::new(rcot_send_kos); + + let rcot_send_ferret = mpz_ot::ferret::Sender::new( mpz_ot::ferret::FerretConfig::builder() .lpn_type(mpz_ot::ferret::LpnType::Regular) .build() .expect("ferret config is valid"), Block::random(&mut rng), - rcot_send, + rcot_send_kos_shared.clone(), ); let rcot_recv = mpz_ot::kos::Receiver::new(mpz_ot::kos::ReceiverConfig::default(), base_ot_send); - let rcot_send = mpz_ot::rcot::shared::SharedRCOTSender::new(rcot_send); let rcot_recv = mpz_ot::rcot::shared::SharedRCOTReceiver::new(rcot_recv); let mpc = Mpc::new(mpz_ot::cot::DerandCOTReceiver::new(rcot_recv.clone())); - let zk = Zk::new(delta, rcot_send.clone()); + let zk = Zk::new(delta, rcot_send_ferret); + let zk_verifier_ot = zk.ot(); let vm = Arc::new(Mutex::new(Deap::new(tlsn_deap::Role::Follower, mpc, zk))); + // A preprocessing future which will be run concurrently with MPC-TLS. + // TODO: does this have to be pin ? + let prepr_fut = Box::pin(async move { + zk_verifier_ot + .try_lock() + .unwrap() + .flush(&mut preprocess_ctx) + .await + .map_err(|e| VerifierError::zk(e)) + }); + ( vm.clone(), MpcTlsFollower::new( config.build_mpc_tls_config(protocol_config), ctx, vm, - rcot_send, + // To minimize latency, the small amount of sent OTs needed by + // MPC-TLS will be provided by KOS rather than Ferret. + rcot_send_kos_shared, (rcot_recv.clone(), rcot_recv.clone(), rcot_recv), ), + prepr_fut, ) }