From 0a9acbc4d92453372764ca983175bae01a7efb35 Mon Sep 17 00:00:00 2001
From: witter-deland <87846830+witter-deland@users.noreply.github.com>
Date: Mon, 22 May 2023 20:40:31 +0800
Subject: [PATCH] fix: incorrect proof hash count causing panic (#28)

---
 src/merkle_proof.rs        |  4 ++++
 tests/merkle_proof_test.rs | 27 +++++++++++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/src/merkle_proof.rs b/src/merkle_proof.rs
index 1a87116..f8a0e49 100644
--- a/src/merkle_proof.rs
+++ b/src/merkle_proof.rs
@@ -228,7 +228,11 @@ impl<T: Hasher> MerkleProof<T> {
         // The next lines copy hashes from proof hashes and group them by layer index
         let mut proof_layers: Vec<Vec<(usize, T::Hash)>> = Vec::with_capacity(tree_depth + 1);
         let mut proof_copy = self.proof_hashes.clone();
+
         for proof_indices in proof_indices_by_layers {
+            if proof_copy.len() < proof_indices.len() {
+                return Err(Error::not_enough_hashes_to_calculate_root());
+            }
             let proof_hashes = proof_copy.splice(0..proof_indices.len(), []);
             proof_layers.push(proof_indices.iter().cloned().zip(proof_hashes).collect());
         }
diff --git a/tests/merkle_proof_test.rs b/tests/merkle_proof_test.rs
index d0bb380..aa32042 100644
--- a/tests/merkle_proof_test.rs
+++ b/tests/merkle_proof_test.rs
@@ -80,6 +80,33 @@ pub mod root {
 
         Ok(())
     }
+
+    #[test]
+    pub fn should_return_error_not_enough_hashes_to_calculate_root() {
+        let test_data = common::setup();
+        let leaf_hashes = &test_data.leaf_hashes;
+        let expected_root = test_data.expected_root_hex.clone();
+        let indices_to_prove = vec![3, 4];
+
+        let leaves_to_prove: Vec<[u8; 32]> = indices_to_prove
+            .iter()
+            .map(|i| leaf_hashes.get(*i).unwrap().clone())
+            .collect();
+
+        let merkle_tree = MerkleTree::<Sha256>::from_leaves(&test_data.leaf_hashes);
+        let proof = merkle_tree.proof(&indices_to_prove);
+        let extracted_root = proof.root(
+            &indices_to_prove,
+            &leaves_to_prove,
+            test_data.leaf_values.len() + 1,
+        );
+
+        assert_eq!(
+            extracted_root.err().unwrap().to_string(),
+            Error::not_enough_hashes_to_calculate_root().to_string(),
+            "Should return error not_enough_hashes_to_calculate_root"
+        );
+    }
 }
 
 pub mod to_bytes {