Pass dataflow map shm to fork server

5 files changed, 53 insertions, 4 deletions

diff --git a/include/config.h b/include/config.h
index d8177a75..8cc344da 100644
--- a/include/config.h
+++ b/include/config.h
@@ -403,6 +403,7 @@
 /* Environment variable used to pass SHM ID to the called program. */
 
 #define SHM_ENV_VAR "__AFL_SHM_ID"
+#define DFG_SHM_ENV_VAR "__AFL_DFG_SHM_ID"
 
 /* Environment variable used to pass SHM FUZZ ID to the called program. */
 
@@ -470,6 +471,9 @@
   #define MAP_INITIAL_SIZE MAP_SIZE
 #endif
 
+#define DFG_MAP_SIZE_POW2 16
+#define DFG_MAP_SIZE (1U << DFG_MAP_SIZE_POW2)
+
 /* Maximum allocator request size (keep well under INT_MAX): */
 
 #define MAX_ALLOC 0x40000000
diff --git a/include/forkserver.h b/include/forkserver.h
index d3d0e086..882f4f01 100644
--- a/include/forkserver.h
+++ b/include/forkserver.h
@@ -103,7 +103,8 @@ typedef struct afl_forkserver {
 
   /* a program that includes afl-forkserver needs to define these */
 
-  u8 *trace_bits;                       /* SHM with instrumentation bitmap  */
+  u8 *trace_bits;                       /* SHM with code coverage bitmap    */
+  u32 *dfg_bits;                        /* SHM with DFG coverage bitmap     */
 
   s32 fsrv_pid,                         /* PID of the fork server           */
       child_pid,                        /* PID of the fuzzed program        */
diff --git a/include/sharedmem.h b/include/sharedmem.h
index 036fa560..ea03205a 100644
--- a/include/sharedmem.h
+++ b/include/sharedmem.h
@@ -38,15 +38,19 @@ typedef struct sharedmem {
   /* ================ Proteas ================ */
   int  g_shm_fd;
   char g_shm_file_path[L_tmpnam];
+  int  dfg_g_shm_fd;
+  char dfg_g_shm_file_path[L_tmpnam];
   int  cmplog_g_shm_fd;
   char cmplog_g_shm_file_path[L_tmpnam];
 /* ========================================= */
 #else
   s32 shm_id;                          /* ID of the SHM region              */
+  s32 dfg_shm_id;                      /* ID of the DFG SHM region          */
   s32 cmplog_shm_id;
 #endif
 
   u8 *map;                                          /* shared memory region */
+  u32 *dfg_map;                                 /* DFG shared memory region */
 
   size_t map_size;                                 /* actual allocated size */
 
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index b7f99ddc..2695adea 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -110,7 +110,8 @@ extern u64 time_spent_working;
 static void at_exit() {
 
   s32   i, pid1 = 0, pid2 = 0, pgrp = -1;
-  char *list[4] = {SHM_ENV_VAR, SHM_FUZZ_ENV_VAR, CMPLOG_SHM_ENV_VAR, NULL};
+  char *list[] = {SHM_ENV_VAR, SHM_FUZZ_ENV_VAR,
+                  CMPLOG_SHM_ENV_VAR, DFG_SHM_ENV_VAR, NULL};
   char *ptr;
 
   ptr = getenv("__AFL_TARGET_PID2");
@@ -2485,6 +2486,7 @@ int main(int argc, char **argv_orig, char **envp) {
   afl->argv = use_argv;
   afl->fsrv.trace_bits =
       afl_shm_init(&afl->shm, afl->fsrv.map_size, afl->non_instrumented_mode);
+  afl->fsrv.dfg_bits = afl->shm.dfg_map;
 
   if (!afl->non_instrumented_mode && !afl->fsrv.qemu_mode &&
       !afl->unicorn_mode && !afl->fsrv.frida_mode && !afl->fsrv.cs_mode &&
diff --git a/src/afl-sharedmem.c b/src/afl-sharedmem.c
index 1dea83f9..f05afd75 100644
--- a/src/afl-sharedmem.c
+++ b/src/afl-sharedmem.c
@@ -75,6 +75,7 @@ void afl_shm_deinit(sharedmem_t *shm) {
   } else {
 
     unsetenv(SHM_ENV_VAR);
+    unsetenv(DFG_SHM_ENV_VAR);
 
   }
 
@@ -142,15 +143,18 @@ void afl_shm_deinit(sharedmem_t *shm) {
 
 u8 *afl_shm_init(sharedmem_t *shm, size_t map_size,
                  unsigned char non_instrumented_mode) {
+  size_t dfg_map_size = sizeof(u32) * DFG_MAP_SIZE;
 
   shm->map_size = 0;
 
   shm->map = NULL;
+  shm->dfg_map = NULL;
   shm->cmp_map = NULL;
 
 #ifdef USEMMAP
 
   shm->g_shm_fd = -1;
+  shm->dfg_g_shm_fd = -1;
   shm->cmplog_g_shm_fd = -1;
 
   const int shmflags = O_RDWR | O_EXCL;
@@ -162,6 +166,7 @@ u8 *afl_shm_init(sharedmem_t *shm, size_t map_size,
   security warning that cannot be suppressed
   so we do this worse workaround */
   snprintf(shm->g_shm_file_path, L_tmpnam, "/afl_%d_%ld", getpid(), random());
+  snprintf(shm->dfg_g_shm_file_path, L_tmpnam, "/afl_%d_%ld", getpid(), random());
 
   #ifdef SHM_LARGEPAGE_ALLOC_DEFAULT
   /* trying to get large memory segment optimised and monitorable separately as
@@ -184,6 +189,12 @@ u8 *afl_shm_init(sharedmem_t *shm, size_t map_size,
 
     }
 
+    for (i = psizes - 1; shm->dfg_g_shm_fd == -1 && i >= 0; --i) {
+      if (sizes[i] == 0 || map_size % sizes[i])
+        continue;
+      shm->dfg_g_shm_fd = shm_create_largepage(shm->dfg_g_shm_file_path,
+        shmflags, i, SHM_LARGEPAGE_ALLOC_DEFAULT, DEFAULT_PERMISSION);
+    }
   }
 
   #endif
@@ -218,14 +229,35 @@ u8 *afl_shm_init(sharedmem_t *shm, size_t map_size,
 
   }
 
+  if (shm->dfg_g_shm_fd == -1)
+    shm->dfg_g_shm_fd = shm_open(shm->dfg_g_shm_file_path,
+      shmflags | O_CREAT, DEFAULT_PERMISSION);
+  if (shm->dfg_g_shm_fd == -1)
+    PFATAL("shm_open() failed");
+  if (ftruncate(shm->g_shm_fd, map_size))
+    PFATAL("setup_shm(): ftruncate() failed");
+  shm->dfg_map = mmap(0, dfg_map_size,
+    PROT_READ | PROT_WRITE, MAP_SHARED, shm->dfg_g_shm_fd, 0);
+  if (shm->dfg_map == MAP_FAILED) {
+    close(shm->dfg_g_shm_fd);
+    shm->dfg_g_shm_fd = -1;
+    shm_unlink(shm->dfg_g_shm_file_path);
+    shm->dfg_g_shm_file_path[0] = 0;
+    PFATAL("mmap() failed");
+  }
+
   /* If somebody is asking us to fuzz instrumented binaries in non-instrumented
      mode, we don't want them to detect instrumentation, since we won't be
      sending fork server commands. This should be replaced with better
      auto-detection later on, perhaps? */
 
-  if (!non_instrumented_mode) setenv(SHM_ENV_VAR, shm->g_shm_file_path, 1);
+  if (!non_instrumented_mode) {
+	setenv(SHM_ENV_VAR, shm->g_shm_file_path, 1);
+	setenv(DFG_SHM_ENV_VAR, shm->dfg_g_shm_file_path, 1);
+  }
 
   if (shm->map == (void *)-1 || !shm->map) PFATAL("mmap() failed");
+  if (shm->dfg_map == (void *)-1 || !shm->dfg_map) PFATAL("mmap() failed");
 
   if (shm->cmplog_mode) {
 
@@ -279,7 +311,10 @@ u8 *afl_shm_init(sharedmem_t *shm, size_t map_size,
   shm->shm_id =
       shmget(IPC_PRIVATE, map_size == MAP_SIZE ? map_size + 8 : map_size,
              IPC_CREAT | IPC_EXCL | DEFAULT_PERMISSION);
-  if (shm->shm_id < 0) {
+  shm->dfg_shm_id =
+      shmget(IPC_PRIVATE, dfg_map_size,
+             IPC_CREAT | IPC_EXCL | DEFAULT_PERMISSION);
+  if (shm->shm_id < 0 || shm->dfg_shm_id < 0) {
 
     PFATAL("shmget() failed, try running afl-system-config");
 
@@ -309,7 +344,10 @@ u8 *afl_shm_init(sharedmem_t *shm, size_t map_size,
        with better auto-detection later on, perhaps? */
 
     setenv(SHM_ENV_VAR, shm_str, 1);
+    ck_free(shm_str);
 
+    shm_str = alloc_printf("%d", shm->dfg_shm_id);
+    setenv(DFG_SHM_ENV_VAR, shm_str, 1);
     ck_free(shm_str);
 
   }