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

With CODE_COVERAGE builds, we need to collect the coverage data of each
iteration in a persistant buffer that has the same size as the regular
trace buffer used for fuzzing. We dump this information at the end of
the run and when combined with pointer data and module info, this can be
used to calculate code coverage.

3 files changed, 51 insertions, 0 deletions

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,