AARCH64 branch suppression

1 files changed, 123 insertions, 0 deletions

diff --git a/frida_mode/src/instrument/instrument_arm64.c b/frida_mode/src/instrument/instrument_arm64.c
index 57b60317..e1dc2441 100644
--- a/frida_mode/src/instrument/instrument_arm64.c
+++ b/frida_mode/src/instrument/instrument_arm64.c
@@ -14,9 +14,12 @@
 
 #define PAGE_MASK (~(GUM_ADDRESS(0xfff)))
 #define PAGE_ALIGNED(x) ((GUM_ADDRESS(x) & PAGE_MASK) == GUM_ADDRESS(x))
+#define GUM_RESTORATION_PROLOG_SIZE 4
 
 #if defined(__aarch64__)
 
+static GHashTable *coverage_blocks = NULL;
+
 __attribute__((aligned(0x1000))) static guint8 area_ptr_dummy[MAP_SIZE];
 
   #pragma pack(push, 1)
@@ -120,6 +123,101 @@ 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_arm64 *   arm64;
+  arm64_cc     cc;
+  gsize        fixup_offset;
+
+  if (from_insn == NULL) { return; }
+
+  arm64 = &from_insn->detail->arm64;
+  cc = arm64->cc;
+
+  if (!g_hash_table_contains(coverage_blocks, GSIZE_TO_POINTER(*target))) {
+
+    return;
+
+  }
+
+  switch (from_insn->id) {
+
+    case ARM64_INS_B:
+    case ARM64_INS_BL:
+      if (cc != ARM64_CC_INVALID) { return; }
+      break;
+
+    case ARM64_INS_RET:
+    case ARM64_INS_RETAA:
+    case ARM64_INS_RETAB:
+      if (arm64->op_count != 0) { return; }
+      break;
+    default:
+      return;
+
+  }
+
+  /*
+   * Since each block is prefixed with a restoration prologue, we need to be
+   * able to begin execution at an offset into the block and execute both this
+   * restoration prologue and the instrumented block without the coverage code.
+   * We therefore layout our block as follows:
+   *
+   *  +-----+--------------------------+------------------+-----+-------------+
+   *  | LDP | COVERAGE INSTRUMENTATION | BR <TARGET CODE> | LDP | TARGET CODE |
+   *  +-----+--------------------------+------------------+-----+-------------+
+   *
+   *  ^     ^                                             ^     ^
+   *  |     |                                             |     |
+   *  A     B                                             C     D
+   *
+   * Without instrumentation suppression, the block is either executed at point
+   * (A) if it is reached by an indirect branch (and registers need to be
+   * restored) or point (B) if it is reached by an direct branch (and hence the
+   * registers don't need restoration). Similarly, we can start execution of the
+   * block at points (C) or (D) to achieve the same functionality, but without
+   * executing the coverage instrumentation.
+   *
+   * In either case, Stalker will call us back with the address of the target
+   * block to be executed as the destination. It is not until later that Stalker
+   * will determine which branch type is required given the location of its
+   * instrumented code and add the `GUM_RESTORATION_PROLOG_SIZE` to the target
+   * address. Therefore, we need to map the address of point (A) to that of
+   * point (C) and also ensure that the offset between (A)->(B) and (C)->(D) is
+   * identical so that if Stalker can use an immediate call, it still branches
+   * to a valid offset.
+   */
+  fixup_offset = GUM_RESTORATION_PROLOG_SIZE +
+                 G_STRUCT_OFFSET(afl_log_code_asm_t, restoration_prolog);
+  *target += fixup_offset;
+
+}
+
+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 gboolean instrument_coverage_in_range(gssize offset) {
 
   return (offset >= G_MININT33 && offset <= G_MAXINT33);
@@ -151,6 +249,7 @@ void instrument_coverage_optimize(const cs_insn *   instr,
 
   afl_log_code    code = {0};
   GumArm64Writer *cw = output->writer.arm64;
+  gpointer   block_start;
   guint64 area_offset = instrument_get_offset_hash(GUM_ADDRESS(instr->address));
   gsize   map_size_pow2;
   gsize   area_offset_ror;
@@ -171,8 +270,32 @@ void instrument_coverage_optimize(const cs_insn *   instr,
 
   // gum_arm64_writer_put_brk_imm(cw, 0x0);
 
+  instrument_coverage_suppress_init();
+
   code_addr = cw->pc;
 
+  /*
+   * On AARCH64, immediate branches can only be encoded with a 28-bit offset. To
+   * make a longer branch, it is necessary to load a register with the target
+   * address, this register must be saved beyond the red-zone before the branch
+   * is taken. To restore this register each block is prefixed by Stalker with
+   * an instruction to load x16,x17 from beyond the red-zone on the stack. A
+   * pair of registers are saved/restored because on AARCH64, the stack pointer
+   * must be 16 byte aligned. This instruction is emitted into the block before
+   * the tranformer (from which we are called) is executed. If is is possible
+   * for Stalker to make a direct branch (the target block is close enough), it
+   * can forego pushing the registers and instead branch at an offset into the
+   * block to skip this restoration prolog.
+   */
+  block_start =
+      GSIZE_TO_POINTER(GUM_ADDRESS(cw->code) - GUM_RESTORATION_PROLOG_SIZE);
+
+  if (!g_hash_table_add(coverage_blocks, block_start)) {
+
+    FATAL("Failed - g_hash_table_add");
+
+  }
+
   code.code = template;
 
   /*