taint integration done

9 files changed, 255 insertions, 64 deletions

diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 328c8405..19807880 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -133,8 +133,10 @@ extern s32
 
 struct queue_entry {
 
-  u8 *fname;                            /* File name for the test case      */
-  u32 len;                              /* Input length                     */
+  u8 *                fname;            /* File name for the test case      */
+  u8 *                fname_taint;      /* File name for taint data         */
+  u32                 len;              /* Input length                     */
+  struct queue_entry *prev;             /* previous queue entry, if any     */
 
   u8 cal_failed,                        /* Calibration failed?              */
       trim_done,                        /* Trimmed?                         */
@@ -148,7 +150,10 @@ struct queue_entry {
       is_ascii;                         /* Is the input just ascii text?    */
 
   u32 bitmap_size,                      /* Number of bits set in bitmap     */
-      fuzz_level;                       /* Number of fuzzing iterations     */
+      fuzz_level,                       /* Number of fuzzing iterations     */
+      taint_bytes_all,                  /* Number of tainted bytes          */
+      taint_bytes_new,                  /* Number of new tainted bytes      */
+      taint_bytes_highest;              /* highest offset in input          */
 
   u64 exec_us,                          /* Execution time (us)              */
       handicap,                         /* Number of queue cycles behind    */
@@ -885,7 +890,7 @@ void   deinit_py(void *);
 void mark_as_det_done(afl_state_t *, struct queue_entry *);
 void mark_as_variable(afl_state_t *, struct queue_entry *);
 void mark_as_redundant(afl_state_t *, struct queue_entry *, u8);
-void add_to_queue(afl_state_t *, u8 *, u32, u8);
+void add_to_queue(afl_state_t *, u8 *, u8 *, u32, struct queue_entry *, u8);
 void destroy_queue(afl_state_t *);
 void update_bitmap_score(afl_state_t *, struct queue_entry *);
 void cull_queue(afl_state_t *);
@@ -975,6 +980,8 @@ void   check_if_tty(afl_state_t *);
 void   setup_signal_handlers(void);
 void   save_cmdline(afl_state_t *, u32, char **);
 void   read_foreign_testcases(afl_state_t *, int);
+void   perform_taint_run(afl_state_t *afl, struct queue_entry *q, u8 *fname,
+                         u8 *mem, u32 len);
 
 /* CmpLog */
 
diff --git a/llvm_mode/afl-llvm-rt.o.c b/llvm_mode/afl-llvm-rt.o.c
index 8cc59cbb..5ffae39c 100644
--- a/llvm_mode/afl-llvm-rt.o.c
+++ b/llvm_mode/afl-llvm-rt.o.c
@@ -831,6 +831,16 @@ void __afl_manual_init(void) {
 
   static u8 init_done;
 
+  if (getenv("AFL_DISABLE_LLVM_INSTRUMENTATION")) {
+
+    init_done = 1;
+    if (getenv("AFL_DEBUG"))
+      fprintf(stderr,
+              "DEBUG: disabled instrumenation because of "
+              "AFL_DISABLE_LLVM_INSTRUMENTATION\n");
+
+  }
+
   if (!init_done) {
 
     __afl_map_shm();
diff --git a/src/afl-common.c b/src/afl-common.c
index c0202821..e01bde3c 100644
--- a/src/afl-common.c
+++ b/src/afl-common.c
@@ -232,7 +232,7 @@ char **get_qemu_argv(u8 *own_loc, u8 **target_path_p, int argc, char **argv) {
 char **get_wine_argv(u8 *own_loc, u8 **target_path_p, int argc, char **argv) {
 
   char **new_argv = ck_alloc(sizeof(char *) * (argc + 3));
-  u8 *   tmp, *cp = NULL, *rsl, *own_copy;
+  u8 *   cp = NULL;
 
   memcpy(&new_argv[2], &argv[1], (int)(sizeof(char *)) * (argc - 1));
   new_argv[argc - 1] = NULL;
diff --git a/src/afl-forkserver.c b/src/afl-forkserver.c
index eeb2f8c3..4dc5e438 100644
--- a/src/afl-forkserver.c
+++ b/src/afl-forkserver.c
@@ -933,7 +933,7 @@ void afl_fsrv_write_to_testcase(afl_forkserver_t *fsrv, u8 *buf, size_t len) {
 
   } else {
 
-    s32 fd = fsrv->out_fd;
+    s32 fd;
 
     if (fsrv->out_file) {
 
@@ -952,6 +952,7 @@ void afl_fsrv_write_to_testcase(afl_forkserver_t *fsrv, u8 *buf, size_t len) {
 
     } else {
 
+      fd = fsrv->out_fd;
       lseek(fd, 0, SEEK_SET);
 
     }
diff --git a/src/afl-fuzz-bitmap.c b/src/afl-fuzz-bitmap.c
index 11a3f121..d4ee36e1 100644
--- a/src/afl-fuzz-bitmap.c
+++ b/src/afl-fuzz-bitmap.c
@@ -648,7 +648,7 @@ u8 save_if_interesting(afl_state_t *afl, void *mem, u32 len, u8 fault) {
 
 #endif                                                    /* ^!SIMPLE_FILES */
 
-    add_to_queue(afl, queue_fn, len, 0);
+    add_to_queue(afl, queue_fn, mem, len, afl->queue_top, 0);
 
     if (hnb == 2) {
 
diff --git a/src/afl-fuzz-init.c b/src/afl-fuzz-init.c
index 350a8599..350a3b4c 100644
--- a/src/afl-fuzz-init.c
+++ b/src/afl-fuzz-init.c
@@ -712,7 +712,7 @@ void read_testcases(afl_state_t *afl) {
 
     if (!access(dfn, F_OK)) { passed_det = 1; }
 
-    add_to_queue(afl, fn2, st.st_size, passed_det);
+    add_to_queue(afl, fn2, NULL, st.st_size, NULL, passed_det);
 
   }
 
@@ -960,6 +960,9 @@ void perform_dry_run(afl_state_t *afl) {
 
     }
 
+    /* perform taint gathering on the input seed */
+    perform_taint_run(afl, q, q->fname, use_mem, q->len);
+
     q = q->next;
 
   }
@@ -1438,6 +1441,10 @@ static void handle_existing_out_dir(afl_state_t *afl) {
 
     u8 *orig_q = alloc_printf("%s/queue", afl->out_dir);
 
+    u8 *fnt = alloc_printf("%s/taint", afl->out_dir);
+    mkdir(fnt, 0755);  // ignore errors
+    ck_free(fnt);
+
     afl->in_dir = alloc_printf("%s/_resume", afl->out_dir);
 
     rename(orig_q, afl->in_dir);                           /* Ignore errors */
@@ -1494,6 +1501,15 @@ static void handle_existing_out_dir(afl_state_t *afl) {
   if (delete_files(fn, CASE_PREFIX)) { goto dir_cleanup_failed; }
   ck_free(fn);
 
+  if (afl->fsrv.taint_mode) {
+
+    fn = alloc_printf("%s/taint", afl->out_dir);
+    mkdir(fn, 0755);  // ignore errors
+    if (delete_files(fn, CASE_PREFIX)) { goto dir_cleanup_failed; }
+    ck_free(fn);
+
+  }
+
   /* All right, let's do <afl->out_dir>/crashes/id:* and
    * <afl->out_dir>/hangs/id:*. */
 
@@ -1721,6 +1737,16 @@ void setup_dirs_fds(afl_state_t *afl) {
   if (mkdir(tmp, 0700)) { PFATAL("Unable to create '%s'", tmp); }
   ck_free(tmp);
 
+  /* Taint directory if taint_mode. */
+
+  if (afl->fsrv.taint_mode) {
+
+    tmp = alloc_printf("%s/taint", afl->out_dir);
+    if (mkdir(tmp, 0700)) { PFATAL("Unable to create '%s'", tmp); }
+    ck_free(tmp);
+
+  }
+
   /* Top-level directory for queue metadata used for session
      resume and related tasks. */
 
diff --git a/src/afl-fuzz-queue.c b/src/afl-fuzz-queue.c
index f35df914..36ec0896 100644
--- a/src/afl-fuzz-queue.c
+++ b/src/afl-fuzz-queue.c
@@ -103,6 +103,139 @@ void mark_as_redundant(afl_state_t *afl, struct queue_entry *q, u8 state) {
 
 }
 
+void perform_taint_run(afl_state_t *afl, struct queue_entry *q, u8 *fname,
+                       u8 *mem, u32 len) {
+
+  u8 *                ptr, *fn = fname;
+  u32                 bytes = 0, plen = len;
+  s32                 fd = -1;
+  struct queue_entry *prev = q->prev;
+
+  if (plen % 4) plen = plen + 4 - (len % 4);
+
+  if ((ptr = strrchr(fname, '/')) != NULL) fn = ptr + 1;
+  q->fname_taint = alloc_printf("%s/taint/%s", afl->out_dir, fn);
+
+  if (q->fname_taint) {
+
+    afl->taint_fsrv.map_size = plen;  // speed :)
+    write_to_testcase(afl, mem, len);
+    if (afl_fsrv_run_target(&afl->taint_fsrv, afl->fsrv.exec_tmout,
+                            &afl->stop_soon) == 0) {
+
+      bytes = count_bytes_len(afl, afl->taint_fsrv.trace_bits, plen);
+      if (afl->debug)
+        fprintf(stderr, "Debug: tainted %u out of %u bytes\n", bytes, len);
+
+      if (bytes) {
+
+        s32 i = len;
+        while (i > 0 && !afl->taint_fsrv.trace_bits[i - 1])
+          i--;
+        q->taint_bytes_highest = i;
+
+      }
+
+    }
+
+    if (((bytes * 100) / len) < 90) {
+
+      // we only use the taint havoc mode if the entry has less than 90% of
+      // overall tainted bytes
+      q->taint_bytes_all = bytes;
+
+      // save the bytes away
+      int w = open(q->fname_taint, O_CREAT | O_WRONLY, 0644);
+      if (w >= 0) {
+
+        ck_write(w, afl->taint_fsrv.trace_bits, plen, q->fname_taint);
+        close(w);
+
+      } else {
+
+        FATAL("could not create %s", q->fname_taint);
+        bytes = 0;
+
+      }
+
+      if (bytes && prev && prev->taint_bytes_all) {
+
+        // check if there are new bytes in the taint vs the previous
+        int r = open(prev->fname_taint, O_RDONLY);
+
+        if (r >= 0) {
+
+          u8 *bufr = mmap(0, prev->len, PROT_READ, MAP_PRIVATE, r, 0);
+
+          if ((size_t)bufr != -1) {
+
+            u32 i;
+            u8 *tmp = ck_maybe_grow(BUF_PARAMS(in_scratch), plen);
+            memset(tmp, 0, plen);
+
+            for (i = 0; i < len; i++)
+              if (afl->taint_fsrv.trace_bits[i] && (i >= prev->len || !bufr[i]))
+                tmp[i] = 1;
+
+            q->taint_bytes_new = count_bytes_len(afl, tmp, plen);
+
+            if (q->taint_bytes_new) {
+
+              u8 *fnw = alloc_printf("%s.new", q->fname_taint);
+              int w = open(fnw, O_CREAT | O_WRONLY, 0644);
+              if (w >= 0) {
+
+                ck_write(w, tmp, plen, fnw);
+                close(w);
+
+              } else {
+
+                q->taint_bytes_new = 0;
+
+              }
+
+              ck_free(fnw);
+
+            }
+
+            munmap(bufr, prev->len);
+
+          }
+
+          close(r);
+
+        }
+
+      }
+
+    } else {
+
+      bytes = 0;
+
+    }
+
+  }
+
+  if (!bytes) {
+
+    q->taint_bytes_highest = q->taint_bytes_all = q->taint_bytes_new = 0;
+
+    if (q->fname_taint) {
+
+      ck_free(q->fname_taint);
+      q->fname_taint = NULL;
+
+    }
+
+  } else {
+
+    if (q->taint_bytes_all && !q->taint_bytes_new)
+      q->taint_bytes_new = q->taint_bytes_all;
+
+  }
+
+}
+
 /* check if ascii or UTF-8 */
 
 static u8 check_if_text(struct queue_entry *q) {
@@ -212,10 +345,12 @@ static u8 check_if_text(struct queue_entry *q) {
 
 /* Append new test case to the queue. */
 
-void add_to_queue(afl_state_t *afl, u8 *fname, u32 len, u8 passed_det) {
+void add_to_queue(afl_state_t *afl, u8 *fname, u8 *mem, u32 len,
+                  struct queue_entry *prev_q, u8 passed_det) {
 
   struct queue_entry *q = ck_alloc(sizeof(struct queue_entry));
 
+  q->prev = prev_q;
   q->fname = fname;
   q->len = len;
   q->depth = afl->cur_depth + 1;
@@ -254,6 +389,17 @@ void add_to_queue(afl_state_t *afl, u8 *fname, u32 len, u8 passed_det) {
 
   afl->last_path_time = get_cur_time();
 
+  /* trigger the tain gathering if this is not a dry run */
+  if (afl->fsrv.taint_mode && mem) {
+
+    perform_taint_run(afl, q, fname, mem, len);
+
+  }
+
+  /* only redqueen currently uses is_ascii */
+  if (afl->shm.cmplog_mode) q->is_ascii = check_if_text(q);
+
+  /* run custom mutators afl_custom_queue_new_entry() */
   if (afl->custom_mutators_count) {
 
     LIST_FOREACH(&afl->custom_mutator_list, struct custom_mutator, {
@@ -273,9 +419,6 @@ void add_to_queue(afl_state_t *afl, u8 *fname, u32 len, u8 passed_det) {
 
   }
 
-  /* only redqueen currently uses is_ascii */
-  if (afl->shm.cmplog_mode) q->is_ascii = check_if_text(q);
-
 }
 
 /* Destroy the entire queue. */
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index 89ae0424..ddbd5524 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -471,24 +471,6 @@ abort_calibration:
   afl->stage_cur = old_sc;
   afl->stage_max = old_sm;
 
-  /* if taint mode was selected, run the taint */
-
-  if (afl->fsrv.taint_mode) {
-
-    write_to_testcase(afl, use_mem, q->len);
-    if (afl_fsrv_run_target(&afl->taint_fsrv, use_tmout, &afl->stop_soon) ==
-        0) {
-
-      u32 len = q->len;
-      if (len % 4) len = len + 4 - (q->len % 4);
-      u32 bytes = count_bytes_len(afl, afl->taint_fsrv.trace_bits, len);
-      if (afl->debug)
-        fprintf(stderr, "Debug: tainted %u out of %u bytes\n", bytes, q->len);
-
-    }
-
-  }
-
   if (!first_run) { show_stats(afl); }
 
   return fault;
@@ -770,56 +752,65 @@ u8 trim_case(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
     while (remove_pos < q->len) {
 
       u32 trim_avail = MIN(remove_len, q->len - remove_pos);
-      u64 cksum;
 
-      write_with_gap(afl, in_buf, q->len, remove_pos, trim_avail);
+      if (likely((!q->taint_bytes_highest) ||
+                 (q->len - trim_avail > q->taint_bytes_highest))) {
 
-      fault = fuzz_run_target(afl, &afl->fsrv, afl->fsrv.exec_tmout);
-      ++afl->trim_execs;
+        u64 cksum;
 
-      if (afl->stop_soon || fault == FSRV_RUN_ERROR) { goto abort_trimming; }
+        write_with_gap(afl, in_buf, q->len, remove_pos, trim_avail);
 
-      /* Note that we don't keep track of crashes or hangs here; maybe TODO?
-       */
+        fault = fuzz_run_target(afl, &afl->fsrv, afl->fsrv.exec_tmout);
+        ++afl->trim_execs;
 
-      cksum = hash64(afl->fsrv.trace_bits, afl->fsrv.map_size, HASH_CONST);
+        if (afl->stop_soon || fault == FSRV_RUN_ERROR) { goto abort_trimming; }
 
-      /* If the deletion had no impact on the trace, make it permanent. This
-         isn't perfect for variable-path inputs, but we're just making a
-         best-effort pass, so it's not a big deal if we end up with false
-         negatives every now and then. */
+        /* Note that we don't keep track of crashes or hangs here; maybe TODO?
+         */
 
-      if (cksum == q->exec_cksum) {
+        cksum = hash64(afl->fsrv.trace_bits, afl->fsrv.map_size, HASH_CONST);
 
-        u32 move_tail = q->len - remove_pos - trim_avail;
+        /* If the deletion had no impact on the trace, make it permanent. This
+           isn't perfect for variable-path inputs, but we're just making a
+           best-effort pass, so it's not a big deal if we end up with false
+           negatives every now and then. */
 
-        q->len -= trim_avail;
-        len_p2 = next_pow2(q->len);
+        if (cksum == q->exec_cksum) {
 
-        memmove(in_buf + remove_pos, in_buf + remove_pos + trim_avail,
-                move_tail);
+          u32 move_tail = q->len - remove_pos - trim_avail;
 
-        /* Let's save a clean trace, which will be needed by
-           update_bitmap_score once we're done with the trimming stuff. */
+          q->len -= trim_avail;
+          len_p2 = next_pow2(q->len);
 
-        if (!needs_write) {
+          memmove(in_buf + remove_pos, in_buf + remove_pos + trim_avail,
+                  move_tail);
 
-          needs_write = 1;
-          memcpy(afl->clean_trace, afl->fsrv.trace_bits, afl->fsrv.map_size);
+          /* Let's save a clean trace, which will be needed by
+             update_bitmap_score once we're done with the trimming stuff. */
+
+          if (!needs_write) {
+
+            needs_write = 1;
+            memcpy(afl->clean_trace, afl->fsrv.trace_bits, afl->fsrv.map_size);
+
+          }
+
+        } else {
+
+          remove_pos += remove_len;
 
         }
 
+        /* Since this can be slow, update the screen every now and then. */
+        if (!(trim_exec++ % afl->stats_update_freq)) { show_stats(afl); }
+        ++afl->stage_cur;
+
       } else {
 
         remove_pos += remove_len;
 
       }
 
-      /* Since this can be slow, update the screen every now and then. */
-
-      if (!(trim_exec++ % afl->stats_update_freq)) { show_stats(afl); }
-      ++afl->stage_cur;
-
     }
 
     remove_len >>= 1;
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index 93ab90e2..6f143db7 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -244,9 +244,10 @@ static int stricmp(char const *a, char const *b) {
 
 int main(int argc, char **argv_orig, char **envp) {
 
-  s32    opt;
-  u64    prev_queued = 0;
-  u32    sync_interval_cnt = 0, seek_to, show_help = 0, map_size = MAP_SIZE;
+  s32 opt;
+  u64 prev_queued = 0;
+  u32 sync_interval_cnt = 0, seek_to, show_help = 0, map_size = MAP_SIZE,
+      real_map_size = 0;
   u8 *   extras_dir = 0;
   u8     mem_limit_given = 0, exit_1 = 0, debug = 0;
   char **use_argv;
@@ -827,6 +828,7 @@ int main(int argc, char **argv_orig, char **envp) {
 
   if (afl->fsrv.taint_mode && afl->fsrv.map_size < MAX_FILE) {
 
+    real_map_size = map_size;
     map_size = afl->fsrv.map_size = afl->shm.map_size = MAX_FILE;
 
   }
@@ -840,8 +842,7 @@ int main(int argc, char **argv_orig, char **envp) {
   OKF("afl++ is open source, get it at "
       "https://github.com/AFLplusplus/AFLplusplus");
   OKF("Power schedules from github.com/mboehme/aflfast");
-  OKF("Python Mutator and llvm_mode instrument file list from "
-      "github.com/choller/afl");
+  OKF("Python Mutator from github.com/choller/afl");
   OKF("MOpt Mutator from github.com/puppet-meteor/MOpt-AFL");
 
   if (afl->sync_id && afl->is_main_node &&
@@ -1088,6 +1089,13 @@ int main(int argc, char **argv_orig, char **envp) {
   memset(afl->virgin_tmout, 255, afl->fsrv.map_size);
   memset(afl->virgin_crash, 255, afl->fsrv.map_size);
 
+  if (map_size != real_map_size) {
+
+    afl->fsrv.map_size = real_map_size;
+    if (afl->cmplog_binary) afl->cmplog_fsrv.map_size;
+
+  }
+
   init_count_class16();
 
   if (afl->is_main_node && check_main_node_exists(afl) == 1) {
@@ -1252,12 +1260,16 @@ int main(int argc, char **argv_orig, char **envp) {
   if (afl->fsrv.taint_mode) {
 
     ACTF("Spawning qemu_taint forkserver");
+    u8 *disable = getenv("AFL_DISABLE_LLVM_INSTRUMENTATION");
+    setenv("AFL_DISABLE_LLVM_INSTRUMENTATION", "1", 0);
     afl_fsrv_init_dup(&afl->taint_fsrv, &afl->fsrv);
     afl->taint_fsrv.qemu_mode = 2;
+    afl->taint_fsrv.taint_mode = 1;
     afl->taint_fsrv.trace_bits = afl->fsrv.trace_bits;
     ck_free(afl->taint_fsrv.target_path);
-    afl->taint_fsrv.target_path = ck_strdup(afl->fsrv.target_path);
     afl->argv_taint = ck_alloc(sizeof(char *) * (argc + 4 - optind));
+    afl->taint_fsrv.target_path =
+        find_binary_own_loc("afl-qemu-taint", argv[0]);
     afl->argv_taint[0] = find_binary_own_loc("afl-qemu-taint", argv[0]);
     if (!afl->argv_taint[0])
       FATAL(
@@ -1278,6 +1290,7 @@ int main(int argc, char **argv_orig, char **envp) {
       setenv("AFL_TAINT_INPUT", afl->fsrv.out_file, 1);
     afl_fsrv_start(&afl->taint_fsrv, afl->argv_taint, &afl->stop_soon,
                    afl->afl_env.afl_debug_child_output);
+    if (!disable) unsetenv("AFL_DISABLE_LLVM_INSTRUMENTATION");
     OKF("Taint forkserver successfully started");
 
   }