1 files changed, 62 insertions, 4 deletions

diff --git a/src/afl-tmin.c b/src/afl-tmin.c
index fc974262..6656712a 100644
--- a/src/afl-tmin.c
+++ b/src/afl-tmin.c
@@ -640,9 +640,11 @@ static void handle_stop_sig(int sig) {
 
 /* Do basic preparations - persistent fds, filenames, etc. */
 
-static void set_up_environment(afl_forkserver_t *fsrv) {
+static void set_up_environment(afl_forkserver_t *fsrv, char **argv) {
 
-  u8 *x;
+  u8 *  x;
+  char *afl_preload;
+  char *frida_afl_preload = NULL;
 
   fsrv->dev_null_fd = open("/dev/null", O_RDWR);
   if (fsrv->dev_null_fd < 0) { PFATAL("Unable to open /dev/null"); }
@@ -712,6 +714,18 @@ static void set_up_environment(afl_forkserver_t *fsrv) {
 
   }
 
+  x = get_afl_env("LSAN_OPTIONS");
+
+  if (x) {
+
+    if (!strstr(x, "symbolize=0")) {
+
+      FATAL("Custom LSAN_OPTIONS set without symbolize=0 - please fix!");
+
+    }
+
+  }
+
   setenv("ASAN_OPTIONS",
          "abort_on_error=1:"
          "detect_leaks=0:"
@@ -749,12 +763,38 @@ static void set_up_environment(afl_forkserver_t *fsrv) {
                          "handle_sigfpe=0:"
                          "handle_sigill=0", 0);
 
+  setenv("LSAN_OPTIONS",
+         "exitcode=" STRINGIFY(LSAN_ERROR) ":"
+         "fast_unwind_on_malloc=0:"
+         "symbolize=0:"
+         "print_suppressions=0",
+         0);
+
   if (get_afl_env("AFL_PRELOAD")) {
 
     if (fsrv->qemu_mode) {
 
       /* afl-qemu-trace takes care of converting AFL_PRELOAD. */
 
+    } else if (fsrv->frida_mode) {
+
+      afl_preload = getenv("AFL_PRELOAD");
+      u8 *frida_binary = find_afl_binary(argv[0], "afl-frida-trace.so");
+      if (afl_preload) {
+
+        frida_afl_preload = alloc_printf("%s:%s", afl_preload, frida_binary);
+
+      } else {
+
+        frida_afl_preload = alloc_printf("%s", frida_binary);
+
+      }
+
+      ck_free(frida_binary);
+
+      setenv("LD_PRELOAD", frida_afl_preload, 1);
+      setenv("DYLD_INSERT_LIBRARIES", frida_afl_preload, 1);
+
     } else {
 
       setenv("LD_PRELOAD", getenv("AFL_PRELOAD"), 1);
@@ -762,8 +802,17 @@ static void set_up_environment(afl_forkserver_t *fsrv) {
 
     }
 
+  } else if (fsrv->frida_mode) {
+
+    u8 *frida_binary = find_afl_binary(argv[0], "afl-frida-trace.so");
+    setenv("LD_PRELOAD", frida_binary, 1);
+    setenv("DYLD_INSERT_LIBRARIES", frida_binary, 1);
+    ck_free(frida_binary);
+
   }
 
+  if (frida_afl_preload) { ck_free(frida_afl_preload); }
+
 }
 
 /* Setup signal handlers, duh. */
@@ -804,6 +853,7 @@ static void usage(u8 *argv0) {
       "  -f file       - input file read by the tested program (stdin)\n"
       "  -t msec       - timeout for each run (%u ms)\n"
       "  -m megs       - memory limit for child process (%u MB)\n"
+      "  -O            - use binary-only instrumentation (FRIDA mode)\n"
       "  -Q            - use binary-only instrumentation (QEMU mode)\n"
       "  -U            - use unicorn-based instrumentation (Unicorn mode)\n"
       "  -W            - use qemu-based instrumentation with Wine (Wine "
@@ -859,7 +909,7 @@ int main(int argc, char **argv_orig, char **envp) {
 
   SAYF(cCYA "afl-tmin" VERSION cRST " by Michal Zalewski\n");
 
-  while ((opt = getopt(argc, argv, "+i:o:f:m:t:B:xeQUWHh")) > 0) {
+  while ((opt = getopt(argc, argv, "+i:o:f:m:t:B:xeOQUWHh")) > 0) {
 
     switch (opt) {
 
@@ -971,6 +1021,14 @@ int main(int argc, char **argv_orig, char **envp) {
 
         break;
 
+      case 'O':                                               /* FRIDA mode */
+
+        if (fsrv->frida_mode) { FATAL("Multiple -O options not supported"); }
+
+        fsrv->frida_mode = 1;
+
+        break;
+
       case 'Q':
 
         if (fsrv->qemu_mode) { FATAL("Multiple -Q options not supported"); }
@@ -1054,7 +1112,7 @@ int main(int argc, char **argv_orig, char **envp) {
   atexit(at_exit_handler);
   setup_signal_handlers();
 
-  set_up_environment(fsrv);
+  set_up_environment(fsrv, argv);
 
   fsrv->target_path = find_binary(argv[optind]);
   fsrv->trace_bits = afl_shm_init(&shm, map_size, 0);