3 files changed, 58 insertions, 43 deletions
diff --git a/unicorn_mode/patches/afl-unicorn-common.h b/unicorn_mode/patches/afl-unicorn-common.h
new file mode 100644
index 00000000..9a1b2a6c
--- /dev/null
+++ b/unicorn_mode/patches/afl-unicorn-common.h
@@ -0,0 +1,50 @@
+/*
+   american fuzzy lop - high-performance binary-only instrumentation
+   -----------------------------------------------------------------
+
+   Written by Andrew Griffiths <agriffiths@google.com> and
+              Michal Zalewski <lcamtuf@google.com>
+
+   TCG instrumentation and block chaining support by Andrea Biondo
+                                      <andrea.biondo965@gmail.com>
+   Adapted for afl-unicorn by Dominik Maier <mail@dmnk.co>
+
+   Idea & design very much by Andrew Griffiths.
+
+   Copyright 2015, 2016 Google Inc. All rights reserved.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at:
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+   This code is a shim patched into the separately-distributed source
+   code of Unicorn 1.0.1. It leverages the built-in QEMU tracing functionality
+   to implement AFL-style instrumentation and to take care of the remaining
+   parts of the AFL fork server logic.
+
+   The resulting QEMU binary is essentially a standalone instrumentation
+   tool; for an example of how to leverage it for other purposes, you can
+   have a look at afl-showmap.c.
+
+ */
+
+#include "../../config.h"
+
+/* NeverZero */ 
+
+#if (defined(__x86_64__) || defined(__i386__)) && defined(AFL_QEMU_NOT_ZERO)
+#  define INC_AFL_AREA(loc) \
+    asm volatile ( \
+      "incb (%0, %1, 1)\n" \
+      "adcb $0, (%0, %1, 1)\n" \
+      : /* no out */ \
+      : "r" (afl_area_ptr), "r" (loc) \
+      : "memory", "eax" \
+    )
+#else
+#  define INC_AFL_AREA(loc) \
+  afl_area_ptr[loc]++
+#endif
+
diff --git a/unicorn_mode/patches/afl-unicorn-cpu-inl.h b/unicorn_mode/patches/afl-unicorn-cpu-inl.h
index ff194696..90937a17 100644
--- a/unicorn_mode/patches/afl-unicorn-cpu-inl.h
+++ b/unicorn_mode/patches/afl-unicorn-cpu-inl.h
@@ -33,7 +33,7 @@
 #include <sys/shm.h>
 #include <sys/types.h>
 #include <sys/wait.h>
-#include "../../config.h"
+#include "afl-unicorn-common.h"
 
 /***************************
  * VARIOUS AUXILIARY STUFF *
@@ -218,17 +218,11 @@ static inline void afl_maybe_log(struct uc_struct* uc, unsigned long cur_loc) {
 
   static __thread unsigned long prev_loc;
 
-  // DEBUG
-  //printf("IN AFL_MAYBE_LOG 0x%lx\n", cur_loc);
+  u8* afl_area_ptr = uc->afl_area_ptr;
 
-  // MODIFIED FOR UNICORN MODE -> We want to log all addresses,
-  // so the checks for 'start < addr < end' are removed
-  if(!uc->afl_area_ptr)
+  if(!afl_area_ptr)
     return;
 
-  // DEBUG
-  //printf("afl_area_ptr = %p\n", afl_area_ptr);
-
   /* Looks like QEMU always maps to fixed locations, so ASAN is not a
      concern. Phew. But instruction addresses may be aligned. Let's mangle
      the value to get something quasi-uniform. */
@@ -239,27 +233,11 @@ static inline void afl_maybe_log(struct uc_struct* uc, unsigned long cur_loc) {
   /* Implement probabilistic instrumentation by looking at scrambled block
      address. This keeps the instrumented locations stable across runs. */
 
-  // DEBUG
-  //printf("afl_inst_rms = 0x%lx\n", afl_inst_rms);
-
   if (cur_loc >= uc->afl_inst_rms) return;
 
-  // DEBUG
-  //printf("cur_loc = 0x%lx\n", cur_loc);  
-
   register uintptr_t afl_idx = cur_loc ^ prev_loc;
 
-#if (defined(__x86_64__) || defined(__i386__)) && defined(AFL_QEMU_NOT_ZERO)
-  asm volatile (
-    "incb (%0, %1, 1)\n"
-    "adcb $0, (%0, %1, 1)\n"
-    : /* no out */
-    : "r" (uc->afl_area_ptr), "r" (afl_idx)
-    : "memory", "eax"
-  );
-#else
-  uc->afl_area_ptr[afl_idx]++;
-#endif
+  INC_AFL_AREA(afl_idx);
 
   prev_loc = cur_loc >> 1;
 
diff --git a/unicorn_mode/patches/afl-unicorn-tcg-runtime-inl.h b/unicorn_mode/patches/afl-unicorn-tcg-runtime-inl.h
index 52cc1afb..0019bbfa 100644
--- a/unicorn_mode/patches/afl-unicorn-tcg-runtime-inl.h
+++ b/unicorn_mode/patches/afl-unicorn-tcg-runtime-inl.h
@@ -31,25 +31,12 @@
  */
 
 #include "uc_priv.h"
-
-#if (defined(__x86_64__) || defined(__i386__)) && defined(AFL_QEMU_NOT_ZERO)
-#  define INC_AFL_AREA(loc) \
-    asm volatile ( \
-      "incb (%0, %1, 1)\n" \
-      "adcb $0, (%0, %1, 1)\n" \
-      : /* no out */ \
-      : "r" (uc->afl_area_ptr), "r" (loc) \
-      : "memory", "eax" \
-    )
-#else
-#  define INC_AFL_AREA(loc) \
-  uc->afl_area_ptr[loc]++
-#endif
+#include "afl-unicorn-common.h"
 
 void HELPER(afl_compcov_log_16)(void* uc_ptr, uint64_t cur_loc, uint64_t arg1,
                                 uint64_t arg2) {
 
-  struct uc_struct* uc = uc_ptr;
+  u8* afl_area_ptr = ((struct uc_struct*)uc_ptr)->afl_area_ptr;
 
   if ((arg1 & 0xff) == (arg2 & 0xff)) {
     INC_AFL_AREA(cur_loc);
@@ -59,7 +46,7 @@ void HELPER(afl_compcov_log_16)(void* uc_ptr, uint64_t cur_loc, uint64_t arg1,
 void HELPER(afl_compcov_log_32)(void* uc_ptr, uint64_t cur_loc, uint64_t arg1,
                                 uint64_t arg2) {
 
-  struct uc_struct* uc = uc_ptr;
+  u8* afl_area_ptr = ((struct uc_struct*)uc_ptr)->afl_area_ptr;
 
   if ((arg1 & 0xff) == (arg2 & 0xff)) {
     INC_AFL_AREA(cur_loc);
@@ -75,7 +62,7 @@ void HELPER(afl_compcov_log_32)(void* uc_ptr, uint64_t cur_loc, uint64_t arg1,
 void HELPER(afl_compcov_log_64)(void* uc_ptr, uint64_t cur_loc, uint64_t arg1,
                                 uint64_t arg2) {
 
-  struct uc_struct* uc = uc_ptr;
+  u8* afl_area_ptr = ((struct uc_struct*)uc_ptr)->afl_area_ptr;
 
   if ((arg1 & 0xff) == (arg2 & 0xff)) {
     INC_AFL_AREA(cur_loc);