refactoring getting the map size

5 files changed, 28 insertions, 44 deletions

diff --git a/src/afl-analyze.c b/src/afl-analyze.c
index 8a84b781..b2c0f841 100644
--- a/src/afl-analyze.c
+++ b/src/afl-analyze.c
@@ -814,7 +814,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, *ptr;
+  char **use_argv;
 
   doc_path = access(DOC_PATH, F_OK) ? "docs" : DOC_PATH;
 
@@ -934,14 +934,7 @@ 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);
-
-  }
+  map_size = get_map_size();
 
   use_hex_offsets = !!get_afl_env("AFL_ANALYZE_HEX");
 
diff --git a/src/afl-common.c b/src/afl-common.c
index 38c19234..2cbf1059 100644
--- a/src/afl-common.c
+++ b/src/afl-common.c
@@ -898,3 +898,21 @@ u32 read_timed(s32 fd, void *buf, size_t len, u32 timeout_ms,
 
 }
 
+u32 get_map_size() {
+
+  uint32_t map_size = MAP_SIZE;
+  char *   ptr;
+
+  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);
+
+  }
+
+  return map_size;
+
+}
+
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index 03726eb0..83e25994 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -234,7 +234,7 @@ 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, map_size = MAP_SIZE;
-  u8 *   extras_dir = 0, *ptr;
+  u8 *   extras_dir = 0;
   u8     mem_limit_given = 0, exit_1 = 0;
   char **use_argv;
 
@@ -247,19 +247,8 @@ int main(int argc, char **argv_orig, char **envp) {
   if (!afl) { FATAL("Could not create afl state"); }
 
   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);
-
-  }
 
+  map_size = get_map_size();
   afl_state_init(afl, map_size);
   afl_fsrv_init(&afl->fsrv);
 
diff --git a/src/afl-showmap.c b/src/afl-showmap.c
index f44b5453..a6adb695 100644
--- a/src/afl-showmap.c
+++ b/src/afl-showmap.c
@@ -536,23 +536,15 @@ 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, *ptr;
+  char **use_argv;
 
   char **argv = argv_cpy_dup(argc, argv_orig);
 
   afl_forkserver_t  fsrv_var = {0};
   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;
-
-  }
+  map_size = get_map_size();
+  fsrv->map_size = map_size;
 
   doc_path = access(DOC_PATH, F_OK) ? "docs" : DOC_PATH;
 
diff --git a/src/afl-tmin.c b/src/afl-tmin.c
index ad7d70c7..dab2a417 100644
--- a/src/afl-tmin.c
+++ b/src/afl-tmin.c
@@ -758,23 +758,15 @@ 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, *ptr;
+  char **use_argv;
 
   char **argv = argv_cpy_dup(argc, argv_orig);
 
   afl_forkserver_t  fsrv_var = {0};
   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;
-
-  }
+  map_size = get_map_size();
+  fsrv->map_size = map_size;
 
   doc_path = access(DOC_PATH, F_OK) ? "docs" : DOC_PATH;