Fixes for arm/arm64

4 files changed, 44 insertions, 111 deletions

diff --git a/frida_mode/GNUmakefile b/frida_mode/GNUmakefile
index c055fcbb..6f58ebbb 100644
--- a/frida_mode/GNUmakefile
+++ b/frida_mode/GNUmakefile
@@ -214,6 +214,9 @@ all: $(FRIDA_TRACE) $(FRIDA_TRACE_LIB) $(AFLPP_FRIDA_DRIVER_HOOK_OBJ) $(AFLPP_QE
 arm:
 	CFLAGS="-marm" LDFLAGS="-marm" ARCH="armhf" TARGET_CC=arm-linux-gnueabihf-gcc TARGET_CXX=arm-linux-gnueabihf-g++ make all
 
+arm64:
+	ARCH="arm64" TARGET_CC=aarch64-linux-gnu-gcc TARGET_CXX=aarch64-linux-gnu-g++ make all
+
 $(BUILD_DIR):
 	mkdir -p $(BUILD_DIR)
 
diff --git a/frida_mode/src/persistent/persistent_arm32.c b/frida_mode/src/persistent/persistent_arm32.c
index dbe51eb5..8790f483 100644
--- a/frida_mode/src/persistent/persistent_arm32.c
+++ b/frida_mode/src/persistent/persistent_arm32.c
@@ -33,7 +33,7 @@
 // r15 - pc
 
 static GumCpuContext saved_regs = {0};
-static gpointer      saved_lr = NULL;
+static gpointer      persistent_loop = NULL;
 
 gboolean persistent_is_supported(void) {
 
@@ -141,17 +141,10 @@ static void instrument_persitent_restore_regs(GumArmWriter  *cw,
 
 }
 
-static void instrument_exit(GumArmWriter *cw) {
+static void instrument_afl_persistent_loop_func(void) {
 
-  gum_arm_writer_put_sub_reg_reg_reg(cw, ARM_REG_R0, ARM_REG_R0, ARM_REG_R0);
-  gum_arm_writer_put_call_address_with_arguments(cw, GUM_ADDRESS(_exit), 1,
-                                                 GUM_ARG_REGISTER, ARM_REG_R0);
+  if (__afl_persistent_loop(persistent_count) == 0) { _exit(0); }
 
-}
-
-static int instrument_afl_persistent_loop_func(void) {
-
-  int ret = __afl_persistent_loop(persistent_count);
   if (instrument_previous_pc_addr == NULL) {
 
     FATAL("instrument_previous_pc_addr uninitialized");
@@ -159,7 +152,6 @@ static int instrument_afl_persistent_loop_func(void) {
   }
 
   *instrument_previous_pc_addr = instrument_hash_zero;
-  return ret;
 
 }
 
@@ -203,7 +195,8 @@ static void instrument_persitent_save_lr(GumArmWriter *cw) {
   gum_arm_writer_put_str_reg_reg_offset(cw, ARM_REG_R0, ARM_REG_SP,
                                         GUM_RED_ZONE_SIZE);
 
-  gum_arm_writer_put_ldr_reg_address(cw, ARM_REG_R0, GUM_ADDRESS(&saved_lr));
+  gum_arm_writer_put_ldr_reg_address(cw, ARM_REG_R0,
+                                     GUM_ADDRESS(&persistent_ret));
   gum_arm_writer_put_str_reg_reg_offset(cw, ARM_REG_LR, ARM_REG_R0, 0);
 
   gum_arm_writer_put_ldr_reg_reg_offset(cw, ARM_REG_R0, ARM_REG_SP,
@@ -214,65 +207,35 @@ static void instrument_persitent_save_lr(GumArmWriter *cw) {
 void persistent_prologue_arch(GumStalkerOutput *output) {
 
   /*
+   *  SAVE RET (Used to write the epilogue if persistent_ret is not set)
    *  SAVE REGS
-   *  SAVE RET
-   *  POP RET
-   * loop:
+   * loop: (Save address of where the eiplogue should jump back to)
    *  CALL instrument_afl_persistent_loop
-   *  TEST EAX, EAX
-   *  JZ end:
-   *  call hook (optionally)
+   *  CALL hook (optionally)
    *  RESTORE REGS
-   *  call original
-   *  jmp loop:
-   *
-   * end:
-   *  JMP SAVED RET
-   *
-   * original:
    *  INSTRUMENTED PERSISTENT FUNC
    */
 
   GumArmWriter *cw = output->writer.arm;
 
-  gconstpointer loop = cw->code + 1;
-
   FVERBOSE("Persistent loop reached");
 
+  if (persistent_ret == 0) { instrument_persitent_save_lr(cw); }
+
+  /* Save the current context */
   instrument_persitent_save_regs(cw, &saved_regs);
 
-  /* loop: */
-  gum_arm_writer_put_label(cw, loop);
+  /* Store a pointer to where we should return for our next iteration */
+  persistent_loop = gum_arm_writer_cur(cw);
 
-  /* call instrument_prologue_func */
+  /* call __afl_persistent_loop and _exit if zero. Also reset our previous_pc */
   instrument_afl_persistent_loop(cw);
 
-  /* jz done */
-  gconstpointer done = cw->code + 1;
-  gum_arm_writer_put_cmp_reg_imm(cw, ARM_REG_R0, 0);
-  gum_arm_writer_put_b_cond_label(cw, ARM_CC_EQ, done);
-
   /* Optionally call the persistent hook */
   persistent_prologue_hook(cw, &saved_regs);
 
+  /* Restore our CPU context before we continue execution */
   instrument_persitent_restore_regs(cw, &saved_regs);
-  gconstpointer original = cw->code + 1;
-  /* call original */
-
-  gum_arm_writer_put_bl_label(cw, original);
-
-  /* jmp loop */
-  gum_arm_writer_put_b_label(cw, loop);
-
-  /* done: */
-  gum_arm_writer_put_label(cw, done);
-
-  instrument_exit(cw);
-
-  /* original: */
-  gum_arm_writer_put_label(cw, original);
-
-  instrument_persitent_save_lr(cw);
 
   if (persistent_debug) { gum_arm_writer_put_breakpoint(cw); }
 
@@ -284,7 +247,8 @@ void persistent_epilogue_arch(GumStalkerOutput *output) {
 
   if (persistent_debug) { gum_arm_writer_put_breakpoint(cw); }
 
-  gum_arm_writer_put_ldr_reg_address(cw, ARM_REG_R0, GUM_ADDRESS(&saved_lr));
+  gum_arm_writer_put_ldr_reg_address(cw, ARM_REG_R0,
+                                     GUM_ADDRESS(&persistent_loop));
 
   gum_arm_writer_put_ldr_reg_reg_offset(cw, ARM_REG_R0, ARM_REG_R0, 0);
 
diff --git a/frida_mode/src/persistent/persistent_arm64.c b/frida_mode/src/persistent/persistent_arm64.c
index 565a2b8c..cfd00b17 100644
--- a/frida_mode/src/persistent/persistent_arm64.c
+++ b/frida_mode/src/persistent/persistent_arm64.c
@@ -16,7 +16,7 @@ typedef struct {
 } persistent_ctx_t;
 
 static persistent_ctx_t saved_regs = {0};
-static gpointer         saved_lr = NULL;
+static gpointer         persistent_loop = NULL;
 
 gboolean persistent_is_supported(void) {
 
@@ -216,17 +216,10 @@ static void instrument_persitent_restore_regs(GumArm64Writer   *cw,
 
 }
 
-static void instrument_exit(GumArm64Writer *cw) {
+static void instrument_afl_persistent_loop_func(void) {
 
-  gum_arm64_writer_put_mov_reg_reg(cw, ARM64_REG_X0, ARM64_REG_XZR);
-  gum_arm64_writer_put_call_address_with_arguments(
-      cw, GUM_ADDRESS(_exit), 1, GUM_ARG_REGISTER, ARM64_REG_X0);
-
-}
-
-static int instrument_afl_persistent_loop_func(void) {
+  if (__afl_persistent_loop(persistent_count) == 0) { _exit(0); }
 
-  int ret = __afl_persistent_loop(persistent_count);
   if (instrument_previous_pc_addr == NULL) {
 
     FATAL("instrument_previous_pc_addr uninitialized");
@@ -234,7 +227,6 @@ static int instrument_afl_persistent_loop_func(void) {
   }
 
   *instrument_previous_pc_addr = instrument_hash_zero;
-  return ret;
 
 }
 
@@ -284,7 +276,7 @@ static void instrument_persitent_save_lr(GumArm64Writer *cw) {
       GUM_INDEX_PRE_ADJUST);
 
   gum_arm64_writer_put_ldr_reg_address(cw, ARM64_REG_X0,
-                                       GUM_ADDRESS(&saved_lr));
+                                       GUM_ADDRESS(&persistent_ret));
 
   gum_arm64_writer_put_str_reg_reg_offset(cw, ARM64_REG_LR, ARM64_REG_X0, 0);
 
@@ -297,65 +289,35 @@ static void instrument_persitent_save_lr(GumArm64Writer *cw) {
 void persistent_prologue_arch(GumStalkerOutput *output) {
 
   /*
+   *  SAVE RET (Used to write the epilogue if persistent_ret is not set)
    *  SAVE REGS
-   *  SAVE RET
-   *  POP RET
-   * loop:
+   * loop: (Save address of where the eiplogue should jump back to)
    *  CALL instrument_afl_persistent_loop
-   *  TEST EAX, EAX
-   *  JZ end:
-   *  call hook (optionally)
+   *  CALL hook (optionally)
    *  RESTORE REGS
-   *  call original
-   *  jmp loop:
-   *
-   * end:
-   *  JMP SAVED RET
-   *
-   * original:
    *  INSTRUMENTED PERSISTENT FUNC
    */
 
   GumArm64Writer *cw = output->writer.arm64;
 
-  gconstpointer loop = cw->code + 1;
-
   FVERBOSE("Persistent loop reached");
 
+  if (persistent_ret == 0) { instrument_persitent_save_lr(cw); }
+
+  /* Save the current context */
   instrument_persitent_save_regs(cw, &saved_regs);
 
-  /* loop: */
-  gum_arm64_writer_put_label(cw, loop);
+  /* Store a pointer to where we should return for our next iteration */
+  persistent_loop = gum_arm64_writer_cur(cw);
 
-  /* call instrument_prologue_func */
+  /* call __afl_persistent_loop and _exit if zero. Also reset our previous_pc */
   instrument_afl_persistent_loop(cw);
 
-  /* jz done */
-  gconstpointer done = cw->code + 1;
-  gum_arm64_writer_put_cmp_reg_reg(cw, ARM64_REG_X0, ARM64_REG_XZR);
-  gum_arm64_writer_put_b_cond_label(cw, ARM64_CC_EQ, done);
-
   /* Optionally call the persistent hook */
   persistent_prologue_hook(cw, &saved_regs);
 
+  /* Restore our CPU context before we continue execution */
   instrument_persitent_restore_regs(cw, &saved_regs);
-  gconstpointer original = cw->code + 1;
-  /* call original */
-
-  gum_arm64_writer_put_bl_label(cw, original);
-
-  /* jmp loop */
-  gum_arm64_writer_put_b_label(cw, loop);
-
-  /* done: */
-  gum_arm64_writer_put_label(cw, done);
-
-  instrument_exit(cw);
-
-  /* original: */
-  gum_arm64_writer_put_label(cw, original);
-
-  instrument_persitent_save_lr(cw);
 
   if (persistent_debug) { gum_arm64_writer_put_brk_imm(cw, 0); }
 
@@ -368,7 +330,7 @@ void persistent_epilogue_arch(GumStalkerOutput *output) {
   if (persistent_debug) { gum_arm64_writer_put_brk_imm(cw, 0); }
 
   gum_arm64_writer_put_ldr_reg_address(cw, ARM64_REG_X0,
-                                       GUM_ADDRESS(&saved_lr));
+                                       GUM_ADDRESS(&persistent_loop));
 
   gum_arm64_writer_put_ldr_reg_reg_offset(cw, ARM64_REG_X0, ARM64_REG_X0, 0);
 
diff --git a/frida_mode/src/persistent/persistent_x86.c b/frida_mode/src/persistent/persistent_x86.c
index af21889c..8950223f 100644
--- a/frida_mode/src/persistent/persistent_x86.c
+++ b/frida_mode/src/persistent/persistent_x86.c
@@ -119,7 +119,7 @@ static void instrument_persitent_restore_regs(GumX86Writer     *cw,
 
 static void instrument_afl_persistent_loop_func(void) {
 
-  if (__afl_persistent_loop(persistent_count) == 0) { _exit(0);};
+  if (__afl_persistent_loop(persistent_count) == 0) { _exit(0); };
 
   if (instrument_previous_pc_addr == NULL) {
 
@@ -135,6 +135,7 @@ static void instrument_afl_persistent_loop(GumX86Writer *cw) {
 
   gum_x86_writer_put_call_address_with_arguments(
       cw, GUM_CALL_CAPI, GUM_ADDRESS(instrument_afl_persistent_loop_func), 0);
+
 }
 
 static void persistent_prologue_hook(GumX86Writer *cw, persistent_ctx_t *regs) {
@@ -167,7 +168,8 @@ static void instrument_persitent_save_ret(GumX86Writer *cw) {
   gum_x86_writer_put_push_reg(cw, GUM_X86_EAX);
   gum_x86_writer_put_push_reg(cw, GUM_X86_EBX);
 
-  gum_x86_writer_put_mov_reg_address(cw, GUM_X86_EAX, GUM_ADDRESS(&persistent_ret));
+  gum_x86_writer_put_mov_reg_address(cw, GUM_X86_EAX,
+                                     GUM_ADDRESS(&persistent_ret));
   gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_X86_EBX, GUM_X86_ESP,
                                             offset);
   gum_x86_writer_put_mov_reg_ptr_reg(cw, GUM_X86_EAX, GUM_X86_EBX);
@@ -218,6 +220,7 @@ void persistent_prologue_arch(GumStalkerOutput *output) {
   if (persistent_debug) { gum_x86_writer_put_breakpoint(cw); }
 
   /* The original instrumented code is emitted here. */
+
 }
 
 void persistent_epilogue_arch(GumStalkerOutput *output) {
@@ -227,10 +230,11 @@ void persistent_epilogue_arch(GumStalkerOutput *output) {
   if (persistent_debug) { gum_x86_writer_put_breakpoint(cw); }
 
   /* The stack should be aligned when we re-enter our loop */
-  gum_x86_writer_put_and_reg_u32 (cw, GUM_X86_ESP, 0xfffffff0);
-  gum_x86_writer_put_sub_reg_imm (cw, GUM_X86_ESP, 0x4);
+  gum_x86_writer_put_and_reg_u32(cw, GUM_X86_ESP, 0xfffffff0);
+  gum_x86_writer_put_sub_reg_imm(cw, GUM_X86_ESP, 0x4);
 
-  gum_x86_writer_put_mov_reg_address(cw, GUM_X86_EAX, GUM_ADDRESS(&persistent_loop));
+  gum_x86_writer_put_mov_reg_address(cw, GUM_X86_EAX,
+                                     GUM_ADDRESS(&persistent_loop));
   gum_x86_writer_put_jmp_reg_ptr(cw, GUM_X86_EAX);
 
 }