From 83f32c5248c8a8a1e69ca2f6f392c27c1736eef1 Mon Sep 17 00:00:00 2001
From: vanhauser-thc <vh@thc.org>
Date: Wed, 8 Jun 2022 10:56:11 +0200
Subject: honor AFL_MAP_SIZE well outside of afl++

---
 instrumentation/afl-compiler-rt.o.c | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

(limited to 'instrumentation/afl-compiler-rt.o.c')

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,
-- 
cgit 1.4.1


From 35d49c7c5c398c6de5d3091fdda298e45726ae1b Mon Sep 17 00:00:00 2001
From: vanhauser-thc <vh@thc.org>
Date: Wed, 8 Jun 2022 12:46:08 +0200
Subject: fix

---
 instrumentation/afl-compiler-rt.o.c | 36 ++++++++++++++++++++++++++----------
 1 file changed, 26 insertions(+), 10 deletions(-)

(limited to 'instrumentation/afl-compiler-rt.o.c')

diff --git a/instrumentation/afl-compiler-rt.o.c b/instrumentation/afl-compiler-rt.o.c
index b94e3dc9..f3a16e95 100644
--- a/instrumentation/afl-compiler-rt.o.c
+++ b/instrumentation/afl-compiler-rt.o.c
@@ -327,7 +327,7 @@ static void __afl_map_shm(void) {
 
   }
 
-  if (!id_str) {
+  if (!id_str && __afl_area_ptr_dummy == __afl_area_initial) {
 
     u32 val = 0;
     u8 *ptr;
@@ -337,7 +337,6 @@ static void __afl_map_shm(void) {
     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) {
 
@@ -348,6 +347,17 @@ static void __afl_map_shm(void) {
 
       }
 
+    } 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);
+
     }
 
   }
@@ -490,20 +500,26 @@ static void __afl_map_shm(void) {
 
     }
 
-  } else if (_is_sancov && __afl_area_ptr != __afl_area_initial &&
+  } else if (__afl_final_loc > __afl_map_size) {
 
-             __afl_area_ptr != __afl_area_ptr_dummy) {
+    if (__afl_area_initial != __afl_area_ptr_dummy) {
 
-    free(__afl_area_ptr);
-    __afl_area_ptr = NULL;
+      free(__afl_area_ptr_dummy);
 
-    if (__afl_final_loc > MAP_INITIAL_SIZE) {
+    }
 
-      __afl_area_ptr = (u8 *)malloc(__afl_final_loc);
+    __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) {
 
-    if (!__afl_area_ptr) { __afl_area_ptr = __afl_area_ptr_dummy; }
+      fprintf(stderr,
+              "Error: AFL++ could not aquire %u bytes of memory, exiting!\n",
+              __afl_final_loc);
+      exit(-1);
+
+    }
 
   }
 
-- 
cgit 1.4.1