Merge pull request #86 from trailofbits/add_input_test_file_option

Add option to run a single file, for test replay and for file-based fuzzing

README.md

@@ -89,6 +89,47 @@ deepstate-angr --num_workers 4 --output_test_dir out $DEEPSTATE/build/examples/I
 
 DeepState consists of a static library, used to write test harnesses, and command-line _executors_ written in Python. At this time, the best documentation is in the [examples](/examples) and in our [paper](https://agroce.github.io/bar18.pdf).
 
+## Fuzzing
+
+DeepState now can be used with a file-based fuzzer (e.g. AFL).  There
+are a few steps to this.  First, compile DeepState itself with any
+needed instrumentation.  E.g., to use it with AFL, you might want to add
+something like:
+
+```
+SET(CMAKE_C_COMPILER /usr/local/bin/afl-gcc)
+SET(CMAKE_CXX_COMPILER /usr/local/bin/afl-g++)
+```
+
+to `deepstate/CMakeLists.txt`.  Second, do the same for your DeepState
+test harness and any code it links to you want instrumented.  Finally, run the fuzzing via the
+interface to replay test files.  For example, to fuzz the `OneOf`
+example, if we were in the `deepstate/build/examples` directory, you
+would do something like:
+
+```shell
+afl-fuzz -d -i corpus -o afl_OneOf -- ./OneOf --input_test_file @@ --abort_on_fail
+```
+
+where `corpus` contains at least one file to start fuzzing from.  The
+file needs to be smaller than the DeepState input size limit, but has
+few other limitations (for AFL it should also not cause test
+failure).  The `abort_on_fail` flag makes DeepState crashes and failed
+tests appear as crashes to the fuzzer.
+
+To replay the tests from AFL:
+
+```shell
+./OneOf --input_test_files_dir afl_OneOf/crashes
+./OneOf --input_test_files_dir afl_OneOf/queue
+```
+
+Finally, if an example has more than one test, you need to specify,
+with a fully qualified name (e.g.,
+`Arithmetic_InvertibleMultiplication_CanFail`), which test to run,
+using the `--input_which_test` flag to the binary.  By
+default, DeepState will run the first test defined.
+
 ## Contributing
 
 All accepted PRs are awarded bounties by Trail of Bits. Join the #deepstate channel on the [Empire Hacking Slack](https://empireslacking.herokuapp.com/) to discuss ongoing development and claim bounties. Check the [good first issue](https://github.com/trailofbits/deepstate/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22) label for suggested contributions.

src/include/deepstate/DeepState.h

@@ -55,9 +55,13 @@
 DEEPSTATE_BEGIN_EXTERN_C
 
 DECLARE_string(input_test_dir);
+DECLARE_string(input_test_file);
+DECLARE_string(input_test_files_dir);
+DECLARE_string(input_which_test);
 DECLARE_string(output_test_dir);
 
 DECLARE_bool(take_over);
+DECLARE_bool(abort_on_fail);
 
 enum {
   DeepState_InputSize = 8192
@@ -518,7 +522,11 @@ DeepState_RunSavedTestCase(struct DeepState_TestInfo *test, const char *dir,
   if (path == NULL) {
     DeepState_Abandon("Error allocating memory");
   }
+  if (strncmp(dir, "", strlen(dir)) != 0) {
     snprintf(path, path_len, "%s/%s", dir, name);
+  } else {
+    snprintf(path, path_len, "%s", name);
+  }
 
   DeepState_InitInputFromFile(path);
 
@@ -585,7 +593,113 @@ static int DeepState_RunSavedCasesForTest(struct DeepState_TestInfo *test) {
   return num_failed_tests;
 }
 
-/* Run tests with saved input from `FLAGS_input_test_dir`.
+/* Run test from `FLAGS_input_test_file`, under `FLAGS_input_which_test`
+ * or first test, if not defined. */
+static int DeepState_RunSingleSavedTestCase(void) {
+  int num_failed_tests = 0;
+  struct DeepState_TestInfo *test = NULL;
+
+  DeepState_Setup();
+
+  for (test = DeepState_FirstTest(); test != NULL; test = test->prev) {
+    if (HAS_FLAG_input_which_test) {
+      if (strncmp(FLAGS_input_which_test, test->test_name, strlen(FLAGS_input_which_test)) == 0) {
+	break;
+      }
+    } else {
+      DeepState_LogFormat(DeepState_LogInfo,
+			  "No test specified, defaulting to first test");
+      break;
+    }
+  }
+
+  if (test == NULL) {
+    DeepState_LogFormat(DeepState_LogInfo,
+                        "Could not find matching test for %s",
+                        FLAGS_input_which_test);
+    return 0;
+  }
+
+  enum DeepState_TestRunResult result =
+    DeepState_RunSavedTestCase(test, "", FLAGS_input_test_file);
+
+  if ((result == DeepState_TestRunFail) || (result == DeepState_TestRunCrash)) {    
+    if (FLAGS_abort_on_fail) {
+      abort();
+    }
+    num_failed_tests++;
+  }
+
+  DeepState_Teardown();
+
+  return num_failed_tests;
+}
+
+/* Run tests from `FLAGS_input_test_files_dir`, under `FLAGS_input_which_test`
+ * or first test, if not defined. */
+static int DeepState_RunSingleSavedTestDir(void) {
+  int num_failed_tests = 0;
+  struct DeepState_TestInfo *test = NULL;  
+
+  DeepState_Setup();
+
+  for (test = DeepState_FirstTest(); test != NULL; test = test->prev) {
+    if (HAS_FLAG_input_which_test) {
+      if (strncmp(FLAGS_input_which_test, test->test_name, strlen(FLAGS_input_which_test)) == 0) {
+	break;
+      }
+    } else {
+      DeepState_LogFormat(DeepState_LogInfo,
+			  "No test specified, defaulting to first test");
+      break;
+    }
+  }
+
+  if (test == NULL) {
+    DeepState_LogFormat(DeepState_LogInfo,
+                        "Could not find matching test for %s",
+                        FLAGS_input_which_test);
+    return 0;
+  }
+
+  struct dirent *dp;
+  DIR *dir_fd;
+
+  struct stat path_stat;
+
+  dir_fd = opendir(FLAGS_input_test_files_dir);
+  if (dir_fd == NULL) {
+    DeepState_LogFormat(DeepState_LogInfo,
+                        "No tests to run.");
+    return 0;
+  }
+
+  /* Read generated test cases and run a test for each file found. */
+  while ((dp = readdir(dir_fd)) != NULL) {
+    size_t path_len = 2 + sizeof(char) * (strlen(FLAGS_input_test_files_dir) + strlen(dp->d_name));
+    char *path = (char *) malloc(path_len);
+    snprintf(path, path_len, "%s/%s", FLAGS_input_test_files_dir, dp->d_name);    
+    stat(path, &path_stat);
+    
+    if (S_ISREG(path_stat.st_mode)) {
+      enum DeepState_TestRunResult result =
+        DeepState_RunSavedTestCase(test, FLAGS_input_test_files_dir, dp->d_name);
+
+      if ((result == DeepState_TestRunFail) || (result == DeepState_TestRunCrash)) {
+	if (FLAGS_abort_on_fail) {
+	  abort();
+	}
+	
+        num_failed_tests++;
+      }
+    }
+  }
+  closedir(dir_fd);
+
+  return num_failed_tests;
+}
+
+/* Run test `FLAGS_input_which_test` with saved input from `FLAGS_input_test_file`.
  *
  * For each test unit and case, see if there are input files in the
  * expected directories. If so, use them to initialize
@@ -615,6 +729,14 @@ static int DeepState_Run(void) {
     return DeepState_RunSavedTestCases();
   }
 
+  if (HAS_FLAG_input_test_file) {
+    return DeepState_RunSingleSavedTestCase();
+  }
+
+  if (HAS_FLAG_input_test_files_dir) {
+    return DeepState_RunSingleSavedTestDir();
+  }  
+
   int num_failed_tests = 0;
   int use_drfuzz = getenv("DYNAMORIO_EXE_PATH") != NULL;
   struct DeepState_TestInfo *test = NULL;

src/include/deepstate/DeepState.hpp

@@ -329,7 +329,7 @@ template <typename... FuncTys>
 inline static void OneOf(FuncTys&&... funcs) {
   std::function<void(void)> func_arr[sizeof...(FuncTys)] = {funcs...};
   unsigned index = DeepState_UIntInRange(
-      0U, static_cast<unsigned>(sizeof...(funcs)));
+      0U, static_cast<unsigned>(sizeof...(funcs))-1);
   func_arr[Pump(index, sizeof...(funcs))]();
 }
 

src/lib/DeepState.c

@@ -29,9 +29,13 @@ DEFINE_uint(num_workers, 1,
             "Number of workers to spawn for testing and test generation.");
 
 DEFINE_string(input_test_dir, "", "Directory of saved tests to run.");
+DEFINE_string(input_which_test, "", "Test to use with --input_test_file or --input_test_files_dir.");
+DEFINE_string(input_test_file, "", "Saved test to run.");
+DEFINE_string(input_test_files_dir, "", "Directory of saved test files to run (flat structure).");
 DEFINE_string(output_test_dir, "", "Directory where tests will be saved.");
 
 DEFINE_bool(take_over, false, "Replay test cases in take-over mode.");
+DEFINE_bool(abort_on_fail, false, "Abort on file replay failure (useful in file fuzzing).");
 
 /* Pointer to the last registers DeepState_TestInfo data structure */
 struct DeepState_TestInfo *DeepState_LastTestInfo = NULL;