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

iota_types/
mock_checkpoint_builder.rs
