Merge pull request #2129 from choller/persist-code-cov

Collect persistent coverage data and dump it at the end of the run

4 files changed, 55 insertions, 0 deletions

diff --git a/include/forkserver.h b/include/forkserver.h
index 593e34a2..3fd813a4 100644
--- a/include/forkserver.h
+++ b/include/forkserver.h
@@ -206,6 +206,10 @@ typedef struct afl_forkserver {
   s32                   nyx_log_fd;
 #endif
 
+#ifdef __AFL_CODE_COVERAGE
+  u8 *persistent_trace_bits;                   /* Persistent copy of bitmap */
+#endif
+
 } afl_forkserver_t;
 
 typedef enum fsrv_run_result {
diff --git a/src/afl-forkserver.c b/src/afl-forkserver.c
index 71d8570d..a998c10f 100644
--- a/src/afl-forkserver.c
+++ b/src/afl-forkserver.c
@@ -252,6 +252,10 @@ void afl_fsrv_init(afl_forkserver_t *fsrv) {
   fsrv->uses_crash_exitcode = false;
   fsrv->uses_asan = false;
 
+#ifdef __AFL_CODE_COVERAGE
+  fsrv->persistent_trace_bits = NULL;
+#endif
+
   fsrv->init_child_func = fsrv_exec_child;
   list_append(&fsrv_list, fsrv);
 
@@ -278,6 +282,10 @@ void afl_fsrv_init_dup(afl_forkserver_t *fsrv_to, afl_forkserver_t *from) {
   fsrv_to->fsrv_kill_signal = from->fsrv_kill_signal;
   fsrv_to->debug = from->debug;
 
+#ifdef __AFL_CODE_COVERAGE
+  fsrv_to->persistent_trace_bits = from->persistent_trace_bits;
+#endif
+
   // These are forkserver specific.
   fsrv_to->out_dir_fd = -1;
   fsrv_to->child_pid = -1;
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index 6a0da6ab..c234fc42 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -60,6 +60,27 @@ fuzz_run_target(afl_state_t *afl, afl_forkserver_t *fsrv, u32 timeout) {
 
   fsrv_run_result_t res = afl_fsrv_run_target(fsrv, timeout, &afl->stop_soon);
 
+#ifdef __AFL_CODE_COVERAGE
+  if (unlikely(!fsrv->persistent_trace_bits)) {
+
+    // On the first run, we allocate the persistent map to collect coverage.
+    fsrv->persistent_trace_bits = (u8 *)malloc(fsrv->map_size);
+    memset(fsrv->persistent_trace_bits, 0, fsrv->map_size);
+
+  }
+
+  for (u32 i = 0; i < fsrv->map_size; ++i) {
+
+    if (fsrv->persistent_trace_bits[i] != 255 && fsrv->trace_bits[i]) {
+
+      fsrv->persistent_trace_bits[i]++;
+
+    }
+
+  }
+
+#endif
+
   /* If post_run() function is defined in custom mutator, the function will be
      called each time after AFL++ executes the target program. */
 
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index a09a53ec..0209e74f 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -3130,6 +3130,28 @@ stop_fuzzing:
   write_bitmap(afl);
   save_auto(afl);
 
+  #ifdef __AFL_CODE_COVERAGE
+  if (afl->fsrv.persistent_trace_bits) {
+
+    char cfn[4096];
+    snprintf(cfn, sizeof(cfn), "%s/covmap.dump", afl->out_dir);
+
+    FILE *cov_fd;
+    if ((cov_fd = fopen(cfn, "w")) == NULL) {
+
+      PFATAL("could not create '%s'", cfn);
+
+    }
+
+    // Write the real map size, as the map size must exactly match the pointer
+    // map in length.
+    fwrite(afl->fsrv.persistent_trace_bits, 1, afl->fsrv.real_map_size, cov_fd);
+    fclose(cov_fd);
+
+  }
+
+  #endif
+
   if (afl->pizza_is_served) {
 
     SAYF(CURSOR_SHOW cLRD "\n\n+++ Baking aborted %s +++\n" cRST,