From 2cb6361242e3e26175ef760bd203cf1c96f3e457 Mon Sep 17 00:00:00 2001 From: Joey Yandle Date: Thu, 2 Apr 2026 12:05:09 +0200 Subject: [PATCH 1/8] add DkgPublicSharesDone and DkgPublicSharesDoneAck messages, to handle the case where in a p2p network the coordinator gets messages from signers before the signers get them from each other --- src/net.rs | 64 +++++++ src/state_machine/coordinator/fire.rs | 234 ++++++++++++++++++++++++- src/state_machine/coordinator/frost.rs | 63 ++++++- src/state_machine/coordinator/mod.rs | 85 ++++++++- src/state_machine/signer/mod.rs | 37 +++- 5 files changed, 468 insertions(+), 15 deletions(-) diff --git a/src/net.rs b/src/net.rs index c6c944f..cded00a 100644 --- a/src/net.rs +++ b/src/net.rs @@ -95,6 +95,10 @@ pub enum Message { DkgBegin(DkgBegin), /// Send DKG public shares DkgPublicShares(DkgPublicShares), + /// Tell signers the coordinator has received all expected public shares + DkgPublicSharesDone(DkgPublicSharesDone), + /// Acknowledge receipt of DkgPublicSharesDone + DkgPublicSharesDoneAck(DkgPublicSharesDoneAck), /// Tell signers to send DKG private shares DkgPrivateBegin(DkgPrivateBegin), /// Send DKG private shares @@ -118,6 +122,8 @@ impl Signable for Message { match self { Message::DkgBegin(msg) => msg.hash(hasher), Message::DkgPublicShares(msg) => msg.hash(hasher), + Message::DkgPublicSharesDone(msg) => msg.hash(hasher), + Message::DkgPublicSharesDoneAck(msg) => msg.hash(hasher), Message::DkgPrivateBegin(msg) => msg.hash(hasher), Message::DkgPrivateShares(msg) => msg.hash(hasher), Message::DkgEndBegin(msg) => msg.hash(hasher), @@ -234,6 +240,42 @@ impl Signable for DkgPublicShares { } } +#[derive(Clone, Serialize, Deserialize, Debug, PartialEq)] +/// DKG public shares done message from coordinator to signers +pub struct DkgPublicSharesDone { + /// DKG round ID + pub dkg_id: u64, + /// Signer IDs that the coordinator received public shares from + pub signer_ids: Vec, +} + +impl Signable for DkgPublicSharesDone { + fn hash(&self, hasher: &mut Sha256) { + hasher.update("DKG_PUBLIC_SHARES_DONE".as_bytes()); + hasher.update(self.dkg_id.to_be_bytes()); + for signer_id in &self.signer_ids { + hasher.update(signer_id.to_be_bytes()); + } + } +} + +#[derive(Clone, Serialize, Deserialize, Debug, PartialEq)] +/// DKG public shares done acknowledgment from signer to coordinator +pub struct DkgPublicSharesDoneAck { + /// DKG round ID + pub dkg_id: u64, + /// Signer ID + pub signer_id: u32, +} + +impl Signable for DkgPublicSharesDoneAck { + fn hash(&self, hasher: &mut Sha256) { + hasher.update("DKG_PUBLIC_SHARES_DONE_ACK".as_bytes()); + hasher.update(self.dkg_id.to_be_bytes()); + hasher.update(self.signer_id.to_be_bytes()); + } +} + #[derive(Clone, Serialize, Deserialize, Debug, PartialEq)] /// DKG private begin message from signer to all signers and coordinator pub struct DkgPrivateBegin { @@ -608,6 +650,28 @@ impl Packet { return false; } } + Message::DkgPublicSharesDone(msg) => { + if !msg.verify(&self.sig, coordinator_public_key) { + warn!("Received a DkgPublicSharesDone message with an invalid signature."); + return false; + } + } + Message::DkgPublicSharesDoneAck(msg) => { + if let Some(public_key) = signers_public_keys.signers.get(&msg.signer_id) { + if !msg.verify(&self.sig, public_key) { + warn!( + "Received a DkgPublicSharesDoneAck message with an invalid signature." + ); + return false; + } + } else { + warn!( + "Received a DkgPublicSharesDoneAck message with an unknown id: {}", + msg.signer_id + ); + return false; + } + } Message::DkgPrivateShares(msg) => { // Private shares have key IDs from [0, N) to reference IDs from [1, N] // in Frost V4 to enable easy indexing hence ID + 1 diff --git a/src/state_machine/coordinator/fire.rs b/src/state_machine/coordinator/fire.rs index fccadaf..5f491df 100644 --- a/src/state_machine/coordinator/fire.rs +++ b/src/state_machine/coordinator/fire.rs @@ -13,8 +13,8 @@ use crate::{ errors::AggregatorError, net::{ DkgBegin, DkgEnd, DkgEndBegin, DkgFailure, DkgPrivateBegin, DkgPrivateShares, - DkgPublicShares, DkgStatus, Message, NonceRequest, NonceResponse, Packet, Signable, - SignatureShareRequest, SignatureType, + DkgPublicShares, DkgPublicSharesDone, DkgStatus, Message, NonceRequest, NonceResponse, + Packet, Signable, SignatureShareRequest, SignatureType, }, state_machine::{ coordinator::{ @@ -95,13 +95,28 @@ impl Coordinator { // we hit the timeout but met the threshold, continue warn!("Timeout gathering DkgPublicShares for dkg round {} signing round {} iteration {}, dkg_threshold was met ({dkg_size}/{}), ", self.current_dkg_id, self.current_sign_id, self.current_sign_iter_id, self.config.dkg_threshold); self.public_shares_gathered()?; - let packet = self.start_private_shares()?; + let packet = self.send_public_shares_done()?; return Ok((Some(packet), None)); } } } } } + State::DkgPublicSharesDoneDistribute => {} + State::DkgPublicSharesDoneGather => { + if let Some(start) = self.dkg_public_start { + if let Some(timeout) = self.config.dkg_public_timeout { + if now.duration_since(start) > timeout { + error!("Timeout gathering DkgPublicSharesDoneAck for dkg round {}, not all signers responded", self.current_dkg_id); + let wait = self.dkg_wait_signer_ids.iter().copied().collect(); + return Ok(( + None, + Some(OperationResult::DkgError(DkgError::DkgPublicTimeout(wait))), + )); + } + } + } + } State::DkgPrivateDistribute => {} State::DkgPrivateGather => { if let Some(start) = self.dkg_private_start { @@ -263,6 +278,17 @@ impl Coordinator { return Ok((None, None)); } } + State::DkgPublicSharesDoneDistribute => { + let packet = self.send_public_shares_done()?; + return Ok((Some(packet), None)); + } + State::DkgPublicSharesDoneGather => { + self.gather_public_shares_done_ack(packet)?; + if self.state == State::DkgPublicSharesDoneGather { + // We need more data + return Ok((None, None)); + } + } State::DkgPrivateDistribute => { let packet = self.start_private_shares()?; return Ok((Some(packet), None)); @@ -528,7 +554,49 @@ impl Coordinator { } fn public_shares_gathered(&mut self) -> Result<(), Error> { - self.move_to(State::DkgPrivateDistribute)?; + self.move_to(State::DkgPublicSharesDoneDistribute)?; + Ok(()) + } + + /// Notify signers that all public shares have been received + pub fn send_public_shares_done(&mut self) -> Result { + let signer_ids: Vec = self.dkg_public_shares.keys().cloned().collect(); + self.dkg_wait_signer_ids = signer_ids.iter().cloned().collect(); + info!(dkg_id = %self.current_dkg_id, "Sending DkgPublicSharesDone"); + let msg = DkgPublicSharesDone { + dkg_id: self.current_dkg_id, + signer_ids, + }; + let packet = Packet { + sig: msg + .sign(&self.config.message_private_key) + .expect("Failed to sign DkgPublicSharesDone"), + msg: Message::DkgPublicSharesDone(msg), + }; + self.move_to(State::DkgPublicSharesDoneGather)?; + self.dkg_public_start = Some(Instant::now()); + Ok(packet) + } + + fn gather_public_shares_done_ack(&mut self, packet: &Packet) -> Result<(), Error> { + if let Message::DkgPublicSharesDoneAck(ack) = &packet.msg { + if ack.dkg_id != self.current_dkg_id { + return Err(Error::BadDkgId(ack.dkg_id, self.current_dkg_id)); + } + if !self.config.public_keys.signers.contains_key(&ack.signer_id) { + warn!(signer_id = %ack.signer_id, "No public key in config"); + return Ok(()); + } + self.dkg_wait_signer_ids.remove(&ack.signer_id); + debug!( + dkg_id = %ack.dkg_id, + signer_id = %ack.signer_id, + "DkgPublicSharesDoneAck received" + ); + } + if self.dkg_wait_signer_ids.is_empty() { + self.move_to(State::DkgPrivateDistribute)?; + } Ok(()) } @@ -1319,7 +1387,12 @@ impl StateMachine for Coordinator { prev_state == &State::DkgPublicDistribute || prev_state == &State::DkgPublicGather } - State::DkgPrivateDistribute => prev_state == &State::DkgPublicGather, + State::DkgPublicSharesDoneDistribute => prev_state == &State::DkgPublicGather, + State::DkgPublicSharesDoneGather => { + prev_state == &State::DkgPublicSharesDoneDistribute + || prev_state == &State::DkgPublicSharesDoneGather + } + State::DkgPrivateDistribute => prev_state == &State::DkgPublicSharesDoneGather, State::DkgPrivateGather => { prev_state == &State::DkgPrivateDistribute || prev_state == &State::DkgPrivateGather } @@ -2133,6 +2206,20 @@ pub mod test { let (outbound_messages, operation_results) = feedback_messages(&mut coordinators, &mut signers, &[message]); assert!(operation_results.is_empty()); + for coordinator in &coordinators { + assert_eq!(coordinator.state, State::DkgPublicSharesDoneGather); + } + + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPublicSharesDone(_)), + "Expected DkgPublicSharesDone message" + ); + + // Send DkgPublicSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); for coordinator in &coordinators { assert_eq!(coordinator.state, State::DkgPrivateGather); } @@ -2141,7 +2228,7 @@ pub mod test { assert_eq!(outbound_messages.len(), 1); assert!( matches!(&outbound_messages[0].msg, Message::DkgPrivateBegin(_)), - "Expected DkgPrviateBegin message" + "Expected DkgPrivateBegin message" ); // Send the DKG Private Begin message to all signers and share their responses with the coordinators and signers let (outbound_messages, operation_results) = @@ -2248,10 +2335,27 @@ pub mod test { assert!(outbound_messages.is_some()); assert!(operation_results.is_none()); + assert_eq!( + minimum_coordinators.first().unwrap().state, + State::DkgPublicSharesDoneGather, + ); + + // Feed DkgPublicSharesDone to signers, get acks back to coordinator + let (outbound_messages, operation_results) = feedback_messages( + &mut minimum_coordinators, + &mut minimum_signers, + &[outbound_messages.unwrap()], + ); + assert!(operation_results.is_empty()); assert_eq!( minimum_coordinators.first().unwrap().state, State::DkgPrivateGather, ); + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPrivateBegin(_)), + "Expected DkgPrivateBegin message" + ); (minimum_coordinators, minimum_signers) } @@ -2332,6 +2436,18 @@ pub mod test { assert!(outbound_messages.is_some()); assert!(operation_results.is_none()); + assert_eq!( + minimum_coordinators.first().unwrap().state, + State::DkgPublicSharesDoneGather, + ); + + // Feed DkgPublicSharesDone to signers, get acks back to coordinator + let (_outbound_messages, operation_results) = feedback_messages( + &mut minimum_coordinators, + &mut minimum_signers, + &[outbound_messages.unwrap()], + ); + assert!(operation_results.is_empty()); assert_eq!( minimum_coordinators.first().unwrap().state, State::DkgPrivateGather, @@ -2345,6 +2461,24 @@ pub mod test { let (outbound_messages, operation_results) = feedback_messages(&mut minimum_coordinators, &mut minimum_signers, &[message]); assert!(operation_results.is_empty()); + assert_eq!( + minimum_coordinators.first().unwrap().state, + State::DkgPublicSharesDoneGather + ); + + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPublicSharesDone(_)), + "Expected DkgPublicSharesDone message" + ); + + // Send DkgPublicSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = feedback_messages( + &mut minimum_coordinators, + &mut minimum_signers, + &outbound_messages, + ); + assert!(operation_results.is_empty()); assert_eq!( minimum_coordinators.first().unwrap().state, State::DkgPrivateGather @@ -2518,6 +2652,24 @@ pub mod test { &[message], ); assert!(operation_results.is_empty()); + assert_eq!( + insufficient_coordinator.first().unwrap().state, + State::DkgPublicSharesDoneGather + ); + + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPublicSharesDone(_)), + "Expected DkgPublicSharesDone message" + ); + + // Send DkgPublicSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = feedback_messages( + &mut insufficient_coordinator, + &mut insufficient_signers, + &outbound_messages, + ); + assert!(operation_results.is_empty()); assert_eq!( insufficient_coordinator.first().unwrap().state, State::DkgPrivateGather @@ -2605,6 +2757,20 @@ pub mod test { let (outbound_messages, operation_results) = feedback_messages(&mut coordinators, &mut signers, &[message]); assert!(operation_results.is_empty()); + for coordinator in &coordinators { + assert_eq!(coordinator.state, State::DkgPublicSharesDoneGather); + } + + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPublicSharesDone(_)), + "Expected DkgPublicSharesDone message" + ); + + // Send DkgPublicSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); for coordinator in &coordinators { assert_eq!(coordinator.state, State::DkgPrivateGather); } @@ -2762,6 +2928,20 @@ pub mod test { }, ); + assert!(operation_results.is_empty()); + for coordinator in &coordinators { + assert_eq!(coordinator.state, State::DkgPublicSharesDoneGather); + } + + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPublicSharesDone(_)), + "Expected DkgPublicSharesDone message" + ); + + // Send DkgPublicSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); assert!(operation_results.is_empty()); for coordinator in &coordinators { assert_eq!(coordinator.state, State::DkgPrivateGather); @@ -3071,6 +3251,20 @@ pub mod test { let (outbound_messages, operation_results) = feedback_messages(&mut coordinators, &mut signers, &[message]); assert!(operation_results.is_empty()); + for coordinator in &coordinators { + assert_eq!(coordinator.state, State::DkgPublicSharesDoneGather); + } + + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPublicSharesDone(_)), + "Expected DkgPublicSharesDone message" + ); + + // Send DkgPublicSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); for coordinator in &coordinators { assert_eq!(coordinator.state, State::DkgPrivateGather); } @@ -3548,6 +3742,20 @@ pub mod test { let (outbound_messages, operation_results) = feedback_messages(&mut coordinators, &mut signers, &[message]); assert!(operation_results.is_empty()); + for coordinator in coordinators.iter() { + assert_eq!(coordinator.get_state(), State::DkgPublicSharesDoneGather); + } + + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPublicSharesDone(_)), + "Expected DkgPublicSharesDone message" + ); + + // Send DkgPublicSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); for coordinator in coordinators.iter() { assert_eq!(coordinator.get_state(), State::DkgPrivateGather); } @@ -3661,6 +3869,20 @@ pub mod test { let (outbound_messages, operation_results) = feedback_messages(&mut coordinators, &mut signers, &[message]); assert!(operation_results.is_empty()); + for coordinator in coordinators.iter() { + assert_eq!(coordinator.get_state(), State::DkgPublicSharesDoneGather); + } + + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPublicSharesDone(_)), + "Expected DkgPublicSharesDone message" + ); + + // Send DkgPublicSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); for coordinator in coordinators.iter() { assert_eq!(coordinator.get_state(), State::DkgPrivateGather); } diff --git a/src/state_machine/coordinator/frost.rs b/src/state_machine/coordinator/frost.rs index 37a0012..837b144 100644 --- a/src/state_machine/coordinator/frost.rs +++ b/src/state_machine/coordinator/frost.rs @@ -8,8 +8,8 @@ use crate::{ curve::{ecdsa, point::Point}, net::{ DkgBegin, DkgEnd, DkgEndBegin, DkgPrivateBegin, DkgPrivateShares, DkgPublicShares, - DkgStatus, Message, NonceRequest, NonceResponse, Packet, Signable, SignatureShareRequest, - SignatureType, + DkgPublicSharesDone, DkgStatus, Message, NonceRequest, NonceResponse, Packet, Signable, + SignatureShareRequest, SignatureType, }, state_machine::{ coordinator::{ @@ -107,6 +107,17 @@ impl Coordinator { return Ok((None, None)); } } + State::DkgPublicSharesDoneDistribute => { + let packet = self.send_public_shares_done()?; + return Ok((Some(packet), None)); + } + State::DkgPublicSharesDoneGather => { + self.gather_public_shares_done_ack(packet)?; + if self.state == State::DkgPublicSharesDoneGather { + // We need more data + return Ok((None, None)); + } + } State::DkgPrivateDistribute => { let packet = self.start_private_shares()?; return Ok((Some(packet), None)); @@ -327,6 +338,47 @@ impl Coordinator { ); } + if self.ids_to_await.is_empty() { + self.move_to(State::DkgPublicSharesDoneDistribute)?; + } + Ok(()) + } + + /// Notify signers that all public shares have been received + pub fn send_public_shares_done(&mut self) -> Result { + let signer_ids: Vec = self.dkg_public_shares.keys().cloned().collect(); + self.ids_to_await = signer_ids.iter().cloned().collect(); + info!(dkg_id = %self.current_dkg_id, "Sending DkgPublicSharesDone"); + let msg = DkgPublicSharesDone { + dkg_id: self.current_dkg_id, + signer_ids, + }; + let packet = Packet { + sig: msg + .sign(&self.config.message_private_key) + .expect("Failed to sign DkgPublicSharesDone"), + msg: Message::DkgPublicSharesDone(msg), + }; + self.move_to(State::DkgPublicSharesDoneGather)?; + Ok(packet) + } + + fn gather_public_shares_done_ack(&mut self, packet: &Packet) -> Result<(), Error> { + if let Message::DkgPublicSharesDoneAck(ack) = &packet.msg { + if ack.dkg_id != self.current_dkg_id { + return Err(Error::BadDkgId(ack.dkg_id, self.current_dkg_id)); + } + if !self.config.public_keys.signers.contains_key(&ack.signer_id) { + warn!(signer_id = %ack.signer_id, "No public key in config"); + return Ok(()); + } + self.ids_to_await.remove(&ack.signer_id); + debug!( + dkg_id = %ack.dkg_id, + signer_id = %ack.signer_id, + "DkgPublicSharesDoneAck received" + ); + } if self.ids_to_await.is_empty() { self.move_to(State::DkgPrivateDistribute)?; } @@ -770,7 +822,12 @@ impl StateMachine for Coordinator { prev_state == &State::DkgPublicDistribute || prev_state == &State::DkgPublicGather } - State::DkgPrivateDistribute => prev_state == &State::DkgPublicGather, + State::DkgPublicSharesDoneDistribute => prev_state == &State::DkgPublicGather, + State::DkgPublicSharesDoneGather => { + prev_state == &State::DkgPublicSharesDoneDistribute + || prev_state == &State::DkgPublicSharesDoneGather + } + State::DkgPrivateDistribute => prev_state == &State::DkgPublicSharesDoneGather, State::DkgPrivateGather => { prev_state == &State::DkgPrivateDistribute || prev_state == &State::DkgPrivateGather } diff --git a/src/state_machine/coordinator/mod.rs b/src/state_machine/coordinator/mod.rs index 9db0cad..48763c6 100644 --- a/src/state_machine/coordinator/mod.rs +++ b/src/state_machine/coordinator/mod.rs @@ -28,6 +28,10 @@ pub enum State { DkgPublicDistribute, /// The coordinator is gathering public shares DkgPublicGather, + /// The coordinator is notifying signers that all public shares are done + DkgPublicSharesDoneDistribute, + /// The coordinator is gathering acknowledgments of DkgPublicSharesDone + DkgPublicSharesDoneGather, /// The coordinator is asking signers to send private shares DkgPrivateDistribute, /// The coordinator is gathering private shares @@ -444,13 +448,36 @@ pub mod test { .is_err()); assert!(coordinator.can_move_to(&State::DkgPublicGather).is_ok()); assert!(coordinator - .can_move_to(&State::DkgPrivateDistribute) + .can_move_to(&State::DkgPublicSharesDoneDistribute) .is_ok()); + assert!(coordinator + .can_move_to(&State::DkgPrivateDistribute) + .is_err()); assert!(coordinator.can_move_to(&State::DkgPrivateGather).is_err()); assert!(coordinator.can_move_to(&State::DkgEndDistribute).is_err()); assert!(coordinator.can_move_to(&State::DkgEndGather).is_err()); assert!(coordinator.can_move_to(&State::Idle).is_ok()); + coordinator + .move_to(State::DkgPublicSharesDoneDistribute) + .unwrap(); + assert!(coordinator + .can_move_to(&State::DkgPublicSharesDoneGather) + .is_ok()); + assert!(coordinator + .can_move_to(&State::DkgPrivateDistribute) + .is_err()); + + coordinator + .move_to(State::DkgPublicSharesDoneGather) + .unwrap(); + assert!(coordinator + .can_move_to(&State::DkgPublicSharesDoneGather) + .is_ok()); + assert!(coordinator + .can_move_to(&State::DkgPrivateDistribute) + .is_ok()); + coordinator.move_to(State::DkgPrivateDistribute).unwrap(); assert!(coordinator .can_move_to(&State::DkgPublicDistribute) @@ -766,6 +793,20 @@ pub mod test { let (outbound_messages, operation_results) = feedback_messages(&mut coordinators, &mut signers, &[message]); assert!(operation_results.is_empty()); + for coordinator in coordinators.iter() { + assert_eq!(coordinator.get_state(), State::DkgPublicSharesDoneGather); + } + + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPublicSharesDone(_)), + "Expected DkgPublicSharesDone message" + ); + + // Send DkgPublicSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); for coordinator in coordinators.iter() { assert_eq!(coordinator.get_state(), State::DkgPrivateGather); } @@ -1053,6 +1094,20 @@ pub mod test { let (outbound_messages, operation_results) = feedback_messages(&mut coordinators, &mut signers, &[message]); assert!(operation_results.is_empty()); + for coordinator in coordinators.iter() { + assert_eq!(coordinator.get_state(), State::DkgPublicSharesDoneGather); + } + + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPublicSharesDone(_)), + "Expected DkgPublicSharesDone message" + ); + + // Send DkgPublicSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); for coordinator in coordinators.iter() { assert_eq!(coordinator.get_state(), State::DkgPrivateGather); } @@ -1731,6 +1786,20 @@ pub mod test { }, ); assert!(operation_results.is_empty()); + for coordinator in coordinators.iter() { + assert_eq!(coordinator.get_state(), State::DkgPublicSharesDoneGather); + } + + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPublicSharesDone(_)), + "Expected DkgPublicSharesDone message" + ); + + // Send DkgPublicSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); for coordinator in coordinators.iter() { assert_eq!(coordinator.get_state(), State::DkgPrivateGather); } @@ -1809,6 +1878,20 @@ pub mod test { let (outbound_messages, operation_results) = feedback_messages(&mut coordinators, &mut signers, &[message]); assert!(operation_results.is_empty()); + for coordinator in coordinators.iter() { + assert_eq!(coordinator.get_state(), State::DkgPublicSharesDoneGather); + } + + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPublicSharesDone(_)), + "Expected DkgPublicSharesDone message" + ); + + // Send DkgPublicSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); for coordinator in coordinators.iter() { assert_eq!(coordinator.get_state(), State::DkgPrivateGather); } diff --git a/src/state_machine/signer/mod.rs b/src/state_machine/signer/mod.rs index 96fcb87..7ccdd92 100644 --- a/src/state_machine/signer/mod.rs +++ b/src/state_machine/signer/mod.rs @@ -20,8 +20,9 @@ use crate::{ errors::{DkgError, EncryptionError}, net::{ BadPrivateShare, DkgBegin, DkgEnd, DkgEndBegin, DkgFailure, DkgPrivateBegin, - DkgPrivateShares, DkgPublicShares, DkgStatus, Message, NonceRequest, NonceResponse, Packet, - SignatureShareRequest, SignatureShareResponse, SignatureType, + DkgPrivateShares, DkgPublicShares, DkgPublicSharesDone, DkgPublicSharesDoneAck, DkgStatus, + Message, NonceRequest, NonceResponse, Packet, SignatureShareRequest, + SignatureShareResponse, SignatureType, }, state_machine::{PublicKeys, StateMachine}, traits::{Signer as SignerTrait, SignerState as SignerSavedState}, @@ -479,6 +480,7 @@ impl Signer { } Message::DkgEndBegin(dkg_end_begin) => self.dkg_end_begin(dkg_end_begin), Message::DkgPublicShares(dkg_public_shares) => self.dkg_public_share(dkg_public_shares), + Message::DkgPublicSharesDone(msg) => self.dkg_public_shares_done(msg), Message::DkgPrivateShares(dkg_private_shares) => { self.dkg_private_shares(dkg_private_shares, rng) } @@ -486,9 +488,10 @@ impl Signer { self.sign_share_request(sign_share_request, rng) } Message::NonceRequest(nonce_request) => self.nonce_request(nonce_request, rng), - Message::DkgEnd(_) | Message::NonceResponse(_) | Message::SignatureShareResponse(_) => { - Ok(vec![]) - } // TODO + Message::DkgEnd(_) + | Message::DkgPublicSharesDoneAck(_) + | Message::NonceResponse(_) + | Message::SignatureShareResponse(_) => Ok(vec![]), // TODO }; match out_msgs { @@ -858,6 +861,30 @@ impl Signer { self.dkg_public_begin(rng) } + fn dkg_public_shares_done(&mut self, msg: &DkgPublicSharesDone) -> Result, Error> { + if msg.dkg_id != self.dkg_id { + warn!( + signer_id = %self.signer_id, + got = %msg.dkg_id, + expected = %self.dkg_id, + "DkgPublicSharesDone dkg_id mismatch" + ); + return Ok(vec![]); + } + if !msg.signer_ids.contains(&self.signer_id) { + warn!( + signer_id = %self.signer_id, + "signer_id not in DkgPublicSharesDone, coordinator did not receive our public shares" + ); + return Ok(vec![]); + } + let ack = DkgPublicSharesDoneAck { + dkg_id: self.dkg_id, + signer_id: self.signer_id, + }; + Ok(vec![Message::DkgPublicSharesDoneAck(ack)]) + } + fn dkg_public_begin( &mut self, rng: &mut R, From b9988efbcca6887463b4c29b5e95bc2c19d628ad Mon Sep 17 00:00:00 2001 From: Joey Yandle Date: Thu, 2 Apr 2026 12:21:39 +0200 Subject: [PATCH 2/8] add DkgPrivateSharesDone and DkgPrivateSharesDoneAck messages and handle them in code and tests --- src/net.rs | 64 ++++++++++ src/state_machine/coordinator/fire.rs | 164 +++++++++++++++++++++++-- src/state_machine/coordinator/frost.rs | 65 +++++++++- src/state_machine/coordinator/mod.rs | 67 +++++++++- src/state_machine/signer/mod.rs | 35 +++++- 5 files changed, 380 insertions(+), 15 deletions(-) diff --git a/src/net.rs b/src/net.rs index cded00a..329494a 100644 --- a/src/net.rs +++ b/src/net.rs @@ -103,6 +103,10 @@ pub enum Message { DkgPrivateBegin(DkgPrivateBegin), /// Send DKG private shares DkgPrivateShares(DkgPrivateShares), + /// Tell signers the coordinator has received all expected private shares + DkgPrivateSharesDone(DkgPrivateSharesDone), + /// Acknowledge receipt of DkgPrivateSharesDone + DkgPrivateSharesDoneAck(DkgPrivateSharesDoneAck), /// Tell signers to compute shares and send DKG end DkgEndBegin(DkgEndBegin), /// Tell coordinator that DKG is complete @@ -126,6 +130,8 @@ impl Signable for Message { Message::DkgPublicSharesDoneAck(msg) => msg.hash(hasher), Message::DkgPrivateBegin(msg) => msg.hash(hasher), Message::DkgPrivateShares(msg) => msg.hash(hasher), + Message::DkgPrivateSharesDone(msg) => msg.hash(hasher), + Message::DkgPrivateSharesDoneAck(msg) => msg.hash(hasher), Message::DkgEndBegin(msg) => msg.hash(hasher), Message::DkgEnd(msg) => msg.hash(hasher), Message::NonceRequest(msg) => msg.hash(hasher), @@ -329,6 +335,42 @@ impl Signable for DkgPrivateShares { } } +#[derive(Clone, Serialize, Deserialize, Debug, PartialEq)] +/// DKG private shares done message from coordinator to signers +pub struct DkgPrivateSharesDone { + /// DKG round ID + pub dkg_id: u64, + /// Signer IDs that the coordinator received private shares from + pub signer_ids: Vec, +} + +impl Signable for DkgPrivateSharesDone { + fn hash(&self, hasher: &mut Sha256) { + hasher.update("DKG_PRIVATE_SHARES_DONE".as_bytes()); + hasher.update(self.dkg_id.to_be_bytes()); + for signer_id in &self.signer_ids { + hasher.update(signer_id.to_be_bytes()); + } + } +} + +#[derive(Clone, Serialize, Deserialize, Debug, PartialEq)] +/// DKG private shares done acknowledgment from signer to coordinator +pub struct DkgPrivateSharesDoneAck { + /// DKG round ID + pub dkg_id: u64, + /// Signer ID + pub signer_id: u32, +} + +impl Signable for DkgPrivateSharesDoneAck { + fn hash(&self, hasher: &mut Sha256) { + hasher.update("DKG_PRIVATE_SHARES_DONE_ACK".as_bytes()); + hasher.update(self.dkg_id.to_be_bytes()); + hasher.update(self.signer_id.to_be_bytes()); + } +} + #[derive(Clone, Serialize, Deserialize, Debug, PartialEq)] /// DKG end begin message from signer to all signers and coordinator pub struct DkgEndBegin { @@ -689,6 +731,28 @@ impl Packet { return false; } } + Message::DkgPrivateSharesDone(msg) => { + if !msg.verify(&self.sig, coordinator_public_key) { + warn!("Received a DkgPrivateSharesDone message with an invalid signature."); + return false; + } + } + Message::DkgPrivateSharesDoneAck(msg) => { + if let Some(public_key) = signers_public_keys.signers.get(&msg.signer_id) { + if !msg.verify(&self.sig, public_key) { + warn!( + "Received a DkgPrivateSharesDoneAck message with an invalid signature." + ); + return false; + } + } else { + warn!( + "Received a DkgPrivateSharesDoneAck message with an unknown id: {}", + msg.signer_id + ); + return false; + } + } Message::NonceRequest(msg) => { if !msg.verify(&self.sig, coordinator_public_key) { warn!("Received a NonceRequest message with an invalid signature."); diff --git a/src/state_machine/coordinator/fire.rs b/src/state_machine/coordinator/fire.rs index 5f491df..3c9fd36 100644 --- a/src/state_machine/coordinator/fire.rs +++ b/src/state_machine/coordinator/fire.rs @@ -13,8 +13,8 @@ use crate::{ errors::AggregatorError, net::{ DkgBegin, DkgEnd, DkgEndBegin, DkgFailure, DkgPrivateBegin, DkgPrivateShares, - DkgPublicShares, DkgPublicSharesDone, DkgStatus, Message, NonceRequest, NonceResponse, - Packet, Signable, SignatureShareRequest, SignatureType, + DkgPrivateSharesDone, DkgPublicShares, DkgPublicSharesDone, DkgStatus, Message, + NonceRequest, NonceResponse, Packet, Signable, SignatureShareRequest, SignatureType, }, state_machine::{ coordinator::{ @@ -138,13 +138,28 @@ impl Coordinator { // we hit the timeout but met the threshold, continue warn!("Timeout gathering DkgPrivateShares for dkg round {} signing round {} iteration {}, dkg_threshold was met ({dkg_size}/{}), ", self.current_dkg_id, self.current_sign_id, self.current_sign_iter_id, self.config.dkg_threshold); self.private_shares_gathered()?; - let packet = self.start_dkg_end()?; + let packet = self.send_private_shares_done()?; return Ok((Some(packet), None)); } } } } } + State::DkgPrivateSharesDoneDistribute => {} + State::DkgPrivateSharesDoneGather => { + if let Some(start) = self.dkg_private_start { + if let Some(timeout) = self.config.dkg_private_timeout { + if now.duration_since(start) > timeout { + error!("Timeout gathering DkgPrivateSharesDoneAck for dkg round {}, not all signers responded", self.current_dkg_id); + let wait = self.dkg_wait_signer_ids.iter().copied().collect(); + return Ok(( + None, + Some(OperationResult::DkgError(DkgError::DkgPrivateTimeout(wait))), + )); + } + } + } + } State::DkgEndDistribute => {} State::DkgEndGather => { if let Some(start) = self.dkg_end_start { @@ -300,6 +315,17 @@ impl Coordinator { return Ok((None, None)); } } + State::DkgPrivateSharesDoneDistribute => { + let packet = self.send_private_shares_done()?; + return Ok((Some(packet), None)); + } + State::DkgPrivateSharesDoneGather => { + self.gather_private_shares_done_ack(packet)?; + if self.state == State::DkgPrivateSharesDoneGather { + // We need more data + return Ok((None, None)); + } + } State::DkgEndDistribute => { let packet = self.start_dkg_end()?; return Ok((Some(packet), None)); @@ -652,7 +678,49 @@ impl Coordinator { } fn private_shares_gathered(&mut self) -> Result<(), Error> { - self.move_to(State::DkgEndDistribute)?; + self.move_to(State::DkgPrivateSharesDoneDistribute)?; + Ok(()) + } + + /// Notify signers that all private shares have been received + pub fn send_private_shares_done(&mut self) -> Result { + let signer_ids: Vec = self.dkg_private_shares.keys().cloned().collect(); + self.dkg_wait_signer_ids = signer_ids.iter().cloned().collect(); + info!(dkg_id = %self.current_dkg_id, "Sending DkgPrivateSharesDone"); + let msg = DkgPrivateSharesDone { + dkg_id: self.current_dkg_id, + signer_ids, + }; + let packet = Packet { + sig: msg + .sign(&self.config.message_private_key) + .expect("Failed to sign DkgPrivateSharesDone"), + msg: Message::DkgPrivateSharesDone(msg), + }; + self.move_to(State::DkgPrivateSharesDoneGather)?; + self.dkg_private_start = Some(Instant::now()); + Ok(packet) + } + + fn gather_private_shares_done_ack(&mut self, packet: &Packet) -> Result<(), Error> { + if let Message::DkgPrivateSharesDoneAck(ack) = &packet.msg { + if ack.dkg_id != self.current_dkg_id { + return Err(Error::BadDkgId(ack.dkg_id, self.current_dkg_id)); + } + if !self.config.public_keys.signers.contains_key(&ack.signer_id) { + warn!(signer_id = %ack.signer_id, "No public key in config"); + return Ok(()); + } + self.dkg_wait_signer_ids.remove(&ack.signer_id); + debug!( + dkg_id = %ack.dkg_id, + signer_id = %ack.signer_id, + "DkgPrivateSharesDoneAck received" + ); + } + if self.dkg_wait_signer_ids.is_empty() { + self.move_to(State::DkgEndDistribute)?; + } Ok(()) } @@ -1396,7 +1464,12 @@ impl StateMachine for Coordinator { prev_state == &State::DkgPrivateDistribute || prev_state == &State::DkgPrivateGather } - State::DkgEndDistribute => prev_state == &State::DkgPrivateGather, + State::DkgPrivateSharesDoneDistribute => prev_state == &State::DkgPrivateGather, + State::DkgPrivateSharesDoneGather => { + prev_state == &State::DkgPrivateSharesDoneDistribute + || prev_state == &State::DkgPrivateSharesDoneGather + } + State::DkgEndDistribute => prev_state == &State::DkgPrivateSharesDoneGather, State::DkgEndGather => prev_state == &State::DkgEndDistribute, State::NonceRequest(signature_type) => { prev_state == &State::Idle @@ -2235,6 +2308,16 @@ pub mod test { feedback_messages(&mut coordinators, &mut signers, &outbound_messages); assert!(operation_results.is_empty()); assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(outbound_messages[0].msg, Message::DkgPrivateSharesDone(_)), + "Expected DkgPrivateSharesDone message" + ); + + // Send DkgPrivateSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); + assert_eq!(outbound_messages.len(), 1); assert!( matches!(outbound_messages[0].msg, Message::DkgEndBegin(_)), "Expected DkgEndBegin message" @@ -2520,8 +2603,25 @@ pub mod test { assert!( matches!( outbound_message.clone().unwrap().msg, - Message::DkgEndBegin(_) + Message::DkgPrivateSharesDone(_) ), + "Expected DkgPrivateSharesDone message" + ); + assert_eq!( + minimum_coordinators.first().unwrap().state, + State::DkgPrivateSharesDoneGather, + ); + + // Send DkgPrivateSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = feedback_messages( + &mut minimum_coordinators, + &mut minimum_signers, + &[outbound_message.unwrap()], + ); + assert!(operation_results.is_empty()); + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(outbound_messages[0].msg, Message::DkgEndBegin(_)), "Expected DkgEndBegin message" ); assert_eq!( @@ -2533,7 +2633,7 @@ pub mod test { let (outbound_messages, operation_results) = feedback_messages( &mut minimum_coordinators, &mut minimum_signers, - &[outbound_message.unwrap()], + &outbound_messages, ); assert!(outbound_messages.is_empty()); assert_eq!(operation_results.len(), 1); @@ -2828,6 +2928,16 @@ pub mod test { ); assert!(operation_results.is_empty()); assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(outbound_messages[0].msg, Message::DkgPrivateSharesDone(_)), + "Expected DkgPrivateSharesDone message" + ); + + // Send DkgPrivateSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); + assert_eq!(outbound_messages.len(), 1); assert!( matches!(outbound_messages[0].msg, Message::DkgEndBegin(_)), "Expected DkgEndBegin message" @@ -2953,6 +3063,16 @@ pub mod test { "Expected DkgPrivateBegin message" ); + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); + assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(outbound_messages[0].msg, Message::DkgPrivateSharesDone(_)), + "Expected DkgPrivateSharesDone message" + ); + + // Send DkgPrivateSharesDone to signers and collect their acks back to the coordinator let (outbound_messages, operation_results) = feedback_messages(&mut coordinators, &mut signers, &outbound_messages); assert!(operation_results.is_empty()); @@ -3280,6 +3400,16 @@ pub mod test { feedback_messages(&mut coordinators, &mut signers, &outbound_messages); assert!(operation_results.is_empty()); assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPrivateSharesDone(_)), + "Expected DkgPrivateSharesDone message" + ); + + // Send DkgPrivateSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); + assert_eq!(outbound_messages.len(), 1); assert!( matches!(&outbound_messages[0].msg, Message::DkgEndBegin(_)), "Expected DkgEndBegin message" @@ -3771,6 +3901,16 @@ pub mod test { feedback_messages(&mut coordinators, &mut signers, &outbound_messages); assert!(operation_results.is_empty()); assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(outbound_messages[0].msg, Message::DkgPrivateSharesDone(_)), + "Expected DkgPrivateSharesDone message" + ); + + // Send DkgPrivateSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); + assert_eq!(outbound_messages.len(), 1); assert!( matches!(outbound_messages[0].msg, Message::DkgEndBegin(_)), "Expected DkgEndBegin message" @@ -3898,6 +4038,16 @@ pub mod test { feedback_messages(&mut coordinators, &mut signers, &outbound_messages); assert!(operation_results.is_empty()); assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPrivateSharesDone(_)), + "Expected DkgPrivateSharesDone message" + ); + + // Send DkgPrivateSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert!(operation_results.is_empty()); + assert_eq!(outbound_messages.len(), 1); assert!( matches!(&outbound_messages[0].msg, Message::DkgEndBegin(_)), "Expected DkgEndBegin message" diff --git a/src/state_machine/coordinator/frost.rs b/src/state_machine/coordinator/frost.rs index 837b144..9fc197c 100644 --- a/src/state_machine/coordinator/frost.rs +++ b/src/state_machine/coordinator/frost.rs @@ -7,9 +7,9 @@ use crate::{ compute, curve::{ecdsa, point::Point}, net::{ - DkgBegin, DkgEnd, DkgEndBegin, DkgPrivateBegin, DkgPrivateShares, DkgPublicShares, - DkgPublicSharesDone, DkgStatus, Message, NonceRequest, NonceResponse, Packet, Signable, - SignatureShareRequest, SignatureType, + DkgBegin, DkgEnd, DkgEndBegin, DkgPrivateBegin, DkgPrivateShares, DkgPrivateSharesDone, + DkgPublicShares, DkgPublicSharesDone, DkgStatus, Message, NonceRequest, NonceResponse, + Packet, Signable, SignatureShareRequest, SignatureType, }, state_machine::{ coordinator::{ @@ -129,6 +129,17 @@ impl Coordinator { return Ok((None, None)); } } + State::DkgPrivateSharesDoneDistribute => { + let packet = self.send_private_shares_done()?; + return Ok((Some(packet), None)); + } + State::DkgPrivateSharesDoneGather => { + self.gather_private_shares_done_ack(packet)?; + if self.state == State::DkgPrivateSharesDoneGather { + // We need more data + return Ok((None, None)); + } + } State::DkgEndDistribute => { let packet = self.start_dkg_end()?; return Ok((Some(packet), None)); @@ -420,6 +431,47 @@ impl Coordinator { ); } + if self.ids_to_await.is_empty() { + self.move_to(State::DkgPrivateSharesDoneDistribute)?; + } + Ok(()) + } + + /// Notify signers that all private shares have been received + pub fn send_private_shares_done(&mut self) -> Result { + let signer_ids: Vec = self.dkg_private_shares.keys().cloned().collect(); + self.ids_to_await = signer_ids.iter().cloned().collect(); + info!(dkg_id = %self.current_dkg_id, "Sending DkgPrivateSharesDone"); + let msg = DkgPrivateSharesDone { + dkg_id: self.current_dkg_id, + signer_ids, + }; + let packet = Packet { + sig: msg + .sign(&self.config.message_private_key) + .expect("Failed to sign DkgPrivateSharesDone"), + msg: Message::DkgPrivateSharesDone(msg), + }; + self.move_to(State::DkgPrivateSharesDoneGather)?; + Ok(packet) + } + + fn gather_private_shares_done_ack(&mut self, packet: &Packet) -> Result<(), Error> { + if let Message::DkgPrivateSharesDoneAck(ack) = &packet.msg { + if ack.dkg_id != self.current_dkg_id { + return Err(Error::BadDkgId(ack.dkg_id, self.current_dkg_id)); + } + if !self.config.public_keys.signers.contains_key(&ack.signer_id) { + warn!(signer_id = %ack.signer_id, "No public key in config"); + return Ok(()); + } + self.ids_to_await.remove(&ack.signer_id); + debug!( + dkg_id = %ack.dkg_id, + signer_id = %ack.signer_id, + "DkgPrivateSharesDoneAck received" + ); + } if self.ids_to_await.is_empty() { self.move_to(State::DkgEndDistribute)?; } @@ -831,7 +883,12 @@ impl StateMachine for Coordinator { prev_state == &State::DkgPrivateDistribute || prev_state == &State::DkgPrivateGather } - State::DkgEndDistribute => prev_state == &State::DkgPrivateGather, + State::DkgPrivateSharesDoneDistribute => prev_state == &State::DkgPrivateGather, + State::DkgPrivateSharesDoneGather => { + prev_state == &State::DkgPrivateSharesDoneDistribute + || prev_state == &State::DkgPrivateSharesDoneGather + } + State::DkgEndDistribute => prev_state == &State::DkgPrivateSharesDoneGather, State::DkgEndGather => prev_state == &State::DkgEndDistribute, State::NonceRequest(_) => { prev_state == &State::Idle || prev_state == &State::DkgEndGather diff --git a/src/state_machine/coordinator/mod.rs b/src/state_machine/coordinator/mod.rs index 48763c6..af351ba 100644 --- a/src/state_machine/coordinator/mod.rs +++ b/src/state_machine/coordinator/mod.rs @@ -36,6 +36,10 @@ pub enum State { DkgPrivateDistribute, /// The coordinator is gathering private shares DkgPrivateGather, + /// The coordinator is notifying signers that all private shares are done + DkgPrivateSharesDoneDistribute, + /// The coordinator is gathering acknowledgments of DkgPrivateSharesDone + DkgPrivateSharesDoneGather, /// The coordinator is asking signers to compute shares and send end DkgEndDistribute, /// The coordinator is gathering DKG End messages @@ -500,10 +504,29 @@ pub mod test { .can_move_to(&State::DkgPrivateDistribute) .is_err()); assert!(coordinator.can_move_to(&State::DkgPrivateGather).is_ok()); - assert!(coordinator.can_move_to(&State::DkgEndDistribute).is_ok()); + assert!(coordinator + .can_move_to(&State::DkgPrivateSharesDoneDistribute) + .is_ok()); + assert!(coordinator.can_move_to(&State::DkgEndDistribute).is_err()); assert!(coordinator.can_move_to(&State::DkgEndGather).is_err()); assert!(coordinator.can_move_to(&State::Idle).is_ok()); + coordinator + .move_to(State::DkgPrivateSharesDoneDistribute) + .unwrap(); + assert!(coordinator + .can_move_to(&State::DkgPrivateSharesDoneGather) + .is_ok()); + assert!(coordinator.can_move_to(&State::DkgEndDistribute).is_err()); + + coordinator + .move_to(State::DkgPrivateSharesDoneGather) + .unwrap(); + assert!(coordinator + .can_move_to(&State::DkgPrivateSharesDoneGather) + .is_ok()); + assert!(coordinator.can_move_to(&State::DkgEndDistribute).is_ok()); + coordinator.move_to(State::DkgEndDistribute).unwrap(); assert!(coordinator.can_move_to(&State::DkgEndGather).is_ok()); @@ -841,6 +864,16 @@ pub mod test { feedback_messages(&mut coordinators, &mut signers, &outbound_messages); assert_eq!(operation_results.len(), 0); assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(outbound_messages[0].msg, Message::DkgPrivateSharesDone(_)), + "Expected DkgPrivateSharesDone message" + ); + + // Send DkgPrivateSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert_eq!(operation_results.len(), 0); + assert_eq!(outbound_messages.len(), 1); assert!( matches!(outbound_messages[0].msg, Message::DkgEndBegin(_)), "Expected DkgEndBegin message" @@ -1125,6 +1158,18 @@ pub mod test { feedback_messages(&mut coordinators, &mut signers, &outbound_messages); assert_eq!(operation_results.len(), 0); assert_eq!(outbound_messages.len(), 1); + match &outbound_messages[0].msg { + Message::DkgPrivateSharesDone(_) => {} + _ => { + panic!("Expected DkgPrivateSharesDone message"); + } + } + + // Send DkgPrivateSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert_eq!(operation_results.len(), 0); + assert_eq!(outbound_messages.len(), 1); match &outbound_messages[0].msg { Message::DkgEndBegin(_) => {} _ => { @@ -1814,6 +1859,16 @@ pub mod test { feedback_messages(&mut coordinators, &mut signers, &outbound_messages); assert_eq!(operation_results.len(), 0); assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(outbound_messages[0].msg, Message::DkgPrivateSharesDone(_)), + "Expected DkgPrivateSharesDone message" + ); + + // Send DkgPrivateSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert_eq!(operation_results.len(), 0); + assert_eq!(outbound_messages.len(), 1); assert!( matches!(outbound_messages[0].msg, Message::DkgEndBegin(_)), "Expected DkgEndBegin message" @@ -1932,6 +1987,16 @@ pub mod test { ); assert_eq!(operation_results.len(), 0); assert_eq!(outbound_messages.len(), 1); + assert!( + matches!(&outbound_messages[0].msg, Message::DkgPrivateSharesDone(_)), + "Expected DkgPrivateSharesDone message" + ); + + // Send DkgPrivateSharesDone to signers and collect their acks back to the coordinator + let (outbound_messages, operation_results) = + feedback_messages(&mut coordinators, &mut signers, &outbound_messages); + assert_eq!(operation_results.len(), 0); + assert_eq!(outbound_messages.len(), 1); assert!( matches!(&outbound_messages[0].msg, Message::DkgEndBegin(_)), "Expected DkgEndBegin message" diff --git a/src/state_machine/signer/mod.rs b/src/state_machine/signer/mod.rs index 7ccdd92..d460ce2 100644 --- a/src/state_machine/signer/mod.rs +++ b/src/state_machine/signer/mod.rs @@ -20,9 +20,9 @@ use crate::{ errors::{DkgError, EncryptionError}, net::{ BadPrivateShare, DkgBegin, DkgEnd, DkgEndBegin, DkgFailure, DkgPrivateBegin, - DkgPrivateShares, DkgPublicShares, DkgPublicSharesDone, DkgPublicSharesDoneAck, DkgStatus, - Message, NonceRequest, NonceResponse, Packet, SignatureShareRequest, - SignatureShareResponse, SignatureType, + DkgPrivateShares, DkgPrivateSharesDone, DkgPrivateSharesDoneAck, DkgPublicShares, + DkgPublicSharesDone, DkgPublicSharesDoneAck, DkgStatus, Message, NonceRequest, + NonceResponse, Packet, SignatureShareRequest, SignatureShareResponse, SignatureType, }, state_machine::{PublicKeys, StateMachine}, traits::{Signer as SignerTrait, SignerState as SignerSavedState}, @@ -484,12 +484,14 @@ impl Signer { Message::DkgPrivateShares(dkg_private_shares) => { self.dkg_private_shares(dkg_private_shares, rng) } + Message::DkgPrivateSharesDone(msg) => self.dkg_private_shares_done(msg), Message::SignatureShareRequest(sign_share_request) => { self.sign_share_request(sign_share_request, rng) } Message::NonceRequest(nonce_request) => self.nonce_request(nonce_request, rng), Message::DkgEnd(_) | Message::DkgPublicSharesDoneAck(_) + | Message::DkgPrivateSharesDoneAck(_) | Message::NonceResponse(_) | Message::SignatureShareResponse(_) => Ok(vec![]), // TODO }; @@ -885,6 +887,33 @@ impl Signer { Ok(vec![Message::DkgPublicSharesDoneAck(ack)]) } + fn dkg_private_shares_done( + &mut self, + msg: &DkgPrivateSharesDone, + ) -> Result, Error> { + if msg.dkg_id != self.dkg_id { + warn!( + signer_id = %self.signer_id, + got = %msg.dkg_id, + expected = %self.dkg_id, + "DkgPrivateSharesDone dkg_id mismatch" + ); + return Ok(vec![]); + } + if !msg.signer_ids.contains(&self.signer_id) { + warn!( + signer_id = %self.signer_id, + "signer_id not in DkgPrivateSharesDone, coordinator did not receive our private shares" + ); + return Ok(vec![]); + } + let ack = DkgPrivateSharesDoneAck { + dkg_id: self.dkg_id, + signer_id: self.signer_id, + }; + Ok(vec![Message::DkgPrivateSharesDoneAck(ack)]) + } + fn dkg_public_begin( &mut self, rng: &mut R, From 93a4da5394e062055aa16b99d6969d526bda76b2 Mon Sep 17 00:00:00 2001 From: Joey Yandle Date: Thu, 2 Apr 2026 17:07:50 +0200 Subject: [PATCH 3/8] fix signer state machine to only accept messages in the proper state --- src/state_machine/coordinator/mod.rs | 22 +++-- src/state_machine/signer/mod.rs | 135 +++++++++++++++++++++------ 2 files changed, 118 insertions(+), 39 deletions(-) diff --git a/src/state_machine/coordinator/mod.rs b/src/state_machine/coordinator/mod.rs index af351ba..d6ef564 100644 --- a/src/state_machine/coordinator/mod.rs +++ b/src/state_machine/coordinator/mod.rs @@ -1586,11 +1586,7 @@ pub mod test { "Expected SignatureShareRequest message" ); - let messages = outbound_messages.clone(); - let result = feedback_messages_with_errors(&mut coordinators, &mut signers, &messages); - assert!(result.is_ok()); - - // test request with no NonceResponses + // test request with no NonceResponses — signers are in SignGather, errors before state change let mut packet = outbound_messages[0].clone(); let Message::SignatureShareRequest(ref mut request) = packet.msg else { panic!("failed to match message"); @@ -1645,6 +1641,11 @@ pub mod test { ), "Should have received signer invalid nonce response error, got {result:?}" ); + + // send valid SSR last — signers are still in SignGather after all error cases + let messages = outbound_messages.clone(); + let result = feedback_messages_with_errors(&mut coordinators, &mut signers, &messages); + assert!(result.is_ok()); } pub fn invalid_nonce( @@ -1685,11 +1686,7 @@ pub mod test { "Expected SignatureShareRequest message" ); - let messages = outbound_messages.clone(); - let result = feedback_messages_with_errors(&mut coordinators, &mut signers, &messages); - assert!(result.is_ok()); - - // test request with NonceResponse having zero nonce + // test request with NonceResponse having zero nonce — signers are in SignGather, errors before state change let mut packet = outbound_messages[0].clone(); let Message::SignatureShareRequest(ref mut request) = packet.msg else { panic!("failed to match message"); @@ -1772,6 +1769,11 @@ pub mod test { ), "Should have received signer invalid nonce response error, got {result:?}" ); + + // send valid SSR last — signers are still in SignGather after all error cases + let messages = outbound_messages.clone(); + let result = feedback_messages_with_errors(&mut coordinators, &mut signers, &messages); + assert!(result.is_ok()); } pub fn empty_public_shares( diff --git a/src/state_machine/signer/mod.rs b/src/state_machine/signer/mod.rs index d460ce2..3a44eba 100644 --- a/src/state_machine/signer/mod.rs +++ b/src/state_machine/signer/mod.rs @@ -41,11 +41,15 @@ pub enum State { DkgPublicDistribute, /// The signer is gathering DKG public shares DkgPublicGather, + /// The signer has acknowledged DkgPublicSharesDone and is waiting for DkgPrivateBegin + DkgPublicSharesDoneAck, /// The signer is distributing DKG private shares DkgPrivateDistribute, /// The signer is gathering DKG private shares DkgPrivateGather, - /// The signer is distributing signature shares + /// The signer has acknowledged DkgPrivateSharesDone and is waiting for DkgEndBegin + DkgPrivateSharesDoneAck, + /// The signer has sent a nonce and is waiting for a signature share request SignGather, } @@ -473,27 +477,60 @@ impl Signer { return Err(Error::InvalidPacketSignature); } } - let out_msgs = match &packet.msg { - Message::DkgBegin(dkg_begin) => self.dkg_begin(dkg_begin, rng), - Message::DkgPrivateBegin(dkg_private_begin) => { - self.dkg_private_begin(dkg_private_begin, rng) + let out_msgs = match (&self.state, &packet.msg) { + // DkgBegin restarts DKG from any state + (_, Message::DkgBegin(msg)) => self.dkg_begin(msg, rng), + // DKG public phase + (State::DkgPublicGather, Message::DkgPublicShares(msg)) => self.dkg_public_share(msg), + // Late-arriving public shares are still stored even after we've acked done + (State::DkgPublicSharesDoneAck, Message::DkgPublicShares(msg)) => { + self.dkg_public_share(msg) } - Message::DkgEndBegin(dkg_end_begin) => self.dkg_end_begin(dkg_end_begin), - Message::DkgPublicShares(dkg_public_shares) => self.dkg_public_share(dkg_public_shares), - Message::DkgPublicSharesDone(msg) => self.dkg_public_shares_done(msg), - Message::DkgPrivateShares(dkg_private_shares) => { - self.dkg_private_shares(dkg_private_shares, rng) + (State::DkgPublicGather, Message::DkgPublicSharesDone(msg)) => { + self.dkg_public_shares_done(msg) } - Message::DkgPrivateSharesDone(msg) => self.dkg_private_shares_done(msg), - Message::SignatureShareRequest(sign_share_request) => { - self.sign_share_request(sign_share_request, rng) + // DKG private phase + (State::DkgPublicSharesDoneAck, Message::DkgPrivateBegin(msg)) => { + self.dkg_private_begin(msg, rng) + } + (State::DkgPrivateGather, Message::DkgPrivateShares(msg)) => { + self.dkg_private_shares(msg, rng) + } + // Late-arriving private shares are still stored even after we've acked done + (State::DkgPrivateSharesDoneAck, Message::DkgPrivateShares(msg)) => { + self.dkg_private_shares(msg, rng) + } + (State::DkgPrivateGather, Message::DkgPrivateSharesDone(msg)) => { + self.dkg_private_shares_done(msg) + } + // DKG end phase + (State::DkgPrivateSharesDoneAck, Message::DkgEndBegin(msg)) => self.dkg_end_begin(msg), + // Signing phase: NonceRequest accepted from Idle or SignGather (coordinator retry) + (State::Idle | State::SignGather, Message::NonceRequest(msg)) => { + self.nonce_request(msg, rng) + } + (State::SignGather, Message::SignatureShareRequest(msg)) => { + self.sign_share_request(msg, rng) + } + // Messages signers never process + ( + _, + Message::DkgEnd(_) + | Message::DkgPublicSharesDoneAck(_) + | Message::DkgPrivateSharesDoneAck(_) + | Message::NonceResponse(_) + | Message::SignatureShareResponse(_), + ) => Ok(vec![]), + // Unexpected state+message combination + (state, msg) => { + warn!( + signer_id = %self.signer_id, + ?state, + msg_type = ?std::mem::discriminant(msg), + "unexpected message in state, dropping" + ); + Ok(vec![]) } - Message::NonceRequest(nonce_request) => self.nonce_request(nonce_request, rng), - Message::DkgEnd(_) - | Message::DkgPublicSharesDoneAck(_) - | Message::DkgPrivateSharesDoneAck(_) - | Message::NonceResponse(_) - | Message::SignatureShareResponse(_) => Ok(vec![]), // TODO }; match out_msgs { @@ -699,7 +736,7 @@ impl Signer { self.dkg_end_begin_msg.is_some(), ); - if self.state == State::DkgPrivateGather { + if self.state == State::DkgPrivateGather || self.state == State::DkgPrivateSharesDoneAck { if let Some(dkg_private_begin) = &self.dkg_private_begin_msg { // need public shares from active signers for signer_id in &dkg_private_begin.signer_ids { @@ -759,6 +796,7 @@ impl Signer { "sending NonceResponse" ); msgs.push(response); + self.move_to(State::SignGather)?; Ok(msgs) } @@ -843,9 +881,11 @@ impl Signer { "sending SignatureShareResponse" ); + self.move_to(State::Idle)?; Ok(vec![Message::SignatureShareResponse(response)]) } else { debug!(signer_id = %self.signer_id, "signer not included in SignatureShareRequest"); + self.move_to(State::Idle)?; Ok(Vec::new()) } } @@ -884,6 +924,7 @@ impl Signer { dkg_id: self.dkg_id, signer_id: self.signer_id, }; + self.move_to(State::DkgPublicSharesDoneAck)?; Ok(vec![Message::DkgPublicSharesDoneAck(ack)]) } @@ -911,6 +952,7 @@ impl Signer { dkg_id: self.dkg_id, signer_id: self.signer_id, }; + self.move_to(State::DkgPrivateSharesDoneAck)?; Ok(vec![Message::DkgPrivateSharesDoneAck(ack)]) } @@ -1229,15 +1271,23 @@ impl StateMachine for Signer let prev_state = &self.state; let accepted = match state { State::Idle => true, - State::DkgPublicDistribute => { - prev_state == &State::Idle - || prev_state == &State::DkgPublicGather - || prev_state == &State::DkgPrivateDistribute - } + // DkgBegin can restart from any state + State::DkgPublicDistribute => matches!( + prev_state, + State::Idle + | State::DkgPublicDistribute + | State::DkgPublicGather + | State::DkgPublicSharesDoneAck + | State::DkgPrivateDistribute + | State::DkgPrivateGather + | State::DkgPrivateSharesDoneAck + ), State::DkgPublicGather => prev_state == &State::DkgPublicDistribute, - State::DkgPrivateDistribute => prev_state == &State::DkgPublicGather, + State::DkgPublicSharesDoneAck => prev_state == &State::DkgPublicGather, + State::DkgPrivateDistribute => prev_state == &State::DkgPublicSharesDoneAck, State::DkgPrivateGather => prev_state == &State::DkgPrivateDistribute, - State::SignGather => prev_state == &State::Idle, + State::DkgPrivateSharesDoneAck => prev_state == &State::DkgPrivateGather, + State::SignGather => prev_state == &State::Idle || prev_state == &State::SignGather, }; if accepted { debug!("state change from {prev_state:?} to {state:?}"); @@ -1258,7 +1308,10 @@ pub mod test { use crate::{ common::PolyCommitment, curve::{ecdsa, scalar::Scalar}, - net::{DkgBegin, DkgEndBegin, DkgPrivateBegin, DkgPublicShares, DkgStatus, Message}, + net::{ + DkgBegin, DkgEndBegin, DkgPrivateBegin, DkgPrivateSharesDone, DkgPublicShares, + DkgPublicSharesDone, DkgStatus, Message, + }, schnorr::ID, state_machine::{ signer::{ConfigError, Error, Signer, State as SignerState}, @@ -1701,6 +1754,18 @@ pub mod test { let _ = signer .process(&dkg_public_shares_packet, &mut rng) .expect("failed to process DkgPublicShares"); + // coordinator signals all public shares received; signer moves to DkgPublicSharesDoneAck + let dkg_public_shares_done = Message::DkgPublicSharesDone(DkgPublicSharesDone { + dkg_id: 1, + signer_ids: vec![0], + }); + let dkg_public_shares_done_packet = Packet { + msg: dkg_public_shares_done, + sig: vec![], + }; + let _ = signer + .process(&dkg_public_shares_done_packet, &mut rng) + .expect("failed to process DkgPublicSharesDone"); let dkg_private_begin = Message::DkgPrivateBegin(DkgPrivateBegin { dkg_id: 1, signer_ids: vec![0], @@ -1712,7 +1777,7 @@ pub mod test { }; let dkg_private_shares = signer .process(&dkg_private_begin_packet, &mut rng) - .expect("failed to process DkgBegin"); + .expect("failed to process DkgPrivateBegin"); let dkg_private_shares_packet = Packet { msg: dkg_private_shares[0].clone(), sig: vec![], @@ -1720,6 +1785,18 @@ pub mod test { let _ = signer .process(&dkg_private_shares_packet, &mut rng) .expect("failed to process DkgPrivateShares"); + // coordinator signals all private shares received; signer moves to DkgPrivateSharesDoneAck + let dkg_private_shares_done = Message::DkgPrivateSharesDone(DkgPrivateSharesDone { + dkg_id: 1, + signer_ids: vec![0], + }); + let dkg_private_shares_done_packet = Packet { + msg: dkg_private_shares_done, + sig: vec![], + }; + let _ = signer + .process(&dkg_private_shares_done_packet, &mut rng) + .expect("failed to process DkgPrivateSharesDone"); let dkg_end_begin = DkgEndBegin { dkg_id: 1, signer_ids: vec![0], From 4600220976675b305e06f3c55ce4ddeecad59379 Mon Sep 17 00:00:00 2001 From: Joey Yandle Date: Thu, 2 Apr 2026 18:18:08 +0200 Subject: [PATCH 4/8] remove late dkg message handling; dont ack Done messages until we have everything we need --- src/state_machine/signer/mod.rs | 130 ++++++++++++++++++++++++-------- 1 file changed, 97 insertions(+), 33 deletions(-) diff --git a/src/state_machine/signer/mod.rs b/src/state_machine/signer/mod.rs index 3a44eba..eb954ad 100644 --- a/src/state_machine/signer/mod.rs +++ b/src/state_machine/signer/mod.rs @@ -163,6 +163,10 @@ pub struct SavedState { pub dkg_private_begin_msg: Option, /// the DKG end begin message received in this round pub dkg_end_begin_msg: Option, + /// pending DkgPublicSharesDone waiting for all public shares to arrive + pending_public_shares_done: Option, + /// pending DkgPrivateSharesDone waiting for all private shares to arrive + pending_private_shares_done: Option, /// whether to verify the signature on Packets pub verify_packet_sigs: bool, /// coordinator public key @@ -246,6 +250,10 @@ pub struct Signer { pub dkg_private_begin_msg: Option, /// the DKG end begin message received in this round pub dkg_end_begin_msg: Option, + /// pending DkgPublicSharesDone waiting for all public shares to arrive + pending_public_shares_done: Option, + /// pending DkgPrivateSharesDone waiting for all private shares to arrive + pending_private_shares_done: Option, /// whether to verify the signature on Packets pub verify_packet_sigs: bool, /// coordinator public key @@ -353,6 +361,8 @@ impl Signer { dkg_private_shares: Default::default(), dkg_private_begin_msg: Default::default(), dkg_end_begin_msg: Default::default(), + pending_public_shares_done: None, + pending_private_shares_done: None, verify_packet_sigs: true, coordinator_public_key: None, kex_private_key: Scalar::random(rng), @@ -384,6 +394,8 @@ impl Signer { dkg_private_shares: state.dkg_private_shares.clone(), dkg_private_begin_msg: state.dkg_private_begin_msg.clone(), dkg_end_begin_msg: state.dkg_end_begin_msg.clone(), + pending_public_shares_done: state.pending_public_shares_done.clone(), + pending_private_shares_done: state.pending_private_shares_done.clone(), verify_packet_sigs: state.verify_packet_sigs, coordinator_public_key: state.coordinator_public_key, kex_private_key: state.kex_private_key, @@ -415,6 +427,8 @@ impl Signer { dkg_private_shares: self.dkg_private_shares.clone(), dkg_private_begin_msg: self.dkg_private_begin_msg.clone(), dkg_end_begin_msg: self.dkg_end_begin_msg.clone(), + pending_public_shares_done: self.pending_public_shares_done.clone(), + pending_private_shares_done: self.pending_private_shares_done.clone(), verify_packet_sigs: self.verify_packet_sigs, coordinator_public_key: self.coordinator_public_key, kex_private_key: self.kex_private_key, @@ -435,6 +449,8 @@ impl Signer { self.dkg_private_shares.clear(); self.dkg_private_begin_msg = None; self.dkg_end_begin_msg = None; + self.pending_public_shares_done = None; + self.pending_private_shares_done = None; self.kex_private_key = Scalar::random(rng); self.kex_public_keys.clear(); self.state = State::Idle; @@ -482,10 +498,6 @@ impl Signer { (_, Message::DkgBegin(msg)) => self.dkg_begin(msg, rng), // DKG public phase (State::DkgPublicGather, Message::DkgPublicShares(msg)) => self.dkg_public_share(msg), - // Late-arriving public shares are still stored even after we've acked done - (State::DkgPublicSharesDoneAck, Message::DkgPublicShares(msg)) => { - self.dkg_public_share(msg) - } (State::DkgPublicGather, Message::DkgPublicSharesDone(msg)) => { self.dkg_public_shares_done(msg) } @@ -496,15 +508,17 @@ impl Signer { (State::DkgPrivateGather, Message::DkgPrivateShares(msg)) => { self.dkg_private_shares(msg, rng) } - // Late-arriving private shares are still stored even after we've acked done - (State::DkgPrivateSharesDoneAck, Message::DkgPrivateShares(msg)) => { - self.dkg_private_shares(msg, rng) - } (State::DkgPrivateGather, Message::DkgPrivateSharesDone(msg)) => { self.dkg_private_shares_done(msg) } - // DKG end phase - (State::DkgPrivateSharesDoneAck, Message::DkgEndBegin(msg)) => self.dkg_end_begin(msg), + // DKG end phase: by the time we reach DkgPrivateSharesDoneAck we have all + // public and private shares, so DkgEndBegin directly triggers dkg_ended + (State::DkgPrivateSharesDoneAck, Message::DkgEndBegin(msg)) => { + let mut out = self.dkg_end_begin(msg)?; + out.push(self.dkg_ended(rng)?); + self.move_to(State::Idle)?; + Ok(out) + } // Signing phase: NonceRequest accepted from Idle or SignGather (coordinator retry) (State::Idle | State::SignGather, Message::NonceRequest(msg)) => { self.nonce_request(msg, rng) @@ -533,17 +547,7 @@ impl Signer { } }; - match out_msgs { - Ok(mut out) => { - if self.can_dkg_end() { - let dkg_end_msgs = self.dkg_ended(rng)?; - out.push(dkg_end_msgs); - self.move_to(State::Idle)?; - } - Ok(out) - } - Err(e) => Err(e), - } + out_msgs } /// DKG is done so compute secrets @@ -920,12 +924,25 @@ impl Signer { ); return Ok(vec![]); } - let ack = DkgPublicSharesDoneAck { - dkg_id: self.dkg_id, - signer_id: self.signer_id, - }; - self.move_to(State::DkgPublicSharesDoneAck)?; - Ok(vec![Message::DkgPublicSharesDoneAck(ack)]) + let have_all = msg + .signer_ids + .iter() + .all(|id| self.dkg_public_shares.contains_key(id)); + if have_all { + let ack = DkgPublicSharesDoneAck { + dkg_id: self.dkg_id, + signer_id: self.signer_id, + }; + self.move_to(State::DkgPublicSharesDoneAck)?; + Ok(vec![Message::DkgPublicSharesDoneAck(ack)]) + } else { + debug!( + signer_id = %self.signer_id, + "DkgPublicSharesDone received but missing some public shares, waiting" + ); + self.pending_public_shares_done = Some(msg.clone()); + Ok(vec![]) + } } fn dkg_private_shares_done( @@ -948,12 +965,25 @@ impl Signer { ); return Ok(vec![]); } - let ack = DkgPrivateSharesDoneAck { - dkg_id: self.dkg_id, - signer_id: self.signer_id, - }; - self.move_to(State::DkgPrivateSharesDoneAck)?; - Ok(vec![Message::DkgPrivateSharesDoneAck(ack)]) + let have_all = msg + .signer_ids + .iter() + .all(|id| self.dkg_private_shares.contains_key(id)); + if have_all { + let ack = DkgPrivateSharesDoneAck { + dkg_id: self.dkg_id, + signer_id: self.signer_id, + }; + self.move_to(State::DkgPrivateSharesDoneAck)?; + Ok(vec![Message::DkgPrivateSharesDoneAck(ack)]) + } else { + debug!( + signer_id = %self.signer_id, + "DkgPrivateSharesDone received but missing some private shares, waiting" + ); + self.pending_private_shares_done = Some(msg.clone()); + Ok(vec![]) + } } fn dkg_public_begin( @@ -1135,6 +1165,23 @@ impl Signer { self.dkg_public_shares .insert(dkg_public_shares.signer_id, dkg_public_shares.clone()); + + // If DkgPublicSharesDone arrived before this share, check if we now have everything + if let Some(pending) = self.pending_public_shares_done.take() { + if pending + .signer_ids + .iter() + .all(|id| self.dkg_public_shares.contains_key(id)) + { + let ack = DkgPublicSharesDoneAck { + dkg_id: self.dkg_id, + signer_id: self.signer_id, + }; + self.move_to(State::DkgPublicSharesDoneAck)?; + return Ok(vec![Message::DkgPublicSharesDoneAck(ack)]); + } + self.pending_public_shares_done = Some(pending); + } Ok(vec![]) } @@ -1219,6 +1266,23 @@ impl Signer { self.decrypted_shares.len(), self.signer.get_num_parties(), ); + + // If DkgPrivateSharesDone arrived before this share, check if we now have everything + if let Some(pending) = self.pending_private_shares_done.take() { + if pending + .signer_ids + .iter() + .all(|id| self.dkg_private_shares.contains_key(id)) + { + let ack = DkgPrivateSharesDoneAck { + dkg_id: self.dkg_id, + signer_id: self.signer_id, + }; + self.move_to(State::DkgPrivateSharesDoneAck)?; + return Ok(vec![Message::DkgPrivateSharesDoneAck(ack)]); + } + self.pending_private_shares_done = Some(pending); + } Ok(vec![]) } From 85ae44caebf18e5e18b0f2e351e988a4a0790bfa Mon Sep 17 00:00:00 2001 From: Joey Yandle Date: Thu, 2 Apr 2026 18:28:02 +0200 Subject: [PATCH 5/8] discard dkg messages not in the coordinator Done messages --- src/state_machine/signer/mod.rs | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/src/state_machine/signer/mod.rs b/src/state_machine/signer/mod.rs index eb954ad..5139124 100644 --- a/src/state_machine/signer/mod.rs +++ b/src/state_machine/signer/mod.rs @@ -924,6 +924,10 @@ impl Signer { ); return Ok(vec![]); } + // Discard any shares already collected from signers not in the coordinator's accepted list + self.dkg_public_shares + .retain(|id, _| msg.signer_ids.contains(id)); + let have_all = msg .signer_ids .iter() @@ -965,6 +969,10 @@ impl Signer { ); return Ok(vec![]); } + // Discard any shares already collected from signers not in the coordinator's accepted list + self.dkg_private_shares + .retain(|id, _| msg.signer_ids.contains(id)); + let have_all = msg .signer_ids .iter() @@ -1139,6 +1147,14 @@ impl Signer { } } + // If we already know which signers the coordinator accepted, discard others + if let Some(pending) = &self.pending_public_shares_done { + if !pending.signer_ids.contains(&signer_id) { + debug!(%signer_id, "discarding DkgPublicShares from signer not in DkgPublicSharesDone"); + return Ok(vec![]); + } + } + let have_shares = self .dkg_public_shares .contains_key(&dkg_public_shares.signer_id); @@ -1215,6 +1231,14 @@ impl Signer { } } + // If we already know which signers the coordinator accepted, discard others + if let Some(pending) = &self.pending_private_shares_done { + if !pending.signer_ids.contains(&src_signer_id) { + debug!(%src_signer_id, "discarding DkgPrivateShares from signer not in DkgPrivateSharesDone"); + return Ok(vec![]); + } + } + if self.dkg_private_shares.contains_key(&src_signer_id) { info!(signer_id = %dkg_private_shares.signer_id, "received duplicate DkgPrivateShares"); return Ok(vec![]); From d0a276ec3e78359ad7f6fdbbe590ea942abfca8f Mon Sep 17 00:00:00 2001 From: Joey Yandle Date: Fri, 3 Apr 2026 08:06:18 +0200 Subject: [PATCH 6/8] add missing dkg_id checks in signer state machine --- src/state_machine/signer/mod.rs | 36 +++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/src/state_machine/signer/mod.rs b/src/state_machine/signer/mod.rs index 5139124..84efbde 100644 --- a/src/state_machine/signer/mod.rs +++ b/src/state_machine/signer/mod.rs @@ -1040,6 +1040,15 @@ impl Signer { dkg_private_begin: &DkgPrivateBegin, rng: &mut R, ) -> Result, Error> { + if dkg_private_begin.dkg_id != self.dkg_id { + warn!( + signer_id = %self.signer_id, + got = %dkg_private_begin.dkg_id, + expected = %self.dkg_id, + "DkgPrivateBegin dkg_id mismatch" + ); + return Ok(vec![]); + } let mut msgs = vec![]; let mut private_shares = DkgPrivateShares { dkg_id: self.dkg_id, @@ -1103,6 +1112,15 @@ impl Signer { /// handle incoming DkgEndBegin pub fn dkg_end_begin(&mut self, dkg_end_begin: &DkgEndBegin) -> Result, Error> { + if dkg_end_begin.dkg_id != self.dkg_id { + warn!( + signer_id = %self.signer_id, + got = %dkg_end_begin.dkg_id, + expected = %self.dkg_id, + "DkgEndBegin dkg_id mismatch" + ); + return Ok(vec![]); + } let msgs = vec![]; self.dkg_end_begin_msg = Some(dkg_end_begin.clone()); @@ -1121,6 +1139,15 @@ impl Signer { &mut self, dkg_public_shares: &DkgPublicShares, ) -> Result, Error> { + if dkg_public_shares.dkg_id != self.dkg_id { + warn!( + signer_id = %self.signer_id, + got = %dkg_public_shares.dkg_id, + expected = %self.dkg_id, + "DkgPublicShares dkg_id mismatch" + ); + return Ok(vec![]); + } debug!( "received DkgPublicShares from signer {} {}/{}", dkg_public_shares.signer_id, @@ -1207,6 +1234,15 @@ impl Signer { dkg_private_shares: &DkgPrivateShares, rng: &mut R, ) -> Result, Error> { + if dkg_private_shares.dkg_id != self.dkg_id { + warn!( + signer_id = %self.signer_id, + got = %dkg_private_shares.dkg_id, + expected = %self.dkg_id, + "DkgPrivateShares dkg_id mismatch" + ); + return Ok(vec![]); + } // go ahead and decrypt here, since we know the signer_id and hence the pubkey of the sender let src_signer_id = dkg_private_shares.signer_id; From c93cf14cfc894f8182743fdfb201882a6653036d Mon Sep 17 00:00:00 2001 From: Joey Yandle Date: Fri, 3 Apr 2026 08:55:17 +0200 Subject: [PATCH 7/8] fix local test flakiness when running tests in parallel by increasing the timeout so we dont hit it accidentally while passing messages around --- src/state_machine/coordinator/fire.rs | 51 +++++++++++++++------------ 1 file changed, 29 insertions(+), 22 deletions(-) diff --git a/src/state_machine/coordinator/fire.rs b/src/state_machine/coordinator/fire.rs index 3c9fd36..14cae10 100644 --- a/src/state_machine/coordinator/fire.rs +++ b/src/state_machine/coordinator/fire.rs @@ -78,12 +78,13 @@ impl Coordinator { State::DkgPublicGather => { if let Some(start) = self.dkg_public_start { if let Some(timeout) = self.config.dkg_public_timeout { - if now.duration_since(start) > timeout { + let elapsed = now.duration_since(start); + if elapsed > timeout { // check dkg_threshold to determine if we can continue let dkg_size = self.compute_dkg_public_size()?; if self.config.dkg_threshold > dkg_size { - error!("Timeout gathering DkgPublicShares for dkg round {} signing round {} iteration {}, dkg_threshold not met ({dkg_size}/{}), unable to continue", self.current_dkg_id, self.current_sign_id, self.current_sign_iter_id, self.config.dkg_threshold); + error!("Timeout gathering DkgPublicShares for dkg round {} signing round {} iteration {}, dkg_threshold not met ({dkg_size}/{}), unable to continue ({:?} > {:?})", self.current_dkg_id, self.current_sign_id, self.current_sign_iter_id, self.config.dkg_threshold, elapsed, timeout); let wait = self.dkg_wait_signer_ids.iter().copied().collect(); return Ok(( None, @@ -93,7 +94,7 @@ impl Coordinator { )); } else { // we hit the timeout but met the threshold, continue - warn!("Timeout gathering DkgPublicShares for dkg round {} signing round {} iteration {}, dkg_threshold was met ({dkg_size}/{}), ", self.current_dkg_id, self.current_sign_id, self.current_sign_iter_id, self.config.dkg_threshold); + warn!("Timeout gathering DkgPublicShares for dkg round {} signing round {} iteration {}, dkg_threshold was met ({dkg_size}/{}), continue ({:?} > {:?})", self.current_dkg_id, self.current_sign_id, self.current_sign_iter_id, self.config.dkg_threshold, elapsed, timeout); self.public_shares_gathered()?; let packet = self.send_public_shares_done()?; return Ok((Some(packet), None)); @@ -106,8 +107,9 @@ impl Coordinator { State::DkgPublicSharesDoneGather => { if let Some(start) = self.dkg_public_start { if let Some(timeout) = self.config.dkg_public_timeout { - if now.duration_since(start) > timeout { - error!("Timeout gathering DkgPublicSharesDoneAck for dkg round {}, not all signers responded", self.current_dkg_id); + let elapsed = now.duration_since(start); + if elapsed > timeout { + error!("Timeout gathering DkgPublicSharesDoneAck for dkg round {}, not all signers responded ({:?} > {:?})", self.current_dkg_id, elapsed, timeout); let wait = self.dkg_wait_signer_ids.iter().copied().collect(); return Ok(( None, @@ -121,12 +123,13 @@ impl Coordinator { State::DkgPrivateGather => { if let Some(start) = self.dkg_private_start { if let Some(timeout) = self.config.dkg_private_timeout { - if now.duration_since(start) > timeout { + let elapsed = now.duration_since(start); + if elapsed > timeout { // check dkg_threshold to determine if we can continue let dkg_size = self.compute_dkg_private_size()?; if self.config.dkg_threshold > dkg_size { - error!("Timeout gathering DkgPrivateShares for dkg round {} signing round {} iteration {}, dkg_threshold not met ({dkg_size}/{}), unable to continue", self.current_dkg_id, self.current_sign_id, self.current_sign_iter_id, self.config.dkg_threshold); + error!("Timeout gathering DkgPrivateShares for dkg round {} signing round {} iteration {}, dkg_threshold not met ({dkg_size}/{}), unable to continue ({:?} > {:?})", self.current_dkg_id, self.current_sign_id, self.current_sign_iter_id, self.config.dkg_threshold, elapsed, timeout); let wait = self.dkg_wait_signer_ids.iter().copied().collect(); return Ok(( None, @@ -136,7 +139,7 @@ impl Coordinator { )); } else { // we hit the timeout but met the threshold, continue - warn!("Timeout gathering DkgPrivateShares for dkg round {} signing round {} iteration {}, dkg_threshold was met ({dkg_size}/{}), ", self.current_dkg_id, self.current_sign_id, self.current_sign_iter_id, self.config.dkg_threshold); + warn!("Timeout gathering DkgPrivateShares for dkg round {} signing round {} iteration {}, dkg_threshold was met ({dkg_size}/{}), continue ({:?} > {:?})", self.current_dkg_id, self.current_sign_id, self.current_sign_iter_id, self.config.dkg_threshold, elapsed, timeout); self.private_shares_gathered()?; let packet = self.send_private_shares_done()?; return Ok((Some(packet), None)); @@ -149,8 +152,9 @@ impl Coordinator { State::DkgPrivateSharesDoneGather => { if let Some(start) = self.dkg_private_start { if let Some(timeout) = self.config.dkg_private_timeout { - if now.duration_since(start) > timeout { - error!("Timeout gathering DkgPrivateSharesDoneAck for dkg round {}, not all signers responded", self.current_dkg_id); + let elapsed = now.duration_since(start); + if elapsed > timeout { + error!("Timeout gathering DkgPrivateSharesDoneAck for dkg round {}, not all signers responded ({:?} > {:?})", self.current_dkg_id, elapsed, timeout); let wait = self.dkg_wait_signer_ids.iter().copied().collect(); return Ok(( None, @@ -164,8 +168,9 @@ impl Coordinator { State::DkgEndGather => { if let Some(start) = self.dkg_end_start { if let Some(timeout) = self.config.dkg_end_timeout { - if now.duration_since(start) > timeout { - error!("Timeout gathering DkgEnd for dkg round {} signing round {} iteration {}, unable to continue", self.current_dkg_id, self.current_sign_id, self.current_sign_iter_id); + let elapsed = now.duration_since(start); + if elapsed > timeout { + error!("Timeout gathering DkgEnd for dkg round {} signing round {} iteration {}, unable to continue ({:?} > {:?})", self.current_dkg_id, self.current_sign_id, self.current_sign_iter_id, elapsed, timeout); let wait = self.dkg_wait_signer_ids.iter().copied().collect(); return Ok(( None, @@ -180,8 +185,9 @@ impl Coordinator { State::NonceGather(_signature_type) => { if let Some(start) = self.nonce_start { if let Some(timeout) = self.config.nonce_timeout { - if now.duration_since(start) > timeout { - error!("Timeout gathering nonces for signing round {} iteration {}, unable to continue", self.current_sign_id, self.current_sign_iter_id); + let elapsed = now.duration_since(start); + if elapsed > timeout { + error!("Timeout gathering nonces for signing round {} iteration {}, unable to continue ({:?} > {:?})", self.current_sign_id, self.current_sign_iter_id, elapsed, timeout); let recv = self .message_nonces .get(&self.message) @@ -204,8 +210,9 @@ impl Coordinator { State::SigShareGather(signature_type) => { if let Some(start) = self.sign_start { if let Some(timeout) = self.config.sign_timeout { - if now.duration_since(start) > timeout { - warn!("Timeout gathering signature shares for signing round {} iteration {}", self.current_sign_id, self.current_sign_iter_id); + let elapsed = now.duration_since(start); + if elapsed > timeout { + warn!("Timeout gathering signature shares for signing round {} iteration {} ({:?} > {:?})", self.current_sign_id, self.current_sign_iter_id, elapsed, timeout); for signer_id in &self .message_nonces .get(&self.message) @@ -2354,8 +2361,8 @@ pub mod test { num_signers: u32, keys_per_signer: u32, ) -> (Vec>, Vec>) { - let timeout = Duration::from_millis(1024); - let expire = Duration::from_millis(1280); + let timeout = Duration::from_millis(2048); + let expire = Duration::from_millis(2222); let (mut coordinators, signers) = setup_with_timeouts::, SignerType>( num_signers, @@ -2457,8 +2464,8 @@ pub mod test { num_signers: u32, keys_per_signer: u32, ) -> (Vec>, Vec>) { - let timeout = Duration::from_millis(1024); - let expire = Duration::from_millis(1280); + let timeout = Duration::from_millis(2048); + let expire = Duration::from_millis(2222); let (coordinators, signers) = setup_with_timeouts::, SignerType>( num_signers, keys_per_signer, @@ -2661,8 +2668,8 @@ pub mod test { } fn insufficient_signers_dkg() { - let timeout = Duration::from_millis(1024); - let expire = Duration::from_millis(1280); + let timeout = Duration::from_millis(2048); + let expire = Duration::from_millis(2222); let num_signers = 10; let keys_per_signer = 2; let (coordinators, signers) = setup_with_timeouts::, Signer>( From 2fc6d03a22c1b6e7e937ca6e35b0fe26d736d038 Mon Sep 17 00:00:00 2001 From: Joey Yandle Date: Fri, 3 Apr 2026 09:59:41 +0200 Subject: [PATCH 8/8] drop Done packets if they contain invalid signer_ids --- src/state_machine/signer/mod.rs | 59 +++++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/src/state_machine/signer/mod.rs b/src/state_machine/signer/mod.rs index 84efbde..6168004 100644 --- a/src/state_machine/signer/mod.rs +++ b/src/state_machine/signer/mod.rs @@ -917,6 +917,21 @@ impl Signer { ); return Ok(vec![]); } + // Reject any unknown signer IDs + let unknown_ids: Vec = msg + .signer_ids + .iter() + .filter(|id| !self.public_keys.signers.contains_key(*id)) + .copied() + .collect(); + if !unknown_ids.is_empty() { + warn!( + signer_id = %self.signer_id, + ?unknown_ids, + "DkgPublicSharesDone contains unknown signer_ids" + ); + return Ok(vec![]); + } if !msg.signer_ids.contains(&self.signer_id) { warn!( signer_id = %self.signer_id, @@ -962,6 +977,21 @@ impl Signer { ); return Ok(vec![]); } + // Reject any unknown signer IDs + let unknown_ids: Vec = msg + .signer_ids + .iter() + .filter(|id| !self.public_keys.signers.contains_key(*id)) + .copied() + .collect(); + if !unknown_ids.is_empty() { + warn!( + signer_id = %self.signer_id, + ?unknown_ids, + "DkgPrivateSharesDone contains unknown signer_ids" + ); + return Ok(vec![]); + } if !msg.signer_ids.contains(&self.signer_id) { warn!( signer_id = %self.signer_id, @@ -1121,6 +1151,35 @@ impl Signer { ); return Ok(vec![]); } + // Reject any unknown signer IDs + let unknown_ids: Vec = dkg_end_begin + .signer_ids + .iter() + .filter(|id| !self.public_keys.signers.contains_key(*id)) + .copied() + .collect(); + if !unknown_ids.is_empty() { + warn!( + signer_id = %self.signer_id, + ?unknown_ids, + "DkgEndBegin contains unknown signer_ids" + ); + return Ok(vec![]); + } + let num_keys: u32 = dkg_end_begin + .signer_ids + .iter() + .filter_map(|id| self.public_keys.signer_key_ids.get(id)) + .map(|key_ids| key_ids.len() as u32) + .sum(); + if num_keys < self.dkg_threshold { + warn!( + signer_id = %self.signer_id, + num_keys, + dkg_threshold = self.dkg_threshold, + "DkgEndBegin below dkg_threshold" + ); + } let msgs = vec![]; self.dkg_end_begin_msg = Some(dkg_end_begin.clone());