From a4feb56a6ca5a0d31fae7b967fedaeb3f351a22c Mon Sep 17 00:00:00 2001 From: LesterEvSe Date: Thu, 16 Jul 2026 16:37:53 +0300 Subject: [PATCH] wip: add smt storage with tests --- simf/lib/storage/smt_storage.simf | 99 +++++++ tests/storage/build_witness.rs | 137 ++++++++++ tests/storage/mod.rs | 435 ++++++++++++++++++++++++++++++ tests/storage/smt.rs | 235 ++++++++++++++++ tests/storage_test.rs | 1 + 5 files changed, 907 insertions(+) create mode 100644 simf/lib/storage/smt_storage.simf create mode 100644 tests/storage/build_witness.rs create mode 100644 tests/storage/mod.rs create mode 100644 tests/storage/smt.rs create mode 100644 tests/storage_test.rs diff --git a/simf/lib/storage/smt_storage.simf b/simf/lib/storage/smt_storage.simf new file mode 100644 index 0000000..4123d30 --- /dev/null +++ b/simf/lib/storage/smt_storage.simf @@ -0,0 +1,99 @@ +/* + * 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_ctx: (u256, Ctx8)) -> (u256, Ctx8) { + let (hash, is_right): (u256, bool) = dbg!(elem); + let (prev_hash, ctx_node): (u256, Ctx8) = prev_hash_ctx; + let ctx: Ctx8 = ctx_node; + + 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), ctx_node) +} + +// Use tag: SMT/1.0/leaf +fn tapleaf_init() -> Ctx8 { + let doubled_tag_hash: [u8; 64] = [ + 0xb6, 0xc5, 0x7e, 0x7a, 0xa0, 0x8c, 0xd9, 0xe8, 0xc2, 0xfb, 0x65, 0xee, 0x31, 0x6d, 0x2f, 0xd9, + 0xa2, 0x4b, 0xf4, 0xf4, 0x00, 0xc8, 0x6d, 0xe5, 0x5d, 0xcd, 0x8c, 0x4d, 0x38, 0x3d, 0x99, 0xd8, + 0xb6, 0xc5, 0x7e, 0x7a, 0xa0, 0x8c, 0xd9, 0xe8, 0xc2, 0xfb, 0x65, 0xee, 0x31, 0x6d, 0x2f, 0xd9, + 0xa2, 0x4b, 0xf4, 0xf4, 0x00, 0xc8, 0x6d, 0xe5, 0x5d, 0xcd, 0x8c, 0x4d, 0x38, 0x3d, 0x99, 0xd8, + ]; + + let ctx: Ctx8 = jet::sha_256_ctx_8_init(); + jet::sha_256_ctx_8_add_64(ctx, doubled_tag_hash) +} + +// Use tag: SMT/1.0/node +fn tapnode_init() -> Ctx8 { + let doubled_tag_hash: [u8; 64] = [ + 0x1f, 0xa6, 0xe8, 0x92, 0x6a, 0x5f, 0x09, 0xa8, 0xf7, 0xb2, 0xe1, 0x57, 0x86, 0x68, 0x2f, 0xaa, + 0xdc, 0xac, 0xee, 0x90, 0x95, 0x2f, 0xe4, 0x61, 0x28, 0x5c, 0x77, 0x8c, 0xb9, 0x5e, 0xb7, 0xb2, + 0x1f, 0xa6, 0xe8, 0x92, 0x6a, 0x5f, 0x09, 0xa8, 0xf7, 0xb2, 0xe1, 0x57, 0x86, 0x68, 0x2f, 0xaa, + 0xdc, 0xac, 0xee, 0x90, 0x95, 0x2f, 0xe4, 0x61, 0x28, 0x5c, 0x77, 0x8c, 0xb9, 0x5e, 0xb7, 0xb2, + ]; + + let ctx: Ctx8 = jet::sha_256_ctx_8_init(); + jet::sha_256_ctx_8_add_64(ctx, doubled_tag_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 = tapleaf_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); + + // Optimize node tag calculation + let ctx_node: Ctx8 = tapnode_init(); + let (computed, _): (u256, Ctx8) = array_fold::( + merkle_data, (hash_leaf, ctx_node) + ); + 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())) + )); +} diff --git a/tests/storage/build_witness.rs b/tests/storage/build_witness.rs new file mode 100644 index 0000000..375e929 --- /dev/null +++ b/tests/storage/build_witness.rs @@ -0,0 +1,137 @@ +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; + +/// Precomputed SHA256 midstate for the `SMT/1.0/leaf` domain separation tag. +/// +/// This is the result of double-hashing `b"SMT/1.0/leaf"` into a fresh SHA256 engine, +/// stored in reversed byte order as required by [`sha256::Midstate`]. +/// +/// The `1.0` denotes the protocol version, allowing future SMT upgrades to safely +/// alter hashing logic without risking backward-compatibility hash collisions. +#[rustfmt::skip] +pub const LEAF_MIDSTATE: [u8; 32] = [ + 0x64, 0x59, 0x25, 0x75, 0x94, 0x38, 0xdd, 0x53, 0xbc, 0x61, 0x54, 0x2a, 0x12, 0x77, 0x85, 0x3e, + 0x62, 0x7e, 0x14, 0x80, 0xd7, 0xad, 0xf0, 0x5b, 0x78, 0x09, 0x93, 0x62, 0x70, 0x98, 0x24, 0xbf, +]; + +/// Precomputed SHA256 midstate for the `SMT/1.0/node` domain separation tag. +/// +/// This is the result of double-hashing `b"SMT/1.0/node"` into a fresh SHA256 engine, +/// stored in reversed byte order as required by [`sha256::Midstate`]. +/// +/// The `1.0` denotes the protocol version, allowing future SMT upgrades to safely +/// alter hashing logic without risking backward-compatibility hash collisions +#[rustfmt::skip] +pub const NODE_MIDSTATE: [u8; 32] = [ + 0xe3, 0x4c, 0x71, 0xa4, 0x67, 0x56, 0x64, 0xaa, 0x91, 0x98, 0x37, 0x94, 0xe6, 0x30, 0x0c, 0xa0, + 0xe2, 0x0d, 0xbb, 0xb2, 0xe6, 0x69, 0x03, 0x3f, 0x2b, 0x10, 0x20, 0xc9, 0x68, 0x55, 0x06, 0xe6, +]; + +#[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 const 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/tests/storage/mod.rs b/tests/storage/mod.rs new file mode 100644 index 0000000..8997f81 --- /dev/null +++ b/tests/storage/mod.rs @@ -0,0 +1,435 @@ +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}; +use simplicityhl::tracker::TrackerLogLevel; +use simplicityhl::{Arguments, CompiledProgram, TemplateProgram}; +use wallet_abi::simplicity_leaf_version; +use wallet_abi::{Network, ProgramError, run_program}; + +mod build_witness; +mod smt; + +pub use build_witness::{DEPTH, SMTWitness, build_smt_storage_witness, u256}; +pub use smt::SparseMerkleTree; + +use crate::smt_storage::build_witness::LEAF_MIDSTATE; +use crate::smt_storage::build_witness::NODE_MIDSTATE; + +#[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()), + simplicity_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") +} + +/// Create a SHA256 context, initialized with a specific `LEAF_TAG` constant and data +/// +/// Based on the C implementation of the `tapdata_init` jet: +/// +#[must_use] +pub fn tap_leaf_hash(data: &[u8]) -> sha256::Hash { + let mut eng = + sha256::HashEngine::from_midstate(sha256::Midstate::from_byte_array(LEAF_MIDSTATE), 64); + eng.input(data); + sha256::Hash::from_engine(eng) +} + +/// Returns a pre-initialized SHA256 engine with the double `NODE_TAG` already processed. +#[must_use] +pub fn node_hash_engine() -> sha256::HashEngine { + sha256::HashEngine::from_midstate(sha256::Midstate::from_byte_array(NODE_MIDSTATE), 64) +} + +/// 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 +/// +/// This implementation employs a dual-layered defense mechanism against **second +/// preimage attacks** (specifically, Merkle substitution attacks) using explicit +/// domain separation and path binding: +/// +/// 1. **Domain Separation (`LEAF_TAG` vs `NODE_TAG`)**: An attacker might try to present +/// an internal node as a leaf, or vice versa. By utilizing distinct tags for leaves +/// and branches, the tree guarantees that a branch hash can never be mathematically +/// parsed as a leaf hash. +/// 2. **Path Binding**: The `raw_path` (bit representation of the path) is included +/// in the initial hash of the leaf alongside the `leaf` data. This strictly binds +/// the data to its exact 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 raw_path: [bool; DEPTH] = std::array::from_fn(|i| path[i].1); + let mut tapdata_input = Vec::with_capacity(leaf.len() + 1); + tapdata_input.extend_from_slice(leaf); + tapdata_input.push(get_path_bits(&raw_path, false)); + let mut current_hash = tap_leaf_hash(&tapdata_input); + + for (hash, is_right_direction) in path { + let mut eng = node_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: Network, + input_index: usize, +) -> Result>, ProgramError> { + let genesis_hash = network.genesis_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: Network, + 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 tap_leaf_midstate() { + let leaf_tag = b"SMT/1.0/leaf"; + let tag = sha256::Hash::hash(leaf_tag); + + let mut eng = sha256::Hash::engine(); + eng.input(tag.as_byte_array()); + eng.input(tag.as_byte_array()); + dbg!(eng.midstate()); + assert_eq!( + eng.midstate(), + sha256::Midstate::from_byte_array(LEAF_MIDSTATE) + ); + } + + #[test] + fn node_midstate() { + let node_tag = b"SMT/1.0/node"; + let tag = sha256::Hash::hash(node_tag); + + let mut eng = sha256::Hash::engine(); + eng.input(tag.as_byte_array()); + eng.input(tag.as_byte_array()); + assert_eq!( + eng.midstate(), + sha256::Midstate::from_byte_array(NODE_MIDSTATE) + ); + } + + #[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, + 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/tests/storage/smt.rs b/tests/storage/smt.rs new file mode 100644 index 0000000..002f4ee --- /dev/null +++ b/tests/storage/smt.rs @@ -0,0 +1,235 @@ +use simplicityhl::simplicity::elements::hashes::HashEngine as _; +use simplicityhl::simplicity::hashes::{Hash, sha256}; + +use crate::smt_storage::{node_hash_engine, tap_leaf_hash}; +use crate::state_management::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 const fn get_hash(&self) -> u256 { + match self { + Self::Leaf { leaf_hash } => *leaf_hash, + Self::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(LEAF_TAG)`, where `LEAF_TAG` is an internal constant. +/// * *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 zero = [0u8; 32]; + precalculate_hashes[0] = *tap_leaf_hash(&zero).as_byte_array(); + + for i in 1..DEPTH { + let mut eng = node_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 using the globally cached double-tag midstate. + fn calculate_hash(left: &TreeNode, right: &TreeNode) -> u256 { + let mut eng = node_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 mut tapdata_input = Vec::with_capacity(leaf.len() + 1); + tapdata_input.extend_from_slice(leaf); + tapdata_input.push(get_path_bits(path, true)); + let leaf_hash = tap_leaf_hash(&tapdata_input); + **root = TreeNode::Leaf { + leaf_hash: *leaf_hash.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/tests/storage_test.rs b/tests/storage_test.rs new file mode 100644 index 0000000..60842b5 --- /dev/null +++ b/tests/storage_test.rs @@ -0,0 +1 @@ +mod storage; \ No newline at end of file