sync afl-gcc-rt.o.c from its more current llvm cousin

1 files changed, 82 insertions, 8 deletions

diff --git a/gcc_plugin/afl-gcc-rt.o.c b/gcc_plugin/afl-gcc-rt.o.c
index 6cf9f167..1dfca9e3 100644
--- a/gcc_plugin/afl-gcc-rt.o.c
+++ b/gcc_plugin/afl-gcc-rt.o.c
@@ -21,12 +21,16 @@
 
 */
 
+#ifdef __ANDROID__
+#include "android-ashmem.h"
+#endif
 #include "../config.h"
 #include "../types.h"
 
 #include <stdlib.h>
 #include <signal.h>
 #include <unistd.h>
+#include <string.h>
 #include <assert.h>
 
 #include <sys/mman.h>
@@ -34,17 +38,21 @@
 #include <sys/wait.h>
 #include <sys/types.h>
 
+#include <sys/mman.h>
+#include <fcntl.h>
+
 /* Globals needed by the injected instrumentation. The __afl_area_initial region
    is used for instrumentation output before __afl_map_shm() has a chance to
    run. It will end up as .comm, so it shouldn't be too wasteful. */
 
 u8  __afl_area_initial[MAP_SIZE];
 u8 *__afl_area_ptr = __afl_area_initial;
-u32 __afl_prev_loc;
-
-/* Running in persistent mode? */
 
-static u8 is_persistent;
+#ifdef __ANDROID__
+u32 __afl_prev_loc;
+#else
+__thread u32 __afl_prev_loc;
+#endif
 
 /* Trace a basic block with some ID */
 void __afl_trace(u32 x) {
@@ -59,6 +67,10 @@ void __afl_trace(u32 x) {
 
 }
 
+/* Running in persistent mode? */
+
+static u8 is_persistent;
+
 /* SHM setup. */
 
 static void __afl_map_shm(void) {
@@ -71,9 +83,38 @@ static void __afl_map_shm(void) {
 
   if (id_str) {
 
+#ifdef USEMMAP
+    const char*    shm_file_path = id_str;
+    int            shm_fd = -1;
+    unsigned char* shm_base = NULL;
+
+    /* create the shared memory segment as if it was a file */
+    shm_fd = shm_open(shm_file_path, O_RDWR, 0600);
+    if (shm_fd == -1) {
+
+      printf("shm_open() failed\n");
+      exit(1);
+
+    }
+
+    /* map the shared memory segment to the address space of the process */
+    shm_base = mmap(0, MAP_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0);
+    if (shm_base == MAP_FAILED) {
+
+      close(shm_fd);
+      shm_fd = -1;
+
+      printf("mmap() failed\n");
+      exit(2);
+
+    }
+
+    __afl_area_ptr = shm_base;
+#else
     u32 shm_id = atoi(id_str);
 
     __afl_area_ptr = shmat(shm_id, NULL, 0);
+#endif
 
     /* Whooooops. */
 
@@ -97,6 +138,8 @@ static void __afl_start_forkserver(void) {
 
   u8 child_stopped = 0;
 
+  void (*old_sigchld_handler)(int) = signal(SIGCHLD, SIG_DFL);
+
   /* Phone home and tell the parent that we're OK. If parent isn't there,
      assume we're not running in forkserver mode and just execute program. */
 
@@ -133,6 +176,8 @@ static void __afl_start_forkserver(void) {
 
       if (!child_pid) {
 
+        signal(SIGCHLD, old_sigchld_handler);
+
         close(FORKSRV_FD);
         close(FORKSRV_FD + 1);
         return;
@@ -178,18 +223,47 @@ int __afl_persistent_loop(unsigned int max_cnt) {
 
   if (first_pass) {
 
+    /* Make sure that every iteration of __AFL_LOOP() starts with a clean slate.
+       On subsequent calls, the parent will take care of that, but on the first
+       iteration, it's our job to erase any trace of whatever happened
+       before the loop. */
+
+    if (is_persistent) {
+
+      memset(__afl_area_ptr, 0, MAP_SIZE);
+      __afl_area_ptr[0] = 1;
+      __afl_prev_loc = 0;
+
+    }
+
     cycle_cnt = max_cnt;
     first_pass = 0;
     return 1;
 
   }
 
-  if (is_persistent && --cycle_cnt) {
+  if (is_persistent) {
 
-    raise(SIGSTOP);
-    return 1;
+    if (--cycle_cnt) {
+
+      raise(SIGSTOP);
 
-  } else
+      __afl_area_ptr[0] = 1;
+      __afl_prev_loc = 0;
+
+      return 1;
+
+    } else {
+
+      /* When exiting __AFL_LOOP(), make sure that the subsequent code that
+         follows the loop is not traced. We do that by pivoting back to the
+         dummy output region. */
+
+      __afl_area_ptr = __afl_area_initial;
+
+    }
+
+  }
 
     return 0;