4 files changed, 47 insertions, 12 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_arm64.c b/frida_mode/src/instrument/instrument_arm64.c
index 39e32b12..4372861d 100644
--- a/frida_mode/src/instrument/instrument_arm64.c
+++ b/frida_mode/src/instrument/instrument_arm64.c
@@ -196,7 +196,15 @@ static void instrument_coverage_switch(GumStalkerObserver *self,
   insn = instrument_disassemble(from_insn);
   deterministic = instrument_is_deterministic(insn);
   cs_free(insn, 1);
-  if (!deterministic) { return; }
+
+  /*
+   * 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
@@ -305,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;
 
@@ -325,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");
+    }
 
   }
 
@@ -363,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 f7b7d6c5..8338f8e7 100644
--- a/frida_mode/src/instrument/instrument_x64.c
+++ b/frida_mode/src/instrument/instrument_x64.c
@@ -380,11 +380,15 @@ void instrument_coverage_optimize(const cs_insn    *instr,
 
   }
 
-  instrument_coverage_suppress_init();
+  if (instrument_suppress) {
 
-  if (!g_hash_table_add(coverage_blocks, GSIZE_TO_POINTER(cw->code))) {
+    instrument_coverage_suppress_init();
 
-    FATAL("Failed - g_hash_table_add");
+    if (!g_hash_table_add(coverage_blocks, GSIZE_TO_POINTER(cw->code))) {
+
+      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 f15893cb..4667ea29 100644
--- a/frida_mode/src/instrument/instrument_x86.c
+++ b/frida_mode/src/instrument/instrument_x86.c
@@ -203,13 +203,17 @@ void instrument_coverage_optimize(const cs_insn    *instr,
 
   code.code = template;
 
-  instrument_coverage_suppress_init();
+  if (instrument_suppress) {
 
-  // gum_x86_writer_put_breakpoint(cw);
+    instrument_coverage_suppress_init();
 
-  if (!g_hash_table_add(coverage_blocks, GSIZE_TO_POINTER(cw->code))) {
+    // gum_x86_writer_put_breakpoint(cw);
 
-    FATAL("Failed - g_hash_table_add");
+    if (!g_hash_table_add(coverage_blocks, GSIZE_TO_POINTER(cw->code))) {
+
+      FATAL("Failed - g_hash_table_add");
+
+    }
 
   }