5 files changed, 64 insertions, 30 deletions
diff --git a/frida_mode/src/entry.c b/frida_mode/src/entry.c
new file mode 100644
index 00000000..e71386a0
--- /dev/null
+++ b/frida_mode/src/entry.c
@@ -0,0 +1,50 @@
+#include "frida-gum.h"
+
+#include "debug.h"
+
+#include "entry.h"
+#include "instrument.h"
+#include "stalker.h"
+#include "util.h"
+
+extern void __afl_manual_init();
+
+guint64 entry_start = 0;
+
+static void entry_launch(void) {
+
+  __afl_manual_init();
+
+  /* Child here */
+  previous_pc = 0;
+
+}
+
+void entry_init(void) {
+
+  entry_start = util_read_address("AFL_ENTRYPOINT");
+  OKF("entry_point: 0x%016" G_GINT64_MODIFIER "X", entry_start);
+
+}
+
+void entry_run(void) {
+
+  if (entry_start == 0) { entry_launch(); }
+
+}
+
+static void entry_callout(GumCpuContext *cpu_context, gpointer user_data) {
+
+  UNUSED_PARAMETER(cpu_context);
+  UNUSED_PARAMETER(user_data);
+  entry_launch();
+
+}
+
+void entry_prologue(GumStalkerIterator *iterator, GumStalkerOutput *output) {
+
+  UNUSED_PARAMETER(output);
+  gum_stalker_iterator_put_callout(iterator, entry_callout, NULL, NULL);
+
+}
+
diff --git a/frida_mode/src/instrument/instrument.c b/frida_mode/src/instrument/instrument.c
index d93f37c7..971f80c0 100644
--- a/frida_mode/src/instrument/instrument.c
+++ b/frida_mode/src/instrument/instrument.c
@@ -5,6 +5,7 @@
 #include "config.h"
 #include "debug.h"
 
+#include "entry.h"
 #include "frida_cmplog.h"
 #include "instrument.h"
 #include "persistent.h"
@@ -79,6 +80,7 @@ static void instr_basic_block(GumStalkerIterator *iterator,
   gboolean       begin = TRUE;
   while (gum_stalker_iterator_next(iterator, &instr)) {
 
+    if (instr->address == entry_start) { entry_prologue(iterator, output); }
     if (instr->address == persistent_start) { persistent_prologue(output); }
 
     if (begin) {
diff --git a/frida_mode/src/main.c b/frida_mode/src/main.c
index 5c64d192..e031dbed 100644
--- a/frida_mode/src/main.c
+++ b/frida_mode/src/main.c
@@ -14,6 +14,7 @@
 #include "config.h"
 #include "debug.h"
 
+#include "entry.h"
 #include "instrument.h"
 #include "interceptor.h"
 #include "lib.h"
@@ -37,8 +38,6 @@ typedef int *(*main_fn_t)(int argc, char **argv, char **envp);
 
 static main_fn_t main_fn = NULL;
 
-extern void __afl_manual_init();
-
 static int on_fork(void) {
 
   prefetch_read();
@@ -79,6 +78,7 @@ static void on_main_os(int argc, char **argv, char **envp) {
 static int *on_main(int argc, char **argv, char **envp) {
 
   void *fork_addr;
+
   on_main_os(argc, argv, envp);
 
   unintercept_self();
@@ -86,6 +86,7 @@ static int *on_main(int argc, char **argv, char **envp) {
   stalker_init();
 
   lib_init();
+  entry_init();
   instrument_init();
   persistent_init();
   prefetch_init();
@@ -95,13 +96,8 @@ static int *on_main(int argc, char **argv, char **envp) {
   intercept(fork_addr, on_fork, NULL);
 
   stalker_start();
-  stalker_pause();
-
-  __afl_manual_init();
+  entry_run();
 
-  /* Child here */
-  previous_pc = 0;
-  stalker_resume();
   return main_fn(argc, argv, envp);
 
 }
diff --git a/frida_mode/src/persistent/persistent_x64.c b/frida_mode/src/persistent/persistent_x64.c
index 5b8493b2..49f1988c 100644
--- a/frida_mode/src/persistent/persistent_x64.c
+++ b/frida_mode/src/persistent/persistent_x64.c
@@ -252,15 +252,20 @@ static void persistent_prologue_hook(GumX86Writer *      cw,
                                         -(GUM_RED_ZONE_SIZE));
 
   gum_x86_writer_put_mov_reg_address(cw, GUM_REG_RCX,
-                                     GUM_ADDRESS(__afl_fuzz_len));
+                                     GUM_ADDRESS(&__afl_fuzz_len));
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_RCX, GUM_REG_RCX, 0);
   gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_RCX, GUM_REG_RCX, 0);
   gum_x86_writer_put_mov_reg_u64(cw, GUM_REG_RDI, 0xffffffff);
   gum_x86_writer_put_and_reg_reg(cw, GUM_REG_RCX, GUM_REG_RDI);
 
+  gum_x86_writer_put_mov_reg_address(cw, GUM_REG_RDX,
+                                     GUM_ADDRESS(&__afl_fuzz_ptr));
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_RDX, GUM_REG_RDX, 0);
+
   gum_x86_writer_put_call_address_with_arguments(
       cw, GUM_CALL_CAPI, GUM_ADDRESS(hook), 4, GUM_ARG_ADDRESS,
-      GUM_ADDRESS(regs), GUM_ARG_ADDRESS, GUM_ADDRESS(0), GUM_ARG_ADDRESS,
-      GUM_ADDRESS(__afl_fuzz_ptr), GUM_ARG_REGISTER, GUM_REG_RCX);
+      GUM_ADDRESS(regs), GUM_ARG_ADDRESS, GUM_ADDRESS(0), GUM_ARG_REGISTER,
+      GUM_REG_RDX, GUM_ARG_REGISTER, GUM_REG_RCX);
 
   gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_REG_RSP, GUM_REG_RSP,
                                         (GUM_RED_ZONE_SIZE));
diff --git a/frida_mode/src/stalker.c b/frida_mode/src/stalker.c
index 5ee519ba..81973e9c 100644
--- a/frida_mode/src/stalker.c
+++ b/frida_mode/src/stalker.c
@@ -21,12 +21,6 @@ GumStalker *stalker_get(void) {
 
 }
 
-__attribute__((noinline)) static void stalker_activation(void) {
-
-  asm volatile("");
-
-}
-
 void stalker_start(void) {
 
   GumStalkerTransformer *transformer = instrument_get_transformer();
@@ -34,16 +28,3 @@ void stalker_start(void) {
 
 }
 
-void stalker_pause(void) {
-
-  gum_stalker_deactivate(stalker);
-
-}
-
-void stalker_resume(void) {
-
-  gum_stalker_activate(stalker, stalker_activation);
-  stalker_activation();
-
-}
-