wsts/state_machine/coordinator/
frost.rs

1use hashbrown::{HashMap, HashSet};
2use std::collections::BTreeMap;
3use tracing::{debug, info, warn};
4
5use crate::{
6    common::{PolyCommitment, PublicNonce, Signature, SignatureShare},
7    compute,
8    curve::{ecdsa, point::Point},
9    net::{
10        DkgBegin, DkgEnd, DkgEndBegin, DkgPrivateBegin, DkgPrivateShares, DkgPrivateSharesDone,
11        DkgPublicShares, DkgPublicSharesDone, DkgStatus, Message, NonceRequest, NonceResponse,
12        Packet, Signable, SignatureShareRequest, SignatureType,
13    },
14    state_machine::{
15        coordinator::{
16            Config, Coordinator as CoordinatorTrait, Error, SavedState, SignRoundInfo, State,
17        },
18        DkgError, OperationResult, SignError, StateMachine,
19    },
20    taproot::SchnorrProof,
21    traits::Aggregator as AggregatorTrait,
22};
23
24/// The coordinator for the FROST algorithm
25#[derive(Clone, Debug, PartialEq)]
26pub struct Coordinator<Aggregator: AggregatorTrait> {
27    /// common config fields
28    config: Config,
29    /// current DKG round ID
30    pub current_dkg_id: u64,
31    /// current signing round ID
32    current_sign_id: u64,
33    /// current signing iteration ID
34    current_sign_iter_id: u64,
35    dkg_public_shares: BTreeMap<u32, DkgPublicShares>,
36    dkg_private_shares: BTreeMap<u32, DkgPrivateShares>,
37    dkg_end_messages: BTreeMap<u32, DkgEnd>,
38    party_polynomials: HashMap<u32, PolyCommitment>,
39    public_nonces: BTreeMap<u32, NonceResponse>,
40    signature_shares: BTreeMap<u32, Vec<SignatureShare>>,
41    /// aggregate public key
42    pub aggregate_public_key: Option<Point>,
43    signature: Option<Signature>,
44    schnorr_proof: Option<SchnorrProof>,
45    /// which signers we're currently waiting on
46    pub ids_to_await: HashSet<u32>,
47    /// the bytes that we're signing
48    pub message: Vec<u8>,
49    /// current state of the state machine
50    pub state: State,
51    /// Aggregator object
52    aggregator: Aggregator,
53    /// coordinator public key
54    pub coordinator_public_key: Option<ecdsa::PublicKey>,
55}
56
57impl<Aggregator: AggregatorTrait> Coordinator<Aggregator> {
58    /// Process the message inside the passed packet
59    pub fn process_message(
60        &mut self,
61        packet: &Packet,
62    ) -> Result<(Option<Packet>, Option<OperationResult>), Error> {
63        if self.config.verify_packet_sigs {
64            let Some(coordinator_public_key) = self.coordinator_public_key else {
65                return Err(Error::MissingCoordinatorPublicKey);
66            };
67            if !packet.verify(&self.config.public_keys, &coordinator_public_key) {
68                return Err(Error::InvalidPacketSignature);
69            }
70        }
71        loop {
72            match self.state.clone() {
73                State::Idle => {
74                    // Did we receive a coordinator message?
75                    if let Message::DkgBegin(dkg_begin) = &packet.msg {
76                        if self.current_dkg_id == dkg_begin.dkg_id {
77                            // We have already processed this DKG round
78                            return Ok((None, None));
79                        }
80                        // use dkg_id from DkgBegin
81                        let packet = self.start_dkg_round(Some(dkg_begin.dkg_id))?;
82                        return Ok((Some(packet), None));
83                    } else if let Message::NonceRequest(nonce_request) = &packet.msg {
84                        if self.current_sign_id == nonce_request.sign_id {
85                            // We have already processed this sign round
86                            return Ok((None, None));
87                        }
88                        self.current_sign_iter_id = nonce_request.sign_iter_id;
89                        // use sign_id from NonceRequest
90                        let packet = self.start_signing_round(
91                            nonce_request.message.as_slice(),
92                            nonce_request.signature_type,
93                            Some(nonce_request.sign_id),
94                        )?;
95                        return Ok((Some(packet), None));
96                    }
97                    return Ok((None, None));
98                }
99                State::DkgPublicDistribute => {
100                    let packet = self.start_public_shares()?;
101                    return Ok((Some(packet), None));
102                }
103                State::DkgPublicGather => {
104                    self.gather_public_shares(packet)?;
105                    if self.state == State::DkgPublicGather {
106                        // We need more data
107                        return Ok((None, None));
108                    }
109                }
110                State::DkgPublicSharesDoneDistribute => {
111                    let packet = self.send_public_shares_done()?;
112                    return Ok((Some(packet), None));
113                }
114                State::DkgPublicSharesDoneGather => {
115                    self.gather_public_shares_done_ack(packet)?;
116                    if self.state == State::DkgPublicSharesDoneGather {
117                        // We need more data
118                        return Ok((None, None));
119                    }
120                }
121                State::DkgPrivateDistribute => {
122                    let packet = self.start_private_shares()?;
123                    return Ok((Some(packet), None));
124                }
125                State::DkgPrivateGather => {
126                    self.gather_private_shares(packet)?;
127                    if self.state == State::DkgPrivateGather {
128                        // We need more data
129                        return Ok((None, None));
130                    }
131                }
132                State::DkgPrivateSharesDoneDistribute => {
133                    let packet = self.send_private_shares_done()?;
134                    return Ok((Some(packet), None));
135                }
136                State::DkgPrivateSharesDoneGather => {
137                    self.gather_private_shares_done_ack(packet)?;
138                    if self.state == State::DkgPrivateSharesDoneGather {
139                        // We need more data
140                        return Ok((None, None));
141                    }
142                }
143                State::DkgEndDistribute => {
144                    let packet = self.start_dkg_end()?;
145                    return Ok((Some(packet), None));
146                }
147                State::DkgEndGather => {
148                    if let Err(error) = self.gather_dkg_end(packet) {
149                        if let Error::DkgFailure {
150                            reported_failures,
151                            malicious_signers,
152                        } = error
153                        {
154                            return Ok((
155                                None,
156                                Some(OperationResult::DkgError(DkgError::DkgEndFailure {
157                                    reported_failures,
158                                    malicious_signers,
159                                })),
160                            ));
161                        } else {
162                            return Err(error);
163                        }
164                    }
165                    if self.state == State::DkgEndGather {
166                        // We need more data
167                        return Ok((None, None));
168                    } else if self.state == State::Idle {
169                        // We are done with the DKG round! Return the operation result
170                        return Ok((
171                            None,
172                            Some(OperationResult::Dkg(
173                                self.aggregate_public_key
174                                    .ok_or(Error::MissingAggregatePublicKey)?,
175                            )),
176                        ));
177                    }
178                }
179                State::NonceRequest(signature_type) => {
180                    let packet = self.request_nonces(signature_type)?;
181                    return Ok((Some(packet), None));
182                }
183                State::NonceGather(signature_type) => {
184                    self.gather_nonces(packet, signature_type)?;
185                    if self.state == State::NonceGather(signature_type) {
186                        // We need more data
187                        return Ok((None, None));
188                    }
189                }
190                State::SigShareRequest(signature_type) => {
191                    let packet = self.request_sig_shares(signature_type)?;
192                    return Ok((Some(packet), None));
193                }
194                State::SigShareGather(signature_type) => {
195                    if let Err(e) = self.gather_sig_shares(packet, signature_type) {
196                        return Ok((
197                            None,
198                            Some(OperationResult::SignError(SignError::Coordinator(e))),
199                        ));
200                    }
201                    if self.state == State::SigShareGather(signature_type) {
202                        // We need more data
203                        return Ok((None, None));
204                    } else if self.state == State::Idle {
205                        // We are done with the DKG round! Return the operation result
206                        if let SignatureType::Taproot(_) = signature_type {
207                            let schnorr_proof = self
208                                .schnorr_proof
209                                .as_ref()
210                                .ok_or(Error::MissingSchnorrProof)?;
211                            return Ok((
212                                None,
213                                Some(OperationResult::SignTaproot(SchnorrProof {
214                                    r: schnorr_proof.r,
215                                    s: schnorr_proof.s,
216                                })),
217                            ));
218                        } else if let SignatureType::Schnorr = signature_type {
219                            let schnorr_proof = self
220                                .schnorr_proof
221                                .as_ref()
222                                .ok_or(Error::MissingSchnorrProof)?;
223                            return Ok((
224                                None,
225                                Some(OperationResult::SignSchnorr(SchnorrProof {
226                                    r: schnorr_proof.r,
227                                    s: schnorr_proof.s,
228                                })),
229                            ));
230                        } else {
231                            let signature =
232                                self.signature.as_ref().ok_or(Error::MissingSignature)?;
233                            return Ok((
234                                None,
235                                Some(OperationResult::Sign(Signature {
236                                    R: signature.R,
237                                    z: signature.z,
238                                })),
239                            ));
240                        }
241                    }
242                }
243            }
244        }
245    }
246
247    /// Ask signers to send DKG public shares
248    pub fn start_public_shares(&mut self) -> Result<Packet, Error> {
249        self.dkg_public_shares.clear();
250        self.party_polynomials.clear();
251        self.ids_to_await = (0..self.config.num_signers).collect();
252        info!(
253            dkg_id = %self.current_dkg_id,
254            "Starting Public Share Distribution"
255        );
256        let dkg_begin = DkgBegin {
257            dkg_id: self.current_dkg_id,
258        };
259
260        let dkg_begin_packet = Packet {
261            sig: dkg_begin
262                .sign(&self.config.message_private_key)
263                .expect("Failed to sign DkgBegin"),
264            msg: Message::DkgBegin(dkg_begin),
265        };
266        self.move_to(State::DkgPublicGather)?;
267        Ok(dkg_begin_packet)
268    }
269
270    /// Ask signers to send DKG private shares
271    pub fn start_private_shares(&mut self) -> Result<Packet, Error> {
272        self.ids_to_await = (0..self.config.num_signers).collect();
273        info!(
274            dkg_id = %self.current_dkg_id,
275            "Starting Private Share Distribution"
276        );
277        let dkg_begin = DkgPrivateBegin {
278            dkg_id: self.current_dkg_id,
279            key_ids: (1..self.config.num_keys + 1).collect(),
280            signer_ids: (0..self.config.num_signers).collect(),
281        };
282        let dkg_private_begin_msg = Packet {
283            sig: dkg_begin
284                .sign(&self.config.message_private_key)
285                .expect("Failed to sign DkgPrivateBegin"),
286            msg: Message::DkgPrivateBegin(dkg_begin),
287        };
288        self.move_to(State::DkgPrivateGather)?;
289        Ok(dkg_private_begin_msg)
290    }
291
292    /// Ask signers to compute secrets and send DKG end
293    pub fn start_dkg_end(&mut self) -> Result<Packet, Error> {
294        self.ids_to_await = (0..self.config.num_signers).collect();
295        info!(
296            dkg_id = %self.current_dkg_id,
297            "Starting DKG End Distribution"
298        );
299        let dkg_begin = DkgEndBegin {
300            dkg_id: self.current_dkg_id,
301            key_ids: (0..self.config.num_keys).collect(),
302            signer_ids: (0..self.config.num_signers).collect(),
303        };
304        let dkg_end_begin_msg = Packet {
305            sig: dkg_begin.sign(&self.config.message_private_key).expect(""),
306            msg: Message::DkgEndBegin(dkg_begin),
307        };
308        self.move_to(State::DkgEndGather)?;
309        Ok(dkg_end_begin_msg)
310    }
311
312    fn gather_public_shares(&mut self, packet: &Packet) -> Result<(), Error> {
313        if let Message::DkgPublicShares(dkg_public_shares) = &packet.msg {
314            if dkg_public_shares.dkg_id != self.current_dkg_id {
315                return Err(Error::BadDkgId(
316                    dkg_public_shares.dkg_id,
317                    self.current_dkg_id,
318                ));
319            }
320
321            // check that the signer_id exists in the config
322            let signer_public_keys = &self.config.public_keys.signers;
323            if !signer_public_keys.contains_key(&dkg_public_shares.signer_id) {
324                warn!(signer_id = %dkg_public_shares.signer_id, "No public key in config");
325                return Ok(());
326            };
327
328            let have_shares = self
329                .dkg_public_shares
330                .contains_key(&dkg_public_shares.signer_id);
331
332            if have_shares {
333                info!(signer_id = %dkg_public_shares.signer_id, "received duplicate DkgPublicShares");
334                return Ok(());
335            }
336
337            self.ids_to_await.remove(&dkg_public_shares.signer_id);
338
339            self.dkg_public_shares
340                .insert(dkg_public_shares.signer_id, dkg_public_shares.clone());
341            for (party_id, comm) in &dkg_public_shares.comms {
342                self.party_polynomials.insert(*party_id, comm.clone());
343            }
344
345            debug!(
346                dkg_id = %dkg_public_shares.dkg_id,
347                signer_id = %dkg_public_shares.signer_id,
348                "DkgPublicShares received"
349            );
350        }
351
352        if self.ids_to_await.is_empty() {
353            self.move_to(State::DkgPublicSharesDoneDistribute)?;
354        }
355        Ok(())
356    }
357
358    /// Notify signers that all public shares have been received
359    pub fn send_public_shares_done(&mut self) -> Result<Packet, Error> {
360        let signer_ids: Vec<u32> = self.dkg_public_shares.keys().cloned().collect();
361        self.ids_to_await = signer_ids.iter().cloned().collect();
362        info!(dkg_id = %self.current_dkg_id, "Sending DkgPublicSharesDone");
363        let msg = DkgPublicSharesDone {
364            dkg_id: self.current_dkg_id,
365            signer_ids,
366        };
367        let packet = Packet {
368            sig: msg
369                .sign(&self.config.message_private_key)
370                .expect("Failed to sign DkgPublicSharesDone"),
371            msg: Message::DkgPublicSharesDone(msg),
372        };
373        self.move_to(State::DkgPublicSharesDoneGather)?;
374        Ok(packet)
375    }
376
377    fn gather_public_shares_done_ack(&mut self, packet: &Packet) -> Result<(), Error> {
378        if let Message::DkgPublicSharesDoneAck(ack) = &packet.msg {
379            if ack.dkg_id != self.current_dkg_id {
380                return Err(Error::BadDkgId(ack.dkg_id, self.current_dkg_id));
381            }
382            if !self.config.public_keys.signers.contains_key(&ack.signer_id) {
383                warn!(signer_id = %ack.signer_id, "No public key in config");
384                return Ok(());
385            }
386            self.ids_to_await.remove(&ack.signer_id);
387            debug!(
388                dkg_id = %ack.dkg_id,
389                signer_id = %ack.signer_id,
390                "DkgPublicSharesDoneAck received"
391            );
392        }
393        if self.ids_to_await.is_empty() {
394            self.move_to(State::DkgPrivateDistribute)?;
395        }
396        Ok(())
397    }
398
399    fn gather_private_shares(&mut self, packet: &Packet) -> Result<(), Error> {
400        if let Message::DkgPrivateShares(dkg_private_shares) = &packet.msg {
401            if dkg_private_shares.dkg_id != self.current_dkg_id {
402                return Err(Error::BadDkgId(
403                    dkg_private_shares.dkg_id,
404                    self.current_dkg_id,
405                ));
406            }
407
408            // check that the signer_id exists in the config
409            let signer_public_keys = &self.config.public_keys.signers;
410            if !signer_public_keys.contains_key(&dkg_private_shares.signer_id) {
411                warn!(signer_id = %dkg_private_shares.signer_id, "No public key in config");
412                return Ok(());
413            };
414
415            let has_received_shares = self
416                .dkg_private_shares
417                .contains_key(&dkg_private_shares.signer_id);
418            if has_received_shares {
419                info!(signer_id = %dkg_private_shares.signer_id, "received duplicate DkgPrivateShares");
420                return Ok(());
421            }
422
423            self.ids_to_await.remove(&dkg_private_shares.signer_id);
424
425            self.dkg_private_shares
426                .insert(dkg_private_shares.signer_id, dkg_private_shares.clone());
427            info!(
428                dkg_id = %dkg_private_shares.dkg_id,
429                signer_id = %dkg_private_shares.signer_id,
430                "DkgPrivateShares received"
431            );
432        }
433
434        if self.ids_to_await.is_empty() {
435            self.move_to(State::DkgPrivateSharesDoneDistribute)?;
436        }
437        Ok(())
438    }
439
440    /// Notify signers that all private shares have been received
441    pub fn send_private_shares_done(&mut self) -> Result<Packet, Error> {
442        let signer_ids: Vec<u32> = self.dkg_private_shares.keys().cloned().collect();
443        self.ids_to_await = signer_ids.iter().cloned().collect();
444        info!(dkg_id = %self.current_dkg_id, "Sending DkgPrivateSharesDone");
445        let msg = DkgPrivateSharesDone {
446            dkg_id: self.current_dkg_id,
447            signer_ids,
448        };
449        let packet = Packet {
450            sig: msg
451                .sign(&self.config.message_private_key)
452                .expect("Failed to sign DkgPrivateSharesDone"),
453            msg: Message::DkgPrivateSharesDone(msg),
454        };
455        self.move_to(State::DkgPrivateSharesDoneGather)?;
456        Ok(packet)
457    }
458
459    fn gather_private_shares_done_ack(&mut self, packet: &Packet) -> Result<(), Error> {
460        if let Message::DkgPrivateSharesDoneAck(ack) = &packet.msg {
461            if ack.dkg_id != self.current_dkg_id {
462                return Err(Error::BadDkgId(ack.dkg_id, self.current_dkg_id));
463            }
464            if !self.config.public_keys.signers.contains_key(&ack.signer_id) {
465                warn!(signer_id = %ack.signer_id, "No public key in config");
466                return Ok(());
467            }
468            self.ids_to_await.remove(&ack.signer_id);
469            debug!(
470                dkg_id = %ack.dkg_id,
471                signer_id = %ack.signer_id,
472                "DkgPrivateSharesDoneAck received"
473            );
474        }
475        if self.ids_to_await.is_empty() {
476            self.move_to(State::DkgEndDistribute)?;
477        }
478        Ok(())
479    }
480
481    fn gather_dkg_end(&mut self, packet: &Packet) -> Result<(), Error> {
482        debug!(
483            dkg_id = %self.current_dkg_id,
484            waiting = ?self.ids_to_await,
485            "Waiting for Dkg End from signers"
486        );
487        if let Message::DkgEnd(dkg_end) = &packet.msg {
488            if dkg_end.dkg_id != self.current_dkg_id {
489                return Err(Error::BadDkgId(dkg_end.dkg_id, self.current_dkg_id));
490            }
491            if self.ids_to_await.contains(&dkg_end.signer_id) {
492                self.ids_to_await.remove(&dkg_end.signer_id);
493                self.dkg_end_messages
494                    .insert(dkg_end.signer_id, dkg_end.clone());
495                debug!(
496                    dkg_id = %dkg_end.dkg_id,
497                    signer_id = %dkg_end.signer_id,
498                    waiting = ?self.ids_to_await,
499                    "DkgEnd received"
500                );
501            }
502        }
503
504        if self.ids_to_await.is_empty() {
505            let mut reported_failures = HashMap::new();
506
507            for (signer_id, dkg_end) in &self.dkg_end_messages {
508                if let DkgStatus::Failure(dkg_failure) = &dkg_end.status {
509                    warn!(%signer_id, ?dkg_failure, "DkgEnd failure");
510                    reported_failures.insert(*signer_id, dkg_failure.clone());
511                }
512            }
513
514            if reported_failures.is_empty() {
515                self.dkg_end_gathered()?;
516            } else {
517                return Err(Error::DkgFailure {
518                    reported_failures,
519                    malicious_signers: Default::default(),
520                });
521            }
522        }
523        Ok(())
524    }
525
526    fn dkg_end_gathered(&mut self) -> Result<(), Error> {
527        // Cache the polynomials used in DKG for the aggregator
528        for signer_id in self.dkg_private_shares.keys() {
529            let Some(dkg_public_shares) = self.dkg_public_shares.get(signer_id) else {
530                warn!(%signer_id, "no DkgPublicShares");
531                return Err(Error::BadStateChange(format!("Should not have transitioned to DkgEndGather since we were missing DkgPublicShares from signer {signer_id}")));
532            };
533            for (party_id, comm) in &dkg_public_shares.comms {
534                self.party_polynomials.insert(*party_id, comm.clone());
535            }
536        }
537
538        // Calculate the aggregate public key
539        let key = self
540            .party_polynomials
541            .iter()
542            .fold(Point::default(), |s, (_, comm)| s + comm.poly[0]);
543
544        info!(
545            %key,
546            "Aggregate public key"
547        );
548        self.aggregate_public_key = Some(key);
549        self.move_to(State::Idle)
550    }
551
552    fn request_nonces(&mut self, signature_type: SignatureType) -> Result<Packet, Error> {
553        self.public_nonces.clear();
554        info!(
555            sign_id = %self.current_sign_id,
556            sign_iter_id = %self.current_sign_iter_id,
557            "Requesting Nonces"
558        );
559        let nonce_request = NonceRequest {
560            dkg_id: self.current_dkg_id,
561            sign_id: self.current_sign_id,
562            sign_iter_id: self.current_sign_iter_id,
563            message: self.message.clone(),
564            signature_type,
565        };
566        let nonce_request_msg = Packet {
567            sig: nonce_request
568                .sign(&self.config.message_private_key)
569                .expect(""),
570            msg: Message::NonceRequest(nonce_request),
571        };
572        self.ids_to_await = (0..self.config.num_signers).collect();
573        self.move_to(State::NonceGather(signature_type))?;
574        Ok(nonce_request_msg)
575    }
576
577    fn gather_nonces(
578        &mut self,
579        packet: &Packet,
580        signature_type: SignatureType,
581    ) -> Result<(), Error> {
582        if let Message::NonceResponse(nonce_response) = &packet.msg {
583            if nonce_response.dkg_id != self.current_dkg_id {
584                return Err(Error::BadDkgId(nonce_response.dkg_id, self.current_dkg_id));
585            }
586            if nonce_response.sign_id != self.current_sign_id {
587                return Err(Error::BadSignId(
588                    nonce_response.sign_id,
589                    self.current_sign_id,
590                ));
591            }
592            if nonce_response.sign_iter_id != self.current_sign_iter_id {
593                return Err(Error::BadSignIterId(
594                    nonce_response.sign_iter_id,
595                    self.current_sign_iter_id,
596                ));
597            }
598
599            // check that the signer_id exists in the config
600            let signer_public_keys = &self.config.public_keys.signers;
601            if !signer_public_keys.contains_key(&nonce_response.signer_id) {
602                warn!(signer_id = %nonce_response.signer_id, "No public key in config");
603                return Ok(());
604            };
605
606            // check that the signer participated in the current DKG round
607            if !self
608                .dkg_end_messages
609                .contains_key(&nonce_response.signer_id)
610            {
611                warn!(signer_id = %nonce_response.signer_id, "NonceResponse rejected: signer not a DKG participant");
612                return Ok(());
613            }
614
615            // check that the key_ids match the config
616            let Some(signer_key_ids) = self
617                .config
618                .public_keys
619                .signer_key_ids
620                .get(&nonce_response.signer_id)
621            else {
622                warn!(signer_id = %nonce_response.signer_id, "No keys IDs configured");
623                return Ok(());
624            };
625
626            let nonce_response_key_ids = nonce_response
627                .key_ids
628                .iter()
629                .cloned()
630                .collect::<HashSet<u32>>();
631            if *signer_key_ids != nonce_response_key_ids {
632                warn!(signer_id = %nonce_response.signer_id, "Nonce response key_ids didn't match config");
633                return Ok(());
634            }
635
636            for nonce in &nonce_response.nonces {
637                if !nonce.is_valid() {
638                    warn!(
639                        sign_id = %nonce_response.sign_id,
640                        sign_iter_id = %nonce_response.sign_iter_id,
641                        signer_id = %nonce_response.signer_id,
642                        "Received invalid nonce in NonceResponse"
643                    );
644                    return Ok(());
645                }
646            }
647
648            let have_nonces = self.public_nonces.contains_key(&nonce_response.signer_id);
649
650            if have_nonces {
651                info!(signer_id = %nonce_response.signer_id, "Received duplicate NonceResponse");
652                return Ok(());
653            }
654
655            self.public_nonces
656                .insert(nonce_response.signer_id, nonce_response.clone());
657            self.ids_to_await.remove(&nonce_response.signer_id);
658            debug!(
659                sign_id = %nonce_response.sign_id,
660                sign_iter_id = %nonce_response.sign_iter_id,
661                signer_id = %nonce_response.signer_id,
662                waiting = ?self.ids_to_await,
663                "NonceResponse received"
664            );
665        }
666        let recv_key_ids = self
667            .public_nonces
668            .values()
669            .flat_map(|nr| nr.key_ids.iter().copied())
670            .collect::<HashSet<u32>>();
671        if recv_key_ids.len() >= self.config.sign_threshold as usize {
672            let aggregate_nonce = self.compute_aggregate_nonce()?;
673            info!(
674                %aggregate_nonce,
675                "Aggregate nonce"
676            );
677
678            // Lock in the set of signers that will participate in the signing round.
679            self.ids_to_await = self.public_nonces.keys().copied().collect();
680            self.move_to(State::SigShareRequest(signature_type))?;
681        }
682        Ok(())
683    }
684
685    fn request_sig_shares(&mut self, signature_type: SignatureType) -> Result<Packet, Error> {
686        self.signature_shares.clear();
687        info!(
688            sign_id = %self.current_sign_id,
689            "Requesting Signature Shares"
690        );
691        let nonce_responses = self
692            .public_nonces
693            .values()
694            .cloned()
695            .collect::<Vec<NonceResponse>>();
696        let sig_share_request = SignatureShareRequest {
697            dkg_id: self.current_dkg_id,
698            sign_id: self.current_sign_id,
699            sign_iter_id: self.current_sign_iter_id,
700            nonce_responses,
701            message: self.message.clone(),
702            signature_type,
703        };
704        let sig_share_request_msg = Packet {
705            sig: sig_share_request
706                .sign(&self.config.message_private_key)
707                .expect(""),
708            msg: Message::SignatureShareRequest(sig_share_request),
709        };
710        self.ids_to_await = self.public_nonces.keys().copied().collect();
711        self.move_to(State::SigShareGather(signature_type))?;
712
713        Ok(sig_share_request_msg)
714    }
715
716    fn gather_sig_shares(
717        &mut self,
718        packet: &Packet,
719        signature_type: SignatureType,
720    ) -> Result<(), Error> {
721        if let Message::SignatureShareResponse(sig_share_response) = &packet.msg {
722            if sig_share_response.dkg_id != self.current_dkg_id {
723                return Err(Error::BadDkgId(
724                    sig_share_response.dkg_id,
725                    self.current_dkg_id,
726                ));
727            }
728            if sig_share_response.sign_id != self.current_sign_id {
729                return Err(Error::BadSignId(
730                    sig_share_response.sign_id,
731                    self.current_sign_id,
732                ));
733            }
734
735            // check that the signer_id exists in the config
736            let signer_public_keys = &self.config.public_keys.signers;
737            if !signer_public_keys.contains_key(&sig_share_response.signer_id) {
738                warn!(signer_id = %sig_share_response.signer_id, "No public key in config");
739                return Ok(());
740            };
741
742            // check that the signer participated in the current DKG round
743            if !self
744                .dkg_end_messages
745                .contains_key(&sig_share_response.signer_id)
746            {
747                warn!(signer_id = %sig_share_response.signer_id, "SignatureShareResponse rejected: signer not a DKG participant");
748                return Ok(());
749            }
750
751            // check that the key_ids match the config
752            let Some(signer_key_ids) = self
753                .config
754                .public_keys
755                .signer_key_ids
756                .get(&sig_share_response.signer_id)
757            else {
758                warn!(signer_id = %sig_share_response.signer_id, "No keys IDs configured");
759                return Ok(());
760            };
761
762            let mut sig_share_response_key_ids = HashSet::new();
763            for sig_share in &sig_share_response.signature_shares {
764                for key_id in &sig_share.key_ids {
765                    sig_share_response_key_ids.insert(*key_id);
766                }
767            }
768
769            if *signer_key_ids != sig_share_response_key_ids {
770                warn!(signer_id = %sig_share_response.signer_id, "SignatureShareResponse key_ids didn't match config");
771                return Ok(());
772            }
773
774            let have_shares = self
775                .signature_shares
776                .contains_key(&sig_share_response.signer_id);
777
778            if have_shares {
779                info!(signer_id = %sig_share_response.signer_id, "received duplicate SignatureShareResponse");
780                return Ok(());
781            }
782
783            self.signature_shares.insert(
784                sig_share_response.signer_id,
785                sig_share_response.signature_shares.clone(),
786            );
787            self.ids_to_await.remove(&sig_share_response.signer_id);
788            debug!(
789                sign_id = %sig_share_response.sign_id,
790                signer_id = %sig_share_response.signer_id,
791                waiting = ?self.ids_to_await,
792                "SignatureShareResponse received"
793            );
794        }
795        if self.ids_to_await.is_empty() {
796            // Calculate the aggregate signature using the signers that participated.
797            let nonce_responses = self
798                .public_nonces
799                .values()
800                .cloned()
801                .collect::<Vec<NonceResponse>>();
802
803            let nonces = nonce_responses
804                .iter()
805                .flat_map(|nr| nr.nonces.clone())
806                .collect::<Vec<PublicNonce>>();
807
808            let key_ids = nonce_responses
809                .iter()
810                .flat_map(|nr| nr.key_ids.clone())
811                .collect::<Vec<u32>>();
812
813            let shares = &self
814                .public_nonces
815                .keys()
816                .flat_map(|i| self.signature_shares[i].clone())
817                .collect::<Vec<SignatureShare>>();
818
819            debug!(
820                nonces_len = %nonces.len(),
821                shares_len = %shares.len(),
822                "aggregator.sign"
823            );
824
825            self.aggregator.init(&self.party_polynomials)?;
826
827            if let SignatureType::Taproot(merkle_root) = signature_type {
828                let schnorr_proof = self.aggregator.sign_taproot(
829                    &self.message,
830                    &nonces,
831                    shares,
832                    &key_ids,
833                    merkle_root,
834                )?;
835                debug!(
836                    r = %schnorr_proof.r,
837                    s = %schnorr_proof.s,
838                    "SchnorrProof"
839                );
840                self.schnorr_proof = Some(schnorr_proof);
841            } else if let SignatureType::Schnorr = signature_type {
842                let schnorr_proof =
843                    self.aggregator
844                        .sign_schnorr(&self.message, &nonces, shares, &key_ids)?;
845                debug!(
846                    r = %schnorr_proof.r,
847                    s = %schnorr_proof.s,
848                    "SchnorrProof"
849                );
850                self.schnorr_proof = Some(schnorr_proof);
851            } else {
852                let signature = self
853                    .aggregator
854                    .sign(&self.message, &nonces, shares, &key_ids)?;
855                debug!(
856                    R = %signature.R,
857                    z = %signature.z,
858                    "Signature"
859                );
860                self.signature = Some(signature);
861            }
862
863            self.move_to(State::Idle)?;
864        }
865        Ok(())
866    }
867
868    #[allow(non_snake_case)]
869    fn compute_aggregate_nonce(&self) -> Result<Point, Error> {
870        // XXX this needs to be key_ids for v1 and signer_ids for v2
871        let party_ids = self
872            .public_nonces
873            .values()
874            .flat_map(|pn| pn.key_ids.clone())
875            .collect::<Vec<u32>>();
876        let nonces = self
877            .public_nonces
878            .values()
879            .flat_map(|pn| pn.nonces.clone())
880            .collect::<Vec<PublicNonce>>();
881        let Some(group_key) = self.aggregate_public_key else {
882            return Err(Error::MissingAggregatePublicKey);
883        };
884        let (_, aggregate_nonce) =
885            compute::intermediate(&self.message, group_key, &party_ids, &nonces);
886
887        Ok(aggregate_nonce)
888    }
889}
890
891impl<Aggregator: AggregatorTrait> StateMachine<State, Error> for Coordinator<Aggregator> {
892    fn move_to(&mut self, state: State) -> Result<(), Error> {
893        self.can_move_to(&state)?;
894        self.state = state;
895        Ok(())
896    }
897
898    fn can_move_to(&self, state: &State) -> Result<(), Error> {
899        let prev_state = &self.state;
900        let accepted = match state {
901            State::Idle => true,
902            State::DkgPublicDistribute => prev_state == &State::Idle,
903            State::DkgPublicGather => {
904                prev_state == &State::DkgPublicDistribute || prev_state == &State::DkgPublicGather
905            }
906            State::DkgPublicSharesDoneDistribute => prev_state == &State::DkgPublicGather,
907            State::DkgPublicSharesDoneGather => {
908                prev_state == &State::DkgPublicSharesDoneDistribute
909                    || prev_state == &State::DkgPublicSharesDoneGather
910            }
911            State::DkgPrivateDistribute => prev_state == &State::DkgPublicSharesDoneGather,
912            State::DkgPrivateGather => {
913                prev_state == &State::DkgPrivateDistribute || prev_state == &State::DkgPrivateGather
914            }
915            State::DkgPrivateSharesDoneDistribute => prev_state == &State::DkgPrivateGather,
916            State::DkgPrivateSharesDoneGather => {
917                prev_state == &State::DkgPrivateSharesDoneDistribute
918                    || prev_state == &State::DkgPrivateSharesDoneGather
919            }
920            State::DkgEndDistribute => prev_state == &State::DkgPrivateSharesDoneGather,
921            State::DkgEndGather => prev_state == &State::DkgEndDistribute,
922            State::NonceRequest(_) => {
923                prev_state == &State::Idle || prev_state == &State::DkgEndGather
924            }
925            State::NonceGather(signature_type) => {
926                prev_state == &State::NonceRequest(*signature_type)
927                    || prev_state == &State::NonceGather(*signature_type)
928            }
929            State::SigShareRequest(signature_type) => {
930                prev_state == &State::NonceGather(*signature_type)
931            }
932            State::SigShareGather(signature_type) => {
933                prev_state == &State::SigShareRequest(*signature_type)
934                    || prev_state == &State::SigShareGather(*signature_type)
935            }
936        };
937        if accepted {
938            debug!("state change from {prev_state:?} to {state:?}");
939            Ok(())
940        } else {
941            Err(Error::BadStateChange(format!(
942                "{prev_state:?} to {state:?}"
943            )))
944        }
945    }
946}
947
948impl<Aggregator: AggregatorTrait> CoordinatorTrait for Coordinator<Aggregator> {
949    /// Create a new coordinator
950    fn new(config: Config) -> Self {
951        Self {
952            aggregator: Aggregator::new(config.num_keys, config.threshold),
953            config,
954            current_dkg_id: 0,
955            current_sign_id: 0,
956            current_sign_iter_id: 0,
957            dkg_public_shares: Default::default(),
958            dkg_private_shares: Default::default(),
959            dkg_end_messages: Default::default(),
960            party_polynomials: Default::default(),
961            public_nonces: Default::default(),
962            signature_shares: Default::default(),
963            aggregate_public_key: None,
964            signature: None,
965            schnorr_proof: None,
966            message: Default::default(),
967            ids_to_await: Default::default(),
968            state: State::Idle,
969            coordinator_public_key: None,
970        }
971    }
972
973    fn load(state: &SavedState) -> Self {
974        Self {
975            aggregator: Aggregator::new(state.config.num_keys, state.config.threshold),
976            config: state.config.clone(),
977            current_dkg_id: state.current_dkg_id,
978            current_sign_id: state.current_sign_id,
979            current_sign_iter_id: state.current_sign_iter_id,
980            dkg_public_shares: state.dkg_public_shares.clone(),
981            dkg_private_shares: state.dkg_private_shares.clone(),
982            dkg_end_messages: state.dkg_end_messages.clone(),
983            party_polynomials: state.party_polynomials.clone(),
984            public_nonces: state.message_nonces[&Vec::new()].public_nonces.clone(),
985            signature_shares: state.signature_shares.clone(),
986            aggregate_public_key: state.aggregate_public_key,
987            signature: state.signature.clone(),
988            schnorr_proof: state.schnorr_proof.clone(),
989            message: state.message.clone(),
990            ids_to_await: state.dkg_wait_signer_ids.clone(),
991            state: state.state.clone(),
992            coordinator_public_key: state.coordinator_public_key,
993        }
994    }
995
996    fn save(&self) -> SavedState {
997        let round_info = SignRoundInfo {
998            public_nonces: self.public_nonces.clone(),
999            nonce_recv_key_ids: Default::default(),
1000            sign_recv_key_ids: Default::default(),
1001            sign_wait_signer_ids: Default::default(),
1002        };
1003        let mut message_nonces = BTreeMap::new();
1004
1005        message_nonces.insert(Vec::new(), round_info);
1006
1007        SavedState {
1008            config: self.config.clone(),
1009            current_dkg_id: self.current_dkg_id,
1010            current_sign_id: self.current_sign_id,
1011            current_sign_iter_id: self.current_sign_iter_id,
1012            dkg_public_shares: self.dkg_public_shares.clone(),
1013            dkg_private_shares: self.dkg_private_shares.clone(),
1014            dkg_end_messages: self.dkg_end_messages.clone(),
1015            party_polynomials: self.party_polynomials.clone(),
1016            message_nonces,
1017            signature_shares: self.signature_shares.clone(),
1018            aggregate_public_key: self.aggregate_public_key,
1019            signature: self.signature.clone(),
1020            schnorr_proof: self.schnorr_proof.clone(),
1021            message: self.message.clone(),
1022            dkg_wait_signer_ids: self.ids_to_await.clone(),
1023            state: self.state.clone(),
1024            dkg_public_start: Default::default(),
1025            dkg_private_start: Default::default(),
1026            dkg_end_start: Default::default(),
1027            nonce_start: Default::default(),
1028            sign_start: Default::default(),
1029            malicious_signer_ids: Default::default(),
1030            malicious_dkg_signer_ids: Default::default(),
1031            coordinator_public_key: self.coordinator_public_key,
1032        }
1033    }
1034
1035    /// Retrieve the config
1036    fn get_config(&self) -> Config {
1037        self.config.clone()
1038    }
1039
1040    #[cfg(any(test, feature = "testing"))]
1041    fn get_config_mut(&mut self) -> &mut Config {
1042        &mut self.config
1043    }
1044
1045    fn set_coordinator_public_key(&mut self, key: Option<ecdsa::PublicKey>) {
1046        self.coordinator_public_key = key;
1047    }
1048
1049    /// Set the aggregate key and polynomial commitments used to form that key.
1050    ///  Check if the polynomial commitments match the key
1051    fn set_key_and_party_polynomials(
1052        &mut self,
1053        aggregate_key: Point,
1054        party_polynomials: Vec<(u32, PolyCommitment)>,
1055    ) -> Result<(), Error> {
1056        let computed_key = party_polynomials
1057            .iter()
1058            .fold(Point::default(), |s, (_, comm)| s + comm.poly[0]);
1059        if computed_key != aggregate_key {
1060            return Err(Error::AggregateKeyPolynomialMismatch(
1061                computed_key,
1062                aggregate_key,
1063            ));
1064        }
1065        let party_polynomials_len = party_polynomials.len();
1066        let party_polynomials = HashMap::from_iter(party_polynomials);
1067        if party_polynomials.len() != party_polynomials_len {
1068            return Err(Error::DuplicatePartyId);
1069        }
1070        self.aggregate_public_key = Some(aggregate_key);
1071        self.party_polynomials = party_polynomials;
1072        Ok(())
1073    }
1074
1075    /// Process the message inside the passed packet
1076    fn process(
1077        &mut self,
1078        packet: &Packet,
1079    ) -> Result<(Option<Packet>, Option<OperationResult>), Error> {
1080        self.process_message(packet)
1081    }
1082
1083    /// Retrieve the aggregate public key
1084    fn get_aggregate_public_key(&self) -> Option<Point> {
1085        self.aggregate_public_key
1086    }
1087
1088    /// Set the aggregate public key
1089    fn set_aggregate_public_key(&mut self, aggregate_public_key: Option<Point>) {
1090        self.aggregate_public_key = aggregate_public_key;
1091    }
1092
1093    /// Retrieve the current message bytes being signed
1094    fn get_message(&self) -> Vec<u8> {
1095        self.message.clone()
1096    }
1097
1098    /// Retrive the current state
1099    fn get_state(&self) -> State {
1100        self.state.clone()
1101    }
1102
1103    /// Start a DKG round
1104    fn start_dkg_round(&mut self, dkg_id: Option<u64>) -> Result<Packet, Error> {
1105        if let Some(id) = dkg_id {
1106            self.current_dkg_id = id;
1107        } else {
1108            self.current_dkg_id = self.current_dkg_id.wrapping_add(1);
1109        }
1110        info!("Starting DKG round {}", self.current_dkg_id);
1111        self.move_to(State::DkgPublicDistribute)?;
1112        self.start_public_shares()
1113    }
1114
1115    /// Start a signing round
1116    fn start_signing_round(
1117        &mut self,
1118        message: &[u8],
1119        signature_type: SignatureType,
1120        sign_id: Option<u64>,
1121    ) -> Result<Packet, Error> {
1122        // We cannot sign if we haven't first set DKG (either manually or via DKG round).
1123        if self.aggregate_public_key.is_none() {
1124            return Err(Error::MissingAggregatePublicKey);
1125        }
1126        self.message = message.to_vec();
1127        if let Some(id) = sign_id {
1128            self.current_sign_id = id;
1129        } else {
1130            self.current_sign_id = self.current_sign_id.wrapping_add(1);
1131        }
1132        info!("Starting signing round {}", self.current_sign_id);
1133        self.move_to(State::NonceRequest(signature_type))?;
1134        self.request_nonces(signature_type)
1135    }
1136
1137    // Reset internal state
1138    fn reset(&mut self) {
1139        self.state = State::Idle;
1140        self.dkg_public_shares.clear();
1141        self.party_polynomials.clear();
1142        self.public_nonces.clear();
1143        self.signature_shares.clear();
1144        self.ids_to_await = (0..self.config.num_signers).collect();
1145    }
1146}
1147
1148#[cfg(test)]
1149/// Test module for coordinator functionality
1150pub mod test {
1151    use super::*;
1152    #[cfg(feature = "with_v1")]
1153    use crate::v1;
1154    use crate::{
1155        curve::scalar::Scalar,
1156        net::{DkgBegin, Message, NonceRequest, Packet, SignatureShareResponse, SignatureType},
1157        schnorr::{self, ID},
1158        state_machine::{
1159            coordinator::{
1160                frost::Coordinator as FrostCoordinator,
1161                test::{
1162                    bad_signature_share_request, btc_sign_verify, check_signature_shares,
1163                    coordinator_state_machine, empty_private_shares, empty_public_shares,
1164                    equal_after_save_load, feedback_messages, invalid_nonce, new_coordinator,
1165                    run_dkg, run_dkg_sign, setup, start_dkg_round, start_signing_round,
1166                    verify_packet_sigs,
1167                },
1168                Config, Coordinator as CoordinatorTrait, State,
1169            },
1170            OperationResult,
1171        },
1172        traits::{Aggregator as AggregatorTrait, Signer as SignerTrait},
1173        util::create_rng,
1174        v2,
1175    };
1176
1177    #[test]
1178    #[cfg(feature = "with_v1")]
1179    fn new_coordinator_v1() {
1180        new_coordinator::<FrostCoordinator<v1::Aggregator>>();
1181    }
1182
1183    #[test]
1184    fn new_coordinator_v2() {
1185        new_coordinator::<FrostCoordinator<v2::Aggregator>>();
1186    }
1187
1188    #[test]
1189    #[cfg(feature = "with_v1")]
1190    fn equal_after_save_load_v1() {
1191        equal_after_save_load::<FrostCoordinator<v1::Aggregator>, v1::Signer>(2, 2);
1192    }
1193
1194    #[test]
1195    fn equal_after_save_load_v2() {
1196        equal_after_save_load::<FrostCoordinator<v2::Aggregator>, v2::Signer>(2, 2);
1197    }
1198
1199    #[test]
1200    #[cfg(feature = "with_v1")]
1201    fn coordinator_state_machine_v1() {
1202        coordinator_state_machine::<FrostCoordinator<v1::Aggregator>>();
1203    }
1204
1205    #[test]
1206    fn coordinator_state_machine_v2() {
1207        coordinator_state_machine::<FrostCoordinator<v2::Aggregator>>();
1208    }
1209
1210    #[test]
1211    #[cfg(feature = "with_v1")]
1212    fn start_dkg_round_v1() {
1213        start_dkg_round::<FrostCoordinator<v1::Aggregator>>(None);
1214        start_dkg_round::<FrostCoordinator<v1::Aggregator>>(Some(12345u64));
1215    }
1216
1217    #[test]
1218    fn start_dkg_round_v2() {
1219        start_dkg_round::<FrostCoordinator<v2::Aggregator>>(None);
1220        start_dkg_round::<FrostCoordinator<v2::Aggregator>>(Some(12345u64));
1221    }
1222
1223    #[test]
1224    #[cfg(feature = "with_v1")]
1225    fn start_signing_round_v1() {
1226        start_signing_round::<FrostCoordinator<v1::Aggregator>>(None);
1227        start_signing_round::<FrostCoordinator<v1::Aggregator>>(Some(12345u64));
1228    }
1229
1230    #[test]
1231    fn start_signing_round_v2() {
1232        start_signing_round::<FrostCoordinator<v2::Aggregator>>(None);
1233        start_signing_round::<FrostCoordinator<v2::Aggregator>>(Some(12345u64));
1234    }
1235
1236    #[test]
1237    #[cfg(feature = "with_v1")]
1238    fn start_public_shares_v1() {
1239        start_public_shares::<v1::Aggregator>();
1240    }
1241
1242    #[test]
1243    fn start_public_shares_v2() {
1244        start_public_shares::<v2::Aggregator>();
1245    }
1246
1247    fn start_public_shares<Aggregator: AggregatorTrait>() {
1248        let mut rng = create_rng();
1249        let config = Config::new(10, 40, 28, Scalar::random(&mut rng));
1250        let mut coordinator = FrostCoordinator::<Aggregator>::new(config);
1251
1252        coordinator.state = State::DkgPublicDistribute; // Must be in this state before calling start public shares
1253
1254        let result = coordinator.start_public_shares().unwrap();
1255
1256        assert!(matches!(result.msg, Message::DkgBegin(_)));
1257        assert_eq!(coordinator.get_state(), State::DkgPublicGather);
1258        assert_eq!(coordinator.current_dkg_id, 0);
1259    }
1260
1261    #[test]
1262    #[cfg(feature = "with_v1")]
1263    fn start_private_shares_v1() {
1264        start_private_shares::<v1::Aggregator>();
1265    }
1266
1267    #[test]
1268    fn start_private_shares_v2() {
1269        start_private_shares::<v2::Aggregator>();
1270    }
1271
1272    fn start_private_shares<Aggregator: AggregatorTrait>() {
1273        let mut rng = create_rng();
1274        let config = Config::new(10, 40, 28, Scalar::random(&mut rng));
1275        let mut coordinator = FrostCoordinator::<Aggregator>::new(config);
1276
1277        coordinator.state = State::DkgPrivateDistribute; // Must be in this state before calling start private shares
1278
1279        let message = coordinator.start_private_shares().unwrap();
1280        assert!(matches!(message.msg, Message::DkgPrivateBegin(_)));
1281        assert_eq!(coordinator.get_state(), State::DkgPrivateGather);
1282        assert_eq!(coordinator.current_dkg_id, 0);
1283    }
1284
1285    #[test]
1286    #[cfg(feature = "with_v1")]
1287    fn dkg_public_share_v1() {
1288        dkg_public_share::<v1::Aggregator, v1::Signer>();
1289    }
1290
1291    #[test]
1292    #[cfg(feature = "with_v1")]
1293    fn dkg_public_share_v2() {
1294        dkg_public_share::<v1::Aggregator, v2::Signer>();
1295    }
1296
1297    /// test basic insertion and detection of duplicates for DkgPublicShares
1298    #[allow(dead_code)]
1299    fn dkg_public_share<Aggregator: AggregatorTrait, Signer: SignerTrait>() {
1300        let ctx = 0u64.to_be_bytes();
1301        let mut rng = create_rng();
1302        let (coordinators, _) = setup::<FrostCoordinator<Aggregator>, Signer>(2, 1);
1303        let mut coordinator: FrostCoordinator<Aggregator> = coordinators[0].clone();
1304
1305        coordinator.start_dkg_round(None).unwrap(); // = State::DkgPublicGather;
1306        let public_shares = DkgPublicShares {
1307            dkg_id: 1,
1308            signer_id: 0,
1309            comms: vec![(
1310                0,
1311                PolyCommitment {
1312                    id: ID::new(&Scalar::new(), &Scalar::new(), &ctx, &mut rng),
1313                    poly: vec![],
1314                },
1315            )],
1316            kex_public_key: Point::from(Scalar::random(&mut rng)),
1317            kex_proof: schnorr::Proof {
1318                R: Point::new(),
1319                s: Scalar::new(),
1320            },
1321        };
1322        let packet = Packet {
1323            msg: Message::DkgPublicShares(public_shares.clone()),
1324            sig: Default::default(),
1325        };
1326        coordinator.gather_public_shares(&packet).unwrap();
1327        assert_eq!(1, coordinator.dkg_public_shares.len());
1328
1329        // check that a duplicate public share is ignored
1330        let dup_public_shares = DkgPublicShares {
1331            dkg_id: 1,
1332            signer_id: 0,
1333            comms: vec![(
1334                0,
1335                PolyCommitment {
1336                    id: ID::new(&Scalar::new(), &Scalar::new(), &ctx, &mut rng),
1337                    poly: vec![],
1338                },
1339            )],
1340            kex_public_key: Point::from(Scalar::random(&mut rng)),
1341            kex_proof: schnorr::Proof {
1342                R: Point::new(),
1343                s: Scalar::new(),
1344            },
1345        };
1346        let dup_packet = Packet {
1347            msg: Message::DkgPublicShares(dup_public_shares.clone()),
1348            sig: Default::default(),
1349        };
1350
1351        coordinator.gather_public_shares(&dup_packet).unwrap();
1352        assert_eq!(1, coordinator.dkg_public_shares.len());
1353        assert_eq!(
1354            public_shares,
1355            coordinator
1356                .dkg_public_shares
1357                .iter()
1358                .next()
1359                .unwrap()
1360                .1
1361                .clone()
1362        );
1363    }
1364
1365    #[test]
1366    #[cfg(feature = "with_v1")]
1367    fn dkg_private_share_v1() {
1368        dkg_private_share::<v1::Aggregator, v1::Signer>();
1369    }
1370
1371    #[test]
1372    fn dkg_private_share_v2() {
1373        dkg_private_share::<v2::Aggregator, v2::Signer>();
1374    }
1375
1376    /// test basic insertion and detection of duplicates for DkgPrivateShares
1377    fn dkg_private_share<Aggregator: AggregatorTrait, Signer: SignerTrait>() {
1378        let (coordinators, _) = setup::<FrostCoordinator<Aggregator>, Signer>(2, 1);
1379        let mut coordinator: FrostCoordinator<Aggregator> = coordinators[0].clone();
1380
1381        coordinator.state = State::DkgPrivateGather;
1382
1383        let private_share = DkgPrivateShares {
1384            dkg_id: 0,
1385            signer_id: 0,
1386            shares: vec![(1, HashMap::new())],
1387        };
1388        let packet = Packet {
1389            msg: Message::DkgPrivateShares(private_share.clone()),
1390            sig: Default::default(),
1391        };
1392        coordinator.gather_private_shares(&packet).unwrap();
1393        assert_eq!(1, coordinator.dkg_private_shares.len());
1394
1395        // check that a duplicate private share is ignored
1396        let dup_private_share = DkgPrivateShares {
1397            dkg_id: 0,
1398            signer_id: 0,
1399            shares: vec![(1, HashMap::new())],
1400        };
1401        let packet = Packet {
1402            msg: Message::DkgPrivateShares(dup_private_share.clone()),
1403            sig: Default::default(),
1404        };
1405        coordinator.gather_private_shares(&packet).unwrap();
1406        assert_eq!(1, coordinator.dkg_private_shares.len());
1407        assert_eq!(
1408            private_share,
1409            coordinator
1410                .dkg_private_shares
1411                .iter()
1412                .next()
1413                .unwrap()
1414                .1
1415                .clone()
1416        );
1417    }
1418
1419    #[test]
1420    #[cfg(feature = "with_v1")]
1421    fn nonce_response_v1() {
1422        nonce_response::<v1::Aggregator, v1::Signer>();
1423    }
1424
1425    #[test]
1426    fn nonce_response_v2() {
1427        nonce_response::<v2::Aggregator, v2::Signer>();
1428    }
1429
1430    /// test basic insertion and detection of duplicates for NonceResponse
1431    fn nonce_response<Aggregator: AggregatorTrait, Signer: SignerTrait>() {
1432        let mut rng = create_rng();
1433        let (coordinators, _) = setup::<FrostCoordinator<Aggregator>, Signer>(2, 1);
1434        let mut coordinator: FrostCoordinator<Aggregator> = coordinators[0].clone();
1435        let signature_type = SignatureType::Frost;
1436        let message = vec![0u8];
1437        coordinator.state = State::NonceGather(signature_type);
1438        coordinator.aggregate_public_key = Some(Point::from(Scalar::random(&mut rng)));
1439        // Populate dkg_end_messages so the DKG participation check passes
1440        for signer_id in 0..coordinator.config.num_signers {
1441            coordinator.dkg_end_messages.insert(
1442                signer_id,
1443                DkgEnd {
1444                    dkg_id: 0,
1445                    signer_id,
1446                    status: DkgStatus::Success,
1447                },
1448            );
1449        }
1450
1451        let nonce_response = NonceResponse {
1452            dkg_id: 0,
1453            sign_id: 0,
1454            sign_iter_id: 0,
1455            signer_id: 0,
1456            key_ids: vec![1u32],
1457            nonces: vec![PublicNonce {
1458                D: Point::from(Scalar::random(&mut rng)),
1459                E: Point::from(Scalar::random(&mut rng)),
1460            }],
1461            message: message.clone(),
1462        };
1463        let packet = Packet {
1464            msg: Message::NonceResponse(nonce_response.clone()),
1465            sig: Default::default(),
1466        };
1467        coordinator.gather_nonces(&packet, signature_type).unwrap();
1468        assert_eq!(1, coordinator.public_nonces.len());
1469
1470        // check that a duplicate private share is ignored
1471        let dup_nonce_response = NonceResponse {
1472            dkg_id: 0,
1473            sign_id: 0,
1474            sign_iter_id: 0,
1475            signer_id: 0,
1476            key_ids: vec![1u32],
1477            nonces: vec![PublicNonce {
1478                D: Point::from(Scalar::random(&mut rng)),
1479                E: Point::from(Scalar::random(&mut rng)),
1480            }],
1481            message: message.clone(),
1482        };
1483        let packet = Packet {
1484            msg: Message::NonceResponse(dup_nonce_response.clone()),
1485            sig: Default::default(),
1486        };
1487        coordinator.gather_nonces(&packet, signature_type).unwrap();
1488        assert_eq!(1, coordinator.public_nonces.len());
1489        assert_eq!(
1490            nonce_response,
1491            coordinator.public_nonces.iter().next().unwrap().1.clone()
1492        );
1493    }
1494
1495    #[test]
1496    #[cfg(feature = "with_v1")]
1497    fn sig_share_v1() {
1498        sig_share::<v1::Aggregator, v1::Signer>();
1499    }
1500
1501    #[test]
1502    #[cfg(feature = "with_v1")]
1503    fn sig_share_v2() {
1504        sig_share::<v2::Aggregator, v1::Signer>();
1505    }
1506
1507    /// test basic insertion and detection of duplicates for SignatureShareResponse
1508    #[allow(dead_code)]
1509    fn sig_share<Aggregator: AggregatorTrait, Signer: SignerTrait>() {
1510        let mut rng = create_rng();
1511        let (coordinators, _) = setup::<FrostCoordinator<Aggregator>, Signer>(2, 1);
1512        let mut coordinator = coordinators[0].clone();
1513        let signature_type = SignatureType::Frost;
1514
1515        coordinator.ids_to_await = (0..coordinator.config.num_signers).collect();
1516        coordinator.state = State::SigShareGather(signature_type);
1517        // Populate dkg_end_messages so the DKG participation check passes
1518        for signer_id in 0..coordinator.config.num_signers {
1519            coordinator.dkg_end_messages.insert(
1520                signer_id,
1521                DkgEnd {
1522                    dkg_id: 0,
1523                    signer_id,
1524                    status: DkgStatus::Success,
1525                },
1526            );
1527        }
1528
1529        let signature_share = SignatureShare {
1530            id: 1,
1531            z_i: Scalar::random(&mut rng),
1532            key_ids: vec![1],
1533        };
1534        let sig_share_response = SignatureShareResponse {
1535            dkg_id: 0,
1536            sign_id: 0,
1537            sign_iter_id: 0,
1538            signer_id: 0,
1539            signature_shares: vec![signature_share.clone()],
1540        };
1541        let packet = Packet {
1542            msg: Message::SignatureShareResponse(sig_share_response.clone()),
1543            sig: Default::default(),
1544        };
1545        coordinator
1546            .gather_sig_shares(&packet, signature_type)
1547            .unwrap();
1548        assert_eq!(1, coordinator.signature_shares.len());
1549
1550        // check that a duplicate private share is ignored
1551        let dup_signature_share = SignatureShare {
1552            id: 1,
1553            z_i: Scalar::random(&mut rng),
1554            key_ids: vec![1],
1555        };
1556        let dup_sig_share_response = SignatureShareResponse {
1557            dkg_id: 0,
1558            sign_id: 0,
1559            sign_iter_id: 0,
1560            signer_id: 0,
1561            signature_shares: vec![dup_signature_share.clone()],
1562        };
1563        let packet = Packet {
1564            msg: Message::SignatureShareResponse(dup_sig_share_response.clone()),
1565            sig: Default::default(),
1566        };
1567        coordinator
1568            .gather_sig_shares(&packet, signature_type)
1569            .unwrap();
1570        assert_eq!(1, coordinator.signature_shares.len());
1571        assert_eq!(
1572            vec![signature_share],
1573            coordinator
1574                .signature_shares
1575                .iter()
1576                .next()
1577                .unwrap()
1578                .1
1579                .clone()
1580        );
1581    }
1582
1583    #[test]
1584    #[cfg(feature = "with_v1")]
1585    fn run_dkg_sign_v1() {
1586        run_dkg_sign::<FrostCoordinator<v1::Aggregator>, v1::Signer>(5, 2);
1587    }
1588
1589    #[test]
1590    fn run_dkg_sign_v2() {
1591        run_dkg_sign::<FrostCoordinator<v2::Aggregator>, v2::Signer>(5, 2);
1592    }
1593
1594    #[test]
1595    #[cfg(feature = "with_v1")]
1596    fn sign_threshold_sign_v1() {
1597        sign_threshold_sign::<v1::Aggregator, v1::Signer>();
1598    }
1599
1600    #[test]
1601    fn sign_threshold_sign_v2() {
1602        sign_threshold_sign::<v2::Aggregator, v2::Signer>();
1603    }
1604
1605    fn sign_threshold_sign<Aggregator: AggregatorTrait, Signer: SignerTrait>() {
1606        let (mut coordinators, mut signers) = run_dkg::<FrostCoordinator<Aggregator>, Signer>(5, 2);
1607
1608        // Require all key_ids to participate in the signing round.
1609        for coordinator in &mut coordinators {
1610            coordinator.config.sign_threshold = coordinator.config.num_keys;
1611        }
1612
1613        let msg = "It was many and many a year ago, in a kingdom by the sea"
1614            .as_bytes()
1615            .to_vec();
1616        let signature_type = SignatureType::Frost;
1617        let message = coordinators
1618            .first_mut()
1619            .unwrap()
1620            .start_signing_round(&msg, signature_type, None)
1621            .unwrap();
1622        assert_eq!(
1623            coordinators.first().unwrap().state,
1624            State::NonceGather(signature_type)
1625        );
1626
1627        // Drive nonce gathering: every signer responds, so we should reach the threshold.
1628        let (outbound_messages, operation_results) =
1629            feedback_messages(&mut coordinators, &mut signers, &[message]);
1630        assert!(operation_results.is_empty());
1631        for coordinator in &coordinators {
1632            assert_eq!(coordinator.state, State::SigShareGather(signature_type));
1633        }
1634
1635        // All signers contributed nonces.
1636        for coordinator in &coordinators {
1637            assert_eq!(coordinator.public_nonces.len(), signers.len());
1638        }
1639
1640        assert_eq!(outbound_messages.len(), 1);
1641        assert!(
1642            matches!(outbound_messages[0].msg, Message::SignatureShareRequest(_)),
1643            "Expected SignatureShareRequest message"
1644        );
1645
1646        // Drive sig share gathering and aggregation.
1647        let (outbound_messages, operation_results) =
1648            feedback_messages(&mut coordinators, &mut signers, &outbound_messages);
1649
1650        for coordinator in &coordinators {
1651            assert_eq!(coordinator.signature_shares.len(), signers.len());
1652        }
1653
1654        assert!(outbound_messages.is_empty());
1655        assert_eq!(operation_results.len(), 1);
1656        let OperationResult::Sign(sig) = &operation_results[0] else {
1657            panic!("Expected Signature Operation result")
1658        };
1659        assert!(sig.verify(
1660            &coordinators
1661                .first()
1662                .unwrap()
1663                .aggregate_public_key
1664                .expect("No aggregate public key set!"),
1665            &msg
1666        ));
1667        for coordinator in &coordinators {
1668            assert_eq!(coordinator.state, State::Idle);
1669        }
1670    }
1671
1672    #[test]
1673    #[cfg(feature = "with_v1")]
1674    fn check_signature_shares_v1() {
1675        check_signature_shares::<FrostCoordinator<v1::Aggregator>, v1::Signer>(
1676            5,
1677            1,
1678            SignatureType::Frost,
1679            vec![1],
1680        );
1681        check_signature_shares::<FrostCoordinator<v1::Aggregator>, v1::Signer>(
1682            5,
1683            1,
1684            SignatureType::Schnorr,
1685            vec![1],
1686        );
1687        check_signature_shares::<FrostCoordinator<v1::Aggregator>, v1::Signer>(
1688            5,
1689            1,
1690            SignatureType::Taproot(None),
1691            vec![1],
1692        );
1693        check_signature_shares::<FrostCoordinator<v1::Aggregator>, v1::Signer>(
1694            5,
1695            1,
1696            SignatureType::Taproot(Some([23u8; 32])),
1697            vec![1],
1698        );
1699    }
1700
1701    #[test]
1702    fn check_signature_shares_v2() {
1703        check_signature_shares::<FrostCoordinator<v2::Aggregator>, v2::Signer>(
1704            5,
1705            2,
1706            SignatureType::Frost,
1707            vec![0],
1708        );
1709        check_signature_shares::<FrostCoordinator<v2::Aggregator>, v2::Signer>(
1710            5,
1711            2,
1712            SignatureType::Schnorr,
1713            vec![0],
1714        );
1715        check_signature_shares::<FrostCoordinator<v2::Aggregator>, v2::Signer>(
1716            5,
1717            2,
1718            SignatureType::Taproot(None),
1719            vec![0],
1720        );
1721        check_signature_shares::<FrostCoordinator<v2::Aggregator>, v2::Signer>(
1722            5,
1723            2,
1724            SignatureType::Taproot(Some([23u8; 32])),
1725            vec![0],
1726        );
1727    }
1728
1729    #[test]
1730    #[cfg(feature = "with_v1")]
1731    fn bad_signature_share_request_v1() {
1732        bad_signature_share_request::<FrostCoordinator<v1::Aggregator>, v1::Signer>(5, 2);
1733    }
1734
1735    #[test]
1736    fn bad_signature_share_request_v2() {
1737        bad_signature_share_request::<FrostCoordinator<v2::Aggregator>, v2::Signer>(5, 2);
1738    }
1739
1740    #[test]
1741    #[cfg(feature = "with_v1")]
1742    fn invalid_nonce_v1() {
1743        invalid_nonce::<FrostCoordinator<v1::Aggregator>, v1::Signer>(5, 2);
1744    }
1745
1746    #[test]
1747    fn invalid_nonce_v2() {
1748        invalid_nonce::<FrostCoordinator<v2::Aggregator>, v2::Signer>(5, 2);
1749    }
1750
1751    #[test]
1752    fn process_inbound_messages_v2() {
1753        run_dkg_sign::<FrostCoordinator<v2::Aggregator>, v2::Signer>(5, 2);
1754    }
1755
1756    #[test]
1757    #[cfg(feature = "with_v1")]
1758    fn old_round_ids_are_ignored_v1() {
1759        old_round_ids_are_ignored::<v1::Aggregator>();
1760    }
1761
1762    #[test]
1763    fn old_round_ids_are_ignored_v2() {
1764        old_round_ids_are_ignored::<v2::Aggregator>();
1765    }
1766
1767    fn old_round_ids_are_ignored<Aggregator: AggregatorTrait>() {
1768        let mut rng = create_rng();
1769        let mut config = Config::new(10, 40, 28, Scalar::random(&mut rng));
1770        config.verify_packet_sigs = false;
1771        let mut coordinator = FrostCoordinator::<Aggregator>::new(config);
1772        let id: u64 = 10;
1773        let old_id = id;
1774        coordinator.current_dkg_id = id;
1775        coordinator.current_sign_id = id;
1776        // Attempt to start an old DKG round
1777        let (packet, result) = coordinator
1778            .process(&Packet {
1779                sig: vec![],
1780                msg: Message::DkgBegin(DkgBegin { dkg_id: old_id }),
1781            })
1782            .unwrap();
1783        assert!(packet.is_none());
1784        assert!(result.is_none());
1785        assert_eq!(coordinator.state, State::Idle);
1786        assert_eq!(coordinator.current_dkg_id, id);
1787
1788        // Attempt to start the same DKG round
1789        let (packet, result) = coordinator
1790            .process(&Packet {
1791                sig: vec![],
1792                msg: Message::DkgBegin(DkgBegin { dkg_id: id }),
1793            })
1794            .unwrap();
1795        assert!(packet.is_none());
1796        assert!(result.is_none());
1797        assert_eq!(coordinator.state, State::Idle);
1798        assert_eq!(coordinator.current_dkg_id, id);
1799
1800        // Attempt to start an old Sign round
1801        let (packet, result) = coordinator
1802            .process(&Packet {
1803                sig: vec![],
1804                msg: Message::NonceRequest(NonceRequest {
1805                    dkg_id: id,
1806                    sign_id: old_id,
1807                    message: vec![],
1808                    sign_iter_id: id,
1809                    signature_type: SignatureType::Frost,
1810                }),
1811            })
1812            .unwrap();
1813        assert!(packet.is_none());
1814        assert!(result.is_none());
1815        assert_eq!(coordinator.state, State::Idle);
1816        assert_eq!(coordinator.current_sign_id, id);
1817
1818        // Attempt to start the same Sign round
1819        let (packet, result) = coordinator
1820            .process(&Packet {
1821                sig: vec![],
1822                msg: Message::NonceRequest(NonceRequest {
1823                    dkg_id: id,
1824                    sign_id: id,
1825                    message: vec![],
1826                    sign_iter_id: id,
1827                    signature_type: SignatureType::Frost,
1828                }),
1829            })
1830            .unwrap();
1831        assert!(packet.is_none());
1832        assert!(result.is_none());
1833        assert_eq!(coordinator.state, State::Idle);
1834        assert_eq!(coordinator.current_sign_id, id);
1835    }
1836
1837    #[test]
1838    #[cfg(feature = "with_v1")]
1839    fn empty_public_shares_v1() {
1840        empty_public_shares::<FrostCoordinator<v1::Aggregator>, v1::Signer>(5, 2);
1841    }
1842
1843    #[test]
1844    fn empty_public_shares_v2() {
1845        empty_public_shares::<FrostCoordinator<v2::Aggregator>, v2::Signer>(5, 2);
1846    }
1847
1848    #[test]
1849    #[cfg(feature = "with_v1")]
1850    fn empty_private_shares_v1() {
1851        empty_private_shares::<FrostCoordinator<v1::Aggregator>, v1::Signer>(5, 2);
1852    }
1853
1854    #[test]
1855    fn empty_private_shares_v2() {
1856        empty_private_shares::<FrostCoordinator<v2::Aggregator>, v2::Signer>(5, 2);
1857    }
1858
1859    #[test]
1860    #[cfg(feature = "with_v1")]
1861    fn verify_packet_sigs_v1() {
1862        verify_packet_sigs::<FrostCoordinator<v1::Aggregator>, v1::Signer>();
1863    }
1864
1865    #[test]
1866    fn verify_packet_sigs_v2() {
1867        verify_packet_sigs::<FrostCoordinator<v2::Aggregator>, v2::Signer>();
1868    }
1869
1870    #[test]
1871    #[cfg(feature = "with_v1")]
1872    fn btc_sign_verify_v1() {
1873        btc_sign_verify::<FrostCoordinator<v1::Aggregator>, v1::Signer>(5, 2);
1874    }
1875
1876    #[test]
1877    fn btc_sign_verify_v2() {
1878        btc_sign_verify::<FrostCoordinator<v2::Aggregator>, v2::Signer>(5, 2);
1879    }
1880}
wsts/state_machine/coordinator/frost.rs

wsts/state_machine/coordinator/
frost.rs
