1 files changed, 64 insertions, 9 deletions

diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index f1f92717..49733594 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -342,7 +342,6 @@ int main(int argc, char **argv_orig, char **envp) {
   afl->debug = debug;
   afl_fsrv_init(&afl->fsrv);
   if (debug) { afl->fsrv.debug = true; }
-
   read_afl_environment(afl, envp);
   if (afl->shm.map_size) { afl->fsrv.map_size = afl->shm.map_size; }
   exit_1 = !!afl->afl_env.afl_bench_just_one;
@@ -702,7 +701,6 @@ int main(int argc, char **argv_orig, char **envp) {
         if (afl->in_bitmap) { FATAL("Multiple -B options not supported"); }
 
         afl->in_bitmap = optarg;
-        read_bitmap(afl->in_bitmap, afl->virgin_bits, afl->fsrv.map_size);
         break;
 
       case 'C':                                               /* crash mode */
@@ -1369,13 +1367,6 @@ int main(int argc, char **argv_orig, char **envp) {
   set_scheduler_mode(SCHEDULER_MODE_LOW_LATENCY);
   #endif
 
-  afl->fsrv.trace_bits =
-      afl_shm_init(&afl->shm, afl->fsrv.map_size, afl->non_instrumented_mode);
-
-  if (!afl->in_bitmap) { memset(afl->virgin_bits, 255, afl->fsrv.map_size); }
-  memset(afl->virgin_tmout, 255, afl->fsrv.map_size);
-  memset(afl->virgin_crash, 255, afl->fsrv.map_size);
-
   init_count_class16();
 
   if (afl->is_main_node && check_main_node_exists(afl) == 1) {
@@ -1542,6 +1533,70 @@ 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);
+
+  if (!afl->non_instrumented_mode) {
+
+    afl->fsrv.map_size = 4194304;  // dummy temporary value
+
+    u32 new_map_size = afl_fsrv_get_mapsize(
+        &afl->fsrv, afl->argv, &afl->stop_soon, afl->afl_env.afl_debug_child);
+
+    if (new_map_size && new_map_size != 4194304) {
+
+      // only reinitialize when it makes sense
+      if (map_size != new_map_size) {
+
+        //      if (map_size < new_map_size ||
+        //          (new_map_size > map_size && new_map_size - map_size >
+        //          MAP_SIZE)) {
+
+        OKF("Re-initializing maps to %u bytes", new_map_size);
+
+        afl->virgin_bits = ck_realloc(afl->virgin_bits, map_size);
+        afl->virgin_tmout = ck_realloc(afl->virgin_tmout, map_size);
+        afl->virgin_crash = ck_realloc(afl->virgin_crash, map_size);
+        afl->var_bytes = ck_realloc(afl->var_bytes, map_size);
+        afl->top_rated = ck_realloc(afl->top_rated, map_size * sizeof(void *));
+        afl->clean_trace = ck_realloc(afl->clean_trace, map_size);
+        afl->clean_trace_custom = ck_realloc(afl->clean_trace_custom, map_size);
+        afl->first_trace = ck_realloc(afl->first_trace, map_size);
+        afl->map_tmp_buf = ck_realloc(afl->map_tmp_buf, map_size);
+
+        afl_shm_deinit(&afl->shm);
+        afl_fsrv_kill(&afl->fsrv);
+        afl->fsrv.map_size = new_map_size;
+        afl->fsrv.trace_bits = afl_shm_init(&afl->shm, afl->fsrv.map_size,
+                                            afl->non_instrumented_mode);
+        setenv("AFL_NO_AUTODICT", "1", 1);  // loaded already
+        afl_fsrv_start(&afl->fsrv, afl->argv, &afl->stop_soon,
+                       afl->afl_env.afl_debug_child);
+
+      }
+
+      map_size = new_map_size;
+
+    }
+
+    afl->fsrv.map_size = map_size;
+
+  }
+
+  // after we have the correct bitmap size we can read the bitmap -B option
+  // and set the virgin maps
+  if (!afl->in_bitmap) {
+
+    memset(afl->virgin_bits, 255, afl->fsrv.map_size);
+
+  } else {
+
+    read_bitmap(afl->in_bitmap, afl->virgin_bits, afl->fsrv.map_size);
+
+  }
+
+  memset(afl->virgin_tmout, 255, afl->fsrv.map_size);
+  memset(afl->virgin_crash, 255, afl->fsrv.map_size);
 
   if (afl->cmplog_binary) {