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

iota_types/
mock_checkpoint_builder.rs
