From 669ebb49fb18724c5c22a5eb416e56cc767b3edc Mon Sep 17 00:00:00 2001
From: "Augusto F. Hack" <hack.augusto@gmail.com>
Date: Thu, 16 Mar 2023 23:26:02 +0100
Subject: [PATCH] bugfix: check if the requested root is in the storage

---
 src/merkle/store.rs | 18 ++++++------------
 1 file changed, 6 insertions(+), 12 deletions(-)

diff --git a/src/merkle/store.rs b/src/merkle/store.rs
index ae2dec9..5b5c88e 100644
--- a/src/merkle/store.rs
+++ b/src/merkle/store.rs
@@ -235,6 +235,12 @@ impl MerkleStore {
     /// This will return `NodeNotInStorage` if the element is not present in the store.
     pub fn get_node(&self, root: Word, index: NodeIndex) -> Result<Word, MerkleError> {
         let mut hash: RpoDigest = root.into();
+
+        // Check the root is in the storage when called with `NodeIndex::root()`
+        self.nodes
+            .get(&hash)
+            .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?;
+
         for bit in index.bit_iterator().rev() {
             let node = self
                 .nodes
@@ -385,18 +391,6 @@ mod test {
         Ok(())
     }
 
-    #[test]
-    fn test_get_node_returns_self_for_root() {
-        let store = MerkleStore::default();
-        let root_idx = NodeIndex::new(0, 0);
-
-        // the root does not need any lookups in the storage itself, so the value is just returned
-        assert_eq!(store.get_node(LEAVES4[0], root_idx).unwrap(), LEAVES4[0]);
-        assert_eq!(store.get_node(LEAVES4[1], root_idx).unwrap(), LEAVES4[1]);
-        assert_eq!(store.get_node(LEAVES4[2], root_idx).unwrap(), LEAVES4[2]);
-        assert_eq!(store.get_node(LEAVES4[3], root_idx).unwrap(), LEAVES4[3]);
-    }
-
     #[test]
     fn test_get_invalid_node() {
         let mut store = MerkleStore::default();