Changes to skip coverage code on deterministic branches on x86

1 files changed, 75 insertions, 0 deletions

diff --git a/frida_mode/src/instrument/instrument_x86.c b/frida_mode/src/instrument/instrument_x86.c
index c2f36995..ad837e2d 100644
--- a/frida_mode/src/instrument/instrument_x86.c
+++ b/frida_mode/src/instrument/instrument_x86.c
@@ -1,10 +1,13 @@
 #include "frida-gumjs.h"
 
 #include "instrument.h"
+#include "stalker.h"
 #include "util.h"
 
 #if defined(__i386__)
 
+static GHashTable *coverage_blocks = NULL;
+
   #pragma pack(push, 1)
 typedef struct {
 
@@ -77,6 +80,70 @@ gboolean instrument_is_coverage_optimize_supported(void) {
 
 }
 
+static void instrument_coverage_switch(GumStalkerObserver *self,
+                                       gpointer            start_address,
+                                       const cs_insn *     from_insn,
+                                       gpointer *          target) {
+
+  UNUSED_PARAMETER(self);
+  UNUSED_PARAMETER(start_address);
+
+  cs_x86 *   x86;
+  cs_x86_op *op;
+  if (from_insn == NULL) { return; }
+
+  x86 = &from_insn->detail->x86;
+  op = x86->operands;
+
+  if (!g_hash_table_contains(coverage_blocks, GSIZE_TO_POINTER(*target))) {
+
+    return;
+
+  }
+
+  switch (from_insn->id) {
+
+    case X86_INS_CALL:
+    case X86_INS_JMP:
+      if (x86->op_count != 1) {
+
+        FATAL("Unexpected operand count: %d", x86->op_count);
+
+      }
+
+      if (op[0].type != X86_OP_IMM) { return; }
+
+      break;
+    case X86_INS_RET:
+      break;
+    default:
+      return;
+
+  }
+
+  *target = (guint8 *)*target + sizeof(afl_log_code);
+
+}
+
+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);
+
+  }
+
+}
+
 void instrument_coverage_optimize(const cs_insn *   instr,
                                   GumStalkerOutput *output) {
 
@@ -88,8 +155,16 @@ void instrument_coverage_optimize(const cs_insn *   instr,
 
   code.code = template;
 
+  instrument_coverage_suppress_init();
+
   // gum_x86_writer_put_breakpoint(cw);
 
+  if (!g_hash_table_add(coverage_blocks, GSIZE_TO_POINTER(cw->code))) {
+
+    FATAL("Failed - g_hash_table_add");
+
+  }
+
   gssize prev_loc_value_offset2 =
       offsetof(afl_log_code, code.mov_eax_prev_loc) +
       sizeof(code.code.mov_eax_prev_loc) - sizeof(gint);