about summary refs log tree commit diff
diff options
context:
space:
mode:
-rw-r--r--frida_mode/GNUmakefile3
-rw-r--r--frida_mode/hook/frida_hook.c4
-rw-r--r--frida_mode/src/persistent/persistent_arm32.c70
-rw-r--r--frida_mode/src/persistent/persistent_arm64.c68
-rw-r--r--frida_mode/src/persistent/persistent_x64.c77
-rw-r--r--frida_mode/src/persistent/persistent_x86.c81
-rw-r--r--frida_mode/test/png/GNUmakefile7
7 files changed, 88 insertions, 222 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/hook/frida_hook.c b/frida_mode/hook/frida_hook.c
index da1a59b2..59a92e7e 100644
--- a/frida_mode/hook/frida_hook.c
+++ b/frida_mode/hook/frida_hook.c
@@ -31,8 +31,8 @@ __attribute__((visibility("default"))) void afl_persistent_hook(
   // do a length check matching the target!
 
   void **esp = (void **)regs->esp;
-  void  *arg1 = esp[0];
-  void **arg2 = &esp[1];
+  void  *arg1 = esp[1];
+  void **arg2 = &esp[2];
   memcpy(arg1, input_buf, input_buf_len);
   *arg2 = (void *)input_buf_len;
 
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_x64.c b/frida_mode/src/persistent/persistent_x64.c
index a8bed7be..fd194c75 100644
--- a/frida_mode/src/persistent/persistent_x64.c
+++ b/frida_mode/src/persistent/persistent_x64.c
@@ -17,7 +17,7 @@ typedef struct {
 } persistent_ctx_t;
 
 static persistent_ctx_t saved_regs = {0};
-static gpointer         saved_ret = NULL;
+static gpointer         persistent_loop = NULL;
 
 gboolean persistent_is_supported(void) {
 
@@ -162,17 +162,10 @@ static void instrument_persitent_restore_regs(GumX86Writer     *cw,
 
 }
 
-static void instrument_exit(GumX86Writer *cw) {
+static void instrument_afl_persistent_loop_func(void) {
 
-  gum_x86_writer_put_mov_reg_address(cw, GUM_X86_RAX, GUM_ADDRESS(_exit));
-  gum_x86_writer_put_mov_reg_u32(cw, GUM_X86_RDI, 0);
-  gum_x86_writer_put_call_reg(cw, GUM_X86_RAX);
+  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");
@@ -180,7 +173,6 @@ static int instrument_afl_persistent_loop_func(void) {
   }
 
   *instrument_previous_pc_addr = instrument_hash_zero;
-  return ret;
 
 }
 
@@ -190,7 +182,6 @@ static void instrument_afl_persistent_loop(GumX86Writer *cw) {
                                         -(GUM_RED_ZONE_SIZE));
   gum_x86_writer_put_call_address_with_arguments(
       cw, GUM_CALL_CAPI, GUM_ADDRESS(instrument_afl_persistent_loop_func), 0);
-  gum_x86_writer_put_test_reg_reg(cw, GUM_X86_RAX, GUM_X86_RAX);
 
   gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_X86_RSP, GUM_X86_RSP,
                                         (GUM_RED_ZONE_SIZE));
@@ -235,7 +226,8 @@ static void instrument_persitent_save_ret(GumX86Writer *cw) {
   gum_x86_writer_put_push_reg(cw, GUM_X86_RAX);
   gum_x86_writer_put_push_reg(cw, GUM_X86_RBX);
 
-  gum_x86_writer_put_mov_reg_address(cw, GUM_X86_RAX, GUM_ADDRESS(&saved_ret));
+  gum_x86_writer_put_mov_reg_address(cw, GUM_X86_RAX,
+                                     GUM_ADDRESS(&persistent_ret));
   gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_X86_RBX, GUM_X86_RSP,
                                             offset);
   gum_x86_writer_put_mov_reg_ptr_reg(cw, GUM_X86_RAX, GUM_X86_RBX);
@@ -252,70 +244,44 @@ static void instrument_persitent_save_ret(GumX86Writer *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
    */
 
   GumX86Writer *cw = output->writer.x86;
 
-  gconstpointer loop = cw->code + 1;
-
   FVERBOSE("Persistent loop reached");
 
-  /* Pop the return value */
-  gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_X86_RSP, GUM_X86_RSP, 8);
+  /*
+   * If we haven't set persistent_ret, then assume that we are dealing with a
+   * function and we should loop when that function returns.
+   */
+  if (persistent_ret == 0) { instrument_persitent_save_ret(cw); }
 
+  /* Save the current context */
   instrument_persitent_save_regs(cw, &saved_regs);
 
-  /* loop: */
-  gum_x86_writer_put_label(cw, loop);
+  /* Store a pointer to where we should return for our next iteration */
+  persistent_loop = gum_x86_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_x86_writer_put_jcc_near_label(cw, X86_INS_JE, done, GUM_UNLIKELY);
-
   /* 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_x86_writer_put_call_near_label(cw, original);
-
-  /* jmp loop */
-  gum_x86_writer_put_jmp_near_label(cw, loop);
-
-  /* done: */
-  gum_x86_writer_put_label(cw, done);
-
-  instrument_exit(cw);
-
-  /* original: */
-  gum_x86_writer_put_label(cw, original);
-
-  instrument_persitent_save_ret(cw);
 
   if (persistent_debug) { gum_x86_writer_put_breakpoint(cw); }
 
+  /* The original instrumented code is emitted here. */
+
 }
 
 void persistent_epilogue_arch(GumStalkerOutput *output) {
@@ -331,7 +297,8 @@ void persistent_epilogue_arch(GumStalkerOutput *output) {
   gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_X86_RSP, GUM_X86_RSP, -8);
   gum_x86_writer_put_label(cw, zero);
 
-  gum_x86_writer_put_mov_reg_address(cw, GUM_X86_RAX, GUM_ADDRESS(&saved_ret));
+  gum_x86_writer_put_mov_reg_address(cw, GUM_X86_RAX,
+                                     GUM_ADDRESS(&persistent_loop));
   gum_x86_writer_put_jmp_reg_ptr(cw, GUM_X86_RAX);
 
 }
diff --git a/frida_mode/src/persistent/persistent_x86.c b/frida_mode/src/persistent/persistent_x86.c
index e9bde3d2..8950223f 100644
--- a/frida_mode/src/persistent/persistent_x86.c
+++ b/frida_mode/src/persistent/persistent_x86.c
@@ -16,8 +16,7 @@ typedef struct {
 } persistent_ctx_t;
 
 static persistent_ctx_t saved_regs = {0};
-
-static gpointer saved_ret = NULL;
+static gpointer         persistent_loop = NULL;
 
 gboolean persistent_is_supported(void) {
 
@@ -118,18 +117,10 @@ static void instrument_persitent_restore_regs(GumX86Writer     *cw,
 
 }
 
-static void instrument_exit(GumX86Writer *cw) {
-
-  gum_x86_writer_put_mov_reg_address(cw, GUM_X86_EAX, GUM_ADDRESS(_exit));
-  gum_x86_writer_put_mov_reg_u32(cw, GUM_X86_EDI, 0);
-  gum_x86_writer_put_push_reg(cw, GUM_X86_EDI);
-  gum_x86_writer_put_call_reg(cw, GUM_X86_EAX);
-
-}
+static void instrument_afl_persistent_loop_func(void) {
 
-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");
@@ -137,7 +128,6 @@ static int instrument_afl_persistent_loop_func(void) {
   }
 
   *instrument_previous_pc_addr = instrument_hash_zero;
-  return ret;
 
 }
 
@@ -145,7 +135,6 @@ 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);
-  gum_x86_writer_put_test_reg_reg(cw, GUM_X86_EAX, GUM_X86_EAX);
 
 }
 
@@ -179,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(&saved_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);
@@ -193,68 +183,44 @@ static void instrument_persitent_save_ret(GumX86Writer *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
    */
 
   GumX86Writer *cw = output->writer.x86;
 
-  gconstpointer loop = cw->code + 1;
-
   FVERBOSE("Persistent loop reached");
 
-  /* Pop the return value */
-  gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_X86_ESP, GUM_X86_ESP, 4);
+  /*
+   * If we haven't set persistent_ret, then assume that we are dealing with a
+   * function and we should loop when that function returns.
+   */
+  if (persistent_ret == 0) { instrument_persitent_save_ret(cw); }
 
+  /* Save the current context */
   instrument_persitent_save_regs(cw, &saved_regs);
 
-  /* loop: */
-  gum_x86_writer_put_label(cw, loop);
+  /* Store a pointer to where we should return for our next iteration */
+  persistent_loop = gum_x86_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_x86_writer_put_jcc_near_label(cw, X86_INS_JE, done, GUM_UNLIKELY);
-
   /* 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_x86_writer_put_call_near_label(cw, original);
-  /* jmp loop */
-  gum_x86_writer_put_jmp_near_label(cw, loop);
-
-  /* done: */
-  gum_x86_writer_put_label(cw, done);
-
-  instrument_exit(cw);
-
-  /* original: */
-  gum_x86_writer_put_label(cw, original);
-
-  instrument_persitent_save_ret(cw);
 
   if (persistent_debug) { gum_x86_writer_put_breakpoint(cw); }
 
+  /* The original instrumented code is emitted here. */
+
 }
 
 void persistent_epilogue_arch(GumStalkerOutput *output) {
@@ -263,7 +229,12 @@ void persistent_epilogue_arch(GumStalkerOutput *output) {
 
   if (persistent_debug) { gum_x86_writer_put_breakpoint(cw); }
 
-  gum_x86_writer_put_mov_reg_address(cw, GUM_X86_EAX, GUM_ADDRESS(&saved_ret));
+  /* 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_mov_reg_address(cw, GUM_X86_EAX,
+                                     GUM_ADDRESS(&persistent_loop));
   gum_x86_writer_put_jmp_reg_ptr(cw, GUM_X86_EAX);
 
 }
diff --git a/frida_mode/test/png/GNUmakefile b/frida_mode/test/png/GNUmakefile
index eccc66f6..4a023a5d 100644
--- a/frida_mode/test/png/GNUmakefile
+++ b/frida_mode/test/png/GNUmakefile
@@ -48,7 +48,7 @@ all: $(TEST_BIN)
 	CFLAGS="-m32" LDFLAGS="-m32" make $(TEST_BIN)
 
 arm:
-	CFLAGS="-marm" LDFLAGS="-marm" CC="arm-linux-gnueabihf-gcc" CXX="arm-linux-gnueabihf-g++" make $(TEST_BIN)
+	CFLAGS="-marm" LDFLAGS="-marm" CC="arm-linux-gnueabihf-gcc" CXX="arm-linux-gnueabihf-g++" PNG_ARCH="arm" make $(TEST_BIN)
 
 $(BUILD_DIR):
 	mkdir -p $@
@@ -93,8 +93,7 @@ $(LIBZ_PC): | $(LIBZ_DIR)
 	cd $(LIBZ_DIR) && \
 	CFLAGS="$(CFLAGS) -fPIC" \
 		./configure \
-			--static \
-			--archs="$(ARCH)"
+			--static 
 
 $(LIBZ_LIB): | $(LIBZ_PC)
 	CFLAGS="$(CFLAGS) -fPIC" \
@@ -120,7 +119,7 @@ $(LIBPNG_MAKEFILE): $(LIBZ_LIB) | $(LIBPNG_DIR)
 		CFLAGS="$(CFLAGS) -I$(LIBZ_DIR)" \
 		LDFLAGS="-L$(LIBZ_DIR)" \
 			./configure \
-				--host="$(ARCH)"
+				--host="$(PNG_ARCH)"
 
 $(LIBPNG_LIB): $(LIBPNG_MAKEFILE)
 	CFLAGS="$(CFLAGS) -I$(LIBZ_DIR)" \