Simon/snapshots

src/test_utils/mockrng.rs

@@ -2,6 +2,7 @@ use rand::{CryptoRng, RngCore};
 use rand_chacha::{rand_core::SeedableRng, ChaCha12Rng};
 
 /// Used for deterministic Rngs and only in testing
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct MockCryptoRng(ChaCha12Rng);
 
 impl MockCryptoRng {
@@ -26,3 +27,41 @@ impl RngCore for MockCryptoRng {
 }
 
 impl CryptoRng for MockCryptoRng {}
+
+#[cfg(test)]
+pub mod test {
+    use crate::participants::Participant;
+    use crate::test_utils::{generate_participants, MockCryptoRng};
+    use rand::RngCore;
+    use rand_core::CryptoRngCore;
+
+    pub fn create_rngs(
+        participants: &[Participant],
+        seed: &mut impl CryptoRngCore,
+    ) -> Vec<MockCryptoRng> {
+        let rngs = participants
+            .iter()
+            .map(|_| MockCryptoRng::seed_from_u64(seed.next_u64()))
+            .collect::<Vec<_>>();
+        rngs
+    }
+
+    #[test]
+    fn test_clone_rngs() {
+        let participants = generate_participants(5);
+        let mut rng = MockCryptoRng::seed_from_u64(42u64);
+        let mut rngs = create_rngs(&participants, &mut rng);
+        // Clone rng
+        let mut clone_rngs = rngs.clone();
+
+        let consumption = rngs.iter_mut().map(RngCore::next_u64).collect::<Vec<_>>();
+        let clone_consumption = clone_rngs
+            .iter_mut()
+            .map(RngCore::next_u64)
+            .collect::<Vec<_>>();
+
+        for (c1, c2) in consumption.iter().zip(clone_consumption.iter()) {
+            assert_eq!(c1, c2);
+        }
+    }
+}

src/test_utils/mod.rs

@@ -9,8 +9,9 @@ mod dkg;
 mod mockrng;
 mod participants;
 mod presign;
-mod run_protocol;
+mod protocol;
 mod sign;
+mod snapshot;
 pub mod test_generators;
 
 use crate::errors::ProtocolError;
@@ -29,8 +30,9 @@ pub use mockrng::MockCryptoRng;
 pub use dkg::{assert_public_key_invariant, run_keygen, run_refresh, run_reshare};
 pub use participants::{generate_participants, generate_participants_with_random_ids};
 pub use presign::ecdsa_generate_rerandpresig_args;
-pub use run_protocol::{run_protocol, run_two_party_protocol};
+pub use protocol::{run_protocol, run_protocol_and_take_snapshots, run_two_party_protocol};
 pub use sign::{check_one_coordinator_output, run_sign};
+pub use snapshot::ProtocolSnapshot;
 pub use test_generators::*;
 
 /// Checks that the list contains all None but one element

src/test_utils/run_protocol.rs → src/test_utils/protocol.rs

@@ -1,6 +1,7 @@
 use crate::errors::ProtocolError;
 use crate::participants::Participant;
 use crate::protocol::{Action, Protocol};
+use crate::test_utils::ProtocolSnapshot;
 use std::collections::HashMap;
 
 // +++++++++++++++++ Any Protocol +++++++++++++++++ //
@@ -51,6 +52,71 @@ pub fn run_protocol<T>(
     Ok(out)
 }
 
+/// Like [`run_protocol()`], except that it snapshots all the communication.
+pub fn run_protocol_and_take_snapshots<T>(
+    mut ps: Vec<(Participant, Box<dyn Protocol<Output = T>>)>,
+) -> Result<(Vec<(Participant, T)>, ProtocolSnapshot), ProtocolError> {
+    // Get the participants
+    let participants: Vec<_> = ps.iter().map(|(p, _)| *p).collect();
+    // Build the snapshot
+    let mut protocol_snapshots = ProtocolSnapshot::new_empty(participants);
+
+    // Compute the participants indices
+    let indices: HashMap<Participant, usize> =
+        ps.iter().enumerate().map(|(i, (p, _))| (*p, i)).collect();
+
+    // run the protocol
+    let size = ps.len();
+    let mut out = Vec::with_capacity(size);
+    while out.len() < size {
+        for i in 0..size {
+            while {
+                let action = ps[i].1.poke()?;
+                match action {
+                    Action::Wait => false,
+                    Action::SendMany(m) => {
+                        for j in 0..size {
+                            if i == j {
+                                continue;
+                            }
+                            let from = ps[i].0;
+                            let to = ps[j].0;
+                            // snapshot the message
+                            protocol_snapshots
+                                .push_message(to, from, m.clone())
+                                .ok_or_else(|| {
+                                    ProtocolError::Other(
+                                        "Participant not found in snapshot".to_string(),
+                                    )
+                                })?;
+                            ps[j].1.message(from, m.clone());
+                        }
+                        true
+                    }
+                    Action::SendPrivate(to, m) => {
+                        let from = ps[i].0;
+                        // snapshot the message
+                        protocol_snapshots
+                            .push_message(to, from, m.clone())
+                            .ok_or_else(|| {
+                                ProtocolError::Other(
+                                    "Participant not found in snapshot".to_string(),
+                                )
+                            })?;
+                        ps[indices[&to]].1.message(from, m);
+                        true
+                    }
+                    Action::Return(r) => {
+                        out.push((ps[i].0, r));
+                        false
+                    }
+                }
+            } {}
+        }
+    }
+    Ok((out, protocol_snapshots))
+}
+
 /// Like [`run_protocol()`], except for just two parties.
 ///
 /// This is more useful for testing two party protocols with assymetric results,

src/test_utils/snapshot.rs

@@ -0,0 +1,236 @@
+use std::collections::HashMap;
+
+use crate::{participants::Participant, protocol::MessageData};
+
+/// A single received message during a protocol run
+#[derive(Debug, PartialEq, Clone)]
+struct ReceivedMessageSnapshot {
+    from: Participant,
+    message: MessageData,
+}
+
+impl ReceivedMessageSnapshot {
+    fn new(from: Participant, message: MessageData) -> Self {
+        Self { from, message }
+    }
+}
+
+/// A participant's view of their received messages
+struct ParticipantSnapshot {
+    // `snaps` is a list of (participant, messages) pairs
+    snaps: Vec<ReceivedMessageSnapshot>,
+    // `read_index` helps knowing which snap has been already
+    // output during the reading queries
+    read_index: usize,
+}
+
+impl ParticipantSnapshot {
+    pub fn new_empty() -> Self {
+        Self {
+            snaps: Vec::new(),
+            read_index: 0,
+        }
+    }
+
+    fn push_received_message_snapshot(&mut self, snap: ReceivedMessageSnapshot) {
+        self.snaps.push(snap);
+    }
+
+    fn push_message(&mut self, from: Participant, message: MessageData) {
+        let snap = ReceivedMessageSnapshot::new(from, message);
+        self.push_received_message_snapshot(snap);
+    }
+
+    fn read_next_message(&mut self) -> Option<(Participant, MessageData)> {
+        if self.read_index >= self.snaps.len() {
+            return None;
+        }
+        let message_snap = &self.snaps[self.read_index];
+        self.read_index += 1;
+        Some((message_snap.from, message_snap.message.clone()))
+    }
+
+    fn refresh_read_all(&mut self) {
+        self.read_index = 0;
+    }
+}
+
+/// Used to store the snapshot of all the messages sent during
+/// the communication rounds of a certain protocol
+pub struct ProtocolSnapshot {
+    snapshots: HashMap<Participant, ParticipantSnapshot>,
+}
+
+impl ProtocolSnapshot {
+    /// Creates an empty snapshot
+    pub fn new_empty(participants: Vec<Participant>) -> Self {
+        let snapshots = participants
+            .into_iter()
+            .map(|p| (p, ParticipantSnapshot::new_empty()))
+            .collect::<HashMap<_, _>>();
+        Self { snapshots }
+    }
+
+    /// Adds a messages sent by a sender and to a receiver to the protocol snapshot
+    pub fn push_message(
+        &mut self,
+        to: Participant,
+        from: Participant,
+        message: MessageData,
+    ) -> Option<()> {
+        self.snapshots
+            .get_mut(&to)
+            .map(|snapshot| snapshot.push_message(from, message))
+    }
+
+    /// Reads the next message stored in the snapshot of a particular participant given as input
+    pub fn read_next_message_for_participant(
+        &mut self,
+        participant: Participant,
+    ) -> Option<(Participant, MessageData)> {
+        self.snapshots
+            .get_mut(&participant)
+            .and_then(ParticipantSnapshot::read_next_message)
+    }
+
+    pub fn refresh_read_all(&mut self) {
+        for snapshot in self.snapshots.values_mut() {
+            snapshot.refresh_read_all();
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::ecdsa::{
+        robust_ecdsa::{presign::presign, PresignArguments, PresignOutput},
+        KeygenOutput, Polynomial,
+    };
+    use crate::test_utils::{
+        generate_participants, run_protocol_and_take_snapshots, GenProtocol, MockCryptoRng,
+    };
+    use crate::SigningShare;
+    use frost_secp256k1::VerifyingKey;
+    use k256::ProjectivePoint;
+    use rand_core::CryptoRngCore;
+
+    fn generate_random_received_snap(rng: &mut impl CryptoRngCore) -> ReceivedMessageSnapshot {
+        let from = Participant::from(rng.next_u32());
+        let mut message: [u8; 32] = [0u8; 32];
+        rng.fill_bytes(&mut message);
+        let message = message.to_vec();
+        ReceivedMessageSnapshot::new(from, message)
+    }
+
+    #[test]
+    fn test_read_next_message() {
+        let mut psnap = ParticipantSnapshot::new_empty();
+        let mut rvec = Vec::new();
+        let mut rng = MockCryptoRng::seed_from_u64(123_123);
+        for _ in 0..50 {
+            let received_snap = generate_random_received_snap(&mut rng);
+            rvec.push(received_snap.clone());
+            psnap.push_received_message_snapshot(received_snap);
+        }
+        for r in rvec {
+            let (from, message) = psnap.read_next_message().unwrap();
+            let read_message = ReceivedMessageSnapshot::new(from, message);
+            assert_eq!(r, read_message);
+        }
+    }
+
+    #[test]
+    fn test_refresh_read_all() {
+        let mut psnap = ParticipantSnapshot::new_empty();
+        let mut rng = MockCryptoRng::seed_from_u64(123_123);
+        for _ in 0..50 {
+            let received_snap = generate_random_received_snap(&mut rng);
+            psnap.push_received_message_snapshot(received_snap);
+        }
+
+        let mut rvec = Vec::new();
+        for _ in 0..50 {
+            let (from, message) = psnap.read_next_message().unwrap();
+            let read_message = ReceivedMessageSnapshot::new(from, message);
+            rvec.push(read_message);
+        }
+        psnap.refresh_read_all();
+        for r in rvec {
+            let (from, message) = psnap.read_next_message().unwrap();
+            let read_message = ReceivedMessageSnapshot::new(from, message);
+            assert_eq!(r, read_message);
+        }
+    }
+
+    fn prepare_keys(p: Participant, f: &Polynomial, big_x: ProjectivePoint) -> KeygenOutput {
+        let private_share = f.eval_at_participant(p).unwrap();
+        let verifying_key = VerifyingKey::new(big_x);
+        KeygenOutput {
+            private_share: SigningShare::new(private_share.0),
+            public_key: verifying_key,
+        }
+    }
+
+    fn compare_presign_outputs(o1: &PresignOutput, o2: &PresignOutput) -> bool {
+        o1.alpha == o2.alpha
+            && o1.beta == o2.beta
+            && o1.c == o2.c
+            && o1.e == o2.e
+            && o1.big_r == o2.big_r
+    }
+
+    #[test]
+    fn ecdsa_presign_should_return_same_snapshot_when_executed_twice() {
+        let max_malicious = 2;
+        let participants = generate_participants(5);
+
+        let mut rng = MockCryptoRng::seed_from_u64(42u64);
+        let f = Polynomial::generate_polynomial(None, max_malicious, &mut rng).unwrap();
+        let big_x = ProjectivePoint::GENERATOR * f.eval_at_zero().unwrap().0;
+
+        // create rngs for first and second snapshots
+        let rngs = crate::test_utils::mockrng::test::create_rngs(&participants, &mut rng);
+
+        let mut results = Vec::new();
+        let mut snapshots = Vec::new();
+
+        // Running the protocol twice
+        for _ in 0..2 {
+            let mut protocols: GenProtocol<PresignOutput> = Vec::with_capacity(participants.len());
+            for (i, p) in participants.iter().enumerate() {
+                // simulating the key packages for each participant
+                let keygen_out = prepare_keys(*p, &f, big_x);
+                let protocol = presign(
+                    &participants[..],
+                    *p,
+                    PresignArguments {
+                        keygen_out,
+                        threshold: max_malicious,
+                    },
+                    rngs[i].clone(),
+                )
+                .unwrap();
+                protocols.push((*p, Box::new(protocol)));
+            }
+            let (result, snapshot) = run_protocol_and_take_snapshots(protocols).unwrap();
+            results.push(result);
+            snapshots.push(snapshot);
+        }
+
+        // Check the results are the same
+        assert!(results[0].iter().all(|(p1, o1)| {
+            results[1]
+                .iter()
+                .any(|(p2, o2)| p1 == p2 && compare_presign_outputs(o1, o2))
+        }));
+
+        // Check the messages sent per participants are the same
+        for p in participants {
+            while let Some((sender1, msg1)) = snapshots[0].read_next_message_for_participant(p) {
+                let (sender2, msg2) = snapshots[1].read_next_message_for_participant(p).unwrap();
+                assert!(sender1 == sender2 && msg1 == msg2);
+            }
+        }
+    }
+}