✨ Additional protection against a second preimage attack

crates/contracts/src/state_management/smt_storage/build_witness.rs

@@ -14,6 +14,24 @@ pub type u256 = [u8; 32];
 /// and cannot dynamically resolve array lengths using `param::LEN`.
 pub const DEPTH: usize = 8;
 
+/// Domain separation tag for Sparse Merkle Tree (SMT) leaves.
+///
+/// This constant is double-hashed and prefixed to the leaf data (and its path)
+/// before the final leaf hash is computed.
+///
+/// The `1.0` denotes the protocol version, allowing future SMT upgrades to safely
+/// alter hashing logic without risking backward-compatibility hash collisions
+pub const LEAF_TAG: &[u8; 12] = b"SMT/1.0/leaf";
+
+/// Domain separation tag for Sparse Merkle Tree (SMT) internal nodes (branches).
+///
+/// This constant is double-hashed and prefixed to the concatenation of two child hashes
+/// to compute their parent node's hash.
+///
+/// The `1.0` denotes the protocol version, allowing future SMT upgrades to safely
+/// alter hashing logic without risking backward-compatibility hash collisions
+pub const NODE_TAG: &[u8; 12] = b"SMT/1.0/node";
+
 #[derive(Debug, Clone, bincode::Encode, bincode::Decode, PartialEq, Eq)]
 pub struct SMTWitness {
     /// The internal public key used for Taproot tweaking.

crates/contracts/src/state_management/smt_storage/mod.rs

@@ -1,4 +1,5 @@
 use std::sync::Arc;
+use std::sync::OnceLock;
 
 use simplicityhl::elements::TxInWitness;
 use simplicityhl::elements::TxOut;
@@ -13,15 +14,18 @@ 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};
-use wallet_abi::{simplicity_leaf_version, tap_data_hash};
 
 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_TAG;
+use crate::smt_storage::build_witness::NODE_TAG;
+
 #[must_use]
 pub fn get_path_bits(path: &[bool], reverse: bool) -> u8 {
     let mut path_bits = 0u8;
@@ -91,20 +95,63 @@ pub fn control_block(cmr: Cmr, spend_info: &TaprootSpendInfo) -> ControlBlock {
         .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:
+/// <https://github.com/BlockstreamResearch/simplicity/blob/d190505509f4c04b1b9193c6739515f9faa18aac/C/jets.c#L1408>
+#[must_use]
+pub fn tap_leaf_hash(data: &[u8]) -> sha256::Hash {
+    static ENGINE: OnceLock<sha256::HashEngine> = OnceLock::new();
+
+    let mut eng = ENGINE
+        .get_or_init(|| {
+            let tag = sha256::Hash::hash(LEAF_TAG);
+            let mut initial_eng = sha256::Hash::engine();
+            initial_eng.input(tag.as_byte_array());
+            initial_eng.input(tag.as_byte_array());
+            initial_eng
+        })
+        .clone();
+
+    eng.input(data);
+
+    sha256::Hash::from_engine(eng)
+}
+
+/// Returns a pre-initialized SHA256 engine with the double `NODE_TAG` already processed.
+/// This runs the tag hashing logic exactly once during the entire program's lifecycle.
+pub fn node_hash_engine() -> sha256::HashEngine {
+    static ENGINE: OnceLock<sha256::HashEngine> = OnceLock::new();
+
+    ENGINE
+        .get_or_init(|| {
+            let mut eng = sha256::Hash::engine();
+            let node_tag = sha256::Hash::hash(NODE_TAG);
+            eng.input(node_tag.as_byte_array());
+            eng.input(node_tag.as_byte_array());
+            eng
+        })
+        .clone()
+}
+
 /// 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 implementation employs a dual-layered defense mechanism against **second
+/// preimage attacks** (specifically, Merkle substitution attacks) using explicit
+/// domain separation and path binding:
 ///
-/// 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.
+/// 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
@@ -125,10 +172,10 @@ pub fn compute_tapdata_tagged_hash_of_the_state(
     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_data_hash(&tapdata_input);
+    let mut current_hash = tap_leaf_hash(&tapdata_input);
 
     for (hash, is_right_direction) in path {
-        let mut eng = sha256::Hash::engine();
+        let mut eng = node_hash_engine();
 
         if *is_right_direction {
             eng.input(hash);

crates/contracts/src/state_management/smt_storage/smt.rs

@@ -1,7 +1,7 @@
 use simplicityhl::simplicity::elements::hashes::HashEngine as _;
 use simplicityhl::simplicity::hashes::{Hash, sha256};
-use wallet_abi::tap_data_hash;
 
+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};
@@ -70,7 +70,7 @@ impl TreeNode {
 ///      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")`.
+///    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.
 ///
@@ -93,13 +93,12 @@ impl SparseMerkleTree {
     #[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();
+        precalculate_hashes[0] = *tap_leaf_hash(&zero).as_byte_array();
 
         for i in 1..DEPTH {
-            let mut eng = sha256::Hash::engine();
+            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();
@@ -113,11 +112,13 @@ impl SparseMerkleTree {
         }
     }
 
-    /// Computes parent hash: `SHA256(left_child_hash || right_child_hash)`.
+    /// Computes parent hash using the globally cached double-tag midstate.
     fn calculate_hash(left: &TreeNode, right: &TreeNode) -> u256 {
-        let mut eng = sha256::Hash::engine();
+        let mut eng = node_hash_engine();
+
         eng.input(&left.get_hash());
         eng.input(&right.get_hash());
+
         *sha256::Hash::from_engine(eng).as_byte_array()
     }
 
@@ -138,7 +139,7 @@ impl SparseMerkleTree {
             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_data_hash(&tapdata_input);
+            let leaf_hash = tap_leaf_hash(&tapdata_input);
             **root = TreeNode::Leaf {
                 leaf_hash: *leaf_hash.as_byte_array(),
             };

crates/contracts/src/state_management/smt_storage/source_simf/smt_storage.simf

@@ -4,9 +4,10 @@
  * 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 {
+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 ctx: Ctx8 = jet::sha_256_ctx_8_init();
+    let (prev_hash, ctx_node): (u256, Ctx8) = prev_hash_ctx;
+    let ctx: Ctx8 = ctx_node;
 
     let new_hash: Ctx8 = match is_right {
         true => {
@@ -19,18 +20,46 @@ fn hash_array_tr_storage_with_update(elem: (u256, bool), prev_hash: u256) -> u25
         }
     };
 
-    jet::sha_256_ctx_8_finalize(new_hash)
+    (jet::sha_256_ctx_8_finalize(new_hash), ctx_node)
+}
+
+// Use tag: SMT/1.0/leaf
+fn tapleaf_init() -> Ctx8 {
+    let ctx: Ctx8 = jet::sha_256_ctx_8_init();
+    let ctx: Ctx8 = jet::sha_256_ctx_8_add_8(ctx, 0x534d542f312e302f);  // "SMT/1.0/"
+    let ctx: Ctx8 = jet::sha_256_ctx_8_add_4(ctx, 0x6c656166);          // "leaf"
+    let tag_hash: u256 = jet::sha_256_ctx_8_finalize(ctx);
+
+    let ctx: Ctx8 = jet::sha_256_ctx_8_init();
+    let ctx: Ctx8 = jet::sha_256_ctx_8_add_32(ctx, tag_hash);
+    jet::sha_256_ctx_8_add_32(ctx, tag_hash)
+}
+
+// Use tag: SMT/1.0/node
+fn tapnode_init() -> Ctx8 {
+    let ctx: Ctx8 = jet::sha_256_ctx_8_init();
+    let ctx: Ctx8 = jet::sha_256_ctx_8_add_8(ctx, 0x534d542f312e302f);  // "SMT/1.0/"
+    let ctx: Ctx8 = jet::sha_256_ctx_8_add_4(ctx, 0x6e6f6465);          // "node"
+    let tag_hash: u256 = jet::sha_256_ctx_8_finalize(ctx);
+
+    let ctx: Ctx8 = jet::sha_256_ctx_8_init();
+    let ctx: Ctx8 = jet::sha_256_ctx_8_add_32(ctx, tag_hash);
+    jet::sha_256_ctx_8_add_32(ctx, 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 = jet::tapdata_init();
+    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);
 
-    let computed: u256 = array_fold::<hash_array_tr_storage_with_update, 8>(merkle_data, hash_leaf);
+    // Optimize node tag calculation
+    let ctx_node: Ctx8 = tapnode_init();
+    let (computed, _): (u256, Ctx8) = array_fold::<hash_array_tr_storage_with_update, 8>(
+        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);