iota_types/
mock_checkpoint_builder.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::mem;
6
7use fastcrypto::traits::Signer;
8use iota_sdk_types::{
9    TransactionEffects,
10    checkpoint::{CheckpointContents, CheckpointSummary, EndOfEpochData},
11    gas::GasCostSummary,
12};
13
14use crate::{
15    base_types::{AuthorityName, VerifiedExecutionData},
16    committee::Committee,
17    crypto::{AuthoritySignInfo, AuthoritySignature, IotaAuthoritySignature},
18    effects::TransactionEffectsAPI,
19    messages_checkpoint::{
20        CertifiedCheckpointSummary, CheckpointContentsExt, CheckpointVersionSpecificData,
21        FullCheckpointContents, VerifiedCheckpoint, VerifiedCheckpointContents,
22    },
23    transaction::VerifiedTransaction,
24};
25
26pub trait ValidatorKeypairProvider {
27    fn get_validator_key(&self, name: &AuthorityName) -> &dyn Signer<AuthoritySignature>;
28    fn get_committee(&self) -> &Committee;
29}
30
31/// A utility to build consecutive checkpoints by adding transactions to the
32/// checkpoint builder. It's mostly used by simulations, tests and benchmarks.
33#[derive(Debug)]
34pub struct MockCheckpointBuilder {
35    previous_checkpoint: VerifiedCheckpoint,
36    transactions: Vec<VerifiedExecutionData>,
37    epoch_rolling_gas_cost_summary: GasCostSummary,
38    epoch: u64,
39}
40
41impl MockCheckpointBuilder {
42    pub fn new(previous_checkpoint: VerifiedCheckpoint) -> Self {
43        let epoch_rolling_gas_cost_summary =
44            previous_checkpoint.epoch_rolling_gas_cost_summary.clone();
45        let epoch = previous_checkpoint.epoch;
46
47        Self {
48            previous_checkpoint,
49            transactions: Vec::new(),
50            epoch_rolling_gas_cost_summary,
51            epoch,
52        }
53    }
54
55    pub fn size(&self) -> usize {
56        self.transactions.len()
57    }
58
59    pub fn epoch_rolling_gas_cost_summary(&self) -> &GasCostSummary {
60        &self.epoch_rolling_gas_cost_summary
61    }
62
63    pub fn push_transaction(
64        &mut self,
65        transaction: VerifiedTransaction,
66        effects: TransactionEffects,
67    ) {
68        self.epoch_rolling_gas_cost_summary += effects.gas_cost_summary();
69
70        self.transactions
71            .push(VerifiedExecutionData::new(transaction, effects))
72    }
73
74    /// Overrides the next checkpoint number indirectly by setting the previous
75    /// checkpoint's number to checkpoint_number - 1. This ensures the next
76    /// generated checkpoint has the exact sequence number provided. This
77    /// can be useful to generate checkpoints with specific sequence
78    /// numbers. Monotonicity of checkpoint numbers is enforced strictly.
79    pub fn override_next_checkpoint_number(
80        &mut self,
81        checkpoint_number: u64,
82        validator_keys: &impl ValidatorKeypairProvider,
83    ) {
84        assert!(
85            checkpoint_number > self.previous_checkpoint.sequence_number,
86            "Checkpoint number must strictly increase."
87        );
88
89        let mut summary = self.previous_checkpoint.data().clone();
90        summary.sequence_number = checkpoint_number - 1;
91        let checkpoint = Self::create_certified_checkpoint(validator_keys, summary);
92        self.previous_checkpoint = checkpoint;
93    }
94
95    /// Builds a checkpoint using internally buffered transactions.
96    pub fn build(
97        &mut self,
98        validator_keys: &impl ValidatorKeypairProvider,
99        timestamp_ms: u64,
100    ) -> (
101        VerifiedCheckpoint,
102        CheckpointContents,
103        VerifiedCheckpointContents,
104    ) {
105        self.build_internal(validator_keys, timestamp_ms, None)
106    }
107
108    pub fn build_end_of_epoch(
109        &mut self,
110        validator_keys: &impl ValidatorKeypairProvider,
111        timestamp_ms: u64,
112        new_epoch: u64,
113        end_of_epoch_data: EndOfEpochData,
114    ) -> (
115        VerifiedCheckpoint,
116        CheckpointContents,
117        VerifiedCheckpointContents,
118    ) {
119        self.build_internal(
120            validator_keys,
121            timestamp_ms,
122            Some((new_epoch, end_of_epoch_data)),
123        )
124    }
125
126    fn build_internal(
127        &mut self,
128        validator_keys: &impl ValidatorKeypairProvider,
129        timestamp_ms: u64,
130        new_epoch_data: Option<(u64, EndOfEpochData)>,
131    ) -> (
132        VerifiedCheckpoint,
133        CheckpointContents,
134        VerifiedCheckpointContents,
135    ) {
136        let contents =
137            CheckpointContents::new_with_causally_ordered_execution_data(self.transactions.iter());
138        let full_contents = VerifiedCheckpointContents::new_unchecked(
139            FullCheckpointContents::new_with_causally_ordered_transactions(
140                mem::take(&mut self.transactions)
141                    .into_iter()
142                    .map(|e| e.into_inner()),
143            ),
144        );
145
146        let (epoch, epoch_rolling_gas_cost_summary, end_of_epoch_data) =
147            if let Some((next_epoch, end_of_epoch_data)) = new_epoch_data {
148                let epoch = std::mem::replace(&mut self.epoch, next_epoch);
149                assert_eq!(next_epoch, epoch + 1);
150                let epoch_rolling_gas_cost_summary =
151                    std::mem::take(&mut self.epoch_rolling_gas_cost_summary);
152
153                (
154                    epoch,
155                    epoch_rolling_gas_cost_summary,
156                    Some(end_of_epoch_data),
157                )
158            } else {
159                (
160                    self.epoch,
161                    self.epoch_rolling_gas_cost_summary.clone(),
162                    None,
163                )
164            };
165
166        let summary = CheckpointSummary {
167            epoch,
168            sequence_number: self
169                .previous_checkpoint
170                .sequence_number
171                .checked_add(1)
172                .expect("checkpoint sequence number overflow"),
173            network_total_transactions: self.previous_checkpoint.network_total_transactions
174                + contents.len() as u64,
175            contents_digest: contents.digest(),
176            previous_digest: Some(*self.previous_checkpoint.digest()),
177            epoch_rolling_gas_cost_summary,
178            end_of_epoch_data,
179            timestamp_ms,
180            version_specific_data: bcs::to_bytes(&CheckpointVersionSpecificData::empty_for_tests())
181                .unwrap(),
182            checkpoint_commitments: Default::default(),
183        };
184
185        let checkpoint = Self::create_certified_checkpoint(validator_keys, summary);
186        self.previous_checkpoint = checkpoint.clone();
187        (checkpoint, contents, full_contents)
188    }
189
190    fn create_certified_checkpoint(
191        validator_keys: &impl ValidatorKeypairProvider,
192        checkpoint: CheckpointSummary,
193    ) -> VerifiedCheckpoint {
194        let signatures = validator_keys
195            .get_committee()
196            .voting_rights
197            .iter()
198            .map(|(name, _)| {
199                let intent_msg = iota_sdk_types::crypto::IntentMessage::new(
200                    iota_sdk_types::crypto::Intent::iota_app(
201                        iota_sdk_types::crypto::IntentScope::CheckpointSummary,
202                    ),
203                    &checkpoint,
204                );
205                let key = validator_keys.get_validator_key(name);
206                let signature = AuthoritySignature::new_secure(&intent_msg, &checkpoint.epoch, key);
207                AuthoritySignInfo {
208                    epoch: checkpoint.epoch,
209                    authority: *name,
210                    signature,
211                }
212            })
213            .collect();
214
215        let checkpoint_cert =
216            CertifiedCheckpointSummary::new(checkpoint, signatures, validator_keys.get_committee())
217                .unwrap();
218        VerifiedCheckpoint::new_unchecked(checkpoint_cert)
219    }
220}
iota_types/mock_checkpoint_builder.rs

iota_types/
mock_checkpoint_builder.rs
