[libFuzzer] improve the afl driver and it's tests. Make it possible to run individual inputs with afl driver

llvm-svn: 303125

3 files changed, 77 insertions, 13 deletions

diff --git a/llvm/lib/Fuzzer/afl/afl_driver.cpp b/llvm/lib/Fuzzer/afl/afl_driver.cpp
index b3a54e57fce..3815ed11cf6 100644
--- a/llvm/lib/Fuzzer/afl/afl_driver.cpp
+++ b/llvm/lib/Fuzzer/afl/afl_driver.cpp
@@ -59,6 +59,11 @@ statistics from the file. If that fails then the process will quit.
 #include <signal.h>
 #include <sys/resource.h>
 #include <sys/time.h>
+
+#include <iostream>
+#include <fstream>
+#include <vector>
+
 // Platform detection. Copied from FuzzerInternal.h
 #ifdef __linux__
 #define LIBFUZZER_LINUX 1
@@ -245,17 +250,39 @@ extern "C" size_t LLVMFuzzerMutate(uint8_t *Data, size_t Size, size_t MaxSize) {
   return 0;
 }
 
+// Execute any files provided as parameters.
+int ExecuteFilesOnyByOne(int argc, char **argv) {
+  for (int i = 1; i < argc; i++) {
+    std::ifstream in(argv[i]);
+    in.seekg(0, in.end);
+    size_t length = in.tellg();
+    in.seekg (0, in.beg);
+    std::cout << "Reading " << length << " bytes from " << argv[i] << std::endl;
+    // Allocate exactly length bytes so that we reliably catch buffer overflows.
+    std::vector<char> bytes(length);
+    in.read(bytes.data(), bytes.size());
+    assert(in);
+    LLVMFuzzerTestOneInput(reinterpret_cast<const uint8_t *>(bytes.data()),
+                           bytes.size());
+    std::cout << "Execution successfull" << std::endl;
+  }
+  return 0;
+}
+
 int main(int argc, char **argv) {
-  fprintf(stderr, "======================= INFO =========================\n"
-                  "This binary is built for AFL-fuzz.\n"
-                  "To run the target function on a single input execute this:\n"
-                  "  %s < INPUT_FILE\n"
-                  "To run the fuzzing execute this:\n"
-                  "  afl-fuzz [afl-flags] %s [N] "
-                  "-- run N fuzzing iterations before "
-                  "re-spawning the process (default: 1000)\n"
-                  "======================================================\n",
-          argv[0], argv[0]);
+  fprintf(stderr,
+      "======================= INFO =========================\n"
+      "This binary is built for AFL-fuzz.\n"
+      "To run the target function on individual input(s) execute this:\n"
+      "  %s < INPUT_FILE\n"
+      "or\n"
+      "  %s INPUT_FILE1 [INPUT_FILE2 ... ]\n"
+      "To fuzz with afl-fuzz execute this:\n"
+      "  afl-fuzz [afl-flags] %s [-N]\n"
+      "afl-fuzz will run N iterations before "
+      "re-spawning the process (default: 1000)\n"
+      "======================================================\n",
+          argv[0], argv[0], argv[0]);
   if (LLVMFuzzerInitialize)
     LLVMFuzzerInitialize(&argc, &argv);
   // Do any other expensive one-time initialization here.
@@ -266,8 +293,14 @@ int main(int argc, char **argv) {
   __afl_manual_init();
 
   int N = 1000;
-  if (argc >= 2)
-    N = atoi(argv[1]);
+  if (argc == 2 && argv[1][0] == '-')
+      N = atoi(argv[1] + 1);
+  else if(argc == 2 && (N = atoi(argv[1])) > 0)
+      fprintf(stderr, "WARNING: using the deprecated call style `%s %d`\n",
+              argv[0], N);
+  else if (argc > 1)
+    return ExecuteFilesOnyByOne(argc, argv);
+
   assert(N > 0);
   time_t unit_time_secs;
   int num_runs = 0;
diff --git a/llvm/lib/Fuzzer/test/AFLDriverTest.cpp b/llvm/lib/Fuzzer/test/AFLDriverTest.cpp
index 3dd0b611730..e3f5f710088 100644
--- a/llvm/lib/Fuzzer/test/AFLDriverTest.cpp
+++ b/llvm/lib/Fuzzer/test/AFLDriverTest.cpp
@@ -4,19 +4,25 @@
 // Contains dummy functions used to avoid dependency on AFL.
 #include <stdint.h>
 #include <stdlib.h>
+#include <stdio.h>
 
 extern "C" void __afl_manual_init() {}
 
-extern "C" int __afl_persistent_loop(unsigned int) {
+extern "C" int __afl_persistent_loop(unsigned int N) {
+  static int Count = N;
+  fprintf(stderr, "__afl_persistent_loop calle, Count = %d\n", Count);
+  if (Count--) return 1;
   return 0;
 }
 
 // This declaration exists to prevent the Darwin linker
 // from complaining about this being a missing weak symbol.
 extern "C" int LLVMFuzzerInitialize(int *argc, char ***argv) {
+  fprintf(stderr, "LLVMFuzzerInitialize called\n");
   return 0;
 }
 
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size) {
+  fprintf(stderr, "LLVMFuzzerTestOneInput called; Size = %zd\n", Size);
   return 0;
 }
diff --git a/llvm/lib/Fuzzer/test/afl-driver.test b/llvm/lib/Fuzzer/test/afl-driver.test
new file mode 100644
index 00000000000..d5472f14655
--- /dev/null
+++ b/llvm/lib/Fuzzer/test/afl-driver.test
@@ -0,0 +1,25 @@
+RUN: echo -n "abc" > %t.file3
+RUN: echo -n "abcd" > %t.file4
+
+RUN: AFLDriverTest < %t.file3 2>&1 | FileCheck %s --check-prefix=CHECK1
+CHECK1: __afl_persistent_loop calle, Count = 1000
+CHECK1: LLVMFuzzerTestOneInput called; Size = 3
+
+
+RUN: AFLDriverTest < %t.file3 -42 2>&1 | FileCheck %s --check-prefix=CHECK2
+CHECK2: __afl_persistent_loop calle, Count = 42
+CHECK2: LLVMFuzzerTestOneInput called; Size = 3
+
+
+RUN: AFLDriverTest < %t.file3 666 2>&1 | FileCheck %s --check-prefix=CHECK3
+CHECK3: WARNING: using the deprecated call style
+CHECK3: __afl_persistent_loop calle, Count = 666
+CHECK3: LLVMFuzzerTestOneInput called; Size = 3
+
+
+RUN: AFLDriverTest %t.file3 2>&1 | FileCheck %s --check-prefix=CHECK4
+CHECK4: LLVMFuzzerTestOneInput called; Size = 3
+
+RUN: AFLDriverTest %t.file3 %t.file4  2>&1 | FileCheck %s --check-prefix=CHECK5
+CHECK5: LLVMFuzzerTestOneInput called; Size = 3
+CHECK5: LLVMFuzzerTestOneInput called; Size = 4