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                                InsertResult::NotEnoughVotes {
184                                    bad_votes,
185                                    bad_authorities,
186                                }
187                            }
188                        }
189                    }
190                    Err(error) => InsertResult::Failed { error },
191                }
192            }
193            // The following is necessary to change the template type of InsertResult.
194            InsertResult::Failed { error } => InsertResult::Failed { error },
195            InsertResult::NotEnoughVotes {
196                bad_votes,
197                bad_authorities,
198            } => InsertResult::NotEnoughVotes {
199                bad_votes,
200                bad_authorities,
201            },
202        }
203    }
204}
205
206pub enum InsertResult<CertT> {
207    QuorumReached(CertT),
208    Failed {
209        error: SuiError,
210    },
211    NotEnoughVotes {
212        bad_votes: u64,
213        bad_authorities: Vec<AuthorityName>,
214    },
215}
216
217impl<CertT> InsertResult<CertT> {
218    pub fn is_quorum_reached(&self) -> bool {
219        matches!(self, Self::QuorumReached(..))
220    }
221}
222
223/// MultiStakeAggregator is a utility data structure that tracks the stake accumulation of
224/// potentially multiple different values (usually due to byzantine/corrupted responses). Each
225/// value is tracked using a StakeAggregator and determine whether it has reached a quorum.
226/// Once quorum is reached, the aggregated signature is returned.
227#[derive(Debug)]
228pub struct MultiStakeAggregator<K, V, const STRENGTH: bool> {
229    committee: Arc<Committee>,
230    stake_maps: HashMap<K, (V, StakeAggregator<AuthoritySignInfo, STRENGTH>)>,
231}
232
233impl<K, V, const STRENGTH: bool> MultiStakeAggregator<K, V, STRENGTH> {
234    pub fn new(committee: Arc<Committee>) -> Self {
235        Self {
236            committee,
237            stake_maps: Default::default(),
238        }
239    }
240
241    pub fn total_votes(&self) -> StakeUnit {
242        let mut voted_authorities = HashSet::new();
243        self.stake_maps.values().for_each(|(_, stake_aggregator)| {
244            stake_aggregator.keys().for_each(|k| {
245                voted_authorities.insert(k);
246            })
247        });
248        voted_authorities
249            .iter()
250            .map(|k| self.committee.weight(k))
251            .sum()
252    }
253
254    #[cfg(test)]
255    pub fn unique_key_count(&self) -> usize {
256        self.stake_maps.len()
257    }
258}
259
260impl<K, V, const STRENGTH: bool> MultiStakeAggregator<K, V, STRENGTH>
261where
262    K: Hash + Eq,
263    V: Message + Serialize + Clone,
264{
265    pub fn insert(
266        &mut self,
267        k: K,
268        envelope: Envelope<V, AuthoritySignInfo>,
269    ) -> InsertResult<AuthorityQuorumSignInfo<STRENGTH>> {
270        if let Some(entry) = self.stake_maps.get_mut(&k) {
271            entry.1.insert(envelope)
272        } else {
273            let mut new_entry = StakeAggregator::new(self.committee.clone());
274            let result = new_entry.insert(envelope.clone());
275            if !matches!(result, InsertResult::Failed { .. }) {
276                // This is very important: ensure that if the insert fails, we don't even add the
277                // new entry to the map.
278                self.stake_maps.insert(k, (envelope.into_data(), new_entry));
279            }
280            result
281        }
282    }
283}
284
285impl<K, V, const STRENGTH: bool> MultiStakeAggregator<K, V, STRENGTH>
286where
287    K: Clone + Ord,
288{
289    pub fn get_all_unique_values(&self) -> BTreeMap<K, (Vec<AuthorityName>, StakeUnit)> {
290        self.stake_maps
291            .iter()
292            .map(|(k, (_, s))| (k.clone(), (s.data.keys().copied().collect(), s.total_votes)))
293            .collect()
294    }
295}
296
297impl<K, V, const STRENGTH: bool> MultiStakeAggregator<K, V, STRENGTH>
298where
299    K: Hash + Eq,
300{
301    #[allow(dead_code)]
302    pub fn authorities_for_key(&self, k: &K) -> Option<impl Iterator<Item = &AuthorityName>> {
303        self.stake_maps.get(k).map(|(_, agg)| agg.keys())
304    }
305
306    /// The sum of all remaining stake, i.e. all stake not yet
307    /// committed by vote to a specific value
308    pub fn uncommitted_stake(&self) -> StakeUnit {
309        self.committee.total_votes() - self.total_votes()
310    }
311
312    /// Total stake of the largest faction
313    pub fn plurality_stake(&self) -> StakeUnit {
314        self.stake_maps
315            .values()
316            .map(|(_, agg)| agg.total_votes())
317            .max()
318            .unwrap_or_default()
319    }
320
321    /// If true, there isn't enough uncommitted stake to reach quorum for any value
322    pub fn quorum_unreachable(&self) -> bool {
323        self.uncommitted_stake() + self.plurality_stake() < self.committee.threshold::<STRENGTH>()
324    }
325}
326
327/// Like MultiStakeAggregator, but for counting votes for a generic value instead of an envelope, in
328/// scenarios where byzantine validators may submit multiple votes for different values.
329pub struct GenericMultiStakeAggregator<K, const STRENGTH: bool> {
330    committee: Arc<Committee>,
331    stake_maps: HashMap<K, StakeAggregator<(), STRENGTH>>,
332    votes_per_authority: HashMap<AuthorityName, u64>,
333}
334
335impl<K, const STRENGTH: bool> GenericMultiStakeAggregator<K, STRENGTH>
336where
337    K: Hash + Eq,
338{
339    pub fn new(committee: Arc<Committee>) -> Self {
340        Self {
341            committee,
342            stake_maps: Default::default(),
343            votes_per_authority: Default::default(),
344        }
345    }
346
347    pub fn insert(
348        &mut self,
349        authority: AuthorityName,
350        k: K,
351    ) -> InsertResult<&HashMap<AuthorityName, ()>> {
352        let agg = self
353            .stake_maps
354            .entry(k)
355            .or_insert_with(|| StakeAggregator::new(self.committee.clone()));
356
357        if !agg.contains_key(&authority) {
358            *self.votes_per_authority.entry(authority).or_default() += 1;
359        }
360
361        agg.insert_generic(authority, ())
362    }
363
364    pub fn has_quorum_for_key(&self, k: &K) -> bool {
365        if let Some(entry) = self.stake_maps.get(k) {
366            entry.has_quorum()
367        } else {
368            false
369        }
370    }
371
372    pub fn votes_for_authority(&self, authority: AuthorityName) -> u64 {
373        self.votes_per_authority
374            .get(&authority)
375            .copied()
376            .unwrap_or_default()
377    }
378}
379
380#[test]
381fn test_votes_per_authority() {
382    let (committee, _) = Committee::new_simple_test_committee();
383    let authorities: Vec<_> = committee.names().copied().collect();
384
385    let mut agg: GenericMultiStakeAggregator<&str, true> =
386        GenericMultiStakeAggregator::new(Arc::new(committee));
387
388    // 1. Inserting an `authority` and a `key`, and then checking the number of votes for that `authority`.
389    let key1: &str = "key1";
390    let authority1 = authorities[0];
391    agg.insert(authority1, key1);
392    assert_eq!(agg.votes_for_authority(authority1), 1);
393
394    // 2. Inserting the same `authority` and `key` pair multiple times to ensure votes aren't incremented incorrectly.
395    agg.insert(authority1, key1);
396    agg.insert(authority1, key1);
397    assert_eq!(agg.votes_for_authority(authority1), 1);
398
399    // 3. Checking votes for an authority that hasn't voted.
400    let authority2 = authorities[1];
401    assert_eq!(agg.votes_for_authority(authority2), 0);
402
403    // 4. Inserting multiple different authorities and checking their vote counts.
404    let key2: &str = "key2";
405    agg.insert(authority2, key2);
406    assert_eq!(agg.votes_for_authority(authority2), 1);
407    assert_eq!(agg.votes_for_authority(authority1), 1);
408
409    // 5. Verifying that inserting different keys for the same authority increments the vote count.
410    let key3: &str = "key3";
411    agg.insert(authority1, key3);
412    assert_eq!(agg.votes_for_authority(authority1), 2);
413}
414
415#[cfg(test)]
416mod multi_stake_aggregator_tests {
417    use super::*;
418    use fastcrypto::hash::{HashFunction, Sha3_256};
419    use shared_crypto::intent::IntentScope;
420
421    #[derive(Clone, Debug, Serialize, PartialEq, Eq, Hash)]
422    struct TestMessage {
423        value: String,
424    }
425
426    impl Message for TestMessage {
427        type DigestType = [u8; 32];
428        const SCOPE: IntentScope = IntentScope::SenderSignedTransaction;
429
430        fn digest(&self) -> Self::DigestType {
431            let mut hasher = Sha3_256::default();
432            hasher.update(self.value.as_bytes());
433            hasher.finalize().digest
434        }
435    }
436
437    #[test]
438    fn test_equivocation_stake_not_double_counted() {
439        let (committee, key_pairs) = Committee::new_simple_test_committee();
440        let committee = Arc::new(committee);
441        let authorities: Vec<_> = committee.names().copied().collect();
442
443        let mut agg: MultiStakeAggregator<String, TestMessage, true> =
444            MultiStakeAggregator::new(committee.clone());
445
446        // Get the actual total stake from the committee
447        let total_stake = committee.total_votes();
448        let num_authorities = authorities.len();
449        let stake_per_authority = total_stake / num_authorities as u64;
450
451        // Simulate equivocation: authority0 signs multiple different values
452        let authority0 = authorities[0];
453        let key0 = &key_pairs[0];
454
455        // First signature for "value1"
456        let msg1 = TestMessage {
457            value: "value1".to_string(),
458        };
459        let envelope1 =
460            <Envelope<TestMessage, AuthoritySignInfo>>::new(0, msg1.clone(), key0, authority0);
461        agg.insert("key1".to_string(), envelope1);
462
463        // Second signature from same authority for "value2" (equivocation)
464        let msg2 = TestMessage {
465            value: "value2".to_string(),
466        };
467        let envelope2 =
468            <Envelope<TestMessage, AuthoritySignInfo>>::new(0, msg2.clone(), key0, authority0);
469        agg.insert("key2".to_string(), envelope2);
470
471        // Third signature from same authority for "value3" (more equivocation)
472        let msg3 = TestMessage {
473            value: "value3".to_string(),
474        };
475        let envelope3 =
476            <Envelope<TestMessage, AuthoritySignInfo>>::new(0, msg3.clone(), key0, authority0);
477        agg.insert("key3".to_string(), envelope3);
478
479        // With the fix: authority0's stake should only be counted once, even though they signed 3 different values
480        let aggregated_votes = agg.total_votes();
481        assert_eq!(aggregated_votes, stake_per_authority);
482
483        // Add more authorities signing different values
484        let authority1 = authorities[1];
485        let key1 = &key_pairs[1];
486        let envelope4 =
487            <Envelope<TestMessage, AuthoritySignInfo>>::new(0, msg1.clone(), key1, authority1);
488        agg.insert("key1".to_string(), envelope4);
489
490        let authority2 = authorities[2];
491        let key2 = &key_pairs[2];
492        let envelope5 = <Envelope<TestMessage, AuthoritySignInfo>>::new(0, msg2, key2, authority2);
493        agg.insert("key2".to_string(), envelope5);
494
495        // Now total_votes() should be stake_per_authority * 3 (3 unique authorities)
496        // NOT stake_per_authority * 5 (which would be if we double-counted authority0)
497        let aggregated_votes = agg.total_votes();
498        assert_eq!(aggregated_votes, stake_per_authority * 3);
499        assert!(aggregated_votes <= total_stake);
500
501        // uncommitted_stake should work without underflow
502        let uncommitted = agg.uncommitted_stake();
503        assert_eq!(uncommitted, stake_per_authority); // Only authority3 hasn't voted
504
505        // Verify we have 3 different keys with votes
506        assert_eq!(agg.unique_key_count(), 3);
507    }
508
509    #[test]
510    fn test_multistake_uncommitted_and_plurality() {
511        let (committee, key_pairs) = Committee::new_simple_test_committee();
512        let committee = Arc::new(committee);
513        let authorities: Vec<_> = committee.names().copied().collect();
514
515        let mut agg: MultiStakeAggregator<String, TestMessage, true> =
516            MultiStakeAggregator::new(committee.clone());
517
518        let total_stake = committee.total_votes();
519        let num_authorities = authorities.len();
520        let stake_per_authority = total_stake / num_authorities as u64;
521
522        // Initially, all stake is uncommitted
523        assert_eq!(agg.uncommitted_stake(), total_stake);
524        assert_eq!(agg.plurality_stake(), 0);
525        assert!(!agg.quorum_unreachable());
526
527        // Add first authority voting for value1
528        let msg1 = TestMessage {
529            value: "value1".to_string(),
530        };
531        let envelope1 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
532            0,
533            msg1.clone(),
534            &key_pairs[0],
535            authorities[0],
536        );
537        agg.insert("key1".to_string(), envelope1);
538
539        assert_eq!(agg.uncommitted_stake(), total_stake - stake_per_authority);
540        assert_eq!(agg.plurality_stake(), stake_per_authority);
541
542        // Add second authority voting for value2
543        let msg2 = TestMessage {
544            value: "value2".to_string(),
545        };
546        let envelope2 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
547            0,
548            msg2.clone(),
549            &key_pairs[1],
550            authorities[1],
551        );
552        agg.insert("key2".to_string(), envelope2);
553
554        assert_eq!(
555            agg.uncommitted_stake(),
556            total_stake - 2 * stake_per_authority
557        );
558        assert_eq!(agg.plurality_stake(), stake_per_authority);
559
560        // Add third authority voting for value1 (now value1 has plurality)
561        let envelope3 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
562            0,
563            msg1.clone(),
564            &key_pairs[2],
565            authorities[2],
566        );
567        agg.insert("key1".to_string(), envelope3);
568
569        assert_eq!(
570            agg.uncommitted_stake(),
571            total_stake - 3 * stake_per_authority
572        );
573        assert_eq!(agg.plurality_stake(), 2 * stake_per_authority);
574    }
575
576    #[test]
577    fn test_multistake_quorum_unreachable() {
578        let (committee, key_pairs) = Committee::new_simple_test_committee();
579        let committee = Arc::new(committee);
580        let authorities: Vec<_> = committee.names().copied().collect();
581
582        let mut agg: MultiStakeAggregator<String, TestMessage, true> =
583            MultiStakeAggregator::new(committee.clone());
584
585        // Split votes evenly so no value can reach quorum
586        // With 4 authorities and strong quorum needing 2f+1, we need at least 3
587        let msg1 = TestMessage {
588            value: "value1".to_string(),
589        };
590        let msg2 = TestMessage {
591            value: "value2".to_string(),
592        };
593
594        // Two authorities vote for value1
595        let envelope1 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
596            0,
597            msg1.clone(),
598            &key_pairs[0],
599            authorities[0],
600        );
601        agg.insert("key1".to_string(), envelope1);
602
603        let envelope2 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
604            0,
605            msg1.clone(),
606            &key_pairs[1],
607            authorities[1],
608        );
609        agg.insert("key1".to_string(), envelope2);
610
611        // Two authorities vote for value2
612        let envelope3 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
613            0,
614            msg2.clone(),
615            &key_pairs[2],
616            authorities[2],
617        );
618        agg.insert("key2".to_string(), envelope3);
619
620        let envelope4 = <Envelope<TestMessage, AuthoritySignInfo>>::new(
621            0,
622            msg2.clone(),
623            &key_pairs[3],
624            authorities[3],
625        );
626        agg.insert("key2".to_string(), envelope4);
627
628        // With evenly split votes, neither can reach quorum now
629        assert!(agg.quorum_unreachable());
630    }
631}
632
633#[cfg(test)]
634mod stake_aggregator_tests {
635    use super::*;
636
637    #[test]
638    fn test_stake_aggregator_strong_quorum() {
639        let (committee, _) = Committee::new_simple_test_committee();
640        let committee = Arc::new(committee);
641        let authorities: Vec<_> = committee.names().copied().collect();
642
643        let mut agg: StakeAggregator<(), true> = StakeAggregator::new(committee.clone());
644
645        let total_stake = committee.total_votes();
646        let num_authorities = authorities.len();
647        let stake_per_authority = total_stake / num_authorities as u64;
648
649        assert_eq!(agg.total_votes(), 0);
650        assert!(!agg.has_quorum());
651        assert_eq!(agg.validator_sig_count(), 0);
652
653        // Add first authority - should not reach quorum yet
654        let result = agg.insert_generic(authorities[0], ());
655        assert!(matches!(result, InsertResult::NotEnoughVotes { .. }));
656        assert_eq!(agg.total_votes(), stake_per_authority);
657        assert!(!agg.has_quorum());
658        assert_eq!(agg.validator_sig_count(), 1);
659
660        // Add second authority - still not enough for strong quorum (2f+1)
661        let result = agg.insert_generic(authorities[1], ());
662        assert!(matches!(result, InsertResult::NotEnoughVotes { .. }));
663        assert_eq!(agg.total_votes(), 2 * stake_per_authority);
664        assert!(!agg.has_quorum());
665
666        // Add third authority - should reach strong quorum
667        let result = agg.insert_generic(authorities[2], ());
668        assert!(result.is_quorum_reached());
669        assert!(agg.has_quorum());
670        assert_eq!(agg.validator_sig_count(), 3);
671    }
672
673    #[test]
674    fn test_stake_aggregator_weak_quorum() {
675        let (committee, _) = Committee::new_simple_test_committee();
676        let committee = Arc::new(committee);
677        let authorities: Vec<_> = committee.names().copied().collect();
678
679        let mut agg: StakeAggregator<(), false> = StakeAggregator::new(committee.clone());
680
681        // Weak quorum (f+1) should be reached faster than strong quorum
682        let result = agg.insert_generic(authorities[0], ());
683        assert!(matches!(result, InsertResult::NotEnoughVotes { .. }));
684        assert!(!agg.has_quorum());
685
686        // Second authority should reach weak quorum
687        let result = agg.insert_generic(authorities[1], ());
688        assert!(result.is_quorum_reached());
689        assert!(agg.has_quorum());
690    }
691
692    #[test]
693    fn test_stake_aggregator_repeated_signer() {
694        let (committee, _) = Committee::new_simple_test_committee();
695        let committee = Arc::new(committee);
696        let authorities: Vec<_> = committee.names().copied().collect();
697
698        let mut agg: StakeAggregator<u32, true> = StakeAggregator::new(committee.clone());
699
700        // Insert first time - should succeed
701        let result = agg.insert_generic(authorities[0], 1);
702        assert!(matches!(result, InsertResult::NotEnoughVotes { .. }));
703
704        // Insert same authority again with same value - should fail
705        let result = agg.insert_generic(authorities[0], 1);
706        assert!(matches!(
707            result,
708            InsertResult::Failed {
709                error
710            } if matches!(error.as_inner(),  SuiErrorKind::StakeAggregatorRepeatedSigner { .. } )
711        ));
712
713        // Insert same authority with different value - should also fail (conflicting signature)
714        let result = agg.insert_generic(authorities[0], 2);
715        let InsertResult::Failed { error } = result else {
716            panic!("Expected StakeAggregatorRepeatedSigner error");
717        };
718        let SuiErrorKind::StakeAggregatorRepeatedSigner {
719            signer,
720            conflicting_sig,
721        } = error.into_inner()
722        else {
723            panic!("Expected StakeAggregatorRepeatedSigner error");
724        };
725        assert_eq!(signer, authorities[0]);
726        assert!(conflicting_sig);
727    }
728
729    #[test]
730    fn test_stake_aggregator_from_iter() {
731        let (committee, _) = Committee::new_simple_test_committee();
732        let committee = Arc::new(committee);
733        let authorities: Vec<_> = committee.names().copied().collect();
734
735        let data = vec![
736            Ok((authorities[0], ())),
737            Ok((authorities[1], ())),
738            Ok((authorities[2], ())),
739        ];
740
741        let agg: StakeAggregator<(), true> =
742            StakeAggregator::from_iter(committee.clone(), data.into_iter()).unwrap();
743
744        assert_eq!(agg.validator_sig_count(), 3);
745        assert!(agg.has_quorum());
746        assert!(agg.contains_key(&authorities[0]));
747        assert!(agg.contains_key(&authorities[1]));
748        assert!(agg.contains_key(&authorities[2]));
749    }
750
751    #[test]
752    fn test_stake_aggregator_from_iter_with_error() {
753        let (committee, _) = Committee::new_simple_test_committee();
754        let committee = Arc::new(committee);
755        let authorities: Vec<_> = committee.names().copied().collect();
756
757        let data: Vec<Result<(AuthorityName, ()), TypedStoreError>> = vec![
758            Ok((authorities[0], ())),
759            Err(TypedStoreError::RocksDBError("test error".to_string())),
760        ];
761
762        let result: SuiResult<StakeAggregator<(), true>> =
763            StakeAggregator::from_iter(committee.clone(), data.into_iter());
764
765        assert!(result.is_err());
766    }
767}
768
769#[cfg(test)]
770mod generic_multi_stake_aggregator_tests {
771    use super::*;
772
773    #[test]
774    fn test_has_quorum_for_key() {
775        let (committee, _) = Committee::new_simple_test_committee();
776        let committee = Arc::new(committee);
777        let authorities: Vec<_> = committee.names().copied().collect();
778
779        let mut agg: GenericMultiStakeAggregator<&str, true> =
780            GenericMultiStakeAggregator::new(committee.clone());
781
782        let key1 = "key1";
783        let key2 = "key2";
784
785        // No quorum initially
786        assert!(!agg.has_quorum_for_key(&key1));
787        assert!(!agg.has_quorum_for_key(&key2));
788
789        // Add votes for key1 until quorum
790        agg.insert(authorities[0], key1);
791        assert!(!agg.has_quorum_for_key(&key1));
792
793        agg.insert(authorities[1], key1);
794        assert!(!agg.has_quorum_for_key(&key1));
795
796        agg.insert(authorities[2], key1);
797        assert!(agg.has_quorum_for_key(&key1));
798        assert!(!agg.has_quorum_for_key(&key2));
799
800        // Add vote for key2, but not enough for quorum
801        agg.insert(authorities[3], key2);
802        assert!(agg.has_quorum_for_key(&key1));
803        assert!(!agg.has_quorum_for_key(&key2));
804    }
805}
sui_core/stake_aggregator.rs

sui_core/
stake_aggregator.rs
