even support dlopen instrumented libs after the forkserver

2 files changed, 37 insertions, 4 deletions

diff --git a/docs/Changelog.md b/docs/Changelog.md
index 6fe3517a..b1c991ff 100644
--- a/docs/Changelog.md
+++ b/docs/Changelog.md
@@ -15,8 +15,9 @@ sending a mail to <afl-users+subscribe@googlegroups.com>.
     - fixed a crash that can occur with ASAN + CMPLOG together plus
       better support for unicode (thanks to @stbergmann for reporting!)
     - handle erroneous setups in which multiple afl-compiler-rt are
-      compiled into the target. This now also supports dlopen instrumented
-      libs loaded before the forkserver.
+      compiled into the target. This now also supports dlopen
+      instrumented libs loaded before the forkserver and even after the
+      forkserver is started (then with collisions though)
     - Renamed CTX to CALLER, added correct/real CTX implemenation to CLASSIC
   - qemu_mode
     - added AFL_QEMU_EXCLUDE_RANGES env by @realmadsci, thanks!
diff --git a/instrumentation/afl-compiler-rt.o.c b/instrumentation/afl-compiler-rt.o.c
index e3aa787f..87bf6486 100644
--- a/instrumentation/afl-compiler-rt.o.c
+++ b/instrumentation/afl-compiler-rt.o.c
@@ -1263,8 +1263,10 @@ void __sanitizer_cov_trace_pc_guard_init(uint32_t *start, uint32_t *stop) {
   if (__afl_debug) {
 
     fprintf(stderr,
-            "Running __sanitizer_cov_trace_pc_guard_init: %p-%p (%lu edges)\n",
-            start, stop, (unsigned long)(stop - start));
+            "Running __sanitizer_cov_trace_pc_guard_init: %p-%p (%lu edges) "
+            "after_fs=%u\n",
+            start, stop, (unsigned long)(stop - start),
+            __afl_already_initialized_forkserver);
 
   }
 
@@ -1280,6 +1282,36 @@ void __sanitizer_cov_trace_pc_guard_init(uint32_t *start, uint32_t *stop) {
 
   }
 
+  /* instrumented code is loaded *after* our forkserver is up. this is a
+     problem. We cannot prevent collisions then :( */
+  if (__afl_already_initialized_forkserver &&
+      __afl_final_loc + 1 + stop - start > __afl_map_size) {
+
+    if (__afl_debug)
+      fprintf(stderr, "Warning: new instrumneted code after the forkserver!\n");
+    __afl_final_loc = 2;
+
+    if (1 + stop - start > __afl_map_size) {
+
+      *(start++) = ++__afl_final_loc;
+
+      while (start < stop) {
+
+        if (R(100) < inst_ratio)
+          *start = ++__afl_final_loc % __afl_map_size;
+        else
+          *start = 0;
+
+        start++;
+
+      }
+
+      return;
+
+    }
+
+  }
+
   /* Make sure that the first element in the range is always set - we use that
      to avoid duplicate calls (which can happen as an artifact of the underlying
      implementation in LLVM). */