Add RandomOT traits and implement them for KOS15 (#91)

* Add RandomOT traits and implement them for KOS15

* Fix trait comments

* Add part of feedback

* Add test for KOS15 RandomOT

* Improve comments 1/2

Adhere to new wording for OT

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

* Improve comments 2/2

Adhere to new wording for OT.

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

---------

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

ot/mpz-ot-core/src/kos/receiver.rs

@@ -581,6 +581,11 @@ impl ReceiverKeys {
             })
             .collect())
     }
+
+    /// Returns the choices and the keys
+    pub fn take_choices_and_keys(self) -> (Vec<bool>, Vec<Block>) {
+        (self.choices, self.keys)
+    }
 }
 
 struct PayloadRecordNoDelta {

ot/mpz-ot-core/src/kos/sender.rs

@@ -448,6 +448,11 @@ impl SenderKeys {
             },
         })
     }
+
+    /// Returns the keys
+    pub fn take_keys(self) -> Vec<[Block; 2]> {
+        self.keys
+    }
 }
 
 /// The sender's state.

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

@@ -53,7 +53,7 @@ mod tests {
 
     use crate::{
         mock::{mock_ot_pair, MockOTReceiver, MockOTSender},
-        OTReceiver, OTSender, OTSetup, VerifiableOTReceiver,
+        OTReceiver, OTSender, OTSetup, RandomOTReceiver, RandomOTSender, VerifiableOTReceiver,
     };
 
     #[fixture]
@@ -146,6 +146,46 @@ mod tests {
         assert_eq!(received, expected);
     }
 
+    #[tokio::test]
+    async fn test_kos_random() {
+        let (sender_channel, receiver_channel) = MemoryDuplex::new();
+
+        let (mut sender_sink, mut sender_stream) = sender_channel.split();
+        let (mut receiver_sink, mut receiver_stream) = receiver_channel.split();
+
+        let (mut sender, mut receiver) = setup(
+            SenderConfig::default(),
+            ReceiverConfig::default(),
+            &mut sender_sink,
+            &mut sender_stream,
+            &mut receiver_sink,
+            &mut receiver_stream,
+            10,
+        )
+        .await;
+
+        let (sender_res, receiver_res) = tokio::join!(
+            RandomOTSender::send_random(&mut sender, &mut sender_sink, &mut sender_stream, 10),
+            RandomOTReceiver::receive_random(
+                &mut receiver,
+                &mut receiver_sink,
+                &mut receiver_stream,
+                10
+            )
+        );
+
+        let sender_output: Vec<[Block; 2]> = sender_res.unwrap();
+        let (choices, receiver_output): (Vec<bool>, Vec<Block>) = receiver_res.unwrap();
+
+        let expected = sender_output
+            .into_iter()
+            .zip(choices)
+            .map(|(output, choice)| output[choice as usize])
+            .collect::<Vec<_>>();
+
+        assert_eq!(receiver_output, expected);
+    }
+
     #[rstest]
     #[tokio::test]
     async fn test_kos_bytes(data: Vec<[Block; 2]>, choices: Vec<bool>) {

ot/mpz-ot/src/kos/receiver.rs

@@ -1,22 +1,25 @@
 use async_trait::async_trait;
 use futures::SinkExt;
 use itybity::{FromBitIterator, IntoBitIterator};
-use mpz_core::{cointoss, Block, ProtocolMessage};
+use mpz_core::{cointoss, prg::Prg, Block, ProtocolMessage};
 use mpz_ot_core::kos::{
     msgs::Message, receiver_state as state, Receiver as ReceiverCore, ReceiverConfig, CSP, SSP,
 };
 
 use enum_try_as_inner::EnumTryAsInner;
 use rand::{thread_rng, Rng};
+use rand_core::{RngCore, SeedableRng};
 use utils_aio::{
     non_blocking_backend::{Backend, NonBlockingBackend},
     sink::IoSink,
     stream::{ExpectStreamExt, IoStream},
 };
 
-use crate::{OTError, OTReceiver, OTSender, OTSetup, VerifiableOTReceiver, VerifiableOTSender};
-
 use super::{into_base_sink, into_base_stream, ReceiverError, ReceiverVerifyError};
+use crate::{
+    OTError, OTReceiver, OTSender, OTSetup, RandomOTReceiver, VerifiableOTReceiver,
+    VerifiableOTSender,
+};
 
 #[derive(Debug, EnumTryAsInner)]
 #[derive_err(Debug)]
@@ -303,6 +306,34 @@ where
     }
 }
 
+#[async_trait]
+impl<BaseOT> RandomOTReceiver<bool, Block> for Receiver<BaseOT>
+where
+    BaseOT: ProtocolMessage + Send,
+{
+    async fn receive_random<
+        Si: IoSink<Message<BaseOT::Msg>> + Send + Unpin,
+        St: IoStream<Message<BaseOT::Msg>> + Send + Unpin,
+    >(
+        &mut self,
+        _sink: &mut Si,
+        _stream: &mut St,
+        count: usize,
+    ) -> Result<(Vec<bool>, Vec<Block>), OTError> {
+        let receiver = self
+            .state
+            .try_as_extension_mut()
+            .map_err(ReceiverError::from)?;
+
+        let (choices, random_outputs) = receiver
+            .keys(count)
+            .map_err(ReceiverError::from)?
+            .take_choices_and_keys();
+
+        Ok((choices, random_outputs))
+    }
+}
+
 #[async_trait]
 impl<const N: usize, BaseOT> OTReceiver<bool, [u8; N]> for Receiver<BaseOT>
 where
@@ -350,6 +381,45 @@ where
     }
 }
 
+#[async_trait]
+impl<const N: usize, BaseOT> RandomOTReceiver<bool, [u8; N]> for Receiver<BaseOT>
+where
+    BaseOT: ProtocolMessage + Send,
+{
+    async fn receive_random<
+        Si: IoSink<Message<BaseOT::Msg>> + Send + Unpin,
+        St: IoStream<Message<BaseOT::Msg>> + Send + Unpin,
+    >(
+        &mut self,
+        _sink: &mut Si,
+        _stream: &mut St,
+        count: usize,
+    ) -> Result<(Vec<bool>, Vec<[u8; N]>), OTError> {
+        let receiver = self
+            .state
+            .try_as_extension_mut()
+            .map_err(ReceiverError::from)?;
+
+        let (choices, random_outputs) = receiver
+            .keys(count)
+            .map_err(ReceiverError::from)?
+            .take_choices_and_keys();
+
+        Ok((
+            choices,
+            random_outputs
+                .into_iter()
+                .map(|block| {
+                    let mut prg = Prg::from_seed(block);
+                    let mut out = [0_u8; N];
+                    prg.fill_bytes(&mut out);
+                    out
+                })
+                .collect(),
+        ))
+    }
+}
+
 #[async_trait]
 impl<BaseOT> VerifiableOTReceiver<bool, Block, [Block; 2]> for Receiver<BaseOT>
 where

ot/mpz-ot/src/kos/sender.rs

@@ -1,25 +1,24 @@
-use crate::{
-    kos::SenderError, CommittedOTReceiver, CommittedOTSender, OTError, OTReceiver, OTSender,
-    OTSetup,
-};
-
 use async_trait::async_trait;
+use enum_try_as_inner::EnumTryAsInner;
 use futures_util::SinkExt;
 use itybity::IntoBits;
-use mpz_core::{cointoss, Block, ProtocolMessage};
+use mpz_core::{cointoss, prg::Prg, Block, ProtocolMessage};
 use mpz_ot_core::kos::{
     msgs::Message, sender_state as state, Sender as SenderCore, SenderConfig, CSP, SSP,
 };
 use rand::{thread_rng, Rng};
+use rand_core::{RngCore, SeedableRng};
 use utils_aio::{
     non_blocking_backend::{Backend, NonBlockingBackend},
     sink::IoSink,
     stream::{ExpectStreamExt, IoStream},
 };
 
-use enum_try_as_inner::EnumTryAsInner;
-
 use super::{into_base_sink, into_base_stream};
+use crate::{
+    kos::SenderError, CommittedOTReceiver, CommittedOTSender, OTError, OTReceiver, OTSender,
+    OTSetup, RandomOTSender,
+};
 
 #[derive(Debug, EnumTryAsInner)]
 #[derive_err(Debug)]
@@ -336,6 +335,30 @@ where
     }
 }
 
+#[async_trait]
+impl<BaseOT> RandomOTSender<[Block; 2]> for Sender<BaseOT>
+where
+    BaseOT: ProtocolMessage + Send,
+{
+    async fn send_random<
+        Si: IoSink<Message<BaseOT::Msg>> + Send + Unpin,
+        St: IoStream<Message<BaseOT::Msg>> + Send + Unpin,
+    >(
+        &mut self,
+        _sink: &mut Si,
+        _stream: &mut St,
+        count: usize,
+    ) -> Result<Vec<[Block; 2]>, OTError> {
+        let sender = self
+            .state
+            .try_as_extension_mut()
+            .map_err(SenderError::from)?;
+
+        let random_outputs = sender.keys(count).map_err(SenderError::from)?;
+        Ok(random_outputs.take_keys())
+    }
+}
+
 #[async_trait]
 impl<const N: usize, BaseOT> OTSender<[[u8; N]; 2]> for Sender<BaseOT>
 where
@@ -375,6 +398,42 @@ where
     }
 }
 
+#[async_trait]
+impl<const N: usize, BaseOT> RandomOTSender<[[u8; N]; 2]> for Sender<BaseOT>
+where
+    BaseOT: ProtocolMessage + Send,
+{
+    async fn send_random<
+        Si: IoSink<Message<BaseOT::Msg>> + Send + Unpin,
+        St: IoStream<Message<BaseOT::Msg>> + Send + Unpin,
+    >(
+        &mut self,
+        _sink: &mut Si,
+        _stream: &mut St,
+        count: usize,
+    ) -> Result<Vec<[[u8; N]; 2]>, OTError> {
+        let sender = self
+            .state
+            .try_as_extension_mut()
+            .map_err(SenderError::from)?;
+
+        let random_outputs = sender.keys(count).map_err(SenderError::from)?;
+
+        let prng = |block| {
+            let mut prg = Prg::from_seed(block);
+            let mut out = [0_u8; N];
+            prg.fill_bytes(&mut out);
+            out
+        };
+
+        Ok(random_outputs
+            .take_keys()
+            .into_iter()
+            .map(|[a, b]| [prng(a), prng(b)])
+            .collect())
+    }
+}
+
 #[async_trait]
 impl<BaseOT> CommittedOTSender<[Block; 2]> for Sender<BaseOT>
 where

ot/mpz-ot/src/lib.rs

@@ -68,6 +68,30 @@ where
     ) -> Result<(), OTError>;
 }
 
+/// A random OT sender.
+#[async_trait]
+pub trait RandomOTSender<T>: ProtocolMessage
+where
+    T: Send + Sync,
+{
+    /// Outputs pairs of random messages.
+    ///
+    /// # Arguments
+    ///
+    /// * `sink` - The IO sink to the receiver.
+    /// * `stream` - The IO stream from the receiver.
+    /// * `count` - The number of pairs of random messages to output.
+    async fn send_random<
+        Si: IoSink<Self::Msg> + Send + Unpin,
+        St: IoStream<Self::Msg> + Send + Unpin,
+    >(
+        &mut self,
+        sink: &mut Si,
+        stream: &mut St,
+        count: usize,
+    ) -> Result<Vec<T>, OTError>;
+}
+
 /// An oblivious transfer receiver.
 #[async_trait]
 pub trait OTReceiver<T, U>: ProtocolMessage
@@ -90,6 +114,31 @@ where
     ) -> Result<Vec<U>, OTError>;
 }
 
+/// A random OT receiver.
+#[async_trait]
+pub trait RandomOTReceiver<T, U>: ProtocolMessage
+where
+    T: Send + Sync,
+    U: Send + Sync,
+{
+    /// Outputs the choice bits and the corresponding messages.
+    ///
+    /// # Arguments
+    ///
+    /// * `sink` - The IO sink to the sender.
+    /// * `stream` - The IO stream from the sender.
+    /// * `count` - The number of random messages to receive.
+    async fn receive_random<
+        Si: IoSink<Self::Msg> + Send + Unpin,
+        St: IoStream<Self::Msg> + Send + Unpin,
+    >(
+        &mut self,
+        sink: &mut Si,
+        stream: &mut St,
+        count: usize,
+    ) -> Result<(Vec<T>, Vec<U>), OTError>;
+}
+
 /// An oblivious transfer sender that is committed to its messages and can reveal them
 /// to the receiver to verify them.
 #[async_trait]