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

iota_types/
mock_checkpoint_builder.rs
