1 files changed, 67 insertions, 3 deletions

diff --git a/qemu_mode/patches/afl-qemu-cpu-inl.h b/qemu_mode/patches/afl-qemu-cpu-inl.h
index ac847371..5e155c74 100644
--- a/qemu_mode/patches/afl-qemu-cpu-inl.h
+++ b/qemu_mode/patches/afl-qemu-cpu-inl.h
@@ -32,10 +32,11 @@
  */
 
 #include <sys/shm.h>
-#include "../../config.h"
 #include "afl-qemu-common.h"
 
-#define PERSISTENT_DEFAULT_MAX_CNT 1000
+#ifndef AFL_QEMU_STATIC_BUILD
+#include <dlfcn.h>
+#endif
 
 /***************************
  * VARIOUS AUXILIARY STUFF *
@@ -81,6 +82,9 @@ u8 afl_compcov_level;
 
 __thread abi_ulong afl_prev_loc;
 
+struct cmp_map *__afl_cmp_map;
+__thread u32    __afl_cmp_counter;
+
 /* Set in the child process in forkserver mode: */
 
 static int forkserver_installed = 0;
@@ -95,6 +99,8 @@ unsigned char persistent_save_gpr;
 target_ulong  persistent_saved_gpr[AFL_REGS_NUM];
 int           persisent_retaddr_offset;
 
+afl_persistent_hook_fn afl_persistent_hook_ptr;
+
 /* Instrumentation ratio: */
 
 unsigned int afl_inst_rms = MAP_SIZE;         /* Exported for afl_gen_trace */
@@ -182,6 +188,22 @@ static void afl_setup(void) {
 
   }
 
+  if (getenv("___AFL_EINS_ZWEI_POLIZEI___")) {  // CmpLog forkserver
+
+    id_str = getenv(CMPLOG_SHM_ENV_VAR);
+
+    if (id_str) {
+
+      u32 shm_id = atoi(id_str);
+
+      __afl_cmp_map = shmat(shm_id, NULL, 0);
+
+      if (__afl_cmp_map == (void *)-1) exit(1);
+
+    }
+
+  }
+
   if (getenv("AFL_INST_LIBS")) {
 
     afl_start_code = 0;
@@ -224,6 +246,43 @@ static void afl_setup(void) {
 
   if (getenv("AFL_QEMU_PERSISTENT_GPR")) persistent_save_gpr = 1;
 
+  if (getenv("AFL_QEMU_PERSISTENT_HOOK")) {
+
+#ifdef AFL_QEMU_STATIC_BUILD
+
+    fprintf(stderr,
+            "[AFL] ERROR: you cannot use AFL_QEMU_PERSISTENT_HOOK when "
+            "afl-qemu-trace is static\n");
+    exit(1);
+
+#else
+
+    persistent_save_gpr = 1;
+
+    void *plib = dlopen(getenv("AFL_QEMU_PERSISTENT_HOOK"), RTLD_NOW);
+    if (!plib) {
+
+      fprintf(stderr, "[AFL] ERROR: invalid AFL_QEMU_PERSISTENT_HOOK=%s\n",
+              getenv("AFL_QEMU_PERSISTENT_HOOK"));
+      exit(1);
+
+    }
+
+    afl_persistent_hook_ptr = dlsym(plib, "afl_persistent_hook");
+    if (!afl_persistent_hook_ptr) {
+
+      fprintf(stderr,
+              "[AFL] ERROR: failed to find the function "
+              "\"afl_persistent_hook\" in %s\n",
+              getenv("AFL_QEMU_PERSISTENT_HOOK"));
+      exit(1);
+
+    }
+
+#endif
+
+  }
+
   if (getenv("AFL_QEMU_PERSISTENT_RETADDR_OFFSET"))
     persisent_retaddr_offset =
         strtoll(getenv("AFL_QEMU_PERSISTENT_RETADDR_OFFSET"), NULL, 0);
@@ -352,8 +411,13 @@ static void afl_forkserver(CPUState *cpu) {
 
     if (WIFSTOPPED(status))
       child_stopped = 1;
-    else if (unlikely(first_run && is_persistent))
+    else if (unlikely(first_run && is_persistent)) {
+
+      fprintf(stderr, "[AFL] ERROR: no persistent iteration executed\n");
       exit(12);  // Persistent is wrong
+
+    }
+
     first_run = 0;
 
     if (write(FORKSRV_FD + 1, &status, 4) != 4) exit(7);