Merge pull request #44 from trailofbits/take-over-api

Take-over API

bin/deepstate/common.py

@@ -121,6 +121,10 @@ class DeepState(object):
         "--output_test_dir", default="out", type=str, required=False,
         help="Directory where tests will be saved.")
 
+    parser.add_argument(
+        "--take_over", action='store_true',
+        help="Explore the program starting at the `TakeOver` hook.")
+
     parser.add_argument(
         "binary", type=str, help="Path to the test binary to run.")
 

bin/deepstate/main_angr.py

@@ -19,7 +19,7 @@ import logging
 import multiprocessing
 import sys
 import traceback
-from .common import DeepState
+from .common import DeepState, TestInfo
 
 L = logging.getLogger("deepstate.angr")
 L.setLevel(logging.INFO)
@@ -261,12 +261,22 @@ class Log(angr.SimProcedure):
     DeepAngr(procedure=self).api_log(level, ea)
 
 
-def do_run_test(project, test, apis, run_state):
+class TakeOver(angr.SimProcedure):
+    def run(self):
+        """Do nothing, returning 1 to indicate that `DeepState_TakeOver()` has
+        been hooked for symbolic execution."""
+        return 1
+
+
+def do_run_test(project, test, apis, run_state, should_call_state):
   """Symbolically executes a single test function."""
 
-  test_state = project.factory.call_state(
-      test.ea,
-      base_state=run_state)
+  if should_call_state:
+    test_state = project.factory.call_state(
+        test.ea,
+        base_state=run_state)
+  else:
+      test_state = run_state
 
   mc = DeepAngr(state=test_state)
   mc.begin_test(test)
@@ -291,10 +301,10 @@ def do_run_test(project, test, apis, run_state):
     da.crash_test()
     da.report()
 
-def run_test(project, test, apis, run_state):
+def run_test(project, test, apis, run_state, should_call_state=True):
   """Symbolically executes a single test function."""
   try:
-    do_run_test(project, test, apis, run_state)
+    do_run_test(project, test, apis, run_state, should_call_state)
   except Exception as e:
     L.error("Uncaught exception: {}\n{}".format(e, traceback.format_exc()))
 
@@ -314,53 +324,8 @@ def find_symbol_ea(project, name):
 
   return 0
 
-def main():
-  """Run DeepState."""
-  args = DeepAngr.parse_args()
-
-  try:
-    project = angr.Project(
-        args.binary,
-        use_sim_procedures=True,
-        translation_cache=True,
-        support_selfmodifying_code=False,
-        auto_load_libs=True,
-        exclude_sim_procedures_list=['printf', '__printf_chk',
-                                     'vprintf', '__vprintf_chk',
-                                     'fprintf', '__fprintf_chk',
-                                     'vfprintf', '__vfprintf_chk',
-                                     'puts', 'abort', '__assert_fail',
-                                     '__stack_chk_fail'])
-  except Exception as e:
-    L.critical("Cannot create Angr instance on binary {}: {}".format(
-        args.binary, e))
-    return 1
-
-  setup_ea = find_symbol_ea(project, 'DeepState_Setup')
-  if not setup_ea:
-    L.critical("Cannot find symbol `DeepState_Setup` in binary `{}`".format(
-        args.binary))
-    return 1
-
-  entry_state = project.factory.entry_state(
-      add_options={angr.options.ZERO_FILL_UNCONSTRAINED_MEMORY,
-                   angr.options.STRICT_PAGE_ACCESS})
-
-  addr_size_bits = entry_state.arch.bits
-
-  # Concretely execute up until `DeepState_Setup`.
-  concrete_manager = angr.SimulationManager(
-        project=project,
-        active_states=[entry_state])
-  concrete_manager.explore(find=setup_ea)
-
-  try:
-    run_state = concrete_manager.found[0]
-  except:
-    L.critical("Execution never hit `DeepState_Setup` in binary `{}`".format(
-        args.binary))
-    return 1
 
+def hook_apis(project, run_state):
   # Read the API table, which will tell us about the location of various
   # symbols. Technically we can look these up with the `labels.lookup` API,
   # but we have the API table for Manticore-compatibility, so we may as well
@@ -392,6 +357,89 @@ def main():
   hook_function(project, apis['ClearStream'], ClearStream)
   hook_function(project, apis['LogStream'], LogStream)
 
+  return mc, apis
+
+
+def main_take_over(args, project):
+  takeover_ea = find_symbol_ea(project, 'DeepState_TakeOver')
+
+  hook_function(project, takeover_ea, TakeOver)
+
+  if not takeover_ea:
+    L.critical("Cannot find symbol `DeepState_TakeOver` in binary `{}`".format(
+        args.binary))
+    return 1
+
+  entry_state = project.factory.entry_state(
+      add_options={angr.options.ZERO_FILL_UNCONSTRAINED_MEMORY,
+                   angr.options.STRICT_PAGE_ACCESS})
+
+  addr_size_bits = entry_state.arch.bits
+
+  # Concretely execute up until `DeepState_TakeOver`.
+  concrete_manager = angr.SimulationManager(
+        project=project,
+        active_states=[entry_state])
+  concrete_manager.explore(find=takeover_ea)
+
+  try:
+    takeover_state = concrete_manager.found[0]
+  except:
+    L.critical("Execution never hit `DeepState_TakeOver` in binary `{}`".format(
+        args.binary))
+    return 1
+
+  try:
+    run_state = takeover_state.step().successors[0]
+  except:
+    L.critical("Unable to exit from `DeepState_TakeOver` in binary `{}`".format(
+        args.binary))
+    return 1
+
+  # Read the API table, which will tell us about the location of various
+  # symbols. Technically we can look these up with the `labels.lookup` API,
+  # but we have the API table for Manticore-compatibility, so we may as well
+  # use it.
+  ea_of_api_table = find_symbol_ea(project, 'DeepState_API')
+  if not ea_of_api_table:
+    L.critical("Could not find API table in binary `{}`".format(args.binary))
+    return 1
+
+  _, apis = hook_apis(project, run_state)
+  fake_test = TestInfo(takeover_ea, '_takeover_test', '_takeover_file', 0)
+
+  return run_test(project, fake_test, apis, run_state, should_call_state=False)
+
+
+def main_unit_test(args, project):
+  setup_ea = find_symbol_ea(project, 'DeepState_Setup')
+  if not setup_ea:
+    L.critical("Cannot find symbol `DeepState_Setup` in binary `{}`".format(
+        args.binary))
+    return 1
+
+  entry_state = project.factory.entry_state(
+      add_options={angr.options.ZERO_FILL_UNCONSTRAINED_MEMORY,
+                   angr.options.STRICT_PAGE_ACCESS})
+
+  addr_size_bits = entry_state.arch.bits
+
+  # Concretely execute up until `DeepState_Setup`.
+  concrete_manager = angr.SimulationManager(
+        project=project,
+        active_states=[entry_state])
+  concrete_manager.explore(find=setup_ea)
+
+  try:
+    run_state = concrete_manager.found[0]
+  except:
+    L.critical("Execution never hit `DeepState_Setup` in binary `{}`".format(
+        args.binary))
+    return 1
+
+  # Hook the DeepState API functions.
+  mc, apis = hook_apis(project, run_state)
+
   # Find the test cases that we want to run.
   tests = mc.find_test_cases()
   del mc
@@ -414,5 +462,34 @@ def main():
 
   return 0
 
+
+def main():
+  """Run DeepState."""
+  args = DeepAngr.parse_args()
+
+  try:
+    project = angr.Project(
+        args.binary,
+        use_sim_procedures=True,
+        translation_cache=True,
+        support_selfmodifying_code=False,
+        auto_load_libs=True,
+        exclude_sim_procedures_list=['printf', '__printf_chk',
+                                     'vprintf', '__vprintf_chk',
+                                     'fprintf', '__fprintf_chk',
+                                     'vfprintf', '__vfprintf_chk',
+                                     'puts', 'abort', '__assert_fail',
+                                     '__stack_chk_fail'])
+  except Exception as e:
+    L.critical("Cannot create Angr instance on binary {}: {}".format(
+        args.binary, e))
+    return 1
+
+  if args.take_over:
+    return main_take_over(args, project)
+  else:
+    return main_unit_test(args, project)
+
+
 if "__main__" == __name__:
   exit(main())

bin/deepstate/main_manticore.py

@@ -21,14 +21,14 @@ import sys
 try:
   import manticore
 except Exception as e:
-  if "Z3NotFoundError" in repr(type(e)):  
+  if "Z3NotFoundError" in repr(type(e)):
     print "Manticore requires Z3 to be installed."
     sys.exit(255)
   else:
     raise
 import multiprocessing
 import traceback
-from .common import DeepState
+from .common import DeepState, TestInfo
 
 from manticore.core.state import TerminateState
 from manticore.utils.helpers import issymbolic
@@ -246,6 +246,12 @@ def hook_Log(state, level, ea):
   DeepManticore(state).api_log(level, ea)
 
 
+def hook_TakeOver(state):
+  """Implements `DeepState_TakeOver`, returning 1 to indicate that it was
+  hooked for symbolic execution."""
+  return 1
+
+
 def hook(func):
   return lambda state: state.invoke_model(func)
 
@@ -332,6 +338,10 @@ def do_run_test(state, apis, test):
   m.add_hook(apis['ClearStream'], hook(hook_ClearStream))
   m.add_hook(apis['LogStream'], hook(hook_LogStream))
 
+  # Here we hook `DeepState_TakeOver()`, even if running unit tests.
+  # In that case, we simply will never hit this hooked function model.
+  m.add_hook(test.ea, hook(hook_TakeOver))
+
   m.subscribe('will_terminate_state', done_test)
   m.run()
 
@@ -364,22 +374,31 @@ def run_tests(args, state, apis):
   exit(0)
 
 
-def main():
-  args = DeepManticore.parse_args()
-
-  try:
-    m = manticore.Manticore(args.binary)
-  except Exception as e:
-    L.critical("Cannot create Manticore instance on binary {}: {}".format(
-        args.binary, e))
+def main_takeover(m, args):
+  takeover_ea = find_symbol_ea(m, 'DeepState_TakeOver')
+  if not takeover_ea:
+    L.critical("Cannot find symbol `DeepState_TakeOver` in binary `{}`".format(
+        args.binary))
     return 1
 
-  m.verbosity(1)
+  takeover_state = m._initial_state
 
-  # Hack to get around current broken _get_symbol_address
-  m._binary_type = 'not elf'
-  m._binary_obj = m._initial_state.platform.elf
+  mc = DeepManticore(takeover_state)
 
+  ea_of_api_table = find_symbol_ea(m, 'DeepState_API')
+  if not ea_of_api_table:
+    L.critical("Could not find API table in binary `{}`".format(args.binary))
+    return 1
+
+  apis = mc.read_api_table(ea_of_api_table)
+  del mc
+
+  fake_test = TestInfo(takeover_ea, '_takeover_test', '_takeover_file', 0)
+  m.add_hook(takeover_ea, lambda state: run_test(state, apis, fake_test))
+  m.run()
+
+
+def main_unit_test(m, args):
   setup_ea = find_symbol_ea(m, 'DeepState_Setup')
   if not setup_ea:
     L.critical("Cannot find symbol `DeepState_Setup` in binary `{}`".format(
@@ -397,9 +416,32 @@ def main():
 
   apis = mc.read_api_table(ea_of_api_table)
   del mc
+
   m.add_hook(setup_ea, lambda state: run_tests(args, state, apis))
   m.run()
 
 
+def main():
+  args = DeepManticore.parse_args()
+
+  try:
+    m = manticore.Manticore(args.binary)
+  except Exception as e:
+    L.critical("Cannot create Manticore instance on binary {}: {}".format(
+        args.binary, e))
+    return 1
+
+  m.verbosity(1)
+
+  # Hack to get around current broken _get_symbol_address
+  m._binary_type = 'not elf'
+  m._binary_obj = m._initial_state.platform.elf
+
+  if args.take_over:
+    return main_takeover(m, args)
+  else:
+    return main_unit_test(m, args)
+
+
 if "__main__" == __name__:
   exit(main())

examples/CMakeLists.txt

@@ -43,3 +43,6 @@ target_link_libraries(StreamingAndFormatting deepstate)
 add_executable(Squares Squares.c)
 target_link_libraries(Squares deepstate)
 set_target_properties(Squares PROPERTIES COMPILE_DEFINITIONS "DEEPSTATE_TEST")
+
+add_executable(TakeOver TakeOver.cpp)
+target_link_libraries(TakeOver deepstate)

examples/TakeOver.cpp

@@ -0,0 +1,45 @@
+/*
+ * Copyright (c) 2018 Trail of Bits, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+
+#include <deepstate/DeepState.hpp>
+
+using namespace deepstate;
+
+DEEPSTATE_NOINLINE void func(uint32_t x) {
+  CHECK_LT(x, 0x1234)
+    << "Found x=" << x << " was not greater than 0x1234.";
+
+  if (x < 0x1234) {
+    printf("hi\n");
+  } else {
+    printf("bye\n");
+  }
+}
+
+int main(int argc, char *argv[]) {
+  DeepState_InitOptions(argc, argv);
+
+  uint32_t x = 123;
+  func(x);  // Unexplored
+
+  DeepState_TakeOver();
+
+  Symbolic<uint32_t> y;
+  Symbolic<uint32_t> z;
+  func(y);  // Explored
+  func(z);  // Explored
+}

src/include/deepstate/DeepState.h

@@ -55,6 +55,7 @@ DEEPSTATE_BEGIN_EXTERN_C
 
 DECLARE_string(input_test_dir);
 DECLARE_string(output_test_dir);
+DECLARE_string(take_over);
 
 enum {
   DeepState_InputSize = 8192
@@ -309,6 +310,8 @@ struct DeepState_TestInfo {
 /* Pointer to the last registered `TestInfo` structure. */
 extern struct DeepState_TestInfo *DeepState_LastTestInfo;
 
+extern int DeepState_TakeOver(void);
+
 /* Defines the entrypoint of a test case. This creates a data structure that
  * contains the information about the test, and then creates an initializer
  * function that runs before `main` that registers the test entrypoint with
@@ -500,8 +503,8 @@ DeepState_ForkAndRunTest(struct DeepState_TestInfo *test) {
 /* Run a single saved test case with input initialized from the file
  * `name` in directory `dir`. */
 static enum DeepState_TestRunResult
-DeepState_DoRunSavedTestCase(struct DeepState_TestInfo *test, const char *dir,
-                             const char *name) {
+DeepState_RunSavedTestCase(struct DeepState_TestInfo *test, const char *dir,
+                           const char *name) {
   size_t path_len = 2 + sizeof(char) * (strlen(dir) + strlen(name));
   char *path = (char *) malloc(path_len);
   if (path == NULL) {
@@ -530,6 +533,50 @@ DeepState_DoRunSavedTestCase(struct DeepState_TestInfo *test, const char *dir,
   return result;
 }
 
+/* Run a single test many times, initialized against each saved test case in
+ * `FLAGS_input_test_dir`. */
+static int DeepState_RunSavedCasesForTest(struct DeepState_TestInfo *test) {
+  int num_failed_tests = 0;
+  const char *test_file_name = basename((char *) test->file_name);
+
+  size_t test_case_dir_len = 3 + strlen(FLAGS_input_test_dir)
+    + strlen(test_file_name) + strlen(test->test_name);
+  char *test_case_dir = (char *) malloc(test_case_dir_len);
+  if (test_case_dir == NULL) {
+    DeepState_Abandon("Error allocating memory");
+  }
+  snprintf(test_case_dir, test_case_dir_len, "%s/%s/%s",
+           FLAGS_input_test_dir, test_file_name, test->test_name);
+
+  struct dirent *dp;
+  DIR *dir_fd;
+
+  dir_fd = opendir(test_case_dir);
+  if (dir_fd == NULL) {
+    DeepState_LogFormat(DeepState_LogInfo,
+                        "Skipping test `%s`, no saved test cases",
+                        test->test_name);
+    free(test_case_dir);
+    return 0;
+  }
+
+  /* Read generated test cases and run a test for each file found. */
+  while ((dp = readdir(dir_fd)) != NULL) {
+    if (IsTestCaseFile(dp->d_name)) {
+      enum DeepState_TestRunResult result =
+        DeepState_RunSavedTestCase(test, test_case_dir, dp->d_name);
+
+      if (result != DeepState_TestRunPass) {
+        num_failed_tests++;
+      }
+    }
+  }
+  closedir(dir_fd);
+  free(test_case_dir);
+
+  return num_failed_tests;
+}
+
 /* Run tests with saved input from `FLAGS_input_test_dir`.
  *
  * For each test unit and case, see if there are input files in the
@@ -542,41 +589,7 @@ static int DeepState_RunSavedTestCases(void) {
   DeepState_Setup();
 
   for (test = DeepState_FirstTest(); test != NULL; test = test->prev) {
-    const char *test_file_name = basename((char *) test->file_name);
-
-    size_t test_case_dir_len = 3 + strlen(FLAGS_input_test_dir)
-                             + strlen(test_file_name) + strlen(test->test_name);
-    char *test_case_dir = (char *) malloc(test_case_dir_len);
-    if (test_case_dir == NULL) {
-      DeepState_Abandon("Error allocating memory");
-    }
-    snprintf(test_case_dir, test_case_dir_len, "%s/%s/%s",
-             FLAGS_input_test_dir, test_file_name, test->test_name);
-
-    struct dirent *dp;
-    DIR *dir_fd;
-
-    dir_fd = opendir(test_case_dir);
-    if (dir_fd == NULL) {
-      DeepState_LogFormat(DeepState_LogInfo,
-                          "Skipping test `%s`, no saved test cases",
-                          test->test_name);
-      continue;
-    }
-
-    /* Read generated test cases and run a test for each file found. */
-    while ((dp = readdir(dir_fd)) != NULL) {
-      if (IsTestCaseFile(dp->d_name)) {
-        enum DeepState_TestRunResult result =
-          DeepState_DoRunSavedTestCase(test, test_case_dir, dp->d_name);
-
-        if (result != DeepState_TestRunPass) {
-          num_failed_tests++;
-        }
-      }
-    }
-    closedir(dir_fd);
-    free(test_case_dir);
+    num_failed_tests += DeepState_RunSavedCasesForTest(test);
   }
 
   DeepState_Teardown();

src/lib/DeepState.c

@@ -30,6 +30,7 @@ DEFINE_uint(num_workers, 1,
 
 DEFINE_string(input_test_dir, "", "Directory of saved tests to run.");
 DEFINE_string(output_test_dir, "", "Directory where tests will be saved.");
+DEFINE_string(take_over, "", "Replay test cases in take-over mode.");
 
 /* Pointer to the last registers DeepState_TestInfo data structure */
 struct DeepState_TestInfo *DeepState_LastTestInfo = NULL;
@@ -64,7 +65,13 @@ void DeepState_Crash(void) {
 DEEPSTATE_NORETURN
 void DeepState_Fail(void) {
   DeepState_TestFailed = 1;
-  longjmp(DeepState_ReturnToRun, 1);
+
+  if (FLAGS_take_over) {
+    // We want to communicate the failure to a parent process, so exit.
+    exit(DeepState_TestRunFail);
+  } else {
+    longjmp(DeepState_ReturnToRun, 1);
+  }
 }
 
 /* Mark this test as passing. */
@@ -364,6 +371,99 @@ void DrMemFuzzFunc(volatile uint8_t *buff, size_t size) {
   }
 }
 
+void DeepState_RunSavedTakeOverCases(jmp_buf env,
+                                     struct DeepState_TestInfo *test) {
+  int num_failed_tests = 0;
+  const char *test_case_dir = FLAGS_input_test_dir;
+
+  DIR *dir_fd = opendir(test_case_dir);
+  if (dir_fd == NULL) {
+    DeepState_LogFormat(DeepState_LogInfo,
+                        "Skipping test `%s`, no saved test cases",
+                        test->test_name);
+    return;
+  }
+
+  struct dirent *dp;
+
+  /* Read generated test cases and run a test for each file found. */
+  while ((dp = readdir(dir_fd)) != NULL) {
+    if (IsTestCaseFile(dp->d_name)) {
+      pid_t case_pid = fork();
+      if (!case_pid) {
+        DeepState_Begin(test);
+
+        size_t path_len = 2 + sizeof(char) * (strlen(test_case_dir) +
+                                              strlen(dp->d_name));
+        char *path = (char *) malloc(path_len);
+        if (path == NULL) {
+          DeepState_Abandon("Error allocating memory");
+        }
+        snprintf(path, path_len, "%s/%s", test_case_dir, dp->d_name);
+        InitializeInputFromFile(path);
+        free(path);
+
+        longjmp(env, 1);
+      }
+
+      int wstatus;
+      waitpid(case_pid, &wstatus, 0);
+
+      /* If we exited normally, the status code tells us if the test passed. */
+      if (WIFEXITED(wstatus)) {
+        uint8_t status = WEXITSTATUS(wstatus);
+
+        switch (status) {
+        case DeepState_TestRunPass:
+          DeepState_LogFormat(DeepState_LogInfo,
+                              "Passed: TakeOver test with data from `%s`",
+                              dp->d_name);
+          break;
+        case DeepState_TestRunFail:
+          DeepState_LogFormat(DeepState_LogError,
+                              "Failed: TakeOver test with data from `%s`",
+                              dp->d_name);
+          break;
+        case DeepState_TestRunAbandon:
+          DeepState_LogFormat(DeepState_LogError,
+                              "Abandoned: TakeOver test with data from `%s`",
+                              dp->d_name);
+          break;
+        default:
+          DeepState_LogFormat(DeepState_LogError,
+                              "Unknown exit code from test with data from `%s`",
+                              dp->d_name);
+        }
+      } else {
+        /* If here, we exited abnormally but didn't catch it in the signal
+         * handler, and thus the test failed due to a crash. */
+        DeepState_LogFormat(DeepState_LogError,
+                            "Crashed: TakeOver test with data from `%s`",
+                            dp->d_name);
+      }
+    }
+  }
+  closedir(dir_fd);
+}
+
+int DeepState_TakeOver(void) {
+  struct DeepState_TestInfo test = {
+    .prev = NULL,
+    .test_func = NULL,
+    .test_name = "__takeover_test",
+    .file_name = "__takeover_file",
+    .line_number = 0,
+  };
+
+  jmp_buf env;
+  if (!setjmp(env)) {
+    DeepState_RunSavedTakeOverCases(env, &test);
+    exit(0);
+  }
+
+  return 0;
+}
+
 /* Notify that we're about to begin a test while running under Dr. Fuzz. */
 void DeepState_BeginDrFuzz(struct DeepState_TestInfo *test) {
   DeepState_DrFuzzTest = test;