Support for AFL_FRIDA_PERSISTENT_RET (#941)

Co-authored-by: Your Name <you@example.com>

7 files changed, 88 insertions, 8 deletions

diff --git a/frida_mode/src/instrument/instrument.c b/frida_mode/src/instrument/instrument.c
index f21849a6..c4f18797 100644
--- a/frida_mode/src/instrument/instrument.c
+++ b/frida_mode/src/instrument/instrument.c
@@ -85,6 +85,7 @@ static void instr_basic_block(GumStalkerIterator *iterator,
 
     if (instr->address == entry_start) { entry_prologue(iterator, output); }
     if (instr->address == persistent_start) { persistent_prologue(output); }
+    if (instr->address == persistent_ret) { persistent_epilogue(output); }
 
     /*
      * Until we reach AFL_ENTRYPOINT (assumed to be main if not specified) or
diff --git a/frida_mode/src/persistent/persistent.c b/frida_mode/src/persistent/persistent.c
index 918ff153..2ec5b9cc 100644
--- a/frida_mode/src/persistent/persistent.c
+++ b/frida_mode/src/persistent/persistent.c
@@ -12,6 +12,9 @@ int                    __afl_sharedmem_fuzzing = 0;
 afl_persistent_hook_fn hook = NULL;
 guint64                persistent_start = 0;
 guint64                persistent_count = 0;
+guint64                persistent_ret = 0;
+guint64                persistent_ret_offset = 0;
+gboolean               persistent_debug = FALSE;
 
 void persistent_init(void) {
 
@@ -19,12 +22,36 @@ void persistent_init(void) {
 
   persistent_start = util_read_address("AFL_FRIDA_PERSISTENT_ADDR");
   persistent_count = util_read_num("AFL_FRIDA_PERSISTENT_CNT");
+  persistent_ret = util_read_address("AFL_FRIDA_PERSISTENT_RET");
+  persistent_ret_offset =
+      util_read_address("AFL_FRIDA_PERSISTENT_RETADDR_OFFSET");
+
+  if (getenv("AFL_FRIDA_PERSISTENT_DEBUG") != NULL) { persistent_debug = TRUE; }
+
+  if (persistent_count != 0 && persistent_start == 0) {
 
-  if (persistent_count != 0 && persistent_start == 0)
     FATAL(
         "AFL_FRIDA_PERSISTENT_ADDR must be specified if "
         "AFL_FRIDA_PERSISTENT_CNT is");
 
+  }
+
+  if (persistent_ret != 0 && persistent_start == 0) {
+
+    FATAL(
+        "AFL_FRIDA_PERSISTENT_ADDR must be specified if "
+        "AFL_FRIDA_PERSISTENT_RET is");
+
+  }
+
+  if (persistent_ret_offset != 0 && persistent_ret == 0) {
+
+    FATAL(
+        "AFL_FRIDA_PERSISTENT_RET must be specified if "
+        "AFL_FRIDA_PERSISTENT_RETADDR_OFFSET is");
+
+  }
+
   if (persistent_start != 0 && persistent_count == 0) persistent_count = 1000;
 
   if (persistent_count != 0 && persistent_count < 100)
@@ -39,6 +66,11 @@ void persistent_init(void) {
       persistent_start == 0 ? ' ' : 'X', persistent_count);
   OKF("Instrumentation - hook [%s]", hook_name);
 
+  OKF("Instrumentation - persistent ret [%c] (0x%016" G_GINT64_MODIFIER "X)",
+      persistent_ret == 0 ? ' ' : 'X', persistent_ret);
+  OKF("Instrumentation - persistent ret offset [%c] (%" G_GINT64_MODIFIER "d)",
+      persistent_ret_offset == 0 ? ' ' : 'X', persistent_ret_offset);
+
   if (hook_name != NULL) {
 
     void *hook_obj = dlopen(hook_name, RTLD_NOW);
diff --git a/frida_mode/src/persistent/persistent_arm32.c b/frida_mode/src/persistent/persistent_arm32.c
index bc021ff3..6a3c06fa 100644
--- a/frida_mode/src/persistent/persistent_arm32.c
+++ b/frida_mode/src/persistent/persistent_arm32.c
@@ -68,5 +68,12 @@ void persistent_prologue(GumStalkerOutput *output) {
 
 }
 
+void persistent_epilogue(GumStalkerOutput *output) {
+
+  UNUSED_PARAMETER(output);
+  FATAL("Persistent mode not supported on this architecture");
+
+}
+
 #endif
 
diff --git a/frida_mode/src/persistent/persistent_arm64.c b/frida_mode/src/persistent/persistent_arm64.c
index c198da69..1215d8da 100644
--- a/frida_mode/src/persistent/persistent_arm64.c
+++ b/frida_mode/src/persistent/persistent_arm64.c
@@ -111,5 +111,12 @@ void persistent_prologue(GumStalkerOutput *output) {
 
 }
 
+void persistent_epilogue(GumStalkerOutput *output) {
+
+  UNUSED_PARAMETER(output);
+  FATAL("Persistent mode not supported on this architecture");
+
+}
+
 #endif
 
diff --git a/frida_mode/src/persistent/persistent_x64.c b/frida_mode/src/persistent/persistent_x64.c
index aa772b7f..4c495d47 100644
--- a/frida_mode/src/persistent/persistent_x64.c
+++ b/frida_mode/src/persistent/persistent_x64.c
@@ -1,9 +1,11 @@
 #include "frida-gum.h"
 
 #include "config.h"
+#include "debug.h"
 
 #include "instrument.h"
 #include "persistent.h"
+#include "util.h"
 
 #if defined(__x86_64__)
 
@@ -264,7 +266,6 @@ void persistent_prologue(GumStalkerOutput *output) {
   GumX86Writer *cw = output->writer.x86;
 
   gconstpointer loop = cw->code + 1;
-  // gum_x86_writer_put_breakpoint(cw);
 
   /* Stack must be 16-byte aligned per ABI */
   instrument_persitent_save_regs(cw, &saved_regs);
@@ -288,7 +289,9 @@ void persistent_prologue(GumStalkerOutput *output) {
   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);
 
@@ -300,9 +303,23 @@ void persistent_prologue(GumStalkerOutput *output) {
   /* original: */
   gum_x86_writer_put_label(cw, original);
 
+  if (persistent_debug) { gum_x86_writer_put_breakpoint(cw); }
+
   gum_x86_writer_flush(cw);
 
 }
 
+void persistent_epilogue(GumStalkerOutput *output) {
+
+  GumX86Writer *cw = output->writer.x86;
+
+  if (persistent_debug) { gum_x86_writer_put_breakpoint(cw); }
+
+  gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_REG_RSP, GUM_REG_RSP,
+                                        persistent_ret_offset);
+  gum_x86_writer_put_ret(cw);
+
+}
+
 #endif
 
diff --git a/frida_mode/src/persistent/persistent_x86.c b/frida_mode/src/persistent/persistent_x86.c
index 20a3dc42..b30dfadf 100644
--- a/frida_mode/src/persistent/persistent_x86.c
+++ b/frida_mode/src/persistent/persistent_x86.c
@@ -244,9 +244,24 @@ void persistent_prologue(GumStalkerOutput *output) {
   /* original: */
   gum_x86_writer_put_label(cw, original);
 
+  if (persistent_debug) { gum_x86_writer_put_breakpoint(cw); }
+
   gum_x86_writer_flush(cw);
 
 }
 
+void persistent_epilogue(GumStalkerOutput *output) {
+
+  GumX86Writer *cw = output->writer.x86;
+
+  if (persistent_debug) { gum_x86_writer_put_breakpoint(cw); }
+
+  gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_REG_ESP, GUM_REG_ESP,
+                                        persistent_ret_offset);
+
+  gum_x86_writer_put_ret(cw);
+
+}
+
 #endif
 
diff --git a/frida_mode/src/util.c b/frida_mode/src/util.c
index 86b94970..09e8a58b 100644
--- a/frida_mode/src/util.c
+++ b/frida_mode/src/util.c
@@ -10,7 +10,7 @@ guint64 util_read_address(char *key) {
 
   if (!g_str_has_prefix(value_str, "0x")) {
 
-    FATAL("Invalid address should have 0x prefix: %s\n", value_str);
+    FATAL("Invalid address should have 0x prefix: %s=%s\n", key, value_str);
 
   }
 
@@ -20,8 +20,8 @@ guint64 util_read_address(char *key) {
 
     if (!g_ascii_isxdigit(*c)) {
 
-      FATAL("Invalid address not formed of hex digits: %s ('%c')\n", value_str,
-            *c);
+      FATAL("Invalid address not formed of hex digits: %s=%s ('%c')\n", key,
+            value_str, *c);
 
     }
 
@@ -30,7 +30,7 @@ guint64 util_read_address(char *key) {
   guint64 value = g_ascii_strtoull(value_str2, NULL, 16);
   if (value == 0) {
 
-    FATAL("Invalid address failed hex conversion: %s\n", value_str2);
+    FATAL("Invalid address failed hex conversion: %s=%s\n", key, value_str2);
 
   }
 
@@ -48,7 +48,8 @@ guint64 util_read_num(char *key) {
 
     if (!g_ascii_isdigit(*c)) {
 
-      FATAL("Invalid address not formed of decimal digits: %s\n", value_str);
+      FATAL("Invalid address not formed of decimal digits: %s=%s\n", key,
+            value_str);
 
     }
 
@@ -57,7 +58,7 @@ guint64 util_read_num(char *key) {
   guint64 value = g_ascii_strtoull(value_str, NULL, 10);
   if (value == 0) {
 
-    FATAL("Invalid address failed numeric conversion: %s\n", value_str);
+    FATAL("Invalid address failed numeric conversion: %s=%s\n", key, value_str);
 
   }