Merge pull request #1438 from AFLplusplus/dev 4.01c

push to stable

1 files changed, 51 insertions, 8 deletions

diff --git a/instrumentation/afl-compiler-rt.o.c b/instrumentation/afl-compiler-rt.o.c
index db7ac7b0..f3a16e95 100644
--- a/instrumentation/afl-compiler-rt.o.c
+++ b/instrumentation/afl-compiler-rt.o.c
@@ -327,6 +327,41 @@ static void __afl_map_shm(void) {
 
   }
 
+  if (!id_str && __afl_area_ptr_dummy == __afl_area_initial) {
+
+    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_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);
+
+      }
+
+    } else {
+
+      __afl_map_size = MAP_INITIAL_SIZE;
+
+    }
+
+    if (__afl_debug) {
+
+      fprintf(stderr, "DEBUG: (0) init map size is %u to %p\n", __afl_map_size,
+              __afl_area_ptr_dummy);
+
+    }
+
+  }
+
   /* 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,18 +500,26 @@ static void __afl_map_shm(void) {
 
     }
 
-  } else if (_is_sancov && __afl_area_ptr != __afl_area_initial) {
-
-    free(__afl_area_ptr);
-    __afl_area_ptr = NULL;
+  } else if (__afl_final_loc > __afl_map_size) {
 
-    if (__afl_final_loc > MAP_INITIAL_SIZE) {
+    if (__afl_area_initial != __afl_area_ptr_dummy) {
 
-      __afl_area_ptr = (u8 *)malloc(__afl_final_loc);
+      free(__afl_area_ptr_dummy);
 
     }
 
-    if (!__afl_area_ptr) { __afl_area_ptr = __afl_area_ptr_dummy; }
+    __afl_area_ptr_dummy = (u8 *)malloc(__afl_final_loc);
+    __afl_area_ptr = __afl_area_ptr_dummy;
+    __afl_map_size = __afl_final_loc;
+
+    if (!__afl_area_ptr_dummy) {
+
+      fprintf(stderr,
+              "Error: AFL++ could not aquire %u bytes of memory, exiting!\n",
+              __afl_final_loc);
+      exit(-1);
+
+    }
 
   }
 
@@ -487,7 +530,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,