shmem support for afl-tmin and afl-showmap

2 files changed, 71 insertions, 0 deletions

diff --git a/src/afl-showmap.c b/src/afl-showmap.c
index 6a26a949..f1926b05 100644
--- a/src/afl-showmap.c
+++ b/src/afl-showmap.c
@@ -566,6 +566,17 @@ static void usage(u8 *argv0) {
 
 }
 
+static sharedmem_t *deinit_shmem(afl_forkserver_t *fsrv,
+                                 sharedmem_t *     shm_fuzz) {
+
+  afl_shm_deinit(shm_fuzz);
+  fsrv->support_shmem_fuzz = 0;
+  fsrv->shmem_fuzz = NULL;
+  ck_free(shm_fuzz);
+  return NULL;
+
+}
+
 /* Main entry point */
 
 int main(int argc, char **argv_orig, char **envp) {
@@ -775,6 +786,17 @@ int main(int argc, char **argv_orig, char **envp) {
 
   check_environment_vars(envp);
 
+  if (getenv("AFL_DEBUG")) {
+
+    SAYF(cMGN "[D]" cRST);
+    for (int i = 0; i < argc; i++)
+      SAYF(" %s", argv[i]);
+    SAYF("\n");
+
+  }
+
+  //  if (afl->shmem_testcase_mode) { setup_testcase_shmem(afl); }
+
   sharedmem_t shm = {0};
   fsrv->trace_bits = afl_shm_init(&shm, map_size, 0);
   setup_signal_handlers();
@@ -829,6 +851,20 @@ int main(int argc, char **argv_orig, char **envp) {
 
   }
 
+  sharedmem_t *shm_fuzz = ck_alloc(sizeof(sharedmem_t));
+  u8 *         map = afl_shm_init(shm_fuzz, MAX_FILE + sizeof(u32), 1);
+  if (!map) { FATAL("BUG: Zero return from afl_shm_init."); }
+#ifdef USEMMAP
+  setenv(SHM_FUZZ_ENV_VAR, shm_fuzz->g_shm_file_path, 1);
+#else
+  u8 *shm_str = alloc_printf("%d", shm_fuzz->shm_id);
+  setenv(SHM_FUZZ_ENV_VAR, shm_str, 1);
+  ck_free(shm_str);
+#endif
+  fsrv->support_shmem_fuzz = 1;
+  fsrv->shmem_fuzz_len = (u32 *)map;
+  fsrv->shmem_fuzz = map + sizeof(u32);
+
   if (in_dir) {
 
     DIR *          dir_in, *dir_out;
@@ -897,6 +933,9 @@ int main(int argc, char **argv_orig, char **envp) {
     afl_fsrv_start(fsrv, use_argv, &stop_soon,
                    get_afl_env("AFL_DEBUG_CHILD_OUTPUT") ? 1 : 0);
 
+    if (fsrv->support_shmem_fuzz && !fsrv->use_shmem_fuzz)
+      shm_fuzz = deinit_shmem(fsrv, shm_fuzz);
+
     while (done == 0 && (dir_ent = readdir(dir_in))) {
 
       if (dir_ent->d_name[0] == '.') {
@@ -966,7 +1005,10 @@ int main(int argc, char **argv_orig, char **envp) {
 
   if (fsrv->target_path) { ck_free(fsrv->target_path); }
 
+  if (fsrv->use_shmem_fuzz) shm_fuzz = deinit_shmem(fsrv, shm_fuzz);
+
   afl_fsrv_deinit(fsrv);
+
   if (stdin_file) { ck_free(stdin_file); }
 
   argv_cpy_free(argv);
diff --git a/src/afl-tmin.c b/src/afl-tmin.c
index 9df5112b..8b028327 100644
--- a/src/afl-tmin.c
+++ b/src/afl-tmin.c
@@ -802,6 +802,17 @@ static void usage(u8 *argv0) {
 
 }
 
+static sharedmem_t *deinit_shmem(afl_forkserver_t *fsrv,
+                                 sharedmem_t *     shm_fuzz) {
+
+  afl_shm_deinit(shm_fuzz);
+  fsrv->support_shmem_fuzz = 0;
+  fsrv->shmem_fuzz = NULL;
+  ck_free(shm_fuzz);
+  return NULL;
+
+}
+
 /* Main entry point */
 
 int main(int argc, char **argv_orig, char **envp) {
@@ -1052,11 +1063,28 @@ int main(int argc, char **argv_orig, char **envp) {
 
   SAYF("\n");
 
+  sharedmem_t *shm_fuzz = ck_alloc(sizeof(sharedmem_t));
+  u8 *         map = afl_shm_init(shm_fuzz, MAX_FILE + sizeof(u32), 1);
+  if (!map) { FATAL("BUG: Zero return from afl_shm_init."); }
+#ifdef USEMMAP
+  setenv(SHM_FUZZ_ENV_VAR, shm_fuzz->g_shm_file_path, 1);
+#else
+  u8 *shm_str = alloc_printf("%d", shm_fuzz->shm_id);
+  setenv(SHM_FUZZ_ENV_VAR, shm_str, 1);
+  ck_free(shm_str);
+#endif
+  fsrv->support_shmem_fuzz = 1;
+  fsrv->shmem_fuzz_len = (u32 *)map;
+  fsrv->shmem_fuzz = map + sizeof(u32);
+
   read_initial_file();
 
   afl_fsrv_start(fsrv, use_argv, &stop_soon,
                  get_afl_env("AFL_DEBUG_CHILD_OUTPUT") ? 1 : 0);
 
+  if (fsrv->support_shmem_fuzz && !fsrv->use_shmem_fuzz)
+    shm_fuzz = deinit_shmem(fsrv, shm_fuzz);
+
   ACTF("Performing dry run (mem limit = %llu MB, timeout = %u ms%s)...",
        fsrv->mem_limit, fsrv->exec_tmout, edges_only ? ", edges only" : "");
 
@@ -1111,6 +1139,7 @@ int main(int argc, char **argv_orig, char **envp) {
   OKF("We're done here. Have a nice day!\n");
 
   afl_shm_deinit(&shm);
+  if (fsrv->use_shmem_fuzz) shm_fuzz = deinit_shmem(fsrv, shm_fuzz);
   afl_fsrv_deinit(fsrv);
   if (fsrv->target_path) { ck_free(fsrv->target_path); }
   if (mask_bitmap) { ck_free(mask_bitmap); }