Skip to main content

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, VersionedFullCheckpointContents,
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();
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_random_checkpoints_with_transactions(
148        &self,
149        number_of_checkpoints: usize,
150        previous_checkpoint: Option<VerifiedCheckpoint>,
151        transactions_per_checkpoint: usize,
152    ) -> MakeCheckpointResults {
153        self.make_checkpoints(number_of_checkpoints, previous_checkpoint, || {
154            VerifiedCheckpointContents::new_unchecked(VersionedFullCheckpointContents::V1(
155                FullCheckpointContents::random_with_transactions_for_testing(
156                    transactions_per_checkpoint,
157                ),
158            ))
159        })
160    }
161
162    pub fn make_empty_checkpoints(
163        &self,
164        number_of_checkpoints: usize,
165        previous_checkpoint: Option<VerifiedCheckpoint>,
166    ) -> MakeCheckpointResults {
167        self.make_checkpoints(number_of_checkpoints, previous_checkpoint, empty_contents)
168    }
169
170    fn make_checkpoints<F: Fn() -> VerifiedCheckpointContents>(
171        &self,
172        number_of_checkpoints: usize,
173        previous_checkpoint: Option<VerifiedCheckpoint>,
174        content_generator: F,
175    ) -> MakeCheckpointResults {
176        // Only skip the first one if it was supplied
177        let skip = previous_checkpoint.is_some() as usize;
178        let first = previous_checkpoint
179            .map(|c| (c, empty_contents()))
180            .unwrap_or_else(|| self.create_root_checkpoint());
181
182        let (ordered_checkpoints, contents): (Vec<_>, Vec<_>) =
183            std::iter::successors(Some(first), |prev| {
184                let contents = content_generator();
185                let contents_digest = *contents
186                    .clone()
187                    .into_inner()
188                    .into_checkpoint_contents()
189                    .digest();
190                let summary = CheckpointSummary {
191                    epoch: self.epoch,
192                    sequence_number: prev.0.sequence_number + 1,
193                    network_total_transactions: prev.0.network_total_transactions
194                        + contents.num_of_transactions() as u64,
195                    content_digest: contents_digest,
196                    previous_digest: Some(*prev.0.digest()),
197                    epoch_rolling_gas_cost_summary: Default::default(),
198                    end_of_epoch_data: None,
199                    timestamp_ms: 0,
200                    version_specific_data: bcs::to_bytes(
201                        &CheckpointVersionSpecificData::empty_for_tests(),
202                    )
203                    .unwrap(),
204                    checkpoint_commitments: Default::default(),
205                };
206
207                let checkpoint = self.create_certified_checkpoint(summary);
208
209                Some((checkpoint, contents))
210            })
211            .skip(skip)
212            .take(number_of_checkpoints)
213            .unzip();
214
215        let (sequence_number_to_digest, checkpoints) = ordered_checkpoints
216            .iter()
217            .cloned()
218            .map(|checkpoint| {
219                let digest = *checkpoint.digest();
220                ((checkpoint.sequence_number, digest), (digest, checkpoint))
221            })
222            .unzip();
223
224        (
225            ordered_checkpoints,
226            contents,
227            sequence_number_to_digest,
228            checkpoints,
229        )
230    }
231
232    pub fn make_end_of_epoch_checkpoint(
233        &self,
234        previous_checkpoint: VerifiedCheckpoint,
235        end_of_epoch_data: Option<EndOfEpochData>,
236    ) -> (
237        CheckpointSequenceNumber,
238        CheckpointDigest,
239        VerifiedCheckpoint,
240    ) {
241        let summary = CheckpointSummary {
242            epoch: self.epoch,
243            sequence_number: previous_checkpoint.sequence_number + 1,
244            network_total_transactions: 0,
245            content_digest: *empty_contents()
246                .into_inner()
247                .into_checkpoint_contents()
248                .digest(),
249            previous_digest: Some(*previous_checkpoint.digest()),
250            epoch_rolling_gas_cost_summary: Default::default(),
251            end_of_epoch_data,
252            timestamp_ms: 0,
253            version_specific_data: bcs::to_bytes(&CheckpointVersionSpecificData::empty_for_tests())
254                .unwrap(),
255            checkpoint_commitments: Default::default(),
256        };
257
258        let checkpoint = self.create_certified_checkpoint(summary);
259
260        (checkpoint.sequence_number, *checkpoint.digest(), checkpoint)
261    }
262}
263
264pub fn empty_contents() -> VerifiedCheckpointContents {
265    VerifiedCheckpointContents::new_unchecked(VersionedFullCheckpointContents::V1(
266        FullCheckpointContents::new_with_causally_ordered_transactions(std::iter::empty()),
267    ))
268}
269
270pub fn random_contents() -> VerifiedCheckpointContents {
271    VerifiedCheckpointContents::new_unchecked(VersionedFullCheckpointContents::V1(
272        FullCheckpointContents::random_for_testing(),
273    ))
274}