forgot MAP_SIZE for afl struct maps

3 files changed, 48 insertions, 38 deletions

diff --git a/src/afl-fuzz-state.c b/src/afl-fuzz-state.c
index 7d068258..476782e0 100644
--- a/src/afl-fuzz-state.c
+++ b/src/afl-fuzz-state.c
@@ -75,12 +75,14 @@ list_t afl_states = {.element_prealloc_count = 0};
 
 /* Initializes an afl_state_t. */
 
-void afl_state_init(afl_state_t *afl) {
+void afl_state_init(afl_state_t *afl, uint32_t map_size) {
 
   /* thanks to this memset, growing vars like out_buf
   and out_size are NULL/0 by default. */
   memset(afl, 0, sizeof(afl_state_t));
 
+  if (!map_size) map_size = MAP_SIZE;
+
   afl->w_init = 0.9;
   afl->w_end = 0.3;
   afl->g_max = 5000;
@@ -97,13 +99,17 @@ void afl_state_init(afl_state_t *afl) {
   afl->cpu_aff = -1;                    /* Selected CPU core                */
 #endif                                                     /* HAVE_AFFINITY */
 
-  afl->fsrv.use_stdin = 1;
-
-  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->virgin_bits = ck_alloc(map_size);
+  afl->virgin_tmout = ck_alloc(map_size);
+  afl->virgin_crash = ck_alloc(map_size);
+  afl->var_bytes = ck_alloc(map_size);
+  afl->top_rated = ck_alloc(map_size);
+  afl->clean_trace = ck_alloc(map_size);
+  afl->clean_trace_custom = ck_alloc(map_size);
+  afl->first_trace = ck_alloc(map_size);
 
+  afl->fsrv.use_stdin = 1;
+  afl->fsrv.map_size = map_size;
   afl->fsrv.function_opt = (u8 *)afl;
   afl->fsrv.function_ptr = &maybe_add_auto;
 
@@ -328,24 +334,6 @@ 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)) {
@@ -386,12 +374,21 @@ void afl_state_deinit(afl_state_t *afl) {
   if (afl->pass_stats) ck_free(afl->pass_stats);
   if (afl->orig_cmp_map) ck_free(afl->orig_cmp_map);
 
-  free(afl->out_buf);
-  free(afl->out_scratch_buf);
-  free(afl->eff_buf);
-  free(afl->in_buf);
-  free(afl->in_scratch_buf);
-  free(afl->ex_buf);
+  if (afl->out_buf) free(afl->out_buf);
+  if (afl->out_scratch_buf) free(afl->out_scratch_buf);
+  if (afl->eff_buf) free(afl->eff_buf);
+  if (afl->in_buf) free(afl->in_buf);
+  if (afl->in_scratch_buf) free(afl->in_scratch_buf);
+  if (afl->ex_buf) free(afl->ex_buf);
+
+  ck_free(afl->virgin_bits);
+  ck_free(afl->virgin_tmout);
+  ck_free(afl->virgin_crash);
+  ck_free(afl->var_bytes);
+  ck_free(afl->top_rated);
+  ck_free(afl->clean_trace);
+  ck_free(afl->clean_trace_custom);
+  ck_free(afl->first_trace);
 
   list_remove(&afl_states, afl);
 
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index 2a1387a9..93e83eed 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -233,8 +233,8 @@ int main(int argc, char **argv_orig, char **envp) {
 
   s32    opt;
   u64    prev_queued = 0;
-  u32    sync_interval_cnt = 0, seek_to, show_help = 0;
-  u8 *   extras_dir = 0;
+  u32    sync_interval_cnt = 0, seek_to, show_help = 0, map_size = MAP_SIZE;
+  u8 *   extras_dir = 0, *ptr;
   u8     mem_limit_given = 0, exit_1 = 0;
   char **use_argv;
 
@@ -246,10 +246,23 @@ int main(int argc, char **argv_orig, char **envp) {
   afl_state_t *afl = calloc(1, sizeof(afl_state_t));
   if (!afl) { FATAL("Could not create afl state"); }
 
-  afl_state_init(afl);
+  if (get_afl_env("AFL_DEBUG")) afl->debug = 1;
+  if ((ptr = get_afl_env("AFL_MAP_SIZE")) ||
+      (ptr = get_afl_env("AFL_MAPSIZE"))) {
+
+    map_size = atoi(ptr);
+    if (map_size < 8 || map_size > (1 << 29))
+      FATAL(
+          "the specified AFL_MAP_SIZE size is illegal and must be between 2^3 "
+          "and 2^30: %u\n",
+          map_size);
+    if (map_size % 8) map_size = (((map_size >> 3) + 1) << 3);
+
+  }
+
+  afl_state_init(afl, map_size);
   afl_fsrv_init(&afl->fsrv);
 
-  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;
diff --git a/src/afl-showmap.c b/src/afl-showmap.c
index a11c128a..0bcb71ed 100644
--- a/src/afl-showmap.c
+++ b/src/afl-showmap.c
@@ -515,7 +515,7 @@ static void usage(u8 *argv0) {
       "For additional help, consult %s/README.md.\n\n"
 
       "Environment variables used:\n"
-      "LD_BIND_LAZY: do not set LD_BIND_NOW env var for target\n",
+      "LD_BIND_LAZY: do not set LD_BIND_NOW env var for target\n"
       "AFL_CMIN_CRASHES_ONLY: (cmin_mode) only write tuples for crashing "
       "inputs\n"
       "AFL_CMIN_ALLOW_ANY: (cmin_mode) write tuples for crashing inputs also\n"
@@ -524,8 +524,8 @@ static void usage(u8 *argv0) {
       "size\n"
       "              the target was compiled for\n"
       "AFL_PRELOAD: LD_PRELOAD / DYLD_INSERT_LIBRARIES settings for target\n"
-      "AFL_QUIET: do not print extra informational output" argv0,
-      MEM_LIMIT, doc_path);
+      "AFL_QUIET: do not print extra informational output",
+      argv0, MEM_LIMIT, doc_path);
 
   exit(1);