Skip to main content

iota_network/randomness/
mod.rs

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