diff --git a/crates/iota-genesis-builder/examples/snapshot_test_outputs.rs b/crates/iota-genesis-builder/examples/snapshot_test_outputs.rs new file mode 100644 index 00000000000..2a0ef37afba --- /dev/null +++ b/crates/iota-genesis-builder/examples/snapshot_test_outputs.rs @@ -0,0 +1,44 @@ +// Copyright (c) 2024 IOTA Stiftung +// SPDX-License-Identifier: Apache-2.0 + +//! Example to add test outputs to a full snapshot. + +use std::{fs::File, path::Path}; + +use iota_genesis_builder::stardust::{ + parse::FullSnapshotParser, test_outputs::add_snapshot_test_outputs, +}; + +fn parse_snapshot>(path: P) -> anyhow::Result<()> { + let file = File::open(path)?; + let parser = FullSnapshotParser::new(file)?; + + println!("Output count: {}", parser.header.output_count()); + + let total_supply_header = parser.total_supply()?; + let total_supply_outputs = parser.outputs().try_fold(0, |acc, output| { + Ok::<_, anyhow::Error>(acc + output?.1.amount()) + })?; + + assert_eq!(total_supply_header, total_supply_outputs); + + println!("Total supply: {total_supply_header}"); + + Ok(()) +} + +fn main() -> anyhow::Result<()> { + let Some(current_path) = std::env::args().nth(1) else { + anyhow::bail!("please provide path to the full-snapshot file"); + }; + let mut new_path = String::from("test-"); + new_path.push_str(¤t_path); + + parse_snapshot(¤t_path)?; + + add_snapshot_test_outputs(¤t_path, &new_path)?; + + parse_snapshot(&new_path)?; + + Ok(()) +} diff --git a/crates/iota-genesis-builder/src/stardust/mod.rs b/crates/iota-genesis-builder/src/stardust/mod.rs index 617b442df70..bd31ffcc634 100644 --- a/crates/iota-genesis-builder/src/stardust/mod.rs +++ b/crates/iota-genesis-builder/src/stardust/mod.rs @@ -9,4 +9,5 @@ pub mod error; pub mod migration; pub mod native_token; pub mod parse; +pub mod test_outputs; pub mod types; diff --git a/crates/iota-genesis-builder/src/stardust/parse.rs b/crates/iota-genesis-builder/src/stardust/parse.rs index aaf481424f8..0e0d47386c3 100644 --- a/crates/iota-genesis-builder/src/stardust/parse.rs +++ b/crates/iota-genesis-builder/src/stardust/parse.rs @@ -48,17 +48,22 @@ impl FullSnapshotParser { self.header.target_milestone_timestamp() } - /// Provide the network main token total supply through the snapshot - /// protocol parameters. - pub fn total_supply(&self) -> Result { + /// Provide the protocol parameters extracted from the snapshot header. + pub fn protocol_parameters(&self) -> Result { if let MilestoneOption::Parameters(params) = self.header.parameters_milestone_option() { let protocol_params = ::unpack_unverified( params.binary_parameters(), ) .expect("invalid protocol params"); - Ok(protocol_params.token_supply()) + Ok(protocol_params) } else { Err(StardustError::HornetSnapshotParametersNotFound.into()) } } + + /// Provide the network main token total supply through the snapshot + /// protocol parameters. + pub fn total_supply(&self) -> Result { + self.protocol_parameters().map(|p| p.token_supply()) + } } diff --git a/crates/iota-genesis-builder/src/stardust/test_outputs/dummy.rs b/crates/iota-genesis-builder/src/stardust/test_outputs/dummy.rs new file mode 100644 index 00000000000..033dea3a28d --- /dev/null +++ b/crates/iota-genesis-builder/src/stardust/test_outputs/dummy.rs @@ -0,0 +1,38 @@ +// Copyright (c) 2024 IOTA Stiftung +// SPDX-License-Identifier: Apache-2.0 + +use std::str::FromStr; + +use iota_sdk::types::block::{ + address::Ed25519Address, + output::{unlock_condition::AddressUnlockCondition, BasicOutputBuilder, Output}, + payload::transaction::TransactionId, +}; + +use crate::stardust::types::snapshot::OutputHeader; + +pub(crate) fn outputs() -> Vec<(OutputHeader, Output)> { + let mut outputs = Vec::new(); + + let output_header = OutputHeader::new_testing( + *TransactionId::from_str( + "0xb191c4bc825ac6983789e50545d5ef07a1d293a98ad974fc9498cb1812345678", + ) + .unwrap(), + rand::random(), + rand::random(), + rand::random(), + ); + let output = Output::from( + BasicOutputBuilder::new_with_amount(1_000_000) + .add_unlock_condition(AddressUnlockCondition::new(Ed25519Address::from( + rand::random::<[u8; Ed25519Address::LENGTH]>(), + ))) + .finish() + .unwrap(), + ); + + outputs.push((output_header, output)); + + outputs +} diff --git a/crates/iota-genesis-builder/src/stardust/test_outputs/mod.rs b/crates/iota-genesis-builder/src/stardust/test_outputs/mod.rs new file mode 100644 index 00000000000..1349f96ae02 --- /dev/null +++ b/crates/iota-genesis-builder/src/stardust/test_outputs/mod.rs @@ -0,0 +1,73 @@ +// Copyright (c) 2024 IOTA Stiftung +// SPDX-License-Identifier: Apache-2.0 + +mod dummy; + +use std::{ + fs::{File, OpenOptions}, + io::BufWriter, + path::Path, +}; + +use iota_sdk::types::block::{ + payload::milestone::{MilestoneOption, ParametersMilestoneOption}, + protocol::ProtocolParameters, +}; +use packable::{packer::IoPacker, Packable, PackableExt}; + +use crate::stardust::parse::FullSnapshotParser; + +/// Adds outputs to test specific and intricate scenario in the full snapshot. +pub fn add_snapshot_test_outputs + core::fmt::Debug>( + current_path: P, + new_path: P, +) -> anyhow::Result<()> { + let current_file = File::open(current_path)?; + let new_file = OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .open(new_path)?; + let mut writer = IoPacker::new(BufWriter::new(new_file)); + let mut parser = FullSnapshotParser::new(current_file)?; + + let new_outputs = dummy::outputs(); + + // Increments the output count according to newly generated outputs. + parser.header.output_count += new_outputs.len() as u64; + + // Creates new protocol parameters to increase the total supply according to newly generated outputs. + let params = parser.protocol_parameters()?; + let new_params = ProtocolParameters::new( + params.protocol_version(), + params.network_name().to_owned(), + params.bech32_hrp(), + params.min_pow_score(), + params.below_max_depth(), + *params.rent_structure(), + params.token_supply() + new_outputs.iter().map(|o| o.1.amount()).sum::(), + )?; + if let MilestoneOption::Parameters(params) = &parser.header.parameters_milestone_option { + parser.header.parameters_milestone_option = + MilestoneOption::Parameters(ParametersMilestoneOption::new( + params.target_milestone_index(), + params.protocol_version(), + new_params.pack_to_vec(), + )?); + } + + // Writes the new header. + parser.header.pack(&mut writer)?; + + // Writes previous and new outputs. + parser + .outputs() + .filter_map(|o| o.ok()) + .chain(new_outputs) + .for_each(|(output_header, output)| { + output_header.pack(&mut writer).unwrap(); + output.pack(&mut writer).unwrap(); + }); + + Ok(()) +} diff --git a/crates/iota-genesis-builder/src/stardust/types/snapshot.rs b/crates/iota-genesis-builder/src/stardust/types/snapshot.rs index b77ef8d2208..33df6f6df33 100644 --- a/crates/iota-genesis-builder/src/stardust/types/snapshot.rs +++ b/crates/iota-genesis-builder/src/stardust/types/snapshot.rs @@ -109,8 +109,8 @@ pub struct FullSnapshotHeader { ledger_milestone_index: MilestoneIndex, treasury_output_milestone_id: MilestoneId, treasury_output_amount: u64, - parameters_milestone_option: MilestoneOption, - output_count: u64, + pub(crate) parameters_milestone_option: MilestoneOption, + pub(crate) output_count: u64, milestone_diff_count: u32, sep_count: u16, }