ensure branch length and depth are same

Signed-off-by: Jun Kimura <[email protected]>
datachainlab · Jan 14, 2024 · 799fec1 · 799fec1
1 parent c288978
commit 799fec1
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/crates/consensus/src/errors.rs b/crates/consensus/src/errors.rs
@@ -29,8 +29,10 @@ pub enum Error {
 pub enum MerkleError {
     /// invalid merkle branch error: leaf={0:?} branch={1:?} index={2:?} root={3:?}
     InvalidMerkleBranch(H256, Vec<H256>, u64, Root),
-    /// too long merkle branch error: leaf={0:?} branch={1:?} index={2:?} root={3:?}
-    TooLongMerkleBranch(H256, Vec<H256>, u64, Root),
+    /// too long merkle branch error: depth={0:?} leaf={1:?} branch={2:?} index={3:?} root={4:?}
+    TooLongMerkleBranchLength(u32, H256, Vec<H256>, u64, Root),
+    /// invalid merkle branch length error: depth={0:?} leaf={1:?} branch={2:?} index={3:?} root={4:?}
+    InvalidMerkleBranchLength(u32, H256, Vec<H256>, u64, Root),
 }
 
 #[cfg(feature = "std")]

diff --git a/crates/consensus/src/merkle.rs b/crates/consensus/src/merkle.rs
@@ -10,6 +10,15 @@ pub fn is_valid_merkle_branch(
     leaf_index: u64,
     root: Root,
 ) -> Result<(), MerkleError> {
+    if depth != branch.len() as u32 {
+        return Err(MerkleError::InvalidMerkleBranchLength(
+            depth,
+            leaf,
+            branch.to_vec(),
+            leaf_index,
+            root,
+        ));
+    }
     let index = 2u64.pow(depth) + leaf_index;
     let mut value = leaf.clone();
     for (i, b) in branch.iter().enumerate() {
@@ -20,7 +29,8 @@ pub fn is_valid_merkle_branch(
                 value = hash([value.as_bytes(), b.as_bytes()].concat());
             }
         } else {
-            return Err(MerkleError::TooLongMerkleBranch(
+            return Err(MerkleError::TooLongMerkleBranchLength(
+                depth,
                 leaf,
                 branch.to_vec(),
                 index,