Merge pull request #426 from AFLplusplus/dev

Dev

1 files changed, 60 insertions, 7 deletions

diff --git a/src/afl-tmin.c b/src/afl-tmin.c
index 091e5177..2db1eae7 100644
--- a/src/afl-tmin.c
+++ b/src/afl-tmin.c
@@ -54,7 +54,9 @@
 
 #include <sys/wait.h>
 #include <sys/time.h>
-#include <sys/shm.h>
+#ifndef USEMMAP
+  #include <sys/shm.h>
+#endif
 #include <sys/stat.h>
 #include <sys/types.h>
 #include <sys/resource.h>
@@ -67,20 +69,27 @@ static u8 *in_file,                    /* Minimizer input test case         */
 static u8 *in_data;                    /* Input data for trimming           */
 
 static u32 in_len,                     /* Input data length                 */
-    orig_cksum,                        /* Original checksum                 */
     missed_hangs,                      /* Misses due to hangs               */
     missed_crashes,                    /* Misses due to crashes             */
     missed_paths,                      /* Misses due to exec path diffs     */
     map_size = MAP_SIZE;
 
+static u64 orig_cksum;                 /* Original checksum                 */
+
 static u8 crash_mode,                  /* Crash-centric mode?               */
     hang_mode,                         /* Minimize as long as it hangs      */
     exit_crash,                        /* Treat non-zero exit as crash?     */
     edges_only,                        /* Ignore hit counts?                */
-    exact_mode;                        /* Require path match for crashes?   */
+    exact_mode,                        /* Require path match for crashes?   */
+    remove_out_file,                   /* remove out_file on exit?          */
+    remove_shm = 1;                    /* remove shmem on exit?             */
 
 static volatile u8 stop_soon;          /* Ctrl-C pressed?                   */
 
+static afl_forkserver_t *fsrv;
+static sharedmem_t       shm;
+static sharedmem_t *     shm_fuzz;
+
 /*
  * forkserver section
  */
@@ -102,6 +111,18 @@ static const u8 count_class_lookup[256] = {
 
 };
 
+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_len = NULL;
+  fsrv->shmem_fuzz = NULL;
+  ck_free(shm_fuzz);
+  return NULL;
+
+}
+
 /* Apply mask to classified bitmap (if set). */
 
 static void apply_mask(u32 *mem, u32 *mask) {
@@ -166,7 +187,15 @@ static inline u8 anything_set(afl_forkserver_t *fsrv) {
 
 static void at_exit_handler(void) {
 
+  if (remove_shm) {
+
+    if (shm.map) afl_shm_deinit(&shm);
+    if (fsrv->use_shmem_fuzz) deinit_shmem(fsrv, shm_fuzz);
+
+  }
+
   afl_fsrv_killall();
+  if (remove_out_file) unlink(out_file);
 
 }
 
@@ -300,7 +329,7 @@ static u8 tmin_run_target(afl_forkserver_t *fsrv, char **argv, u8 *mem, u32 len,
 
   if (ret == FSRV_RUN_NOINST) { FATAL("Binary not instrumented?"); }
 
-  u32 cksum = hash32(fsrv->trace_bits, fsrv->map_size, HASH_CONST);
+  u64 cksum = hash64(fsrv->trace_bits, fsrv->map_size, HASH_CONST);
 
   if (first_run) { orig_cksum = cksum; }
 
@@ -620,6 +649,7 @@ static void set_up_environment(afl_forkserver_t *fsrv) {
     }
 
     out_file = alloc_printf("%s/.afl-tmin-temp-%u", use_dir, (u32)getpid());
+    remove_out_file = 1;
 
   }
 
@@ -809,8 +839,8 @@ int main(int argc, char **argv_orig, char **envp) {
 
   char **argv = argv_cpy_dup(argc, argv_orig);
 
-  afl_forkserver_t  fsrv_var = {0};
-  afl_forkserver_t *fsrv = &fsrv_var;
+  afl_forkserver_t fsrv_var = {0};
+  fsrv = &fsrv_var;
   afl_fsrv_init(fsrv);
   map_size = get_map_size();
   fsrv->map_size = map_size;
@@ -1007,7 +1037,8 @@ int main(int argc, char **argv_orig, char **envp) {
 
   check_environment_vars(envp);
 
-  sharedmem_t shm = {0};
+  /* initialize cmplog_mode */
+  shm.cmplog_mode = 0;
   fsrv->trace_bits = afl_shm_init(&shm, map_size, 0);
 
   atexit(at_exit_handler);
@@ -1049,11 +1080,31 @@ int main(int argc, char **argv_orig, char **envp) {
 
   SAYF("\n");
 
+  shm_fuzz = ck_alloc(sizeof(sharedmem_t));
+
+  /* initialize cmplog_mode */
+  shm_fuzz->cmplog_mode = 0;
+  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" : "");
 
@@ -1107,7 +1158,9 @@ int main(int argc, char **argv_orig, char **envp) {
 
   OKF("We're done here. Have a nice day!\n");
 
+  remove_shm = 0;
   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); }