network: mock: Add mock that allows unit testing MPC circuits

Author: joeykraut

src/algebra/authenticated_scalar.rs

@@ -664,7 +664,7 @@ impl Sub<&AuthenticatedScalarResult> for &AuthenticatedScalarResult {
         AuthenticatedScalarResult {
             share: &self.share - &rhs.share,
             mac: &self.mac - &rhs.mac,
-            public_modifier: self.public_modifier.clone(),
+            public_modifier: self.public_modifier.clone() - rhs.public_modifier.clone(),
         }
     }
 }
@@ -1029,3 +1029,23 @@ pub mod test_helpers {
         val.public_modifier = val.fabric().allocate_scalar(new_value)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{algebra::scalar::Scalar, test_helpers::execute_mock_mpc};
+
+    /// Test a simple `XOR` circuit
+    #[tokio::test]
+    async fn test_xor_circuit() {
+        let (res, _) = execute_mock_mpc(|fabric| async move {
+            let a = &fabric.zero_authenticated();
+            let b = &fabric.zero_authenticated();
+            let res = a + b - Scalar::from(2u64) * a * b;
+
+            res.open_authenticated().await
+        })
+        .await;
+
+        assert_eq!(res.unwrap(), 0.into());
+    }
+}

src/beaver.rs

@@ -56,37 +56,46 @@ pub trait SharedValueSource: Send + Sync {
         (a_vals, b_vals, c_vals)
     }
 }
-
-/// A dummy value source that outputs only ones
-/// Used for testing
+/// An implementation of a beaver value source that returns
+/// beaver triples (0, 0, 0) for party 0 and (1, 1, 1) for party 1
 #[cfg(any(feature = "test_helpers", test))]
 #[derive(Clone, Debug, Default)]
-pub struct DummySharedScalarSource;
+pub struct PartyIDBeaverSource {
+    /// The ID of the local party
+    party_id: u64,
+}
 
 #[cfg(any(feature = "test_helpers", test))]
-#[allow(dead_code)]
-impl DummySharedScalarSource {
-    /// Constructor
-    pub fn new() -> Self {
-        Self
+impl PartyIDBeaverSource {
+    /// Create a new beaver source given the local party_id
+    pub fn new(party_id: u64) -> Self {
+        Self { party_id }
     }
 }
 
+/// The PartyIDBeaverSource returns beaver triplets split statically between the
+/// parties. We assume a = 2, b = 3 ==> c = 6. [a] = (1, 1); [b] = (3, 0) [c] = (2, 4)
 #[cfg(any(feature = "test_helpers", test))]
-impl SharedValueSource for DummySharedScalarSource {
+impl SharedValueSource for PartyIDBeaverSource {
     fn next_shared_bit(&mut self) -> Scalar {
-        Scalar::one()
+        // Simply output partyID, assume partyID \in {0, 1}
+        assert!(self.party_id == 0 || self.party_id == 1);
+        Scalar::from(self.party_id)
     }
 
-    fn next_shared_value(&mut self) -> Scalar {
-        Scalar::one()
+    fn next_triplet(&mut self) -> (Scalar, Scalar, Scalar) {
+        if self.party_id == 0 {
+            (Scalar::from(1u64), Scalar::from(3u64), Scalar::from(2u64))
+        } else {
+            (Scalar::from(1u64), Scalar::from(0u64), Scalar::from(4u64))
+        }
     }
 
     fn next_shared_inverse_pair(&mut self) -> (Scalar, Scalar) {
-        (Scalar::one(), Scalar::one())
+        (Scalar::from(self.party_id), Scalar::from(self.party_id))
     }
 
-    fn next_triplet(&mut self) -> (Scalar, Scalar, Scalar) {
-        (Scalar::one(), Scalar::one(), Scalar::one())
+    fn next_shared_value(&mut self) -> Scalar {
+        Scalar::from(self.party_id)
     }
 }

src/lib.rs

@@ -67,13 +67,57 @@ pub type BeaverSource<S: SharedValueSource> = Rc<RefCell<S>>;
 
 #[cfg(test)]
 pub(crate) mod test_helpers {
-    use crate::{beaver::DummySharedScalarSource, network::NoRecvNetwork, MpcFabric};
+    use futures::Future;
+
+    use crate::{
+        beaver::PartyIDBeaverSource,
+        network::{MockNetwork, NoRecvNetwork, UnboundedDuplexStream},
+        MpcFabric, PARTY0, PARTY1,
+    };
 
     /// Create a mock fabric
     pub fn mock_fabric() -> MpcFabric {
         let network = NoRecvNetwork::default();
-        let beaver_source = DummySharedScalarSource::new();
+        let beaver_source = PartyIDBeaverSource::default();
 
         MpcFabric::new(network, beaver_source)
     }
+
+    /// Run a mock MPC connected by a duplex stream as the mock network
+    ///
+    /// This will spawn two tasks to execute either side of the MPC
+    ///
+    /// Returns the outputs of both parties
+    pub async fn execute_mock_mpc<T, S, F>(mut f: F) -> (T, T)
+    where
+        T: Send + 'static,
+        S: Future<Output = T> + Send + 'static,
+        F: FnMut(MpcFabric) -> S,
+    {
+        // Build a duplex stream to broker communication between the two parties
+        let (party0_stream, party1_stream) = UnboundedDuplexStream::new_duplex_pair();
+        let party0_fabric = MpcFabric::new(
+            MockNetwork::new(PARTY0, party0_stream),
+            PartyIDBeaverSource::new(PARTY0),
+        );
+        let party1_fabric = MpcFabric::new(
+            MockNetwork::new(PARTY1, party1_stream),
+            PartyIDBeaverSource::new(PARTY1),
+        );
+
+        // Spawn two tasks to execute the MPC
+        let fabric0 = party0_fabric.clone();
+        let fabric1 = party1_fabric.clone();
+        let party0_task = tokio::spawn(f(fabric0));
+        let party1_task = tokio::spawn(f(fabric1));
+
+        let party0_output = party0_task.await.unwrap();
+        let party1_output = party1_task.await.unwrap();
+
+        // Shutdown the fabrics
+        party0_fabric.shutdown();
+        party1_fabric.shutdown();
+
+        (party0_output, party1_output)
+    }
 }

src/network.rs

@@ -5,7 +5,7 @@ mod config;
 mod mock;
 
 #[cfg(any(feature = "test_helpers", test))]
-pub use mock::NoRecvNetwork;
+pub use mock::{NoRecvNetwork, UnboundedDuplexStream, MockNetwork};
 
 use async_trait::async_trait;
 use quinn::{Endpoint, RecvStream, SendStream};

src/network/mock.rs

@@ -2,6 +2,7 @@
 
 use async_trait::async_trait;
 use futures::future::pending;
+use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender};
 
 use crate::{error::MpcNetworkError, PARTY0};
 
@@ -36,3 +37,93 @@ impl MpcNetwork for NoRecvNetwork {
         Ok(())
     }
 }
+
+/// A dummy MPC network that operates over a duplex channel instead of a network connection/// An unbounded duplex channel used to mock a network connection
+pub struct UnboundedDuplexStream {
+    /// The send side of the stream
+    send: UnboundedSender<NetworkOutbound>,
+    /// The receive side of the stream
+    recv: UnboundedReceiver<NetworkOutbound>,
+}
+
+impl UnboundedDuplexStream {
+    /// Create a new pair of duplex streams
+    pub fn new_duplex_pair() -> (Self, Self) {
+        let (send1, recv1) = unbounded_channel();
+        let (send2, recv2) = unbounded_channel();
+
+        (
+            Self {
+                send: send1,
+                recv: recv2,
+            },
+            Self {
+                send: send2,
+                recv: recv1,
+            },
+        )
+    }
+
+    /// Send a message on the stream
+    pub fn send(&mut self, msg: NetworkOutbound) {
+        self.send.send(msg).unwrap();
+    }
+
+    /// Recv a message from the stream
+    pub async fn recv(&mut self) -> NetworkOutbound {
+        self.recv.recv().await.unwrap()
+    }
+}
+
+/// A dummy network implementation used for unit testing
+pub struct MockNetwork {
+    /// The ID of the local party
+    party_id: PartyId,
+    /// The underlying mock network connection
+    mock_conn: UnboundedDuplexStream,
+}
+
+impl MockNetwork {
+    /// Create a new mock network from one half of a duplex stream
+    pub fn new(party_id: PartyId, stream: UnboundedDuplexStream) -> Self {
+        Self {
+            party_id,
+            mock_conn: stream,
+        }
+    }
+}
+
+#[async_trait]
+impl MpcNetwork for MockNetwork {
+    fn party_id(&self) -> PartyId {
+        self.party_id
+    }
+
+    async fn send_message(&mut self, message: NetworkOutbound) -> Result<(), MpcNetworkError> {
+        self.mock_conn.send(message);
+        Ok(())
+    }
+
+    async fn receive_message(&mut self) -> Result<NetworkOutbound, MpcNetworkError> {
+        let msg = self.mock_conn.recv().await;
+        Ok(msg)
+    }
+
+    async fn exchange_messages(
+        &mut self,
+        message: NetworkOutbound,
+    ) -> Result<NetworkOutbound, MpcNetworkError> {
+        if self.party_id() == PARTY0 {
+            self.send_message(message).await?;
+            self.receive_message().await
+        } else {
+            let res = self.receive_message().await?;
+            self.send_message(message).await?;
+            Ok(res)
+        }
+    }
+
+    async fn close(&mut self) -> Result<(), MpcNetworkError> {
+        Ok(())
+    }
+}