add extended forkserver feature to gcc_plugin and qemu_mode

2 files changed, 18 insertions, 0 deletions

diff --git a/gcc_plugin/afl-gcc-rt.o.c b/gcc_plugin/afl-gcc-rt.o.c
index b157b50f..0a2246e7 100644
--- a/gcc_plugin/afl-gcc-rt.o.c
+++ b/gcc_plugin/afl-gcc-rt.o.c
@@ -139,6 +139,7 @@ static void __afl_map_shm(void) {
 static void __afl_start_forkserver(void) {
 
   u8  tmp[4] = {0, 0, 0, 0};
+  u32 map_size = MAP_SIZE;
   s32 child_pid;
 
   u8 child_stopped = 0;
@@ -148,6 +149,13 @@ static void __afl_start_forkserver(void) {
   /* 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. */
 
+  if (MAP_SIZE <= 0x800000) {
+
+    map_size = (FS_OPT_ENABLED | FS_OPT_MAPSIZE | FS_OPT_SET_MAPSIZE(MAP_SIZE));
+    memcpy(tmp, &map_size, 4);
+
+  }
+
   if (write(FORKSRV_FD + 1, tmp, 4) != 4) return;
 
   while (1) {
diff --git a/qemu_mode/patches/afl-qemu-cpu-inl.h b/qemu_mode/patches/afl-qemu-cpu-inl.h
index d73566fc..ee3ec44e 100644
--- a/qemu_mode/patches/afl-qemu-cpu-inl.h
+++ b/qemu_mode/patches/afl-qemu-cpu-inl.h
@@ -293,6 +293,7 @@ static void print_mappings(void) {
 
 void afl_forkserver(CPUState *cpu) {
 
+  u32                  map_size = 0;
   static unsigned char tmp[4] = {0, 0, 0, 0};
 
   if (forkserver_installed == 1) return;
@@ -306,6 +307,15 @@ void afl_forkserver(CPUState *cpu) {
   int   t_fd[2];
   u8    child_stopped = 0;
 
+  // if in the future qemu has non-collding coverage then switch MAP_SIZE
+  // with the max ID value
+  if (MAP_SIZE <= 0x800000) {
+
+    map_size = (FS_OPT_ENABLED | FS_OPT_MAPSIZE | FS_OPT_SET_MAPSIZE(MAP_SIZE));
+    memcpy(tmp, &map_size, 4);
+
+  }
+
   /* Tell the parent that we're alive. If the parent doesn't want
      to talk, assume that we're not running in forkserver mode. */