Pass dataflow map shm to fork server

2 files changed, 43 insertions, 3 deletions

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);
 
   }