diff --git a/base_layer/mmr/src/sparse_merkle_tree/bit_utils.rs b/base_layer/mmr/src/sparse_merkle_tree/bit_utils.rs index 77620c48aa..4471bd3214 100644 --- a/base_layer/mmr/src/sparse_merkle_tree/bit_utils.rs +++ b/base_layer/mmr/src/sparse_merkle_tree/bit_utils.rs @@ -14,24 +14,9 @@ pub(crate) fn get_bit(data: &[u8], position: usize) -> usize { 0 } -/// Sets the n-th bit in a byte array to the given value. `position` 0 is the most significant bit. -/// -/// # Panics -/// `set_bit` will panic if -/// * `position` is out of bounds -/// * `value` is not 0 or 1 -#[inline] -pub(crate) fn set_bit(data: &mut [u8], position: usize, value: usize) { - match value { - 0 => data[position / 8] &= !(1 << (7 - (position % 8))), - 1 => data[position / 8] |= 1 << (7 - (position % 8)), - _ => panic!("Invalid bit value"), - } -} - /// Given two node keys, this function returns the number of bits that are common to both keys, starting from the most /// significant bit. This function is used to tell you the height at which two node keys would diverge in the sparse -/// merkle tree. For example, key 0110 and 0101 would diverge at height 2, because the first two bits are the same. +/// Merkle& tree. For example, key 0110 and 0101 would diverge at height 2, because the first two bits are the same. #[inline] pub(crate) fn count_common_prefix(a: &NodeKey, b: &NodeKey) -> usize { let mut offset = 0; @@ -58,7 +43,7 @@ pub fn height_key(key: &NodeKey, height: usize) -> NodeKey { let mut result = NodeKey::default(); // Keep the first `height` bits and ignore the rest let key = key.as_slice(); - let bytes = result.as_mut_slice(); + let bytes = result.as_slice_mut(); // First height/8 bytes are the same, so just copy bytes[0..height / 8].copy_from_slice(&key[0..height / 8]); // The height/8th byte is only partially copied, so mask the byte & 11100000, where the number of 1s is @@ -67,21 +52,12 @@ pub fn height_key(key: &NodeKey, height: usize) -> NodeKey { result } -pub fn path_matches_key(key: &NodeKey, path: &[TraverseDirection]) -> bool { - let height = path.len(); - - let prefix = path - .iter() - .enumerate() - .fold(NodeKey::default(), |mut prefix, (i, dir)| { - let bit = match dir { - TraverseDirection::Left => 0, - TraverseDirection::Right => 1, - }; - set_bit(prefix.as_mut_slice(), i, bit); - prefix - }); - count_common_prefix(key, &prefix) >= height +pub const fn bit_to_dir(bit: usize) -> TraverseDirection { + match bit { + 0 => TraverseDirection::Left, + 1 => TraverseDirection::Right, + _ => panic!("Invalid bit"), + } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -116,11 +92,8 @@ pub fn traverse_direction( }); } - match get_bit(child_key.as_slice(), parent_height) { - 0 => Ok(TraverseDirection::Left), - 1 => Ok(TraverseDirection::Right), - _ => unreachable!(), - } + let dir = bit_to_dir(get_bit(child_key.as_slice(), parent_height)); + Ok(dir) } #[cfg(test)] @@ -234,43 +207,6 @@ mod test { assert_eq!(hkey, expected); } - #[test] - fn set_bits() { - let mut val = [0b10101010, 0b01010101]; - set_bit(&mut val, 0, 1); - assert_eq!(val, [0b10101010, 0b01010101]); - set_bit(&mut val, 1, 1); - assert_eq!(val, [0b11101010, 0b01010101]); - set_bit(&mut val, 2, 1); - assert_eq!(val, [0b11101010, 0b01010101]); - set_bit(&mut val, 3, 1); - assert_eq!(val, [0b11111010, 0b01010101]); - set_bit(&mut val, 4, 1); - assert_eq!(val, [0b11111010, 0b01010101]); - set_bit(&mut val, 5, 1); - assert_eq!(val, [0b11111110, 0b01010101]); - set_bit(&mut val, 6, 1); - assert_eq!(val, [0b11111110, 0b01010101]); - set_bit(&mut val, 7, 1); - assert_eq!(val, [0b11111111, 0b01010101]); - set_bit(&mut val, 8, 0); - assert_eq!(val, [0b11111111, 0b01010101]); - set_bit(&mut val, 9, 0); - assert_eq!(val, [0b11111111, 0b00010101]); - set_bit(&mut val, 10, 0); - assert_eq!(val, [0b11111111, 0b00010101]); - set_bit(&mut val, 11, 0); - assert_eq!(val, [0b11111111, 0b00000101]); - set_bit(&mut val, 12, 0); - assert_eq!(val, [0b11111111, 0b00000101]); - set_bit(&mut val, 13, 0); - assert_eq!(val, [0b11111111, 0b00000001]); - set_bit(&mut val, 14, 0); - assert_eq!(val, [0b11111111, 0b00000001]); - set_bit(&mut val, 15, 0); - assert_eq!(val, [0b11111111, 0b00000000]); - } - #[test] fn get_bits() { let val = [0b10101010, 0b10101010, 0b00000000, 0b11111111]; diff --git a/base_layer/mmr/src/sparse_merkle_tree/error.rs b/base_layer/mmr/src/sparse_merkle_tree/error.rs index 6ebdc6f121..ce4791fbf0 100644 --- a/base_layer/mmr/src/sparse_merkle_tree/error.rs +++ b/base_layer/mmr/src/sparse_merkle_tree/error.rs @@ -29,4 +29,8 @@ pub enum SMTError { IllegalKey(String), #[error("The hash for the tree needs to be recalculated before calling this function")] StaleHash, + #[error( + "Cannot construct a proof. Either the key exists for an exclusion proof, or it does not for an inclusion proof" + )] + NonViableProof, } diff --git a/base_layer/mmr/src/sparse_merkle_tree/mod.rs b/base_layer/mmr/src/sparse_merkle_tree/mod.rs index cefad3af9b..5cd6ab7281 100644 --- a/base_layer/mmr/src/sparse_merkle_tree/mod.rs +++ b/base_layer/mmr/src/sparse_merkle_tree/mod.rs @@ -41,7 +41,7 @@ //! r###" │A│ │B│ │D│ │C│ "### //! r###" └─┘ └─┘ └─┘ └─┘ "### //! -//! The merkle root is calculated by hashing nodes in the familiar way. +//! The Merkle& root is calculated by hashing nodes in the familiar way. //! ```rust //! use blake2::Blake2b; //! use digest::consts::U32; @@ -84,5 +84,5 @@ mod tree; pub use error::SMTError; pub use node::{BranchNode, EmptyNode, LeafNode, Node, NodeHash, NodeKey, ValueHash, EMPTY_NODE_HASH}; -pub use proofs::MerkleProof; +pub use proofs::{ExclusionProof, InclusionProof}; pub use tree::{SparseMerkleTree, UpdateResult}; diff --git a/base_layer/mmr/src/sparse_merkle_tree/node.rs b/base_layer/mmr/src/sparse_merkle_tree/node.rs index 81b15d606a..e07985c59b 100644 --- a/base_layer/mmr/src/sparse_merkle_tree/node.rs +++ b/base_layer/mmr/src/sparse_merkle_tree/node.rs @@ -10,38 +10,37 @@ use std::{ use digest::{consts::U32, Digest}; use crate::sparse_merkle_tree::{ - bit_utils::{count_common_prefix, get_bit, height_key, TraverseDirection}, + bit_utils::{bit_to_dir, count_common_prefix, get_bit, height_key, TraverseDirection}, Node::*, SMTError, }; +pub const KEY_LENGTH: usize = 32; + macro_rules! hash_type { ($name: ident) => { /// A wrapper around a 32-byte hash value. Provides convenience functions to display as hex or binary - #[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] - pub struct $name([u8; 32]); + #[derive(Clone, Debug, PartialEq, Eq, PartialOrd)] + pub struct $name([u8; KEY_LENGTH]); + #[allow(clippy::len_without_is_empty)] impl $name { pub fn as_slice(&self) -> &[u8] { &self.0 } - pub fn as_mut_slice(&mut self) -> &mut [u8] { + pub fn as_slice_mut(&mut self) -> &mut [u8] { &mut self.0 } pub fn len(&self) -> usize { self.0.len() } - - pub fn is_empty(&self) -> bool { - self.0.is_empty() - } } impl Default for $name { fn default() -> Self { - Self([0; 32]) + Self([0; KEY_LENGTH]) } } @@ -49,23 +48,23 @@ macro_rules! hash_type { type Error = SMTError; fn try_from(value: &[u8]) -> Result { - if value.len() < 32 { + if value.len() < KEY_LENGTH { return Err(SMTError::ArrayTooShort(value.len())); } - let mut bytes = [0u8; 32]; + let mut bytes = [0u8; KEY_LENGTH]; bytes.copy_from_slice(value); Ok(Self(bytes)) } } - impl From<[u8; 32]> for $name { - fn from(arr: [u8; 32]) -> Self { + impl From<[u8; KEY_LENGTH]> for $name { + fn from(arr: [u8; KEY_LENGTH]) -> Self { Self(arr) } } - impl From<&[u8; 32]> for $name { - fn from(arr: &[u8; 32]) -> Self { + impl From<&[u8; KEY_LENGTH]> for $name { + fn from(arr: &[u8; KEY_LENGTH]) -> Self { Self(arr.clone()) } } @@ -110,7 +109,72 @@ hash_type!(NodeHash); hash_type!(ValueHash); hash_type!(NodeKey); -pub const EMPTY_NODE_HASH: NodeHash = NodeHash([0; 32]); +impl NodeKey { + pub fn as_directions(&self) -> PathIterator { + PathIterator::new(self) + } +} + +pub const EMPTY_NODE_HASH: NodeHash = NodeHash([0; KEY_LENGTH]); + +pub struct PathIterator<'a> { + cursor_front: usize, + // position *after* next bit when going backwards + cursor_back: usize, + key: &'a NodeKey, +} + +impl PathIterator<'_> { + pub fn new(key: &NodeKey) -> PathIterator { + PathIterator { + cursor_front: 0, + // KEY_LENGTH is currently 32 bytes, so this will not overflow + cursor_back: KEY_LENGTH * 8, + key, + } + } +} + +impl<'a> Iterator for PathIterator<'a> { + type Item = TraverseDirection; + + fn next(&mut self) -> Option { + if self.cursor_front >= self.cursor_back { + return None; + } + let bit = get_bit(self.key.as_slice(), self.cursor_front); + self.cursor_front += 1; + Some(bit_to_dir(bit)) + } + + // This must be overridden, otherwise iterator connectors don't work + fn size_hint(&self) -> (usize, Option) { + let len = self.cursor_back.saturating_sub(self.cursor_front); + (len, Some(len)) + } +} + +impl<'a> DoubleEndedIterator for PathIterator<'a> { + fn next_back(&mut self) -> Option { + if self.cursor_front >= self.cursor_back { + return None; + } + self.cursor_back -= 1; + let bit = get_bit(self.key.as_slice(), self.cursor_back); + Some(bit_to_dir(bit)) + } + + fn nth_back(&mut self, n: usize) -> Option { + self.cursor_back -= n; + self.next_back() + } +} + +impl<'a> ExactSizeIterator for PathIterator<'a> { + fn len(&self) -> usize { + self.cursor_back - self.cursor_front + } +} #[derive(Debug)] pub enum Node { @@ -186,15 +250,6 @@ impl Node { _ => Err(SMTError::UnexpectedNodeType), } } - - /// Indicates whether the node is semi-terminal, i.e. whether it is a leaf or empty node, or if a branch, if it is - /// the last branch in the sub-tree. - pub fn is_semi_terminal(&self) -> bool { - match self { - Leaf(_) | Empty(_) => true, - Branch(n) => !n.left.is_branch() && !n.right.is_branch(), - } - } } impl> Node { @@ -289,7 +344,7 @@ impl> LeafNode { .chain_update(key.as_slice()) .chain_update(value.as_slice()) .finalize(); - let mut result = [0; 32]; + let mut result = [0; KEY_LENGTH]; result.copy_from_slice(hash.as_slice()); result.into() } @@ -462,9 +517,10 @@ mod test { use rand::{self, RngCore}; use super::*; + use crate::sparse_merkle_tree::bit_utils::TraverseDirection::{Left, Right}; - fn random_arr() -> [u8; 32] { - let mut result = [0; 32]; + fn random_arr() -> [u8; KEY_LENGTH] { + let mut result = [0; KEY_LENGTH]; rand::thread_rng().fill_bytes(&mut result); result } @@ -501,14 +557,24 @@ mod test { let left = Node::Empty(EmptyNode {}); let right = Node::Leaf(LeafNode::>::new(random_key(), random_value_hash())); let branch = BranchNode::>::new(0, random_key(), left, right); + let exp_msg = "A branch node cannot an empty node and leaf node as children"; // Should not be allowed - since this can be represented as a leaf node - assert!(matches!(branch, Err(SMTError::InvalidBranch(_)))); + assert!(matches!(branch, Err(SMTError::InvalidBranch(msg)) if msg == exp_msg)); let left = Node::Leaf(LeafNode::>::new(random_key(), random_value_hash())); let right = Node::Empty(EmptyNode {}); let branch = BranchNode::>::new(0, random_key(), left, right); // Should not be allowed - since this can be represented as a leaf node - assert!(matches!(branch, Err(SMTError::InvalidBranch(_)))); + assert!(matches!(branch, Err(SMTError::InvalidBranch(msg)) if msg == exp_msg)); + } + + #[test] + fn cannot_create_branch_with_empty_nodes() { + let left = Node::Empty(EmptyNode {}); + let right = Node::Empty(EmptyNode {}); + let branch = BranchNode::>::new(0, random_key(), left, right); + // Should not be allowed - since this can be represented as a leaf node + assert!(matches!(branch, Err(SMTError::InvalidBranch(msg)) if msg == "Both left and right nodes are empty")); } #[test] @@ -527,4 +593,99 @@ mod test { .finalize(); assert_eq!(branch.hash().as_slice(), expected.as_slice()); } + + #[test] + fn path_iterator_default() { + let key = NodeKey::from(&[0; KEY_LENGTH]); + let path = key.as_directions().collect::>(); + assert_eq!(path.len(), 256); + assert_eq!(path, [TraverseDirection::Left; 256]); + } + + #[test] + fn path_iterator_connectors() { + let key = NodeKey::from(&[0; KEY_LENGTH]); + let iter = key.as_directions(); + assert_eq!(iter.len(), 256); + assert_eq!(iter.take(14).len(), 14); + let iter = key.as_directions(); + assert_eq!(iter.rev().take(18).len(), 18); + } + + #[test] + fn path_iterator() { + let mut key = [0u8; KEY_LENGTH]; + key[0] = 0b1101_1011; + let key = NodeKey::from(key); + let dirs = key.as_directions().take(8).collect::>(); + assert_eq!(dirs, [Right, Right, Left, Right, Right, Left, Right, Right]); + } + + #[test] + fn path_iterator_rev_iter() { + let key = NodeKey::default(); + let mut dirs = key.as_directions().skip(256); + assert_eq!(dirs.next(), None); + assert_eq!(dirs.next(), None); + } + + #[test] + fn path_iterator_iter() { + let mut key = [0u8; KEY_LENGTH]; + key[0] = 0b1101_1011; + let key = NodeKey::from(key); + let mut dirs = key.as_directions().skip(255); + assert_eq!(dirs.next(), Some(Left)); + assert_eq!(dirs.next(), None); + } + + #[test] + fn path_iterator_rev() { + let mut key = [0u8; KEY_LENGTH]; + key[0] = 0b0011_1000; + key[31] = 0b1110_0011; + let key = NodeKey::from(key); + let dirs = key.as_directions().rev().take(8).collect::>(); + assert_eq!(dirs, [Right, Right, Left, Left, Left, Right, Right, Right]); + let dirs = key.as_directions().take(8).rev().collect::>(); + assert_eq!(dirs, [Left, Left, Left, Right, Right, Right, Left, Left]); + } + + #[test] + fn hash_type_from_slice() { + let arr = vec![1u8; 32]; + assert!(matches!(NodeKey::try_from(&arr[..3]), Err(SMTError::ArrayTooShort(3)))); + assert!(NodeKey::try_from(&arr[..]).is_ok()); + assert!(matches!( + ValueHash::try_from(&arr[..4]), + Err(SMTError::ArrayTooShort(4)) + )); + assert!(ValueHash::try_from(&arr[..]).is_ok()); + assert!(matches!( + NodeHash::try_from(&arr[..16]), + Err(SMTError::ArrayTooShort(16)) + )); + assert!(NodeHash::try_from(&arr[..]).is_ok()); + } + + #[test] + fn hash_type_display() { + let key = NodeKey::from(&[ + 1u8, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, + 29, 30, 31, 32, + ]); + let bin = "0000000100000010000000110000010000000101000001100000011100001000\ + 0000100100001010000010110000110000001101000011100000111100010000\ + 0001000100010010000100110001010000010101000101100001011100011000\ + 0001100100011010000110110001110000011101000111100001111100100000"; + let lower_hex = "0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20"; + assert_eq!(format!("{key:x}"), lower_hex); + assert_eq!(format!("{key}"), lower_hex); + assert_eq!( + format!("{key:X}"), + "0102030405060708090A0B0C0D0E0F101112131415161718191A1B1C1D1E1F20" + ); + assert_eq!(format!("{key:b}"), bin); + assert_eq!(format!("{key:#}"), bin); + } } diff --git a/base_layer/mmr/src/sparse_merkle_tree/proofs.rs b/base_layer/mmr/src/sparse_merkle_tree/proofs.rs index 2dfe2b2046..43760ad980 100644 --- a/base_layer/mmr/src/sparse_merkle_tree/proofs.rs +++ b/base_layer/mmr/src/sparse_merkle_tree/proofs.rs @@ -6,57 +6,121 @@ use std::marker::PhantomData; use digest::{consts::U32, Digest}; use crate::sparse_merkle_tree::{ - bit_utils::{height_key, path_matches_key, TraverseDirection}, + bit_utils::{height_key, TraverseDirection}, BranchNode, + EmptyNode, LeafNode, NodeHash, NodeKey, SMTError, SparseMerkleTree, ValueHash, - EMPTY_NODE_HASH, }; -pub struct MerkleProof { - path: Vec, +/// An inclusion proof for a key-value pair in a sparse Merkle& tree. +/// +/// Given a sparse Merkle& tree, `tree`, you can create a proof that a certain key-value pair exists by calling +/// [`InclusionProof::from_tree`], for example: +/// +/// ``` +/// # use blake2::Blake2b; +/// # use digest::consts::U32; +/// # use tari_mmr::sparse_merkle_tree::{ExclusionProof, InclusionProof, NodeKey, SparseMerkleTree, ValueHash}; +/// let key = NodeKey::from([64u8; 32]); +/// let value = ValueHash::from([128u8; 32]); +/// +/// let mut tree = SparseMerkleTree::>::default(); +/// tree.upsert(key.clone(), value.clone()).unwrap(); +/// let hash = tree.hash().clone(); +/// +/// let in_proof = InclusionProof::from_tree(&tree, &key, &value).unwrap(); +/// assert!(in_proof.validate(&key, &value, &hash)); +/// ``` +/// +/// If you try to create an inclusion proof that is invalid, such as using the wrong value, or a key that is not in +/// the tree, `from_tree` will return a `NonViableProof` error. +/// +/// ``` +/// # use blake2::Blake2b; +/// # use digest::consts::U32; +/// # use tari_mmr::sparse_merkle_tree::{ExclusionProof, InclusionProof, NodeKey, SparseMerkleTree, ValueHash, SMTError}; +/// let key = NodeKey::from([64u8; 32]); +/// let non_existent_key = NodeKey::from([65u8; 32]); +/// let value = ValueHash::from([128u8; 32]); +/// let wrong_value = ValueHash::from([127u8; 32]); +/// +/// let mut tree = SparseMerkleTree::>::default(); +/// tree.upsert(key.clone(), value.clone()).unwrap(); +/// let root = tree.hash().clone(); +/// let in_proof = InclusionProof::from_tree(&tree, &non_existent_key, &value); +/// assert!(matches!(in_proof, Err(SMTError::NonViableProof))); +/// let in_proof = InclusionProof::from_tree(&tree, &key, &wrong_value); +/// assert!(matches!(in_proof, Err(SMTError::NonViableProof))); +/// ``` +pub struct InclusionProof { siblings: Vec, - key: NodeKey, - value: Option, phantom: std::marker::PhantomData, } -impl> MerkleProof { - pub fn new(path: Vec, siblings: Vec, key: NodeKey, value: Option) -> Self { - Self { - path, - siblings, - key, - value, - phantom: PhantomData::, - } - } +/// An exclusion proof for a key in a sparse Merkle& tree. +/// +/// Given a sparse Merkle& tree, `tree`, you can create a proof that a certain key does *not exist* in the tree by +/// calling [`ExclusionProof::from_tree`]. For example: +/// +/// ``` +/// # use blake2::Blake2b; +/// # use digest::consts::U32; +/// # use tari_mmr::sparse_merkle_tree::{ExclusionProof, InclusionProof, NodeKey, SparseMerkleTree, ValueHash}; +/// let key = NodeKey::from([64u8; 32]); +/// let value = ValueHash::from([128u8; 32]); +/// let non_existent_key = NodeKey::from([65u8; 32]); +/// let mut tree = SparseMerkleTree::>::default(); +/// tree.upsert(key, value).unwrap(); +/// let hash = tree.hash().clone(); +/// let ex_proof = ExclusionProof::from_tree(&tree, &non_existent_key).unwrap(); +/// assert!(ex_proof.validate(&non_existent_key, &hash)); +/// ``` +/// +/// As with [`InclusionProof`], if you try to create an exclusion proof that is invalid, such as using a key that is +/// in the tree, `from_tree` will return a `NonViableProof` error. For example, using the same tree from the last +/// example, +/// ``` +/// # use blake2::Blake2b; +/// # use digest::consts::U32; +/// # use tari_mmr::sparse_merkle_tree::{ExclusionProof, InclusionProof, NodeKey, SparseMerkleTree, ValueHash, SMTError}; +/// # let key = NodeKey::from([64u8; 32]); +/// # let value = ValueHash::from([128u8; 32]); +/// # let non_existent_key = NodeKey::from([65u8; 32]); +/// # let mut tree = SparseMerkleTree::>::default(); +/// # tree.upsert(key.clone(), value).unwrap(); +/// let ex_proof = ExclusionProof::from_tree(&tree, &key); +/// assert!(matches!(ex_proof, Err(SMTError::NonViableProof))); +/// ``` +pub struct ExclusionProof { + siblings: Vec, + // The terminal node of the tree proof, or `None` if the the node is `Empty`. + leaf: Option>, + phantom: std::marker::PhantomData, +} - pub fn from_tree(tree: &SparseMerkleTree, key: &NodeKey) -> Result { - tree.build_proof(key) - } +trait MerkleProofDigest> { + /// Returns an array to the vector of sibling hashes along the path to the key's leaf node for this proof. + fn siblings(&self) -> &[NodeHash]; - fn calculate_root_hash(&self) -> NodeHash { - let node_hash = match self.value.as_ref() { - Some(v) => LeafNode::::hash_value(&self.key, v), - None => EMPTY_NODE_HASH, - }; - let n = self.siblings.len(); - let hash = self.siblings.iter().zip(self.path.iter()).rev().enumerate().fold( - node_hash, + /// Calculate the Merkle& tree root for this proof, given the key and value hash. + fn calculate_root_hash(&self, key: &NodeKey, leaf_hash: NodeHash) -> NodeHash { + let n = self.siblings().len(); + let dirs = key.as_directions().take(n); + let hash = self.siblings().iter().zip(dirs).rev().enumerate().fold( + leaf_hash, |current, (i, (sibling_hash, direction))| { let height = n - i - 1; - match direction { TraverseDirection::Left => { - BranchNode::::branch_hash(height, &height_key(&self.key, height), ¤t, sibling_hash) + BranchNode::::branch_hash(height, &height_key(key, height), ¤t, sibling_hash) }, TraverseDirection::Right => { - BranchNode::::branch_hash(height, &height_key(&self.key, height), sibling_hash, ¤t) + BranchNode::::branch_hash(height, &height_key(key, height), sibling_hash, ¤t) }, } }, @@ -65,105 +129,234 @@ impl> MerkleProof { result.copy_from_slice(hash.as_slice()); result.into() } +} - pub fn validate_inclusion_proof( - &self, - expected_key: &NodeKey, - expected_value: &ValueHash, - expected_root: &NodeHash, - ) -> bool { - expected_key == &self.key && - Some(expected_value) == self.value.as_ref() && - expected_root == &self.calculate_root_hash() +impl> InclusionProof { + /// Construct an inclusion proof using the vector of siblings provided. Usually you will not use this method, but + /// will generate the proof using [`InclusionProof::from_tree`] instead. + pub fn new(siblings: Vec) -> Self { + Self { + siblings, + phantom: PhantomData::, + } } - pub fn validate_exclusion_proof(&self, expected_key: &NodeKey, expected_root: &NodeHash) -> bool { - path_matches_key(expected_key, &self.path) && - expected_root == &self.calculate_root_hash() && - (self.value.is_none() || &self.key != expected_key) + /// Generates an inclusion proof for the given key and value hash from the given tree. If the key does not exist in + /// tree, or the key does exist, but the value hash does not match, then `from_tree` will return a + /// `NonViableProof` error. + pub fn from_tree(tree: &SparseMerkleTree, key: &NodeKey, value_hash: &ValueHash) -> Result { + let proof = tree.build_proof_candidate(key)?; + match proof.leaf { + Some(leaf) => { + let node_hash = LeafNode::::hash_value(key, value_hash); + if leaf.hash() != &node_hash { + return Err(SMTError::NonViableProof); + } + }, + None => return Err(SMTError::NonViableProof), + } + Ok(Self::new(proof.siblings)) } - pub fn key(&self) -> &NodeKey { - &self.key + /// Validates the inclusion proof against the given key, value hash and root hash. + /// The function reconstructs the tree using the expected key and value hash, and then calculates the root hash. + /// Validation succeeds if the calculated root hash matches the given root hash. + pub fn validate(&self, expected_key: &NodeKey, expected_value: &ValueHash, expected_root: &NodeHash) -> bool { + // calculate expected leaf node hash + let leaf_hash = LeafNode::::hash_value(expected_key, expected_value); + let calculated_root = self.calculate_root_hash(expected_key, leaf_hash); + calculated_root == *expected_root } +} - pub fn value_hash(&self) -> Option<&ValueHash> { - self.value.as_ref() +impl> MerkleProofDigest for InclusionProof { + fn siblings(&self) -> &[NodeHash] { + &self.siblings } } -#[cfg(test)] -mod test { - use blake2::Blake2b; - use digest::consts::U32; - use rand::{RngCore, SeedableRng}; - - use super::*; +impl> ExclusionProof { + /// Construct an exclusion proof using the vector of siblings and the existing leaf node provided. Usually you will + /// not use this method, but will generate the proof using [`ExclusionProof::from_tree`] instead. + pub fn new(siblings: Vec, leaf: Option>) -> Self { + Self { + siblings, + leaf, + phantom: PhantomData::, + } + } - fn random_arr(n: usize, seed: u64) -> Vec<[u8; 32]> { - let mut rng = rand::rngs::StdRng::seed_from_u64(seed); - (0..n) - .map(|_| { - let mut key = [0u8; 32]; - rng.fill_bytes(&mut key); - key - }) - .collect() + /// Generates an exclusion proof for the given key from the given tree. If the key exists in the tree then + /// `from_tree` will return a `NonViableProof` error. + pub fn from_tree(tree: &SparseMerkleTree, key: &NodeKey) -> Result { + let proof = tree.build_proof_candidate(key)?; + // If the keys match, then we cannot provide an exclusion proof, since the key *is* in the tree + if let Some(leaf) = &proof.leaf { + if leaf.key() == key { + return Err(SMTError::NonViableProof); + } + } + Ok(proof) } - fn random_keys(n: usize, seed: u64) -> Vec { - random_arr(n, seed).into_iter().map(|k| k.into()).collect() + /// Validates the exclusion proof against the given key and root hash. The function reconstructs the tree using the + /// expected key and places the leaf node provided in the proof at the terminal position. It then calculates the + /// root hash. Validation succeeds if the calculated root hash matches the given root hash, and the leaf node is + /// empty, or the existing leaf node has a different key to the expected key. + pub fn validate(&self, expected_key: &NodeKey, expected_root: &NodeHash) -> bool { + let leaf_hash = match &self.leaf { + Some(leaf) => leaf.hash().clone(), + None => (EmptyNode {}).hash().clone(), + }; + let root = self.calculate_root_hash(expected_key, leaf_hash); + // For exclusion proof, roots must match AND existing leaf must be empty, or keys must not match + root == *expected_root && + match &self.leaf { + Some(leaf) => leaf.key() != expected_key, + None => true, + } } +} - fn random_values(n: usize, seed: u64) -> Vec { - random_arr(n, seed).into_iter().map(|k| k.into()).collect() +impl> MerkleProofDigest for ExclusionProof { + fn siblings(&self) -> &[NodeHash] { + &self.siblings } +} + +#[cfg(test)] +mod test { + use blake2::Blake2b; + use digest::consts::U32; + + use super::*; #[test] fn root_proof() { let key = NodeKey::from([64u8; 32]); + let key2 = NodeKey::from([65u8; 32]); let value = ValueHash::from([128u8; 32]); let mut tree = SparseMerkleTree::>::default(); let hash = tree.hash().clone(); - let proof = tree.build_proof(&key).unwrap(); + let in_proof = InclusionProof::from_tree(&tree, &key, &value); + assert!(matches!(in_proof, Err(SMTError::NonViableProof))); + let ex_proof = ExclusionProof::from_tree(&tree, &key).unwrap(); + assert!(ex_proof.validate(&key, &hash)); + + tree.upsert(key.clone(), value.clone()).unwrap(); + let hash2 = tree.hash().clone(); + + let in_proof = InclusionProof::from_tree(&tree, &key, &value).unwrap(); + let ex_proof = ExclusionProof::from_tree(&tree, &key); + assert!(matches!(ex_proof, Err(SMTError::NonViableProof))); + + assert!(in_proof.validate(&key, &value, &hash2)); + // correct key, wrong value + assert!(!in_proof.validate(&key, &ValueHash::from([1u8; 32]), &hash2),); + // incorrect key, correct value + assert!(!in_proof.validate(&key2, &value, &hash2)); + // correct key, wrong hash + assert!(!in_proof.validate(&key, &value, &hash)); - assert!(!proof.validate_inclusion_proof(&key, &value, &hash)); - assert!(proof.validate_exclusion_proof(&key, &hash)); + // exclusion proof assertions + let ex_proof = ExclusionProof::from_tree(&tree, &key2).unwrap(); + assert!(!ex_proof.validate(&key, &hash2)); + assert!(!ex_proof.validate(&key, &hash)); + assert!(!ex_proof.validate(&key2, &hash)); + assert!(ex_proof.validate(&key2, &hash2)); + } + + #[test] + fn non_viable_inclusion_proof() { + let mut tree = SparseMerkleTree::>::default(); + let key = NodeKey::from([64u8; 32]); + let value = ValueHash::from([128u8; 32]); + // Inclusion proof on empty tree + let proof = InclusionProof::from_tree(&tree, &key, &value); + assert!(matches!(proof, Err(SMTError::NonViableProof))); + tree.upsert(key, value.clone()).unwrap(); + // Inclusion proof on non-existent key + let proof = InclusionProof::from_tree(&tree, &NodeKey::from([65u8; 32]), &value); + assert!(matches!(proof, Err(SMTError::NonViableProof))); + // Existing key, wrong value + let proof = InclusionProof::from_tree(&tree, &NodeKey::from([64u8; 32]), &ValueHash::from([0u8; 32])); + assert!(matches!(proof, Err(SMTError::NonViableProof))); + } + #[test] + fn proof_with_stale_hash() { + let mut tree = SparseMerkleTree::>::default(); + tree.upsert(NodeKey::from([64u8; 32]), ValueHash::from([128u8; 32])) + .unwrap(); + tree.upsert(NodeKey::from([155u8; 32]), ValueHash::from([128u8; 32])) + .unwrap(); + let key = NodeKey::from([65u8; 32]); + let value = ValueHash::from([65u8; 32]); tree.upsert(key.clone(), value.clone()).unwrap(); - let hash = tree.hash().clone(); - let proof = tree.build_proof(&key).unwrap(); + let proof = InclusionProof::from_tree(&tree, &key, &value); + assert!(matches!(proof, Err(SMTError::StaleHash))); + } - assert!(proof.validate_inclusion_proof(&key, &value, &hash)); - assert!(!proof.validate_inclusion_proof(&key, &ValueHash::from([1u8; 32]), &hash),); - assert!(!proof.validate_exclusion_proof(&key, &hash)); + #[test] + fn deep_inclusion_proof() { + let key1 = NodeKey::from([64u8; 32]); + let mut key2 = key1.clone(); + key2.as_slice_mut()[31] = 65; + let value = ValueHash::from([128u8; 32]); + let mut tree = SparseMerkleTree::>::default(); + tree.upsert(key1, ValueHash::from([42u8; 32])).unwrap(); + tree.upsert(key2.clone(), value.clone()).unwrap(); + tree.hash(); + let proof = InclusionProof::from_tree(&tree, &key2, &value).unwrap(); + assert!(proof.validate(&key2, &value, tree.hash())); } #[test] - fn merkle_proofs() { - let n = 15; - let keys = random_keys(n, 420); - let values = random_values(n, 1420); + fn exclusion_proofs() { let mut tree = SparseMerkleTree::>::default(); - (0..n).for_each(|i| { - let _ = tree.upsert(keys[i].clone(), values[i].clone()).unwrap(); - }); - let root_hash = tree.hash().clone(); - (0..n).for_each(|i| { - let proof = tree.build_proof(&keys[i]).unwrap(); - // Validate the proof with correct key / value - assert!(proof.validate_inclusion_proof(&keys[i], &values[i], &root_hash)); - // Show that incorrect value for existing key fails - assert!(!proof.validate_inclusion_proof(&keys[i], &values[(i + 3) % n], &root_hash),); - // Exclusion proof fails - assert!(!proof.validate_exclusion_proof(&keys[i], &root_hash)); - }); - // Test exclusion proof - let unused_keys = random_keys(n, 72); - (0..n).for_each(|i| { - let proof = tree.build_proof(&unused_keys[i]).unwrap(); - assert!(proof.validate_exclusion_proof(&unused_keys[i], &root_hash)); - assert!(!proof.validate_inclusion_proof(&unused_keys[i], &values[i], &root_hash),); - }); + let proof = ExclusionProof::from_tree(&tree, &NodeKey::from([64; 32])).unwrap(); + // Assert that key does not exist + assert!(proof.validate(&NodeKey::from([64u8; 32]), tree.hash())); + // Validation with a non-existent key, but different root should fail + assert!(!proof.validate(&NodeKey::from([64u8; 32]), &NodeHash::from([1; 32]))); + + // Tree with single node + tree.upsert(NodeKey::from([64u8; 32]), ValueHash::from([42u8; 32])) + .unwrap(); + // Cannot create an exclusion proof for a key that exists + let proof = ExclusionProof::from_tree(&tree, &NodeKey::from([64u8; 32])); + assert!(matches!(proof, Err(SMTError::NonViableProof))); + // A valid exclusion proof + let proof = ExclusionProof::from_tree(&tree, &NodeKey::from([65u8; 32])).unwrap(); + assert!(proof.validate(&NodeKey::from([65u8; 32]), tree.hash())); + // Does not validate against invalid root hash + assert!(!proof.validate(&NodeKey::from([65u8; 32]), &NodeHash::default())); + // All key paths for exclusion proofs are identical right now, so any non-existent key will validate. (bug or + // feature?) + assert!(proof.validate(&NodeKey::from([67u8; 32]), tree.hash())); + // Use an exclusion proof to try and give an invalid validation for a key that does exist + assert_eq!(proof.validate(&NodeKey::from([64u8; 32]), tree.hash()), false); + + // A tree with branches + tree.upsert(NodeKey::from([96u8; 32]), ValueHash::from([1u8; 32])) + .unwrap(); + tree.upsert(NodeKey::from([222u8; 32]), ValueHash::from([2u8; 32])) + .unwrap(); + tree.hash(); + // Cannot create an exclusion proof for a key that exists + let proof = ExclusionProof::from_tree(&tree, &NodeKey::from([222u8; 32])); + assert!(matches!(proof, Err(SMTError::NonViableProof))); + // A valid exclusion proof + let proof = ExclusionProof::from_tree(&tree, &NodeKey::from([99u8; 32])).unwrap(); + assert!(proof.validate(&NodeKey::from([99u8; 32]), tree.hash())); + // Does not validate against invalid root hash + assert!(!proof.validate(&NodeKey::from([99u8; 32]), &NodeHash::default())); + // Not all non-existent keys will validate. (bug or feature?) + assert_eq!(proof.validate(&NodeKey::from([65; 32]), tree.hash()), false); + // But any keys with prefix 011 (that is not 96) will validate + assert!(proof.validate(&NodeKey::from([0b0110_0011; 32]), tree.hash())); + assert!(proof.validate(&NodeKey::from([0b0111_0001; 32]), tree.hash())); + // The key 96 does exist.. + assert_eq!(proof.validate(&NodeKey::from([0b0110_0000; 32]), tree.hash()), false); } } diff --git a/base_layer/mmr/src/sparse_merkle_tree/tree.rs b/base_layer/mmr/src/sparse_merkle_tree/tree.rs index 4d2cdd12d2..9f74a048c9 100644 --- a/base_layer/mmr/src/sparse_merkle_tree/tree.rs +++ b/base_layer/mmr/src/sparse_merkle_tree/tree.rs @@ -8,8 +8,8 @@ use digest::{consts::U32, Digest}; use crate::sparse_merkle_tree::{ bit_utils::{traverse_direction, TraverseDirection}, EmptyNode, + ExclusionProof, LeafNode, - MerkleProof, Node, Node::{Branch, Empty, Leaf}, NodeHash, @@ -172,14 +172,14 @@ impl<'a, H: Digest> TerminalBranch<'a, H> { impl> SparseMerkleTree { /// Lazily returns the hash of the Sparse Merkle tree. This function requires a mutable reference to `self` in - /// case the root node needs to be updated. If you are absolutely sure that the merkle root is correct and want a + /// case the root node needs to be updated. If you are absolutely sure that the Merkle& root is correct and want a /// non-mutable reference, use [`SparseMerkleTree::unsafe_hash()`] instead. pub fn hash(&mut self) -> &NodeHash { self.root.hash() } /// Returns the hash of the Sparse Merkle tree. This function does not require a mutable reference to `self` but - /// should only be used if you are absolutely sure that the merkle root is correct. Otherwise, use + /// should only be used if you are absolutely sure that the Merkle& root is correct. Otherwise, use /// [`SparseMerkleTree::hash()`] instead. pub fn unsafe_hash(&self) -> &NodeHash { self.root.unsafe_hash() @@ -271,9 +271,10 @@ impl> SparseMerkleTree { Ok(node.map(|n| n.as_leaf().unwrap().value())) } - /// Constructs a Merkle proof for the value at location `key`. - pub fn build_proof(&self, key: &NodeKey) -> Result, SMTError> { - let mut path = Vec::new(); + /// Construct the data structures needed to generate the Merkle& proofs. Although this function returns a struct + /// of type `ExclusionProof` it is not really a valid (exclusion) proof. The constructors do additional + /// validation before passing the structure on. For this reason, this method is `private` outside of the module. + pub(crate) fn build_proof_candidate(&self, key: &NodeKey) -> Result, SMTError> { let mut siblings = Vec::new(); let mut current_node = &self.root; while current_node.is_branch() { @@ -282,7 +283,6 @@ impl> SparseMerkleTree { return Err(SMTError::StaleHash); } let dir = traverse_direction(branch.height(), branch.key(), key)?; - path.push(dir); current_node = match dir { TraverseDirection::Left => { siblings.push(branch.right().unsafe_hash().clone()); @@ -294,13 +294,8 @@ impl> SparseMerkleTree { }, }; } - let (key, value) = match current_node { - Branch(_) => return Err(SMTError::UnexpectedNodeType), - Leaf(leaf) => (leaf.key().clone(), Some(leaf.value().clone())), - Empty(_) => (key.clone(), None), - }; - siblings.iter().for_each(|s| println!("Sibling: {s:x}")); - let proof = MerkleProof::new(path, siblings, key, value); + let leaf = current_node.as_leaf().cloned(); + let proof = ExclusionProof::new(siblings, leaf); Ok(proof) } @@ -544,6 +539,13 @@ mod test { assert_eq!(right.key(), &key2); // Hash is e3f62f1bfccca2e03e3238cf22748d6a39a7e5eee1dd4b78e2fdd04b5c47d303 assert_eq!(right.hash().to_string(), format!("{right_hash:x}")); + + // Update a key-value + let old_hash = tree.unsafe_hash().to_string(); + let res = tree.upsert(key1, value2).unwrap(); + assert_eq!(tree.size(), 2); + assert!(matches!(res, UpdateResult::Updated(v) if v == value1)); + assert_ne!(tree.hash().to_string(), old_hash); } #[test] @@ -766,6 +768,12 @@ mod test { // │A│ │B│ // └─┘ └─┘ // Root hash is e693520b5ba4ff8b1e37ae4feabcb54701f32efd6bc4b78db356fa9baa64ca99 + + // Deleting a key that does not exist is ok. + let res = tree.delete(&short_key(5)); + assert!(matches!(res, Ok(DeleteResult::KeyNotFound))); + + // Delete an existing key let res = tree.delete(&short_key(224)).unwrap(); assert_eq!(res, DeleteResult::Deleted(ValueHash::from([4u8; 32]))); assert_eq!( diff --git a/scripts/install_ubuntu_dependencies.sh b/scripts/install_ubuntu_dependencies.sh index ea5ca5c54f..99e786d1d8 100755 --- a/scripts/install_ubuntu_dependencies.sh +++ b/scripts/install_ubuntu_dependencies.sh @@ -3,7 +3,6 @@ apt-get -y install \ libssl-dev \ pkg-config \ libsqlite3-dev \ - clang-10 \ git \ cmake \ dh-autoreconf \