COT traits (#94)

* add 'owned' COT traits

* add 'shared' COT traits

* relocate ideal modules

* add ideal COT 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

@@ -3,5 +3,10 @@
 mod owned;
 mod shared;
 
-pub use owned::{ideal_ot_pair, IdealOTReceiver, IdealOTSender};
-pub use shared::{ideal_ot_shared_pair, IdealSharedOTReceiver, IdealSharedOTSender};
+pub use owned::{
+    ideal_cot_pair, ideal_ot_pair, IdealCOTReceiver, IdealCOTSender, IdealOTReceiver, IdealOTSender,
+};
+pub use shared::{
+    ideal_cot_shared_pair, ideal_ot_shared_pair, IdealSharedCOTReceiver, IdealSharedCOTSender,
+    IdealSharedOTReceiver, IdealSharedOTSender,
+};

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

@@ -0,0 +1,159 @@
+use crate::{COTReceiver, COTSender, OTError, OTSetup};
+use async_trait::async_trait;
+use futures::{channel::mpsc, StreamExt};
+use mpz_core::{Block, ProtocolMessage};
+use utils_aio::{sink::IoSink, stream::IoStream};
+
+/// Ideal OT sender.
+#[derive(Debug)]
+pub struct IdealCOTSender<T = Block> {
+    sender: mpsc::Sender<Vec<[T; 2]>>,
+    delta: Block,
+}
+
+/// Ideal OT receiver.
+#[derive(Debug)]
+pub struct IdealCOTReceiver<T = Block> {
+    receiver: mpsc::Receiver<Vec<[T; 2]>>,
+}
+
+impl<T> ProtocolMessage for IdealCOTSender<T> {
+    type Msg = ();
+}
+
+impl<T> ProtocolMessage for IdealCOTReceiver<T> {
+    type Msg = ();
+}
+
+/// Creates a pair of ideal COT sender and receiver.
+pub fn ideal_cot_pair<T: Send + Sync + 'static>(
+    delta: Block,
+) -> (IdealCOTSender<T>, IdealCOTReceiver<T>) {
+    let (sender, receiver) = mpsc::channel(10);
+
+    (
+        IdealCOTSender { sender, delta },
+        IdealCOTReceiver { receiver },
+    )
+}
+
+#[async_trait]
+impl<T> OTSetup for IdealCOTSender<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 COTSender<Block> for IdealCOTSender<Block> {
+    async fn send_correlated<Si: IoSink<()> + Send + Unpin, St: IoStream<()> + Send + Unpin>(
+        &mut self,
+        _sink: &mut Si,
+        _stream: &mut St,
+        msgs: &[Block],
+    ) -> Result<(), OTError> {
+        self.sender
+            .try_send(
+                msgs.iter()
+                    .map(|msg| [*msg, *msg ^ self.delta])
+                    .collect::<Vec<_>>(),
+            )
+            .expect("IdealCOTSender should be able to send");
+
+        Ok(())
+    }
+}
+
+#[async_trait]
+impl<T> OTSetup for IdealCOTReceiver<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 COTReceiver<bool, Block> for IdealCOTReceiver<Block> {
+    async fn receive_correlated<Si: IoSink<()> + Send + Unpin, St: IoStream<()> + Send + Unpin>(
+        &mut self,
+        _sink: &mut Si,
+        _stream: &mut St,
+        choices: &[bool],
+    ) -> Result<Vec<Block>, OTError> {
+        let payload = self
+            .receiver
+            .next()
+            .await
+            .expect("IdealCOTSender should send a value");
+
+        Ok(payload
+            .into_iter()
+            .zip(choices)
+            .map(|(v, c)| {
+                let [low, high] = v;
+                if *c {
+                    high
+                } else {
+                    low
+                }
+            })
+            .collect())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use itybity::IntoBits;
+    use rand::Rng;
+    use rand_chacha::ChaCha12Rng;
+    use rand_core::SeedableRng;
+    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_cot_owned() {
+        let mut rng = ChaCha12Rng::seed_from_u64(0);
+        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 values = Block::random_vec(&mut rng, 8);
+        let choices = rng.gen::<u8>().into_lsb0_vec();
+        let delta = Block::from([42u8; 16]);
+        let (mut sender, mut receiver) = ideal_cot_pair::<Block>(delta);
+
+        sender
+            .send_correlated(&mut send_sink, &mut send_stream, &values)
+            .await
+            .unwrap();
+
+        let received = receiver
+            .receive_correlated(&mut recv_sink, &mut recv_stream, &choices)
+            .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/owned/mod.rs

@@ -0,0 +1,5 @@
+mod cot;
+mod ot;
+
+pub use cot::{ideal_cot_pair, IdealCOTReceiver, IdealCOTSender};
+pub use ot::{ideal_ot_pair, IdealOTReceiver, IdealOTSender};

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

@@ -0,0 +1,137 @@
+use std::{
+    any::Any,
+    collections::HashMap,
+    sync::{Arc, Mutex},
+};
+
+use async_trait::async_trait;
+use futures::channel::oneshot;
+use mpz_core::Block;
+
+use crate::{COTReceiverShared, COTSenderShared, OTError};
+
+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 correlated ot sender and receiver pair.
+pub fn ideal_cot_shared_pair(delta: Block) -> (IdealSharedCOTSender, IdealSharedCOTReceiver) {
+    let sender_buffer = Arc::new(Mutex::new(HashMap::new()));
+    let receiver_buffer = Arc::new(Mutex::new(HashMap::new()));
+
+    let sender = IdealSharedCOTSender {
+        delta,
+        sender_buffer: sender_buffer.clone(),
+        receiver_buffer: receiver_buffer.clone(),
+    };
+
+    let receiver = IdealSharedCOTReceiver {
+        sender_buffer,
+        receiver_buffer,
+    };
+
+    (sender, receiver)
+}
+
+/// An ideal correlated oblivious transfer sender.
+#[derive(Clone, Debug)]
+pub struct IdealSharedCOTSender {
+    delta: Block,
+    sender_buffer: SenderBuffer,
+    receiver_buffer: ReceiverBuffer,
+}
+
+#[async_trait]
+impl COTSenderShared<Block> for IdealSharedCOTSender {
+    async fn send_correlated(&self, id: &str, msgs: &[Block]) -> Result<(), OTError> {
+        let msgs = Box::new(
+            msgs.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(())
+    }
+}
+
+/// An ideal correlated oblivious transfer receiver.
+#[derive(Clone, Debug)]
+pub struct IdealSharedCOTReceiver {
+    sender_buffer: SenderBuffer,
+    receiver_buffer: ReceiverBuffer,
+}
+
+#[async_trait]
+impl COTReceiverShared<bool, Block> for IdealSharedCOTReceiver {
+    async fn receive_correlated(&self, id: &str, choices: &[bool]) -> Result<Vec<Block>, OTError> {
+        if let Some(value) = self.sender_buffer.lock().unwrap().remove(id) {
+            let value = *value
+                .downcast::<Vec<[Block; 2]>>()
+                .expect("value type should be consistent");
+
+            return Ok(value
+                .into_iter()
+                .zip(choices)
+                .map(|(v, c)| v[*c as usize])
+                .collect::<Vec<_>>());
+        }
+
+        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");
+
+        Ok(values
+            .into_iter()
+            .zip(choices)
+            .map(|(v, c)| v[*c as usize])
+            .collect::<Vec<_>>())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use itybity::IntoBits;
+    use rand::Rng;
+    use rand_chacha::ChaCha12Rng;
+    use rand_core::SeedableRng;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn test_ideal_cot_shared() {
+        let mut rng = ChaCha12Rng::seed_from_u64(0);
+
+        let values = Block::random_vec(&mut rng, 8);
+        let choices = rng.gen::<u8>().into_lsb0_vec();
+        let delta = Block::from([42u8; 16]);
+        let (sender, receiver) = ideal_cot_shared_pair(delta);
+
+        sender.send_correlated("", &values).await.unwrap();
+
+        let received = receiver.receive_correlated("", &choices).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

@@ -0,0 +1,5 @@
+mod cot;
+mod ot;
+
+pub use cot::{ideal_cot_shared_pair, IdealSharedCOTReceiver, IdealSharedCOTSender};
+pub use ot::{ideal_ot_shared_pair, IdealSharedOTReceiver, IdealSharedOTSender};