Keep track of nodes that have sent us successful messages

Otherwise, we reject all nodes as faulty.
poanetwork · Sep 18, 2018 · 2895c55 · 2895c55
1 parent 8c27895
commit 2895c55
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 12 deletions.
diff --git a/src/honey_badger/epoch_state.rs b/src/honey_badger/epoch_state.rs
@@ -117,7 +117,7 @@ pub struct EpochState<C, N: Rand> {
     /// The status of threshold decryption, by proposer.
     decryption: BTreeMap<N, DecryptionState<N>>,
     /// N seen so far
-    map: BTreeMap<N, Vec<u8>>,
+    nodes_found_so_far: BTreeSet<N>,
     _phantom: PhantomData<C>,
 }
 
@@ -134,7 +134,7 @@ where
             netinfo,
             subset: SubsetState::Ongoing(cs),
             decryption: BTreeMap::default(),
-            map: Default::default(),
+            nodes_found_so_far: Default::default(),
             _phantom: PhantomData,
         })
     }
@@ -190,8 +190,11 @@ where
         let plaintexts: BTreeMap<N, &[u8]> = self
             .decryption
             .iter()
-            .flat_map(|(id, dec_state)| dec_state.plaintext().map(|pt| (id.clone(), pt)))
-            .collect();
+            .flat_map(|(id, dec_state)| {
+                dec_state
+                    .plaintext()
+                    .nodes_found_so_far(|pt| (id.clone(), pt))
+            }).collect();
         if !proposer_ids.iter().eq(plaintexts.keys()) {
             return None; // Not all accepted contributions are decrypted yet.
         }
@@ -228,15 +231,16 @@ where
         for cs_output in cs_outputs {
             match cs_output {
                 SubsetOutput::Contribution(k, v) => {
-                    step.extend(self.send_decryption_share(k, &v)?);
+                    step.extend(self.send_decryption_share(k.clone(), &v)?);
+                    self.nodes_found_so_far.insert(k);
                 }
                 SubsetOutput::Done => {
-                    self.subset = SubsetState::Complete(self.map.keys().cloned().collect());
+                    self.subset = SubsetState::Complete(self.nodes_found_so_far.cloned().collect());
 
                     let faulty_shares: Vec<_> = self
                         .decryption
                         .keys()
-                        .filter(|id| !self.map.contains_key(id))
+                        .filter(|id| !self.nodes_found_so_far.contains_key(id))
                         .cloned()
                         .collect();
                     for id in faulty_shares {

diff --git a/tests/subset.rs b/tests/subset.rs
@@ -43,7 +43,6 @@ fn test_subset<A: Adversary<Subset<NodeId>>>(
         network.step();
     }
     // Verify that all instances output the same set.
-    let mut expected = None;
     for node in network.nodes.values() {
         let outputs = node.outputs();
         let mut actual = BTreeMap::default();
@@ -61,10 +60,9 @@ fn test_subset<A: Adversary<Subset<NodeId>>>(
         // The Subset algorithm guarantees that more than two thirds of the proposed elements
         // are in the set.
         assert!(actual.len() * 3 > inputs.len() * 2);
-        match expected {
-            None => expected = Some(actual),
-            Some(ref set) => assert_eq!(&actual, set),
-        };
+        for (id, value) in actual {
+            assert_eq!(&inputs[&id], value);
+        }
     }
 }