From 2a91a8d0cc3414f27c92824525c1f9a7dc2485f9 Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Wed, 25 Sep 2019 17:02:43 +0200
Subject: [PATCH] Add convenience init and test macro to create fuzzing tests

---
 tests/fuzzing/discovery/fuzz.nim | 58 ++++++++------------------------
 tests/fuzzing/fuzz_helpers.nim   | 44 +++++++++++++++++++++---
 2 files changed, 54 insertions(+), 48 deletions(-)

diff --git a/tests/fuzzing/discovery/fuzz.nim b/tests/fuzzing/discovery/fuzz.nim
index 214df7e..ba45379 100644
--- a/tests/fuzzing/discovery/fuzz.nim
+++ b/tests/fuzzing/discovery/fuzz.nim
@@ -5,7 +5,7 @@ import
 const DefaultListeningPort = 30303
 var targetNode: DiscoveryProtocol
 
-proc init() =
+init:
   # Set up a discovery node, this is the node we target when fuzzing
   var
     targetNodeKey = initPrivateKey("a2b50376a79b1a8c8a3296485572bdfbf54708bb46d3c25d73d2723aaaf6a617")
@@ -14,51 +14,21 @@ proc init() =
   # Create the transport as else replies on the messages send will fail.
   targetNode.open()
 
-proc test(payload: seq[byte]) =
+test:
   var
     msg: seq[byte]
     address: Address
 
-  fuzz:
-    # Sending raw payload is possible but won't find us much. We need a hash and
-    # a signature, and without it there is a big chance it will always result in
-    # "Wrong msg mac from" error.
-    let nodeKey = initPrivateKey("a2b50376a79b1a8c8a3296485572bdfbf54708bb46d3c25d73d2723aaaf6a618")
-    msg = packData(payload, nodeKey)
-    address = localAddress(DefaultListeningPort + 1)
+  # Sending raw payload is possible but won't find us much. We need a hash and
+  # a signature, and without it there is a big chance it will always result in
+  # "Wrong msg mac from" error.
+  let nodeKey = initPrivateKey("a2b50376a79b1a8c8a3296485572bdfbf54708bb46d3c25d73d2723aaaf6a618")
+  msg = packData(payload, nodeKey)
+  address = localAddress(DefaultListeningPort + 1)
 
-    try:
-      targetNode.receive(address, msg)
-    # These errors are also catched in `processClient` in discovery.nim
-    # TODO: move them a layer down in discovery so we can do a cleaner test there?
-    except RlpError, DiscProtocolError:
-      debug "Receive failed", err = getCurrentExceptionMsg()
-
-proc NimMain() {.importc: "NimMain".}
-
-proc fuzzerInit(): cint {.exportc: "LLVMFuzzerInitialize".} =
-  NimMain()
-
-  init()
-
-  return 0
-
-template `+`*[T](p: ptr T, off: int): ptr T =
-  cast[ptr type(p[])](cast[ByteAddress](p) +% off * sizeof(p[]))
-
-proc fuzzerCall(data: ptr byte, len: csize): cint {.exportc: "LLVMFuzzerTestOneInput".} =
-  if len > 0:
-    var input: seq[byte]
-
-    # TODO: something better to get this data in the seq?
-    newSeq(input, len)
-    for i in 0..<len:
-      input[i] = (data + i)[]
-
-    test(input)
-
-  return 0
-
-when defined(afl):
-  init()
-  test(readStdin())
\ No newline at end of file
+  try:
+    targetNode.receive(address, msg)
+  # These errors are also catched in `processClient` in discovery.nim
+  # TODO: move them a layer down in discovery so we can do a cleaner test there?
+  except RlpError, DiscProtocolError:
+    debug "Receive failed", err = getCurrentExceptionMsg()
\ No newline at end of file
diff --git a/tests/fuzzing/fuzz_helpers.nim b/tests/fuzzing/fuzz_helpers.nim
index 5ad3497..6020270 100644
--- a/tests/fuzzing/fuzz_helpers.nim
+++ b/tests/fuzzing/fuzz_helpers.nim
@@ -1,6 +1,6 @@
-import streams, posix, sequtils, strutils, chronicles
+import streams, posix, strutils, chronicles, macros
 
-template fuzz*(body) =
+template fuzz(body) =
   # For code we want to fuzz, SIGSEGV is needed on unwanted exceptions.
   # However, this is only needed when fuzzing with afl.
   when defined(afl):
@@ -23,5 +23,41 @@ proc readStdin*(): seq[byte] =
   s.close()
   # Remove newline if it is there
   input.removeSuffix
-  # TODO: is there a better/faster way?
-  result = input.mapIt(it.byte)
\ No newline at end of file
+  result = cast[seq[byte]](input)
+
+proc NimMain() {.importc: "NimMain".}
+
+template `+`*[T](p: ptr T, off: int): ptr T =
+  cast[ptr type(p[])](cast[ByteAddress](p) +% off * sizeof(p[]))
+
+macro test*(body: untyped): untyped =
+  when defined(afl):
+    result = quote do:
+      var payload {.inject.} = readStdin()
+
+      fuzz: `body`
+  else:
+    result = quote do:
+      proc fuzzerCall(data: ptr byte, len: csize):
+          cint {.exportc: "LLVMFuzzerTestOneInput".} =
+        var payload {.inject.} : seq[byte]
+        if len > 0:
+          # TODO: something better to get this data in the seq?
+          newSeq(payload, len)
+          for i in 0..<len:
+            payload[i] = (data + i)[]
+
+        `body`
+
+macro init*(body: untyped): untyped =
+  when defined(afl):
+    result = quote do:
+      fuzz: `body`
+  else:
+    result = quote do:
+      proc fuzzerInit(): cint {.exportc: "LLVMFuzzerInitialize".} =
+        NimMain()
+
+        `body`
+
+        return 0