feat: implement DKG message bouncing

src/consensus/dkg.rs

@@ -345,6 +345,20 @@ impl DkgVoter {
             session.timer_token = token;
         }
     }
+
+    // If this node participating in a DKG session with the given key?
+    pub fn is_participating(&self, dkg_key: &DkgKey) -> bool {
+        self.participant
+            .as_ref()
+            .map(|session| session.dkg_key == *dkg_key)
+            .unwrap_or(false)
+    }
+
+    pub fn observing_elders_info(&self, dkg_key: &DkgKey) -> Option<&EldersInfo> {
+        self.observers
+            .get(dkg_key)
+            .map(|session| &session.elders_info)
+    }
 }
 
 // Data for a DKG participant.

src/routing/approved.rs

@@ -421,6 +421,11 @@ impl Approved {
                     return Ok(MessageStatus::Unknown);
                 }
             }
+            Variant::DKGMessage { dkg_key, .. } => {
+                if !self.dkg_voter.is_participating(dkg_key) {
+                    return Ok(MessageStatus::Unknown);
+                }
+            }
             Variant::NodeApproval(_) | Variant::BootstrapResponse(_) => {
                 // Skip validation of these. We will validate them inside the bootstrap task.
                 return Ok(MessageStatus::Useful);
@@ -448,8 +453,7 @@ impl Approved {
             | Variant::Relocate(_)
             | Variant::BootstrapRequest(_)
             | Variant::BouncedUntrustedMessage(_)
-            | Variant::BouncedUnknownMessage { .. }
-            | Variant::DKGMessage { .. } => (),
+            | Variant::BouncedUnknownMessage { .. } => (),
         }
 
         if self.verify_message(msg)? {
@@ -701,6 +705,62 @@ impl Approved {
         sender: Peer,
         bounced_msg_bytes: Bytes,
         sender_last_key: &bls::PublicKey,
+    ) -> Result<Vec<Command>> {
+        let bounced_msg = Message::from_bytes(&bounced_msg_bytes)?;
+
+        if let Variant::DKGMessage { dkg_key, .. } = bounced_msg.variant() {
+            self.handle_bounced_unknown_dkg_message(sender, bounced_msg_bytes, dkg_key)
+        } else {
+            self.handle_bounced_unknown_other_message(sender, bounced_msg_bytes, sender_last_key)
+        }
+    }
+
+    fn handle_bounced_unknown_dkg_message(
+        &self,
+        sender: Peer,
+        bounced_msg_bytes: Bytes,
+        dkg_key: &DkgKey,
+    ) -> Result<Vec<Command>> {
+        let elders_info = if let Some(elders_info) = self.dkg_voter.observing_elders_info(dkg_key) {
+            trace!(
+                "Received BouncedUnknownMessage(DKGMessage) from {:?} \
+                     - resending with DKGStart",
+                sender
+            );
+            elders_info
+        } else {
+            trace!(
+                "Received BouncedUnknownMessage(DKGMessage) from {:?} \
+                     - peer is not a DKG participant, discarding",
+                sender
+            );
+            return Ok(vec![]);
+        };
+
+        // Normally a peer would bounce a DKG message if they don't have the corresponding DKG
+        // session which means they haven't yet reached the `DKGStart` consensus. Let's send the
+        // `DKGStart` again in case the previous one got lost and then resend the original message.
+
+        let variant = Variant::DKGStart {
+            dkg_key: *dkg_key,
+            elders_info: elders_info.clone(),
+        };
+        let vote = self.create_send_message_vote(DstLocation::Direct, variant, None)?;
+
+        let mut commands = self.send_vote(slice::from_ref(&sender), vote)?;
+        commands.push(Command::send_message_to_target(
+            sender.addr(),
+            bounced_msg_bytes,
+        ));
+
+        Ok(commands)
+    }
+
+    fn handle_bounced_unknown_other_message(
+        &self,
+        sender: Peer,
+        bounced_msg_bytes: Bytes,
+        sender_last_key: &bls::PublicKey,
     ) -> Result<Vec<Command>> {
         if !self.section.chain().has_key(sender_last_key)
             || sender_last_key == self.section.chain().last_key()