sui_swarm_config/
test_utils.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::network_config::NetworkConfig;
5use shared_crypto::intent::{Intent, IntentMessage, IntentScope};
6use std::collections::HashMap;
7use sui_types::{
8    base_types::AuthorityName,
9    committee::{Committee, EpochId, StakeUnit},
10    crypto::{
11        AuthorityKeyPair, AuthoritySignInfo, AuthoritySignature, KeypairTraits,
12        SuiAuthoritySignature,
13    },
14    messages_checkpoint::{
15        CertifiedCheckpointSummary, CheckpointDigest, CheckpointSequenceNumber, CheckpointSummary,
16        CheckpointVersionSpecificData, EndOfEpochData, FullCheckpointContents, VerifiedCheckpoint,
17        VerifiedCheckpointContents,
18    },
19};
20
21pub struct CommitteeFixture {
22    epoch: EpochId,
23    validators: HashMap<AuthorityName, (AuthorityKeyPair, StakeUnit)>,
24    committee: Committee,
25}
26
27type MakeCheckpointResults = (
28    Vec<VerifiedCheckpoint>,
29    Vec<VerifiedCheckpointContents>,
30    HashMap<CheckpointSequenceNumber, CheckpointDigest>,
31    HashMap<CheckpointDigest, VerifiedCheckpoint>,
32);
33
34impl CommitteeFixture {
35    pub fn generate<R: ::rand::RngCore + ::rand::CryptoRng>(
36        mut rng: R,
37        epoch: EpochId,
38        committee_size: usize,
39    ) -> Self {
40        let validators = (0..committee_size)
41            .map(|_| sui_types::crypto::get_key_pair_from_rng::<AuthorityKeyPair, _>(&mut rng).1)
42            .map(|keypair| (keypair.public().into(), (keypair, 1)))
43            .collect::<HashMap<_, _>>();
44
45        let committee = Committee::new_for_testing_with_normalized_voting_power(
46            epoch,
47            validators
48                .iter()
49                .map(|(name, (_, stake))| (*name, *stake))
50                .collect(),
51        );
52
53        Self {
54            epoch,
55            validators,
56            committee,
57        }
58    }
59
60    pub fn from_network_config(network_config: &NetworkConfig) -> Self {
61        let committee = network_config.genesis.committee().unwrap();
62        Self {
63            epoch: committee.epoch,
64            validators: committee
65                .members()
66                .map(|(name, stake)| {
67                    (
68                        *name,
69                        (
70                            network_config
71                                .validator_configs()
72                                .iter()
73                                .find(|config| config.protocol_public_key() == *name)
74                                .unwrap()
75                                .protocol_key_pair()
76                                .copy(),
77                            *stake,
78                        ),
79                    )
80                })
81                .collect(),
82            committee,
83        }
84    }
85
86    pub fn committee(&self) -> &Committee {
87        &self.committee
88    }
89
90    fn create_root_checkpoint(&self) -> (VerifiedCheckpoint, VerifiedCheckpointContents) {
91        assert_eq!(self.epoch, 0, "root checkpoint must be epoch 0");
92        let checkpoint = CheckpointSummary {
93            epoch: 0,
94            sequence_number: 0,
95            network_total_transactions: 0,
96            content_digest: *empty_contents()
97                .into_inner()
98                .into_checkpoint_contents()
99                .digest(),
100            previous_digest: None,
101            epoch_rolling_gas_cost_summary: Default::default(),
102            end_of_epoch_data: None,
103            timestamp_ms: 0,
104            version_specific_data: bcs::to_bytes(&CheckpointVersionSpecificData::empty_for_tests())
105                .unwrap(),
106            checkpoint_commitments: Default::default(),
107        };
108
109        (
110            self.create_certified_checkpoint(checkpoint),
111            empty_contents(),
112        )
113    }
114
115    fn create_certified_checkpoint(&self, checkpoint: CheckpointSummary) -> VerifiedCheckpoint {
116        let signatures = self
117            .validators
118            .iter()
119            .map(|(name, (key, _))| {
120                let intent_msg = IntentMessage::new(
121                    Intent::sui_app(IntentScope::CheckpointSummary),
122                    checkpoint.clone(),
123                );
124                let signature = AuthoritySignature::new_secure(&intent_msg, &checkpoint.epoch, key);
125                AuthoritySignInfo {
126                    epoch: checkpoint.epoch,
127                    authority: *name,
128                    signature,
129                }
130            })
131            .collect();
132
133        CertifiedCheckpointSummary::new(checkpoint, signatures, self.committee())
134            .unwrap()
135            .try_into_verified(self.committee())
136            .unwrap()
137    }
138
139    pub fn make_random_checkpoints(
140        &self,
141        number_of_checkpoints: usize,
142        previous_checkpoint: Option<VerifiedCheckpoint>,
143    ) -> MakeCheckpointResults {
144        self.make_checkpoints(number_of_checkpoints, previous_checkpoint, random_contents)
145    }
146
147    pub fn make_empty_checkpoints(
148        &self,
149        number_of_checkpoints: usize,
150        previous_checkpoint: Option<VerifiedCheckpoint>,
151    ) -> MakeCheckpointResults {
152        self.make_checkpoints(number_of_checkpoints, previous_checkpoint, empty_contents)
153    }
154
155    fn make_checkpoints<F: Fn() -> VerifiedCheckpointContents>(
156        &self,
157        number_of_checkpoints: usize,
158        previous_checkpoint: Option<VerifiedCheckpoint>,
159        content_generator: F,
160    ) -> MakeCheckpointResults {
161        // Only skip the first one if it was supplied
162        let skip = previous_checkpoint.is_some() as usize;
163        let first = previous_checkpoint
164            .map(|c| (c, empty_contents()))
165            .unwrap_or_else(|| self.create_root_checkpoint());
166
167        let (ordered_checkpoints, contents): (Vec<_>, Vec<_>) =
168            std::iter::successors(Some(first), |prev| {
169                let contents = content_generator();
170                let contents_digest = *contents
171                    .clone()
172                    .into_inner()
173                    .into_checkpoint_contents()
174                    .digest();
175                let summary = CheckpointSummary {
176                    epoch: self.epoch,
177                    sequence_number: prev.0.sequence_number + 1,
178                    network_total_transactions: prev.0.network_total_transactions
179                        + contents.num_of_transactions() as u64,
180                    content_digest: contents_digest,
181                    previous_digest: Some(*prev.0.digest()),
182                    epoch_rolling_gas_cost_summary: Default::default(),
183                    end_of_epoch_data: None,
184                    timestamp_ms: 0,
185                    version_specific_data: bcs::to_bytes(
186                        &CheckpointVersionSpecificData::empty_for_tests(),
187                    )
188                    .unwrap(),
189                    checkpoint_commitments: Default::default(),
190                };
191
192                let checkpoint = self.create_certified_checkpoint(summary);
193
194                Some((checkpoint, contents))
195            })
196            .skip(skip)
197            .take(number_of_checkpoints)
198            .unzip();
199
200        let (sequence_number_to_digest, checkpoints) = ordered_checkpoints
201            .iter()
202            .cloned()
203            .map(|checkpoint| {
204                let digest = *checkpoint.digest();
205                ((checkpoint.sequence_number, digest), (digest, checkpoint))
206            })
207            .unzip();
208
209        (
210            ordered_checkpoints,
211            contents,
212            sequence_number_to_digest,
213            checkpoints,
214        )
215    }
216
217    pub fn make_end_of_epoch_checkpoint(
218        &self,
219        previous_checkpoint: VerifiedCheckpoint,
220        end_of_epoch_data: Option<EndOfEpochData>,
221    ) -> (
222        CheckpointSequenceNumber,
223        CheckpointDigest,
224        VerifiedCheckpoint,
225    ) {
226        let summary = CheckpointSummary {
227            epoch: self.epoch,
228            sequence_number: previous_checkpoint.sequence_number + 1,
229            network_total_transactions: 0,
230            content_digest: *empty_contents()
231                .into_inner()
232                .into_checkpoint_contents()
233                .digest(),
234            previous_digest: Some(*previous_checkpoint.digest()),
235            epoch_rolling_gas_cost_summary: Default::default(),
236            end_of_epoch_data,
237            timestamp_ms: 0,
238            version_specific_data: bcs::to_bytes(&CheckpointVersionSpecificData::empty_for_tests())
239                .unwrap(),
240            checkpoint_commitments: Default::default(),
241        };
242
243        let checkpoint = self.create_certified_checkpoint(summary);
244
245        (checkpoint.sequence_number, *checkpoint.digest(), checkpoint)
246    }
247}
248
249pub fn empty_contents() -> VerifiedCheckpointContents {
250    VerifiedCheckpointContents::new_unchecked(
251        FullCheckpointContents::new_with_causally_ordered_transactions(std::iter::empty()),
252    )
253}
254
255pub fn random_contents() -> VerifiedCheckpointContents {
256    VerifiedCheckpointContents::new_unchecked(FullCheckpointContents::random_for_testing())
257}