From 57e9bea4eca4a73e2bb0cb34b8f59dd048c67361 Mon Sep 17 00:00:00 2001
From: jangko <jangko128@gmail.com>
Date: Sun, 24 May 2020 11:41:20 +0700
Subject: [PATCH] add more test for invalid input

---
 ...h-6081fe3734250348623b03c1d3acb35d90d2ecfc | Bin 0 -> 13 bytes
 ...h-a865d42d79ebbfd42b86d11e2e068643cf0fd503 | Bin 0 -> 28 bytes
 ...h-a8a8b9592a31539c0318ae6a1cd416d5644cb74a | Bin 0 -> 9 bytes
 ...h-d755f044304f53abf28a735763865a449a198235 | Bin 0 -> 198 bytes
 ...h-eb262bbb7d5a0b11e1a8112ea592eb67cc618c9b | Bin 0 -> 137 bytes
 ...h-f9a002e389b120714a2324487a6e473b0fedc12e | Bin 0 -> 202 bytes
 stateless/test_witness_keys.nim               |  15 +++++++++-
 stateless/tree_from_witness.nim               |  26 +++++++++++++-----
 8 files changed, 33 insertions(+), 8 deletions(-)
 create mode 100644 stateless/invalidInput/crash-6081fe3734250348623b03c1d3acb35d90d2ecfc
 create mode 100644 stateless/invalidInput/crash-a865d42d79ebbfd42b86d11e2e068643cf0fd503
 create mode 100644 stateless/invalidInput/crash-a8a8b9592a31539c0318ae6a1cd416d5644cb74a
 create mode 100644 stateless/invalidInput/crash-d755f044304f53abf28a735763865a449a198235
 create mode 100644 stateless/invalidInput/crash-eb262bbb7d5a0b11e1a8112ea592eb67cc618c9b
 create mode 100644 stateless/invalidInput/crash-f9a002e389b120714a2324487a6e473b0fedc12e

diff --git a/stateless/invalidInput/crash-6081fe3734250348623b03c1d3acb35d90d2ecfc b/stateless/invalidInput/crash-6081fe3734250348623b03c1d3acb35d90d2ecfc
new file mode 100644
index 0000000000000000000000000000000000000000..1bf239ad0608a0689781f6fb1e365acae3dfc632
GIT binary patch
literal 13
UcmZQ%VEF%^k%5VUiBXFY026xxF#rGn

literal 0
HcmV?d00001

diff --git a/stateless/invalidInput/crash-a865d42d79ebbfd42b86d11e2e068643cf0fd503 b/stateless/invalidInput/crash-a865d42d79ebbfd42b86d11e2e068643cf0fd503
new file mode 100644
index 0000000000000000000000000000000000000000..09e441efd8166e54223de47e1dd6921a11cc5e2f
GIT binary patch
literal 28
ccmZQ%U}R)q=mruj3_u`we}--Xhyes;03<&HAOHXW

literal 0
HcmV?d00001

diff --git a/stateless/invalidInput/crash-a8a8b9592a31539c0318ae6a1cd416d5644cb74a b/stateless/invalidInput/crash-a8a8b9592a31539c0318ae6a1cd416d5644cb74a
new file mode 100644
index 0000000000000000000000000000000000000000..d5d64275bee6b2cc40100ed9c83037fbee401f44
GIT binary patch
literal 9
QcmZQ%U}Ruo)N+aj00D#m%>V!Z

literal 0
HcmV?d00001

diff --git a/stateless/invalidInput/crash-d755f044304f53abf28a735763865a449a198235 b/stateless/invalidInput/crash-d755f044304f53abf28a735763865a449a198235
new file mode 100644
index 0000000000000000000000000000000000000000..bece4c4b9b5b817cfe1fc23eaba2716d439af173
GIT binary patch
literal 198
zcmZQ%V6<XnU|<B|ZV;g<2qJAA7#W#C0E>7xSS7O}Gc$uG16(y&^gjgHVb{5y6u=0z
Y24oKcKrH}LKrzMy2B5<k7&w@j0kIW17XSbN

literal 0
HcmV?d00001

diff --git a/stateless/invalidInput/crash-eb262bbb7d5a0b11e1a8112ea592eb67cc618c9b b/stateless/invalidInput/crash-eb262bbb7d5a0b11e1a8112ea592eb67cc618c9b
new file mode 100644
index 0000000000000000000000000000000000000000..5d6b306ebe4cc80d2bb246ae6f83251b102a075b
GIT binary patch
literal 137
hcmZQ%U}R(f;UC?M_6!j4|Cj>{-2fwFHzUygG5|Oz27CYj

literal 0
HcmV?d00001

diff --git a/stateless/invalidInput/crash-f9a002e389b120714a2324487a6e473b0fedc12e b/stateless/invalidInput/crash-f9a002e389b120714a2324487a6e473b0fedc12e
new file mode 100644
index 0000000000000000000000000000000000000000..83050f8a5a177523a3d503ddd10beffa81afac9f
GIT binary patch
literal 202
zcmZQ%U}DsAiWZ#P{j2}>jG*s9-&ZlO`Om<h|Av8qf$={OpefK_wdRdC3iyx23-+?<
zah!dD`$|V6P%DzThCprVugW%p?9l|;1LA)uQ$$j)uG17-9(rj<htukrFRs58QV95G
b^}gq+bD40<YGfBeOai(N1bBfaL0Aj`Mci5+

literal 0
HcmV?d00001

diff --git a/stateless/test_witness_keys.nim b/stateless/test_witness_keys.nim
index 93e160916..95904c0ac 100644
--- a/stateless/test_witness_keys.nim
+++ b/stateless/test_witness_keys.nim
@@ -1,5 +1,5 @@
 import
-  randutils, random, unittest2, stew/byteutils,
+  randutils, random, unittest2, stew/byteutils, os,
   eth/[common, rlp], eth/trie/[hexary, db, trie_defs, nibbles],
   faststreams/inputs, nimcrypto/sysrand,
   ../stateless/[witness_from_tree, tree_from_witness],
@@ -124,6 +124,14 @@ proc initMultiKeys(keys: openArray[string]): MultikeysRef =
       hash: hexToByteArray[32](x)
     )
 
+proc parseInvalidInput(payload: openArray[byte]): bool =
+  var db = newMemoryDB()
+  try:
+    var tb = initTreeBuilder(payload, db, {wfEIP170})
+    discard tb.buildTree()
+  except ParsingError, ContractCodeError:
+    result = true
+
 proc witnessKeysMain*() =
   suite "random keys block witness roundtrip test":
     randomize()
@@ -223,5 +231,10 @@ proc witnessKeysMain*() =
         mg.group.first == 2
         mg.group.last == 3
 
+    test "parse invalid input":
+      for x in walkDirRec("stateless" / "invalidInput"):
+        let z = readFile(x)
+        check parseInvalidInput(z.toOpenArrayByte(0, z.len-1))
+
 when isMainModule:
   witnessKeysMain()
diff --git a/stateless/tree_from_witness.nim b/stateless/tree_from_witness.nim
index d0ba39070..cc75c2bd5 100644
--- a/stateless/tree_from_witness.nim
+++ b/stateless/tree_from_witness.nim
@@ -218,7 +218,9 @@ proc buildForest*(t: var TreeBuilder): seq[KeccakHash]
     result.add KeccakHash(data: res.data)
 
 proc treeNode(t: var TreeBuilder, depth: int, storageMode = false): NodeKey =
-  assert(depth < 64)
+  if depth >= 64:
+    raise newException(ParsingError, "invalid trie structure")
+
   let nodeType = safeReadEnum(t, TrieNodeType)
   case nodeType
   of BranchNodeType: result = t.branchNode(depth, storageMode)
@@ -235,7 +237,9 @@ proc treeNode(t: var TreeBuilder, depth: int, storageMode = false): NodeKey =
     result = t.forceSmallNodeKeyToHash(result)
 
 proc branchNode(t: var TreeBuilder, depth: int, storageMode: bool): NodeKey =
-  assert(depth < 64)
+  if depth >= 64:
+    raise newException(ParsingError, "invalid trie structure")
+
   let mask = constructBranchMask(t.safeReadByte, t.safeReadByte)
 
   when defined(debugDepth):
@@ -305,9 +309,13 @@ func hexPrefixLeaf(r: var RlpWriter, x: openArray[byte], depth: int) =
   r.append toOpenArray(bytes, 0, nibblesLen div 2)
 
 proc extensionNode(t: var TreeBuilder, depth: int, storageMode: bool): NodeKey =
-  assert(depth < 63)
+  if depth >= 63:
+    raise newException(ParsingError, "invalid trie structure")
+
   let nibblesLen = t.safeReadByte().int
-  assert(nibblesLen < 65)
+  if nibblesLen > 64 or nibblesLen < 1:
+    raise newException(ParsingError, "nibblesLen should between 1..64")
+
   var r = initRlpList(2)
   let pathLen = nibblesLen div 2 + nibblesLen mod 2
   safeReadBytes(t, pathLen):
@@ -321,7 +329,9 @@ proc extensionNode(t: var TreeBuilder, depth: int, storageMode: bool): NodeKey =
     var hash: NodeKey
     toKeccak(hash, t.read(32))
 
-  assert(depth + nibblesLen < 65)
+  if nibblesLen + depth > 64 or nibblesLen + depth < 1:
+    raise newException(ParsingError, "depth should between 1..64")
+
   let nodeType = safeReadEnum(t, TrieNodeType)
   case nodeType
   of BranchNodeType: r.append t.branchNode(depth + nibblesLen, storageMode)
@@ -375,7 +385,8 @@ proc readByteCode(t: var TreeBuilder, acc: var Account) =
     acc.codeHash.data = codeHash.data
 
 proc accountNode(t: var TreeBuilder, depth: int): NodeKey =
-  assert(depth < 65)
+  if depth >= 65:
+    raise newException(ParsingError, "invalid trie structure")
 
   when defined(debugHash):
     let len = t.safeReadU32().int
@@ -443,7 +454,8 @@ proc readStorageSlot(t: var TreeBuilder): Hash256 =
     t.keys[^1].slots.add slot
 
 proc accountStorageLeafNode(t: var TreeBuilder, depth: int): NodeKey =
-  assert(depth < 65)
+  if depth >= 65:
+    raise newException(ParsingError, "invalid trie structure")
 
   when defined(debugHash):
     let len = t.safeReadU32().int