feat(state-tree)!: implement verify for Merkle in/exclusion proofs (#…

…1216)

Description
---
feat(state-tree): implement verify for Merkle in/exclusion proofs
fix(consensus)!: include substate version in JMT node hashes

Motivation and Context
---
Implement inclusion and exclusion proofs for JMT. This is useful for
command inclusion proofs e.g. eviction proof as well as inclusion proofs
for substates.

Usage:

```rust
let (key, proof_value, proof) = tree.get_proof(current_version, &substate_id).unwrap();
let value_hash = hash_substate(substate);
proof.verify_inclusion(&root_hash, &key, &value_hash).unwrap();
proof.verify_exclusion(&root_hash, &some_non_existing_key).unwrap();
```

How Has This Been Tested?
---
Unit tests, manually

What process can a PR reviewer use to test or verify this change?
---


Breaking Changes
---

- [ ] None
- [x] Requires data directory to be deleted
- [ ] Other - Please specify

BREAKING CHANGE: substate changes for JMT state tree include the
substate version in the key hash

applications/tari_swarm_daemon/src/main.rs

@@ -197,7 +197,8 @@ async fn start(cli: &Cli) -> anyhow::Result<()> {
     let lock_file = config.base_dir.join("tari_swarm.pid");
     let _pid = lockfile::Lockfile::create(&lock_file).with_context(|| {
         anyhow!(
-            "Failed to acquire lockfile at {}. Is another instance already running?",
+            "Failed to acquire lockfile at '{}'. Is another instance already running? If not, swarm may have \
+             previously crashed and you may remove the lockfile.",
             lock_file.display()
         )
     })?;

dan_layer/rpc_state_sync/src/manager.rs

@@ -207,11 +207,11 @@ where TConsensusSpec: ConsensusSpec<Addr = PeerAddress>
 
                     let change = match &transition.update {
                         SubstateUpdate::Create(create) => SubstateTreeChange::Up {
-                            id: create.substate.substate_id.clone(),
+                            id: create.substate.to_versioned_substate_id(),
                             value_hash: hash_substate(&create.substate.substate_value, create.substate.version),
                         },
                         SubstateUpdate::Destroy(destroy) => SubstateTreeChange::Down {
-                            id: destroy.substate_id.clone(),
+                            id: destroy.to_versioned_substate_id()
                         },
                     };
 

dan_layer/state_tree/src/bit_iter.rs

@@ -0,0 +1,54 @@
+//   Copyright 2024 The Tari Project
+//   SPDX-License-Identifier: BSD-3-Clause
+
+use std::ops::Range;
+
+/// An iterator over a hash value that generates one bit for each iteration.
+pub struct BitIterator<'a> {
+    /// The reference to the bytes that represent the `HashValue`.
+    bytes: &'a [u8],
+    pos: Range<usize>,
+    // invariant hash_bytes.len() == HashValue::LENGTH;
+    // invariant pos.end == hash_bytes.len() * 8;
+}
+
+impl<'a> BitIterator<'a> {
+    /// Constructs a new `BitIterator` using given `HashValue`.
+    pub fn new(bytes: &'a [u8]) -> Self {
+        BitIterator {
+            bytes,
+            pos: 0..bytes.len() * 8,
+        }
+    }
+
+    /// Returns the `index`-th bit in the bytes.
+    fn get_bit(&self, index: usize) -> bool {
+        // MIRAI annotations - important?
+        // assume!(index < self.pos.end); // assumed precondition
+        // assume!(self.hash_bytes.len() == 32); // invariant
+        // assume!(self.pos.end == self.hash_bytes.len() * 8); // invariant
+        let pos = index / 8;
+        let bit = 7 - index % 8;
+        (self.bytes[pos] >> bit) & 1 != 0
+    }
+}
+
+impl<'a> Iterator for BitIterator<'a> {
+    type Item = bool;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.pos.next().map(|x| self.get_bit(x))
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.pos.size_hint()
+    }
+}
+
+impl<'a> DoubleEndedIterator for BitIterator<'a> {
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.pos.next_back().map(|x| self.get_bit(x))
+    }
+}
+
+impl<'a> ExactSizeIterator for BitIterator<'a> {}

dan_layer/state_tree/src/jellyfish/error.rs

@@ -0,0 +1,30 @@
+//   Copyright 2024 The Tari Project
+//   SPDX-License-Identifier: BSD-3-Clause
+
+use crate::{Hash, LeafKey};
+
+#[derive(Debug, thiserror::Error)]
+pub enum JmtProofVerifyError {
+    #[error("Sparse Merkle Tree proof has more than 256 ({num_siblings}) siblings.")]
+    TooManySiblings { num_siblings: usize },
+    #[error("Keys do not match. Key in proof: {actual_key}. Expected key: {expected_key}.")]
+    KeyMismatch { actual_key: LeafKey, expected_key: LeafKey },
+    #[error("Value hashes do not match. Value hash in proof: {actual}. Expected value hash: {expected}.")]
+    ValueMismatch { actual: Hash, expected: Hash },
+    #[error("Expected inclusion proof. Found non-inclusion proof.")]
+    ExpectedInclusionProof,
+    #[error("Expected non-inclusion proof, but key exists in proof.")]
+    ExpectedNonInclusionProof,
+    #[error(
+        "Key would not have ended up in the subtree where the provided key in proof is the only existing  key, if it \
+         existed. So this is not a valid non-inclusion proof."
+    )]
+    InvalidNonInclusionProof,
+    #[error(
+        "Root hashes do not match. Actual root hash: {actual_root_hash}. Expected root hash: {expected_root_hash}."
+    )]
+    RootHashMismatch {
+        actual_root_hash: Hash,
+        expected_root_hash: Hash,
+    },
+}

dan_layer/state_tree/src/jellyfish/mod.rs

@@ -10,6 +10,7 @@ pub use tree::*;
 mod types;
 pub use types::*;
 
+mod error;
 mod store;
 
 pub use store::*;

dan_layer/state_tree/src/jellyfish/tree.rs

@@ -375,7 +375,7 @@ impl<'a, R: 'a + TreeStoreReader<P>, P: Clone> JellyfishMerkleTree<'a, R, P> {
 
         if kvs.len() == 1 && kvs[0].0 == existing_leaf_key {
             if let (key, Some((value_hash, payload))) = kvs[0] {
-                let new_leaf_node = Node::new_leaf(key.clone(), *value_hash, payload.clone(), version);
+                let new_leaf_node = Node::new_leaf(*key, *value_hash, payload.clone(), version);
                 Ok(Some(new_leaf_node))
             } else {
                 Ok(None)
@@ -454,7 +454,7 @@ impl<'a, R: 'a + TreeStoreReader<P>, P: Clone> JellyfishMerkleTree<'a, R, P> {
     ) -> Result<Option<Node<P>>, JmtStorageError> {
         if kvs.len() == 1 {
             if let (key, Some((value_hash, payload))) = kvs[0] {
-                let new_leaf_node = Node::new_leaf(key.clone(), *value_hash, payload.clone(), version);
+                let new_leaf_node = Node::new_leaf(*key, *value_hash, payload.clone(), version);
                 Ok(Some(new_leaf_node))
             } else {
                 Ok(None)

dan_layer/state_tree/src/jellyfish/types.rs

@@ -81,7 +81,7 @@
 // Copyright (c) Aptos
 // SPDX-License-Identifier: Apache-2.0
 
-use std::{fmt, io, ops::Range};
+use std::{fmt, fmt::Display, io, ops::Range};
 
 use blake2::{digest::consts::U32, Blake2b};
 use indexmap::IndexMap;
@@ -94,7 +94,10 @@ use tari_crypto::{
 use tari_dan_common_types::optional::IsNotFoundError;
 use tari_engine_types::serde_with;
 
-use crate::jellyfish::store::TreeStoreReader;
+use crate::{
+    bit_iter::BitIterator,
+    jellyfish::{error::JmtProofVerifyError, store::TreeStoreReader},
+};
 
 pub type Hash = tari_common_types::types::FixedHash;
 
@@ -146,7 +149,7 @@ pub struct SparseMerkleProofExt {
 
 impl SparseMerkleProofExt {
     /// Constructs a new `SparseMerkleProofExt` using leaf and a list of sibling nodes.
-    pub fn new(leaf: Option<SparseMerkleLeafNode>, siblings: Vec<NodeInProof>) -> Self {
+    pub(crate) fn new(leaf: Option<SparseMerkleLeafNode>, siblings: Vec<NodeInProof>) -> Self {
         Self { leaf, siblings }
     }
 
@@ -159,6 +162,106 @@ impl SparseMerkleProofExt {
     pub fn siblings(&self) -> &[NodeInProof] {
         &self.siblings
     }
+
+    /// Verifies an element whose key is `element_key` and value is `element_value` exists in the Sparse Merkle Tree
+    /// using the provided proof
+    pub fn verify_inclusion(
+        &self,
+        expected_root_hash: &Hash,
+        element_key: &LeafKey,
+        element_value_hash: &Hash,
+    ) -> Result<(), JmtProofVerifyError> {
+        self.verify(expected_root_hash, element_key, Some(element_value_hash))
+    }
+
+    /// Verifies the proof is a valid non-inclusion proof that shows this key doesn't exist in the tree.
+    pub fn verify_exclusion(
+        &self,
+        expected_root_hash: &Hash,
+        element_key: &LeafKey,
+    ) -> Result<(), JmtProofVerifyError> {
+        self.verify(expected_root_hash, element_key, None)
+    }
+
+    /// If `element_value` is present, verifies an element whose key is `element_key` and value is
+    /// `element_value` exists in the Sparse Merkle Tree using the provided proof. Otherwise,
+    /// verifies the proof is a valid non-inclusion proof that shows this key doesn't exist in the
+    /// tree.
+    fn verify(
+        &self,
+        expected_root_hash: &Hash,
+        element_key: &LeafKey,
+        element_value: Option<&Hash>,
+    ) -> Result<(), JmtProofVerifyError> {
+        if self.siblings.len() > 256 {
+            return Err(JmtProofVerifyError::TooManySiblings {
+                num_siblings: self.siblings.len(),
+            });
+        }
+
+        match (element_value, &self.leaf) {
+            (Some(value_hash), Some(leaf)) => {
+                // This is an inclusion proof, so the key and value hash provided in the proof
+                // should match element_key and element_value_hash. `siblings` should prove the
+                // route from the leaf node to the root.
+                if element_key != leaf.key() {
+                    return Err(JmtProofVerifyError::KeyMismatch {
+                        actual_key: *leaf.key(),
+                        expected_key: *element_key,
+                    });
+                }
+                if *value_hash != leaf.value_hash {
+                    return Err(JmtProofVerifyError::ValueMismatch {
+                        actual: leaf.value_hash,
+                        expected: *value_hash,
+                    });
+                }
+            },
+            (Some(_), None) => return Err(JmtProofVerifyError::ExpectedInclusionProof),
+            (None, Some(leaf)) => {
+                // This is a non-inclusion proof. The proof intends to show that if a leaf node
+                // representing `element_key` is inserted, it will break a currently existing leaf
+                // node represented by `proof_key` into a branch. `siblings` should prove the
+                // route from that leaf node to the root.
+                if element_key == leaf.key() {
+                    return Err(JmtProofVerifyError::ExpectedNonInclusionProof);
+                }
+                if element_key.common_prefix_bits_len(leaf.key()) < self.siblings.len() {
+                    return Err(JmtProofVerifyError::InvalidNonInclusionProof);
+                }
+            },
+            (None, None) => {
+                // This is a non-inclusion proof. The proof intends to show that if a leaf node
+                // representing `element_key` is inserted, it will show up at a currently empty
+                // position. `sibling` should prove the route from this empty position to the root.
+            },
+        }
+
+        let current_hash = self
+            .leaf
+            .clone()
+            .map_or(SPARSE_MERKLE_PLACEHOLDER_HASH, |leaf| leaf.hash());
+        let actual_root_hash = self
+            .siblings
+            .iter()
+            .zip(element_key.iter_bits().rev().skip(256 - self.siblings.len()))
+            .fold(current_hash, |hash, (sibling_node, bit)| {
+                if bit {
+                    SparseMerkleInternalNode::new(sibling_node.hash(), hash).hash()
+                } else {
+                    SparseMerkleInternalNode::new(hash, sibling_node.hash()).hash()
+                }
+            });
+
+        if actual_root_hash != *expected_root_hash {
+            return Err(JmtProofVerifyError::RootHashMismatch {
+                actual_root_hash,
+                expected_root_hash: *expected_root_hash,
+            });
+        }
+
+        Ok(())
+    }
 }
 
 impl From<SparseMerkleProofExt> for SparseMerkleProof {
@@ -552,8 +655,8 @@ impl NibblePath {
     }
 
     /// Get a bit iterator iterates over the whole nibble path.
-    pub fn bits(&self) -> BitIterator {
-        BitIterator {
+    pub fn bits(&self) -> NibbleBitIterator {
+        NibbleBitIterator {
             nibble_path: self,
             pos: (0..self.num_nibbles * 4),
         }
@@ -599,12 +702,12 @@ pub trait Peekable: Iterator {
 }
 
 /// BitIterator iterates a nibble path by bit.
-pub struct BitIterator<'a> {
+pub struct NibbleBitIterator<'a> {
     nibble_path: &'a NibblePath,
     pos: Range<usize>,
 }
 
-impl<'a> Peekable for BitIterator<'a> {
+impl<'a> Peekable for NibbleBitIterator<'a> {
     /// Returns the `next()` value without advancing the iterator.
     fn peek(&self) -> Option<Self::Item> {
         if self.pos.start < self.pos.end {
@@ -616,7 +719,7 @@ impl<'a> Peekable for BitIterator<'a> {
 }
 
 /// BitIterator spits out a boolean each time. True/false denotes 1/0.
-impl<'a> Iterator for BitIterator<'a> {
+impl<'a> Iterator for NibbleBitIterator<'a> {
     type Item = bool;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -625,7 +728,7 @@ impl<'a> Iterator for BitIterator<'a> {
 }
 
 /// Support iterating bits in reversed order.
-impl<'a> DoubleEndedIterator for BitIterator<'a> {
+impl<'a> DoubleEndedIterator for NibbleBitIterator<'a> {
     fn next_back(&mut self) -> Option<Self::Item> {
         self.pos.next_back().map(|i| self.nibble_path.get_bit(i))
     }
@@ -690,8 +793,8 @@ impl<'a> NibbleIterator<'a> {
     }
 
     /// Turn it into a `BitIterator`.
-    pub fn bits(&self) -> BitIterator<'a> {
-        BitIterator {
+    pub fn bits(&self) -> NibbleBitIterator<'a> {
+        NibbleBitIterator {
             nibble_path: self.nibble_path,
             pos: (self.pos.start * 4..self.pos.end * 4),
         }
@@ -717,7 +820,7 @@ impl<'a> NibbleIterator<'a> {
 
 // INITIAL-MODIFICATION: We will use this type (instead of `Hash`) to allow for arbitrary key length
 /// A leaf key (i.e. a complete nibble path).
-#[derive(Clone, Debug, Hash, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
+#[derive(Clone, Debug, Copy, Hash, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
 pub struct LeafKey {
     /// The underlying bytes.
     /// All leaf keys of the same tree must be of the same length - otherwise the tree's behavior
@@ -737,6 +840,23 @@ impl LeafKey {
     pub fn as_ref(&self) -> LeafKeyRef<'_> {
         LeafKeyRef::new(self.bytes.as_slice())
     }
+
+    pub fn iter_bits(&self) -> BitIterator<'_> {
+        BitIterator::new(self.bytes.as_slice())
+    }
+
+    pub fn common_prefix_bits_len(&self, other: &LeafKey) -> usize {
+        self.iter_bits()
+            .zip(other.iter_bits())
+            .take_while(|(x, y)| x == y)
+            .count()
+    }
+}
+
+impl Display for LeafKey {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.bytes.fmt(f)
+    }
 }
 
 // INITIAL-MODIFICATION: We will use this type (instead of `Hash`) to allow for arbitrary key length

dan_layer/state_tree/src/key_mapper.rs

@@ -1,7 +1,7 @@
 //   Copyright 2024 The Tari Project
 //   SPDX-License-Identifier: BSD-3-Clause
 
-use tari_engine_types::substate::SubstateId;
+use tari_dan_common_types::VersionedSubstateId;
 
 use crate::{jellyfish::LeafKey, Hash};
 
@@ -11,8 +11,8 @@ pub trait DbKeyMapper<T> {
 
 pub struct SpreadPrefixKeyMapper;
 
-impl DbKeyMapper<SubstateId> for SpreadPrefixKeyMapper {
-    fn map_to_leaf_key(id: &SubstateId) -> LeafKey {
+impl DbKeyMapper<VersionedSubstateId> for SpreadPrefixKeyMapper {
+    fn map_to_leaf_key(id: &VersionedSubstateId) -> LeafKey {
         let hash = crate::jellyfish::jmt_node_hash(id);
         LeafKey::new(hash)
     }

dan_layer/state_tree/src/lib.rs

@@ -12,5 +12,7 @@ pub mod memory_store;
 mod staged_store;
 pub use staged_store::*;
 
+mod bit_iter;
 mod tree;
+
 pub use tree::*;