Merge remote-tracking branch 'origin/dev' into atnwalk

# Conflicts:
#	include/afl-fuzz.h
#	src/afl-fuzz-run.c

5 files changed, 224 insertions, 36 deletions

diff --git a/frida_mode/src/instrument/instrument.c b/frida_mode/src/instrument/instrument.c
index e1e4ac22..a6aac666 100644
--- a/frida_mode/src/instrument/instrument.c
+++ b/frida_mode/src/instrument/instrument.c
@@ -27,6 +27,7 @@ gboolean instrument_optimize = false;
 gboolean instrument_unique = false;
 guint64  instrument_hash_zero = 0;
 guint64  instrument_hash_seed = 0;
+gboolean instrument_suppress = false;
 
 gboolean instrument_use_fixed_seed = FALSE;
 guint64  instrument_fixed_seed = 0;
@@ -290,6 +291,7 @@ void instrument_config(void) {
       (getenv("AFL_FRIDA_INST_UNSTABLE_COVERAGE_FILE"));
   instrument_coverage_insn = (getenv("AFL_FRIDA_INST_INSN") != NULL);
   instrument_regs_filename = getenv("AFL_FRIDA_INST_REGS_FILE");
+  instrument_suppress = (getenv("AFL_FRIDA_INST_NO_SUPPRESS") == NULL);
 
   instrument_debug_config();
   instrument_coverage_config();
@@ -321,6 +323,9 @@ void instrument_init(void) {
   FOKF(cBLU "Instrumentation" cRST " - " cGRN "instructions:" cYEL " [%c]",
        instrument_coverage_insn ? 'X' : ' ');
 
+  FOKF(cBLU "Instrumentation" cRST " - " cGRN "suppression:" cYEL " [%c]",
+       instrument_suppress ? 'X' : ' ');
+
   if (instrument_tracing && instrument_optimize) {
 
     WARNF("AFL_FRIDA_INST_TRACE implies AFL_FRIDA_INST_NO_OPTIMIZE");
diff --git a/frida_mode/src/instrument/instrument_arm32.c b/frida_mode/src/instrument/instrument_arm32.c
index f2e825ee..51f78a35 100644
--- a/frida_mode/src/instrument/instrument_arm32.c
+++ b/frida_mode/src/instrument/instrument_arm32.c
@@ -1,6 +1,7 @@
 #include "frida-gumjs.h"
 
 #include "instrument.h"
+#include "stalker.h"
 #include "util.h"
 
 #if defined(__arm__)
@@ -8,8 +9,9 @@
   #define PAGE_MASK (~(GUM_ADDRESS(0xfff)))
   #define PAGE_ALIGNED(x) ((GUM_ADDRESS(x) & PAGE_MASK) == GUM_ADDRESS(x))
 
-gboolean instrument_cache_enabled = FALSE;
-gsize    instrument_cache_size = 0;
+gboolean           instrument_cache_enabled = FALSE;
+gsize              instrument_cache_size = 0;
+static GHashTable *coverage_blocks = NULL;
 
 extern __thread guint64 instrument_previous_pc;
 
@@ -22,7 +24,24 @@ typedef struct {
   // shared_mem[cur_location ^ prev_location]++;
   // prev_location = cur_location >> 1;
 
-  /* We can remove this branch when we add support for branch suppression */
+  // str     r0, [sp, #-128] ; 0xffffff80
+  // str     r1, [sp, #-132] ; 0xffffff7c
+  // ldr     r0, [pc, #-20]  ; 0xf691b29c
+  // ldrh    r1, [r0]
+  // movw    r0, #33222      ; 0x81c6
+  // eor     r0, r0, r1
+  // ldr     r1, [pc, #-40]  ; 0xf691b298
+  // add     r1, r1, r0
+  // ldrb    r0, [r1]
+  // add     r0, r0, #1
+  // add     r0, r0, r0, lsr #8
+  // strb    r0, [r1]
+  // movw    r0, #49379      ; 0xc0e3
+  // ldr     r1, [pc, #-64]  ; 0xf691b29c
+  // strh    r0, [r1]
+  // ldr     r1, [sp, #-132] ; 0xffffff7c
+  // ldr     r0, [sp, #-128] ; 0xffffff80
+
   uint32_t  b_code;                                                /* b imm */
   uint8_t  *shared_mem;
   uint64_t *prev_location;
@@ -115,6 +134,46 @@ gboolean instrument_is_coverage_optimize_supported(void) {
 
 }
 
+static void instrument_coverage_switch(GumStalkerObserver *self,
+                                       gpointer            from_address,
+                                       gpointer start_address, void *from_insn,
+                                       gpointer *target) {
+
+  UNUSED_PARAMETER(self);
+  UNUSED_PARAMETER(from_address);
+  UNUSED_PARAMETER(start_address);
+  UNUSED_PARAMETER(from_insn);
+
+  if (!g_hash_table_contains(coverage_blocks, GSIZE_TO_POINTER(*target))) {
+
+    return;
+
+  }
+
+  *target =
+      (guint8 *)*target + G_STRUCT_OFFSET(afl_log_code_asm_t, str_r0_sp_rz);
+
+}
+
+static void instrument_coverage_suppress_init(void) {
+
+  static gboolean initialized = false;
+  if (initialized) { return; }
+  initialized = true;
+
+  GumStalkerObserver          *observer = stalker_get_observer();
+  GumStalkerObserverInterface *iface = GUM_STALKER_OBSERVER_GET_IFACE(observer);
+  iface->switch_callback = instrument_coverage_switch;
+
+  coverage_blocks = g_hash_table_new(g_direct_hash, g_direct_equal);
+  if (coverage_blocks == NULL) {
+
+    FATAL("Failed to g_hash_table_new, errno: %d", errno);
+
+  }
+
+}
+
 static void patch_t3_insn(uint32_t *insn, uint16_t val) {
 
   uint32_t orig = GUINT32_FROM_LE(*insn);
@@ -135,14 +194,17 @@ void instrument_coverage_optimize(const cs_insn    *instr,
   guint64 area_offset = instrument_get_offset_hash(GUM_ADDRESS(instr->address));
   gsize   map_size_pow2;
   gsize   area_offset_ror;
-  GumAddress code_addr = 0;
 
-  // gum_arm64_writer_put_brk_imm(cw, 0x0);
-
-  code_addr = cw->pc;
+  instrument_coverage_suppress_init();
 
   block_start = GSIZE_TO_POINTER(GUM_ADDRESS(cw->code));
 
+  if (!g_hash_table_add(coverage_blocks, block_start)) {
+
+    FATAL("Failed - g_hash_table_add");
+
+  }
+
   code.code = template;
 
   g_assert(PAGE_ALIGNED(__afl_area_ptr));
@@ -211,7 +273,19 @@ void instrument_flush(GumStalkerOutput *output) {
 
 gpointer instrument_cur(GumStalkerOutput *output) {
 
-  return gum_arm_writer_cur(output->writer.arm);
+  gpointer curr = NULL;
+
+  if (output->encoding == GUM_INSTRUCTION_SPECIAL) {
+
+    curr = gum_thumb_writer_cur(output->writer.thumb);
+
+  } else {
+
+    curr = gum_arm_writer_cur(output->writer.arm);
+
+  }
+
+  return curr;
 
 }
 
diff --git a/frida_mode/src/instrument/instrument_arm64.c b/frida_mode/src/instrument/instrument_arm64.c
index 87811b38..4372861d 100644
--- a/frida_mode/src/instrument/instrument_arm64.c
+++ b/frida_mode/src/instrument/instrument_arm64.c
@@ -156,26 +156,55 @@ static gboolean instrument_is_deterministic(const cs_insn *from_insn) {
 
 }
 
+cs_insn *instrument_disassemble(gconstpointer address) {
+
+  csh      capstone;
+  cs_insn *insn = NULL;
+
+  cs_open(CS_ARCH_ARM64, GUM_DEFAULT_CS_ENDIAN, &capstone);
+  cs_option(capstone, CS_OPT_DETAIL, CS_OPT_ON);
+
+  cs_disasm(capstone, address, 16, GPOINTER_TO_SIZE(address), 1, &insn);
+
+  cs_close(&capstone);
+
+  return insn;
+
+}
+
 static void instrument_coverage_switch(GumStalkerObserver *self,
                                        gpointer            from_address,
-                                       gpointer            start_address,
-                                       const cs_insn      *from_insn,
-                                       gpointer           *target) {
+                                       gpointer start_address, void *from_insn,
+                                       gpointer *target) {
 
   UNUSED_PARAMETER(self);
   UNUSED_PARAMETER(from_address);
   UNUSED_PARAMETER(start_address);
 
-  gsize fixup_offset;
+  cs_insn *insn = NULL;
+  gboolean deterministic = FALSE;
+  gsize    fixup_offset;
 
   if (!g_hash_table_contains(coverage_blocks, GSIZE_TO_POINTER(*target)) &&
-      !g_hash_table_contains(coverage_blocks, GSIZE_TO_POINTER(*target + 4))) {
+      !g_hash_table_contains(coverage_blocks,
+                             GSIZE_TO_POINTER((guint8 *)*target + 4))) {
 
     return;
 
   }
 
-  if (instrument_is_deterministic(from_insn)) { return; }
+  insn = instrument_disassemble(from_insn);
+  deterministic = instrument_is_deterministic(insn);
+  cs_free(insn, 1);
+
+  /*
+   * If the branch is deterministic, then we should start execution at the
+   * begining of the block. From here, we will branch and skip the coverage
+   * code and jump right to the target code of the instrumented block.
+   * Otherwise, if the branch is non-deterministic, then we need to branch
+   * part way into the block to where the coverage instrumentation starts.
+   */
+  if (deterministic) { return; }
 
   /*
    * Since each block is prefixed with a restoration prologue, we need to be
@@ -208,7 +237,7 @@ static void instrument_coverage_switch(GumStalkerObserver *self,
    */
   fixup_offset = GUM_RESTORATION_PROLOG_SIZE +
                  G_STRUCT_OFFSET(afl_log_code_asm_t, restoration_prolog);
-  *target += fixup_offset;
+  *target = (guint8 *)*target + fixup_offset;
 
 }
 
@@ -284,7 +313,7 @@ void instrument_coverage_optimize(const cs_insn    *instr,
 
   // gum_arm64_writer_put_brk_imm(cw, 0x0);
 
-  instrument_coverage_suppress_init();
+  if (instrument_suppress) { instrument_coverage_suppress_init(); }
 
   code_addr = cw->pc;
 
@@ -304,9 +333,13 @@ void instrument_coverage_optimize(const cs_insn    *instr,
   block_start =
       GSIZE_TO_POINTER(GUM_ADDRESS(cw->code) - GUM_RESTORATION_PROLOG_SIZE);
 
-  if (!g_hash_table_add(coverage_blocks, block_start)) {
+  if (instrument_suppress) {
+
+    if (!g_hash_table_add(coverage_blocks, block_start)) {
 
-    FATAL("Failed - g_hash_table_add");
+      FATAL("Failed - g_hash_table_add");
+
+    }
 
   }
 
@@ -342,7 +375,17 @@ void instrument_coverage_optimize(const cs_insn    *instr,
 
   code.code.mov_x1_curr_loc_shr_1 |= (area_offset_ror << 5);
 
-  gum_arm64_writer_put_bytes(cw, code.bytes, sizeof(afl_log_code));
+  if (instrument_suppress) {
+
+    gum_arm64_writer_put_bytes(cw, code.bytes, sizeof(afl_log_code));
+
+  } else {
+
+    size_t offset = offsetof(afl_log_code, code.stp_x0_x1);
+    gum_arm64_writer_put_bytes(cw, &code.bytes[offset],
+                               sizeof(afl_log_code) - offset);
+
+  }
 
 }
 
diff --git a/frida_mode/src/instrument/instrument_x64.c b/frida_mode/src/instrument/instrument_x64.c
index 13ced4a3..8338f8e7 100644
--- a/frida_mode/src/instrument/instrument_x64.c
+++ b/frida_mode/src/instrument/instrument_x64.c
@@ -171,11 +171,11 @@ void instrument_coverage_optimize_init(void) {
 
 }
 
-static void instrument_coverage_switch(GumStalkerObserver *self,
-                                       gpointer            from_address,
-                                       gpointer            start_address,
-                                       const cs_insn      *from_insn,
-                                       gpointer           *target) {
+static void instrument_coverage_switch_insn(GumStalkerObserver *self,
+                                            gpointer            from_address,
+                                            gpointer            start_address,
+                                            const cs_insn      *from_insn,
+                                            gpointer           *target) {
 
   UNUSED_PARAMETER(self);
   UNUSED_PARAMETER(from_address);
@@ -224,6 +224,35 @@ static void instrument_coverage_switch(GumStalkerObserver *self,
 
 }
 
+cs_insn *instrument_disassemble(gconstpointer address) {
+
+  csh      capstone;
+  cs_insn *insn = NULL;
+
+  cs_open(CS_ARCH_X86, GUM_CPU_MODE, &capstone);
+  cs_option(capstone, CS_OPT_DETAIL, CS_OPT_ON);
+
+  cs_disasm(capstone, address, 16, GPOINTER_TO_SIZE(address), 1, &insn);
+
+  cs_close(&capstone);
+
+  return insn;
+
+}
+
+static void instrument_coverage_switch(GumStalkerObserver *self,
+                                       gpointer            from_address,
+                                       gpointer start_address, void *from_insn,
+                                       gpointer *target) {
+
+  if (from_insn == NULL) { return; }
+  cs_insn *insn = instrument_disassemble(from_insn);
+  instrument_coverage_switch_insn(self, from_address, start_address, insn,
+                                  target);
+  cs_free(insn, 1);
+
+}
+
 static void instrument_coverage_suppress_init(void) {
 
   static gboolean initialized = false;
@@ -351,11 +380,15 @@ void instrument_coverage_optimize(const cs_insn    *instr,
 
   }
 
-  instrument_coverage_suppress_init();
+  if (instrument_suppress) {
+
+    instrument_coverage_suppress_init();
+
+    if (!g_hash_table_add(coverage_blocks, GSIZE_TO_POINTER(cw->code))) {
 
-  if (!g_hash_table_add(coverage_blocks, GSIZE_TO_POINTER(cw->code))) {
+      FATAL("Failed - g_hash_table_add");
 
-    FATAL("Failed - g_hash_table_add");
+    }
 
   }
 
diff --git a/frida_mode/src/instrument/instrument_x86.c b/frida_mode/src/instrument/instrument_x86.c
index eabd5be4..4667ea29 100644
--- a/frida_mode/src/instrument/instrument_x86.c
+++ b/frida_mode/src/instrument/instrument_x86.c
@@ -83,11 +83,11 @@ gboolean instrument_is_coverage_optimize_supported(void) {
 
 }
 
-static void instrument_coverage_switch(GumStalkerObserver *self,
-                                       gpointer            from_address,
-                                       gpointer            start_address,
-                                       const cs_insn      *from_insn,
-                                       gpointer           *target) {
+static void instrument_coverage_switch_insn(GumStalkerObserver *self,
+                                            gpointer            from_address,
+                                            gpointer            start_address,
+                                            const cs_insn      *from_insn,
+                                            gpointer           *target) {
 
   UNUSED_PARAMETER(self);
   UNUSED_PARAMETER(from_address);
@@ -130,6 +130,35 @@ static void instrument_coverage_switch(GumStalkerObserver *self,
 
 }
 
+cs_insn *instrument_disassemble(gconstpointer address) {
+
+  csh      capstone;
+  cs_insn *insn = NULL;
+
+  cs_open(CS_ARCH_X86, GUM_CPU_MODE, &capstone);
+  cs_option(capstone, CS_OPT_DETAIL, CS_OPT_ON);
+
+  cs_disasm(capstone, address, 16, GPOINTER_TO_SIZE(address), 1, &insn);
+
+  cs_close(&capstone);
+
+  return insn;
+
+}
+
+static void instrument_coverage_switch(GumStalkerObserver *self,
+                                       gpointer            from_address,
+                                       gpointer start_address, void *from_insn,
+                                       gpointer *target) {
+
+  if (from_insn == NULL) { return; }
+  cs_insn *insn = instrument_disassemble(from_insn);
+  instrument_coverage_switch_insn(self, from_address, start_address, insn,
+                                  target);
+  cs_free(insn, 1);
+
+}
+
 static void instrument_coverage_suppress_init(void) {
 
   static gboolean initialized = false;
@@ -174,13 +203,17 @@ void instrument_coverage_optimize(const cs_insn    *instr,
 
   code.code = template;
 
-  instrument_coverage_suppress_init();
+  if (instrument_suppress) {
+
+    instrument_coverage_suppress_init();
+
+    // gum_x86_writer_put_breakpoint(cw);
 
-  // gum_x86_writer_put_breakpoint(cw);
+    if (!g_hash_table_add(coverage_blocks, GSIZE_TO_POINTER(cw->code))) {
 
-  if (!g_hash_table_add(coverage_blocks, GSIZE_TO_POINTER(cw->code))) {
+      FATAL("Failed - g_hash_table_add");
 
-    FATAL("Failed - g_hash_table_add");
+    }
 
   }