sui_core/
stake_aggregator.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use serde::Serialize;
5use shared_crypto::intent::Intent;
6use std::collections::hash_map::Entry;
7use std::collections::{BTreeMap, HashMap, HashSet};
8use std::hash::Hash;
9use std::sync::Arc;
10use sui_types::base_types::AuthorityName;
11use sui_types::base_types::ConciseableName;
12use sui_types::committee::{Committee, CommitteeTrait, StakeUnit};
13use sui_types::crypto::{AuthorityQuorumSignInfo, AuthoritySignInfo, AuthoritySignInfoTrait};
14use sui_types::error::{SuiError, SuiErrorKind, SuiResult};
15use sui_types::message_envelope::{Envelope, Message};
16use tracing::warn;
17use typed_store::TypedStoreError;
18
19/// StakeAggregator allows us to keep track of the total stake of a set of validators.
20/// STRENGTH indicates whether we want a strong quorum (2f+1) or a weak quorum (f+1).
21#[derive(Debug)]
22pub struct StakeAggregator<S, const STRENGTH: bool> {
23    data: HashMap<AuthorityName, S>,
24    total_votes: StakeUnit,
25    committee: Arc<Committee>,
26}
27
28/// StakeAggregator is a utility data structure that allows us to aggregate a list of validator
29/// signatures over time. A committee is used to determine whether we have reached sufficient
30/// quorum (defined based on `STRENGTH`). The generic implementation does not require `S` to be
31/// an actual signature, but just an indication that a specific validator has voted. A specialized
32/// implementation for `AuthoritySignInfo` is followed below.
33impl<S: Clone + Eq, const STRENGTH: bool> StakeAggregator<S, STRENGTH> {
34    pub fn new(committee: Arc<Committee>) -> Self {
35        Self {
36            data: Default::default(),
37            total_votes: Default::default(),
38            committee,
39        }
40    }
41
42    pub fn from_iter<I: Iterator<Item = Result<(AuthorityName, S), TypedStoreError>>>(
43        committee: Arc<Committee>,
44        data: I,
45    ) -> SuiResult<Self> {
46        let mut this = Self::new(committee);
47        for item in data {
48            let (authority, s) = item?;
49            this.insert_generic(authority, s);
50        }
51        Ok(this)
52    }
53
54    /// A generic version of inserting arbitrary type of V (e.g. void type).
55    /// If V is AuthoritySignInfo, the `insert` function should be used instead since it does extra
56    /// checks and aggregations in the end.
57    /// Returns Map authority -> S, without aggregating it.
58    /// If you want to get an aggregated signature instead, use `StakeAggregator::insert`
59    pub fn insert_generic(
60        &mut self,
61        authority: AuthorityName,
62        s: S,
63    ) -> InsertResult<&HashMap<AuthorityName, S>> {
64        match self.data.entry(authority) {
65            Entry::Occupied(oc) => {
66                return InsertResult::Failed {
67                    error: SuiErrorKind::StakeAggregatorRepeatedSigner {
68                        signer: authority,
69                        conflicting_sig: oc.get() != &s,
70                    }
71                    .into(),
72                };
73            }
74            Entry::Vacant(va) => {
75                va.insert(s);
76            }
77        }
78        let votes = self.committee.weight(&authority);
79        if votes > 0 {
80            self.total_votes += votes;
81            if self.total_votes >= self.committee.threshold::<STRENGTH>() {
82                InsertResult::QuorumReached(&self.data)
83            } else {
84                InsertResult::NotEnoughVotes {
85                    bad_votes: 0,
86                    bad_authorities: vec![],
87                }
88            }
89        } else {
90            InsertResult::Failed {
91                error: SuiErrorKind::InvalidAuthenticator.into(),
92            }
93        }
94    }
95
96    pub fn contains_key(&self, authority: &AuthorityName) -> bool {
97        self.data.contains_key(authority)
98    }
99
100    pub fn keys(&self) -> impl Iterator<Item = &AuthorityName> {
101        self.data.keys()
102    }
103
104    pub fn committee(&self) -> &Committee {
105        &self.committee
106    }
107
108    pub fn total_votes(&self) -> StakeUnit {
109        self.total_votes
110    }
111
112    pub fn has_quorum(&self) -> bool {
113        self.total_votes >= self.committee.threshold::<STRENGTH>()
114    }
115
116    #[cfg(test)]
117    pub fn validator_sig_count(&self) -> usize {
118        self.data.len()
119    }
120}
121
122impl<const STRENGTH: bool> StakeAggregator<AuthoritySignInfo, STRENGTH> {
123    /// Insert an authority signature. This is the primary way to use the aggregator and a few
124    /// dedicated checks are performed to make sure things work.
125    /// If quorum is reached, we return AuthorityQuorumSignInfo directly.
126    pub fn insert<T: Message + Serialize>(
127        &mut self,
128        envelope: Envelope<T, AuthoritySignInfo>,
129    ) -> InsertResult<AuthorityQuorumSignInfo<STRENGTH>> {
130        let (data, sig) = envelope.into_data_and_sig();
131        if self.committee.epoch != sig.epoch {
132            return InsertResult::Failed {
133                error: SuiErrorKind::WrongEpoch {
134                    expected_epoch: self.committee.epoch,
135                    actual_epoch: sig.epoch,
136                }
137                .into(),
138            };
139        }
140        match self.insert_generic(sig.authority, sig) {
141            InsertResult::QuorumReached(_) => {
142                match AuthorityQuorumSignInfo::<STRENGTH>::new_from_auth_sign_infos(
143                    self.data.values().cloned().collect(),
144                    self.committee(),
145                ) {
146                    Ok(aggregated) => {
147                        match aggregated.verify_secure(
148                            &data,
149                            Intent::sui_app(T::SCOPE),
150                            self.committee(),
151                        ) {
152                            // In the happy path, the aggregated signature verifies ok and no need to verify
153                            // individual.
154                            Ok(_) => InsertResult::QuorumReached(aggregated),
155                            Err(_) => {
156                                // If the aggregated signature fails to verify, fallback to iterating through
157                                // all signatures and verify individually. Decrement total votes and continue
158                                // to find new authority for signature to reach the quorum.
159                                //
160                                // TODO(joyqvq): It is possible for the aggregated signature to fail every time
161                                // when the latest one single signature fails to verify repeatedly, and trigger
162                                // this for loop to run. This can be optimized by caching single sig verification
163                                // result only verify the net new ones.
164                                let mut bad_votes = 0;
165                                let mut bad_authorities = vec![];
166                                for (name, sig) in &self.data.clone() {
167                                    if let Err(err) = sig.verify_secure(
168                                        &data,
169                                        Intent::sui_app(T::SCOPE),
170                                        self.committee(),
171                                    ) {
172                                        // TODO(joyqvq): Currently, the aggregator cannot do much with an authority that
173                                        // always returns an invalid signature other than saving to errors in state. It
174                                        // is possible to add the authority to a denylist or  punish the byzantine authority.
175                                        warn!(name=?name.concise(), "Bad stake from validator: {:?}", err);
176                                        self.data.remove(name);
177                                        let votes = self.committee.weight(name);
178                                        self.total_votes -= votes;
179                                        bad_votes += votes;
180                                        bad_authorities.push(*name);
181                                    }
182                                }
183                                // After evicting invalid sigs, the remaining valid sigs may
184                                // still constitute a quorum on their own.
185                                if self.total_votes >= self.committee.threshold::<STRENGTH>() {
186                                    match AuthorityQuorumSignInfo::<STRENGTH>::new_from_auth_sign_infos(
187                                        self.data.values().cloned().collect(),
188                                        self.committee(),
189                                    ) {
190                                        Ok(aggregated) => InsertResult::QuorumReached(aggregated),
191                                        Err(error) => InsertResult::Failed { error },
192                                    }
193                                } else {
194                                    InsertResult::NotEnoughVotes {
195                                        bad_votes,
196                                        bad_authorities,
197                                    }
198                                }
199                            }
200                        }
201                    }
202                    Err(error) => InsertResult::Failed { error },
203                }
204            }
205            // The following is necessary to change the template type of InsertResult.
206            InsertResult::Failed { error } => InsertResult::Failed { error },
207            InsertResult::NotEnoughVotes {
208                bad_votes,
209                bad_authorities,
210            } => InsertResult::NotEnoughVotes {
211                bad_votes,
212                bad_authorities,
213            },
214        }
215    }
216}
217
218pub enum InsertResult<CertT> {
219    QuorumReached(CertT),
220    Failed {
221        error: SuiError,
222    },
223    NotEnoughVotes {
224        bad_votes: u64,
225        bad_authorities: Vec<AuthorityName>,
226    },
227}
228
229impl<CertT> InsertResult<CertT> {
230    pub fn is_quorum_reached(&self) -> bool {
231        matches!(self, Self::QuorumReached(..))
232    }
233}
234
235/// MultiStakeAggregator is a utility data structure that tracks the stake accumulation of
236/// potentially multiple different values (usually due to byzantine/corrupted responses). Each
237/// value is tracked using a StakeAggregator and determine whether it has reached a quorum.
238/// Once quorum is reached, the aggregated signature is returned.
239#[derive(Debug)]
240pub struct MultiStakeAggregator<K, V, const STRENGTH: bool> {
241    committee: Arc<Committee>,
242    stake_maps: HashMap<K, (V, StakeAggregator<AuthoritySignInfo, STRENGTH>)>,
243}
244
245impl<K, V, const STRENGTH: bool> MultiStakeAggregator<K, V, STRENGTH> {
246    pub fn new(committee: Arc<Committee>) -> Self {
247        Self {
248            committee,
249            stake_maps: Default::default(),
250        }
251    }
252
253    pub fn total_votes(&self) -> StakeUnit {
254        let mut voted_authorities = HashSet::new();
255        self.stake_maps.values().for_each(|(_, stake_aggregator)| {
256            stake_aggregator.keys().for_each(|k| {
257                voted_authorities.insert(k);
258            })
259        });
260        voted_authorities
261            .iter()
262            .map(|k| self.committee.weight(k))
263            .sum()
264    }
265
266    #[cfg(test)]
267    pub fn unique_key_count(&self) -> usize {
268        self.stake_maps.len()
269    }
270}
271
272impl<K, V, const STRENGTH: bool> MultiStakeAggregator<K, V, STRENGTH>
273where
274    K: Hash + Eq,
275    V: Message + Serialize + Clone,
276{
277    pub fn insert(
278        &mut self,
279        k: K,
280        envelope: Envelope<V, AuthoritySignInfo>,
281    ) -> InsertResult<AuthorityQuorumSignInfo<STRENGTH>> {
282        if let Some(entry) = self.stake_maps.get_mut(&k) {
283            entry.1.insert(envelope)
284        } else {
285            let mut new_entry = StakeAggregator::new(self.committee.clone());
286            let result = new_entry.insert(envelope.clone());
287            if !matches!(result, InsertResult::Failed { .. }) {
288                // This is very important: ensure that if the insert fails, we don't even add the
289                // new entry to the map.
290                self.stake_maps.insert(k, (envelope.into_data(), new_entry));
291            }
292            result
293        }
294    }
295}
296
297impl<K, V, const STRENGTH: bool> MultiStakeAggregator<K, V, STRENGTH>
298where
299    K: Clone + Ord,
300{
301    pub fn get_all_unique_values(&self) -> BTreeMap<K, (Vec<AuthorityName>, StakeUnit)> {
302        self.stake_maps
303            .iter()
304            .map(|(k, (_, s))| (k.clone(), (s.data.keys().copied().collect(), s.total_votes)))
305            .collect()
306    }
307}
308
309impl<K, V, const STRENGTH: bool> MultiStakeAggregator<K, V, STRENGTH>
310where
311    K: Hash + Eq,
312{
313    #[allow(dead_code)]
314    pub fn authorities_for_key(&self, k: &K) -> Option<impl Iterator<Item = &AuthorityName>> {
315        self.stake_maps.get(k).map(|(_, agg)| agg.keys())
316    }
317
318    /// The sum of all remaining stake, i.e. all stake not yet
319    /// committed by vote to a specific value
320    pub fn uncommitted_stake(&self) -> StakeUnit {
321        self.committee.total_votes() - self.total_votes()
322    }
323
324    /// Total stake of the largest faction
325    pub fn plurality_stake(&self) -> StakeUnit {
326        self.stake_maps
327            .values()
328            .map(|(_, agg)| agg.total_votes())
329            .max()
330            .unwrap_or_default()
331    }
332
333    /// If true, there isn't enough uncommitted stake to reach quorum for any value
334    pub fn quorum_unreachable(&self) -> bool {
335        self.uncommitted_stake() + self.plurality_stake() < self.committee.threshold::<STRENGTH>()
336    }
337}
338
339/// Like MultiStakeAggregator, but for counting votes for a generic value instead of an envelope, in
340/// scenarios where byzantine validators may submit multiple votes for different values.
341pub struct GenericMultiStakeAggregator<K, const STRENGTH: bool> {
342    committee: Arc<Committee>,
343    stake_maps: HashMap<K, StakeAggregator<(), STRENGTH>>,
344    votes_per_authority: HashMap<AuthorityName, u64>,
345}
346
347impl<K, const STRENGTH: bool> GenericMultiStakeAggregator<K, STRENGTH>
348where
349    K: Hash + Eq,
350{
351    pub fn new(committee: Arc<Committee>) -> Self {
352        Self {
353            committee,
354            stake_maps: Default::default(),
355            votes_per_authority: Default::default(),
356        }
357    }
358
359    pub fn insert(
360        &mut self,
361        authority: AuthorityName,
362        k: K,
363    ) -> InsertResult<&HashMap<AuthorityName, ()>> {
364        let agg = self
365            .stake_maps
366            .entry(k)
367            .or_insert_with(|| StakeAggregator::new(self.committee.clone()));
368
369        if !agg.contains_key(&authority) {
370            *self.votes_per_authority.entry(authority).or_default() += 1;
371        }
372
373        agg.insert_generic(authority, ())
374    }
375
376    pub fn has_quorum_for_key(&self, k: &K) -> bool {
377        if let Some(entry) = self.stake_maps.get(k) {
378            entry.has_quorum()
379        } else {
380            false
381        }
382    }
383
384    pub fn votes_for_authority(&self, authority: AuthorityName) -> u64 {
385        self.votes_per_authority
386            .get(&authority)
387            .copied()
388            .unwrap_or_default()
389    }
390}
391
392#[test]
393fn test_votes_per_authority() {
394    let (committee, _) = Committee::new_simple_test_committee();
395    let authorities: Vec<_> = committee.names().copied().collect();
396
397    let mut agg: GenericMultiStakeAggregator<&str, true> =
398        GenericMultiStakeAggregator::new(Arc::new(committee));
399
400    // 1. Inserting an `authority` and a `key`, and then checking the number of votes for that `authority`.
401    let key1: &str = "key1";
402    let authority1 = authorities[0];
403    agg.insert(authority1, key1);
404    assert_eq!(agg.votes_for_authority(authority1), 1);
405
406    // 2. Inserting the same `authority` and `key` pair multiple times to ensure votes aren't incremented incorrectly.
407    agg.insert(authority1, key1);
408    agg.insert(authority1, key1);
409    assert_eq!(agg.votes_for_authority(authority1), 1);
410
411    // 3. Checking votes for an authority that hasn't voted.
412    let authority2 = authorities[1];
413    assert_eq!(agg.votes_for_authority(authority2), 0);
414
415    // 4. Inserting multiple different authorities and checking their vote counts.
416    let key2: &str = "key2";
417    agg.insert(authority2, key2);
418    assert_eq!(agg.votes_for_authority(authority2), 1);
419    assert_eq!(agg.votes_for_authority(authority1), 1);
420
421    // 5. Verifying that inserting different keys for the same authority increments the vote count.
422    let key3: &str = "key3";
423    agg.insert(authority1, key3);
424    assert_eq!(agg.votes_for_authority(authority1), 2);
425}
426
427#[cfg(test)]
428mod multi_stake_aggregator_tests {
429    use super::*;
430    use fastcrypto::hash::{HashFunction, Sha3_256};
431    use shared_crypto::intent::IntentScope;
432
433    #[derive(Clone, Debug, Serialize, PartialEq, Eq, Hash)]
434    struct TestMessage {
435        value: String,
436    }
437
438    impl Message for TestMessage {
439        type DigestType = [u8; 32];
440        const SCOPE: IntentScope = IntentScope::SenderSignedTransaction;
441
442        fn digest(&self) -> Self::DigestType {
443            let mut hasher = Sha3_256::default();
444            hasher.update(self.value.as_bytes());
445            hasher.finalize().digest
446        }
447    }
448
449    #[test]
450    fn test_equivocation_stake_not_double_counted() {
451        let (committee, key_pairs) = Committee::new_simple_test_committee();
452        let committee = Arc::new(committee);
453        let authorities: Vec<_> = committee.names().copied().collect();
454
455        let mut agg: MultiStakeAggregator<String, TestMessage, true> =
456            MultiStakeAggregator::new(committee.clone());
457
458        // Get the actual total stake from the committee
459        let total_stake = committee.total_votes();
460        let num_authorities = authorities.len();
461        let stake_per_authority = total_stake / num_authorities as u64;
462
463        // Simulate equivocation: authority0 signs multiple different values
464        let authority0 = authorities[0];
465        let key0 = &key_pairs[0];
466
467        // First signature for "value1"
468        let msg1 = TestMessage {
469            value: "value1".to_string(),
470        };
471        let envelope1 =
472            <Envelope<TestMessage, AuthoritySignInfo>>::new(0, msg1.clone(), key0, authority0);
473        agg.insert("key1".to_string(), envelope1);
474
475        // Second signature from same authority for "value2" (equivocation)
476        let msg2 = TestMessage {
477            value: "value2".to_string(),
478        };
479        let envelope2 =
480            <Envelope<TestMessage, AuthoritySignInfo>>::new(0, msg2.clone(), key0, authority0);
481        agg.insert("key2".to_string(), envelope2);
482
483        // Third signature from same authority for "value3" (more equivocation)
484        let msg3 = TestMessage {
485            value: "value3".to_string(),
486        };
487        let envelope3 =
488            <Envelope<TestMessage, AuthoritySignInfo>>::new(0, msg3.clone(), key0, authority0);
489        agg.insert("key3".to_string(), envelope3);
490
491        // With the fix: authority0's stake should only be counted once, even though they signed 3 different values
492        let aggregated_votes = agg.total_votes();
493        assert_eq!(aggregated_votes, stake_per_authority);
494
495        // Add more authorities signing different values
496        let authority1 = authorities[1];
497        let key1 = &key_pairs[1];
498        let envelope4 =
499            <Envelope<TestMessage, AuthoritySignInfo>>::new(0, msg1.clone(), key1, authority1);
500        agg.insert("key1".to_string(), envelope4);
501
502        let authority2 = authorities[2];
503        let key2 = &key_pairs[2];
504        let envelope5 = <Envelope<TestMessage, AuthoritySignInfo>>::new(0, msg2, key2, authority2);
505        agg.insert("key2".to_string(), envelope5);
506
507        // Now total_votes() should be stake_per_authority * 3 (3 unique authorities)
508        // NOT stake_per_authority * 5 (which would be if we double-counted authority0)
509        let aggregated_votes = agg.total_votes();
510        assert_eq!(aggregated_votes, stake_per_authority * 3);
511        assert!(aggregated_votes <= total_stake);
512
513        // uncommitted_stake should work without underflow
514        let uncommitted = agg.uncommitted_stake();
515        assert_eq!(uncommitted, stake_per_authority); // Only authority3 hasn't voted
516
517        // Verify we have 3 different keys with votes
518        assert_eq!(agg.unique_key_count(), 3);
519    }
520
521    #[test]
522    fn test_multistake_uncommitted_and_plurality() {
523        let (committee, key_pairs) = Committee::new_simple_test_committee();
524        let committee = Arc::new(committee);
525        let authorities: Vec<_> = committee.names().copied().collect();
526
527        let mut agg: MultiStakeAggregator<String, TestMessage, true> =
528            MultiStakeAggregator::new(committee.clone());
529
530        let total_stake = committee.total_votes();
531        let num_authorities = authorities.len();
532        let stake_per_authority = total_stake / num_authorities as u64;
533
534        // Initially, all stake is uncommitted
535        assert_eq!(agg.uncommitted_stake(), total_stake);
536        assert_eq!(agg.plurality_stake(), 0);
537        assert!(!agg.quorum_unreachable());
538
539        // Add first authority voting for value1
540        let msg1 = TestMessage {
541            value: "value1".to_string(),
542        };
543        let envelope1 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
544            0,
545            msg1.clone(),
546            &key_pairs[0],
547            authorities[0],
548        );
549        agg.insert("key1".to_string(), envelope1);
550
551        assert_eq!(agg.uncommitted_stake(), total_stake - stake_per_authority);
552        assert_eq!(agg.plurality_stake(), stake_per_authority);
553
554        // Add second authority voting for value2
555        let msg2 = TestMessage {
556            value: "value2".to_string(),
557        };
558        let envelope2 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
559            0,
560            msg2.clone(),
561            &key_pairs[1],
562            authorities[1],
563        );
564        agg.insert("key2".to_string(), envelope2);
565
566        assert_eq!(
567            agg.uncommitted_stake(),
568            total_stake - 2 * stake_per_authority
569        );
570        assert_eq!(agg.plurality_stake(), stake_per_authority);
571
572        // Add third authority voting for value1 (now value1 has plurality)
573        let envelope3 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
574            0,
575            msg1.clone(),
576            &key_pairs[2],
577            authorities[2],
578        );
579        agg.insert("key1".to_string(), envelope3);
580
581        assert_eq!(
582            agg.uncommitted_stake(),
583            total_stake - 3 * stake_per_authority
584        );
585        assert_eq!(agg.plurality_stake(), 2 * stake_per_authority);
586    }
587
588    #[test]
589    fn test_multistake_quorum_unreachable() {
590        let (committee, key_pairs) = Committee::new_simple_test_committee();
591        let committee = Arc::new(committee);
592        let authorities: Vec<_> = committee.names().copied().collect();
593
594        let mut agg: MultiStakeAggregator<String, TestMessage, true> =
595            MultiStakeAggregator::new(committee.clone());
596
597        // Split votes evenly so no value can reach quorum
598        // With 4 authorities and strong quorum needing 2f+1, we need at least 3
599        let msg1 = TestMessage {
600            value: "value1".to_string(),
601        };
602        let msg2 = TestMessage {
603            value: "value2".to_string(),
604        };
605
606        // Two authorities vote for value1
607        let envelope1 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
608            0,
609            msg1.clone(),
610            &key_pairs[0],
611            authorities[0],
612        );
613        agg.insert("key1".to_string(), envelope1);
614
615        let envelope2 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
616            0,
617            msg1.clone(),
618            &key_pairs[1],
619            authorities[1],
620        );
621        agg.insert("key1".to_string(), envelope2);
622
623        // Two authorities vote for value2
624        let envelope3 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
625            0,
626            msg2.clone(),
627            &key_pairs[2],
628            authorities[2],
629        );
630        agg.insert("key2".to_string(), envelope3);
631
632        let envelope4 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
633            0,
634            msg2.clone(),
635            &key_pairs[3],
636            authorities[3],
637        );
638        agg.insert("key2".to_string(), envelope4);
639
640        // With evenly split votes, neither can reach quorum now
641        assert!(agg.quorum_unreachable());
642    }
643}
644
645#[cfg(test)]
646mod stake_aggregator_insert_tests {
647    use super::*;
648    use fastcrypto::hash::{HashFunction, Sha3_256};
649    use shared_crypto::intent::IntentScope;
650
651    #[derive(Clone, Debug, Serialize, PartialEq, Eq, Hash)]
652    struct TestMessage {
653        value: String,
654    }
655
656    impl Message for TestMessage {
657        type DigestType = [u8; 32];
658        const SCOPE: IntentScope = IntentScope::SenderSignedTransaction;
659
660        fn digest(&self) -> Self::DigestType {
661            let mut hasher = Sha3_256::default();
662            hasher.update(self.value.as_bytes());
663            hasher.finalize().digest
664        }
665    }
666
667    /// Regression test for a bug where evicting a bad sig after batch verification failure
668    /// incorrectly returns NotEnoughVotes even when the remaining valid sigs form a quorum.
669    ///
670    /// Scenario: insert a bad sig (weight < Q), then a valid sig (weight > Q).
671    /// Batch verification fails; bad sig is evicted; valid sig alone has weight > Q.
672    /// The expected result is QuorumReached, but the bug causes NotEnoughVotes.
673    #[test]
674    fn test_quorum_not_lost_after_bad_sig_eviction() {
675        // Two-validator committee: auth0 has ~7000 weight (> QUORUM_THRESHOLD 6667),
676        // auth1 has ~3000 weight. So auth0 alone can form a strong quorum.
677        let (committee, key_pairs) =
678            Committee::new_simple_test_committee_with_normalized_voting_power(vec![7, 3]);
679        let committee = Arc::new(committee);
680        // committee.names() is sorted by AuthorityName (== public key bytes),
681        // matching the sort applied to key_pairs in the constructor above.
682        let authorities: Vec<_> = committee.names().copied().collect();
683        let (auth0, key0) = (authorities[0], &key_pairs[0]);
684        let (auth1, key1) = (authorities[1], &key_pairs[1]);
685
686        let mut agg: StakeAggregator<AuthoritySignInfo, true> =
687            StakeAggregator::new(committee.clone());
688
689        let msg = TestMessage {
690            value: "real".to_string(),
691        };
692        let msg_bad = TestMessage {
693            value: "wrong".to_string(),
694        };
695
696        // auth1 signs the wrong message. Its sig will be stored but will fail
697        // verify_secure when checked against `msg` during individual verification.
698        let envelope_bad = Envelope::<TestMessage, AuthoritySignInfo>::new(0, msg_bad, key1, auth1);
699        let result = agg.insert(envelope_bad);
700        // auth1 weight (~3000) < QUORUM_THRESHOLD (6667): no quorum yet.
701        assert!(matches!(result, InsertResult::NotEnoughVotes { .. }));
702
703        // auth0 signs the real message. Total stored weight is ~10000 >= quorum,
704        // so insert triggers the batch verification path. The batch fails because
705        // auth1's sig was for msg_bad. Individual verification evicts auth1.
706        // After eviction, auth0's weight alone (~7000) still exceeds the threshold,
707        // so the result must be QuorumReached.
708        let envelope_good = Envelope::<TestMessage, AuthoritySignInfo>::new(0, msg, key0, auth0);
709        let result = agg.insert(envelope_good);
710        assert!(
711            result.is_quorum_reached(),
712            "valid sig with weight > quorum threshold must yield QuorumReached after bad sig is evicted"
713        );
714    }
715}
716
717#[cfg(test)]
718mod stake_aggregator_tests {
719    use super::*;
720
721    #[test]
722    fn test_stake_aggregator_strong_quorum() {
723        let (committee, _) = Committee::new_simple_test_committee();
724        let committee = Arc::new(committee);
725        let authorities: Vec<_> = committee.names().copied().collect();
726
727        let mut agg: StakeAggregator<(), true> = StakeAggregator::new(committee.clone());
728
729        let total_stake = committee.total_votes();
730        let num_authorities = authorities.len();
731        let stake_per_authority = total_stake / num_authorities as u64;
732
733        assert_eq!(agg.total_votes(), 0);
734        assert!(!agg.has_quorum());
735        assert_eq!(agg.validator_sig_count(), 0);
736
737        // Add first authority - should not reach quorum yet
738        let result = agg.insert_generic(authorities[0], ());
739        assert!(matches!(result, InsertResult::NotEnoughVotes { .. }));
740        assert_eq!(agg.total_votes(), stake_per_authority);
741        assert!(!agg.has_quorum());
742        assert_eq!(agg.validator_sig_count(), 1);
743
744        // Add second authority - still not enough for strong quorum (2f+1)
745        let result = agg.insert_generic(authorities[1], ());
746        assert!(matches!(result, InsertResult::NotEnoughVotes { .. }));
747        assert_eq!(agg.total_votes(), 2 * stake_per_authority);
748        assert!(!agg.has_quorum());
749
750        // Add third authority - should reach strong quorum
751        let result = agg.insert_generic(authorities[2], ());
752        assert!(result.is_quorum_reached());
753        assert!(agg.has_quorum());
754        assert_eq!(agg.validator_sig_count(), 3);
755    }
756
757    #[test]
758    fn test_stake_aggregator_weak_quorum() {
759        let (committee, _) = Committee::new_simple_test_committee();
760        let committee = Arc::new(committee);
761        let authorities: Vec<_> = committee.names().copied().collect();
762
763        let mut agg: StakeAggregator<(), false> = StakeAggregator::new(committee.clone());
764
765        // Weak quorum (f+1) should be reached faster than strong quorum
766        let result = agg.insert_generic(authorities[0], ());
767        assert!(matches!(result, InsertResult::NotEnoughVotes { .. }));
768        assert!(!agg.has_quorum());
769
770        // Second authority should reach weak quorum
771        let result = agg.insert_generic(authorities[1], ());
772        assert!(result.is_quorum_reached());
773        assert!(agg.has_quorum());
774    }
775
776    #[test]
777    fn test_stake_aggregator_repeated_signer() {
778        let (committee, _) = Committee::new_simple_test_committee();
779        let committee = Arc::new(committee);
780        let authorities: Vec<_> = committee.names().copied().collect();
781
782        let mut agg: StakeAggregator<u32, true> = StakeAggregator::new(committee.clone());
783
784        // Insert first time - should succeed
785        let result = agg.insert_generic(authorities[0], 1);
786        assert!(matches!(result, InsertResult::NotEnoughVotes { .. }));
787
788        // Insert same authority again with same value - should fail
789        let result = agg.insert_generic(authorities[0], 1);
790        assert!(matches!(
791            result,
792            InsertResult::Failed {
793                error
794            } if matches!(error.as_inner(),  SuiErrorKind::StakeAggregatorRepeatedSigner { .. } )
795        ));
796
797        // Insert same authority with different value - should also fail (conflicting signature)
798        let result = agg.insert_generic(authorities[0], 2);
799        let InsertResult::Failed { error } = result else {
800            panic!("Expected StakeAggregatorRepeatedSigner error");
801        };
802        let SuiErrorKind::StakeAggregatorRepeatedSigner {
803            signer,
804            conflicting_sig,
805        } = error.into_inner()
806        else {
807            panic!("Expected StakeAggregatorRepeatedSigner error");
808        };
809        assert_eq!(signer, authorities[0]);
810        assert!(conflicting_sig);
811    }
812
813    #[test]
814    fn test_stake_aggregator_from_iter() {
815        let (committee, _) = Committee::new_simple_test_committee();
816        let committee = Arc::new(committee);
817        let authorities: Vec<_> = committee.names().copied().collect();
818
819        let data = vec![
820            Ok((authorities[0], ())),
821            Ok((authorities[1], ())),
822            Ok((authorities[2], ())),
823        ];
824
825        let agg: StakeAggregator<(), true> =
826            StakeAggregator::from_iter(committee.clone(), data.into_iter()).unwrap();
827
828        assert_eq!(agg.validator_sig_count(), 3);
829        assert!(agg.has_quorum());
830        assert!(agg.contains_key(&authorities[0]));
831        assert!(agg.contains_key(&authorities[1]));
832        assert!(agg.contains_key(&authorities[2]));
833    }
834
835    #[test]
836    fn test_stake_aggregator_from_iter_with_error() {
837        let (committee, _) = Committee::new_simple_test_committee();
838        let committee = Arc::new(committee);
839        let authorities: Vec<_> = committee.names().copied().collect();
840
841        let data: Vec<Result<(AuthorityName, ()), TypedStoreError>> = vec![
842            Ok((authorities[0], ())),
843            Err(TypedStoreError::RocksDBError("test error".to_string())),
844        ];
845
846        let result: SuiResult<StakeAggregator<(), true>> =
847            StakeAggregator::from_iter(committee.clone(), data.into_iter());
848
849        assert!(result.is_err());
850    }
851}
852
853#[cfg(test)]
854mod generic_multi_stake_aggregator_tests {
855    use super::*;
856
857    #[test]
858    fn test_has_quorum_for_key() {
859        let (committee, _) = Committee::new_simple_test_committee();
860        let committee = Arc::new(committee);
861        let authorities: Vec<_> = committee.names().copied().collect();
862
863        let mut agg: GenericMultiStakeAggregator<&str, true> =
864            GenericMultiStakeAggregator::new(committee.clone());
865
866        let key1 = "key1";
867        let key2 = "key2";
868
869        // No quorum initially
870        assert!(!agg.has_quorum_for_key(&key1));
871        assert!(!agg.has_quorum_for_key(&key2));
872
873        // Add votes for key1 until quorum
874        agg.insert(authorities[0], key1);
875        assert!(!agg.has_quorum_for_key(&key1));
876
877        agg.insert(authorities[1], key1);
878        assert!(!agg.has_quorum_for_key(&key1));
879
880        agg.insert(authorities[2], key1);
881        assert!(agg.has_quorum_for_key(&key1));
882        assert!(!agg.has_quorum_for_key(&key2));
883
884        // Add vote for key2, but not enough for quorum
885        agg.insert(authorities[3], key2);
886        assert!(agg.has_quorum_for_key(&key1));
887        assert!(!agg.has_quorum_for_key(&key2));
888    }
889}
sui_core/stake_aggregator.rs

sui_core/
stake_aggregator.rs
