From 0f1f6aee34dadab57191fc7de02011501093cdd9 Mon Sep 17 00:00:00 2001 From: RolandSherwin Date: Mon, 26 Sep 2022 12:17:35 +0530 Subject: [PATCH 1/3] tests(sdkg): add test helpers --- Cargo.toml | 1 + src/sdkg.rs | 196 ++++++++++++++++++++++++++++++++-------------------- 2 files changed, 121 insertions(+), 76 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 99d14a9..083614e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,5 +23,6 @@ log = "0.4.13" [dev-dependencies] env_logger = "0.8" +eyre = "~0.6.5" # quickcheck = "1" # quickcheck_macros = "1" diff --git a/src/sdkg.rs b/src/sdkg.rs index 372d20f..0c463cf 100644 --- a/src/sdkg.rs +++ b/src/sdkg.rs @@ -502,9 +502,10 @@ pub enum PartFault { #[cfg(test)] mod tests { - use super::{AckOutcome, PartOutcome, SyncKeyGen}; + use super::{Ack, AckOutcome, Part, PartOutcome, SyncKeyGen}; use bls::{PublicKey, PublicKeySet, SecretKey, SecretKeyShare, SignatureShare}; - use std::collections::BTreeMap; + use eyre::{eyre, Result}; + use std::collections::{BTreeMap, BTreeSet}; #[test] fn test_sdkg() { @@ -606,121 +607,164 @@ mod tests { assert!(pub_key_set.public_key().verify(&sig, msg)); } - fn gen_dkg_key( - threshold: usize, - node_num: usize, - ) -> (Vec<(usize, SecretKeyShare)>, PublicKeySet) { - // Use the OS random number generator for any randomness: - // let mut rng = bls::rand::rngs::OsRng::fill_bytes([0u8; 16]); - let mut rng = bls::rand::rngs::OsRng; + #[test] + fn test_threshold() -> Result<()> { + for nodes_num in 2..10 { + // for threshold in 1..((nodes_num-1)/2+1) { + for threshold in 1..nodes_num { + println!("Testing for threshold {}/{}...", threshold, nodes_num); - // Generate individual key pairs for encryption. These are not suitable for threshold schemes. - let sec_keys: Vec = (0..node_num).map(|_| bls::rand::random()).collect(); + let (secret_key_shares, pub_key_set) = simulate_dkg_round(nodes_num, threshold)?; + let msg = "signed message"; + + // check threshold + 1 sigs matches master key + let mut sig_shares: BTreeMap = BTreeMap::new(); + for (id, sks) in &secret_key_shares[0..threshold + 1] { + let sig_share = sks.sign(msg); + let pks = pub_key_set.public_key_share(id); + assert!(pks.verify(&sig_share, msg)); + sig_shares.insert(*id, sig_share); + } + let sig = pub_key_set + .combine_signatures(&sig_shares) + .map_err(|err| eyre!("The shares can be combined: {err:?}"))?; + assert!(pub_key_set.public_key().verify(&sig, msg)); + + // check threshold sigs are not enough to match master key + let mut sig_shares: BTreeMap = BTreeMap::new(); + for (id, sks) in &secret_key_shares[0..threshold] { + let sig_share = sks.sign(msg); + let pks = pub_key_set.public_key_share(id); + assert!(pks.verify(&sig_share, msg)); + sig_shares.insert(*id, sig_share); + } + let _sig = pub_key_set.combine_signatures(&sig_shares).is_err(); + } + } + Ok(()) + } + + // Test helpers + fn init_nodes( + num_nodes: usize, + threshold: usize, + rng: &mut R, + ) -> Result<(BTreeMap>, Vec<(usize, Part)>)> { + let sec_keys: Vec = (0..num_nodes).map(|_| bls::rand::random()).collect(); let pub_keys: BTreeMap = sec_keys .iter() .map(SecretKey::public_key) .enumerate() .collect(); - // Create the `SyncKeyGen` instances. The constructor also outputs the part that needs to - // be sent to all other participants, so we save the parts together with their sender ID. let mut nodes = BTreeMap::new(); let mut parts = Vec::new(); - for (id, sk) in sec_keys.into_iter().enumerate() { + for (id, sk) in sec_keys.clone().into_iter().enumerate() { let (sync_key_gen, opt_part) = - SyncKeyGen::new(id, sk, pub_keys.clone(), threshold, &mut rng).unwrap_or_else( - |_| panic!("Failed to create `SyncKeyGen` instance for node #{}", id), - ); + SyncKeyGen::new(id, sk, pub_keys.clone(), threshold, rng)?; nodes.insert(id, sync_key_gen); parts.push((id, opt_part.unwrap())); // Would be `None` for observer nodes. } - // All nodes now handle the parts and send the resulting `Ack` messages. + Ok((nodes, parts)) + } + + fn handle_parts( + nodes: &mut BTreeMap>, + parts: &Vec<(usize, Part)>, + rng: &mut R, + ) -> Result> { let mut acks = Vec::new(); for (sender_id, part) in parts { - for (&id, node) in &mut nodes { - match node - .handle_part(&sender_id, part.clone(), &mut rng) - .expect("Failed to handle Part") - { + for (&id, node) in nodes.iter_mut() { + match node.handle_part(&sender_id, part.clone(), rng)? { PartOutcome::Valid(Some(ack)) => acks.push((id, ack)), - PartOutcome::Invalid(fault) => panic!("Invalid Part: {:?}", fault), - PartOutcome::Valid(None) => { - panic!("We are not an observer, so we should send Ack.") - } + _ => return Err(eyre!("We are an observer/invalid part")), } } } + Ok(acks) + } - // Finally, we handle all the `Ack`s. + fn handle_acks( + nodes: &mut BTreeMap>, + acks: &Vec<(usize, Ack)>, + ) -> Result<()> { for (sender_id, ack) in acks { for node in nodes.values_mut() { match node .handle_ack(&sender_id, ack.clone()) - .expect("Failed to handle Ack") + .map_err(|err| eyre!("Failed to handle Ack {err:?}"))? { AckOutcome::Valid => (), - AckOutcome::Invalid(fault) => panic!("Invalid Ack: {:?}", fault), + AckOutcome::Invalid(fault) => return Err(eyre!("Invalid Ack {fault:?}")), } } } + Ok(()) + } + + fn gen_key_share( + nodes: &mut BTreeMap>, + ) -> Result<(Vec<(usize, SecretKeyShare)>, PublicKeySet)> { + let mut pk_set = BTreeSet::new(); - // We have all the information and can generate the key sets. - // Generate the public key set; which is identical for all nodes. - let pub_key_set = nodes[&0] - .generate() - .expect("Failed to create `PublicKeySet` from node #0") - .0; let mut secret_key_shares = Vec::new(); - for (&id, node) in &mut nodes { - assert!(node.is_ready()); - let (pks, opt_sks) = node.generate().unwrap_or_else(|_| { - panic!( - "Failed to create `PublicKeySet` and `SecretKeyShare` for node #{}", - id - ) - }); - assert_eq!(pks, pub_key_set); // All nodes now know the public keys and public key shares. - let sks = opt_sks.expect("Not an observer node: We receive a secret key share."); + for (&id, node) in nodes { + if !node.is_ready() { + return Err(eyre!("Node: {id} is not ready")); + } + let (pks, opt_sks) = node.generate()?; + let sks = opt_sks.ok_or_else(|| eyre!("Node: {id} is an observer"))?; + pk_set.insert(pks); secret_key_shares.push((id, sks)); } - (secret_key_shares, pub_key_set) + // verify that they produced a single pks + if pk_set.len() != 1 { + return Err(eyre!("The pub_key_set is not the same for all the nodes")); + } + let pk_set = Vec::from_iter(pk_set.into_iter()); + + Ok((secret_key_shares, pk_set[0].clone())) } - #[test] - fn test_threshold() { - for nodes_num in 2..=7 { - // for threshold in 1..((nodes_num-1)/2+1) { - for threshold in 1..nodes_num { - println!("Testing for threshold {}/{}...", threshold, nodes_num); + fn simulate_dkg_round( + num_nodes: usize, + threshold: usize, + ) -> Result<(Vec<(usize, SecretKeyShare)>, PublicKeySet)> { + let mut rng = bls::rand::rngs::OsRng; - let (secret_key_shares, pub_key_set) = gen_dkg_key(threshold, nodes_num); - let msg = "signed message"; + let (mut nodes, parts) = init_nodes(num_nodes, threshold, &mut rng)?; + let acks = handle_parts(&mut nodes, &parts, &mut rng)?; + handle_acks(&mut nodes, &acks)?; + gen_key_share(&mut nodes) + } - // check threshold + 1 sigs matches master key - let mut sig_shares: BTreeMap = BTreeMap::new(); - for (id, sks) in &secret_key_shares[0..threshold + 1] { - let sig_share = sks.sign(msg); - let pks = pub_key_set.public_key_share(id); - assert!(pks.verify(&sig_share, msg)); - sig_shares.insert(*id, sig_share); - } - let sig = pub_key_set - .combine_signatures(&sig_shares) - .expect("The shares can be combined."); - assert!(pub_key_set.public_key().verify(&sig, msg)); + #[allow(dead_code)] + fn verify_threshold( + threshold: usize, + sk_shares: &Vec<(usize, SecretKeyShare)>, + pk_set: &PublicKeySet, + ) -> Result<()> { + let msg = "verify threshold"; + let mut sig_shares: BTreeMap = BTreeMap::new(); - // check threshold sigs are not enough to match master key - let mut sig_shares: BTreeMap = BTreeMap::new(); - for (id, sks) in &secret_key_shares[0..threshold] { - let sig_share = sks.sign(msg); - let pks = pub_key_set.public_key_share(id); - assert!(pks.verify(&sig_share, msg)); - sig_shares.insert(*id, sig_share); - } - let _sig = pub_key_set.combine_signatures(&sig_shares).is_err(); + for (id, sks) in sk_shares.iter().take(threshold + 1) { + let sig_share = sks.sign(msg); + let pks = pk_set.public_key_share(id); + if !pks.verify(&sig_share, msg) { + return Err(eyre!("The pub_key_share cannot verify the sig")); } + sig_shares.insert(*id, sig_share); } + + let sig = pk_set.combine_signatures(&sig_shares)?; + + if !pk_set.public_key().verify(&sig, msg) { + return Err(eyre!("The pub_key_set cannot verify the sig")); + } + + Ok(()) } } From 7ca647bf6af6b7415e14d37b81d122b9185dd0ca Mon Sep 17 00:00:00 2001 From: RolandSherwin Date: Thu, 29 Sep 2022 14:34:16 +0530 Subject: [PATCH 2/3] tests(state): fuzz test to ensure DKG completes regardless of the order of votes --- src/sdkg.rs | 14 +- src/state.rs | 517 +++++++++++++++++++++++++++++++++++++++++++++++++++ src/vote.rs | 22 +++ 3 files changed, 546 insertions(+), 7 deletions(-) diff --git a/src/sdkg.rs b/src/sdkg.rs index 0c463cf..f97b963 100644 --- a/src/sdkg.rs +++ b/src/sdkg.rs @@ -501,7 +501,7 @@ pub enum PartFault { } #[cfg(test)] -mod tests { +pub(crate) mod tests { use super::{Ack, AckOutcome, Part, PartOutcome, SyncKeyGen}; use bls::{PublicKey, PublicKeySet, SecretKey, SecretKeyShare, SignatureShare}; use eyre::{eyre, Result}; @@ -645,6 +645,7 @@ mod tests { } // Test helpers + #[allow(clippy::type_complexity)] fn init_nodes( num_nodes: usize, threshold: usize, @@ -659,7 +660,7 @@ mod tests { let mut nodes = BTreeMap::new(); let mut parts = Vec::new(); - for (id, sk) in sec_keys.clone().into_iter().enumerate() { + for (id, sk) in sec_keys.into_iter().enumerate() { let (sync_key_gen, opt_part) = SyncKeyGen::new(id, sk, pub_keys.clone(), threshold, rng)?; nodes.insert(id, sync_key_gen); @@ -677,7 +678,7 @@ mod tests { let mut acks = Vec::new(); for (sender_id, part) in parts { for (&id, node) in nodes.iter_mut() { - match node.handle_part(&sender_id, part.clone(), rng)? { + match node.handle_part(sender_id, part.clone(), rng)? { PartOutcome::Valid(Some(ack)) => acks.push((id, ack)), _ => return Err(eyre!("We are an observer/invalid part")), } @@ -693,7 +694,7 @@ mod tests { for (sender_id, ack) in acks { for node in nodes.values_mut() { match node - .handle_ack(&sender_id, ack.clone()) + .handle_ack(sender_id, ack.clone()) .map_err(|err| eyre!("Failed to handle Ack {err:?}"))? { AckOutcome::Valid => (), @@ -741,10 +742,9 @@ mod tests { gen_key_share(&mut nodes) } - #[allow(dead_code)] - fn verify_threshold( + pub(crate) fn verify_threshold( threshold: usize, - sk_shares: &Vec<(usize, SecretKeyShare)>, + sk_shares: &[(usize, SecretKeyShare)], pk_set: &PublicKeySet, ) -> Result<()> { let msg = "verify threshold"; diff --git a/src/state.rs b/src/state.rs index 2a0367d..0791d65 100644 --- a/src/state.rs +++ b/src/state.rs @@ -339,6 +339,8 @@ impl DkgState { #[cfg(test)] mod tests { use super::*; + use crate::{sdkg::tests::verify_threshold, vote::test_utils::*}; + use bls::rand::{rngs::StdRng, seq::IteratorRandom, thread_rng, Rng, RngCore, SeedableRng}; #[test] fn test_recursive_handle_vote() { @@ -365,4 +367,519 @@ mod tests { assert!(matches!(res[2], VoteResponse::DkgComplete(_, _))); assert_eq!(res.len(), 3); } + + #[test] + fn fuzz_test() -> Result<()> { + let mut fuzz_count = 20; + let mut rng_for_seed = thread_rng(); + let num_nodes = 7; + let threshold = 4; + + while fuzz_count != 0 { + let seed = rng_for_seed.gen(); + println!(" SEED {seed:?} => count_remaining: {fuzz_count}"); + let mut rng = StdRng::seed_from_u64(seed); + + let mut nodes = generate_nodes(num_nodes, threshold, &mut rng)?; + let mut parts: BTreeMap = BTreeMap::new(); + let mut acks: BTreeMap = BTreeMap::new(); + let mut all_acks: BTreeMap = BTreeMap::new(); + let mut sk_shares: BTreeMap = BTreeMap::new(); + let mut pk_set: BTreeSet = BTreeSet::new(); + + for node in nodes.iter_mut() { + parts.insert(node.id() as usize, node.first_vote()?); + } + + for cmd in fuzz_commands(num_nodes, seed) { + // println!("==> {cmd:?}"); + match cmd { + SendVote::Parts(from_node, to_nodes) => { + for (to, expt_resp) in to_nodes { + let actual_resp = nodes[to] + .handle_signed_vote(parts[&from_node].clone(), &mut rng)?; + assert_eq!(expt_resp.len(), actual_resp.len()); + expt_resp.into_iter().zip(actual_resp.into_iter()).for_each( + |(exp, actual)| { + assert!(exp.match_resp( + actual, + &mut acks, + &mut all_acks, + &mut sk_shares, + &mut pk_set, + to + )); + }, + ) + } + } + SendVote::Acks(from_node, to_nodes) => { + for (to, expt_resp) in to_nodes { + let actual_resp = + nodes[to].handle_signed_vote(acks[&from_node].clone(), &mut rng)?; + assert_eq!(expt_resp.len(), actual_resp.len()); + expt_resp.into_iter().zip(actual_resp.into_iter()).for_each( + |(exp, actual)| { + assert!(exp.match_resp( + actual, + &mut acks, + &mut all_acks, + &mut sk_shares, + &mut pk_set, + to + )); + }, + ) + } + } + SendVote::AllAcks(from_node, to_nodes) => { + for (to, expt_resp) in to_nodes { + let actual_resp = nodes[to] + .handle_signed_vote(all_acks[&from_node].clone(), &mut rng)?; + assert_eq!(expt_resp.len(), actual_resp.len()); + expt_resp.into_iter().zip(actual_resp.into_iter()).for_each( + |(exp, actual)| { + assert!(exp.match_resp( + actual, + &mut acks, + &mut all_acks, + &mut sk_shares, + &mut pk_set, + to + )); + }, + ) + } + } + } + } + + assert_eq!(pk_set.len(), 1); + let pk_set = pk_set.into_iter().collect::>()[0].clone(); + let sk_shares: Vec<_> = sk_shares.into_iter().collect(); + + assert!(verify_threshold(threshold, &sk_shares, &pk_set).is_ok()); + fuzz_count -= 1; + } + Ok(()) + } + + // Returns a list of `SendVote` which when executed in that order will simulate a DKG round from start to completion + // for each node + fn fuzz_commands(num_nodes: usize, seed: u64) -> Vec { + let mut rng = StdRng::seed_from_u64(seed); + let mut nodes = MockNode::new(num_nodes); + // probability for a node to resend vote to another node which has already handled it. + let resend_probability = Some((1, 10)); + // these nodes are required to help other nodes terminate + let mut active_nodes = MockNode::active_nodes(&nodes); + let mut commands = Vec::new(); + + while !active_nodes.is_empty() { + // get a random active node + let current_node = active_nodes[rng.gen::() % active_nodes.len()]; + + // check if current_node can send part/acks/all_acks. + let parts = nodes[current_node].can_send_parts(&nodes, resend_probability, &mut rng); + let acks = nodes[current_node].can_send_acks(&nodes, resend_probability, &mut rng); + let all_acks = + nodes[current_node].can_send_all_acks(&nodes, resend_probability, &mut rng); + + // continue if current_node cant progress + if parts.is_empty() && acks.is_empty() && all_acks.is_empty() { + continue; + } + + // randomly send out part/acks/all_acks + match rng.gen::() % 3 { + 0 if !parts.is_empty() => { + let to_nodes = MockNode::sample_nodes(&parts, &mut rng); + + // update each `to` node and get its (id, response) + let to_nodes_resp = to_nodes + .into_iter() + .map(|to| { + let mut resp = Vec::new(); + // skip if already handled + if let Some(val) = nodes[to].handled_parts.get(¤t_node) { + if *val { + return (to, resp); + } + } + + if let Some(val) = nodes[to].handled_parts.insert(current_node, true) { + if nodes[to].parts_done() { + resp.push(MockVoteResponse::BroadcastVote( + MockDkgVote::SingleAck, + )); + // if we have handled the all the `Acks` before the parts + if nodes[to].acks_done() { + resp.push(MockVoteResponse::BroadcastVote( + MockDkgVote::AllAcks, + )); + } + } else { + // if false, we need more votes + if !val { + resp.push(MockVoteResponse::WaitingForMoreVotes) + } + } + } + + (to, resp) + }) + .collect(); + + commands.push(SendVote::Parts(current_node, to_nodes_resp)); + } + 1 if !acks.is_empty() => { + let to_nodes = MockNode::sample_nodes(&acks, &mut rng); + + let to_nodes_resp = to_nodes + .into_iter() + .map(|to| { + let mut resp = Vec::new(); + // skip if already handled + if let Some(val) = nodes[to].handled_acks.get(¤t_node) { + if *val { + return (to, resp); + } + } + let res = nodes[to].handled_acks.insert(current_node, true); + // if our parts are not done, we will not understand this vote + if !nodes[to].parts_done() { + resp.push(MockVoteResponse::RequestAntiEntropy) + } else if let Some(val) = res { + if nodes[to].acks_done() { + resp.push(MockVoteResponse::BroadcastVote( + MockDkgVote::AllAcks, + )); + // if we have handled the all the `AllAcks` before the Acks + if nodes[to].all_acks_done() { + resp.push(MockVoteResponse::DkgComplete); + } + } else { + // if false, we need more votes + if !val { + resp.push(MockVoteResponse::WaitingForMoreVotes) + } + } + }; + + (to, resp) + }) + .collect(); + + commands.push(SendVote::Acks(current_node, to_nodes_resp)); + } + 2 if !all_acks.is_empty() => { + let to_nodes = MockNode::sample_nodes(&all_acks, &mut rng); + + let to_nodes_resp = to_nodes + .into_iter() + .map(|to| { + let mut resp = Vec::new(); + // skip if already handled + if let Some(val) = nodes[to].handled_all_acks.get(¤t_node) { + if *val { + return (to, resp); + } + } + let res = nodes[to].handled_all_acks.insert(current_node, true); + + // if our Acks are not done, we will not understand this vote + if !nodes[to].acks_done() { + resp.push(MockVoteResponse::RequestAntiEntropy); + } else if let Some(val) = res { + if nodes[to].all_acks_done() { + resp.push(MockVoteResponse::DkgComplete) + } else { + // if false, we need more votes + if !val { + resp.push(MockVoteResponse::WaitingForMoreVotes) + } + } + }; + (to, resp) + }) + .collect(); + + commands.push(SendVote::AllAcks(current_node, to_nodes_resp)); + } + _ => {} + } + + active_nodes = MockNode::active_nodes(&nodes); + } + commands + } + + // Test helpers + fn generate_nodes( + num_nodes: usize, + threshold: usize, + mut rng: &mut R, + ) -> Result> { + let secret_keys: Vec = (0..num_nodes).map(|_| bls::rand::random()).collect(); + let pub_keys: BTreeMap<_, _> = secret_keys + .iter() + .enumerate() + .map(|(id, sk)| (id as u8, sk.public_key())) + .collect(); + secret_keys + .iter() + .enumerate() + .map(|(id, sk)| { + DkgState::new(id as u8, sk.clone(), pub_keys.clone(), threshold, &mut rng) + }) + .collect() + } + + #[derive(Debug)] + enum SendVote { + // from_node, list of (to_node, vec of response when handled) + Parts(usize, Vec<(usize, Vec)>), + Acks(usize, Vec<(usize, Vec)>), + AllAcks(usize, Vec<(usize, Vec)>), + } + + #[derive(Debug)] + enum MockVoteResponse { + WaitingForMoreVotes, + BroadcastVote(MockDkgVote), + RequestAntiEntropy, + DkgComplete, + } + + impl MockVoteResponse { + pub fn match_resp( + &self, + actual_resp: VoteResponse, + update_acks: &mut BTreeMap, + update_all_acks: &mut BTreeMap, + update_sk: &mut BTreeMap, + update_pk: &mut BTreeSet, + id: usize, + ) -> bool { + if (matches!(self, Self::WaitingForMoreVotes) + && matches!(actual_resp, VoteResponse::WaitingForMoreVotes)) + || (matches!(self, Self::RequestAntiEntropy) + && matches!(actual_resp, VoteResponse::RequestAntiEntropy)) + { + true + } else if matches!(self, Self::BroadcastVote(MockDkgVote::SingleAck)) { + match actual_resp { + VoteResponse::BroadcastVote(vote) + if matches!(vote.mock(), MockDkgVote::SingleAck) => + { + update_acks.insert(id, *vote); + true + } + _ => false, + } + } else if matches!(self, Self::BroadcastVote(MockDkgVote::AllAcks)) { + match actual_resp { + VoteResponse::BroadcastVote(vote) + if matches!(vote.mock(), MockDkgVote::AllAcks) => + { + update_all_acks.insert(id, *vote); + true + } + _ => false, + } + } else if matches!(self, Self::DkgComplete) { + if let VoteResponse::DkgComplete(pk, sk) = actual_resp { + update_pk.insert(pk); + update_sk.insert(id, sk); + true + } else { + false + } + } else { + false + } + } + } + + #[derive(Debug)] + struct MockNode { + id: usize, + // Has the current node handled parts, acks, all_acks from another node? + handled_parts: BTreeMap, + handled_acks: BTreeMap, + handled_all_acks: BTreeMap, + } + + impl MockNode { + pub fn new(num_nodes: usize) -> Vec { + let mut status: BTreeMap = BTreeMap::new(); + (0..num_nodes).for_each(|id| { + let _ = status.insert(id, false); + }); + (0..num_nodes) + .map(|id| { + // we have handled our parts/acks/all_acks by default + let mut our_status = status.clone(); + our_status.insert(id, true); + MockNode { + id, + handled_parts: our_status.clone(), + handled_acks: our_status.clone(), + handled_all_acks: our_status, + } + }) + .collect() + } + + // return the node IDs that have not handled self's part; Also choose nodes which have already handled + // self's part with a probability of (num/den) + pub fn can_send_parts( + &self, + nodes: &[MockNode], + resend_probability: Option<(u32, u32)>, + rng: &mut R, + ) -> Vec { + nodes + .iter() + .filter_map(|node| { + // if node has not handled self's part + if !node.handled_parts[&self.id] { + Some(node.id) + } else { + // resend to the node which has already handled self's part with the provided probability + if let Some((num, den)) = resend_probability { + if rng.gen_ratio(num, den) { + Some(node.id) + } else { + None + } + } else { + None + } + } + }) + .collect() + } + + pub fn can_send_acks( + &self, + nodes: &[MockNode], + resend_probability: Option<(u32, u32)>, + rng: &mut R, + ) -> Vec { + // if self has not handled the parts from other nodes, then it cant produce an ack + if !self.parts_done() { + return Vec::new(); + } + // the other node should not have handled self's ack + nodes + .iter() + .filter_map(|node| { + if !node.handled_acks[&self.id] { + Some(node.id) + } else { + // resend to the node which has already handled self's ack with the provided probability + if let Some((num, den)) = resend_probability { + if rng.gen_ratio(num, den) { + Some(node.id) + } else { + None + } + } else { + None + } + } + }) + .collect() + } + + pub fn can_send_all_acks( + &self, + nodes: &[MockNode], + resend_probability: Option<(u32, u32)>, + rng: &mut R, + ) -> Vec { + // self should've handled all the acks/parts (except self's) + if !self.parts_done() { + return Vec::new(); + } + if !self.acks_done() { + return Vec::new(); + } + // the other node should not have handled self's all_ack + nodes + .iter() + .filter_map(|node| { + if !node.handled_all_acks[&self.id] { + Some(node.id) + } else if let Some((num, den)) = resend_probability { + if rng.gen_ratio(num, den) { + Some(node.id) + } else { + None + } + } else { + None + } + }) + .collect() + } + + // returns true if self has received/handled all the parts (except itself) + pub fn parts_done(&self) -> bool { + self.handled_parts + .iter() + .filter(|(&id, _)| id != self.id) + .all(|(_, &val)| val) + } + + pub fn acks_done(&self) -> bool { + self.handled_acks + .iter() + .filter(|(&id, _)| id != self.id) + .all(|(_, &val)| val) + } + + pub fn all_acks_done(&self) -> bool { + // check if current_node has completed the dkg round; i.e., it has handled all_acks from all other nodes + self.handled_all_acks + .iter() + .filter(|(&id, _)| id != self.id) + .all(|(_, &val)| val) + } + + pub fn active_nodes(nodes: &[MockNode]) -> Vec { + // a node is active if any of the other node still requires votes from the current node + // filter out current node as we don't necessarily have to deal with our votes to move forward + let mut active_nodes = BTreeSet::new(); + nodes.iter().for_each(|node| { + // check parts + node.handled_parts.iter().for_each(|(&id, &val)| { + // if current node has not handled a part from another node (i.e. false), we need the other node + if id != node.id && !val { + active_nodes.insert(id); + }; + }); + + node.handled_acks.iter().for_each(|(&id, &val)| { + if id != node.id && !val { + active_nodes.insert(id); + }; + }); + + node.handled_all_acks.iter().for_each(|(&id, &val)| { + if id != node.id && !val { + active_nodes.insert(id); + }; + }); + }); + active_nodes.into_iter().collect() + } + + // select a subset of node i's from the given list + pub fn sample_nodes(nodes: &Vec, rng: &mut R) -> Vec { + let sample_n_nodes = (rng.gen::() % nodes.len()) + 1; + nodes.iter().cloned().choose_multiple(rng, sample_n_nodes) + } + } } diff --git a/src/vote.rs b/src/vote.rs index d75e907..14a70dd 100644 --- a/src/vote.rs +++ b/src/vote.rs @@ -75,3 +75,25 @@ impl DkgSignedVote { matches!(self.vote, DkgVote::AllAcks(_)) } } + +#[cfg(test)] +pub(crate) mod test_utils { + use super::{DkgSignedVote, DkgVote}; + + #[derive(Debug)] + pub(crate) enum MockDkgVote { + SinglePart, + SingleAck, + AllAcks, + } + + impl DkgSignedVote { + pub(crate) fn mock(&self) -> MockDkgVote { + match self.vote { + DkgVote::SinglePart(_) => MockDkgVote::SinglePart, + DkgVote::SingleAck(_) => MockDkgVote::SingleAck, + DkgVote::AllAcks(_) => MockDkgVote::AllAcks, + } + } + } +} From 844ea9cb7dae78f648ccc51dd65231bb2ddb0249 Mon Sep 17 00:00:00 2001 From: RolandSherwin Date: Fri, 7 Oct 2022 13:24:59 +0530 Subject: [PATCH 3/3] chore(tests): implement ParitalEq for the Mock types --- src/state.rs | 371 +++++++++++++++++++++++++-------------------------- src/vote.rs | 14 +- 2 files changed, 190 insertions(+), 195 deletions(-) diff --git a/src/state.rs b/src/state.rs index 0791d65..aeb438f 100644 --- a/src/state.rs +++ b/src/state.rs @@ -338,9 +338,12 @@ impl DkgState { #[cfg(test)] mod tests { + use super::*; use crate::{sdkg::tests::verify_threshold, vote::test_utils::*}; use bls::rand::{rngs::StdRng, seq::IteratorRandom, thread_rng, Rng, RngCore, SeedableRng}; + use eyre::{eyre, Result}; + use std::env; #[test] fn test_recursive_handle_vote() { @@ -370,7 +373,11 @@ mod tests { #[test] fn fuzz_test() -> Result<()> { - let mut fuzz_count = 20; + let mut fuzz_count = if let Ok(count) = env::var("FUZZ_TEST_COUNT") { + count.parse::().map_err(|err| eyre!("{err}"))? + } else { + 20 + }; let mut rng_for_seed = thread_rng(); let num_nodes = 7; let threshold = 4; @@ -393,64 +400,28 @@ mod tests { for cmd in fuzz_commands(num_nodes, seed) { // println!("==> {cmd:?}"); - match cmd { - SendVote::Parts(from_node, to_nodes) => { - for (to, expt_resp) in to_nodes { - let actual_resp = nodes[to] - .handle_signed_vote(parts[&from_node].clone(), &mut rng)?; - assert_eq!(expt_resp.len(), actual_resp.len()); - expt_resp.into_iter().zip(actual_resp.into_iter()).for_each( - |(exp, actual)| { - assert!(exp.match_resp( - actual, - &mut acks, - &mut all_acks, - &mut sk_shares, - &mut pk_set, - to - )); - }, - ) - } - } - SendVote::Acks(from_node, to_nodes) => { - for (to, expt_resp) in to_nodes { - let actual_resp = - nodes[to].handle_signed_vote(acks[&from_node].clone(), &mut rng)?; - assert_eq!(expt_resp.len(), actual_resp.len()); - expt_resp.into_iter().zip(actual_resp.into_iter()).for_each( - |(exp, actual)| { - assert!(exp.match_resp( - actual, - &mut acks, - &mut all_acks, - &mut sk_shares, - &mut pk_set, - to - )); - }, - ) - } - } - SendVote::AllAcks(from_node, to_nodes) => { - for (to, expt_resp) in to_nodes { - let actual_resp = nodes[to] - .handle_signed_vote(all_acks[&from_node].clone(), &mut rng)?; - assert_eq!(expt_resp.len(), actual_resp.len()); - expt_resp.into_iter().zip(actual_resp.into_iter()).for_each( - |(exp, actual)| { - assert!(exp.match_resp( - actual, - &mut acks, - &mut all_acks, - &mut sk_shares, - &mut pk_set, - to - )); - }, - ) - } - } + let (to_nodes, vote) = match cmd { + SendVote::Parts(from, to_nodes) => (to_nodes, parts[&from].clone()), + SendVote::Acks(from, to_nodes) => (to_nodes, acks[&from].clone()), + SendVote::AllAcks(from, to_nodes) => (to_nodes, all_acks[&from].clone()), + }; + // send the vote to each `to` node + for (to, expt_resp) in to_nodes { + let actual_resp = nodes[to].handle_signed_vote(vote.clone(), &mut rng)?; + assert_eq!(expt_resp.len(), actual_resp.len()); + expt_resp + .into_iter() + .zip(actual_resp.into_iter()) + .for_each(|(exp, actual)| { + assert!(exp.match_resp( + actual, + &mut acks, + &mut all_acks, + &mut sk_shares, + &mut pk_set, + to + )); + }) } } @@ -470,7 +441,7 @@ mod tests { let mut rng = StdRng::seed_from_u64(seed); let mut nodes = MockNode::new(num_nodes); // probability for a node to resend vote to another node which has already handled it. - let resend_probability = Some((1, 10)); + let resend_probability = Some((1, 5)); // these nodes are required to help other nodes terminate let mut active_nodes = MockNode::active_nodes(&nodes); let mut commands = Vec::new(); @@ -490,123 +461,132 @@ mod tests { continue; } + let mut done = false; // randomly send out part/acks/all_acks - match rng.gen::() % 3 { - 0 if !parts.is_empty() => { - let to_nodes = MockNode::sample_nodes(&parts, &mut rng); - - // update each `to` node and get its (id, response) - let to_nodes_resp = to_nodes - .into_iter() - .map(|to| { - let mut resp = Vec::new(); - // skip if already handled - if let Some(val) = nodes[to].handled_parts.get(¤t_node) { - if *val { - return (to, resp); + while !done { + match rng.gen::() % 3 { + 0 if !parts.is_empty() => { + let to_nodes = MockNode::sample_nodes(&parts, &mut rng); + + // update each `to` node and get its (id, response) + let to_nodes_resp = to_nodes + .into_iter() + .map(|to| { + let mut resp = Vec::new(); + // skip if already handled + if let Some(val) = nodes[to].handled_parts.get(¤t_node) { + if *val { + return (to, resp); + } } - } - if let Some(val) = nodes[to].handled_parts.insert(current_node, true) { - if nodes[to].parts_done() { - resp.push(MockVoteResponse::BroadcastVote( - MockDkgVote::SingleAck, - )); - // if we have handled the all the `Acks` before the parts - if nodes[to].acks_done() { + if let Some(val) = + nodes[to].handled_parts.insert(current_node, true) + { + if nodes[to].parts_done() { resp.push(MockVoteResponse::BroadcastVote( - MockDkgVote::AllAcks, + MockDkgVote::SingleAck, )); - } - } else { - // if false, we need more votes - if !val { - resp.push(MockVoteResponse::WaitingForMoreVotes) + // if we have handled the all the `Acks` before the parts + if nodes[to].acks_done() { + resp.push(MockVoteResponse::BroadcastVote( + MockDkgVote::AllAcks, + )); + } + } else { + // if false, we need more votes + if !val { + resp.push(MockVoteResponse::WaitingForMoreVotes) + } } } - } - (to, resp) - }) - .collect(); - - commands.push(SendVote::Parts(current_node, to_nodes_resp)); - } - 1 if !acks.is_empty() => { - let to_nodes = MockNode::sample_nodes(&acks, &mut rng); + (to, resp) + }) + .collect(); - let to_nodes_resp = to_nodes - .into_iter() - .map(|to| { - let mut resp = Vec::new(); - // skip if already handled - if let Some(val) = nodes[to].handled_acks.get(¤t_node) { - if *val { - return (to, resp); - } - } - let res = nodes[to].handled_acks.insert(current_node, true); - // if our parts are not done, we will not understand this vote - if !nodes[to].parts_done() { - resp.push(MockVoteResponse::RequestAntiEntropy) - } else if let Some(val) = res { - if nodes[to].acks_done() { - resp.push(MockVoteResponse::BroadcastVote( - MockDkgVote::AllAcks, - )); - // if we have handled the all the `AllAcks` before the Acks - if nodes[to].all_acks_done() { - resp.push(MockVoteResponse::DkgComplete); - } - } else { - // if false, we need more votes - if !val { - resp.push(MockVoteResponse::WaitingForMoreVotes) + commands.push(SendVote::Parts(current_node, to_nodes_resp)); + done = true; + } + 1 if !acks.is_empty() => { + let to_nodes = MockNode::sample_nodes(&acks, &mut rng); + + let to_nodes_resp = to_nodes + .into_iter() + .map(|to| { + let mut resp = Vec::new(); + // skip if already handled + if let Some(val) = nodes[to].handled_acks.get(¤t_node) { + if *val { + return (to, resp); } } - }; - - (to, resp) - }) - .collect(); + let res = nodes[to].handled_acks.insert(current_node, true); + // if our parts are not done, we will not understand this vote + if !nodes[to].parts_done() { + resp.push(MockVoteResponse::RequestAntiEntropy) + } else if let Some(val) = res { + if nodes[to].acks_done() { + resp.push(MockVoteResponse::BroadcastVote( + MockDkgVote::AllAcks, + )); + // if we have handled the all the `AllAcks` before the Acks + if nodes[to].all_acks_done() { + resp.push(MockVoteResponse::DkgComplete); + } + } else { + // if false, we need more votes + if !val { + resp.push(MockVoteResponse::WaitingForMoreVotes) + } + } + }; - commands.push(SendVote::Acks(current_node, to_nodes_resp)); - } - 2 if !all_acks.is_empty() => { - let to_nodes = MockNode::sample_nodes(&all_acks, &mut rng); + (to, resp) + }) + .collect(); - let to_nodes_resp = to_nodes - .into_iter() - .map(|to| { - let mut resp = Vec::new(); - // skip if already handled - if let Some(val) = nodes[to].handled_all_acks.get(¤t_node) { - if *val { - return (to, resp); - } - } - let res = nodes[to].handled_all_acks.insert(current_node, true); - - // if our Acks are not done, we will not understand this vote - if !nodes[to].acks_done() { - resp.push(MockVoteResponse::RequestAntiEntropy); - } else if let Some(val) = res { - if nodes[to].all_acks_done() { - resp.push(MockVoteResponse::DkgComplete) - } else { - // if false, we need more votes - if !val { - resp.push(MockVoteResponse::WaitingForMoreVotes) + commands.push(SendVote::Acks(current_node, to_nodes_resp)); + done = true + } + 2 if !all_acks.is_empty() => { + let to_nodes = MockNode::sample_nodes(&all_acks, &mut rng); + + let to_nodes_resp = to_nodes + .into_iter() + .map(|to| { + let mut resp = Vec::new(); + // skip if already handled + if let Some(val) = nodes[to].handled_all_acks.get(¤t_node) { + if *val { + return (to, resp); } } - }; - (to, resp) - }) - .collect(); + let res = nodes[to].handled_all_acks.insert(current_node, true); - commands.push(SendVote::AllAcks(current_node, to_nodes_resp)); + // if our Acks are not done, we will not understand this vote + if !nodes[to].acks_done() { + resp.push(MockVoteResponse::RequestAntiEntropy); + } else if let Some(val) = res { + if nodes[to].all_acks_done() { + resp.push(MockVoteResponse::DkgComplete) + } else { + // if false, we need more votes + if !val { + resp.push(MockVoteResponse::WaitingForMoreVotes) + } + } + }; + (to, resp) + }) + .collect(); + + commands.push(SendVote::AllAcks(current_node, to_nodes_resp)); + done = true; + } + // happens if the rng lands on a vote list (e.g., all_acks) that is empty + _ => {} } - _ => {} } active_nodes = MockNode::active_nodes(&nodes); @@ -631,6 +611,7 @@ mod tests { .enumerate() .map(|(id, sk)| { DkgState::new(id as u8, sk.clone(), pub_keys.clone(), threshold, &mut rng) + .map_err(|err| eyre!("{err}")) }) .collect() } @@ -651,6 +632,37 @@ mod tests { DkgComplete, } + impl PartialEq for MockVoteResponse { + fn eq(&self, other: &VoteResponse) -> bool { + match self { + MockVoteResponse::WaitingForMoreVotes + if matches!(other, VoteResponse::WaitingForMoreVotes) => + { + true + } + MockVoteResponse::BroadcastVote(mock_vote) => { + if let VoteResponse::BroadcastVote(signed_vote) = other { + *mock_vote == **signed_vote + } else { + false + } + } + + MockVoteResponse::RequestAntiEntropy + if matches!(other, VoteResponse::RequestAntiEntropy) => + { + true + } + MockVoteResponse::DkgComplete + if matches!(other, VoteResponse::DkgComplete(_, _)) => + { + true + } + _ => false, + } + } + } + impl MockVoteResponse { pub fn match_resp( &self, @@ -661,40 +673,21 @@ mod tests { update_pk: &mut BTreeSet, id: usize, ) -> bool { - if (matches!(self, Self::WaitingForMoreVotes) - && matches!(actual_resp, VoteResponse::WaitingForMoreVotes)) - || (matches!(self, Self::RequestAntiEntropy) - && matches!(actual_resp, VoteResponse::RequestAntiEntropy)) - { - true - } else if matches!(self, Self::BroadcastVote(MockDkgVote::SingleAck)) { + if *self == actual_resp { match actual_resp { - VoteResponse::BroadcastVote(vote) - if matches!(vote.mock(), MockDkgVote::SingleAck) => - { + VoteResponse::BroadcastVote(vote) if MockDkgVote::SingleAck == *vote => { update_acks.insert(id, *vote); - true } - _ => false, - } - } else if matches!(self, Self::BroadcastVote(MockDkgVote::AllAcks)) { - match actual_resp { - VoteResponse::BroadcastVote(vote) - if matches!(vote.mock(), MockDkgVote::AllAcks) => - { + VoteResponse::BroadcastVote(vote) if MockDkgVote::AllAcks == *vote => { update_all_acks.insert(id, *vote); - true } - _ => false, - } - } else if matches!(self, Self::DkgComplete) { - if let VoteResponse::DkgComplete(pk, sk) = actual_resp { - update_pk.insert(pk); - update_sk.insert(id, sk); - true - } else { - false + VoteResponse::DkgComplete(pk, sk) => { + update_pk.insert(pk); + update_sk.insert(id, sk); + } + _ => {} } + true } else { false } @@ -799,7 +792,7 @@ mod tests { resend_probability: Option<(u32, u32)>, rng: &mut R, ) -> Vec { - // self should've handled all the acks/parts (except self's) + // // self should've handled all the acks/parts (except self's) if !self.parts_done() { return Vec::new(); } diff --git a/src/vote.rs b/src/vote.rs index 14a70dd..15d90af 100644 --- a/src/vote.rs +++ b/src/vote.rs @@ -80,6 +80,7 @@ impl DkgSignedVote { pub(crate) mod test_utils { use super::{DkgSignedVote, DkgVote}; + #[allow(dead_code)] #[derive(Debug)] pub(crate) enum MockDkgVote { SinglePart, @@ -87,12 +88,13 @@ pub(crate) mod test_utils { AllAcks, } - impl DkgSignedVote { - pub(crate) fn mock(&self) -> MockDkgVote { - match self.vote { - DkgVote::SinglePart(_) => MockDkgVote::SinglePart, - DkgVote::SingleAck(_) => MockDkgVote::SingleAck, - DkgVote::AllAcks(_) => MockDkgVote::AllAcks, + impl PartialEq for MockDkgVote { + fn eq(&self, other: &DkgSignedVote) -> bool { + match self { + MockDkgVote::SinglePart if matches!(other.vote, DkgVote::SinglePart(_)) => true, + MockDkgVote::SingleAck if matches!(other.vote, DkgVote::SingleAck(_)) => true, + MockDkgVote::AllAcks if matches!(other.vote, DkgVote::AllAcks(_)) => true, + _ => false, } } }