From 03ae1a23aea1f8524b261ac4b59c0b4faedfdb35 Mon Sep 17 00:00:00 2001
From: Richard Pringle <richard.pringle@avalabs.org>
Date: Mon, 20 Nov 2023 18:11:43 -0500
Subject: [PATCH] Add proof verification test (#359)

---
 firewood/src/merkle.rs      | 32 ++++++++++++++++++++++++++++++++
 firewood/src/merkle/node.rs |  1 +
 firewood/src/merkle_util.rs |  2 +-
 firewood/src/proof.rs       | 19 +++++++++++++++----
 firewood/tests/db.rs        |  2 +-
 5 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/firewood/src/merkle.rs b/firewood/src/merkle.rs
index 30c683f24..4f2f77c71 100644
--- a/firewood/src/merkle.rs
+++ b/firewood/src/merkle.rs
@@ -1837,4 +1837,36 @@ mod tests {
         assert_eq!(rangeproof.first_key_proof.0, rangeproof.last_key_proof.0);
         assert_eq!(rangeproof.middle.len(), 0);
     }
+
+    #[test]
+    fn shared_path_proof() {
+        let mut merkle = create_test_merkle();
+        let root = merkle.init_root().unwrap();
+
+        let key1 = b"key1";
+        let value1 = b"1";
+        merkle.insert(key1, value1.to_vec(), root).unwrap();
+
+        let key2 = b"key2";
+        let value2 = b"2";
+        merkle.insert(key2, value2.to_vec(), root).unwrap();
+
+        let root_hash = merkle.root_hash(root).unwrap();
+
+        let verified = {
+            let key = key1;
+            let proof = merkle.prove(key, root).unwrap();
+            proof.verify(key, root_hash.0).unwrap()
+        };
+
+        assert_eq!(verified, Some(value1.to_vec()));
+
+        let verified = {
+            let key = key2;
+            let proof = merkle.prove(key, root).unwrap();
+            proof.verify(key, root_hash.0).unwrap()
+        };
+
+        assert_eq!(verified, Some(value2.to_vec()));
+    }
 }
diff --git a/firewood/src/merkle/node.rs b/firewood/src/merkle/node.rs
index b048bb2a5..604b38d2e 100644
--- a/firewood/src/merkle/node.rs
+++ b/firewood/src/merkle/node.rs
@@ -376,6 +376,7 @@ impl NodeType {
         match items.len() {
             EXT_NODE_SIZE => {
                 let mut items = items.into_iter();
+
                 let decoded_key: Vec<u8> = items.next().unwrap().decode()?;
 
                 let decoded_key_nibbles = Nibbles::<0>::new(&decoded_key);
diff --git a/firewood/src/merkle_util.rs b/firewood/src/merkle_util.rs
index 3e9fcd3a1..f7bad0a89 100644
--- a/firewood/src/merkle_util.rs
+++ b/firewood/src/merkle_util.rs
@@ -100,7 +100,7 @@ impl<S: ShaleStore<Node> + Send + Sync> MerkleSetup<S> {
     ) -> Result<Option<Vec<u8>>, DataStoreError> {
         let hash: [u8; 32] = *self.root_hash()?;
         proof
-            .verify_proof(key, hash)
+            .verify(key, hash)
             .map_err(|_err| DataStoreError::ProofVerificationError)
     }
 
diff --git a/firewood/src/proof.rs b/firewood/src/proof.rs
index 8a566a737..fa50404a6 100644
--- a/firewood/src/proof.rs
+++ b/firewood/src/proof.rs
@@ -92,6 +92,8 @@ impl From<DbError> for ProofError {
 /// to a single proof step. If reaches an end step during proof verification,
 /// the hash value will be none, and the encoded value will be the value of the
 /// node.
+
+#[derive(Debug)]
 struct SubProof {
     encoded: Vec<u8>,
     hash: Option<[u8; 32]>,
@@ -103,7 +105,7 @@ impl<N: AsRef<[u8]> + Send> Proof<N> {
     /// proof contains invalid trie nodes or the wrong value.
     ///
     /// The generic N represents the storage for the node data
-    pub fn verify_proof<K: AsRef<[u8]>>(
+    pub fn verify<K: AsRef<[u8]>>(
         &self,
         key: K,
         root_hash: [u8; 32],
@@ -113,24 +115,33 @@ impl<N: AsRef<[u8]> + Send> Proof<N> {
         let mut cur_hash = root_hash;
         let proofs_map = &self.0;
 
-        loop {
+        let value_node = loop {
             let cur_proof = proofs_map
                 .get(&cur_hash)
                 .ok_or(ProofError::ProofNodeMissing)?;
+
             let node = NodeType::decode(cur_proof.as_ref())?;
             let (sub_proof, traversed_nibbles) = locate_subproof(key_nibbles, node)?;
             key_nibbles = traversed_nibbles;
 
             cur_hash = match sub_proof {
                 // Return when reaching the end of the key.
-                Some(p) if key_nibbles.size_hint().0 == 0 => return Ok(Some(p.encoded)),
+                Some(p) if key_nibbles.size_hint().0 == 0 => break p.encoded,
                 // The trie doesn't contain the key.
                 Some(SubProof {
                     hash: Some(hash), ..
                 }) => hash,
                 _ => return Ok(None),
             };
-        }
+        };
+
+        let value = match NodeType::decode(&value_node) {
+            Ok(NodeType::Branch(branch)) => branch.value().as_ref().map(|v| v.to_vec()),
+            Ok(NodeType::Leaf(leaf)) => leaf.data().to_vec().into(),
+            _ => return Ok(value_node.into()),
+        };
+
+        Ok(value)
     }
 
     pub fn concat_proofs(&mut self, other: Proof<N>) {
diff --git a/firewood/tests/db.rs b/firewood/tests/db.rs
index 2fcc0e20f..83524b136 100644
--- a/firewood/tests/db.rs
+++ b/firewood/tests/db.rs
@@ -221,7 +221,7 @@ async fn create_db_issue_proof() {
 
     match rev.single_key_proof(key).await {
         Ok(proof) => {
-            let verification = proof.unwrap().verify_proof(key, root_hash).unwrap();
+            let verification = proof.unwrap().verify(key, root_hash).unwrap();
             assert!(verification.is_some());
         }
         Err(e) => {