honor AFL_MAP_SIZE well outside of afl++

4 files changed, 41 insertions, 10 deletions

diff --git a/instrumentation/afl-compiler-rt.o.c b/instrumentation/afl-compiler-rt.o.c
index db7ac7b0..b94e3dc9 100644
--- a/instrumentation/afl-compiler-rt.o.c
+++ b/instrumentation/afl-compiler-rt.o.c
@@ -327,6 +327,31 @@ static void __afl_map_shm(void) {
 
   }
 
+  if (!id_str) {
+
+    u32 val = 0;
+    u8 *ptr;
+
+    if ((ptr = getenv("AFL_MAP_SIZE")) != NULL) val = atoi(ptr);
+
+    if (val > MAP_INITIAL_SIZE) {
+
+      __afl_map_size = val;
+      __afl_final_loc = val;
+      __afl_area_ptr_dummy = malloc(__afl_map_size);
+      if (!__afl_area_ptr_dummy) {
+
+        fprintf(stderr,
+                "Error: AFL++ could not aquire %u bytes of memory, exiting!\n",
+                __afl_map_size);
+        exit(-1);
+
+      }
+
+    }
+
+  }
+
   /* If we're running under AFL, attach to the appropriate region, replacing the
      early-stage __afl_area_initial region that is needed to allow some really
      hacky .init code to work correctly in projects such as OpenSSL. */
@@ -465,7 +490,9 @@ static void __afl_map_shm(void) {
 
     }
 
-  } else if (_is_sancov && __afl_area_ptr != __afl_area_initial) {
+  } else if (_is_sancov && __afl_area_ptr != __afl_area_initial &&
+
+             __afl_area_ptr != __afl_area_ptr_dummy) {
 
     free(__afl_area_ptr);
     __afl_area_ptr = NULL;
@@ -487,7 +514,7 @@ static void __afl_map_shm(void) {
     fprintf(stderr,
             "DEBUG: (2) id_str %s, __afl_area_ptr %p, __afl_area_initial %p, "
             "__afl_area_ptr_dummy %p, __afl_map_addr 0x%llx, MAP_SIZE "
-            "%u, __afl_final_loc %u, __afl_map_size %u,"
+            "%u, __afl_final_loc %u, __afl_map_size %u, "
             "max_size_forkserver %u/0x%x\n",
             id_str == NULL ? "<null>" : id_str, __afl_area_ptr,
             __afl_area_initial, __afl_area_ptr_dummy, __afl_map_addr, MAP_SIZE,
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index 09e773f0..5703a66a 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -130,11 +130,7 @@ write_to_testcase(afl_state_t *afl, void **mem, u32 len, u32 fix) {
 
     }
 
-    if (new_mem != *mem) {
-
-      *mem = new_mem;
-
-    }
+    if (new_mem != *mem) { *mem = new_mem; }
 
     /* everything as planned. use the potentially new data. */
     afl_fsrv_write_to_testcase(&afl->fsrv, *mem, new_size);
diff --git a/src/afl-fuzz-state.c b/src/afl-fuzz-state.c
index cbe32c75..8334af75 100644
--- a/src/afl-fuzz-state.c
+++ b/src/afl-fuzz-state.c
@@ -526,11 +526,18 @@ void read_afl_environment(afl_state_t *afl, char **envp) {
 
             int time = atoi((u8 *)get_afl_env(afl_environment_variables[i]));
             if (time > 0) {
-                afl->sync_time = time * (60 * 1000LL);
+
+              afl->sync_time = time * (60 * 1000LL);
+
             } else {
-              WARNF("incorrect value for AFL_SYNC_TIME environment variable, "
-                    "used default value %lld instead.", afl->sync_time / 60 / 1000);
+
+              WARNF(
+                  "incorrect value for AFL_SYNC_TIME environment variable, "
+                  "used default value %lld instead.",
+                  afl->sync_time / 60 / 1000);
+
             }
+
           }
 
         } else {
diff --git a/utils/libdislocator/libdislocator.so.c b/utils/libdislocator/libdislocator.so.c
index fecf3bc6..c821a8f7 100644
--- a/utils/libdislocator/libdislocator.so.c
+++ b/utils/libdislocator/libdislocator.so.c
@@ -526,6 +526,7 @@ size_t malloc_good_size(size_t len) {
   return (len & ~(ALLOC_ALIGN_SIZE - 1)) + ALLOC_ALIGN_SIZE;
 
 }
+
 #endif
 
 __attribute__((constructor)) void __dislocator_init(void) {