Optimize AARCH64 inline assembly

2 files changed, 177 insertions, 57 deletions

diff --git a/frida_mode/src/instrument/instrument.c b/frida_mode/src/instrument/instrument.c
index e3f3717e..bf102a82 100644
--- a/frida_mode/src/instrument/instrument.c
+++ b/frida_mode/src/instrument/instrument.c
@@ -32,7 +32,7 @@ char *   instrument_coverage_unstable_filename = NULL;
 
 static GumStalkerTransformer *transformer = NULL;
 
-__thread guint64 instrument_previous_pc = 0;
+__attribute__((aligned(0x1000))) __thread guint64 instrument_previous_pc = 0;
 
 static GumAddress previous_rip = 0;
 static GumAddress previous_end = 0;
diff --git a/frida_mode/src/instrument/instrument_arm64.c b/frida_mode/src/instrument/instrument_arm64.c
index 0f635458..27142f1f 100644
--- a/frida_mode/src/instrument/instrument_arm64.c
+++ b/frida_mode/src/instrument/instrument_arm64.c
@@ -1,50 +1,121 @@
+#include <stddef.h>
+
 #include "frida-gumjs.h"
 
 #include "config.h"
 
 #include "instrument.h"
+#include "ranges.h"
+#include "stalker.h"
+#include "util.h"
+
+#define G_MININT33 ((gssize)0xffffffff00000000)
+#define G_MAXINT33 ((gssize)0x00000000ffffffff)
+
+#define PAGE_MASK (~(GUM_ADDRESS(0xfff)))
+#define PAGE_ALIGNED(x) ((GUM_ADDRESS(x) & PAGE_MASK) == GUM_ADDRESS(x))
 
 #if defined(__aarch64__)
 
-static GumAddress current_log_impl = GUM_ADDRESS(0);
+__attribute__((aligned(0x1000))) static guint8 area_ptr_dummy[MAP_SIZE];
+
+  #pragma pack(push, 1)
+typedef struct {
+
+  // cur_location = (block_address >> 4) ^ (block_address << 8);
+  // shared_mem[cur_location ^ prev_location]++;
+  // prev_location = cur_location >> 1;
+
+  // stp     x0, x1, [sp, #-160]
+  // adrp    x0, 0x7fb7ff4000
+  // ldr     x1, [x0]
+  // mov     x0, #0x18b8
+  // eor     x0, x1, x0
+  // adrp    x1, 0x7fb7d73000
+  // add     x0, x1, x0
+  // ldrb    w1, [x0]
+  // add     w1, w1, #0x1
+  // tbz     w1, #8, 0x7fb6f0dec8
+  // mov     w1, #0x1
+  // strb    w1, [x0]
+  // adrp    x0, 0x7fb7ff4000
+  // mov     x1, #0xc5c
+  // str     x1, [x0]
+  // ldp     x0, x1, [sp, #-160]
+  // b       0x7fb6f0dee4
+  // ldp     x16, x17, [sp], #144
+
+  uint32_t stp_x0_x1;                           /* stp x0, x1, [sp, #-0xa0] */
+
+  uint32_t adrp_x0_prev_loc1;                           /* adrp x0, #0xXXXX */
+  uint32_t ldr_x1_ptr_x0;                                   /* ldr x1, [x0] */
+
+  uint32_t mov_x0_curr_loc;                             /* movz x0, #0xXXXX */
+  uint32_t eor_x0_x1_x0;                                  /* eor x0, x1, x0 */
+  uint32_t adrp_x1_area_ptr;                            /* adrp x1, #0xXXXX */
+  uint32_t add_x0_x1_x0;                                  /* add x0, x1, x0 */
+
+  uint32_t ldrb_w1_x0;                                     /* ldrb w1, [x0] */
+  uint32_t add_w1_w1_1;                                   /* add w1, w1, #1 */
+  uint32_t tbz_w1_8_8;                                    /* tbz w1, #8, #8 */
+  uint32_t mov_w1_1;                                          /* mov w1, #1 */
+
+  uint32_t strb_w1_ptr_x0;                                 /* strb w1, [x0] */
+
+  uint32_t adrp_x0_prev_loc2;                           /* adrp x0, #0xXXXX */
+  uint32_t mov_x1_curr_loc_shr_1;                       /* movz x1, #0xXXXX */
+  uint32_t str_x1_ptr_x0;                                   /* str x1, [x0] */
+
+  uint32_t ldp_x0_x1;                           /* ldp x0, x1, [sp, #-0xa0] */
 
-static const guint8 afl_log_code[] = {
+  uint32_t b_imm8;                                                 /* br #8 */
+  uint32_t restoration_prolog;                 /* ldp x16, x17, [sp], #0x90 */
 
-    // __afl_area_ptr[current_pc ^ previous_pc]++;
-    // previous_pc = current_pc ROR 1;
-    0xE1, 0x0B, 0xBF, 0xA9,  // stp x1, x2, [sp, -0x10]!
-    0xE3, 0x13, 0xBF, 0xA9,  // stp x3, x4, [sp, -0x10]!
+} afl_log_code_asm_t;
 
-    // x0 = current_pc
-    0x21, 0x02, 0x00, 0x58,  // ldr x1, #0x44, =&__afl_area_ptr
-    0x21, 0x00, 0x40, 0xf9,  // ldr x1, [x1] (=__afl_area_ptr)
+  #pragma pack(pop)
 
-    0x22, 0x02, 0x00, 0x58,  // ldr x2, #0x44, =&previous_pc
-    0x42, 0x00, 0x40, 0xf9,  // ldr x2, [x2] (=previous_pc)
+typedef union {
 
-    // __afl_area_ptr[current_pc ^ previous_pc]++;
-    0x42, 0x00, 0x00, 0xca,  // eor x2, x2, x0
-    0x23, 0x68, 0x62, 0xf8,  // ldr x3, [x1, x2]
-    0x63, 0x04, 0x00, 0x91,  // add x3, x3, #1
-    0x63, 0x00, 0x1f, 0x9a,  // adc x3, x3, xzr
-    0x23, 0x68, 0x22, 0xf8,  // str x3, [x1, x2]
+  afl_log_code_asm_t code;
+  uint8_t            bytes[0];
 
-    // previous_pc = current_pc ROR 1;
-    0xe4, 0x07, 0x40, 0x8b,  // add x4, xzr, x0, LSR #1
-    0xe0, 0xff, 0x00, 0x8b,  // add x0, xzr, x0, LSL #63
-    0x80, 0xc0, 0x40, 0x8b,  // add x0, x4, x0, LSR #48
+} afl_log_code;
 
-    0xe2, 0x00, 0x00, 0x58,  // ldr x2, #0x1c, =&previous_pc
-    0x40, 0x00, 0x00, 0xf9,  // str x0, [x2]
+static const afl_log_code_asm_t template =
+    {
 
-    0xE3, 0x13, 0xc1, 0xA8,  // ldp x3, x4, [sp], #0x10
-    0xE1, 0x0B, 0xc1, 0xA8,  // ldp x1, x2, [sp], #0x10
-    0xC0, 0x03, 0x5F, 0xD6,  // ret
+        .stp_x0_x1 = 0xa93607e0,
 
-    // &afl_area_ptr_ptr
-    // &afl_prev_loc_ptr
+        .adrp_x0_prev_loc1 = 0xb0000000,
+        .ldr_x1_ptr_x0 = 0xf9400001,
 
-};
+        .mov_x0_curr_loc = 0xd2800000,
+        .eor_x0_x1_x0 = 0xca000020,
+
+        .adrp_x1_area_ptr = 0x90000001,
+        .add_x0_x1_x0 = 0x8b000020,
+
+        .ldrb_w1_x0 = 0x39400001,
+
+        .add_w1_w1_1 = 0x11000421,
+        .tbz_w1_8_8 = 0x36400041,
+        .mov_w1_1 = 0x52800021,
+
+        .strb_w1_ptr_x0 = 0x39000001,
+
+        .adrp_x0_prev_loc2 = 0xb0000000,
+        .mov_x1_curr_loc_shr_1 = 0xd2800001,
+        .str_x1_ptr_x0 = 0xf9000001,
+
+        .ldp_x0_x1 = 0xa97607e0,
+
+        .b_imm8 = 0x14000002,
+        .restoration_prolog = 0xa8c947f0,
+
+}
+
+;
 
 gboolean instrument_is_coverage_optimize_supported(void) {
 
@@ -52,50 +123,99 @@ gboolean instrument_is_coverage_optimize_supported(void) {
 
 }
 
+static gboolean instrument_coverage_in_range(gssize offset) {
+
+  return (offset >= G_MININT33 && offset <= G_MAXINT33);
+
+}
+
+static void instrument_patch_ardp(guint32 *patch, GumAddress insn,
+                                  GumAddress target) {
+
+  if (!PAGE_ALIGNED(target)) { FATAL("Target not page aligned"); }
+
+  gssize distance = target - (GUM_ADDRESS(insn) & PAGE_MASK);
+  if (!instrument_coverage_in_range(distance)) {
+
+    FATAL("Patch out of range 0x%016lX->0x%016lX = 0x%016lX", insn, target,
+          distance);
+
+  }
+
+  guint32 imm_low = ((distance >> 12) & 0x3) << 29;
+  guint32 imm_high = ((distance >> 14) & 0x7FFFF) << 5;
+  *patch |= imm_low;
+  *patch |= imm_high;
+
+}
+
 void instrument_coverage_optimize(const cs_insn *   instr,
                                   GumStalkerOutput *output) {
 
-  guint64 current_pc = instr->address;
-  guint64 area_offset = instrument_get_offset_hash(GUM_ADDRESS(instr->address));
+  afl_log_code    code = {0};
   GumArm64Writer *cw = output->writer.arm64;
+  guint64 area_offset = instrument_get_offset_hash(GUM_ADDRESS(instr->address));
+  gsize      map_size_pow2;
+  gsize      area_offset_ror;
+  GumAddress code_addr = 0;
 
-  if (current_log_impl == 0 ||
-      !gum_arm64_writer_can_branch_directly_between(cw, cw->pc,
-                                                    current_log_impl) ||
-      !gum_arm64_writer_can_branch_directly_between(cw, cw->pc + 128,
-                                                    current_log_impl)) {
+  // gum_arm64_writer_put_brk_imm(cw, 0x0);
 
-    gconstpointer after_log_impl = cw->code + 1;
+  code_addr = cw->pc;
 
-    gum_arm64_writer_put_b_label(cw, after_log_impl);
+  code.code = template;
 
-    current_log_impl = cw->pc;
-    gum_arm64_writer_put_bytes(cw, afl_log_code, sizeof(afl_log_code));
+  /*
+   * Given our map is allocated on a 64KB boundary and our map is a multiple of
+   * 64KB in size, then it should also end on a 64 KB boundary. It is followed
+   * by our previous_pc, so this too should be 64KB aligned.
+   */
+  g_assert(PAGE_ALIGNED(&instrument_previous_pc));
+  g_assert(PAGE_ALIGNED(__afl_area_ptr));
 
-    uint8_t **afl_area_ptr_ptr = &__afl_area_ptr;
-    uint64_t *afl_prev_loc_ptr = &instrument_previous_pc;
-    gum_arm64_writer_put_bytes(cw, (const guint8 *)&afl_area_ptr_ptr,
-                               sizeof(afl_area_ptr_ptr));
-    gum_arm64_writer_put_bytes(cw, (const guint8 *)&afl_prev_loc_ptr,
-                               sizeof(afl_prev_loc_ptr));
+  instrument_patch_ardp(
+      &code.code.adrp_x0_prev_loc1,
+      code_addr + offsetof(afl_log_code, code.adrp_x0_prev_loc1),
+      GUM_ADDRESS(&instrument_previous_pc));
 
-    gum_arm64_writer_put_label(cw, after_log_impl);
+  code.code.mov_x0_curr_loc |= area_offset << 5;
 
-  }
+  instrument_patch_ardp(
+      &code.code.adrp_x1_area_ptr,
+      code_addr + offsetof(afl_log_code, code.adrp_x1_area_ptr),
+      GUM_ADDRESS(__afl_area_ptr));
+
+  map_size_pow2 = util_log2(__afl_map_size);
+  area_offset_ror = util_rotate(area_offset, 1, map_size_pow2);
 
-  gum_arm64_writer_put_stp_reg_reg_reg_offset(
-      cw, ARM64_REG_LR, ARM64_REG_X0, ARM64_REG_SP, -(16 + GUM_RED_ZONE_SIZE),
-      GUM_INDEX_PRE_ADJUST);
-  gum_arm64_writer_put_ldr_reg_u64(cw, ARM64_REG_X0, area_offset);
-  gum_arm64_writer_put_bl_imm(cw, current_log_impl);
-  gum_arm64_writer_put_ldp_reg_reg_reg_offset(
-      cw, ARM64_REG_LR, ARM64_REG_X0, ARM64_REG_SP, 16 + GUM_RED_ZONE_SIZE,
-      GUM_INDEX_POST_ADJUST);
+  instrument_patch_ardp(
+      &code.code.adrp_x0_prev_loc2,
+      code_addr + offsetof(afl_log_code, code.adrp_x0_prev_loc2),
+      GUM_ADDRESS(&instrument_previous_pc));
+
+  code.code.mov_x1_curr_loc_shr_1 |= (area_offset_ror << 5);
+
+  gum_arm64_writer_put_bytes(cw, code.bytes, sizeof(afl_log_code));
 
 }
 
 void instrument_coverage_optimize_init(void) {
 
+  char *shm_env = getenv(SHM_ENV_VAR);
+  FVERBOSE("SHM_ENV_VAR: %s", shm_env);
+
+  if (shm_env == NULL) {
+
+    FWARNF("SHM_ENV_VAR not set, using dummy for debugging purposes");
+
+    __afl_area_ptr = area_ptr_dummy;
+    memset(area_ptr_dummy, '\0', sizeof(area_ptr_dummy));
+
+  }
+
+  FVERBOSE("__afl_area_ptr: %p", __afl_area_ptr);
+  FVERBOSE("instrument_previous_pc: %p", &instrument_previous_pc);
+
 }
 
 void instrument_flush(GumStalkerOutput *output) {