From 343b3810be72524e12388af443931c9b5db301cd Mon Sep 17 00:00:00 2001
From: Jordy Romuald <87231934+JordyRo1@users.noreply.github.com>
Date: Tue, 16 Jul 2024 18:30:48 +0200
Subject: [PATCH] Audit fixed size for merkle tree construction (#83)

* fix: `build_tree` iteration condition change

* fix: fix related test
---
 .../contracts/hooks/merkle_tree_hook.cairo    |  4 ++--
 .../tests/hooks/test_merkle_tree_hook.cairo   | 23 +++++++++++--------
 2 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/contracts/src/contracts/hooks/merkle_tree_hook.cairo b/contracts/src/contracts/hooks/merkle_tree_hook.cairo
index 7bd8be3..1e8b88c 100644
--- a/contracts/src/contracts/hooks/merkle_tree_hook.cairo
+++ b/contracts/src/contracts/hooks/merkle_tree_hook.cairo
@@ -303,8 +303,8 @@ pub mod merkle_tree_hook {
             let mut cur_idx = 0;
             let mut tree = array![];
             loop {
-                if (cur_idx == self.count.read() - 1) {
-                    break ();
+                if (cur_idx >= TREE_DEPTH) {
+                    break;
                 }
                 tree.append(self.tree.read(cur_idx));
                 cur_idx += 1;
diff --git a/contracts/src/tests/hooks/test_merkle_tree_hook.cairo b/contracts/src/tests/hooks/test_merkle_tree_hook.cairo
index 98f3d0b..a398d6b 100644
--- a/contracts/src/tests/hooks/test_merkle_tree_hook.cairo
+++ b/contracts/src/tests/hooks/test_merkle_tree_hook.cairo
@@ -164,15 +164,20 @@ fn test_insert_node_into_merkle_tree_hook() {
     assert_eq!(state.tree.read(0), ByteData { value: node_3, size: 6 });
     assert_eq!(state.tree.read(1), ByteData { value: expected_hash, size: HASH_SIZE });
     assert_eq!(state.tree.read(2), ByteData { value: expected_hash_2, size: HASH_SIZE });
-    assert(
-        state
-            ._build_tree() == array![
-                ByteData { value: node_3, size: 6 },
-                ByteData { value: expected_hash, size: HASH_SIZE },
-                ByteData { value: expected_hash_2, size: HASH_SIZE }
-            ],
-        'build tree failed'
-    );
+    let mut expected_result = array![
+        ByteData { value: node_3, size: 6 },
+        ByteData { value: expected_hash, size: HASH_SIZE },
+        ByteData { value: expected_hash_2, size: HASH_SIZE },
+    ];
+    let mut cur_idx = 0;
+    loop {
+        if (cur_idx >= merkle_tree_hook::TREE_DEPTH - 3) {
+            break;
+        }
+        expected_result.append(ByteData { value: 0, size: 0 });
+        cur_idx += 1;
+    };
+    assert(state._build_tree() == expected_result, 'build tree failed');
     assert(state._root() != 0, 'root computation failed');
     let (root, count) = state.latest_checkpoint();
     assert_eq!(root, state._root());