4 files changed, 37 insertions, 33 deletions
diff --git a/frida_mode/include/entry.h b/frida_mode/include/entry.h
index edc41467..949fab71 100644
--- a/frida_mode/include/entry.h
+++ b/frida_mode/include/entry.h
@@ -14,8 +14,6 @@ void entry_init(void);
 
 void entry_start(void);
 
-void entry_prologue(GumStalkerIterator *iterator, GumStalkerOutput *output);
-
 void entry_on_fork(void);
 
 #endif
diff --git a/frida_mode/src/entry.c b/frida_mode/src/entry.c
index 05af7ebb..fc49b4d7 100644
--- a/frida_mode/src/entry.c
+++ b/frida_mode/src/entry.c
@@ -78,31 +78,7 @@ void entry_init(void) {
 
 void entry_start(void) {
 
-  if (persistent_start == 0) {
-
-    ranges_exclude();
-    stalker_trust();
-
-  }
-
-  if (entry_point == 0) { entry_launch(); }
-
-}
-
-static void entry_callout(GumCpuContext *cpu_context, gpointer user_data) {
-
-  UNUSED_PARAMETER(cpu_context);
-  UNUSED_PARAMETER(user_data);
-  entry_compiled = TRUE;
-  entry_launch();
-
-}
-
-void entry_prologue(GumStalkerIterator *iterator, GumStalkerOutput *output) {
-
-  UNUSED_PARAMETER(output);
   FVERBOSE("AFL_ENTRYPOINT reached");
-
   if (persistent_start == 0) {
 
     ranges_exclude();
@@ -110,7 +86,7 @@ void entry_prologue(GumStalkerIterator *iterator, GumStalkerOutput *output) {
 
   }
 
-  gum_stalker_iterator_put_callout(iterator, entry_callout, NULL, NULL);
+  entry_launch();
 
 }
 
diff --git a/frida_mode/src/instrument/instrument.c b/frida_mode/src/instrument/instrument.c
index a6aac666..db73d845 100644
--- a/frida_mode/src/instrument/instrument.c
+++ b/frida_mode/src/instrument/instrument.c
@@ -169,7 +169,6 @@ static void instrument_basic_block(GumStalkerIterator *iterator,
 
     if (unlikely(begin)) { instrument_debug_start(instr->address, output); }
 
-    if (instr->address == entry_point) { entry_prologue(iterator, output); }
     if (instr->address == persistent_start) { persistent_prologue(output); }
     if (instr->address == persistent_ret) { persistent_epilogue(output); }
 
diff --git a/frida_mode/src/main.c b/frida_mode/src/main.c
index c8c50b37..f11c4b25 100644
--- a/frida_mode/src/main.c
+++ b/frida_mode/src/main.c
@@ -197,7 +197,7 @@ static void afl_print_env(void) {
 
 }
 
-__attribute__((visibility("default"))) void afl_frida_start(void) {
+void afl_frida_config(void) {
 
   FOKF(cRED "**********************");
   FOKF(cRED "* " cYEL "******************" cRED " *");
@@ -225,9 +225,7 @@ __attribute__((visibility("default"))) void afl_frida_start(void) {
 
   js_start();
 
-  /* Initialize */
   output_init();
-
   embedded_init();
   entry_init();
   instrument_init();
@@ -240,12 +238,35 @@ __attribute__((visibility("default"))) void afl_frida_start(void) {
   ranges_init();
   stats_init();
 
-  /* Start */
+}
+
+void afl_frida_run(void) {
+
   stalker_start();
   entry_start();
 
 }
 
+__attribute__((visibility("default"))) void afl_frida_start(void) {
+
+  afl_frida_config();
+  afl_frida_run();
+
+}
+
+typedef void *(*entry_func_t)(size_t a1, size_t a2, size_t a3, size_t a4,
+                              size_t a5, size_t a6);
+
+static void *on_entry(size_t a1, size_t a2, size_t a3, size_t a4, size_t a5,
+                      size_t a6) {
+
+  intercept_unhook(GSIZE_TO_POINTER(entry_point));
+  afl_frida_run();
+  entry_func_t entry = (entry_func_t)entry_point;
+  return entry(a1, a2, a3, a4, a5, a6);
+
+}
+
 static int on_main(int argc, char **argv, char **envp) {
 
   int ret;
@@ -254,7 +275,17 @@ static int on_main(int argc, char **argv, char **envp) {
 
   intercept_unhook_self();
 
-  afl_frida_start();
+  afl_frida_config();
+
+  if (entry_point == 0) {
+
+    afl_frida_run();
+
+  } else {
+
+    intercept_hook(GSIZE_TO_POINTER(entry_point), on_entry, NULL);
+
+  }
 
   if (js_main_hook != NULL) {