Merge pull request #116 from ralexstokes/merkle-api

refactor interface for `is_valid_merkle_branch`
ralexstokes · Nov 8, 2023 · 5f1ec83 · 5f1ec83
2 parents 7e6986d + 0fe3842
commit 5f1ec83
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 21 deletions.
diff --git a/ssz-rs/src/merkleization/mod.rs b/ssz-rs/src/merkleization/mod.rs
@@ -31,6 +31,8 @@ pub enum MerkleizationError {
     SerializationError(SerializeError),
     /// More data was provided than expected
     InputExceedsLimit(usize),
+    /// Proof verification failed
+    InvalidProof,
 }
 
 impl From<SerializeError> for MerkleizationError {
@@ -46,6 +48,7 @@ impl Display for MerkleizationError {
                 write!(f, "failed to serialize value: {err}")
             }
             Self::InputExceedsLimit(size) => write!(f, "data exceeds the declared limit {size}"),
+            Self::InvalidProof => write!(f, "merkle proof verification failed"),
         }
     }
 }

diff --git a/ssz-rs/src/merkleization/proofs.rs b/ssz-rs/src/merkleization/proofs.rs
@@ -1,33 +1,40 @@
-use crate::merkleization::Node;
+use crate::merkleization::{MerkleizationError as Error, Node};
 use sha2::{Digest, Sha256};
 
 /// `is_valid_merkle_branch` verifies the Merkle proof
 /// against the `root` given the other metadata.
-pub fn is_valid_merkle_branch<'a>(
-    leaf: &Node,
-    mut branch: impl Iterator<Item = &'a Node>,
+pub fn is_valid_merkle_branch<T: AsRef<[u8]>>(
+    leaf: Node,
+    branch: &[T],
     depth: usize,
     index: usize,
-    root: &Node,
-) -> bool {
-    let mut value = *leaf;
+    root: Node,
+) -> Result<(), Error> {
+    if branch.len() != depth {
+        return Err(Error::InvalidProof)
+    }
 
+    let mut derived_root = leaf;
     let mut hasher = Sha256::new();
-    for i in 0..depth {
-        let next_node = match branch.next() {
-            Some(node) => node,
-            None => return false,
-        };
+
+    for (i, node) in branch.iter().enumerate() {
+        let node = Node::try_from(node.as_ref()).map_err(|_| Error::InvalidProof)?;
+
         if (index / 2usize.pow(i as u32)) % 2 != 0 {
-            hasher.update(next_node.as_ref());
-            hasher.update(value.as_ref());
+            hasher.update(node.as_ref());
+            hasher.update(derived_root.as_ref());
         } else {
-            hasher.update(value.as_ref());
-            hasher.update(next_node.as_ref());
+            hasher.update(derived_root.as_ref());
+            hasher.update(node.as_ref());
         }
-        value.as_mut().copy_from_slice(&hasher.finalize_reset());
+        derived_root.copy_from_slice(&hasher.finalize_reset());
+    }
+
+    if derived_root == root {
+        Ok(())
+    } else {
+        Err(Error::InvalidProof)
     }
-    value == *root
 }
 
 #[cfg(test)]
@@ -48,17 +55,16 @@ mod tests {
             "8f594dbb4f4219ad4967f86b9cccdb26e37e44995a291582a431eef36ecba45c",
             "f8c2ed25e9c31399d4149dcaa48c51f394043a6a1297e65780a5979e3d7bb77c",
             "382ba9638ce263e802593b387538faefbaed106e9f51ce793d405f161b105ee6",
-            "c78009fdf07fc56a11f122370658a353aaa542ed63e44c4bc15ff4cd105ab33c",
         ]
         .into_iter()
-        .map(decode_node_from_hex)
+        .map(|str| hex::decode(str).expect("is valid"))
         .collect::<Vec<_>>();
         let depth = 3;
         let index = 2;
         let root = decode_node_from_hex(
             "27097c728aade54ff1376d5954681f6d45c282a81596ef19183148441b754abb",
         );
 
-        assert!(is_valid_merkle_branch(&leaf, branch.iter(), depth, index, &root))
+        assert!(is_valid_merkle_branch(leaf, &branch, depth, index, root).is_ok());
     }
 }