Fix aristo proof verification - handle embedded trie nodes

execution_chain/db/aristo/aristo_proof.nim

@@ -199,6 +199,12 @@ proc makeMultiProof*(
 
   ok()
 
+template rlpNodeToBytes(node: Rlp): seq[byte] =
+  if node.isList():
+    node.rawData.toSeq()
+  else:
+    node.toBytes()
+
 proc trackRlpNodes(
     chain: openArray[seq[byte]];
     nextIndex: int;
@@ -233,14 +239,14 @@ proc trackRlpNodes(
   of 2:
     let (isLeaf, segm) = NibblesBuf.fromHexPrefix rlpNode.listElem(0).toBytes
     nChewOff = sharedPrefixLen(path, segm)
-    link = rlpNode.listElem(1).toBytes # link or payload
+    link = rlpNode.listElem(1).rlpNodeToBytes() # link or payload
     if isLeaf:
       if nChewOff == path.len:
         return ok(link)
       return err(PartTrkLeafPfxMismatch)
   of 17:
     nChewOff = 1
-    link = rlpNode.listElem(path[0].int).toBytes
+    link = rlpNode.listElem(path[0].int).rlpNodeToBytes()
   else:
     return err(PartTrkGarbledNode)
 
@@ -292,14 +298,14 @@ proc trackRlpNodes(
   of 2:
     let (isLeaf, segm) = NibblesBuf.fromHexPrefix rlpNode.listElem(0).toBytes
     nChewOff = sharedPrefixLen(path, segm)
-    link = rlpNode.listElem(1).toBytes # link or payload
+    link = rlpNode.listElem(1).rlpNodeToBytes() # link or payload
     if isLeaf:
       if nChewOff == path.len:
         return ok(link)
       return err(PartTrkLeafPfxMismatch)
   of 17:
     nChewOff = 1
-    link = rlpNode.listElem(path[0].int).toBytes
+    link = rlpNode.listElem(path[0].int).rlpNodeToBytes()
   else:
     return err(PartTrkGarbledNode)
 

tests/test_aristo_proof.nim

@@ -14,13 +14,14 @@ import
   unittest2,
   stint,
   results,
+  eth/rlp,
   stew/byteutils,
   eth/common/hashes,
   eth/trie/[hexary, db, trie_defs],
   ../execution_chain/db/aristo/aristo_proof
 
 template getBytes(i: int): seq[byte] =
-  @(u256(i).toBytesBE())
+  rlp.encode(u256(i))
 
 func toNodesTable(proofNodes: openArray[seq[byte]]): Table[Hash32, seq[byte]] =
   var nodes: Table[Hash32, seq[byte]]
@@ -29,7 +30,7 @@ func toNodesTable(proofNodes: openArray[seq[byte]]): Table[Hash32, seq[byte]] =
   nodes
 
 suite "Aristo proof verification":
-  const numValues = 1000
+  const numValues = 10000
 
   setup:
     let db = newMemoryDB()
@@ -151,6 +152,41 @@ suite "Aristo proof verification":
     check:
       leafValue.isNone()
 
+  # By using a small value (1 in this case) and storing a large number of keys
+  # this test reproduces the scenario where leaf trie nodes get embedded into
+  # the parent node because the len of the rlp encoded node is less than 32.
+  test "Validate proof for existing value - embedded leafs":
+    const
+      iterations = 100000
+      leaf = getBytes(1)
+
+    for i in 1..iterations:
+      let
+        indexBytes = getBytes(i)
+        key = keccak256(indexBytes)
+        value = indexBytes
+      trie.put(key.data, leaf)
+
+    for i in 1..iterations:
+      let
+        indexBytes = getBytes(i)
+        key = keccak256(indexBytes)
+        value = indexBytes
+        root = trie.rootHash()
+        proof = trie.getBranch(key.data)
+
+      block:
+        let leafValue = verifyProof(proof, root, key).expect("valid proof")
+        check:
+          leafValue.isSome()
+          leafValue.get() == leaf
+
+      block:
+        let leafValue = verifyProof(toNodesTable(proof), root, key).expect("valid proof")
+        check:
+          leafValue.isSome()
+          leafValue.get() == leaf
+
   # The following test cases were copied from the Rust hexary trie implementation.
   # See here: https://github.com/citahub/cita_trie/blob/master/src/tests/mod.rs#L554
   test "Validate proof for empty trie":