Add correlated OT

robinhundt · robinhundt · commit 183f6a5d4c53 · 2025-02-23T23:56:24.000+01:00
diff --git a/cryprot-core/src/buf.rs b/cryprot-core/src/buf.rs
@@ -14,6 +14,7 @@ use crate::alloc::{HugePageMemory, allocate_zeroed_vec};
 pub trait Buf<T>:
     Default + Debug + Deref<Target = [T]> + DerefMut + Send + Sync + 'static + private::Sealed
 {
+    type BufKind<E>: Buf<E> where E: Zeroable + Clone + Default + Debug + Send + Sync + 'static;
     /// Create a new `Buf` of length `len` with all elements set to zero.
     ///
     /// Implementations of this directly allocate zeroed memory and do not write
@@ -41,6 +42,8 @@ pub trait Buf<T>:
 }
 
 impl<T: Zeroable + Clone + Default + Debug + Send + Sync + 'static> Buf<T> for Vec<T> {
+    type BufKind<E> = Vec<E> where E: Zeroable + Clone + Default + Debug + Send + Sync + 'static;
+    
     fn zeroed(len: usize) -> Self {
         allocate_zeroed_vec(len)
     }
@@ -62,10 +65,12 @@ impl<T: Zeroable + Clone + Default + Debug + Send + Sync + 'static> Buf<T> for V
 
     fn grow_zeroed(&mut self, new_size: usize) {
         self.resize(new_size, T::zeroed());
-    }
+    }    
 }
 
 impl<T: Zeroable + Clone + Default + Debug + Send + Sync + 'static> Buf<T> for HugePageMemory<T> {
+    type BufKind<E> = HugePageMemory<E> where E: Zeroable + Clone + Default + Debug + Send + Sync + 'static;
+
     fn zeroed(len: usize) -> Self {
         HugePageMemory::zeroed(len)
     }
diff --git a/cryprot-ot/benches/bench.rs b/cryprot-ot/benches/bench.rs
@@ -7,17 +7,13 @@ use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
 use cryprot_core::{Block, alloc::HugePageMemory};
 use cryprot_net::testing::{init_bench_tracing, local_conn};
 use cryprot_ot::{
-    RotReceiver, RotSender,
-    base::SimplestOt,
-    extension::{
+    base::SimplestOt, extension::{
         MaliciousOtExtensionReceiver, MaliciousOtExtensionSender, SemiHonestOtExtensionReceiver,
         SemiHonestOtExtensionSender,
-    },
-    random_choices,
-    silent_ot::{
+    }, random_choices, silent_ot::{
         MaliciousSilentOtReceiver, MaliciousSilentOtSender, SemiHonestSilentOtReceiver,
         SemiHonestSilentOtSender,
-    },
+    }, CotReceiver, CotSender, RotReceiver, RotSender
 };
 use rand::{SeedableRng, rngs::StdRng};
 use tokio::runtime::{self, Runtime};
@@ -99,7 +95,7 @@ fn criterion_benchmark(c: &mut Criterion) {
                         }),
                         tokio::spawn(async move {
                             receiver
-                                .receive_into(&choices, &mut receiver_ots)
+                                .receive_into(&mut receiver_ots, &choices)
                                 .await
                                 .unwrap();
                             receiver_ots
@@ -113,64 +109,50 @@ fn criterion_benchmark(c: &mut Criterion) {
         })
     });
 
-    g.bench_function(format!("2 parallel 2**{p} extension OTs"), |b| {
+    g.bench_function(format!("2**{p} correlated extension OTs"), |b| {
         b.to_async(&rt).iter_custom(|iters| {
             let mut c11 = c1.sub_connection();
             let mut c22 = c2.sub_connection();
+
             async move {
                 let mut duration = Duration::ZERO;
+                let mut sender_ots = HugePageMemory::zeroed(count);
+                let mut receiver_ots = HugePageMemory::zeroed(count);
                 for _ in 0..iters {
-                    let (
-                        mut sender1,
-                        mut receiver1,
-                        mut sender2,
-                        mut receiver2,
-                        choices1,
-                        choices2,
-                    ) = {
+                    // setup not included in duration
+                    let (mut sender, mut receiver, choices) = {
                         let mut rng1 = StdRng::seed_from_u64(42);
-                        let mut rng2 = StdRng::seed_from_u64(42 * 42);
-                        let choices1 = random_choices(count, &mut rng1);
-                        let choices2 = random_choices(count, &mut rng2);
-                        let mut sender1 = SemiHonestOtExtensionSender::new_with_rng(
-                            c11.sub_connection(),
-                            rng1.clone(),
-                        );
-                        let mut receiver1 = SemiHonestOtExtensionReceiver::new_with_rng(
-                            c22.sub_connection(),
-                            rng2.clone(),
-                        );
-
-                        let mut sender2 =
+                        let rng2 = StdRng::seed_from_u64(42 * 42);
+                        let choices = random_choices(count, &mut rng1);
+                        let mut sender =
                             SemiHonestOtExtensionSender::new_with_rng(c11.sub_connection(), rng1);
-                        let mut receiver2 =
+                        let mut receiver =
                             SemiHonestOtExtensionReceiver::new_with_rng(c22.sub_connection(), rng2);
-
-                        tokio::try_join!(
-                            sender1.do_base_ots(),
-                            receiver1.do_base_ots(),
-                            sender2.do_base_ots(),
-                            receiver2.do_base_ots()
-                        )
-                        .unwrap();
-                        (sender1, receiver1, sender2, receiver2, choices1, choices2)
+                        tokio::try_join!(sender.do_base_ots(), receiver.do_base_ots()).unwrap();
+                        (sender, receiver, choices)
                     };
                     let now = Instant::now();
-                    let jh1 = tokio::spawn(async move { sender1.send(count).await });
-                    let jh2 = tokio::spawn(async move { receiver1.receive(&choices1).await });
-                    let jh3 = tokio::spawn(async move { sender2.send(count).await });
-                    let jh4 = tokio::spawn(async move { receiver2.receive(&choices2).await });
-                    let (ot1, ot2, ot3, ot4) = tokio::try_join!(jh1, jh2, jh3, jh4).unwrap();
+                    (sender_ots, receiver_ots) = tokio::try_join!(
+                        tokio::spawn(async move {
+                            sender.correlated_send_into(&mut sender_ots, |_, b| b ^ Block::ONES).await.unwrap();
+                            sender_ots
+                        }),
+                        tokio::spawn(async move {
+                            receiver
+                                .correlated_receive_into(&mut receiver_ots, &choices)
+                                .await
+                                .unwrap();
+                            receiver_ots
+                        })
+                    )
+                    .unwrap();
                     duration += now.elapsed();
-                    ot1.unwrap();
-                    ot2.unwrap();
-                    ot3.unwrap();
-                    ot4.unwrap();
                 }
                 duration
             }
         })
     });
+
     g.finish();
 
     let mut g = c.benchmark_group("malicious OT extension");
@@ -206,7 +188,7 @@ fn criterion_benchmark(c: &mut Criterion) {
                         }),
                         tokio::spawn(async move {
                             receiver
-                                .receive_into(&choices, &mut receiver_ots)
+                                .receive_into(&mut receiver_ots, &choices)
                                 .await
                                 .unwrap();
                             receiver_ots
diff --git a/cryprot-ot/src/adapter.rs b/cryprot-ot/src/adapter.rs
@@ -31,8 +31,8 @@ impl<R: RandChoiceRotReceiver> RotReceiver for ChosenChoice<R> {
 
     async fn receive_into(
         &mut self,
-        choices: &[subtle::Choice],
         ots: &mut impl cryprot_core::buf::Buf<cryprot_core::Block>,
+        choices: &[subtle::Choice],
     ) -> Result<(), Self::Error> {
         let mut rand_choices = self
             .0
diff --git a/cryprot-ot/src/base.rs b/cryprot-ot/src/base.rs
@@ -114,8 +114,8 @@ impl RotReceiver for SimplestOt {
     #[tracing::instrument(target = "cryprot_metrics", level = Level::TRACE, skip_all, fields(phase = phase::BASE_OT))]
     async fn receive_into(
         &mut self,
-        choices: &[Choice],
         ots: &mut impl Buf<Block>,
+        choices: &[Choice],
     ) -> Result<(), Self::Error> {
         assert_eq!(choices.len(), ots.len());
         let (mut send, mut recv) = self.conn.byte_stream().await?;
diff --git a/cryprot-ot/src/extension.rs b/cryprot-ot/src/extension.rs
@@ -36,8 +36,8 @@ use tokio::{
 use tracing::Level;
 
 use crate::{
-    Connected, Malicious, MaliciousMarker, RotReceiver, RotSender, Security, SemiHonest,
-    SemiHonestMarker,
+    Connected, CotReceiver, CotSender, Malicious, MaliciousMarker, RotReceiver, RotSender,
+    Security, SemiHonest, SemiHonestMarker,
     base::{self, SimplestOt},
     phase, random_choices,
 };
@@ -448,8 +448,8 @@ impl<S: Security> RotReceiver for OtExtensionReceiver<S> {
     #[tracing::instrument(target = "cryprot_metrics", level = Level::TRACE, skip_all, fields(phase = phase::OT_EXTENSION))]
     async fn receive_into(
         &mut self,
-        choices: &[Choice],
         ots: &mut impl Buf<Block>,
+        choices: &[Choice],
     ) -> Result<(), Self::Error> {
         assert_eq!(choices.len(), ots.len());
         assert_eq!(
@@ -621,6 +621,81 @@ impl<S: Security> RotReceiver for OtExtensionReceiver<S> {
     }
 }
 
+// should fit in one jumbo frame
+const COR_CHUNK_SIZE: usize = 8500 / Block::BYTES;
+
+impl<S: Security> CotSender for OtExtensionSender<S> {
+    type Error = Error;
+
+    async fn correlated_send_into<B, F>(
+        &mut self,
+        ots: &mut B,
+        mut correlation: F,
+    ) -> Result<(), Self::Error>
+    where
+        B: Buf<Block>,
+        F: FnMut(usize, Block) -> Block + Send,
+    {
+        let mut r_ots: B::BufKind<[Block; 2]> = B::BufKind::zeroed(ots.len());
+        self.send_into(&mut r_ots).await?;
+        let mut send_buf: Vec<Block> = Vec::zeroed(COR_CHUNK_SIZE);
+        let (mut tx, _) = self.conn.byte_stream().await?;
+        // TODO benchmark if this should use spawn_compute, I suspect not, since we're
+        // not doing that much work per byte sent
+        for (chunk_idx, (ot_chunk, rot_chunk)) in ots
+            .chunks_mut(send_buf.len())
+            .zip(r_ots.chunks(send_buf.len()))
+            .enumerate()
+        {
+            for (idx, ((ot, r_ot), correction)) in ot_chunk
+                .iter_mut()
+                .zip(rot_chunk)
+                .zip(&mut send_buf)
+                .enumerate()
+            {
+                *ot = r_ot[0];
+                *correction = r_ot[1] ^ correlation(chunk_idx * COR_CHUNK_SIZE + idx, r_ot[0]);
+            }
+            tx.write_all(bytemuck::must_cast_slice_mut(
+                &mut send_buf[..ot_chunk.len()],
+            ))
+            .await?;
+        }
+        Ok(())
+    }
+}
+
+impl<S: Security> CotReceiver for OtExtensionReceiver<S> {
+    type Error = Error;
+
+    async fn correlated_receive_into<B>(
+        &mut self,
+        ots: &mut B,
+        choices: &[Choice],
+    ) -> Result<(), Self::Error>
+    where
+        B: Buf<Block>,
+    {
+        self.receive_into(ots, choices).await?;
+        let mut recv_buf: Vec<Block> = Vec::zeroed(COR_CHUNK_SIZE);
+        let (_, mut rx) = self.conn.byte_stream().await?;
+        for (ot_chunk, choice_chunk) in ots
+            .chunks_mut(COR_CHUNK_SIZE)
+            .zip(choices.chunks(COR_CHUNK_SIZE))
+        {
+            rx.read_exact(bytemuck::must_cast_slice_mut(
+                &mut recv_buf[..ot_chunk.len()],
+            ))
+            .await?;
+            for ((ot, correction), choice) in ot_chunk.iter_mut().zip(&recv_buf).zip(choice_chunk) {
+                let use_correction = Block::conditional_select(&Block::ZERO, &Block::ONES, *choice);
+                *ot ^= use_correction & *correction;
+            }
+        }
+        Ok(())
+    }
+}
+
 fn commit(b: Block) -> random_oracle::Hash {
     random_oracle::hash(b.as_bytes())
 }
@@ -651,11 +726,12 @@ impl<T> From<tokio::sync::mpsc::error::SendError<T>> for Error {
 #[cfg(test)]
 mod tests {
 
+    use cryprot_core::Block;
     use cryprot_net::testing::{init_tracing, local_conn};
     use rand::{SeedableRng, rngs::StdRng};
 
     use crate::{
-        MaliciousMarker, RotReceiver, RotSender,
+        CotReceiver, CotSender, MaliciousMarker, RotReceiver, RotSender,
         extension::{
             DEFAULT_OT_BATCH_SIZE, OtExtensionReceiver, OtExtensionSender,
             SemiHonestOtExtensionReceiver, SemiHonestOtExtensionSender,
@@ -730,4 +806,28 @@ mod tests {
             assert_eq!(r, s[c.unwrap_u8() as usize]);
         }
     }
+
+    #[tokio::test]
+    async fn test_correlated_extension() {
+        let _g = init_tracing();
+        const COUNT: usize = 128;
+        let (c1, c2) = local_conn().await.unwrap();
+        let rng1 = StdRng::seed_from_u64(42);
+        let mut rng2 = StdRng::seed_from_u64(24);
+        let choices = random_choices(COUNT, &mut rng2);
+        let mut sender = SemiHonestOtExtensionSender::new_with_rng(c1, rng1);
+        let mut receiver = SemiHonestOtExtensionReceiver::new_with_rng(c2, rng2);
+        let (send_ots, recv_ots) = tokio::try_join!(
+            sender.correlated_send(COUNT, |_, b| b ^ Block::ONES),
+            receiver.correlated_receive(&choices)
+        )
+        .unwrap();
+        for (i, ((r, s), c)) in recv_ots.into_iter().zip(send_ots).zip(choices).enumerate() {
+            if bool::from(c) {
+                assert_eq!(r ^ Block::ONES, s, "Block {i}");
+            } else {
+                assert_eq!(r, s, "Block {i}")
+            }
+        }
+    }
 }
diff --git a/cryprot-ot/src/lib.rs b/cryprot-ot/src/lib.rs

Original file line number	Diff line number	Diff line change
`@@ -14,6 +14,7 @@ use crate::alloc::{HugePageMemory, allocate_zeroed_vec};`
`14`	`14`	`pub trait Buf<T>:`
`15`	`15`	`Default + Debug + Deref<Target = [T]> + DerefMut + Send + Sync + 'static + private::Sealed`
`16`	`16`	`{`
	`17`	`+ type BufKind<E>: Buf<E> where E: Zeroable + Clone + Default + Debug + Send + Sync + 'static;`
`17`	`18`	/// Create a new `Buf` of length `len` with all elements set to zero.
`18`	`19`	`///`
`19`	`20`	`/// Implementations of this directly allocate zeroed memory and do not write`
`@@ -41,6 +42,8 @@ pub trait Buf<T>:`
`41`	`42`	`}`
`42`	`43`
`43`	`44`	`impl<T: Zeroable + Clone + Default + Debug + Send + Sync + 'static> Buf<T> for Vec<T> {`
	`45`	`+ type BufKind<E> = Vec<E> where E: Zeroable + Clone + Default + Debug + Send + Sync + 'static;`
	`46`	`+`
`44`	`47`	`fn zeroed(len: usize) -> Self {`
`45`	`48`	`allocate_zeroed_vec(len)`
`46`	`49`	`}`
`@@ -62,10 +65,12 @@ impl<T: Zeroable + Clone + Default + Debug + Send + Sync + 'static> Buf<T> for V`
`62`	`65`
`63`	`66`	`fn grow_zeroed(&mut self, new_size: usize) {`
`64`	`67`	`self.resize(new_size, T::zeroed());`
`65`		`- }`
	`68`	`+ }`
`66`	`69`	`}`
`67`	`70`
`68`	`71`	`impl<T: Zeroable + Clone + Default + Debug + Send + Sync + 'static> Buf<T> for HugePageMemory<T> {`
	`72`	`+ type BufKind<E> = HugePageMemory<E> where E: Zeroable + Clone + Default + Debug + Send + Sync + 'static;`
	`73`	`+`
`69`	`74`	`fn zeroed(len: usize) -> Self {`
`70`	`75`	`HugePageMemory::zeroed(len)`
`71`	`76`	`}`