iota_swarm_config/
test_utils.rs

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