From af7b25cee4b77a67a987227efeb07afd0c9cd6ec Mon Sep 17 00:00:00 2001
From: Pierre Krieger <pierre.krieger1708@gmail.com>
Date: Mon, 28 Nov 2022 20:23:58 +0100
Subject: [PATCH] Fix Merkle proof whose trie root node is inferior to 32 bytes
 being considered as invalid (#3046)

Right now, we either hash each proof entry or take its inline value if
it is too small.

This isn't correct, for two reasons:

- The root node is always hashed, meaning that right now if the root
node of a proof is < 32 bytes, its hash won't be found in our array.
- Nodes that are < 32 bytes should be inlined. There's no situation
where a node whose length is < 32 bytes has its own proof entry.

Consequently, we don't care about inline values and just hash the node
entries.
We also check whether hashed nodes shouldn't have been inlined, although
this is very secondary.

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
---
 bin/wasm-node/CHANGELOG.md |  4 +++
 src/trie/proof_decode.rs   | 50 +++++++++++++++++++++++++++++---------
 2 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/bin/wasm-node/CHANGELOG.md b/bin/wasm-node/CHANGELOG.md
index bbbc3c833e..2c6ae370e8 100644
--- a/bin/wasm-node/CHANGELOG.md
+++ b/bin/wasm-node/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Fixed
+
+- Fix Merkle proofs whose trie root node has a size inferior to 32 bytes being considered as invalid. ([#3046](https://github.com/paritytech/smoldot/pull/3046))
+
 ## 0.7.9 - 2022-11-28
 
 ### Fixed
diff --git a/src/trie/proof_decode.rs b/src/trie/proof_decode.rs
index 636eadde81..c1b5b696b8 100644
--- a/src/trie/proof_decode.rs
+++ b/src/trie/proof_decode.rs
@@ -96,16 +96,15 @@ where
             .copied()
             .enumerate()
             .map(
-                |(proof_entry_num, proof_entry)| -> (arrayvec::ArrayVec<u8, 32>, (usize, ops::Range<usize>)) {
-                    let hash = if proof_entry.len() >= 32 {
-                        blake2_rfc::blake2b::blake2b(32, &[], proof_entry)
-                            .as_bytes()
-                            .iter()
-                            .copied()
-                            .collect()
-                    } else {
-                        proof_entry.iter().copied().collect()
-                    };
+                |(proof_entry_num, proof_entry)| -> ([u8; 32], (usize, ops::Range<usize>)) {
+                    // The merkle value of a trie node is normally either its hash or the node
+                    // itself if its length is < 32. In the context of a proof, however, nodes
+                    // whose length is < 32 aren't supposed to be their own entry. For this reason,
+                    // we only hash each entry.
+                    let hash = *<&[u8; 32]>::try_from(
+                        blake2_rfc::blake2b::blake2b(32, &[], proof_entry).as_bytes(),
+                    )
+                    .unwrap();
 
                     let proof_entry_offset = if proof_entry.len() == 0 {
                         0
@@ -115,7 +114,10 @@ where
 
                     (
                         hash,
-                        (proof_entry_num, proof_entry_offset..(proof_entry_offset + proof_entry.len())),
+                        (
+                            proof_entry_num,
+                            proof_entry_offset..(proof_entry_offset + proof_entry.len()),
+                        ),
                     )
                 },
             )
@@ -221,6 +223,12 @@ where
                     if let Some((child_position, child_entry_range)) =
                         merkle_values.get(child_node_value)
                     {
+                        // If the node value of the child is less than 32 bytes long, it should
+                        // have been inlined instead of given separately.
+                        if child_entry_range.end - child_entry_range.start < 32 {
+                            return Err(Error::UnexpectedHashedNode);
+                        }
+
                         // Remove the entry from `unvisited_proof_entries`.
                         // Note that it is questionable what to do if the same entry is visited
                         // multiple times. In case where multiple storage branches are identical,
@@ -600,6 +608,9 @@ pub enum Error {
     UnusedProofEntry,
     /// The same entry has been found multiple times in the proof.
     DuplicateProofEntry,
+    /// A node has been passed separately and referred to by its hash, while its length is inferior
+    /// to 32 bytes.
+    UnexpectedHashedNode,
 }
 
 /// Information about an entry in the proof.
@@ -857,4 +868,21 @@ mod tests {
 
         assert_eq!(obtained, Some(&[80, 82, 127, 41, 119, 1, 0, 0][..]));
     }
+
+    #[test]
+    fn very_small_root_node_decodes() {
+        // Checks that a proof with one root node whose length is < 32 bytes properly verifies.
+        let proof = vec![
+            4, 64, 66, 3, 52, 120, 31, 215, 222, 245, 16, 76, 51, 181, 0, 245, 192, 194,
+        ];
+
+        super::decode_and_verify_proof(super::Config {
+            proof,
+            trie_root_hash: &[
+                83, 2, 191, 235, 8, 252, 233, 114, 129, 199, 229, 115, 221, 238, 15, 205, 193, 110,
+                145, 107, 12, 3, 10, 145, 117, 211, 203, 151, 182, 147, 221, 178,
+            ],
+        })
+        .unwrap();
+    }
 }