fix account node with 'CodeUntouched' type

stateless/tree_from_witness.nim

@@ -13,6 +13,7 @@ type
 
   AccountAndSlots* = object
     address*: EthAddress
+    codeLen*: int
     slots*: seq[StorageSlot]
 
   TreeBuilder = object
@@ -190,7 +191,7 @@ proc buildForest*(t: var TreeBuilder): seq[KeccakHash]
 
     result.add KeccakHash(data: res.data)
 
-proc treeNode(t: var TreeBuilder, depth: int = 0, storageMode = false): NodeKey =
+proc treeNode(t: var TreeBuilder, depth: int, storageMode = false): NodeKey =
   assert(depth < 64)
   let nodeType = safeReadEnum(t, TrieNodeType)
   case nodeType
@@ -303,6 +304,19 @@ proc readAddress(t: var TreeBuilder) =
   safeReadBytes(t, 20):
     t.keys.add AccountAndSlots(address: toAddress(t.read(20)))
 
+proc readCodeLen(t: var TreeBuilder): int =
+  let codeLen = t.safeReadU32()
+  if wfEIP170 in t.flags and codeLen > EIP170_CODE_SIZE_LIMIT:
+    raise newException(ContractCodeError, "code len exceed EIP170 code size limit")
+  t.keys[^1].codeLen = codeLen.int
+  result = codeLen.int
+
+proc readHashNode(t: var TreeBuilder): NodeKey =
+  let nodeType = safeReadEnum(t, TrieNodeType)
+  if nodeType != HashNodeType:
+    raise newException(ParsingError, "hash node expected but got " & $nodeType)
+  result = t.hashNode()
+
 proc accountNode(t: var TreeBuilder, depth: int): NodeKey =
   assert(depth < 65)
 
@@ -337,12 +351,9 @@ proc accountNode(t: var TreeBuilder, depth: int): NodeKey =
       acc.codeHash = blankStringHash
       acc.storageRoot = emptyRlpHash
     of ExtendedAccountType:
-      let codeLen = t.safeReadU32()
-      if wfEIP170 in t.flags and codeLen > EIP170_CODE_SIZE_LIMIT:
-        raise newException(ContractCodeError, "code len exceed EIP170 code size limit")
-
-      safeReadBytes(t, codeLen.int):
-        acc.codeHash = t.writeCode(t.read(codeLen.int))
+      let codeLen = t.readCodeLen()
+      safeReadBytes(t, codeLen):
+        acc.codeHash = t.writeCode(t.read(codeLen))
 
       # switch to account storage parsing mode
       # and reset the depth
@@ -350,11 +361,12 @@ proc accountNode(t: var TreeBuilder, depth: int): NodeKey =
       doAssert(storageRoot.usedBytes == 32)
       acc.storageRoot.data = storageRoot.data
     of CodeUntouched:
-      # TODO: should we check it is a hashNode ?
-      let codeHash = t.treeNode(depth)
+      let codeHash = t.readHashNode()
       doAssert(codeHash.usedBytes == 32)
       acc.codeHash.data = codeHash.data
 
+      discard t.readCodeLen()
+
       let storageRoot = t.treeNode(0, storageMode = true)
       doAssert(storageRoot.usedBytes == 32)
       acc.storageRoot.data = storageRoot.data

stateless/witness_from_tree.nim

@@ -105,7 +105,7 @@ proc writeHashNode(wb: var WitnessBuilder, node: openArray[byte]) =
   wb.output.append(HashNodeType.byte)
   wb.output.append(node)
 
-proc getBranchRecurseAux(wb: var WitnessBuilder, z: var StackElem) {.raises: [ContractCodeError, IOError, Defect, CatchableError, Exception].}
+proc getBranchRecurse(wb: var WitnessBuilder, z: var StackElem) {.raises: [ContractCodeError, IOError, Defect, CatchableError, Exception].}
 
 proc writeAccountNode(wb: var WitnessBuilder, kd: KeyData, acc: Account, nibbles: NibblesSeq,
   node: openArray[byte], depth: int) {.raises: [ContractCodeError, IOError, Defect, CatchableError, Exception].} =
@@ -136,6 +136,11 @@ proc writeAccountNode(wb: var WitnessBuilder, kd: KeyData, acc: Account, nibbles
   if accountType != SimpleAccountType:
     if not kd.codeTouched:
       wb.writeHashNode(acc.codeHash.data)
+      let code = get(wb.db, contractHashKey(acc.codeHash).toOpenArray)
+      if wfEIP170 in wb.flags and code.len > EIP170_CODE_SIZE_LIMIT:
+        raise newException(ContractCodeError, "code len exceed EIP170 code size limit")
+      wb.writeU32(code.len.uint32)
+      # no code here
     elif acc.codeHash != blankStringHash:
       let code = get(wb.db, contractHashKey(acc.codeHash).toOpenArray)
       if wfEIP170 in wb.flags and code.len > EIP170_CODE_SIZE_LIMIT:
@@ -156,12 +161,13 @@ proc writeAccountNode(wb: var WitnessBuilder, kd: KeyData, acc: Account, nibbles
         depth: 0,          # reset depth
         storageMode: true  # switch to storage mode
       )
-      getBranchRecurseAux(wb, zz)
+      getBranchRecurse(wb, zz)
     else:
       wb.writeHashNode(emptyRlpHash.data)
 
   #0x00 pathnibbles:<Nibbles(64-d)> address:<Address> balance:<Bytes32> nonce:<Bytes32>
   #0x01 pathnibbles:<Nibbles(64-d)> address:<Address> balance:<Bytes32> nonce:<Bytes32> bytecode:<Bytecode> storage:<Tree_Node(0,1)>
+  #0x02 pathnibbles:<Nibbles(64-d)> address:<Address> balance:<Bytes32> nonce:<Bytes32> codehash:<Bytes32> codesize:<U32> storage:<Account_Storage_Tree_Node(0)>
 
 proc writeAccountStorageLeafNode(wb: var WitnessBuilder, key: openArray[byte], val: UInt256, nibbles: NibblesSeq, node: openArray[byte], depth: int) =
   wb.output.append(StorageLeafNodeType.byte)
@@ -181,7 +187,7 @@ proc writeAccountStorageLeafNode(wb: var WitnessBuilder, key: openArray[byte], v
 
   #<Storage_Leaf_Node(d<65)> := pathnibbles:<Nibbles(64-d))> key:<Bytes32> val:<Bytes32>
 
-proc getBranchRecurseAux(wb: var WitnessBuilder, z: var StackElem) =
+proc getBranchRecurse(wb: var WitnessBuilder, z: var StackElem) =
   if z.node.len == 0: return
   var nodeRlp = rlpFromBytes z.node
 
@@ -189,8 +195,12 @@ proc getBranchRecurseAux(wb: var WitnessBuilder, z: var StackElem) =
   of 2:
     let (isLeaf, k) = nodeRlp.extensionNodeKey
     var match = false
+    # only zero or one group can match the path
+    # but if there is a match, it can be in any position
+    # 1st, 2nd, or max 3rd position
+    # recursion will go deeper depend on the common-prefix length nibbles
     for mg in groups(z.keys, z.depth, k, z.parentGroup):
-      if mg.match:
+      if not mg.match: continue
       doAssert(match == false) # should be only one match
       match = true
       let value = nodeRlp.listElem(1)
@@ -204,7 +214,7 @@ proc getBranchRecurseAux(wb: var WitnessBuilder, z: var StackElem) =
           depth: z.depth + k.len,
           storageMode: z.storageMode
         )
-          getBranchRecurseAux(wb, zz)
+        getBranchRecurse(wb, zz)
       else:
         # this should be only one match
         # if there is more than one match
@@ -225,9 +235,12 @@ proc getBranchRecurseAux(wb: var WitnessBuilder, z: var StackElem) =
     writeBranchNode(wb, branchMask, z.depth, z.node)
     let path = groups(z.keys, z.parentGroup, z.depth)
 
+    # if there is a match in any branch elem
+    # 1st to 16th, the recursion will go deeper
+    # by one nibble
     let notLeaf = z.depth != 63 # path.len == 0
     for i in 0..<16:
-      if branchMask.branchMaskBitIsSet(i):
+      if not branchMask.branchMaskBitIsSet(i): continue
       var branch = nodeRlp.listElem(i)
       if notLeaf and branchMaskBitIsSet(path.mask, i):
         var zz = StackElem(
@@ -237,7 +250,7 @@ proc getBranchRecurseAux(wb: var WitnessBuilder, z: var StackElem) =
           depth: z.depth + 1,
           storageMode: z.storageMode
         )
-          getBranchRecurseAux(wb, zz)
+        getBranchRecurse(wb, zz)
       else:
         if branch.isList:
           doAssert(false, "Short node should not exist in block witness")
@@ -251,7 +264,7 @@ proc getBranchRecurseAux(wb: var WitnessBuilder, z: var StackElem) =
     raise newException(CorruptedTrieDatabase,
                        "HexaryTrie node with an unexpected number of children")
 
-proc buildWitness*(wb: var WitnessBuilder, keys: MultikeysRef, withVersion: bool = true): seq[byte]
+proc buildWitness*(wb: var WitnessBuilder, keys: MultikeysRef): seq[byte]
   {.raises: [ContractCodeError, IOError, Defect, CatchableError, Exception].} =
 
   # witness version
@@ -269,7 +282,7 @@ proc buildWitness*(wb: var WitnessBuilder, keys: MultikeysRef, withVersion: bool
     depth: 0,
     storageMode: false
   )
-  getBranchRecurseAux(wb, z)
+  getBranchRecurse(wb, z)
 
   # result
   result = wb.output.getOutput(seq[byte])