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

sui_core/
stake_aggregator.rs
