perf: pre-garble key exchange and PRF (#371)

* prf pre-garble

* fix

* update mpz version to 7669232

* fix ValueId dependency

* PR feedback

* bump mpz to 1ac6779

components/aead/Cargo.toml

@@ -25,8 +25,8 @@ tracing = [
 tlsn-block-cipher = { path = "../cipher/block-cipher" }
 tlsn-stream-cipher = { path = "../cipher/stream-cipher" }
 tlsn-universal-hash = { path = "../universal-hash" }
-mpz-core = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "1f2c922" }
-mpz-garble = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "1f2c922" }
+mpz-core = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "1ac6779" }
+mpz-garble = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "1ac6779" }
 tlsn-utils-aio = { git = "https://github.com/tlsnotary/tlsn-utils", rev = "8d8ffe1" }
 
 async-trait = "0.1"

components/aead/src/aes_gcm/mod.rs

@@ -17,7 +17,7 @@ use futures::{SinkExt, StreamExt, TryFutureExt};
 
 use block_cipher::{Aes128, BlockCipher};
 use mpz_core::commit::HashCommit;
-use mpz_garble::ValueRef;
+use mpz_garble::value::ValueRef;
 use tlsn_stream_cipher::{Aes128Ctr, StreamCipher};
 use tlsn_universal_hash::UniversalHash;
 use utils_aio::expect_msg_or_err;
@@ -365,15 +365,25 @@ mod tests {
 
         let leader_thread = leader_vm.new_thread("test_thread").await.unwrap();
         let leader_key = leader_thread
-            .new_public_array_input("key", key.clone())
+            .new_public_array_input::<u8>("key", key.len())
             .unwrap();
         let leader_iv = leader_thread
-            .new_public_array_input("iv", iv.clone())
+            .new_public_array_input::<u8>("iv", iv.len())
             .unwrap();
 
+        leader_thread.assign(&leader_key, key.clone()).unwrap();
+        leader_thread.assign(&leader_iv, iv.clone()).unwrap();
+
         let follower_thread = follower_vm.new_thread("test_thread").await.unwrap();
-        let follower_key = follower_thread.new_public_array_input("key", key).unwrap();
-        let follower_iv = follower_thread.new_public_array_input("iv", iv).unwrap();
+        let follower_key = follower_thread
+            .new_public_array_input::<u8>("key", key.len())
+            .unwrap();
+        let follower_iv = follower_thread
+            .new_public_array_input::<u8>("iv", iv.len())
+            .unwrap();
+
+        follower_thread.assign(&follower_key, key.clone()).unwrap();
+        follower_thread.assign(&follower_iv, iv.clone()).unwrap();
 
         let leader_config = AesGcmConfigBuilder::default()
             .id("test".to_string())

components/aead/src/lib.rs

@@ -19,7 +19,7 @@ pub use msg::AeadMessage;
 
 use async_trait::async_trait;
 
-use mpz_garble::ValueRef;
+use mpz_garble::value::ValueRef;
 use utils_aio::duplex::Duplex;
 
 /// A channel for sending and receiving AEAD messages.

components/cipher/Cargo.toml

@@ -4,8 +4,8 @@ resolver = "2"
 
 [workspace.dependencies]
 # tlsn
-mpz-circuits = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "1f2c922" }
-mpz-garble = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "1f2c922" }
+mpz-circuits = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "1ac6779" }
+mpz-garble = { git = "https://github.com/privacy-scaling-explorations/mpz", rev = "1ac6779" }
 tlsn-utils = { git = "https://github.com/tlsnotary/tlsn-utils", rev = "8d8ffe1" }
 
 # crypto

components/cipher/block-cipher/src/cipher.rs

@@ -2,7 +2,7 @@ use std::marker::PhantomData;
 
 use async_trait::async_trait;
 
-use mpz_garble::{Decode, DecodePrivate, Execute, Memory, ValueRef};
+use mpz_garble::{value::ValueRef, Decode, DecodePrivate, Execute, Memory};
 use utils::id::NestedId;
 
 use crate::{BlockCipher, BlockCipherCircuit, BlockCipherConfig, BlockCipherError};
@@ -80,11 +80,13 @@ where
 
         let msg = self
             .executor
-            .new_private_input::<C::BLOCK>(&format!("{}/msg", &id), Some(block))?;
+            .new_private_input::<C::BLOCK>(&format!("{}/msg", &id))?;
         let ciphertext = self
             .executor
             .new_output::<C::BLOCK>(&format!("{}/ciphertext", &id))?;
 
+        self.executor.assign(&msg, block)?;
+
         self.executor
             .execute(C::circuit(), &[key, msg], &[ciphertext.clone()])
             .await?;
@@ -115,7 +117,7 @@ where
 
         let msg = self
             .executor
-            .new_private_input::<C::BLOCK>(&format!("{}/msg", &id), None)?;
+            .new_blind_input::<C::BLOCK>(&format!("{}/msg", &id))?;
         let ciphertext = self
             .executor
             .new_output::<C::BLOCK>(&format!("{}/ciphertext", &id))?;
@@ -155,11 +157,13 @@ where
 
         let msg = self
             .executor
-            .new_public_input::<C::BLOCK>(&format!("{}/msg", &id), block)?;
+            .new_public_input::<C::BLOCK>(&format!("{}/msg", &id))?;
         let ciphertext = self
             .executor
             .new_output::<C::BLOCK>(&format!("{}/ciphertext", &id))?;
 
+        self.executor.assign(&msg, block)?;
+
         self.executor
             .execute(C::circuit(), &[key, msg], &[ciphertext.clone()])
             .await?;

components/cipher/block-cipher/src/lib.rs

@@ -12,7 +12,7 @@ mod config;
 
 use async_trait::async_trait;
 
-use mpz_garble::ValueRef;
+use mpz_garble::value::ValueRef;
 
 pub use crate::{
     cipher::MpcBlockCipher,
@@ -96,8 +96,11 @@ mod tests {
         let follower_thread = follower_vm.new_thread("test").await.unwrap();
 
         // Key is public just for this test, typically it is private
-        let leader_key = leader_thread.new_public_input("key", key).unwrap();
-        let follower_key = follower_thread.new_public_input("key", key).unwrap();
+        let leader_key = leader_thread.new_public_input::<[u8; 16]>("key").unwrap();
+        let follower_key = follower_thread.new_public_input::<[u8; 16]>("key").unwrap();
+
+        leader_thread.assign(&leader_key, key).unwrap();
+        follower_thread.assign(&follower_key, key).unwrap();
 
         let mut leader = MpcBlockCipher::<Aes128, _>::new(leader_config, leader_thread);
         leader.set_key(leader_key);
@@ -131,8 +134,11 @@ mod tests {
         let follower_thread = follower_vm.new_thread("test").await.unwrap();
 
         // Key is public just for this test, typically it is private
-        let leader_key = leader_thread.new_public_input("key", key).unwrap();
-        let follower_key = follower_thread.new_public_input("key", key).unwrap();
+        let leader_key = leader_thread.new_public_input::<[u8; 16]>("key").unwrap();
+        let follower_key = follower_thread.new_public_input::<[u8; 16]>("key").unwrap();
+
+        leader_thread.assign(&leader_key, key).unwrap();
+        follower_thread.assign(&follower_key, key).unwrap();
 
         let mut leader = MpcBlockCipher::<Aes128, _>::new(leader_config, leader_thread);
         leader.set_key(leader_key);

components/cipher/stream-cipher/benches/mock.rs

@@ -1,18 +1,24 @@
-use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};
+use criterion::{criterion_group, criterion_main, Criterion, Throughput};
 
 use mpz_garble::{protocol::deap::mock::create_mock_deap_vm, Memory, Vm};
 use tlsn_stream_cipher::{Aes128Ctr, MpcStreamCipher, StreamCipher, StreamCipherConfigBuilder};
 
-async fn bench_stream_cipher_public_encrypt(thread_count: usize, len: usize) {
+async fn bench_stream_cipher_encrypt(thread_count: usize, len: usize) {
     let (mut leader_vm, mut follower_vm) = create_mock_deap_vm("test").await;
 
     let leader_thread = leader_vm.new_thread("key_config").await.unwrap();
-    let leader_key = leader_thread.new_public_input("key", [0u8; 16]).unwrap();
-    let leader_iv = leader_thread.new_public_input("iv", [0u8; 4]).unwrap();
+    let leader_key = leader_thread.new_public_input::<[u8; 16]>("key").unwrap();
+    let leader_iv = leader_thread.new_public_input::<[u8; 4]>("iv").unwrap();
+
+    leader_thread.assign(&leader_key, [0u8; 16]).unwrap();
+    leader_thread.assign(&leader_iv, [0u8; 4]).unwrap();
 
     let follower_thread = follower_vm.new_thread("key_config").await.unwrap();
-    let follower_key = follower_thread.new_public_input("key", [0u8; 16]).unwrap();
-    let follower_iv = follower_thread.new_public_input("iv", [0u8; 4]).unwrap();
+    let follower_key = follower_thread.new_public_input::<[u8; 16]>("key").unwrap();
+    let follower_iv = follower_thread.new_public_input::<[u8; 4]>("iv").unwrap();
+
+    follower_thread.assign(&follower_key, [0u8; 16]).unwrap();
+    follower_thread.assign(&follower_iv, [0u8; 4]).unwrap();
 
     let leader_thread_pool = leader_vm
         .new_thread_pool("mock", thread_count)
@@ -60,9 +66,8 @@ fn criterion_benchmark(c: &mut Criterion) {
 
     group.throughput(Throughput::Bytes(len as u64));
     group.bench_function(format!("{}", len), |b| {
-        b.to_async(&rt).iter(|| async {
-            black_box(bench_stream_cipher_public_encrypt(thread_count, len).await)
-        })
+        b.to_async(&rt)
+            .iter(|| async { bench_stream_cipher_encrypt(thread_count, len).await })
     });
 
     drop(group);

components/cipher/stream-cipher/src/config.rs

@@ -1,7 +1,7 @@
 use std::marker::PhantomData;
 
 use derive_builder::Builder;
-use mpz_garble::ValueRef;
+use mpz_garble::value::ValueRef;
 use std::fmt::Debug;
 
 use crate::CtrCircuit;

components/cipher/stream-cipher/src/lib.rs

@@ -24,7 +24,7 @@ pub use config::{StreamCipherConfig, StreamCipherConfigBuilder, StreamCipherConf
 pub use stream_cipher::MpcStreamCipher;
 
 use async_trait::async_trait;
-use mpz_garble::ValueRef;
+use mpz_garble::value::ValueRef;
 
 /// Error that can occur when using a stream cipher
 #[derive(Debug, thiserror::Error)]
@@ -208,12 +208,18 @@ mod tests {
         let (mut leader_vm, mut follower_vm) = create_mock_deap_vm("test").await;
 
         let leader_thread = leader_vm.new_thread("key_config").await.unwrap();
-        let leader_key = leader_thread.new_public_input("key", key).unwrap();
-        let leader_iv = leader_thread.new_public_input("iv", iv).unwrap();
+        let leader_key = leader_thread.new_public_input::<[u8; 16]>("key").unwrap();
+        let leader_iv = leader_thread.new_public_input::<[u8; 4]>("iv").unwrap();
+
+        leader_thread.assign(&leader_key, key).unwrap();
+        leader_thread.assign(&leader_iv, iv).unwrap();
 
         let follower_thread = follower_vm.new_thread("key_config").await.unwrap();
-        let follower_key = follower_thread.new_public_input("key", key).unwrap();
-        let follower_iv = follower_thread.new_public_input("iv", iv).unwrap();
+        let follower_key = follower_thread.new_public_input::<[u8; 16]>("key").unwrap();
+        let follower_iv = follower_thread.new_public_input::<[u8; 4]>("iv").unwrap();
+
+        follower_thread.assign(&follower_key, key).unwrap();
+        follower_thread.assign(&follower_iv, iv).unwrap();
 
         let leader_thread_pool = leader_vm
             .new_thread_pool("mock", thread_count)