Limit the number of buffered key gen messages.

src/dynamic_honey_badger/dynamic_honey_badger.rs

@@ -234,19 +234,11 @@ where
             }
         };
 
-        // If the joining node is correct, it will send at most (N + 1)² + 1 key generation
-        // messages.
-        if Some(sender_id) == kgs.change.candidate() {
-            let n = self.netinfo.num_nodes() + 1;
-            if kgs.candidate_msg_count > n * n {
-                info!(
-                    "Too many key gen messages from candidate {:?}: {:?}.",
-                    sender_id, kg_msg
-                );
-                let fault_kind = FaultKind::TooManyCandidateKeyGenMessages;
-                return Ok(Fault::new(sender_id.clone(), fault_kind).into());
-            }
-            kgs.candidate_msg_count += 1;
+        // If the sender is correct, it will send at most _N + 1_ key generation messages:
+        // one `Part`, and for each validator an `Ack`. _N_ is the node number _after_ the change.
+        if kgs.count_messages(sender_id) > kgs.key_gen.num_nodes() + 1 {
+            let fault_kind = FaultKind::TooManyKeyGenMessages;
+            return Ok(Fault::new(sender_id.clone(), fault_kind).into());
         }
 
         let tx = SignedKeyGenMsg(self.start_epoch, sender_id.clone(), kg_msg, sig);

src/dynamic_honey_badger/mod.rs

@@ -150,9 +150,9 @@ struct KeyGenState<N> {
     key_gen: SyncKeyGen<N>,
     /// The change for which key generation is performed.
     change: Change<N>,
-    /// The number of key generation messages received from the candidate. At most _N² + 1_ are
+    /// The number of key generation messages received from each peer. At most _N + 1_ are
     /// accepted.
-    candidate_msg_count: usize,
+    msg_count: BTreeMap<N, usize>,
 }
 
 impl<N: NodeIdT> KeyGenState<N> {
@@ -160,7 +160,7 @@ impl<N: NodeIdT> KeyGenState<N> {
         KeyGenState {
             key_gen,
             change,
-            candidate_msg_count: 0,
+            msg_count: BTreeMap::new(),
         }
     }
 
@@ -178,6 +178,13 @@ impl<N: NodeIdT> KeyGenState<N> {
             Change::Add(_, _) | Change::Remove(_) => None,
         }
     }
+
+    /// Increments the message count for the given node, and returns the new count.
+    fn count_messages(&mut self, node_id: &N) -> usize {
+        let count = self.msg_count.entry(node_id.clone()).or_insert(0);
+        *count += 1;
+        *count
+    }
 }
 
 /// The contribution for the internal `HoneyBadger` instance: this includes a user-defined

src/fault_log.rs

@@ -53,8 +53,8 @@ pub enum FaultKind {
     UnexpectedKeyGenAck,
     /// `DynamicHoneyBadger` received a signed `Part` when no key generation in progress.
     UnexpectedKeyGenPart,
-    /// `DynamicHoneyBadger` received more key generation messages from the candidate than expected.
-    TooManyCandidateKeyGenMessages,
+    /// `DynamicHoneyBadger` received more key generation messages from the peer than expected.
+    TooManyKeyGenMessages,
     /// `DynamicHoneyBadger` received a message (Accept, Propose, or Change)
     /// with an invalid signature.
     IncorrectPayloadSignature,

src/sync_key_gen.rs

@@ -275,7 +275,7 @@ pub struct SyncKeyGen<N> {
     our_idx: Option<u64>,
     /// Our secret key.
     sec_key: SecretKey,
-    /// The public keys of all nodes, by node index.
+    /// The public keys of all nodes, by node ID.
     pub_keys: BTreeMap<N, PublicKey>,
     /// Proposed bivariate polynomials.
     parts: BTreeMap<u64, ProposalState>,
@@ -461,6 +461,11 @@ impl<N: NodeIdT> SyncKeyGen<N> {
         Ok(netinfo)
     }
 
+    /// Returns the number of nodes participating in the key generation.
+    pub fn num_nodes(&self) -> usize {
+        self.pub_keys.len()
+    }
+
     /// Handles an `Ack` message or returns an error string.
     fn handle_ack_or_err(
         &mut self,