implement covmap in forkserver

include/forkserver.h

@@ -151,6 +151,8 @@ typedef struct afl_forkserver {
 
   bool uses_asan;                       /* Target uses ASAN?                */
 
+  bool keep_coverage;                   /* use coverage feature             */
+
   bool debug;                           /* debug mode?                      */
 
   bool uses_crash_exitcode;             /* Custom crash exitcode specified? */
@@ -160,6 +162,8 @@ typedef struct afl_forkserver {
 
   u8 *shmem_fuzz;                       /* allocated memory for fuzzing     */
 
+  u8 *coverage_map;                     /* for coverage feature             */
+
   char *cmplog_binary;                  /* the name of the cmplog binary    */
 
   /* persistent mode replay functionality */

instrumentation/afl-compiler-rt.o.c

@@ -14,15 +14,13 @@
 
 */
 
-#ifdef __AFL_CODE_COVERAGE
-  #ifndef _GNU_SOURCE
+#ifndef _GNU_SOURCE
   #define _GNU_SOURCE
-  #endif
-  #ifndef __USE_GNU
-    #define __USE_GNU
-  #endif
-  #include <dlfcn.h>
 #endif
+#ifndef __USE_GNU
+  #define __USE_GNU
+#endif
+#include <dlfcn.h>
 
 #ifdef __ANDROID__
   #include "android-ashmem.h"
@@ -115,7 +113,6 @@ u32 __afl_dictionary_len;
 u64 __afl_map_addr;
 u32 __afl_first_final_loc;
 
-#ifdef __AFL_CODE_COVERAGE
 typedef struct afl_module_info_t afl_module_info_t;
 
 struct afl_module_info_t {
@@ -149,9 +146,16 @@ typedef struct {
 
 afl_module_info_t *__afl_module_info = NULL;
 
-u32        __afl_pcmap_size = 0;
+u32        __afl_pcmap_size = 0, __afl_pcmap_shmem = 1;
 uintptr_t *__afl_pcmap_ptr = NULL;
-#endif  // __AFL_CODE_COVERAGE
+
+// Maximum path length on Linux
+#ifndef PATH_MAX
+  #define PATH_MAX 4096
+#endif
+
+// Maximum length of an uint32_t as string
+#define START_STOP_MAX 10
 
 /* 1 if we are running in afl, and the forkserver was started, else 0 */
 u32 __afl_connected = 0;
@@ -727,7 +731,6 @@ static void __afl_map_shm(void) {
 
   }
 
-#ifdef __AFL_CODE_COVERAGE
   char *pcmap_id_str = getenv("__AFL_PCMAP_SHM_ID");
 
   if (pcmap_id_str) {
@@ -744,9 +747,15 @@ static void __afl_map_shm(void) {
 
     }
 
-  }
+  } else if (getenv("__AFL_CODE_COVERAGE") ||
 
-#endif  // __AFL_CODE_COVERAGE
+             getenv("AFL_DUMP_CODE_COVERAGE")) {
+
+    __afl_pcmap_size = __afl_map_size * sizeof(void *);
+    __afl_pcmap_ptr = (uintptr_t *)malloc(__afl_pcmap_size);
+    __afl_pcmap_shmem = 0;
+
+  }
 
 }
 
@@ -756,17 +765,15 @@ static void __afl_unmap_shm(void) {
 
   if (!__afl_already_initialized_shm) return;
 
-#ifdef __AFL_CODE_COVERAGE
   if (__afl_pcmap_size) {
 
-    shmdt((void *)__afl_pcmap_ptr);
+    if (__afl_pcmap_shmem) { shmdt((void *)__afl_pcmap_ptr); }
     __afl_pcmap_ptr = NULL;
     __afl_pcmap_size = 0;
+    __afl_pcmap_shmem = 1;
 
   }
 
-#endif  // __AFL_CODE_COVERAGE
-
   char *id_str = getenv(SHM_ENV_VAR);
 
   if (id_str) {
@@ -1082,9 +1089,88 @@ static void __afl_start_snapshots(void) {
 
 static void __afl_start_forkserver(void) {
 
-  if (__afl_already_initialized_forkserver) return;
+  if (__afl_already_initialized_forkserver) { return; }
   __afl_already_initialized_forkserver = 1;
 
+  if (getenv("AFL_DUMP_CODE_COVERAGE")) {
+
+    if (__afl_module_info) {
+
+      int32_t            cnt = 0;
+      afl_module_info_t *start = __afl_module_info;
+
+      while (start) {
+
+        ++cnt;
+        start = start->next;
+
+      }
+
+      // Allocate per entry enough space for:
+      //
+      //   1. One path
+      //   2. Two pcguard start/stop offsets
+      //   3. Two spaces and a trailing newline
+      //
+      // This is a very conservative allocation so we can just YOLO the rest.
+      size_t bufsize = (PATH_MAX + START_STOP_MAX * 2 + 2 + 1) * cnt + 1;
+      char  *buf = malloc(bufsize);
+      char  *cur = buf;
+
+      if (!buf) { perror("malloc"); };
+
+      start = __afl_module_info;
+
+      while (start) {
+
+        size_t namelen = strlen(start->name);
+
+        memcpy(cur, start->name, namelen);
+        cur += namelen;
+        *cur = ' ';
+        cur += 1;
+        cur += sprintf(cur, "%u %u", start->start, start->stop);
+        *cur = '\n';
+        cur += 1;
+
+        start = start->next;
+
+      }
+
+      *cur = '\0';
+
+      FILE *f = fopen("modinfo.txt", "w");
+      if (!f) {
+
+        fprintf(stderr, "Error: Could not create modinfo.txt!");
+        exit(-1);
+
+      }
+
+      fprintf(f, "%s\n", buf);
+      fclose(f);
+
+      f = fopen("pcmap.dump", "w");
+      if (!f) {
+
+        fprintf(stderr, "Error: Could not create pcmap.dump!");
+        exit(-1);
+
+      }
+
+      fwrite(__afl_pcmap_ptr, __afl_map_size * sizeof(void *), 1, f);
+      fclose(f);
+
+      fprintf(stderr,
+              "[+] Created modinfo.txt and pcmap.dump for coverage analysis "
+              "purposes. Now run afl-showmap with '-V -b -o covmap.dump'.\n");
+
+    }
+
+    exit(-1);
+
+  }
+
   struct sigaction orig_action;
   sigaction(SIGTERM, NULL, &orig_action);
   old_sigterm_handler = orig_action.sa_handler;
@@ -1588,7 +1674,6 @@ void __sanitizer_cov_trace_pc_guard(uint32_t *guard) {
 
 }
 
-#ifdef __AFL_CODE_COVERAGE
 void __sanitizer_cov_pcs_init(const uintptr_t *pcs_beg,
                               const uintptr_t *pcs_end) {
 
@@ -1644,15 +1729,17 @@ void __sanitizer_cov_pcs_init(const uintptr_t *pcs_beg,
 
     u32 in_module_index = 0;
 
-    while (start < end) {
+    if (mod_info->start - in_module_index >= __afl_map_size) {
 
-      if (mod_info->start + in_module_index >= __afl_map_size) {
-
-        fprintf(stderr, "ERROR: __sanitizer_cov_pcs_init out of bounds?!\n");
+      fprintf(stderr,
+              "ERROR: __sanitizer_cov_pcs_init out of bounds?! (%u >= %u)\n",
+              mod_info->start, __afl_map_size);
       abort();
 
     }
 
+    while (start < end) {
+
       uintptr_t PC = start->PC;
 
       // This is what `GetPreviousInstructionPc` in sanitizer runtime does
@@ -1682,8 +1769,6 @@ void __sanitizer_cov_pcs_init(const uintptr_t *pcs_beg,
 
 }
 
-#endif  // __AFL_CODE_COVERAGE
-
 /* Init callback. Populates instrumentation IDs. Note that we're using
    ID of 0 as a special value to indicate non-instrumented bits. That may
    still touch the bitmap, but in a fairly harmless way. */
@@ -1715,11 +1800,13 @@ void __sanitizer_cov_trace_pc_guard_init(uint32_t *start, uint32_t *stop) {
 
   if (start == stop || *start) { return; }
 
-#ifdef __AFL_CODE_COVERAGE
   u32               *orig_start = start;
   afl_module_info_t *mod_info = NULL;
 
+  if (getenv("__AFL_CODE_COVERAGE") || getenv("AFL_DUMP_CODE_COVERAGE")) {
+
     Dl_info dlinfo;
+
     if (dladdr(__builtin_return_address(0), &dlinfo)) {
 
       if (__afl_already_initialized_forkserver) {
@@ -1758,8 +1845,12 @@ void __sanitizer_cov_trace_pc_guard_init(uint32_t *start, uint32_t *stop) {
 
         }
 
-      fprintf(stderr, "[pcmap] Module: %s Base Address: %p\n", dlinfo.dli_fname,
-              dlinfo.dli_fbase);
+        if (__afl_debug) {
+
+          fprintf(stderr, "[pcmap] Module: %s Base Address: %p\n",
+                  dlinfo.dli_fname, dlinfo.dli_fbase);
+
+        }
 
       }
 
@@ -1769,7 +1860,7 @@ void __sanitizer_cov_trace_pc_guard_init(uint32_t *start, uint32_t *stop) {
 
     }
 
-#endif  // __AFL_CODE_COVERAGE
+  }
 
   x = getenv("AFL_INST_RATIO");
   if (x) {
@@ -1858,7 +1949,6 @@ void __sanitizer_cov_trace_pc_guard_init(uint32_t *start, uint32_t *stop) {
 
   }
 
-#ifdef __AFL_CODE_COVERAGE
   if (mod_info) {
 
     mod_info->start = *orig_start;
@@ -1872,8 +1962,6 @@ void __sanitizer_cov_trace_pc_guard_init(uint32_t *start, uint32_t *stop) {
 
   }
 
-#endif  // __AFL_CODE_COVERAGE
-
   if (__afl_debug) {
 
     fprintf(stderr,

src/afl-cc.c

@@ -1868,6 +1868,7 @@ int main(int argc, char **argv, char **envp) {
 
           instrument_mode = INSTRUMENT_LLVMNATIVE;
           instrument_opt_mode |= INSTRUMENT_OPT_CODECOV;
+          setenv("AFL_DONT_OPTIMIZE", "1", 1);
 
         } else {
 

src/afl-forkserver.c

@@ -226,6 +226,8 @@ void afl_fsrv_init(afl_forkserver_t *fsrv) {
   /* Settings */
   fsrv->use_stdin = true;
   fsrv->no_unlink = false;
+  fsrv->keep_coverage = false;
+  fsrv->coverage_map = NULL;
   fsrv->exec_tmout = EXEC_TIMEOUT;
   fsrv->init_tmout = EXEC_TIMEOUT * FORK_WAIT_MULT;
   fsrv->mem_limit = MEM_LIMIT;
@@ -276,6 +278,9 @@ void afl_fsrv_init_dup(afl_forkserver_t *fsrv_to, afl_forkserver_t *from) {
   fsrv_to->init_child_func = from->init_child_func;
   // Note: do not copy ->add_extra_func or ->persistent_record*
 
+  fsrv_to->keep_coverage = false;
+  fsrv_to->coverage_map = NULL;
+
   list_append(&fsrv_list, fsrv_to);
 
 }
@@ -1175,6 +1180,12 @@ void afl_fsrv_start(afl_forkserver_t *fsrv, char **argv,
 
     }
 
+    if (unlikely(fsrv->keep_coverage)) {
+
+      fsrv->coverage_map = ck_alloc(fsrv->map_size);
+
+    }
+
     return;
 
   }
@@ -1407,6 +1418,19 @@ void afl_fsrv_kill(afl_forkserver_t *fsrv) {
   afl_nyx_runner_kill(fsrv);
 #endif
 
+  if (unlikely(fsrv->keep_coverage)) {
+
+    // TODO maybe we should write that to out_dir
+    FILE *f = fopen("covmap.dump", "w");
+    if (!f) { PFATAL("creating covmap.dump failed"); }
+    fwrite(fsrv->coverage_map, fsrv->map_size, 1, f);
+    fclose(f);
+    OKF("Wrote coverage map to covmap.dump");
+    ck_free(fsrv->coverage_map);
+    fsrv->coverage_map = NULL;
+
+  }
+
 }
 
 /* Get the map size from the target forkserver */
@@ -1565,7 +1589,22 @@ afl_fsrv_run_target(afl_forkserver_t *fsrv, u32 timeout,
     enum NyxReturnValue ret_val =
         fsrv->nyx_handlers->nyx_exec(fsrv->nyx_runner);
 
-    fsrv->total_execs++;
+    ++fsrv->total_execs;
+
+    if (unlikely(fsrv->keep_coverage)) {
+
+      for (u32 i = 0; i < fsrv->map_size; ++i) {
+
+        if (unlikely(fsrv->trace_bits[i]) &&
+            likely(fsrv->coverage_map[i] < 255)) {
+
+          ++fsrv->coverage_map[i];
+
+        }
+
+      }
+
+    }
 
     switch (ret_val) {
 

src/afl-showmap.c

@@ -75,8 +75,9 @@ static u8 *in_dir = NULL,              /* input folder                      */
 
 static u8 outfile[PATH_MAX];
 
-static u8 *in_data,                    /* Input data                        */
-    *coverage_map;                     /* Coverage map                      */
+static u8 *in_data;                    /* Input data                        */
+
+static u8 *coverage_map;               /* Coverage map                      */
 
 static u64 total;                      /* tuple content information         */
 static u32 tcnt, highest;              /* tuple content information         */
@@ -97,7 +98,8 @@ static bool quiet_mode,                /* Hide non-essential messages?      */
     no_classify,                       /* do not classify counts            */
     debug,                             /* debug mode                        */
     print_filenames,                   /* print the current filename        */
-    wait_for_gdb;
+    wait_for_gdb,                      /* wait for gdb to allow attaching   */
+    code_cov;                          /* code coverage cmdline parameter   */
 
 static volatile u8 stop_soon,          /* Ctrl-C pressed?                   */
     child_crashed;                     /* Child crashed?                    */
@@ -244,13 +246,20 @@ static void analyze_results(afl_forkserver_t *fsrv) {
 
       total += fsrv->trace_bits[i];
       if (fsrv->trace_bits[i] > highest) highest = fsrv->trace_bits[i];
-      // if (!coverage_map[i]) { coverage_map[i] = 1; }
+      if (code_cov) {
+
+        if (coverage_map[i] < 255 && fsrv->trace_bits[i]) { coverage_map[i]++; }
+
+      } else {
+
         coverage_map[i] |= fsrv->trace_bits[i];
 
       }
 
     }
 
+  }
+
 }
 
 /* Write results. */
@@ -348,6 +357,76 @@ static u32 write_results_to_file(afl_forkserver_t *fsrv, u8 *outfile) {
 
 }
 
+#if 0
+static u32 write_results_to_file32(u32 *map, u8 *outfile) {
+
+  s32 fd;
+  u32 i, ret = 0;
+
+  if (!outfile || !*outfile) {
+
+    FATAL("Output filename not set (Bug in AFL++?)");
+
+  }
+
+  if (!strncmp(outfile, "/dev/", 5)) {
+
+    fd = open(outfile, O_WRONLY);
+
+    if (fd < 0) { PFATAL("Unable to open '%s'", out_file); }
+
+  } else if (!strcmp(outfile, "-")) {
+
+    fd = dup(1);
+    if (fd < 0) { PFATAL("Unable to open stdout"); }
+
+  } else {
+
+    unlink(outfile);                                       /* Ignore errors */
+    fd = open(outfile, O_WRONLY | O_CREAT | O_EXCL, DEFAULT_PERMISSION);
+    if (fd < 0) { PFATAL("Unable to create '%s'", outfile); }
+
+  }
+
+  if (binary_mode) {
+
+    for (i = 0; i < map_size; i++) {
+
+      if (map[i]) { ret++; }
+
+    }
+
+    ck_write(fd, map, map_size, outfile);
+    close(fd);
+
+  } else {
+
+    FILE *f = fdopen(fd, "w");
+
+    if (!f) { PFATAL("fdopen() failed"); }
+
+    for (i = 0; i < map_size; i++) {
+
+      if (!map[i]) { continue; }
+      ret++;
+
+      total += map[i];
+      if (highest < map[i]) { highest = map[i]; }
+
+      fprintf(f, "%06u:%u\n", i, map[i]);
+
+    }
+
+    fclose(f);
+
+  }
+
+  return ret;
+
+}
+
+#endif
+
 void pre_afl_fsrv_write_to_testcase(afl_forkserver_t *fsrv, u8 *mem, u32 len) {
 
   static u8 buf[MAX_FILE];
@@ -1096,10 +1175,17 @@ int main(int argc, char **argv_orig, char **envp) {
 
   if (getenv("AFL_QUIET") != NULL) { be_quiet = true; }
 
-  while ((opt = getopt(argc, argv, "+i:I:o:f:m:t:AeqCZOH:QUWbcrshXY")) > 0) {
+  while ((opt = getopt(argc, argv, "+i:I:o:f:m:t:AeqCZOH:QUWbcrshVXY")) > 0) {
 
     switch (opt) {
 
+      case 'V':
+        code_cov = true;
+        collect_coverage = true;
+        quiet_mode = true;
+        setenv("__AFL_CODE_COVERAGE", "1", 1);
+        break;
+
       case 's':
         no_classify = true;
         break;
@@ -1673,8 +1759,12 @@ int main(int argc, char **argv_orig, char **envp) {
 
     } else {
 
-      if ((coverage_map = (u8 *)malloc(map_size + 64)) == NULL)
+      if ((coverage_map = (u8 *)malloc((map_size + 64))) == NULL) {
+
         FATAL("coult not grab memory");
+
+      }
+
       edges_only = false;
       raw_instr_output = true;