Skip to main content

sui_network/randomness/
mod.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use self::{auth::AllowedPeersUpdatable, metrics::Metrics};
5use anemo::PeerId;
6use anyhow::Result;
7use fastcrypto::groups::bls12381;
8use fastcrypto_tbls::{
9    dkg_v1,
10    nodes::PartyId,
11    tbls::ThresholdBls,
12    types::{ShareIndex, ThresholdBls12381MinSig},
13};
14use mysten_common::ZipDebugEqIteratorExt;
15use mysten_metrics::spawn_monitored_task;
16use mysten_network::anemo_ext::NetworkExt;
17use serde::{Deserialize, Serialize};
18use std::{
19    collections::{HashMap, HashSet, btree_map::BTreeMap},
20    ops::Bound,
21    sync::Arc,
22    time::{self, Duration},
23};
24use sui_config::p2p::RandomnessConfig;
25use sui_macros::fail_point_if;
26use sui_types::{
27    base_types::AuthorityName,
28    committee::EpochId,
29    crypto::{RandomnessPartialSignature, RandomnessRound, RandomnessSignature},
30};
31use tokio::sync::{
32    OnceCell, {mpsc, oneshot},
33};
34use tracing::{debug, error, info, instrument, warn};
35
36mod auth;
37mod builder;
38mod generated {
39    include!(concat!(env!("OUT_DIR"), "/sui.Randomness.rs"));
40}
41mod metrics;
42mod server;
43#[cfg(test)]
44mod tests;
45
46pub use builder::{Builder, UnstartedRandomness};
47pub use generated::{
48    randomness_client::RandomnessClient,
49    randomness_server::{Randomness, RandomnessServer},
50};
51
52#[derive(Clone, Debug, Serialize, Deserialize)]
53pub struct SendSignaturesRequest {
54    epoch: EpochId,
55    round: RandomnessRound,
56    // BCS-serialized `RandomnessPartialSignature` values. We store raw bytes here to enable
57    // defenses against too-large messages.
58    // The protocol requires the signatures to be ordered by share index (as provided by fastcrypto).
59    partial_sigs: Vec<Vec<u8>>,
60    // If peer already has a full signature available for the round, it's provided here in lieu
61    // of partial sigs.
62    sig: Option<RandomnessSignature>,
63}
64
65/// A handle to the Randomness network subsystem.
66///
67/// This handle can be cloned and shared. Once all copies of a Randomness system's Handle have been
68/// dropped, the Randomness system will be gracefully shutdown.
69#[derive(Clone, Debug)]
70pub struct Handle {
71    sender: mpsc::Sender<RandomnessMessage>,
72}
73
74impl Handle {
75    /// Transitions the Randomness system to a new epoch. Cancels all partial signature sends for
76    /// prior epochs.
77    pub fn update_epoch(
78        &self,
79        new_epoch: EpochId,
80        authority_info: HashMap<AuthorityName, (PeerId, PartyId)>,
81        dkg_output: dkg_v1::Output<bls12381::G2Element, bls12381::G2Element>,
82        aggregation_threshold: u16,
83        recovered_last_completed_round: Option<RandomnessRound>, // set to None if not starting up mid-epoch
84    ) {
85        self.sender
86            .try_send(RandomnessMessage::UpdateEpoch(
87                new_epoch,
88                authority_info,
89                dkg_output,
90                aggregation_threshold,
91                recovered_last_completed_round,
92            ))
93            .expect("RandomnessEventLoop mailbox should not overflow or be closed")
94    }
95
96    /// Begins transmitting partial signatures for the given epoch and round until completed.
97    pub fn send_partial_signatures(&self, epoch: EpochId, round: RandomnessRound) {
98        self.sender
99            .try_send(RandomnessMessage::SendPartialSignatures(epoch, round))
100            .expect("RandomnessEventLoop mailbox should not overflow or be closed")
101    }
102
103    /// Records the given round as complete, stopping any partial signature sends.
104    pub fn complete_round(&self, epoch: EpochId, round: RandomnessRound) {
105        self.sender
106            .try_send(RandomnessMessage::CompleteRound(epoch, round))
107            .expect("RandomnessEventLoop mailbox should not overflow or be closed")
108    }
109
110    /// Admin interface handler: generates partial signatures for the given round at the
111    /// current epoch.
112    pub fn admin_get_partial_signatures(
113        &self,
114        round: RandomnessRound,
115        tx: oneshot::Sender<Vec<u8>>,
116    ) {
117        self.sender
118            .try_send(RandomnessMessage::AdminGetPartialSignatures(round, tx))
119            .expect("RandomnessEventLoop mailbox should not overflow or be closed")
120    }
121
122    /// Admin interface handler: injects partial signatures for the given round at the
123    /// current epoch, skipping validity checks.
124    pub fn admin_inject_partial_signatures(
125        &self,
126        authority_name: AuthorityName,
127        round: RandomnessRound,
128        sigs: Vec<RandomnessPartialSignature>,
129        result_channel: oneshot::Sender<Result<()>>,
130    ) {
131        self.sender
132            .try_send(RandomnessMessage::AdminInjectPartialSignatures(
133                authority_name,
134                round,
135                sigs,
136                result_channel,
137            ))
138            .expect("RandomnessEventLoop mailbox should not overflow or be closed")
139    }
140
141    /// Admin interface handler: injects full signature for the given round at the
142    /// current epoch, skipping validity checks.
143    pub fn admin_inject_full_signature(
144        &self,
145        round: RandomnessRound,
146        sig: RandomnessSignature,
147        result_channel: oneshot::Sender<Result<()>>,
148    ) {
149        self.sender
150            .try_send(RandomnessMessage::AdminInjectFullSignature(
151                round,
152                sig,
153                result_channel,
154            ))
155            .expect("RandomnessEventLoop mailbox should not overflow or be closed")
156    }
157
158    // For testing.
159    pub fn new_stub() -> Self {
160        let (sender, mut receiver) = mpsc::channel(100);
161        // Keep receiver open until all senders are closed.
162        tokio::spawn(async move {
163            loop {
164                tokio::select! {
165                    m = receiver.recv() => {
166                        if m.is_none() {
167                            break;
168                        }
169                    },
170                }
171            }
172        });
173        Self { sender }
174    }
175}
176
177#[derive(Debug)]
178enum RandomnessMessage {
179    UpdateEpoch(
180        EpochId,
181        HashMap<AuthorityName, (PeerId, PartyId)>,
182        dkg_v1::Output<bls12381::G2Element, bls12381::G2Element>,
183        u16,                     // aggregation_threshold
184        Option<RandomnessRound>, // recovered_highest_completed_round
185    ),
186    SendPartialSignatures(EpochId, RandomnessRound),
187    CompleteRound(EpochId, RandomnessRound),
188    ReceiveSignatures(
189        PeerId,
190        EpochId,
191        RandomnessRound,
192        Vec<Vec<u8>>,
193        Option<RandomnessSignature>,
194    ),
195    MaybeIgnoreByzantinePeer(EpochId, PeerId),
196    AdminGetPartialSignatures(RandomnessRound, oneshot::Sender<Vec<u8>>),
197    AdminInjectPartialSignatures(
198        AuthorityName,
199        RandomnessRound,
200        Vec<RandomnessPartialSignature>,
201        oneshot::Sender<Result<()>>,
202    ),
203    AdminInjectFullSignature(
204        RandomnessRound,
205        RandomnessSignature,
206        oneshot::Sender<Result<()>>,
207    ),
208}
209
210struct RandomnessEventLoop {
211    name: AuthorityName,
212    config: RandomnessConfig,
213    mailbox: mpsc::Receiver<RandomnessMessage>,
214    mailbox_sender: mpsc::WeakSender<RandomnessMessage>,
215    network: anemo::Network,
216    allowed_peers: AllowedPeersUpdatable,
217    allowed_peers_set: HashSet<PeerId>,
218    metrics: Metrics,
219    randomness_tx: mpsc::Sender<(EpochId, RandomnessRound, Vec<u8>)>,
220
221    epoch: EpochId,
222    authority_info: Arc<HashMap<AuthorityName, (PeerId, PartyId)>>,
223    peer_share_ids: Option<HashMap<PeerId, Vec<ShareIndex>>>,
224    blocked_share_id_count: usize,
225    dkg_output: Option<dkg_v1::Output<bls12381::G2Element, bls12381::G2Element>>,
226    aggregation_threshold: u16,
227    highest_requested_round: BTreeMap<EpochId, RandomnessRound>,
228    send_tasks: BTreeMap<
229        RandomnessRound,
230        (
231            tokio::task::JoinHandle<()>,
232            Arc<OnceCell<RandomnessSignature>>,
233        ),
234    >,
235    round_request_time: BTreeMap<(EpochId, RandomnessRound), time::Instant>,
236    future_epoch_partial_sigs: BTreeMap<(EpochId, RandomnessRound, PeerId), Vec<Vec<u8>>>,
237    received_partial_sigs: BTreeMap<(RandomnessRound, PeerId), Vec<RandomnessPartialSignature>>,
238    completed_sigs: BTreeMap<RandomnessRound, RandomnessSignature>,
239    highest_completed_round: BTreeMap<EpochId, RandomnessRound>,
240}
241
242impl RandomnessEventLoop {
243    pub async fn start(mut self) {
244        info!("Randomness network event loop started");
245
246        loop {
247            tokio::select! {
248                maybe_message = self.mailbox.recv() => {
249                    // Once all handles to our mailbox have been dropped this
250                    // will yield `None` and we can terminate the event loop.
251                    if let Some(message) = maybe_message {
252                        self.handle_message(message);
253                    } else {
254                        break;
255                    }
256                },
257            }
258        }
259
260        info!("Randomness network event loop ended");
261    }
262
263    fn handle_message(&mut self, message: RandomnessMessage) {
264        match message {
265            RandomnessMessage::UpdateEpoch(
266                epoch,
267                authority_info,
268                dkg_output,
269                aggregation_threshold,
270                recovered_highest_completed_round,
271            ) => {
272                if let Err(e) = self.update_epoch(
273                    epoch,
274                    authority_info,
275                    dkg_output,
276                    aggregation_threshold,
277                    recovered_highest_completed_round,
278                ) {
279                    error!("BUG: failed to update epoch in RandomnessEventLoop: {e:?}");
280                }
281            }
282            RandomnessMessage::SendPartialSignatures(epoch, round) => {
283                self.send_partial_signatures(epoch, round)
284            }
285            RandomnessMessage::CompleteRound(epoch, round) => self.complete_round(epoch, round),
286            RandomnessMessage::ReceiveSignatures(peer_id, epoch, round, partial_sigs, sig) => {
287                if let Some(sig) = sig {
288                    self.receive_full_signature(peer_id, epoch, round, sig)
289                } else {
290                    self.receive_partial_signatures(peer_id, epoch, round, partial_sigs)
291                }
292            }
293            RandomnessMessage::MaybeIgnoreByzantinePeer(epoch, peer_id) => {
294                self.maybe_ignore_byzantine_peer(epoch, peer_id)
295            }
296            RandomnessMessage::AdminGetPartialSignatures(round, tx) => {
297                self.admin_get_partial_signatures(round, tx)
298            }
299            RandomnessMessage::AdminInjectPartialSignatures(
300                authority_name,
301                round,
302                sigs,
303                result_channel,
304            ) => {
305                let _ = result_channel.send(self.admin_inject_partial_signatures(
306                    authority_name,
307                    round,
308                    sigs,
309                ));
310            }
311            RandomnessMessage::AdminInjectFullSignature(round, sig, result_channel) => {
312                let _ = result_channel.send(self.admin_inject_full_signature(round, sig));
313            }
314        }
315    }
316
317    #[instrument(level = "debug", skip_all, fields(?new_epoch))]
318    fn update_epoch(
319        &mut self,
320        new_epoch: EpochId,
321        authority_info: HashMap<AuthorityName, (PeerId, PartyId)>,
322        dkg_output: dkg_v1::Output<bls12381::G2Element, bls12381::G2Element>,
323        aggregation_threshold: u16,
324        recovered_highest_completed_round: Option<RandomnessRound>,
325    ) -> Result<()> {
326        assert!(self.dkg_output.is_none() || new_epoch > self.epoch);
327
328        debug!("updating randomness network loop to new epoch");
329
330        self.peer_share_ids = Some(authority_info.iter().try_fold(
331            HashMap::new(),
332            |mut acc, (_name, (peer_id, party_id))| -> Result<_> {
333                let ids = dkg_output
334                    .nodes
335                    .share_ids_of(*party_id)
336                    .expect("party_id should be valid");
337                acc.insert(*peer_id, ids);
338                Ok(acc)
339            },
340        )?);
341        self.allowed_peers_set = authority_info
342            .values()
343            .map(|(peer_id, _)| *peer_id)
344            .collect();
345        self.allowed_peers
346            .update(Arc::new(self.allowed_peers_set.clone()));
347        self.epoch = new_epoch;
348        self.authority_info = Arc::new(authority_info);
349        self.dkg_output = Some(dkg_output);
350        self.aggregation_threshold = aggregation_threshold;
351        if let Some(round) = recovered_highest_completed_round {
352            self.highest_completed_round
353                .entry(new_epoch)
354                .and_modify(|r| *r = std::cmp::max(*r, round))
355                .or_insert(round);
356        }
357        for (_, (task, _)) in std::mem::take(&mut self.send_tasks) {
358            task.abort();
359        }
360        self.metrics.set_epoch(new_epoch);
361
362        // Throw away info from old epochs.
363        self.highest_requested_round = self.highest_requested_round.split_off(&new_epoch);
364        self.round_request_time = self
365            .round_request_time
366            .split_off(&(new_epoch, RandomnessRound(0)));
367        self.received_partial_sigs.clear();
368        self.completed_sigs.clear();
369        self.highest_completed_round = self.highest_completed_round.split_off(&new_epoch);
370
371        // Start any pending tasks for the new epoch.
372        self.maybe_start_pending_tasks();
373
374        // Aggregate any sigs received early from the new epoch.
375        // (We can't call `maybe_aggregate_partial_signatures` directly while iterating,
376        // because it takes `&mut self`, so we store in a Vec first.)
377        for ((epoch, round, peer_id), sig_bytes) in
378            std::mem::take(&mut self.future_epoch_partial_sigs)
379        {
380            // We can fully validate these now that we have current epoch DKG output.
381            self.receive_partial_signatures(peer_id, epoch, round, sig_bytes);
382        }
383        let rounds_to_aggregate: Vec<_> =
384            self.received_partial_sigs.keys().map(|(r, _)| *r).collect();
385        for round in rounds_to_aggregate {
386            self.maybe_aggregate_partial_signatures(new_epoch, round);
387        }
388
389        Ok(())
390    }
391
392    #[instrument(level = "debug", skip_all, fields(?epoch, ?round))]
393    fn send_partial_signatures(&mut self, epoch: EpochId, round: RandomnessRound) {
394        if epoch < self.epoch {
395            error!(
396                "BUG: skipping sending partial sigs, we are already up to epoch {}",
397                self.epoch
398            );
399            debug_assert!(
400                false,
401                "skipping sending partial sigs, we are already up to higher epoch"
402            );
403            return;
404        }
405        if epoch == self.epoch
406            && let Some(highest_completed_round) = self.highest_completed_round.get(&epoch)
407            && round <= *highest_completed_round
408        {
409            info!("skipping sending partial sigs, we already have completed this round");
410            return;
411        }
412
413        self.highest_requested_round
414            .entry(epoch)
415            .and_modify(|r| *r = std::cmp::max(*r, round))
416            .or_insert(round);
417        self.round_request_time
418            .insert((epoch, round), time::Instant::now());
419        self.maybe_start_pending_tasks();
420    }
421
422    #[instrument(level = "debug", skip_all, fields(?epoch, ?round))]
423    fn complete_round(&mut self, epoch: EpochId, round: RandomnessRound) {
424        debug!("completing randomness round");
425        let new_highest_round = *self
426            .highest_completed_round
427            .entry(epoch)
428            .and_modify(|r| *r = std::cmp::max(*r, round))
429            .or_insert(round);
430        if round != new_highest_round {
431            // This round completion came out of order, and we're already ahead. Nothing more
432            // to do in that case.
433            return;
434        }
435
436        self.round_request_time = self.round_request_time.split_off(&(epoch, round + 1));
437
438        if epoch == self.epoch {
439            self.remove_partial_sigs_in_range((
440                Bound::Included((RandomnessRound(0), PeerId([0; 32]))),
441                Bound::Excluded((round + 1, PeerId([0; 32]))),
442            ));
443            self.completed_sigs = self.completed_sigs.split_off(&(round + 1));
444            for (_, (task, _)) in self.send_tasks.iter().take_while(|(r, _)| **r <= round) {
445                task.abort();
446            }
447            self.send_tasks = self.send_tasks.split_off(&(round + 1));
448            self.maybe_start_pending_tasks();
449        }
450
451        self.update_rounds_pending_metric();
452    }
453
454    #[instrument(level = "debug", skip_all, fields(?peer_id, ?epoch, ?round))]
455    fn receive_partial_signatures(
456        &mut self,
457        peer_id: PeerId,
458        epoch: EpochId,
459        round: RandomnessRound,
460        sig_bytes: Vec<Vec<u8>>,
461    ) {
462        // Basic validity checks.
463        if epoch < self.epoch {
464            debug!(
465                "skipping received partial sigs, we are already up to epoch {}",
466                self.epoch
467            );
468            return;
469        }
470        if epoch > self.epoch + 1 {
471            debug!(
472                "skipping received partial sigs, we are still on epoch {}",
473                self.epoch
474            );
475            return;
476        }
477        if epoch == self.epoch && self.completed_sigs.contains_key(&round) {
478            debug!("skipping received partial sigs, we already have completed this sig");
479            return;
480        }
481        let highest_completed_round = self.highest_completed_round.get(&epoch).copied();
482        if let Some(highest_completed_round) = &highest_completed_round
483            && *highest_completed_round >= round
484        {
485            debug!("skipping received partial sigs, we already have completed this round");
486            return;
487        }
488
489        // If sigs are for a future epoch, we can't fully verify them without DKG output.
490        // Save them for later use.
491        if epoch != self.epoch || self.peer_share_ids.is_none() {
492            if round.0 >= self.config.max_partial_sigs_rounds_ahead() {
493                debug!("skipping received partial sigs for future epoch, round too far ahead",);
494                return;
495            }
496
497            debug!("saving partial sigs from future epoch for later use");
498            self.future_epoch_partial_sigs
499                .insert((epoch, round, peer_id), sig_bytes);
500            return;
501        }
502
503        // Verify shape of sigs matches what we expect for the peer.
504        let peer_share_ids = self.peer_share_ids.as_ref().expect("checked above");
505        let expected_share_ids = if let Some(expected_share_ids) = peer_share_ids.get(&peer_id) {
506            expected_share_ids
507        } else {
508            debug!("received partial sigs from unknown peer");
509            return;
510        };
511        if sig_bytes.len() != expected_share_ids.len() as usize {
512            warn!(
513                "received partial sigs with wrong share ids count: expected {}, got {}",
514                expected_share_ids.len(),
515                sig_bytes.len(),
516            );
517            return;
518        }
519
520        // Accept partial signatures up to `max_partial_sigs_rounds_ahead` past the round of the
521        // last completed signature, or the highest completed round, whichever is greater.
522        let last_completed_signature = self.completed_sigs.last_key_value().map(|(r, _)| *r);
523        let last_completed_round = std::cmp::max(last_completed_signature, highest_completed_round)
524            .unwrap_or(RandomnessRound(0));
525        if round.0
526            >= last_completed_round
527                .0
528                .saturating_add(self.config.max_partial_sigs_rounds_ahead())
529        {
530            debug!(
531                "skipping received partial sigs, most recent round we completed was only {last_completed_round}",
532            );
533            return;
534        }
535
536        // Deserialize the partial sigs.
537        let partial_sigs =
538            match sig_bytes
539                .iter()
540                .try_fold(Vec::new(), |mut acc, bytes| -> Result<_> {
541                    let sig: RandomnessPartialSignature = bcs::from_bytes(bytes)?;
542                    acc.push(sig);
543                    Ok(acc)
544                }) {
545                Ok(partial_sigs) => partial_sigs,
546                Err(e) => {
547                    warn!("failed to deserialize partial sigs: {e:?}");
548                    return;
549                }
550            };
551        // Verify we received the expected share IDs (to protect against a validator that sends
552        // valid signatures of other peers which will be successfully verified below).
553        let received_share_ids = partial_sigs.iter().map(|s| s.index);
554        if received_share_ids
555            .zip_debug_eq(expected_share_ids.iter())
556            .any(|(a, b)| a != *b)
557        {
558            let received_share_ids = partial_sigs.iter().map(|s| s.index).collect::<Vec<_>>();
559            warn!(
560                "received partial sigs with wrong share ids: expected {expected_share_ids:?}, received {received_share_ids:?}"
561            );
562            return;
563        }
564
565        // We passed all the checks, save the partial sigs.
566        debug!("recording received partial signatures");
567        self.received_partial_sigs
568            .insert((round, peer_id), partial_sigs);
569
570        self.maybe_aggregate_partial_signatures(epoch, round);
571    }
572
573    #[instrument(level = "debug", skip_all, fields(?epoch, ?round))]
574    fn maybe_aggregate_partial_signatures(&mut self, epoch: EpochId, round: RandomnessRound) {
575        if let Some(highest_completed_round) = self.highest_completed_round.get(&epoch)
576            && round <= *highest_completed_round
577        {
578            info!("skipping aggregation for already-completed round");
579            return;
580        }
581
582        let highest_requested_round = self.highest_requested_round.get(&epoch);
583        if highest_requested_round.is_none() || round > *highest_requested_round.unwrap() {
584            // We have to wait here, because even if we have enough information from other nodes
585            // to complete the signature, local shared object versions are not set until consensus
586            // finishes processing the corresponding commit. This function will be called again
587            // after maybe_start_pending_tasks begins this round locally.
588            debug!(
589                "waiting to aggregate randomness partial signatures until local consensus catches up"
590            );
591            return;
592        }
593
594        if epoch != self.epoch {
595            debug!(
596                "waiting to aggregate randomness partial signatures until DKG completes for epoch"
597            );
598            return;
599        }
600
601        if self.completed_sigs.contains_key(&round) {
602            info!("skipping aggregation for already-completed signature");
603            return;
604        }
605
606        let vss_pk = {
607            let Some(dkg_output) = &self.dkg_output else {
608                debug!("called maybe_aggregate_partial_signatures before DKG completed");
609                return;
610            };
611            &dkg_output.vss_pk
612        };
613
614        let sig_bounds = (
615            Bound::Included((round, PeerId([0; 32]))),
616            Bound::Excluded((round + 1, PeerId([0; 32]))),
617        );
618
619        // If we have enough partial signatures, aggregate them.
620        let sig_range = self
621            .received_partial_sigs
622            .range(sig_bounds)
623            .flat_map(|(_, sigs)| sigs);
624        let mut sig =
625            match ThresholdBls12381MinSig::aggregate(self.aggregation_threshold, sig_range) {
626                Ok(sig) => sig,
627                Err(fastcrypto::error::FastCryptoError::NotEnoughInputs) => return, // wait for more input
628                Err(e) => {
629                    error!("error while aggregating randomness partial signatures: {e:?}");
630                    return;
631                }
632            };
633
634        // Try to verify the aggregated signature all at once. (Should work in the happy path.)
635        if ThresholdBls12381MinSig::verify(&vss_pk.c0(), &round.signature_message(), &sig).is_err()
636        {
637            // If verifiation fails, some of the inputs must be invalid. We have to go through
638            // one-by-one to find which.
639            // TODO: add test for individual sig verification.
640            self.received_partial_sigs
641                .retain(|&(r, peer_id), partial_sigs| {
642                    if round != r {
643                        return true;
644                    }
645                    if ThresholdBls12381MinSig::partial_verify_batch(
646                        vss_pk,
647                        &round.signature_message(),
648                        partial_sigs.iter(),
649                        &mut rand::thread_rng(),
650                    )
651                    .is_err()
652                    {
653                        warn!(
654                            "received invalid partial signatures from possibly-Byzantine peer {peer_id}"
655                        );
656                        if let Some(sender) = self.mailbox_sender.upgrade() {
657                            sender.try_send(RandomnessMessage::MaybeIgnoreByzantinePeer(
658                                epoch,
659                                peer_id,
660                            ))
661                            .expect("RandomnessEventLoop mailbox should not overflow or be closed");
662                        }
663                        return false;
664                    }
665                    true
666                });
667            let sig_range = self
668                .received_partial_sigs
669                .range(sig_bounds)
670                .flat_map(|(_, sigs)| sigs);
671            sig = match ThresholdBls12381MinSig::aggregate(self.aggregation_threshold, sig_range) {
672                Ok(sig) => sig,
673                Err(fastcrypto::error::FastCryptoError::NotEnoughInputs) => return, // wait for more input
674                Err(e) => {
675                    error!("error while aggregating randomness partial signatures: {e:?}");
676                    return;
677                }
678            };
679            if let Err(e) =
680                ThresholdBls12381MinSig::verify(&vss_pk.c0(), &round.signature_message(), &sig)
681            {
682                error!(
683                    "error while verifying randomness partial signatures after removing invalid partials: {e:?}"
684                );
685                debug_assert!(
686                    false,
687                    "error while verifying randomness partial signatures after removing invalid partials"
688                );
689                return;
690            }
691        }
692
693        debug!("successfully generated randomness full signature");
694        self.process_valid_full_signature(epoch, round, sig);
695    }
696
697    #[instrument(level = "debug", skip_all, fields(?peer_id, ?epoch, ?round))]
698    fn receive_full_signature(
699        &mut self,
700        peer_id: PeerId,
701        epoch: EpochId,
702        round: RandomnessRound,
703        sig: RandomnessSignature,
704    ) {
705        let vss_pk = {
706            let Some(dkg_output) = &self.dkg_output else {
707                debug!("called receive_full_signature before DKG completed");
708                return;
709            };
710            &dkg_output.vss_pk
711        };
712
713        // Basic validity checks.
714        if epoch != self.epoch {
715            debug!("skipping received full sig, we are on epoch {}", self.epoch);
716            return;
717        }
718        if self.completed_sigs.contains_key(&round) {
719            debug!("skipping received full sigs, we already have completed this sig");
720            return;
721        }
722        let highest_completed_round = self.highest_completed_round.get(&epoch).copied();
723        if let Some(highest_completed_round) = &highest_completed_round
724            && *highest_completed_round >= round
725        {
726            debug!("skipping received full sig, we already have completed this round");
727            return;
728        }
729
730        let highest_requested_round = self.highest_requested_round.get(&epoch);
731        if highest_requested_round.is_none() || round > *highest_requested_round.unwrap() {
732            // Wait for local consensus to catch up if necessary.
733            debug!(
734                "skipping received full signature, local consensus is not caught up to its round"
735            );
736            return;
737        }
738
739        if let Err(e) =
740            ThresholdBls12381MinSig::verify(&vss_pk.c0(), &round.signature_message(), &sig)
741        {
742            info!("received invalid full signature from peer {peer_id}: {e:?}");
743            if let Some(sender) = self.mailbox_sender.upgrade() {
744                sender
745                    .try_send(RandomnessMessage::MaybeIgnoreByzantinePeer(epoch, peer_id))
746                    .expect("RandomnessEventLoop mailbox should not overflow or be closed");
747            }
748            return;
749        }
750
751        debug!("received valid randomness full signature");
752        self.process_valid_full_signature(epoch, round, sig);
753    }
754
755    fn process_valid_full_signature(
756        &mut self,
757        epoch: EpochId,
758        round: RandomnessRound,
759        sig: RandomnessSignature,
760    ) {
761        assert_eq!(epoch, self.epoch);
762
763        if let Some((_, full_sig_cell)) = self.send_tasks.get(&round) {
764            full_sig_cell
765                .set(sig)
766                .expect("full signature should never be processed twice");
767        }
768        self.completed_sigs.insert(round, sig);
769        self.remove_partial_sigs_in_range((
770            Bound::Included((round, PeerId([0; 32]))),
771            Bound::Excluded((round + 1, PeerId([0; 32]))),
772        ));
773        self.metrics.record_completed_round(round);
774        if let Some(start_time) = self.round_request_time.get(&(epoch, round))
775            && let Some(metric) = self.metrics.round_generation_latency_metric()
776        {
777            metric.observe(start_time.elapsed().as_secs_f64());
778        }
779
780        let sig_bytes = bcs::to_bytes(&sig).expect("signature serialization should not fail");
781        if let Err(e) = self.randomness_tx.try_send((epoch, round, sig_bytes)) {
782            match e {
783                // Receiver is torn down during node shutdown; dropping the round is harmless.
784                mpsc::error::TrySendError::Closed(_) => {
785                    info!("dropping completed randomness round {round}: receiver channel closed");
786                }
787                // Mailbox capacity is huge (default 1M); a full mailbox means a real bug.
788                mpsc::error::TrySendError::Full(_) => {
789                    panic!("RandomnessRoundReceiver mailbox should not overflow");
790                }
791            }
792        }
793    }
794
795    fn maybe_ignore_byzantine_peer(&mut self, epoch: EpochId, peer_id: PeerId) {
796        if epoch != self.epoch {
797            return; // make sure we're still on the same epoch
798        }
799        let Some(dkg_output) = &self.dkg_output else {
800            return; // can't ignore a peer if we haven't finished DKG
801        };
802        if !self.allowed_peers_set.contains(&peer_id) {
803            return; // peer is already disallowed
804        }
805        let Some(peer_share_ids) = &self.peer_share_ids else {
806            return; // can't ignore a peer if we haven't finished DKG
807        };
808        let Some(peer_shares) = peer_share_ids.get(&peer_id) else {
809            warn!("can't ignore unknown byzantine peer {peer_id:?}");
810            return;
811        };
812        let max_ignored_shares = (self.config.max_ignored_peer_weight_factor()
813            * (dkg_output.nodes.total_weight() as f64)) as usize;
814        if self.blocked_share_id_count + peer_shares.len() > max_ignored_shares {
815            warn!(
816                "ignoring byzantine peer {peer_id:?} with {} shares would exceed max ignored peer weight {max_ignored_shares}",
817                peer_shares.len()
818            );
819            return;
820        }
821
822        warn!(
823            "ignoring byzantine peer {peer_id:?} with {} shares",
824            peer_shares.len()
825        );
826        self.blocked_share_id_count += peer_shares.len();
827        self.allowed_peers_set.remove(&peer_id);
828        self.allowed_peers
829            .update(Arc::new(self.allowed_peers_set.clone()));
830        self.metrics.inc_num_ignored_byzantine_peers();
831    }
832
833    fn maybe_start_pending_tasks(&mut self) {
834        let dkg_output = if let Some(dkg_output) = &self.dkg_output {
835            dkg_output
836        } else {
837            return; // wait for DKG
838        };
839        let shares = if let Some(shares) = &dkg_output.shares {
840            shares
841        } else {
842            return; // can't participate in randomness generation without shares
843        };
844        let highest_requested_round =
845            if let Some(highest_requested_round) = self.highest_requested_round.get(&self.epoch) {
846                highest_requested_round
847            } else {
848                return; // no rounds to start
849            };
850        // Begin from the next round after the most recent one we've started (or, if none are running,
851        // after the highest completed round in the epoch).
852        let start_round = std::cmp::max(
853            if let Some(highest_completed_round) = self.highest_completed_round.get(&self.epoch) {
854                highest_completed_round.checked_add(1).unwrap()
855            } else {
856                RandomnessRound(0)
857            },
858            self.send_tasks
859                .last_key_value()
860                .map(|(r, _)| r.checked_add(1).unwrap())
861                .unwrap_or(RandomnessRound(0)),
862        );
863
864        let mut rounds_to_aggregate = Vec::new();
865        for round in start_round.0..=highest_requested_round.0 {
866            let round = RandomnessRound(round);
867
868            if self.send_tasks.len() >= self.config.max_partial_sigs_concurrent_sends() {
869                break; // limit concurrent tasks
870            }
871
872            let full_sig_cell = Arc::new(OnceCell::new());
873            self.send_tasks.entry(round).or_insert_with(|| {
874                let name = self.name;
875                let network = self.network.clone();
876                let retry_interval = self.config.partial_signature_retry_interval();
877                let metrics = self.metrics.clone();
878                let authority_info = self.authority_info.clone();
879                let epoch = self.epoch;
880                let partial_sigs = ThresholdBls12381MinSig::partial_sign_batch(
881                    shares.iter(),
882                    &round.signature_message(),
883                );
884                let full_sig_cell_clone = full_sig_cell.clone();
885
886                // Record own partial sigs.
887                if !self.completed_sigs.contains_key(&round) {
888                    self.received_partial_sigs
889                        .insert((round, self.network.peer_id()), partial_sigs.clone());
890                    rounds_to_aggregate.push((epoch, round));
891                }
892
893                debug!("sending partial sigs for epoch {epoch}, round {round}");
894                (
895                    spawn_monitored_task!(RandomnessEventLoop::send_signatures_task(
896                        name,
897                        network,
898                        retry_interval,
899                        metrics,
900                        authority_info,
901                        epoch,
902                        round,
903                        partial_sigs,
904                        full_sig_cell_clone,
905                    )),
906                    full_sig_cell,
907                )
908            });
909        }
910
911        self.update_rounds_pending_metric();
912
913        // After starting a round, we have generated our own partial sigs. Check if that's
914        // enough for us to aggregate already.
915        for (epoch, round) in rounds_to_aggregate {
916            self.maybe_aggregate_partial_signatures(epoch, round);
917        }
918    }
919
920    #[allow(clippy::type_complexity)]
921    fn remove_partial_sigs_in_range(
922        &mut self,
923        range: (
924            Bound<(RandomnessRound, PeerId)>,
925            Bound<(RandomnessRound, PeerId)>,
926        ),
927    ) {
928        let keys_to_remove: Vec<_> = self
929            .received_partial_sigs
930            .range(range)
931            .map(|(key, _)| *key)
932            .collect();
933        for key in keys_to_remove {
934            // Have to remove keys one-by-one because BTreeMap does not support range-removal.
935            self.received_partial_sigs.remove(&key);
936        }
937    }
938
939    async fn send_signatures_task(
940        name: AuthorityName,
941        network: anemo::Network,
942        retry_interval: Duration,
943        metrics: Metrics,
944        authority_info: Arc<HashMap<AuthorityName, (PeerId, PartyId)>>,
945        epoch: EpochId,
946        round: RandomnessRound,
947        partial_sigs: Vec<RandomnessPartialSignature>,
948        full_sig: Arc<OnceCell<RandomnessSignature>>,
949    ) {
950        // For simtests, we may test not sending partial signatures.
951        #[allow(unused_mut)]
952        let mut fail_point_skip_sending = false;
953        fail_point_if!("rb-send-partial-signatures", || {
954            fail_point_skip_sending = true;
955        });
956        if fail_point_skip_sending {
957            warn!("skipping sending partial sigs due to simtest fail point");
958            return;
959        }
960
961        let _metrics_guard = metrics
962            .round_observation_latency_metric()
963            .map(|metric| metric.start_timer());
964
965        let peers: HashMap<_, _> = authority_info
966            .iter()
967            .map(|(name, (peer_id, _party_id))| (name, network.waiting_peer(*peer_id)))
968            .collect();
969        let partial_sigs: Vec<_> = partial_sigs
970            .iter()
971            .map(|sig| bcs::to_bytes(sig).expect("message serialization should not fail"))
972            .collect();
973
974        loop {
975            let mut requests = Vec::new();
976            for (peer_name, peer) in &peers {
977                if name == **peer_name {
978                    continue; // don't send partial sigs to self
979                }
980                let mut client = RandomnessClient::new(peer.clone());
981                const SEND_PARTIAL_SIGNATURES_TIMEOUT: Duration = Duration::from_secs(10);
982                let full_sig = full_sig.get().cloned();
983                let request = anemo::Request::new(SendSignaturesRequest {
984                    epoch,
985                    round,
986                    partial_sigs: if full_sig.is_none() {
987                        partial_sigs.clone()
988                    } else {
989                        Vec::new()
990                    },
991                    sig: full_sig,
992                })
993                .with_timeout(SEND_PARTIAL_SIGNATURES_TIMEOUT);
994                requests.push(async move {
995                    let result = client.send_signatures(request).await;
996                    if let Err(_error) = result {
997                        // TODO: add Display impl to anemo::rpc::Status, log it here
998                        debug!("failed to send partial signatures to {peer_name}");
999                    }
1000                });
1001            }
1002
1003            // Process all requests.
1004            futures::future::join_all(requests).await;
1005
1006            // Keep retrying send to all peers until task is aborted via external message.
1007            tokio::time::sleep(retry_interval).await;
1008        }
1009    }
1010
1011    fn update_rounds_pending_metric(&self) {
1012        let highest_requested_round = self
1013            .highest_requested_round
1014            .get(&self.epoch)
1015            .map(|r| r.0)
1016            .unwrap_or(0);
1017        let highest_completed_round = self
1018            .highest_completed_round
1019            .get(&self.epoch)
1020            .map(|r| r.0)
1021            .unwrap_or(0);
1022        let num_rounds_pending =
1023            highest_requested_round.saturating_sub(highest_completed_round) as i64;
1024        let prev_value = self.metrics.num_rounds_pending().unwrap_or_default();
1025        if num_rounds_pending / 100 > prev_value / 100 {
1026            warn!(
1027                // Recording multiples of 100 so tests can match on the log message.
1028                "RandomnessEventLoop randomness generation backlog: over {} rounds are pending (oldest is {:?})",
1029                (num_rounds_pending / 100) * 100,
1030                highest_completed_round + 1,
1031            );
1032        }
1033        self.metrics.set_num_rounds_pending(num_rounds_pending);
1034    }
1035
1036    fn admin_get_partial_signatures(&self, round: RandomnessRound, tx: oneshot::Sender<Vec<u8>>) {
1037        let shares = if let Some(shares) = self.dkg_output.as_ref().and_then(|d| d.shares.as_ref())
1038        {
1039            shares
1040        } else {
1041            let _ = tx.send(Vec::new()); // no error handling needed if receiver is already dropped
1042            return;
1043        };
1044
1045        let partial_sigs =
1046            ThresholdBls12381MinSig::partial_sign_batch(shares.iter(), &round.signature_message());
1047        // no error handling needed if receiver is already dropped
1048        let _ = tx.send(bcs::to_bytes(&partial_sigs).expect("serialization should not fail"));
1049    }
1050
1051    fn admin_inject_partial_signatures(
1052        &mut self,
1053        authority_name: AuthorityName,
1054        round: RandomnessRound,
1055        sigs: Vec<RandomnessPartialSignature>,
1056    ) -> Result<()> {
1057        let peer_id = self
1058            .authority_info
1059            .get(&authority_name)
1060            .map(|(peer_id, _)| *peer_id)
1061            .ok_or(anyhow::anyhow!("unknown AuthorityName {authority_name:?}"))?;
1062        self.received_partial_sigs.insert((round, peer_id), sigs);
1063        self.maybe_aggregate_partial_signatures(self.epoch, round);
1064        Ok(())
1065    }
1066
1067    fn admin_inject_full_signature(
1068        &mut self,
1069        round: RandomnessRound,
1070        sig: RandomnessSignature,
1071    ) -> Result<()> {
1072        let vss_pk = {
1073            let Some(dkg_output) = &self.dkg_output else {
1074                return Err(anyhow::anyhow!(
1075                    "called admin_inject_full_signature before DKG completed"
1076                ));
1077            };
1078            &dkg_output.vss_pk
1079        };
1080
1081        ThresholdBls12381MinSig::verify(&vss_pk.c0(), &round.signature_message(), &sig)
1082            .map_err(|e| anyhow::anyhow!("invalid full signature: {e:?}"))?;
1083
1084        self.process_valid_full_signature(self.epoch, round, sig);
1085        Ok(())
1086    }
1087}