1 files changed, 67 insertions, 37 deletions

diff --git a/frida_mode/src/main.c b/frida_mode/src/main.c
index 7505c2f9..1ab9993f 100644
--- a/frida_mode/src/main.c
+++ b/frida_mode/src/main.c
@@ -1,3 +1,4 @@
+#include <errno.h>
 #include <unistd.h>
 #include <sys/types.h>
 
@@ -10,13 +11,21 @@
 #endif
 
 #include "frida-gum.h"
+
 #include "config.h"
 #include "debug.h"
 
-#include "interceptor.h"
+#include "entry.h"
 #include "instrument.h"
+#include "interceptor.h"
+#include "lib.h"
+#include "output.h"
+#include "persistent.h"
 #include "prefetch.h"
 #include "ranges.h"
+#include "stalker.h"
+#include "stats.h"
+#include "util.h"
 
 #ifdef __APPLE__
 extern mach_port_t mach_task_self();
@@ -30,16 +39,11 @@ extern int  __libc_start_main(int *(main)(int, char **, char **), int argc,
 
 typedef int *(*main_fn_t)(int argc, char **argv, char **envp);
 
-static main_fn_t      main_fn = NULL;
-static GumStalker *   stalker = NULL;
-static GumMemoryRange code_range = {0};
-
-extern void              __afl_manual_init();
-extern __thread uint64_t previous_pc;
+static main_fn_t main_fn = NULL;
 
-static int on_fork() {
+static int on_fork(void) {
 
-  prefetch_read(stalker);
+  prefetch_read();
   return fork();
 
 }
@@ -47,14 +51,20 @@ static int on_fork() {
 #ifdef __APPLE__
 static void on_main_os(int argc, char **argv, char **envp) {
 
+  UNUSED_PARAMETER(argc);
+  UNUSED_PARAMETER(argv);
+  UNUSED_PARAMETER(envp);
+
 }
 
 #else
 static void on_main_os(int argc, char **argv, char **envp) {
 
+  UNUSED_PARAMETER(argc);
   /* Personality doesn't affect the current process, it only takes effect on
    * evec */
   int persona = personality(ADDR_NO_RANDOMIZE);
+  if (persona == -1) { WARNF("Failed to set ADDR_NO_RANDOMIZE: %d", errno); }
   if ((persona & ADDR_NO_RANDOMIZE) == 0) { execvpe(argv[0], argv, envp); }
 
   GumInterceptor *interceptor = gum_interceptor_obtain();
@@ -68,39 +78,64 @@ static void on_main_os(int argc, char **argv, char **envp) {
 
 #endif
 
-static int *on_main(int argc, char **argv, char **envp) {
+static void embedded_init() {
 
-  on_main_os(argc, argv, envp);
+  static gboolean initialized = false;
+  if (!initialized) {
 
-  stalker = gum_stalker_new();
-  if (stalker == NULL) { FATAL("Failed to initialize stalker"); }
+    gum_init_embedded();
+    initialized = true;
 
-  gum_stalker_set_trust_threshold(stalker, 0);
+  }
+
+}
 
-  GumStalkerTransformer *transformer =
-      gum_stalker_transformer_make_from_callback(instr_basic_block, NULL, NULL);
+void afl_frida_start() {
 
+  embedded_init();
+  stalker_init();
+  lib_init();
+  entry_init();
   instrument_init();
+  output_init();
+  persistent_init();
   prefetch_init();
-  ranges_init(stalker);
+  ranges_init();
+  stats_init();
+
+  void *fork_addr =
+      GSIZE_TO_POINTER(gum_module_find_export_by_name(NULL, "fork"));
+  intercept(fork_addr, on_fork, NULL);
+
+  stalker_start();
+  entry_run();
+
+}
 
-  intercept(fork, on_fork, stalker);
+static int *on_main(int argc, char **argv, char **envp) {
+
+  on_main_os(argc, argv, envp);
 
-  gum_stalker_follow_me(stalker, transformer, NULL);
-  gum_stalker_deactivate(stalker);
+  unintercept_self();
 
-  __afl_manual_init();
+  afl_frida_start();
 
-  /* Child here */
-  previous_pc = 0;
-  prefetch_start(stalker);
-  main_fn(argc, argv, envp);
-  _exit(0);
+  return main_fn(argc, argv, envp);
 
 }
 
-#ifdef __APPLE__
-static void intercept_main() {
+#if defined(EMBEDDED)
+extern int *main(int argc, char **argv, char **envp);
+
+static void intercept_main(void) {
+
+  main_fn = main;
+  intercept(main, on_main, NULL);
+
+}
+
+#elif defined(__APPLE__)
+static void intercept_main(void) {
 
   mach_port_t task = mach_task_self();
   OKF("Task Id: %u", task);
@@ -119,13 +154,14 @@ static int on_libc_start_main(int *(main)(int, char **, char **), int argc,
                               void(*stack_end)) {
 
   main_fn = main;
+  unintercept_self();
   intercept(main, on_main, NULL);
   return __libc_start_main(main, argc, ubp_av, init, fini, rtld_fini,
                            stack_end);
 
 }
 
-static void intercept_main() {
+static void intercept_main(void) {
 
   intercept(__libc_start_main, on_libc_start_main, NULL);
 
@@ -133,15 +169,9 @@ static void intercept_main() {
 
 #endif
 
-__attribute__((constructor)) static void init() {
+__attribute__((constructor)) static void init(void) {
 
-  gum_init_embedded();
-  if (!gum_stalker_is_supported()) {
-
-    gum_deinit_embedded();
-    FATAL("Failed to initialize embedded");
-
-  }
+  embedded_init();
 
   intercept_main();