ARM branch suppression

1 files changed, 58 insertions, 5 deletions

diff --git a/frida_mode/src/instrument/instrument_arm32.c b/frida_mode/src/instrument/instrument_arm32.c
index f2e825ee..5b6ddf09 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__)
@@ -10,6 +11,7 @@
 
 gboolean instrument_cache_enabled = FALSE;
 gsize    instrument_cache_size = 0;
+static GHashTable *coverage_blocks = NULL;
 
 extern __thread guint64 instrument_previous_pc;
 
@@ -22,8 +24,25 @@ 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 */
-  uint32_t  b_code;                                                /* b imm */
+  // 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,15 +134,45 @@ gboolean instrument_is_coverage_optimize_supported(void) {
 
 }
 
-static void patch_t3_insn(uint32_t *insn, uint16_t val) {
+static void instrument_coverage_switch(GumStalkerObserver *self,
+                                       gpointer            from_address,
+                                       gpointer            start_address,
+                                       const cs_insn *     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 += 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);
   uint32_t imm12 = (val & 0xfff);
   uint32_t imm4 = (val >> 12);
   orig |= imm12;
   orig |= (imm4 << 16);
   *insn = GUINT32_TO_LE(orig);
-
 }
 
 void instrument_coverage_optimize(const cs_insn    *instr,
@@ -137,12 +186,16 @@ void instrument_coverage_optimize(const cs_insn    *instr,
   gsize   area_offset_ror;
   GumAddress code_addr = 0;
 
-  // gum_arm64_writer_put_brk_imm(cw, 0x0);
+  instrument_coverage_suppress_init();
 
   code_addr = cw->pc;
 
   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));