Random correlated OT traits + ideal impls (#95)

* feat: random correlated OT traits + ideal impls

* Apply suggestions from code review

Co-authored-by: dan <themighty1@users.noreply.github.com>

---------

Co-authored-by: dan <themighty1@users.noreply.github.com>

ot/mpz-ot/src/ideal/mod.rs

@@ -4,9 +4,11 @@ mod owned;
 mod shared;
 
 pub use owned::{
-    ideal_cot_pair, ideal_ot_pair, IdealCOTReceiver, IdealCOTSender, IdealOTReceiver, IdealOTSender,
+    ideal_cot_pair, ideal_ot_pair, ideal_random_cot_pair, IdealCOTReceiver, IdealCOTSender,
+    IdealOTReceiver, IdealOTSender, IdealRandomCOTReceiver, IdealRandomCOTSender,
 };
 pub use shared::{
-    ideal_cot_shared_pair, ideal_ot_shared_pair, IdealSharedCOTReceiver, IdealSharedCOTSender,
-    IdealSharedOTReceiver, IdealSharedOTSender,
+    ideal_cot_shared_pair, ideal_ot_shared_pair, ideal_random_cot_shared_pair,
+    IdealSharedCOTReceiver, IdealSharedCOTSender, IdealSharedOTReceiver, IdealSharedOTSender,
+    IdealSharedRandomCOTReceiver, IdealSharedRandomCOTSender,
 };

ot/mpz-ot/src/ideal/owned/mod.rs

@@ -1,5 +1,7 @@
 mod cot;
 mod ot;
+mod rcot;
 
 pub use cot::{ideal_cot_pair, IdealCOTReceiver, IdealCOTSender};
 pub use ot::{ideal_ot_pair, IdealOTReceiver, IdealOTSender};
+pub use rcot::{ideal_random_cot_pair, IdealRandomCOTReceiver, IdealRandomCOTSender};

ot/mpz-ot/src/ideal/owned/rcot.rs

@@ -0,0 +1,181 @@
+use crate::{OTError, OTSetup, RandomCOTReceiver, RandomCOTSender};
+use async_trait::async_trait;
+use futures::{channel::mpsc, StreamExt};
+use mpz_core::{Block, ProtocolMessage};
+use rand::Rng;
+use rand_chacha::ChaCha12Rng;
+use rand_core::SeedableRng;
+use utils_aio::{sink::IoSink, stream::IoStream};
+
+/// Ideal random OT sender.
+#[derive(Debug)]
+pub struct IdealRandomCOTSender<T = Block> {
+    sender: mpsc::Sender<Vec<[T; 2]>>,
+    delta: Block,
+    rng: ChaCha12Rng,
+}
+
+/// Ideal random OT receiver.
+#[derive(Debug)]
+pub struct IdealRandomCOTReceiver<T = Block> {
+    receiver: mpsc::Receiver<Vec<[T; 2]>>,
+    rng: ChaCha12Rng,
+}
+
+impl<T> ProtocolMessage for IdealRandomCOTSender<T> {
+    type Msg = ();
+}
+
+impl<T> ProtocolMessage for IdealRandomCOTReceiver<T> {
+    type Msg = ();
+}
+
+/// Creates a pair of ideal random COT sender and receiver.
+pub fn ideal_random_cot_pair<T: Send + Sync + 'static>(
+    seed: [u8; 32],
+    delta: Block,
+) -> (IdealRandomCOTSender<T>, IdealRandomCOTReceiver<T>) {
+    let (sender, receiver) = mpsc::channel(10);
+
+    (
+        IdealRandomCOTSender {
+            sender,
+            delta,
+            rng: ChaCha12Rng::from_seed(seed),
+        },
+        IdealRandomCOTReceiver {
+            receiver,
+            rng: ChaCha12Rng::from_seed(seed),
+        },
+    )
+}
+
+#[async_trait]
+impl<T> OTSetup for IdealRandomCOTSender<T>
+where
+    T: Send + Sync,
+{
+    async fn setup<Si: IoSink<()> + Send + Unpin, St: IoStream<()> + Send + Unpin>(
+        &mut self,
+        _sink: &mut Si,
+        _stream: &mut St,
+    ) -> Result<(), OTError> {
+        Ok(())
+    }
+}
+
+#[async_trait]
+impl RandomCOTSender<Block> for IdealRandomCOTSender<Block> {
+    async fn send_random_correlated<
+        Si: IoSink<()> + Send + Unpin,
+        St: IoStream<()> + Send + Unpin,
+    >(
+        &mut self,
+        _sink: &mut Si,
+        _stream: &mut St,
+        count: usize,
+    ) -> Result<Vec<Block>, OTError> {
+        let low = (0..count)
+            .map(|_| Block::random(&mut self.rng))
+            .collect::<Vec<_>>();
+
+        self.sender
+            .try_send(
+                low.iter()
+                    .map(|msg| [*msg, *msg ^ self.delta])
+                    .collect::<Vec<_>>(),
+            )
+            .expect("IdealRandomCOTSender should be able to send");
+
+        Ok(low)
+    }
+}
+
+#[async_trait]
+impl<T> OTSetup for IdealRandomCOTReceiver<T>
+where
+    T: Send + Sync,
+{
+    async fn setup<Si: IoSink<()> + Send + Unpin, St: IoStream<()> + Send + Unpin>(
+        &mut self,
+        _sink: &mut Si,
+        _stream: &mut St,
+    ) -> Result<(), OTError> {
+        Ok(())
+    }
+}
+
+#[async_trait]
+impl RandomCOTReceiver<bool, Block> for IdealRandomCOTReceiver<Block> {
+    async fn receive_random_correlated<
+        Si: IoSink<()> + Send + Unpin,
+        St: IoStream<()> + Send + Unpin,
+    >(
+        &mut self,
+        _sink: &mut Si,
+        _stream: &mut St,
+        count: usize,
+    ) -> Result<(Vec<bool>, Vec<Block>), OTError> {
+        let payload = self
+            .receiver
+            .next()
+            .await
+            .expect("IdealRandomCOTSender should send a value");
+
+        assert_eq!(payload.len(), count);
+
+        let choices = (0..count).map(|_| self.rng.gen()).collect::<Vec<bool>>();
+        let payload = payload
+            .into_iter()
+            .zip(&choices)
+            .map(|(v, c)| {
+                let [low, high] = v;
+                if *c {
+                    high
+                } else {
+                    low
+                }
+            })
+            .collect();
+
+        Ok((choices, payload))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use utils_aio::duplex::MemoryDuplex;
+
+    use super::*;
+
+    // Test that the sender and receiver can be used to send and receive values
+    #[tokio::test]
+    async fn test_ideal_random_cot_owned() {
+        let seed = [0u8; 32];
+        let (send_channel, recv_channel) = MemoryDuplex::<()>::new();
+
+        let (mut send_sink, mut send_stream) = send_channel.split();
+        let (mut recv_sink, mut recv_stream) = recv_channel.split();
+
+        let delta = Block::from([42u8; 16]);
+        let (mut sender, mut receiver) = ideal_random_cot_pair::<Block>(seed, delta);
+
+        let values = sender
+            .send_random_correlated(&mut send_sink, &mut send_stream, 8)
+            .await
+            .unwrap();
+
+        let (choices, received) = receiver
+            .receive_random_correlated(&mut recv_sink, &mut recv_stream, 8)
+            .await
+            .unwrap();
+
+        let expected = values
+            .into_iter()
+            .zip(choices)
+            .map(|(v, c)| if c { v ^ delta } else { v })
+            .collect::<Vec<_>>();
+
+        assert_eq!(received, expected);
+    }
+}

ot/mpz-ot/src/ideal/shared/mod.rs

@@ -1,5 +1,9 @@
 mod cot;
 mod ot;
+mod rcot;
 
 pub use cot::{ideal_cot_shared_pair, IdealSharedCOTReceiver, IdealSharedCOTSender};
 pub use ot::{ideal_ot_shared_pair, IdealSharedOTReceiver, IdealSharedOTSender};
+pub use rcot::{
+    ideal_random_cot_shared_pair, IdealSharedRandomCOTReceiver, IdealSharedRandomCOTSender,
+};

ot/mpz-ot/src/ideal/shared/rcot.rs

@@ -0,0 +1,151 @@
+use std::{
+    any::Any,
+    collections::HashMap,
+    sync::{Arc, Mutex},
+};
+
+use async_trait::async_trait;
+use futures::channel::oneshot;
+use mpz_core::Block;
+use rand::Rng;
+use rand_chacha::ChaCha12Rng;
+use rand_core::SeedableRng;
+
+use crate::{OTError, RandomCOTReceiverShared, RandomCOTSenderShared};
+
+type SenderBuffer = Arc<Mutex<HashMap<String, Box<dyn Any + Send + 'static>>>>;
+type ReceiverBuffer = Arc<Mutex<HashMap<String, oneshot::Sender<Box<dyn Any + Send + 'static>>>>>;
+
+/// Creates an ideal random cot sender and receiver pair.
+pub fn ideal_random_cot_shared_pair(
+    seed: [u8; 32],
+    delta: Block,
+) -> (IdealSharedRandomCOTSender, IdealSharedRandomCOTReceiver) {
+    let sender_buffer: Arc<Mutex<HashMap<String, Box<dyn Any + Send>>>> =
+        Arc::new(Mutex::new(HashMap::new()));
+    let receiver_buffer = Arc::new(Mutex::new(HashMap::new()));
+
+    let sender = IdealSharedRandomCOTSender {
+        rng: Arc::new(Mutex::new(ChaCha12Rng::from_seed(seed))),
+        delta,
+        sender_buffer: sender_buffer.clone(),
+        receiver_buffer: receiver_buffer.clone(),
+    };
+
+    let receiver = IdealSharedRandomCOTReceiver {
+        rng: Arc::new(Mutex::new(ChaCha12Rng::from_seed(seed))),
+        sender_buffer,
+        receiver_buffer,
+    };
+
+    (sender, receiver)
+}
+
+/// An ideal random correlated oblivious transfer sender.
+#[derive(Clone, Debug)]
+pub struct IdealSharedRandomCOTSender {
+    delta: Block,
+    rng: Arc<Mutex<ChaCha12Rng>>,
+    sender_buffer: SenderBuffer,
+    receiver_buffer: ReceiverBuffer,
+}
+
+#[async_trait]
+impl RandomCOTSenderShared<Block> for IdealSharedRandomCOTSender {
+    async fn send_random_correlated(&self, id: &str, count: usize) -> Result<Vec<Block>, OTError> {
+        let low = Block::random_vec(&mut (*self.rng.lock().unwrap()), count);
+        let msgs = Box::new(
+            low.iter()
+                .map(|msg| [*msg, *msg ^ self.delta])
+                .collect::<Vec<_>>(),
+        );
+        if let Some(sender) = self.receiver_buffer.lock().unwrap().remove(id) {
+            sender
+                .send(msgs)
+                .expect("IdealCOTSenderControl should be able to send");
+        } else {
+            self.sender_buffer
+                .lock()
+                .unwrap()
+                .insert(id.to_string(), msgs);
+        }
+        Ok(low)
+    }
+}
+
+/// An ideal random correlated oblivious transfer receiver.
+#[derive(Clone, Debug)]
+pub struct IdealSharedRandomCOTReceiver {
+    rng: Arc<Mutex<ChaCha12Rng>>,
+    sender_buffer: SenderBuffer,
+    receiver_buffer: ReceiverBuffer,
+}
+
+#[async_trait]
+impl RandomCOTReceiverShared<bool, Block> for IdealSharedRandomCOTReceiver {
+    async fn receive_random_correlated(
+        &self,
+        id: &str,
+        count: usize,
+    ) -> Result<(Vec<bool>, Vec<Block>), OTError> {
+        let choices = (0..count)
+            .map(|_| (*self.rng.lock().unwrap()).gen())
+            .collect::<Vec<bool>>();
+        if let Some(value) = self.sender_buffer.lock().unwrap().remove(id) {
+            let values = *value
+                .downcast::<Vec<[Block; 2]>>()
+                .expect("value type should be consistent");
+
+            let value = values
+                .into_iter()
+                .zip(&choices)
+                .map(|([low, high], c)| if *c { high } else { low })
+                .collect::<Vec<_>>();
+
+            return Ok((choices, value));
+        }
+
+        let (sender, receiver) = oneshot::channel();
+        self.receiver_buffer
+            .lock()
+            .unwrap()
+            .insert(id.to_string(), sender);
+
+        let values = receiver.await.unwrap();
+
+        let values = *values
+            .downcast::<Vec<[Block; 2]>>()
+            .expect("value type should be consistent");
+
+        let values = values
+            .into_iter()
+            .zip(&choices)
+            .map(|([low, high], c)| if *c { high } else { low })
+            .collect::<Vec<_>>();
+
+        Ok((choices, values))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_ideal_random_cot_shared() {
+        let delta = Block::from([42u8; 16]);
+        let (sender, receiver) = ideal_random_cot_shared_pair([0u8; 32], delta);
+
+        let values = sender.send_random_correlated("", 8).await.unwrap();
+
+        let (choices, received) = receiver.receive_random_correlated("", 8).await.unwrap();
+
+        let expected = values
+            .into_iter()
+            .zip(choices)
+            .map(|(v, c)| if c { v ^ delta } else { v })
+            .collect::<Vec<_>>();
+
+        assert_eq!(received, expected);
+    }
+}