From a52d92cffc39e24c4abc5b6f120aa8531ce9b2d1 Mon Sep 17 00:00:00 2001
From: Tobi Demeco <tdemeco@itba.edu.ar>
Date: Thu, 14 Nov 2024 15:39:34 -0300
Subject: [PATCH] fix: :bug: add checks to avoid panicking when decoding proofs
 with `sp-trie`

---
 primitives/file-key-verifier/src/lib.rs   |  8 ++++++++
 primitives/file-key-verifier/src/types.rs |  8 ++++++++
 primitives/forest-verifier/src/lib.rs     | 16 ++++++++++++++++
 3 files changed, 32 insertions(+)

diff --git a/primitives/file-key-verifier/src/lib.rs b/primitives/file-key-verifier/src/lib.rs
index 8d4c27d54..3e84bb42d 100644
--- a/primitives/file-key-verifier/src/lib.rs
+++ b/primitives/file-key-verifier/src/lib.rs
@@ -80,6 +80,14 @@ where
             .try_into()
             .map_err(|_| "Failed to convert fingerprint to a hasher output.")?;
 
+        // Basic sanity check to make sure the received compact proof won't panic
+        // TODO: remove this after [this PR](https://github.com/paritytech/polkadot-sdk/pull/6486) is merged.
+        for encoded_node in proof.proof.encoded_nodes.iter() {
+            if encoded_node.len() < 2 {
+                return Err("Invalid encoded node in the proof.".into());
+            }
+        }
+
         // This generates a partial trie based on the proof and checks that the root hash matches the `expected_root`.
         let (memdb, root) = proof
             .proof
diff --git a/primitives/file-key-verifier/src/types.rs b/primitives/file-key-verifier/src/types.rs
index 2057ce288..1d652bb16 100644
--- a/primitives/file-key-verifier/src/types.rs
+++ b/primitives/file-key-verifier/src/types.rs
@@ -62,6 +62,14 @@ impl<const H_LENGTH: usize, const CHUNK_SIZE: u64, const SIZE_TO_CHALLENGES: u64
             .try_into()
             .map_err(|_| ProvenFileKeyError::FingerprintAndTrieHashMismatch)?;
 
+        // Basic sanity check to make sure the received compact proof won't panic
+        // TODO: remove this after [this PR](https://github.com/paritytech/polkadot-sdk/pull/6486) is merged.
+        for encoded_node in self.proof.encoded_nodes.iter() {
+            if encoded_node.len() < 2 {
+                return Err(ProvenFileKeyError::FailedToDecodeChunkFromProof);
+            }
+        }
+
         // This generates a partial trie based on the proof and checks that the root hash matches the `expected_root`.
         let (memdb, root) = self
             .proof
diff --git a/primitives/forest-verifier/src/lib.rs b/primitives/forest-verifier/src/lib.rs
index 9513970c5..bdcf50782 100644
--- a/primitives/forest-verifier/src/lib.rs
+++ b/primitives/forest-verifier/src/lib.rs
@@ -40,6 +40,14 @@ where
             return Err("No challenges provided.".into());
         }
 
+        // Basic sanity check to make sure the received compact proof won't panic
+        // TODO: remove this after [this PR](https://github.com/paritytech/polkadot-sdk/pull/6486) is merged.
+        for encoded_node in proof.encoded_nodes.iter() {
+            if encoded_node.len() < 2 {
+                return Err("Invalid encoded node in the proof.".into());
+            }
+        }
+
         // This generates a partial trie based on the proof and checks that the root hash matches the `expected_root`.
         let (memdb, root) = proof.to_memory_db(Some(root.into())).map_err(|_| {
             "Failed to convert proof to memory DB, root doesn't match with expected."
@@ -235,6 +243,14 @@ where
             return Err("Root is empty.".into());
         }
 
+        // Basic sanity check to make sure the received compact proof won't panic
+        // TODO: remove this after [this PR](https://github.com/paritytech/polkadot-sdk/pull/6486) is merged.
+        for encoded_node in proof.encoded_nodes.iter() {
+            if encoded_node.len() < 2 {
+                return Err("Invalid encoded node in the proof.".into());
+            }
+        }
+
         // TODO: Understand why `CompactProof` cannot be used directly to construct memdb and modify a partial trie. (it fails with error IncompleteDatabase)
         // Convert compact proof to `sp_trie::StorageProof` in order to access the trie nodes.
         let (storage_proof, mut root) = proof