diff --git a/crates/cli/src/commands/mod.rs b/crates/cli/src/commands/mod.rs index 71e1ffa..b4d90c3 100644 --- a/crates/cli/src/commands/mod.rs +++ b/crates/cli/src/commands/mod.rs @@ -2,5 +2,6 @@ use simplicityhl_core::SimplicityNetwork; pub mod basic; pub mod options; +pub mod smt_storage; const NETWORK: SimplicityNetwork = SimplicityNetwork::LiquidTestnet; diff --git a/crates/cli/src/commands/smt_storage.rs b/crates/cli/src/commands/smt_storage.rs new file mode 100644 index 0000000..614dc31 --- /dev/null +++ b/crates/cli/src/commands/smt_storage.rs @@ -0,0 +1,217 @@ +use crate::commands::NETWORK; +use crate::explorer::{broadcast_tx, fetch_utxo}; +use crate::modules::utils::derive_keypair; +use clap::Subcommand; +use contracts::smt_storage::{ + DEPTH, SMTWitness, SparseMerkleTree, finalize_get_storage_transaction, get_path_bits, + get_smt_storage_compiled_program, smt_storage_taproot_spend_info, +}; +use simplicityhl::elements::pset::serialize::Serialize; +use simplicityhl::simplicity::elements::OutPoint; +use simplicityhl::simplicity::elements::Script; +use simplicityhl::simplicity::elements::taproot::TaprootSpendInfo; +use simplicityhl::simplicity::hex::DisplayHex; +use simplicityhl::tracker::TrackerLogLevel; +use simplicityhl_core::{create_p2pk_signature, finalize_p2pk_transaction, hash_script}; + +fn parse_hex_32(s: &str) -> Result<[u8; 32], String> { + let bytes = hex::decode(s).map_err(|_| "Invalid hex string".to_string())?; + + if bytes.len() != 32 { + return Err(format!( + "Expected 32 bytes (64 hex characters), got {}", + bytes.len() + )); + } + + let mut array = [0u8; 32]; + array.copy_from_slice(&bytes); + Ok(array) +} + +fn parse_bit_path(s: &str) -> Result<[bool; DEPTH], String> { + if s.len() != DEPTH { + return Err(format!("Expected 7 bits, got {}", s.len())); + } + + let mut path = [false; DEPTH]; + + for (ind, char) in s.char_indices() { + if char == 'r' { + path[ind] = true; + } else if char == 'l' { + path[ind] = false; + } else { + return Err(String::from( + "Expected only 'r' and 'l' symbols, got something else.", + )); + } + } + + Ok(path) +} + +/// SMT Storage contract utilities +#[derive(Subcommand, Debug)] +pub enum SMTStorage { + /// Lock collateral on the storage contract (Mint/Fund operation) + GetStorageAddress { + /// The initial 32-byte data payload to store in the tree at the specified path + #[arg(long = "storage-bytes", value_parser = parse_hex_32)] + storage_bytes: [u8; 32], + /// The path in the Merkle Tree use for the contract logic (e.g., "rrll...") + #[arg(long = "path", value_parser = parse_bit_path)] + path: [bool; DEPTH], + + /// Account that will pay for transaction fees and that owns a tokens to send + #[arg(long = "account-index", default_value_t = 0)] + account_index: u32, + }, + /// Build tx transferring an asset UTXO to recipient (LBTC UTXO pays fees) and updating state + TransferFromStorageAddress { + /// Transaction id (hex) and output index (vout) of the ASSET UTXO you will spend + #[arg(long = "storage-utxo")] + storage_utxo: OutPoint, + /// Transaction id (hex) and output index (vout) of the LBTC UTXO used to pay fees (P2PK) + #[arg(long = "fee-utxo")] + fee_utxo: OutPoint, + /// Miner fee in satoshis (LBTC) + #[arg(long = "fee-sats")] + fee_amount: u64, + + /// The current 32-byte data payload stored in the contract (Pre-state) + #[arg(long = "storage-bytes", value_parser = parse_hex_32)] + storage_bytes: [u8; 32], + /// The new 32-byte data payload to replace the old one (Post-state) + #[arg(long = "changed-bytes", value_parser = parse_hex_32)] + changed_bytes: [u8; 32], + /// The Merkle path used to generate the witness for the state transition + #[arg(long = "path", value_parser = parse_bit_path)] + path: [bool; DEPTH], + + /// Account that will pay for transaction fees and that owns a tokens to send + #[arg(long = "account-index", default_value_t = 0)] + account_index: u32, + /// When set, broadcast the built transaction via Esplora and print txid + #[arg(long = "broadcast")] + broadcast: bool, + }, +} + +impl SMTStorage { + /// Handle basic CLI subcommand execution. + /// + /// # Errors + /// Returns error if the subcommand operation fails. + /// + /// # Panics + /// Panics if asset entropy conversion fails. + pub async fn handle(&self) -> anyhow::Result<()> { + match self { + Self::GetStorageAddress { + storage_bytes, + path, + account_index, + } => { + let keypair = derive_keypair(*account_index); + let public_key = keypair.x_only_public_key().0; + + let address = contracts::sdk::get_storage_address( + &public_key, + storage_bytes, + *path, + NETWORK, + )?; + + let mut script_hash: [u8; 32] = hash_script(&address.script_pubkey()); + script_hash.reverse(); + + println!("X Only Public Key: {public_key}"); + println!("P2PK Address: {address}"); + println!("Script hash: {}", hex::encode(script_hash)); + + Ok(()) + } + Self::TransferFromStorageAddress { + storage_utxo: storage_utxo_outpoint, + fee_utxo: fee_utxo_outpoint, + fee_amount, + storage_bytes, + changed_bytes, + path, + account_index, + broadcast, + } => { + let keypair = derive_keypair(*account_index); + let public_key = keypair.x_only_public_key().0; + + let storage_tx_out = fetch_utxo(*storage_utxo_outpoint).await?; + let fee_tx_out = fetch_utxo(*fee_utxo_outpoint).await?; + + let mut smt = SparseMerkleTree::new(); + let merkle_hashes = smt.update(storage_bytes, *path); + + let merkle_data = + std::array::from_fn(|i| (merkle_hashes[DEPTH - i - 1], path[DEPTH - i - 1])); + + let witness = SMTWitness::new( + &public_key.serialize(), + storage_bytes, + get_path_bits(path, true), + &merkle_data, + ); + smt.update(changed_bytes, *path); + + let program = get_smt_storage_compiled_program(); + let cmr = program.commit().cmr(); + + let old_spend_info: TaprootSpendInfo = + smt_storage_taproot_spend_info(public_key, storage_bytes, &merkle_data, cmr); + + let new_spend_info = + smt_storage_taproot_spend_info(public_key, changed_bytes, &merkle_data, cmr); + let new_script_pubkey = Script::new_v1_p2tr_tweaked(new_spend_info.output_key()); + + let pst = contracts::sdk::transfer_asset_with_storage( + (*storage_utxo_outpoint, storage_tx_out.clone()), + (*fee_utxo_outpoint, fee_tx_out.clone()), + *fee_amount, + &new_script_pubkey, + )?; + + let tx = pst.extract_tx()?; + let utxos = vec![storage_tx_out, fee_tx_out]; + + let tx = finalize_get_storage_transaction( + tx, + &old_spend_info, + &witness, + &program, + &utxos, + 0, + NETWORK, + TrackerLogLevel::None, + )?; + + let signature = create_p2pk_signature(&tx, &utxos, &keypair, 1, NETWORK)?; + let tx = finalize_p2pk_transaction( + tx, + &utxos, + &public_key, + &signature, + 1, + NETWORK, + TrackerLogLevel::None, + )?; + + if *broadcast { + println!("Broadcasted txid: {}", broadcast_tx(&tx).await?); + } else { + println!("{}", tx.serialize().to_lower_hex_string()); + } + + Ok(()) + } + } + } +} diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index b5b3b12..2df1fe2 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -15,6 +15,7 @@ use clap::{Parser, Subcommand}; use crate::commands::basic::Basic; use crate::commands::options::Options; +use crate::commands::smt_storage::SMTStorage; /// Command-line entrypoint for the Simplicity helper CLI. #[derive(Parser, Debug)] @@ -41,6 +42,11 @@ enum Commands { #[command(subcommand)] options: Box, }, + /// Storage utilities + Storage { + #[command(subcommand)] + storage: Box, + }, } #[tokio::main] @@ -48,5 +54,6 @@ async fn main() -> Result<()> { match Cli::parse().command { Commands::Basic { basic } => basic.handle().await, Commands::Options { options } => options.handle().await, + Commands::Storage { storage } => storage.handle().await, } } diff --git a/crates/contracts/Cargo.toml b/crates/contracts/Cargo.toml index 0d4aeb8..cabbc19 100644 --- a/crates/contracts/Cargo.toml +++ b/crates/contracts/Cargo.toml @@ -10,7 +10,7 @@ keywords = ["simplicity", "liquid", "bitcoin", "elements", "contracts"] categories = ["cryptography::cryptocurrencies"] [features] -default = ["sdk-basic", "finance-options", "finance-dcd", "finance-option-offer", "simple-storage", "bytes32-tr-storage", "array-tr-storage"] +default = ["sdk-basic", "finance-options", "finance-dcd", "finance-option-offer", "simple-storage", "bytes32-tr-storage", "array-tr-storage", "smt-storage"] sdk-basic = [] finance-options = [] finance-dcd = [] @@ -18,6 +18,7 @@ finance-option-offer = [] simple-storage = [] bytes32-tr-storage = [] array-tr-storage = [] +smt-storage = [] [lints] workspace = true diff --git a/crates/contracts/src/lib.rs b/crates/contracts/src/lib.rs index 2f1ebf1..75535c6 100644 --- a/crates/contracts/src/lib.rs +++ b/crates/contracts/src/lib.rs @@ -18,7 +18,8 @@ pub mod bytes32_tr_storage; pub mod finance; #[cfg(feature = "simple-storage")] pub mod simple_storage; - +#[cfg(feature = "smt-storage")] +pub mod smt_storage; #[cfg(feature = "finance-dcd")] pub use finance::dcd; #[cfg(feature = "finance-option-offer")] diff --git a/crates/contracts/src/sdk/mod.rs b/crates/contracts/src/sdk/mod.rs index c5fa6cd..1736ec7 100644 --- a/crates/contracts/src/sdk/mod.rs +++ b/crates/contracts/src/sdk/mod.rs @@ -2,6 +2,8 @@ mod basic; #[cfg(any(feature = "finance-option-offer", feature = "finance-options"))] mod finance; +#[cfg(feature = "smt-storage")] +mod storage; pub mod taproot_pubkey_gen; pub mod validation; @@ -10,3 +12,5 @@ pub mod validation; pub use basic::*; #[cfg(any(feature = "finance-option-offer", feature = "finance-options"))] pub use finance::*; +#[cfg(feature = "smt-storage")] +pub use storage::*; diff --git a/crates/contracts/src/sdk/storage/get_storage_address.rs b/crates/contracts/src/sdk/storage/get_storage_address.rs new file mode 100644 index 0000000..005e2d6 --- /dev/null +++ b/crates/contracts/src/sdk/storage/get_storage_address.rs @@ -0,0 +1,48 @@ +use simplicityhl::elements::schnorr::XOnlyPublicKey; +use simplicityhl::simplicity::elements::Address; +use simplicityhl::simplicity::elements::Script; +use simplicityhl::simplicity::elements::taproot::TaprootSpendInfo; +use simplicityhl_core::SimplicityNetwork; + +use crate::error::TransactionBuildError; +use crate::smt_storage::{ + DEPTH, SparseMerkleTree, get_smt_storage_compiled_program, smt_storage_taproot_spend_info, +}; + +/// Derives the Taproot address for the SMT storage contract based on its initial state. +/// +/// This function calculates the script pubkey by committing to the Simplicity program +/// configured with the provided `storage_bytes` (root hash) and `path`. It then +/// encodes this script into a network-specific address. +/// +/// # Errors +/// +/// Returns an error if: +/// - The function signature requires a `Result` for consistency with the builder API, +/// though the current implementation primarily panics on failure rather than returning `Err`. +/// +/// # Panics +/// +/// Panics if: +/// - The generated script is invalid for address creation (e.g., invalid witness program). +pub fn get_storage_address( + storage_key: &XOnlyPublicKey, + storage_bytes: &[u8; 32], + path: [bool; DEPTH], + network: SimplicityNetwork, +) -> Result { + let mut smt = SparseMerkleTree::new(); + let merkle_hashes = smt.update(storage_bytes, path); + + let merkle_data = std::array::from_fn(|i| (merkle_hashes[DEPTH - i - 1], path[DEPTH - i - 1])); + + let program = get_smt_storage_compiled_program(); + let cmr = program.commit().cmr(); + + let mint_spend_info: TaprootSpendInfo = + smt_storage_taproot_spend_info(*storage_key, storage_bytes, &merkle_data, cmr); + + let mint_script_pubkey = Script::new_v1_p2tr_tweaked(mint_spend_info.output_key()); + + Ok(Address::from_script(&mint_script_pubkey, None, network.address_params()).unwrap()) +} diff --git a/crates/contracts/src/sdk/storage/mod.rs b/crates/contracts/src/sdk/storage/mod.rs new file mode 100644 index 0000000..41df9ef --- /dev/null +++ b/crates/contracts/src/sdk/storage/mod.rs @@ -0,0 +1,5 @@ +mod get_storage_address; +mod transfer_from_storage_address; + +pub use get_storage_address::*; +pub use transfer_from_storage_address::*; diff --git a/crates/contracts/src/sdk/storage/transfer_from_storage_address.rs b/crates/contracts/src/sdk/storage/transfer_from_storage_address.rs new file mode 100644 index 0000000..9ea8251 --- /dev/null +++ b/crates/contracts/src/sdk/storage/transfer_from_storage_address.rs @@ -0,0 +1,85 @@ +use simplicityhl::elements::bitcoin::secp256k1; +use simplicityhl::elements::pset::{Input, Output, PartiallySignedTransaction}; +use simplicityhl::simplicity::elements::OutPoint; +use simplicityhl::simplicity::elements::Script; +use simplicityhl::simplicity::elements::TxOut; + +use crate::error::TransactionBuildError; +use crate::sdk::validation::TxOutExt as _; + +/// Derives the Taproot address for the SMT storage contract. +/// +/// This function constructs the Simplicity program committed to the provided `storage_bytes` +/// (initial state) and `path`. It then calculates the Taproot script pubkey by tweaking +/// the `storage_key` with the program's commitment (CMR) and converts it into a +/// human-readable address for the specified network. +/// +/// Use this address to fund (mint) the contract by sending assets to it. +/// +/// # Arguments +/// +/// * `storage_key` - The internal X-only public key used for Taproot tweaking (usually an unspendable key). +/// * `storage_bytes` - The 32-byte data payload (SMT root hash) representing the initial state of the contract. +/// * `path` - The binary path in the Sparse Merkle Tree used to generate the witness data. +/// * `network` - The network parameters (e.g., Liquid Testnet, Mainnet) used to format the address. +/// +/// +/// # Errors +/// +/// This function returns a `Result` to maintain consistency with the builder API, +/// though the current implementation is unlikely to return an `Err` variant unless +/// address generation logic changes. +/// +/// # Panics +/// +/// Panics if the generated script is not a valid witness program (this should theoretically +/// never happen with a valid `new_v1_p2tr_tweaked` script). +pub fn transfer_asset_with_storage( + storage_utxo: (OutPoint, TxOut), + fee_utxo: (OutPoint, TxOut), + fee_amount: u64, + new_script_pubkey: &Script, +) -> Result { + let (storage_out_point, storage_tx_out) = storage_utxo; + let (fee_out_point, fee_tx_out) = fee_utxo; + + let (storage_asset_id, total_input_storage_amount) = storage_tx_out.explicit()?; + let (fee_asset_id, change_amount) = ( + fee_tx_out.explicit_asset()?, + fee_tx_out.validate_amount(fee_amount)?, + ); + + let mut pst = PartiallySignedTransaction::new_v2(); + let change_recipient_script = fee_tx_out.script_pubkey.clone(); + + let mut storage_input = Input::from_prevout(storage_out_point); + storage_input.witness_utxo = Some(storage_tx_out.clone()); + pst.add_input(storage_input); + + let mut fee_input = Input::from_prevout(fee_out_point); + fee_input.witness_utxo = Some(fee_tx_out.clone()); + pst.add_input(fee_input); + + pst.add_output(Output::new_explicit( + new_script_pubkey.clone(), + total_input_storage_amount, + storage_asset_id, + None, + )); + + if change_amount > 0 { + pst.add_output(Output::new_explicit( + change_recipient_script.clone(), + change_amount, + fee_asset_id, + None, + )); + } + + pst.add_output(Output::from_txout(TxOut::new_fee(fee_amount, fee_asset_id))); + + pst.extract_tx()? + .verify_tx_amt_proofs(secp256k1::SECP256K1, &[storage_tx_out, fee_tx_out])?; + + Ok(pst) +} diff --git a/crates/contracts/src/sdk/validation.rs b/crates/contracts/src/sdk/validation.rs index f6130dd..c4fdf2d 100644 --- a/crates/contracts/src/sdk/validation.rs +++ b/crates/contracts/src/sdk/validation.rs @@ -72,6 +72,7 @@ impl TxOutExt for TxOut { }) } + // TODO: Change this validation to another func fn validate_amount(&self, required: u64) -> Result { let available = self.explicit_value()?; diff --git a/crates/contracts/src/smt_storage/build_witness.rs b/crates/contracts/src/smt_storage/build_witness.rs new file mode 100644 index 0000000..16e177d --- /dev/null +++ b/crates/contracts/src/smt_storage/build_witness.rs @@ -0,0 +1,111 @@ +use std::collections::HashMap; + +use simplicityhl::num::U256; +use simplicityhl::types::{ResolvedType, TypeConstructible, UIntType}; +use simplicityhl::value::{UIntValue, ValueConstructible}; +use simplicityhl::{WitnessValues, str::WitnessName}; + +#[allow(non_camel_case_types)] +pub type u256 = [u8; 32]; + +/// The fixed depth of the Sparse Merkle Tree (SMT). +/// +/// This is set to 8 because Simplicity currently requires fixed-length arrays +/// and cannot dynamically resolve array lengths using `param::LEN`. +pub const DEPTH: usize = 8; + +#[derive(Debug, Clone, bincode::Encode, bincode::Decode, PartialEq, Eq)] +pub struct SMTWitness { + /// The internal public key used for Taproot tweaking. + /// + /// This corresponds to the `key` parameter in the Simplicity expression: + /// `let tweaked_key: u256 = jet::build_taptweak(key, tap_node);`. + key: u256, + + /// The leaf node (value) being stored or verified in the tree. + leaf: u256, + + /// A bitwise representation of the tree traversal path. + /// + /// Since `DEPTH` is 8, the path fits into a single `u8`. + /// * `1` (or `true`) represents a move to the **Right**. + /// * `0` (or `false`) represents a move to the **Left**. + /// + /// **Note:** The bits are ordered from the **leaf up to the root**. + /// This order is chosen to simplify bitwise processing within the Simplicity contract. + path_bits: u8, + + /// The sibling nodes required to reconstruct the Merkle path. + /// + /// Each element is a tuple containing the sibling's hash and a boolean direction. + /// Like `path_bits`, this array is ordered from the **leaf up to the root** + /// to facilitate efficient processing in the Simplicity loop. + merkle_data: [(u256, bool); DEPTH], +} + +impl SMTWitness { + #[must_use] + pub fn new( + key: &u256, + leaf: &u256, + path_bits: u8, + merkle_data: &[(u256, bool); DEPTH], + ) -> Self { + Self { + key: *key, + leaf: *leaf, + path_bits, + merkle_data: *merkle_data, + } + } +} + +impl Default for SMTWitness { + fn default() -> Self { + Self { + key: [0u8; 32], + leaf: [0u8; 32], + path_bits: 0, + merkle_data: [([0u8; 32], false); DEPTH], + } + } +} + +#[must_use] +pub fn build_smt_storage_witness(witness: &SMTWitness) -> WitnessValues { + let values: Vec = witness + .merkle_data + .iter() + .map(|(value, is_right)| { + let hash_val = + simplicityhl::Value::from(UIntValue::U256(U256::from_byte_array(*value))); + let direction_val = simplicityhl::Value::from(*is_right); + + simplicityhl::Value::product(hash_val, direction_val) + }) + .collect(); + + let element_type = simplicityhl::types::TypeConstructible::product( + UIntType::U256.into(), + ResolvedType::boolean(), + ); + + simplicityhl::WitnessValues::from(HashMap::from([ + ( + WitnessName::from_str_unchecked("KEY"), + simplicityhl::Value::from(UIntValue::U256(U256::from_byte_array(witness.key))), + ), + ( + WitnessName::from_str_unchecked("LEAF"), + simplicityhl::Value::from(UIntValue::U256(U256::from_byte_array(witness.leaf))), + ), + ( + WitnessName::from_str_unchecked("PATH_BITS"), + simplicityhl::Value::from(UIntValue::U8(witness.path_bits)), + ), + ( + WitnessName::from_str_unchecked("MERKLE_DATA"), + simplicityhl::Value::array(values, element_type), + ), + ])) +} diff --git a/crates/contracts/src/smt_storage/mod.rs b/crates/contracts/src/smt_storage/mod.rs new file mode 100644 index 0000000..94e4993 --- /dev/null +++ b/crates/contracts/src/smt_storage/mod.rs @@ -0,0 +1,381 @@ +use std::sync::Arc; + +use simplicityhl::elements::TxInWitness; +use simplicityhl::elements::TxOut; +use simplicityhl::elements::taproot::ControlBlock; +use simplicityhl::simplicity::bitcoin::secp256k1; +use simplicityhl::simplicity::elements::hashes::HashEngine as _; +use simplicityhl::simplicity::elements::taproot::{LeafVersion, TaprootBuilder, TaprootSpendInfo}; +use simplicityhl::simplicity::elements::{Script, Transaction}; +use simplicityhl::simplicity::hashes::{Hash, sha256}; +use simplicityhl::simplicity::jet::Elements; +use simplicityhl::simplicity::jet::elements::{ElementsEnv, ElementsUtxo}; +use simplicityhl::simplicity::{Cmr, RedeemNode, leaf_version}; +use simplicityhl::tracker::TrackerLogLevel; +use simplicityhl::{Arguments, CompiledProgram, TemplateProgram}; +use simplicityhl_core::{ProgramError, SimplicityNetwork, run_program}; + +mod build_witness; +mod smt; + +pub use build_witness::{DEPTH, SMTWitness, build_smt_storage_witness, u256}; +pub use smt::SparseMerkleTree; + +#[must_use] +pub fn get_path_bits(path: &[bool], reverse: bool) -> u8 { + let mut path_bits = 0u8; + for (i, direction) in path.iter().enumerate().take(DEPTH) { + let shift = if reverse { DEPTH - i - 1 } else { i }; + path_bits |= u8::from(*direction) << shift; + } + path_bits +} + +pub const SMT_STORAGE_SOURCE: &str = include_str!("source_simf/smt_storage.simf"); + +/// Get the storage template program for instantiation. +/// +/// # Panics +/// +/// Panics if the embedded source fails to compile (should never happen). +#[must_use] +pub fn get_smt_storage_template_program() -> TemplateProgram { + TemplateProgram::new(SMT_STORAGE_SOURCE).expect("INTERNAL: expected to compile successfully.") +} + +/// Get compiled storage program, panicking on failure. +/// +/// # Panics +/// +/// Panics if program instantiation fails. +#[must_use] +pub fn get_smt_storage_compiled_program() -> CompiledProgram { + let program = get_smt_storage_template_program(); + + program.instantiate(Arguments::default(), true).unwrap() +} + +/// Execute storage program with new state. +/// +/// # Errors +/// Returns error if program execution fails. +pub fn execute_smt_storage_program( + witness: &SMTWitness, + compiled_program: &CompiledProgram, + env: &ElementsEnv>, + runner_log_level: TrackerLogLevel, +) -> Result>, ProgramError> { + let witness_values = build_smt_storage_witness(witness); + Ok(run_program(compiled_program, witness_values, env, runner_log_level)?.0) +} + +#[must_use] +pub fn smt_storage_script_ver(cmr: Cmr) -> (Script, LeafVersion) { + (Script::from(cmr.as_ref().to_vec()), leaf_version()) +} + +/// Computes the control block for the given CMR and spend info. +/// +/// # Panics +/// +/// Panics if the control block cannot be retrieved. This typically happens if the +/// provided `cmr` corresponds to a script that is not present in the `spend_info` tree. +#[must_use] +pub fn control_block(cmr: Cmr, spend_info: &TaprootSpendInfo) -> ControlBlock { + spend_info + .control_block(&smt_storage_script_ver(cmr)) + .expect("must get control block") +} + +/// Computes the TapData-tagged hash of the Simplicity state (SMT Root). +/// +/// This involves hashing the tag "`TapData`" twice, followed by the leaf value +/// and the path bits, and finally performing the Merkle proof hashing up to the root. +/// +/// # Security Note: Second Preimage Resistance +/// +/// The `raw_path` (bit representation of the path) is included in the initial hash of the leaf +/// alongside the `leaf` data. +/// +/// This is a defense mechanism against **second preimage attacks** (specifically, Merkle substitution attacks). +/// In Merkle trees (especially those with variable depth), an attacker might try to present +/// an internal node as a leaf, or vice versa. By including the path in the leaf's hash, +/// we strictly bind the data to its specific position in the tree hierarchy. +/// +/// Although `DEPTH` is currently fixed (which mitigates some of these risks naturally), +/// this explicit domain separation ensures that a valid proof for a leaf at one position +/// cannot be reused or confused with a node at another level or branch, ensuring future +/// safety even if depth constraints change. +/// +/// # Panics +/// +/// This function **does not panic**. +/// All hashing operations (`sha256::Hash::engine`, `input`, `from_engine`) are +/// infallible, and iterating over the state limbs is safe. +#[must_use] +pub fn compute_tapdata_tagged_hash_of_the_state( + leaf: &u256, + path: &[(u256, bool); DEPTH], +) -> sha256::Hash { + let tag = sha256::Hash::hash(b"TapData"); + let mut eng = sha256::Hash::engine(); + eng.input(tag.as_byte_array()); + eng.input(tag.as_byte_array()); + eng.input(leaf); + + let raw_path: [bool; DEPTH] = std::array::from_fn(|i| path[i].1); + eng.input(&[get_path_bits(&raw_path, false)]); + + let mut current_hash = sha256::Hash::from_engine(eng); + + for (hash, is_right_direction) in path { + let mut eng = sha256::Hash::engine(); + + if *is_right_direction { + eng.input(hash); + eng.input(¤t_hash.to_byte_array()); + } else { + eng.input(¤t_hash.to_byte_array()); + eng.input(hash); + } + + current_hash = sha256::Hash::from_engine(eng); + } + current_hash +} + +/// Given a Simplicity CMR and an internal key, computes the [`TaprootSpendInfo`] +/// for a Taptree with this CMR as its single leaf. +/// +/// # Panics +/// +/// This function **panics** if building the taproot tree fails (the calls to +/// `TaprootBuilder::add_leaf_with_ver` or `.add_hidden` return `Err`) or if +/// finalizing the builder fails. Those panics come from the `.expect(...)` +/// calls on the builder methods. +#[must_use] +pub fn smt_storage_taproot_spend_info( + internal_key: secp256k1::XOnlyPublicKey, + leaf: &u256, + path: &[(u256, bool); DEPTH], + cmr: Cmr, +) -> TaprootSpendInfo { + let (script, version) = smt_storage_script_ver(cmr); + let state_hash = compute_tapdata_tagged_hash_of_the_state(leaf, path); + + // Build taproot tree with hidden leaf + let builder = TaprootBuilder::new() + .add_leaf_with_ver(1, script, version) + .expect("tap tree should be valid") + .add_hidden(1, state_hash) + .expect("tap tree should be valid"); + + builder + .finalize(secp256k1::SECP256K1, internal_key) + .expect("tap tree should be valid") +} + +/// Constructs and verifies the Simplicity environment for the SMT storage execution. +/// +/// # Errors +/// +/// Returns an error if: +/// - The `input_index` is out of bounds for the provided `utxos`. +/// - The script pubkey of the UTXO at `input_index` does not match the expected SMT storage address. +pub fn get_and_verify_env( + tx: &Transaction, + program: &CompiledProgram, + spend_info: &TaprootSpendInfo, + utxos: &[TxOut], + network: SimplicityNetwork, + input_index: usize, +) -> Result>, ProgramError> { + let genesis_hash = network.genesis_block_hash(); + let cmr = program.commit().cmr(); + + if utxos.len() <= input_index { + return Err(ProgramError::UtxoIndexOutOfBounds { + input_index, + utxo_count: utxos.len(), + }); + } + + let target_utxo = &utxos[input_index]; + let script_pubkey = Script::new_v1_p2tr_tweaked(spend_info.output_key()); + + if target_utxo.script_pubkey != script_pubkey { + return Err(ProgramError::ScriptPubkeyMismatch { + expected_hash: script_pubkey.script_hash().to_string(), + actual_hash: target_utxo.script_pubkey.script_hash().to_string(), + }); + } + + Ok(ElementsEnv::new( + Arc::new(tx.clone()), + utxos + .iter() + .map(|utxo| ElementsUtxo { + script_pubkey: utxo.script_pubkey.clone(), + asset: utxo.asset, + value: utxo.value, + }) + .collect(), + u32::try_from(input_index)?, + cmr, + control_block(cmr, spend_info), + None, + genesis_hash, + )) +} + +/// Finalizes the SMT storage transaction by executing the program and attaching the witness. +/// +/// # Errors +/// +/// Returns an error if: +/// - The environment verification fails (e.g., mismatched UTXOs or script pubkeys). +/// - The SMT storage program execution fails during the simulation. +#[allow(clippy::too_many_arguments)] +pub fn finalize_get_storage_transaction( + mut tx: Transaction, + spend_info: &TaprootSpendInfo, + witness: &SMTWitness, + storage_program: &CompiledProgram, + utxos: &[TxOut], + input_index: usize, + network: SimplicityNetwork, + log_level: TrackerLogLevel, +) -> Result { + let env = get_and_verify_env( + &tx, + storage_program, + spend_info, + utxos, + network, + input_index, + )?; + + let pruned = execute_smt_storage_program(witness, storage_program, &env, log_level)?; + + let (simplicity_program_bytes, simplicity_witness_bytes) = pruned.to_vec_with_witness(); + let cmr = pruned.cmr(); + + tx.input[input_index].witness = TxInWitness { + amount_rangeproof: None, + inflation_keys_rangeproof: None, + script_witness: vec![ + simplicity_witness_bytes, + simplicity_program_bytes, + cmr.as_ref().to_vec(), + control_block(cmr, spend_info).serialize(), + ], + pegin_witness: vec![], + }; + + Ok(tx) +} + +#[cfg(test)] +mod smt_storage_tests { + use super::*; + use anyhow::Result; + use simplicityhl::elements::secp256k1_zkp::rand::{Rng, thread_rng}; + use std::sync::Arc; + + use simplicityhl::elements::confidential::{Asset, Value}; + use simplicityhl::elements::pset::{Input, Output, PartiallySignedTransaction}; + use simplicityhl::elements::{AssetId, BlockHash, OutPoint, Script, Txid}; + use simplicityhl::simplicity::jet::elements::{ElementsEnv, ElementsUtxo}; + + fn add_elements(smt: &mut SparseMerkleTree, num: u64) -> (u256, [u256; DEPTH], [bool; DEPTH]) { + let mut rng = thread_rng(); + + let mut leaf = [0u8; 32]; + let mut merkle_hashes = [[0u8; 32]; DEPTH]; + let mut path = [false; DEPTH]; + + for _ in 0..num { + leaf = rng.r#gen(); + path = std::array::from_fn(|_| rng.r#gen()); + merkle_hashes = smt.update(&leaf, path); + } + + (leaf, merkle_hashes, path) + } + + #[rustfmt::skip] // mangles byte vectors + fn unspendable_internal_key() -> secp256k1::XOnlyPublicKey { + secp256k1::XOnlyPublicKey::from_slice(&[ + 0x50, 0x92, 0x9b, 0x74, 0xc1, 0xa0, 0x49, 0x54, 0xb7, 0x8b, 0x4b, 0x60, 0x35, 0xe9, 0x7a, 0x5e, + 0x07, 0x8a, 0x5a, 0x0f, 0x28, 0xec, 0x96, 0xd5, 0x47, 0xbf, 0xee, 0x9a, 0xce, 0x80, 0x3a, 0xc0, + ]) + .expect("key should be valid") + } + + #[test] + fn test_smt_storage_mint_path() -> Result<()> { + let mut smt = SparseMerkleTree::new(); + let (old_leaf, merkle_hashes, path) = add_elements(&mut smt, 1); + + let merkle_data = + std::array::from_fn(|i| (merkle_hashes[DEPTH - i - 1], path[DEPTH - i - 1])); + + let internal_key = unspendable_internal_key(); + let witness = SMTWitness::new( + &internal_key.serialize(), + &old_leaf, + get_path_bits(&path, true), + &merkle_data, + ); + + // Set last leaf qword to 1 + let mut new_leaf = old_leaf; + for byte in new_leaf.iter_mut().skip(24) { + *byte = 0; + } + new_leaf[31] = 1; + smt.update(&new_leaf, path); + + let program = get_smt_storage_compiled_program(); + let cmr = program.commit().cmr(); + + let old_spend_info: TaprootSpendInfo = + smt_storage_taproot_spend_info(internal_key, &old_leaf, &merkle_data, cmr); + let old_script_pubkey = Script::new_v1_p2tr_tweaked(old_spend_info.output_key()); + + let new_spend_info = + smt_storage_taproot_spend_info(internal_key, &new_leaf, &merkle_data, cmr); + let new_script_pubkey = Script::new_v1_p2tr_tweaked(new_spend_info.output_key()); + + let mut pst = PartiallySignedTransaction::new_v2(); + let outpoint0 = OutPoint::new(Txid::from_slice(&[0; 32])?, 0); + pst.add_input(Input::from_prevout(outpoint0)); + pst.add_output(Output::new_explicit( + new_script_pubkey.clone(), + 0, + AssetId::default(), + None, + )); + + let env = ElementsEnv::new( + Arc::new(pst.extract_tx()?), + vec![ElementsUtxo { + script_pubkey: old_script_pubkey, + asset: Asset::default(), + value: Value::default(), + }], + 0, + cmr, + control_block(cmr, &old_spend_info), + None, + BlockHash::all_zeros(), + ); + + assert!( + execute_smt_storage_program(&witness, &program, &env, TrackerLogLevel::Trace).is_ok(), + "expected success mint path" + ); + + Ok(()) + } +} diff --git a/crates/contracts/src/smt_storage/smt.rs b/crates/contracts/src/smt_storage/smt.rs new file mode 100644 index 0000000..8c20c5d --- /dev/null +++ b/crates/contracts/src/smt_storage/smt.rs @@ -0,0 +1,236 @@ +use simplicityhl::simplicity::elements::hashes::HashEngine as _; +use simplicityhl::simplicity::hashes::{Hash, sha256}; + +use crate::smt_storage::get_path_bits; + +use super::build_witness::{DEPTH, u256}; + +/// Represents a node within the Sparse Merkle Tree. +/// +/// The tree is structured as a recursive binary tree where: +/// - [`TreeNode::Leaf`] represents the bottom-most layer containing the actual data hash. +/// - [`TreeNode::Branch`] represents an internal node containing the combined hash of its children. +#[derive(Debug, Clone, bincode::Encode, bincode::Decode, PartialEq, Eq)] +enum TreeNode { + /// A leaf node at the bottom of the tree. + /// + /// Contains the `leaf_hash` which is the hash of the stored value (or a default empty value). + Leaf { leaf_hash: u256 }, + /// An internal branch node. + /// + /// Contains pointers to the `left` and `right` child nodes and their combined `hash`. + /// The `hash` is typically calculated as `Hash(Left_Child_Hash || Right_Child_Hash)`. + Branch { + hash: u256, + left: Box, + right: Box, + }, +} + +impl TreeNode { + pub fn get_hash(&self) -> u256 { + match self { + TreeNode::Leaf { leaf_hash } => *leaf_hash, + TreeNode::Branch { hash, .. } => *hash, + } + } +} + +/// An implementation of a Sparse Merkle Tree (SMT) with fixed depth. +/// +/// Functionally, this structure acts as a **Key-Value store**: +/// - **Key**: The path from the root to the leaf. +/// - **Value**: The data hash stored at that specific leaf. +/// +/// A Sparse Merkle Tree is a perfectly balanced binary tree where most leaves are empty (contain default values). +/// Instead of storing every node of the massive tree (which would be impossible for depths like 256), +/// this implementation stores only the non-empty branches. +/// +/// # Optimization: Precalculated Hashes +/// +/// To efficiently handle the "sparse" nature of the tree, we utilize a `precalculate_hashes` array. +/// This array stores the default hash values for empty subtrees at each height level. +/// - `precalculate_hashes[0]` is the hash of an empty leaf. +/// - `precalculate_hashes[1]` is the hash of a branch connecting two empty leaves. +/// - ...and so on. +/// +/// This allows getting the hash of an empty branch at any level in **O(1)** time without recomputing it. +/// +/// # Security & Attack Mitigation +/// +/// This implementation explicitly guards against **Second Preimage Attacks** (specifically +/// Merkle Substitution or Length Extension attacks) using the following techniques: +/// +/// 1. **Path Binding (Position Binding)**: +/// The `raw_path` (bit representation of the tree path) is mixed into the initial leaf hash via +/// `eng.input(&[get_path_bits(...)])`. +/// * *Why?* This binds the data to a specific location in the tree. It prevents an attacker from +/// taking a valid internal node hash (from a deeper level) and presenting it as a valid leaf +/// at a higher level. Even if the data matches, the path/position will differ, changing the hash. +/// +/// 2. **Domain Separation**: +/// The function initializes with `Hash(b"TapData")`. +/// * *Why?* This ensures that hashes generated for this SMT state cannot be confused with other +/// Bitcoin/Elements hashes (like `TapLeaf` or `TapBranch` hashes), preventing cross-context collisions. +/// +/// # See Also +/// +/// * [What is a Sparse Merkle Tree?](https://medium.com/@kelvinfichter/whats-a-sparse-merkle-tree-acda70aeb837) +/// * [Merkle Tree Concepts](https://en.wikipedia.org/wiki/Merkle_tree) +pub struct SparseMerkleTree { + /// The root node of the tree, initialized to a leaf containing `precalculate_hashes[0]` by default. + root: Box, + /// Cache of default hashes for empty subtrees at each depth level [0..DEPTH]. + precalculate_hashes: [u256; DEPTH], +} + +impl SparseMerkleTree { + /// Initializes a new SMT with precalculated default hashes. + /// + /// Computes hashes for empty subtrees at all depths (0..DEPTH) to optimize + /// calculation. The tree starts with a root pointing to the default empty leaf (`precalculate_hashes[0]`). + #[must_use] + pub fn new() -> Self { + let mut precalculate_hashes = [[0u8; 32]; DEPTH]; + let mut eng = sha256::Hash::engine(); + let zero = [0u8; 32]; + eng.input(&zero); + precalculate_hashes[0] = *sha256::Hash::from_engine(eng).as_byte_array(); + + for i in 1..DEPTH { + let mut eng = sha256::Hash::engine(); + eng.input(&precalculate_hashes[i - 1]); + eng.input(&precalculate_hashes[i - 1]); + precalculate_hashes[i] = *sha256::Hash::from_engine(eng).as_byte_array(); + } + + Self { + root: Box::new(TreeNode::Leaf { + leaf_hash: precalculate_hashes[0], + }), + precalculate_hashes, + } + } + + /// Computes parent hash: `SHA256(left_child_hash || right_child_hash)`. + fn calculate_hash(left: &mut TreeNode, right: &mut TreeNode) -> u256 { + let mut eng = sha256::Hash::engine(); + eng.input(&left.get_hash()); + eng.input(&right.get_hash()); + *sha256::Hash::from_engine(eng).as_byte_array() + } + + /// Internal recursive DFS helper to insert or update a node. + /// + /// Navigates down based on `path`. Expands `Leaf` nodes into `Branch` nodes + /// when descending. Collects sibling hashes into `hashes` and recalculates + /// branch hashes on the return path. + fn traverse( + defaults: &[u256], + leaf: &u256, + path: &[bool], + ind: usize, + root: &mut Box, + hashes: &mut [u256], + ) { + if ind >= DEPTH { + let tag = sha256::Hash::hash(b"TapData"); + let mut eng = sha256::Hash::engine(); + eng.input(tag.as_byte_array()); + eng.input(tag.as_byte_array()); + eng.input(leaf); + eng.input(&[get_path_bits(path, true)]); + + **root = TreeNode::Leaf { + leaf_hash: *sha256::Hash::from_engine(eng).as_byte_array(), + }; + return; + } + + let (child_zero, remaining_defaults) = defaults + .split_last() + .expect("Defaults length must match path length"); + + if matches!(**root, TreeNode::Leaf { .. }) { + let new_branch = Box::new(TreeNode::Branch { + hash: [0u8; 32], + left: Box::new(TreeNode::Leaf { + leaf_hash: *child_zero, + }), + right: Box::new(TreeNode::Leaf { + leaf_hash: *child_zero, + }), + }); + + *root = new_branch; + } + + let (current_hash_slot, remaining_hashes) = hashes + .split_first_mut() + .expect("Hashes length must match path length"); + + if let TreeNode::Branch { + ref mut left, + ref mut right, + ref mut hash, + } = **root + { + if path[ind] { + *current_hash_slot = left.get_hash(); + Self::traverse( + remaining_defaults, + leaf, + path, + ind + 1, + right, + remaining_hashes, + ); + } else { + *current_hash_slot = right.get_hash(); + Self::traverse( + remaining_defaults, + leaf, + path, + ind + 1, + left, + remaining_hashes, + ); + } + + *hash = Self::calculate_hash(left, right); + } else { + unreachable!("Should be a branch at this point"); + } + } + + /// Inserts or updates a leaf at the specified path. + /// + /// Traverses the tree, modifying the target leaf and recalculating the root. + /// + /// # Arguments + /// + /// * `leaf` - The 32-byte value to be stored at the target position. + /// * `path` - The navigation path represented as a fixed-size boolean array. + /// The order of bits is from **Root to Leaf** (index 0 is the first step from the root). + /// + /// # Returns + /// An array of sibling hashes (Merkle path) collected from the root down to the leaf. + pub fn update(&mut self, leaf: &u256, path: [bool; DEPTH]) -> [u256; DEPTH] { + let mut hashes = self.precalculate_hashes; + Self::traverse( + &self.precalculate_hashes, + leaf, + &path, + 0, + &mut self.root, + &mut hashes, + ); + hashes + } +} + +impl Default for SparseMerkleTree { + fn default() -> Self { + Self::new() + } +} diff --git a/crates/contracts/src/smt_storage/source_simf/smt_storage.simf b/crates/contracts/src/smt_storage/source_simf/smt_storage.simf new file mode 100644 index 0000000..cd10da2 --- /dev/null +++ b/crates/contracts/src/smt_storage/source_simf/smt_storage.simf @@ -0,0 +1,68 @@ +/* + * Extends `bytes32_tr_storage` using `array_fold` for larger buffers. + * Optimized for small, fixed-size states where linear hashing is more efficient + * than Merkle Trees. By avoiding proof overhead like sibling hashes, we reduce + * witness size and simplify contract logic for small N. + */ +fn hash_array_tr_storage_with_update(elem: (u256, bool), prev_hash: u256) -> u256 { + let (hash, is_right): (u256, bool) = dbg!(elem); + let ctx: Ctx8 = jet::sha_256_ctx_8_init(); + + let new_hash: Ctx8 = match is_right { + true => { + let ctx: Ctx8 = jet::sha_256_ctx_8_add_32(ctx, hash); + jet::sha_256_ctx_8_add_32(ctx, prev_hash) + }, + false => { + let ctx: Ctx8 = jet::sha_256_ctx_8_add_32(ctx, prev_hash); + jet::sha_256_ctx_8_add_32(ctx, hash) + } + }; + + jet::sha_256_ctx_8_finalize(new_hash) +} + +fn script_hash_for_input_script(key: u256, leaf: u256, path_bits: u8, merkle_data: [(u256, bool); 8]) -> u256 { + let tap_leaf: u256 = jet::tapleaf_hash(); + let ctx: Ctx8 = jet::tapdata_init(); + let ctx: Ctx8 = jet::sha_256_ctx_8_add_32(ctx, leaf); + let ctx: Ctx8 = jet::sha_256_ctx_8_add_1(ctx, path_bits); + + let hash_leaf: u256 = jet::sha_256_ctx_8_finalize(ctx); + + let computed: u256 = array_fold::(merkle_data, hash_leaf); + let tap_node: u256 = jet::build_tapbranch(tap_leaf, computed); + + let tweaked_key: u256 = jet::build_taptweak(key, tap_node); + + let hash_ctx1: Ctx8 = jet::sha_256_ctx_8_init(); + let hash_ctx2: Ctx8 = jet::sha_256_ctx_8_add_2(hash_ctx1, 0x5120); // Segwit v1, length 32 + let hash_ctx3: Ctx8 = jet::sha_256_ctx_8_add_32(hash_ctx2, tweaked_key); + jet::sha_256_ctx_8_finalize(hash_ctx3) +} + +fn main() { + let key: u256 = witness::KEY; + let leaf_data: u256 = witness::LEAF; + let path_bits: u8 = witness::PATH_BITS; + + // Path and hash + let merkle_data: [(u256, bool); 8] = witness::MERKLE_DATA; + let (leaf1, leaf2, leaf3, leaf4): (u64, u64, u64, u64) = ::into(leaf_data); + + // Load + assert!(jet::eq_256( + script_hash_for_input_script(key, leaf_data, path_bits, merkle_data), + unwrap(jet::input_script_hash(jet::current_index())) + )); + + // There may be arbitrary logic here + let new_leaf4: u64 = 1; + let new_leaf: u256 = <(u64, u64, u64, u64)>::into((leaf1, leaf2, leaf3, new_leaf4)); + + // Store + assert!(jet::eq_256( + script_hash_for_input_script(key, new_leaf, path_bits, merkle_data), + unwrap(jet::output_script_hash(jet::current_index())) + )); +} \ No newline at end of file