Merge pull request #1759 from AFLplusplus/dev

Dev

1 files changed, 36 insertions, 5 deletions

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) {