treebuilder will raise exception on parsing error

stateless/randutils.nim

@@ -1,4 +1,4 @@
-import random, sets, nimcrypto/[utils, sysrand]
+import random, sets, nimcrypto/sysrand
 
 type
   RandGen*[T] = object

stateless/test_witness_keys.nim

@@ -1,7 +1,7 @@
 import
   randutils, stew/byteutils, random,
   eth/[common, rlp], eth/trie/[hexary, db, trie_defs],
-  faststreams/input_stream, nimcrypto/[utils, sysrand],
+  faststreams/input_stream, nimcrypto/sysrand,
   ../stateless/[witness_from_tree, tree_from_witness],
   ../nimbus/db/storage_types, ./witness_types
 

stateless/tree_from_witness.nim

@@ -46,10 +46,6 @@ else:
 func rootHash*(t: TreeBuilder): KeccakHash {.inline.} =
   t.root
 
-proc writeNode(t: var TreeBuilder, n: openArray[byte]): KeccakHash =
-  result = keccak(n)
-  t.db.put(result.data, n)
-
 when defined(useInputStream):
   template readByte(t: var TreeBuilder): byte =
     t.input.read
@@ -62,6 +58,10 @@ when defined(useInputStream):
 
   template readable(t: var TreeBuilder): bool =
     t.input.readable
+
+  template readable(t: var TreeBuilder, len: int): bool =
+    t.input.readable(len)
+
 else:
   template readByte(t: var TreeBuilder): byte =
     let pos = t.pos
@@ -74,38 +74,44 @@ else:
   template readable(t: var TreeBuilder): bool =
     t.pos < t.input.len
 
+  template readable(t: var TreeBuilder, length: int): bool =
+    t.pos + length <= t.input.len
+
   template read(t: var TreeBuilder, len: int): auto =
     let pos = t.pos
     inc(t.pos, len)
     toOpenArray(t.input, pos, pos+len-1)
 
-proc readU32(t: var TreeBuilder): uint32 =
+proc safeReadByte(t: var TreeBuilder): byte =
-  result = fromBytesBE(uint32, t.read(4))
+  if t.readable:
+    result = t.readByte()
+  else:
+    raise newException(IOError, "Cannot read byte from input stream")
 
-proc toAddress(r: var EthAddress, x: openArray[byte]) {.inline.} =
+proc safeReadU32(t: var TreeBuilder): uint32 =
-  r[0..19] = x[0..19]
+  if t.readable(4):
+    result = fromBytesBE(uint32, t.read(4))
+  else:
+    raise newException(IOError, "Cannot read U32 from input stream")
 
 proc toKeccak(r: var NodeKey, x: openArray[byte]) {.inline.} =
   r.data[0..31] = x[0..31]
   r.usedBytes = 32
 
-proc toKeccak(x: openArray[byte]): NodeKey {.inline.} =
-  result.data[0..31] = x[0..31]
-  result.usedBytes = 32
-
 proc append(r: var RlpWriter, n: NodeKey) =
   if n.usedBytes < 32:
     r.append rlpFromBytes(n.data.toOpenArray(0, n.usedBytes-1))
   else:
     r.append n.data.toOpenArray(0, n.usedBytes-1)
 
-proc toNodeKey(z: openArray[byte]): NodeKey =
+proc toNodeKey(t: var TreeBuilder, z: openArray[byte]): NodeKey =
   if z.len < 32:
     result.usedBytes = z.len
     result.data[0..z.len-1] = z[0..z.len-1]
   else:
     result.data = keccak(z).data
     result.usedBytes = 32
+    t.db.put(result.data, z)
 
 proc writeCode(t: var TreeBuilder, code: openArray[byte]): Hash256 =
   result = keccak(code)
@@ -119,14 +125,14 @@ proc hashNode(t: var TreeBuilder): NodeKey
 proc treeNode(t: var TreeBuilder, depth: int = 0, storageMode = false): NodeKey
 
 proc buildTree*(t: var TreeBuilder): KeccakHash =
-  let version = t.readByte().int
+  let version = t.safeReadByte().int
   if version != BlockWitnessVersion.int:
     raise newException(ParsingError, "Wrong block witness version")
 
   # one or more trees
 
   # we only parse one tree here
-  let metadataType = t.readByte().int
+  let metadataType = t.safeReadByte().int
   if metadataType != MetadataNothing.int:
     raise newException(ParsingError, "This tree builder support no metadata")
 
@@ -137,12 +143,12 @@ proc buildTree*(t: var TreeBuilder): KeccakHash =
   result.data = res.data
 
 proc buildForest*(t: var TreeBuilder): seq[KeccakHash] =
-  let version = t.readByte().int
+  let version = t.safeReadByte().int
   if version != BlockWitnessVersion.int:
     raise newException(ParsingError, "Wrong block witness version")
 
   while t.readable:
-    let metadataType = t.readByte().int
+    let metadataType = t.safeReadByte().int
     if metadataType != MetadataNothing.int:
       raise newException(ParsingError, "This tree builder support no metadata")
 
@@ -154,8 +160,11 @@ proc buildForest*(t: var TreeBuilder): seq[KeccakHash] =
 
 proc treeNode(t: var TreeBuilder, depth: int = 0, storageMode = false): NodeKey =
   assert(depth < 64)
-  let nodeType = TrieNodeType(t.readByte)
+  let typ = t.safeReadByte().int
+  if typ < low(TrieNodeType).int or typ > high(TrieNodeType).int:
+    raise newException(ParsingError, "Wrong trie node type")
 
+  let nodeType = TrieNodeType(typ)
   case nodeType
   of BranchNodeType: result = t.branchNode(depth, storageMode)
   of ExtensionNodeType: result = t.extensionNode(depth, storageMode)
@@ -168,15 +177,18 @@ proc treeNode(t: var TreeBuilder, depth: int = 0, storageMode = false): NodeKey
   of HashNodeType: result = t.hashNode()
 
   if depth == 0 and result.usedBytes < 32:
-    result.data = keccak(result.data.toOpenArray(0, result.usedBytes-1)).data
+    let hash = keccak(result.data.toOpenArray(0, result.usedBytes-1))
+    t.db.put(hash.data, result.data.toOpenArray(0, result.usedBytes-1))
+    result.data = hash.data
     result.usedBytes = 32
 
+
 proc branchNode(t: var TreeBuilder, depth: int, storageMode: bool): NodeKey =
   assert(depth < 64)
-  let mask = constructBranchMask(t.readByte, t.readByte)
+  let mask = constructBranchMask(t.safeReadByte, t.safeReadByte)
 
   when defined(debugDepth):
-    let readDepth = t.readByte.int
+    let readDepth = t.safeReadByte().int
     doAssert(readDepth == depth, "branchNode " & $readDepth & " vs. " & $depth)
 
   when defined(debugHash):
@@ -190,10 +202,13 @@ proc branchNode(t: var TreeBuilder, depth: int, storageMode: bool): NodeKey =
     else:
       r.append ""
 
+  if branchMaskBitIsSet(mask, 16):
+    raise newException(ParsingError, "The 17th elem of branh node should empty")
+
   # 17th elem should always empty
   r.append ""
 
-  result = toNodeKey(r.finish)
+  result = t.toNodeKey(r.finish)
 
   when defined(debugHash):
     if result != hash:
@@ -201,6 +216,7 @@ proc branchNode(t: var TreeBuilder, depth: int, storageMode: bool): NodeKey =
       debugEcho "result: ", result.data.toHex, " vs. ", hash.data.toHex
 
 func hexPrefix(r: var RlpWriter, x: openArray[byte], nibblesLen: int, isLeaf: static[bool] = false) =
+  doAssert(nibblesLen <= 64)
   var bytes: array[33, byte]
   if (nibblesLen mod 2) == 0: # even
     when isLeaf:
@@ -224,27 +240,34 @@ func hexPrefix(r: var RlpWriter, x: openArray[byte], nibblesLen: int, isLeaf: st
 
 proc extensionNode(t: var TreeBuilder, depth: int, storageMode: bool): NodeKey =
   assert(depth < 63)
-  let nibblesLen = int(t.readByte)
+  let nibblesLen = t.safeReadByte().int
   assert(nibblesLen < 65)
   var r = initRlpList(2)
-  r.hexPrefix(t.read(nibblesLen div 2 + nibblesLen mod 2), nibblesLen)
+  let pathLen = nibblesLen div 2 + nibblesLen mod 2
+  if t.readable(pathLen):
+    r.hexPrefix(t.read(pathLen), nibblesLen)
+  else:
+    raise newException(ParsingError, "Failed when read nibbles path")
 
   when defined(debugDepth):
-    let readDepth = t.readByte.int
+    let readDepth = t.safeReadByte().int
     doAssert(readDepth == depth, "extensionNode " & $readDepth & " vs. " & $depth)
 
   when defined(debugHash):
     let hash = toKeccak(t.read(32))
 
   assert(depth + nibblesLen < 65)
-  let nodeType = TrieNodeType(t.readByte)
+  let typ = t.safeReadByte.int
+  if typ < low(TrieNodeType).int or typ > high(TrieNodeType).int:
+    raise newException(ParsingError, "Wrong trie node type")
 
+  let nodeType = TrieNodeType(typ)
   case nodeType
   of BranchNodeType: r.append t.branchNode(depth + nibblesLen, storageMode)
   of HashNodeType: r.append t.hashNode()
   else: raise newException(ValueError, "wrong type during parsing child of extension node")
 
-  result = toNodeKey(r.finish)
+  result = t.toNodeKey(r.finish)
 
   when defined(debugHash):
     if result != hash:
@@ -255,22 +278,34 @@ proc accountNode(t: var TreeBuilder, depth: int): NodeKey =
   assert(depth < 65)
 
   when defined(debugHash):
-    let len = t.readU32().int
+    let len = t.safeReadU32().int
     let node = @(t.read(len))
-    let nodeKey = toNodeKey(node)
+    let nodeKey = t.toNodeKey(node)
 
   when defined(debugDepth):
-    let readDepth = t.readByte.int
+    let readDepth = t.safeReadByte().int
     doAssert(readDepth == depth, "accountNode " & $readDepth & " vs. " & $depth)
 
-  let accountType = AccountType(t.readByte)
+  let typ = t.safeReadByte.int
+  if typ < low(AccountType).int or typ > high(AccountType).int:
+    raise newException(ParsingError, "Wrong account type")
+
+  let accountType = AccountType(typ)
   let nibblesLen = 64 - depth
   var r = initRlpList(2)
-  r.hexPrefix(t.read(nibblesLen div 2 + nibblesLen mod 2), nibblesLen, true)
+
+  let pathLen = nibblesLen div 2 + nibblesLen mod 2
+  if t.readable(pathLen):
+    r.hexPrefix(t.read(pathLen), nibblesLen, true)
+  else:
+    raise newException(ParsingError, "Failed when read nibbles path")
 
   # TODO: parse address
   # let address = toAddress(t.read(20))
 
+  if not t.readable(64):
+    raise newException(ParsingError, "Not enough bytes to read account")
+
   var acc = Account(
     balance: UInt256.fromBytesBE(t.read(32), false),
     # TODO: why nonce must be 32 bytes, isn't 64 bit uint  enough?
@@ -281,10 +316,14 @@ proc accountNode(t: var TreeBuilder, depth: int): NodeKey =
     acc.codeHash = blankStringHash
     acc.storageRoot = emptyRlpHash
   else:
-    let codeLen = t.readU32()
+    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")
+
+    if t.readable(codeLen.int):
       acc.codeHash = t.writeCode(t.read(codeLen.int))
+    else:
+      raise newException(ParsingError, "not enought bytes to read contract code")
 
     # switch to account storage parsing mode
     # and reset the depth
@@ -293,15 +332,15 @@ proc accountNode(t: var TreeBuilder, depth: int): NodeKey =
     acc.storageRoot.data = storageRoot.data
 
   r.append rlp.encode(acc)
-  let noderes = r.finish
+  let nodeRes = r.finish
-  result = toNodeKey(noderes)
+  result = t.toNodeKey(nodeRes)
 
   when defined(debugHash):
     if result != nodeKey:
       debugEcho "result.usedBytes: ", result.usedBytes
       debugEcho "nodeKey.usedBytes: ", nodeKey.usedBytes
       var rlpa = rlpFromBytes(node)
-      var rlpb = rlpFromBytes(noderes)
+      var rlpb = rlpFromBytes(nodeRes)
       debugEcho "Expected: ", inspect(rlpa)
       debugEcho "Actual: ", inspect(rlpb)
       var a = rlpa.listElem(1).toBytes.decode(Account)
@@ -315,13 +354,23 @@ proc accountStorageLeafNode(t: var TreeBuilder, depth: int): NodeKey =
   assert(depth < 65)
   let nibblesLen = 64 - depth
   var r = initRlpList(2)
-  r.hexPrefix(t.read(nibblesLen div 2 + nibblesLen mod 2), nibblesLen, true)
+  let pathLen = nibblesLen div 2 + nibblesLen mod 2
+  if not t.readable(pathLen):
+    raise newException(ParsingError, "Enot enough bytes to read path nibbles")
+
+  r.hexPrefix(t.read(pathLen), nibblesLen, true)
   # TODO: parse key
   # let key = @(t.read(32))
   # UInt256 -> BytesBE -> keccak
+
+  if not t.readable(32):
+    raise newException(ParsingError, "Not enough bytes to read path storage value")
+
   let val = UInt256.fromBytesBE(t.read(32))
   r.append rlp.encode(val)
-  result = toNodeKey(r.finish)
+  result = t.toNodeKey(r.finish)
 
 proc hashNode(t: var TreeBuilder): NodeKey =
+  if not t.readable(32):
+    raise newException(ParsingError, "Not enough bytes to read hash node")
   result.toKeccak(t.read(32))

stateless/witness_types.nim

@@ -1,4 +1,4 @@
-import nimcrypto/[keccak, hash], stew/bitops2
+import nimcrypto/hash, stew/bitops2
 
 type
   TrieNodeType* = enum