added AFL_MAP_SIZE (wip)

15 files changed, 116 insertions, 76 deletions

diff --git a/docs/Changelog.md b/docs/Changelog.md
index 60d83508..3ad80b7b 100644
--- a/docs/Changelog.md
+++ b/docs/Changelog.md
@@ -26,6 +26,7 @@ sending a mail to <afl-users+subscribe@googlegroups.com>.
     static global and local variable comparisons (cannot find all though)
   - extended forkserver: map_size and more information is communicated to
     afl-fuzz (and afl-fuzz acts accordingly)
+  - new environment variable: AFL_MAP_SIZE to specify the size of the shared map
   - if AFL_CC/AFL_CXX is set but empty afl compilers did fail, fixed
     (this bug is in vanilla afl too)
   - added NO_PYTHON flag to disable python support when building afl-fuzz
diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 363776cb..88cacc4f 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -325,6 +325,8 @@ typedef struct afl_env_vars {
       *afl_python_module, *afl_path, *afl_hang_tmout, *afl_skip_crashes,
       *afl_preload;
 
+  uint32_t map_size;
+
 } afl_env_vars_t;
 
 struct afl_pass_stat {
diff --git a/include/config.h b/include/config.h
index f0274fd3..fae97a42 100644
--- a/include/config.h
+++ b/include/config.h
@@ -407,8 +407,7 @@
 #define FS_OPT_SNAPSHOT 0x20000000
 #define FS_OPT_AUTODICT 0x10000000
 #define FS_OPT_GET_MAPSIZE(x) (((x & 0x00fffffe) >> 1) + 1)
-#define FS_OPT_SET_MAPSIZE(x) \
-  (x <= 1 || x > MAP_SIZE || x > 0x1000000 ? 0 : ((x - 1) << 1))
+#define FS_OPT_SET_MAPSIZE(x) (x <= 1 || x > 0x1000000 ? 0 : ((x - 1) << 1))
 
 #endif                                                  /* ! _HAVE_CONFIG_H */
 
diff --git a/include/forkserver.h b/include/forkserver.h
index ac89b681..d76dfc7a 100644
--- a/include/forkserver.h
+++ b/include/forkserver.h
@@ -61,7 +61,7 @@ typedef struct afl_forkserver {
   u64 total_execs;                      /* How often run_target was called  */
 
   u8 *out_file,                         /* File to fuzz, if any             */
-      *target_path;                                   /* Path of the target */
+      *target_path;                     /* Path of the target               */
 
   FILE *plot_file;                      /* Gnuplot output file              */
 
diff --git a/src/afl-analyze.c b/src/afl-analyze.c
index 6f946ed5..3d86efb1 100644
--- a/src/afl-analyze.c
+++ b/src/afl-analyze.c
@@ -84,6 +84,7 @@ static volatile u8 stop_soon,          /* Ctrl-C pressed?                   */
 
 static u8 *target_path;
 static u8  qemu_mode;
+static u32 map_size = MAP_SIZE;
 
 /* Constants used for describing byte behavior. */
 
@@ -115,7 +116,7 @@ static u8 count_class_lookup[256] = {
 
 static void classify_counts(u8 *mem) {
 
-  u32 i = MAP_SIZE;
+  u32 i = map_size;
 
   if (edges_only) {
 
@@ -144,7 +145,7 @@ static void classify_counts(u8 *mem) {
 static inline u8 anything_set(void) {
 
   u32 *ptr = (u32 *)trace_bits;
-  u32  i = (MAP_SIZE >> 2);
+  u32  i = (map_size >> 2);
 
   while (i--)
     if (*(ptr++)) return 1;
@@ -217,7 +218,7 @@ static u32 analyze_run_target(char **argv, u8 *mem, u32 len, u8 first_run) {
   s32 prog_in_fd;
   u32 cksum;
 
-  memset(trace_bits, 0, MAP_SIZE);
+  memset(trace_bits, 0, map_size);
   MEM_BARRIER();
 
   prog_in_fd = write_to_file(prog_in, mem, len);
@@ -311,7 +312,7 @@ static u32 analyze_run_target(char **argv, u8 *mem, u32 len, u8 first_run) {
 
   }
 
-  cksum = hash32(trace_bits, MAP_SIZE, HASH_CONST);
+  cksum = hash32(trace_bits, map_size, HASH_CONST);
 
   /* We don't actually care if the target is crashing or not,
      except that when it does, the checksum should be different. */
@@ -811,7 +812,7 @@ int main(int argc, char **argv, char **envp) {
 
   s32    opt;
   u8     mem_limit_given = 0, timeout_given = 0, unicorn_mode = 0, use_wine = 0;
-  char **use_argv;
+  char **use_argv, *ptr;
 
   doc_path = access(DOC_PATH, F_OK) ? "docs" : DOC_PATH;
 
@@ -931,12 +932,21 @@ int main(int argc, char **argv, char **envp) {
 
   if (optind == argc || !in_file) usage(argv[0]);
 
+  if ((ptr = getenv("AFL_MAP_SIZE")) || (ptr = getenv("AFL_MAPSIZE"))) {
+
+    map_size = atoi(ptr);
+    if (map_size < 8 || map_size > (1 << 29))
+      FATAL("illegal AFL_MAP_SIZE %u, must be between 2^3 and 2^30", map_size);
+    if (map_size % 8) map_size = (((map_size >> 3) + 1) << 3);
+
+  }
+
   use_hex_offsets = !!get_afl_env("AFL_ANALYZE_HEX");
 
   check_environment_vars(envp);
 
   sharedmem_t shm = {0};
-  trace_bits = afl_shm_init(&shm, MAP_SIZE, 0);
+  trace_bits = afl_shm_init(&shm, map_size, 0);
   atexit(at_exit_handler);
   setup_signal_handlers();
 
diff --git a/src/afl-common.c b/src/afl-common.c
index c9f09d38..3210ee97 100644
--- a/src/afl-common.c
+++ b/src/afl-common.c
@@ -72,7 +72,7 @@ char *afl_environment_variables[] = {
     "AFL_LLVM_LTO_DONTWRITEID", "AFL_NO_ARITH", "AFL_NO_BUILTIN",
     "AFL_NO_CPU_RED", "AFL_NO_FORKSRV", "AFL_NO_UI",
     "AFL_NO_X86",  // not really an env but we dont want to warn on it
-    "AFL_PATH", "AFL_PERFORMANCE_FILE",
+    "AFL_MAP_SIZE", "AFL_MAPSIZE", "AFL_PATH", "AFL_PERFORMANCE_FILE",
     //"AFL_PERSISTENT", // not implemented anymore, so warn additionally
     "AFL_POST_LIBRARY", "AFL_PRELOAD", "AFL_PYTHON_MODULE", "AFL_QEMU_COMPCOV",
     "AFL_QEMU_COMPCOV_DEBUG", "AFL_QEMU_DEBUG_MAPS", "AFL_QEMU_DISABLE_CACHE",
diff --git a/src/afl-forkserver.c b/src/afl-forkserver.c
index 006764d9..9b915a7a 100644
--- a/src/afl-forkserver.c
+++ b/src/afl-forkserver.c
@@ -407,21 +407,26 @@ void afl_fsrv_start(afl_forkserver_t *fsrv, char **argv,
 
       if ((status & FS_OPT_MAPSIZE) == FS_OPT_MAPSIZE) {
 
-        fsrv->map_size = FS_OPT_GET_MAPSIZE(status);
-        if (unlikely(fsrv->map_size % 8)) {
+        u32 tmp_map_size = FS_OPT_GET_MAPSIZE(status);
+
+        if (!fsrv->map_size) fsrv->map_size = MAP_SIZE;
+
+        if (unlikely(tmp_map_size % 8)) {
 
           // should not happen
-          WARNF("Target reported non-aligned map size of %ud", fsrv->map_size);
-          fsrv->map_size = (((fsrv->map_size + 8) >> 3) << 3);
+          WARNF("Target reported non-aligned map size of %ud", tmp_map_size);
+          tmp_map_size = (((tmp_map_size + 8) >> 3) << 3);
 
         }
 
-        if (!be_quiet) ACTF("Target map size: %u", fsrv->map_size);
-        if (fsrv->map_size > MAP_SIZE)
+        if (!be_quiet) ACTF("Target map size: %u", tmp_map_size);
+        if (tmp_map_size > fsrv->map_size)
           FATAL(
               "Target's coverage map size of %u is larger than the one this "
-              "afl++ is compiled with (%u) (change MAP_SIZE and recompile)\n",
-              fsrv->map_size, MAP_SIZE);
+              "afl++ is set with (%u) (change MAP_SIZE_POW2 in config.h and "
+              "recompile or set AFL_MAP_SIZE)\n",
+              tmp_map_size, fsrv->map_size);
+        fsrv->map_size = tmp_map_size;
 
       }
 
diff --git a/src/afl-fuzz-bitmap.c b/src/afl-fuzz-bitmap.c
index be8f504e..0823deed 100644
--- a/src/afl-fuzz-bitmap.c
+++ b/src/afl-fuzz-bitmap.c
@@ -43,7 +43,7 @@ void write_bitmap(afl_state_t *afl) {
 
   if (fd < 0) PFATAL("Unable to open '%s'", fname);
 
-  ck_write(fd, afl->virgin_bits, MAP_SIZE, fname);
+  ck_write(fd, afl->virgin_bits, afl->fsrv.map_size, fname);
 
   close(fd);
 
@@ -145,8 +145,6 @@ u32 count_bits(afl_state_t *afl, u8 *mem) {
   u32  i = (afl->fsrv.map_size >> 2);
   u32  ret = 0;
 
-  if (i == 0) i = 1;
-
   while (i--) {
 
     u32 v = *(ptr++);
@@ -181,8 +179,6 @@ u32 count_bytes(afl_state_t *afl, u8 *mem) {
   u32  i = (afl->fsrv.map_size >> 2);
   u32  ret = 0;
 
-  if (i == 0) i = 1;
-
   while (i--) {
 
     u32 v = *(ptr++);
@@ -208,8 +204,6 @@ u32 count_non_255_bytes(afl_state_t *afl, u8 *mem) {
   u32  i = (afl->fsrv.map_size >> 2);
   u32  ret = 0;
 
-  if (i == 0) i = 1;
-
   while (i--) {
 
     u32 v = *(ptr++);
@@ -246,8 +240,6 @@ void simplify_trace(afl_state_t *afl, u64 *mem) {
 
   u32 i = (afl->fsrv.map_size >> 3);
 
-  if (i == 0) i = 1;
-
   while (i--) {
 
     /* Optimize for sparse bitmaps. */
@@ -281,8 +273,6 @@ void simplify_trace(afl_state_t *afl, u32 *mem) {
 
   u32 i = (afl->fsrv.map_size >> 2);
 
-  if (i == 0) i = 1;
-
   while (i--) {
 
     /* Optimize for sparse bitmaps. */
@@ -347,8 +337,6 @@ void classify_counts(afl_forkserver_t *fsrv) {
 
   u32 i = (fsrv->map_size >> 3);
 
-  if (i == 0) i = 1;
-
   while (i--) {
 
     /* Optimize for sparse bitmaps. */
@@ -378,8 +366,6 @@ void classify_counts(afl_forkserver_t *fsrv) {
 
   u32 i = (fsrv->map_size >> 2);
 
-  if (i == 0) i = 1;
-
   while (i--) {
 
     /* Optimize for sparse bitmaps. */
diff --git a/src/afl-fuzz-init.c b/src/afl-fuzz-init.c
index 7131ceed..3da348d2 100644
--- a/src/afl-fuzz-init.c
+++ b/src/afl-fuzz-init.c
@@ -442,23 +442,6 @@ void read_testcases(afl_state_t *afl) {
 
 }
 
-/* Examine map coverage. Called once, for first test case. */
-
-static void check_map_coverage(afl_state_t *afl) {
-
-  u32 i;
-
-  if (count_bytes(afl, afl->fsrv.trace_bits) < 100) return;
-
-  for (i = (1 << (MAP_SIZE_POW2 - 1)); i < MAP_SIZE; ++i)
-    if (afl->fsrv.trace_bits[i]) return;
-
-  if (afl->fsrv.map_size != MAP_SIZE) return;
-
-  WARNF("Recompile binary with newer version of afl to improve coverage!");
-
-}
-
 /* Perform dry run of all test cases to confirm that the app is working as
    expected. This is done only for the initial inputs, and only once. */
 
@@ -501,8 +484,6 @@ void perform_dry_run(afl_state_t *afl) {
 
       case FSRV_RUN_OK:
 
-        if (q == afl->queue) check_map_coverage(afl);
-
         if (afl->crash_mode) FATAL("Test case '%s' does *NOT* crash", fn);
 
         break;
diff --git a/src/afl-fuzz-queue.c b/src/afl-fuzz-queue.c
index d05eee08..373f12d8 100644
--- a/src/afl-fuzz-queue.c
+++ b/src/afl-fuzz-queue.c
@@ -249,7 +249,6 @@ void update_bitmap_score(afl_state_t *afl, struct queue_entry *q) {
       if (!q->trace_mini) {
 
         u32 len = (afl->fsrv.map_size >> 3);
-        if (len == 0) len = 1;
         q->trace_mini = ck_alloc(len);
         minimize_bits(afl, q->trace_mini, afl->fsrv.trace_bits);
 
@@ -272,12 +271,12 @@ void cull_queue(afl_state_t *afl) {
   struct queue_entry *q;
   u32                 len = (afl->fsrv.map_size >> 3);
   u32                 i;
-  u8                  temp_v[MAP_SIZE >> 3];
-
-  if (len == 0) len = 1;
+  u8 *                temp_v;
 
   if (afl->dumb_mode || !afl->score_changed) return;
 
+  temp_v = ck_alloc(afl->fsrv.map_size >> 3);
+
   afl->score_changed = 0;
 
   memset(temp_v, 255, len);
@@ -325,6 +324,8 @@ void cull_queue(afl_state_t *afl) {
 
   }
 
+  ck_free(temp_v);
+
 }
 
 /* Calculate case desirability score to adjust the length of havoc fuzzing.
diff --git a/src/afl-fuzz-state.c b/src/afl-fuzz-state.c
index 7664c521..7d068258 100644
--- a/src/afl-fuzz-state.c
+++ b/src/afl-fuzz-state.c
@@ -99,7 +99,11 @@ void afl_state_init(afl_state_t *afl) {
 
   afl->fsrv.use_stdin = 1;
 
-  afl->fsrv.map_size = MAP_SIZE;
+  if (afl->afl_env.map_size > 8 && afl->afl_env.map_size <= (1 << 29))
+    afl->fsrv.map_size = afl->afl_env.map_size;
+  else
+    afl->fsrv.map_size = MAP_SIZE;
+
   afl->fsrv.function_opt = (u8 *)afl;
   afl->fsrv.function_ptr = &maybe_add_auto;
 
@@ -324,6 +328,24 @@ void read_afl_environment(afl_state_t *afl, char **envp) {
             afl->afl_env.afl_path =
                 (u8 *)get_afl_env(afl_environment_variables[i]);
 
+          } else if (!strncmp(env, "AFL_MAP_SIZE",
+
+                              afl_environment_variable_len) ||
+                     !strncmp(env, "AFL_MAPSIZE",
+                              afl_environment_variable_len)) {
+
+            afl->afl_env.map_size =
+                atoi((u8 *)get_afl_env(afl_environment_variables[i]));
+
+            if (afl->afl_env.map_size < 8 || afl->afl_env.map_size > (1 << 29))
+              FATAL(
+                  "the specified AFL_MAP_SIZE size is illegal and must be "
+                  "between 2^3 and 2^30: %u\n",
+                  afl->afl_env.map_size);
+
+            if (afl->afl_env.map_size % 8)
+              afl->afl_env.map_size = (((afl->afl_env.map_size >> 3) + 1) << 3);
+
           } else if (!strncmp(env, "AFL_PRELOAD",
 
                               afl_environment_variable_len)) {
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index 925dbb1a..3cf57f86 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -249,6 +249,7 @@ int main(int argc, char **argv_orig, char **envp) {
 
   if (get_afl_env("AFL_DEBUG")) afl->debug = 1;
   read_afl_environment(afl, envp);
+  if (afl->afl_env.map_size) afl->fsrv.map_size = afl->afl_env.map_size;
   exit_1 = !!afl->afl_env.afl_bench_just_one;
 
   SAYF(cCYA "afl-fuzz" VERSION cRST
@@ -476,7 +477,7 @@ 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, MAP_SIZE);
+        read_bitmap(afl->in_bitmap, afl->virgin_bits, afl->fsrv.map_size);
         break;
 
       case 'C':                                               /* crash mode */
@@ -910,13 +911,14 @@ int main(int argc, char **argv_orig, char **envp) {
   check_crash_handling();
   check_cpu_governor(afl);
 
-  afl->fsrv.trace_bits = afl_shm_init(&afl->shm, MAP_SIZE, afl->dumb_mode);
+  afl->fsrv.trace_bits =
+      afl_shm_init(&afl->shm, afl->fsrv.map_size, afl->dumb_mode);
 
   setup_post(afl);
 
-  if (!afl->in_bitmap) memset(afl->virgin_bits, 255, MAP_SIZE);
-  memset(afl->virgin_tmout, 255, MAP_SIZE);
-  memset(afl->virgin_crash, 255, MAP_SIZE);
+  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();
 
diff --git a/src/afl-gcc.c b/src/afl-gcc.c
index 32cd36cb..86a88014 100644
--- a/src/afl-gcc.c
+++ b/src/afl-gcc.c
@@ -411,6 +411,14 @@ int main(int argc, char **argv) {
 
   }
 
+  u8 *ptr;
+  if ((ptr = getenv("AFL_MAP_SIZE")) || (ptr = getenv("AFL_MAPSIZE"))) {
+
+    u32 map_size = atoi(ptr);
+    if (map_size != MAP_SIZE) FATAL("AFL_MAP_SIZE is not supported by afl-gcc");
+
+  }
+
   find_as(argv[0]);
 
   edit_params(argc, argv);
diff --git a/src/afl-showmap.c b/src/afl-showmap.c
index 21e18061..c1561b4c 100644
--- a/src/afl-showmap.c
+++ b/src/afl-showmap.c
@@ -72,6 +72,8 @@ static u32 total, highest;             /* tuple content information         */
 static u32 in_len,                     /* Input data length                 */
     arg_offset;                        /* Total number of execs             */
 
+static u32 map_size = MAP_SIZE;
+
 static u8 quiet_mode,                  /* Hide non-essential messages?      */
     edges_only,                        /* Ignore hit counts?                */
     raw_instr_output,                  /* Do not apply AFL filters          */
@@ -112,7 +114,7 @@ static void classify_counts(afl_forkserver_t *fsrv) {
   u8 *      mem = fsrv->trace_bits;
   const u8 *map = binary_mode ? count_class_binary : count_class_human;
 
-  u32 i = MAP_SIZE;
+  u32 i = map_size;
 
   if (edges_only) {
 
@@ -175,10 +177,10 @@ static u32 write_results_to_file(afl_forkserver_t *fsrv, u8 *outfile) {
 
   if (binary_mode) {
 
-    for (i = 0; i < MAP_SIZE; i++)
+    for (i = 0; i < map_size; i++)
       if (fsrv->trace_bits[i]) ret++;
 
-    ck_write(fd, fsrv->trace_bits, MAP_SIZE, outfile);
+    ck_write(fd, fsrv->trace_bits, map_size, outfile);
     close(fd);
 
   } else {
@@ -187,7 +189,7 @@ static u32 write_results_to_file(afl_forkserver_t *fsrv, u8 *outfile) {
 
     if (!f) PFATAL("fdopen() failed");
 
-    for (i = 0; i < MAP_SIZE; i++) {
+    for (i = 0; i < map_size; i++) {
 
       if (!fsrv->trace_bits[i]) continue;
       ret++;
@@ -535,7 +537,7 @@ int main(int argc, char **argv_orig, char **envp) {
   s32    opt, i;
   u8     mem_limit_given = 0, timeout_given = 0, unicorn_mode = 0, use_wine = 0;
   u32    tcnt = 0;
-  char **use_argv;
+  char **use_argv, *ptr;
 
   char **argv = argv_cpy_dup(argc, argv_orig);
 
@@ -543,6 +545,16 @@ int main(int argc, char **argv_orig, char **envp) {
   afl_forkserver_t *fsrv = &fsrv_var;
   afl_fsrv_init(fsrv);
 
+  if ((ptr = getenv("AFL_MAP_SIZE")) || (ptr = getenv("AFL_MAPSIZE"))) {
+
+    map_size = atoi(ptr);
+    if (map_size < 8 || map_size > (1 << 29))
+      FATAL("illegal AFL_MAP_SIZE %u, must be between 2^3 and 2^30", map_size);
+    if (map_size % 8) map_size = (((map_size >> 3) + 1) << 3);
+    fsrv->map_size = map_size;
+
+  }
+
   doc_path = access(DOC_PATH, F_OK) ? "docs" : DOC_PATH;
 
   if (getenv("AFL_QUIET") != NULL) be_quiet = 1;
@@ -715,7 +727,7 @@ int main(int argc, char **argv_orig, char **envp) {
   check_environment_vars(envp);
 
   sharedmem_t shm = {0};
-  fsrv->trace_bits = afl_shm_init(&shm, MAP_SIZE, 0);
+  fsrv->trace_bits = afl_shm_init(&shm, map_size, 0);
   setup_signal_handlers();
 
   set_up_environment(fsrv);
diff --git a/src/afl-tmin.c b/src/afl-tmin.c
index 0a462e9a..e366d260 100644
--- a/src/afl-tmin.c
+++ b/src/afl-tmin.c
@@ -70,7 +70,8 @@ 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     */
+    missed_paths,                      /* Misses due to exec path diffs     */
+    map_size = MAP_SIZE;
 
 static u8 crash_mode,                  /* Crash-centric mode?               */
     hang_mode,                         /* Minimize as long as it hangs      */
@@ -105,7 +106,7 @@ static const u8 count_class_lookup[256] = {
 
 static void apply_mask(u32 *mem, u32 *mask) {
 
-  u32 i = (MAP_SIZE >> 2);
+  u32 i = (map_size >> 2);
 
   if (!mask) return;
 
@@ -122,7 +123,7 @@ static void apply_mask(u32 *mem, u32 *mask) {
 static void classify_counts(afl_forkserver_t *fsrv) {
 
   u8 *mem = fsrv->trace_bits;
-  u32 i = MAP_SIZE;
+  u32 i = map_size;
 
   if (edges_only) {
 
@@ -151,7 +152,7 @@ static void classify_counts(afl_forkserver_t *fsrv) {
 static inline u8 anything_set(afl_forkserver_t *fsrv) {
 
   u32 *ptr = (u32 *)fsrv->trace_bits;
-  u32  i = (MAP_SIZE >> 2);
+  u32  i = (map_size >> 2);
 
   while (i--)
     if (*(ptr++)) return 1;
@@ -755,7 +756,7 @@ int main(int argc, char **argv_orig, char **envp) {
 
   s32    opt;
   u8     mem_limit_given = 0, timeout_given = 0, unicorn_mode = 0, use_wine = 0;
-  char **use_argv;
+  char **use_argv, *ptr;
 
   char **argv = argv_cpy_dup(argc, argv_orig);
 
@@ -763,6 +764,16 @@ int main(int argc, char **argv_orig, char **envp) {
   afl_forkserver_t *fsrv = &fsrv_var;
   afl_fsrv_init(fsrv);
 
+  if ((ptr = getenv("AFL_MAP_SIZE")) || (ptr = getenv("AFL_MAPSIZE"))) {
+
+    map_size = atoi(ptr);
+    if (map_size < 8 || map_size > (1 << 29))
+      FATAL("illegal AFL_MAP_SIZE %u, must be between 2^3 and 2^30", map_size);
+    if (map_size % 8) map_size = (((map_size >> 3) + 1) << 3);
+    fsrv->map_size = map_size;
+
+  }
+
   doc_path = access(DOC_PATH, F_OK) ? "docs" : DOC_PATH;
 
   SAYF(cCYA "afl-tmin" VERSION cRST " by Michal Zalewski\n");
@@ -910,8 +921,8 @@ int main(int argc, char **argv_orig, char **envp) {
            to be useful. */
 
         if (mask_bitmap) FATAL("Multiple -B options not supported");
-        mask_bitmap = ck_alloc(MAP_SIZE);
-        read_bitmap(optarg, mask_bitmap, MAP_SIZE);
+        mask_bitmap = ck_alloc(map_size);
+        read_bitmap(optarg, mask_bitmap, map_size);
         break;
 
       case 'h':
@@ -928,7 +939,7 @@ int main(int argc, char **argv_orig, char **envp) {
   check_environment_vars(envp);
 
   sharedmem_t shm = {0};
-  fsrv->trace_bits = afl_shm_init(&shm, MAP_SIZE, 0);
+  fsrv->trace_bits = afl_shm_init(&shm, map_size, 0);
 
   atexit(at_exit_handler);
   setup_signal_handlers();