Improve JS bindings for hooking functions

10 files changed, 337 insertions, 119 deletions

diff --git a/frida_mode/Scripting.md b/frida_mode/Scripting.md
index 691b03d1..2ee0c858 100644
--- a/frida_mode/Scripting.md
+++ b/frida_mode/Scripting.md
@@ -246,7 +246,7 @@ FRIDA mode supports the replacement of any function, with an implementation
 generated by CModule. This allows for a bespoke harness to be written as
 follows:
 
-```
+```js
 const slow = DebugSymbol.fromName('slow').address;
 Afl.print(`slow: ${slow}`);
 
@@ -281,13 +281,90 @@ Afl.done();
 Here, we replace the function `slow` with our own code. This code is then
 selected as the entry point as well as the persistent loop address.
 
-**WARNING** There are two key limitations in replacing a function in this way:
-- The function which is to be replaced must not be `main` this is because this
-is the point at which FRIDA mode is initialized and at the point the the JS has
-been run, the start of the `main` function has already been instrumented and
-cached.
-- The replacement function must not call itself. e.g. in this example we
-couldn't replace `LLVMFuzzerTestOneInput` and call itself.
+## Replacing LLVMFuzzerTestOneInput
+The function `LLVMFuzzerTestOneInput` can be replaced just like any other. Also
+any replaced function can also call itself. In the example below, we replace
+`LLVMFuzzerTestOneInput` with `My_LLVMFuzzerTestOneInput` which ignores the
+parameters `buf` and `len` and then calls the original `LLVMFuzzerTestOneInput`
+with the paramaters `__afl_fuzz_ptr` and `__afl_fuzz_len`. This allows us to
+carry out in-memory fuzzing without the need for any hook function. It should be
+noted that the replacement function and the original can *NOT* share the same
+name, since otherwise the `C` code in the `CModule` will not compile due to a
+symbol name collision.
+
+```js
+const LLVMFuzzerTestOneInput = DebugSymbol.fromName('LLVMFuzzerTestOneInput').address;
+Afl.print(`LLVMFuzzerTestOneInput: ${LLVMFuzzerTestOneInput}`);
+
+const cm = new CModule(`
+
+    extern unsigned char * __afl_fuzz_ptr;
+    extern unsigned int * __afl_fuzz_len;
+    extern void LLVMFuzzerTestOneInput(char *buf, int len);
+
+    void My_LLVMFuzzerTestOneInput(char *buf, int len) {
+
+      LLVMFuzzerTestOneInput(__afl_fuzz_ptr, *__afl_fuzz_len);
+
+    }
+    `,
+    {
+        LLVMFuzzerTestOneInput: LLVMFuzzerTestOneInput,
+        __afl_fuzz_ptr: Afl.getAflFuzzPtr(),
+        __afl_fuzz_len: Afl.getAflFuzzLen()
+    });
+
+Afl.setEntryPoint(cm.My_LLVMFuzzerTestOneInput);
+Afl.setPersistentAddress(cm.My_LLVMFuzzerTestOneInput);
+Afl.setInMemoryFuzzing();
+Interceptor.replace(LLVMFuzzerTestOneInput, cm.My_LLVMFuzzerTestOneInput);
+```
+
+## Hooking `main`
+Lastly, it should be noted that using FRIDA mode's scripting support to hook
+the `main` function is a special case. This is because the `main` function is
+already hooked by the FRIDA mode engine itself and hence the function `main` (or
+at least the first basic block already been compiled by Stalker ready for
+execution). Hence any attempt to use `Interceptor.replace` like in the example
+above will not work. Instead the JS bindings provide a function `setJsMainHook`
+for just this scenario as demonstrated in the example below.
+
+```js
+const main = DebugSymbol.fromName('main').address;
+Afl.print(`main: ${main}`);
+
+const LLVMFuzzerTestOneInput = DebugSymbol.fromName('LLVMFuzzerTestOneInput').address;
+Afl.print(`LLVMFuzzerTestOneInput: ${LLVMFuzzerTestOneInput}`);
+
+const cm = new CModule(`
+
+    extern unsigned char * __afl_fuzz_ptr;
+    extern unsigned int * __afl_fuzz_len;
+    extern void LLVMFuzzerTestOneInput(char *buf, int len);
+
+    int main(int argc, char **argv)  {
+
+      LLVMFuzzerTestOneInput(__afl_fuzz_ptr, *__afl_fuzz_len);
+
+    }
+    `,
+    {
+        LLVMFuzzerTestOneInput: LLVMFuzzerTestOneInput,
+        __afl_fuzz_ptr: Afl.getAflFuzzPtr(),
+        __afl_fuzz_len: Afl.getAflFuzzLen()
+    });
+
+Afl.setEntryPoint(cm.main);
+Afl.setPersistentAddress(cm.main);
+Afl.setInMemoryFuzzing();
+Afl.setJsMainHook(cm.main);
+```
+## Library Fuzzing
+
+It doesn't take too much imagination to see that the above example can be
+extended to use FRIDA's `Module.load` API so that the replaced `main` function
+can then call an arbitrary function. In this way, if we have a library which we
+wish to fuzz rather than an execuatble, then a surrogate executable can be used.
 
 # Patching
 Consider the [following](test/js/test2.c) test code...
@@ -620,41 +697,31 @@ value of the `-t` flag passed to `afl-fuzz`.
 # API
 ```js
 class Afl {
-
-  /**
-   * Field containing the `Module` object for `afl-frida-trace.so` (the FRIDA mode
-   * implementation).
-   */
-  public static module: Module = Process.getModuleByName("afl-frida-trace.so");
-
   /**
    * This is equivalent to setting a value in `AFL_FRIDA_EXCLUDE_RANGES`,
    * it takes as arguments a `NativePointer` and a `number`. It can be
    * called multiple times to exclude several ranges.
    */
-  public static addExcludedRange(addressess: NativePointer, size: number): void {
-    Afl.jsApiAddExcludeRange(addressess, size);
+  static addExcludedRange(addressess, size) {
+      Afl.jsApiAddExcludeRange(addressess, size);
   }
-
   /**
    * This is equivalent to setting a value in `AFL_FRIDA_INST_RANGES`,
    * it takes as arguments a `NativePointer` and a `number`. It can be
    * called multiple times to include several ranges.
    */
-  public static addIncludedRange(addressess: NativePointer, size: number): void {
-    Afl.jsApiAddIncludeRange(addressess, size);
+  static addIncludedRange(addressess, size) {
+      Afl.jsApiAddIncludeRange(addressess, size);
   }
-
   /**
    * This must always be called at the end of your script. This lets
    * FRIDA mode know that your configuration is finished and that
    * execution has reached the end of your script. Failure to call
    * this will result in a fatal error.
    */
-  public static done(): void {
-    Afl.jsApiDone();
+  static done() {
+      Afl.jsApiDone();
   }
-
   /**
    * This function can be called within your script to cause FRIDA
    * mode to trigger a fatal error. This is useful if for example you
@@ -662,49 +729,48 @@ class Afl {
    * stop. The user will need to enable `AFL_DEBUG_CHILD=1` to view
    * this error message.
    */
-  public static error(msg: string): void {
-    const buf = Memory.allocUtf8String(msg);
-    Afl.jsApiError(buf);
+  static error(msg) {
+      const buf = Memory.allocUtf8String(msg);
+      Afl.jsApiError(buf);
   }
-
   /**
    * Function used to provide access to `__afl_fuzz_ptr`, which contains the length of
    * fuzzing data when using in-memory test case fuzzing.
    */
-  public static getAflFuzzLen(): NativePointer {
-
-    return Afl.jsApiGetSymbol("__afl_fuzz_len");
+  static getAflFuzzLen() {
+      return Afl.jsApiGetSymbol("__afl_fuzz_len");
   }
-
   /**
    * Function used to provide access to `__afl_fuzz_ptr`, which contains the fuzzing
    * data when using in-memory test case fuzzing.
    */
-  public static getAflFuzzPtr(): NativePointer {
-
-    return Afl.jsApiGetSymbol("__afl_fuzz_ptr");
+  static getAflFuzzPtr() {
+      return Afl.jsApiGetSymbol("__afl_fuzz_ptr");
   }
-
   /**
    * Print a message to the STDOUT. This should be preferred to
    * FRIDA's `console.log` since FRIDA will queue it's log messages.
    * If `console.log` is used in a callback in particular, then there
    * may no longer be a thread running to service this queue.
    */
-  public static print(msg: string): void {
-    const STDOUT_FILENO = 2;
-    const log = `${msg}\n`;
-    const buf = Memory.allocUtf8String(log);
-    Afl.jsApiWrite(STDOUT_FILENO, buf, log.length);
+  static print(msg) {
+      const STDOUT_FILENO = 2;
+      const log = `${msg}\n`;
+      const buf = Memory.allocUtf8String(log);
+      Afl.jsApiWrite(STDOUT_FILENO, buf, log.length);
+  }
+  /**
+   * See `AFL_FRIDA_INST_NO_BACKPATCH`.
+   */
+  static setBackpatchDisable() {
+      Afl.jsApiSetBackpatchDisable();
   }
-
   /**
    * See `AFL_FRIDA_DEBUG_MAPS`.
    */
-  public static setDebugMaps(): void {
-    Afl.jsApiSetDebugMaps();
+  static setDebugMaps() {
+      Afl.jsApiSetDebugMaps();
   }
-
   /**
    * This has the same effect as setting `AFL_ENTRYPOINT`, but has the
    * convenience of allowing you to use FRIDAs APIs to determine the
@@ -713,143 +779,198 @@ class Afl {
    * function should be called with a `NativePointer` as its
    * argument.
    */
-  public static setEntryPoint(address: NativePointer): void {
-    Afl.jsApiSetEntryPoint(address);
+  static setEntryPoint(address) {
+      Afl.jsApiSetEntryPoint(address);
   }
-
   /**
    * Function used to enable in-memory test cases for fuzzing.
    */
-  public static setInMemoryFuzzing(): void {
-    Afl.jsApiAflSharedMemFuzzing.writeInt(1);
+  static setInMemoryFuzzing() {
+      Afl.jsApiAflSharedMemFuzzing.writeInt(1);
+  }
+  /**
+   * See `AFL_FRIDA_INST_COVERAGE_FILE`. This function takes a single `string`
+   * as an argument.
+   */
+  static setInstrumentCoverageFile(file) {
+      const buf = Memory.allocUtf8String(file);
+      Afl.jsApiSetInstrumentCoverageFile(buf);
   }
-
   /**
    * See `AFL_FRIDA_INST_DEBUG_FILE`. This function takes a single `string` as
    * an argument.
    */
-  public static setInstrumentDebugFile(file: string): void {
-    const buf = Memory.allocUtf8String(file);
-    Afl.jsApiSetInstrumentDebugFile(buf);
+  static setInstrumentDebugFile(file) {
+      const buf = Memory.allocUtf8String(file);
+      Afl.jsApiSetInstrumentDebugFile(buf);
   }
-
   /**
    * See `AFL_FRIDA_INST_TRACE`.
    */
-  public static setInstrumentEnableTracing(): void {
-    Afl.jsApiSetInstrumentTrace();
+  static setInstrumentEnableTracing() {
+      Afl.jsApiSetInstrumentTrace();
+  }
+  /**
+   * See `AFL_FRIDA_INST_JIT`.
+   */
+  static setInstrumentJit() {
+      Afl.jsApiSetInstrumentJit();
   }
-
   /**
    * See `AFL_INST_LIBS`.
    */
-  public static setInstrumentLibraries(): void {
-    Afl.jsApiSetInstrumentLibraries();
+  static setInstrumentLibraries() {
+      Afl.jsApiSetInstrumentLibraries();
   }
-
   /**
    * See `AFL_FRIDA_INST_NO_OPTIMIZE`
    */
-  public static setInstrumentNoOptimize(): void {
-    Afl.jsApiSetInstrumentNoOptimize();
+  static setInstrumentNoOptimize() {
+      Afl.jsApiSetInstrumentNoOptimize();
+  }
+  /*
+    * See `AFL_FRIDA_INST_SEED`
+    */
+  static setInstrumentSeed(seed) {
+      Afl.jsApiSetInstrumentSeed(seed);
   }
-
   /**
    * See `AFL_FRIDA_INST_TRACE_UNIQUE`.
    */
-  public static setInstrumentTracingUnique(): void {
-    Afl.jsApiSetInstrumentTraceUnique();
+  static setInstrumentTracingUnique() {
+      Afl.jsApiSetInstrumentTraceUnique();
+  }
+  /**
+   * See `AFL_FRIDA_INST_UNSTABLE_COVERAGE_FILE`. This function takes a single
+   * `string` as an argument.
+   */
+  static setInstrumentUnstableCoverageFile(file) {
+      const buf = Memory.allocUtf8String(file);
+      Afl.jsApiSetInstrumentUnstableCoverageFile(buf);
+  }
+  /*
+    * Set a callback to be called in place of the usual `main` function. This see
+    * `Scripting.md` for details.
+    */
+  static setJsMainHook(address) {
+      Afl.jsApiSetJsMainHook(address);
   }
-
   /**
    * This is equivalent to setting `AFL_FRIDA_PERSISTENT_ADDR`, again a
    * `NativePointer` should be provided as it's argument.
    */
-  public static setPersistentAddress(address: NativePointer): void {
-    Afl.jsApiSetPersistentAddress(address);
+  static setPersistentAddress(address) {
+      Afl.jsApiSetPersistentAddress(address);
   }
-
   /**
    * This is equivalent to setting `AFL_FRIDA_PERSISTENT_CNT`, a
    * `number` should be provided as it's argument.
    */
-  public static setPersistentCount(count: number): void {
-    Afl.jsApiSetPersistentCount(count);
+  static setPersistentCount(count) {
+      Afl.jsApiSetPersistentCount(count);
   }
-
   /**
    * See `AFL_FRIDA_PERSISTENT_DEBUG`.
    */
-  public static setPersistentDebug(): void {
-    Afl.jsApiSetPersistentDebug();
+  static setPersistentDebug() {
+      Afl.jsApiSetPersistentDebug();
   }
-
   /**
    * See `AFL_FRIDA_PERSISTENT_ADDR`. This function takes a NativePointer as an
    * argument. See above for examples of use.
    */
-  public static setPersistentHook(address: NativePointer): void {
-    Afl.jsApiSetPersistentHook(address);
+  static setPersistentHook(address) {
+      Afl.jsApiSetPersistentHook(address);
   }
-
   /**
    * This is equivalent to setting `AFL_FRIDA_PERSISTENT_RET`, again a
    * `NativePointer` should be provided as it's argument.
    */
-  public static setPersistentReturn(address: NativePointer): void {
-    Afl.jsApiSetPersistentReturn(address);
+  static setPersistentReturn(address) {
+      Afl.jsApiSetPersistentReturn(address);
+  }
+  /**
+   * See `AFL_FRIDA_INST_NO_PREFETCH_BACKPATCH`.
+   */
+  static setPrefetchBackpatchDisable() {
+      Afl.jsApiSetPrefetchBackpatchDisable();
   }
-
   /**
    * See `AFL_FRIDA_INST_NO_PREFETCH`.
    */
-  public static setPrefetchDisable(): void {
-    Afl.jsApiSetPrefetchDisable();
+  static setPrefetchDisable() {
+      Afl.jsApiSetPrefetchDisable();
+  }
+  /**
+   * See `AFL_FRIDA_SECCOMP_FILE`. This function takes a single `string` as
+   * an argument.
+   */
+  static setSeccompFile(file) {
+      const buf = Memory.allocUtf8String(file);
+      Afl.jsApiSetSeccompFile(buf);
+  }
+  /**
+   * See `AFL_FRIDA_STALKER_ADJACENT_BLOCKS`.
+   */
+  static setStalkerAdjacentBlocks(val) {
+      Afl.jsApiSetStalkerAdjacentBlocks(val);
   }
-
   /*
-   * Set a function to be called for each instruction which is instrumented
-   * by AFL FRIDA mode.
+    * Set a function to be called for each instruction which is instrumented
+    * by AFL FRIDA mode.
+    */
+  static setStalkerCallback(callback) {
+      Afl.jsApiSetStalkerCallback(callback);
+  }
+  /**
+   * See `AFL_FRIDA_STALKER_IC_ENTRIES`.
    */
-  public static setStalkerCallback(callback: NativePointer): void {
-    Afl.jsApiSetStalkerCallback(callback);
+  static setStalkerIcEntries(val) {
+      Afl.jsApiSetStalkerIcEntries(val);
   }
-
   /**
    * See `AFL_FRIDA_STATS_FILE`. This function takes a single `string` as
    * an argument.
    */
-  public static setStatsFile(file: string): void {
-    const buf = Memory.allocUtf8String(file);
-    Afl.jsApiSetStatsFile(buf);
+  static setStatsFile(file) {
+      const buf = Memory.allocUtf8String(file);
+      Afl.jsApiSetStatsFile(buf);
   }
-
   /**
    * See `AFL_FRIDA_STATS_INTERVAL`. This function takes a `number` as an
    * argument
    */
-  public static setStatsInterval(interval: number): void {
-    Afl.jsApiSetStatsInterval(interval);
+  static setStatsInterval(interval) {
+      Afl.jsApiSetStatsInterval(interval);
   }
-
   /**
    * See `AFL_FRIDA_OUTPUT_STDERR`. This function takes a single `string` as
    * an argument.
    */
-  public static setStdErr(file: string): void {
-    const buf = Memory.allocUtf8String(file);
-    Afl.jsApiSetStdErr(buf);
+  static setStdErr(file) {
+      const buf = Memory.allocUtf8String(file);
+      Afl.jsApiSetStdErr(buf);
   }
-
   /**
    * See `AFL_FRIDA_OUTPUT_STDOUT`. This function takes a single `string` as
    * an argument.
    */
-  public static setStdOut(file: string): void {
-    const buf = Memory.allocUtf8String(file);
-    Afl.jsApiSetStdOut(buf);
+  static setStdOut(file) {
+      const buf = Memory.allocUtf8String(file);
+      Afl.jsApiSetStdOut(buf);
+  }
+  /**
+   * See `AFL_FRIDA_TRACEABLE`.
+   */
+  static setTraceable() {
+      Afl.jsApiSetTraceable();
+  }
+  static jsApiGetFunction(name, retType, argTypes) {
+      const addr = Afl.module.getExportByName(name);
+      return new NativeFunction(addr, retType, argTypes);
+  }
+  static jsApiGetSymbol(name) {
+      return Afl.module.getExportByName(name);
   }
-
 }
-
 ```
diff --git a/frida_mode/frida.map b/frida_mode/frida.map
index 61eb19ee..5276db91 100644
--- a/frida_mode/frida.map
+++ b/frida_mode/frida.map
@@ -20,6 +20,7 @@
     js_api_set_instrument_trace;
     js_api_set_instrument_trace_unique;
     js_api_set_instrument_unstable_coverage_file;
+    js_api_set_js_main_hook;
     js_api_set_persistent_address;
     js_api_set_persistent_count;
     js_api_set_persistent_debug;
diff --git a/frida_mode/include/js.h b/frida_mode/include/js.h
index a5ecb712..39aa0573 100644
--- a/frida_mode/include/js.h
+++ b/frida_mode/include/js.h
@@ -7,11 +7,14 @@ typedef gboolean (*js_api_stalker_callback_t)(const cs_insn *insn,
                                               gboolean begin, gboolean excluded,
                                               GumStalkerOutput *output);
 
+typedef int (*js_main_hook_t)(int argc, char **argv, char **envp);
+
 extern unsigned char api_js[];
 extern unsigned int  api_js_len;
 
 extern gboolean                  js_done;
 extern js_api_stalker_callback_t js_user_callback;
+extern js_main_hook_t            js_main_hook;
 
 /* Frida Mode */
 
diff --git a/frida_mode/src/js/api.js b/frida_mode/src/js/api.js
index 8e810d09..215fbdaf 100644
--- a/frida_mode/src/js/api.js
+++ b/frida_mode/src/js/api.js
@@ -151,6 +151,13 @@ class Afl {
         const buf = Memory.allocUtf8String(file);
         Afl.jsApiSetInstrumentUnstableCoverageFile(buf);
     }
+    /*
+     * Set a callback to be called in place of the usual `main` function. This see
+     * `Scripting.md` for details.
+     */
+    static setJsMainHook(address) {
+        Afl.jsApiSetJsMainHook(address);
+    }
     /**
      * This is equivalent to setting `AFL_FRIDA_PERSISTENT_ADDR`, again a
      * `NativePointer` should be provided as it's argument.
@@ -291,6 +298,7 @@ Afl.jsApiSetInstrumentSeed = Afl.jsApiGetFunction("js_api_set_instrument_seed",
 Afl.jsApiSetInstrumentTrace = Afl.jsApiGetFunction("js_api_set_instrument_trace", "void", []);
 Afl.jsApiSetInstrumentTraceUnique = Afl.jsApiGetFunction("js_api_set_instrument_trace_unique", "void", []);
 Afl.jsApiSetInstrumentUnstableCoverageFile = Afl.jsApiGetFunction("js_api_set_instrument_unstable_coverage_file", "void", ["pointer"]);
+Afl.jsApiSetJsMainHook = Afl.jsApiGetFunction("js_api_set_js_main_hook", "void", ["pointer"]);
 Afl.jsApiSetPersistentAddress = Afl.jsApiGetFunction("js_api_set_persistent_address", "void", ["pointer"]);
 Afl.jsApiSetPersistentCount = Afl.jsApiGetFunction("js_api_set_persistent_count", "void", ["uint64"]);
 Afl.jsApiSetPersistentDebug = Afl.jsApiGetFunction("js_api_set_persistent_debug", "void", []);
@@ -299,8 +307,8 @@ Afl.jsApiSetPersistentReturn = Afl.jsApiGetFunction("js_api_set_persistent_retur
 Afl.jsApiSetPrefetchBackpatchDisable = Afl.jsApiGetFunction("js_api_set_prefetch_backpatch_disable", "void", []);
 Afl.jsApiSetPrefetchDisable = Afl.jsApiGetFunction("js_api_set_prefetch_disable", "void", []);
 Afl.jsApiSetSeccompFile = Afl.jsApiGetFunction("js_api_set_seccomp_file", "void", ["pointer"]);
-Afl.jsApiSetStalkerCallback = Afl.jsApiGetFunction("js_api_set_stalker_callback", "void", ["pointer"]);
 Afl.jsApiSetStalkerAdjacentBlocks = Afl.jsApiGetFunction("js_api_set_stalker_adjacent_blocks", "void", ["uint32"]);
+Afl.jsApiSetStalkerCallback = Afl.jsApiGetFunction("js_api_set_stalker_callback", "void", ["pointer"]);
 Afl.jsApiSetStalkerIcEntries = Afl.jsApiGetFunction("js_api_set_stalker_ic_entries", "void", ["uint32"]);
 Afl.jsApiSetStatsFile = Afl.jsApiGetFunction("js_api_set_stats_file", "void", ["pointer"]);
 Afl.jsApiSetStatsInterval = Afl.jsApiGetFunction("js_api_set_stats_interval", "void", ["uint64"]);
diff --git a/frida_mode/src/js/js.c b/frida_mode/src/js/js.c
index 37cd377b..5f477388 100644
--- a/frida_mode/src/js/js.c
+++ b/frida_mode/src/js/js.c
@@ -3,10 +3,11 @@
 #include "js.h"
 #include "util.h"
 
-static char *             js_script = NULL;
 gboolean                  js_done = FALSE;
 js_api_stalker_callback_t js_user_callback = NULL;
+js_main_hook_t            js_main_hook = NULL;
 
+static char *              js_script = NULL;
 static gchar *             filename = "afl.js";
 static gchar *             contents;
 static GumScriptBackend *  backend;
diff --git a/frida_mode/src/js/js_api.c b/frida_mode/src/js/js_api.c
index 4221fb80..5021b531 100644
--- a/frida_mode/src/js/js_api.c
+++ b/frida_mode/src/js/js_api.c
@@ -255,3 +255,10 @@ __attribute__((visibility("default"))) void js_api_set_stalker_adjacent_blocks(
 
 }
 
+__attribute__((visibility("default"))) void js_api_set_js_main_hook(
+    const js_main_hook_t hook) {
+
+  js_main_hook = hook;
+
+}
+
diff --git a/frida_mode/src/main.c b/frida_mode/src/main.c
index cb88eabe..913e3a46 100644
--- a/frida_mode/src/main.c
+++ b/frida_mode/src/main.c
@@ -36,13 +36,13 @@
 extern mach_port_t mach_task_self();
 extern GumAddress  gum_darwin_find_entrypoint(mach_port_t task);
 #else
-extern int  __libc_start_main(int *(main)(int, char **, char **), int argc,
+extern int  __libc_start_main(int (*main)(int, char **, char **), int argc,
                               char **ubp_av, void (*init)(void),
                               void (*fini)(void), void (*rtld_fini)(void),
                               void(*stack_end));
 #endif
 
-typedef int *(*main_fn_t)(int argc, char **argv, char **envp);
+typedef int (*main_fn_t)(int argc, char **argv, char **envp);
 
 static main_fn_t main_fn = NULL;
 
@@ -217,7 +217,7 @@ __attribute__((visibility("default"))) void afl_frida_start(void) {
 
 }
 
-static int *on_main(int argc, char **argv, char **envp) {
+static int on_main(int argc, char **argv, char **envp) {
 
   on_main_os(argc, argv, envp);
 
@@ -225,12 +225,20 @@ static int *on_main(int argc, char **argv, char **envp) {
 
   afl_frida_start();
 
-  return main_fn(argc, argv, envp);
+  if (js_main_hook != NULL) {
+
+    return js_main_hook(argc, argv, envp);
+
+  } else {
+
+    return main_fn(argc, argv, envp);
+
+  }
 
 }
 
 #if defined(EMBEDDED)
-extern int *main(int argc, char **argv, char **envp);
+extern int main(int argc, char **argv, char **envp);
 
 static void intercept_main(void) {
 
@@ -253,7 +261,7 @@ static void intercept_main(void) {
 }
 
 #else
-static int on_libc_start_main(int *(main)(int, char **, char **), int argc,
+static int on_libc_start_main(int (*main)(int, char **, char **), int argc,
                               char **ubp_av, void (*init)(void),
                               void (*fini)(void), void (*rtld_fini)(void),
                               void(*stack_end)) {
diff --git a/frida_mode/test/js/GNUmakefile b/frida_mode/test/js/GNUmakefile
index ccd990c0..c702ad98 100644
--- a/frida_mode/test/js/GNUmakefile
+++ b/frida_mode/test/js/GNUmakefile
@@ -47,6 +47,18 @@ $(AFLPP_DRIVER_DUMMY_INPUT): | $(BUILD_DIR)
 clean:
 	rm -rf $(BUILD_DIR)
 
+frida_js_main: $(TESTINSTBIN) $(TEST_DATA_FILE) $(AFLPP_DRIVER_DUMMY_INPUT)
+	AFL_PRELOAD=$(AFL_PRELOAD) \
+	AFL_FRIDA_JS_SCRIPT=main.js \
+	$(ROOT)afl-fuzz \
+		-D \
+		-O \
+		-i $(TEST_DATA_DIR) \
+		-o $(FRIDA_OUT) \
+		-t 10000+ \
+		-- \
+			$(TESTINSTBIN) $(AFLPP_DRIVER_DUMMY_INPUT)
+
 frida_js_fuzz: $(TESTINSTBIN) $(TEST_DATA_FILE) $(AFLPP_DRIVER_DUMMY_INPUT)
 	AFL_PRELOAD=$(AFL_PRELOAD) \
 	AFL_FRIDA_JS_SCRIPT=fuzz.js \
diff --git a/frida_mode/test/js/main.js b/frida_mode/test/js/main.js
new file mode 100644
index 00000000..06306fc4
--- /dev/null
+++ b/frida_mode/test/js/main.js
@@ -0,0 +1,44 @@
+Afl.print('******************');
+Afl.print('* AFL FRIDA MODE *');
+Afl.print('******************');
+Afl.print('');
+
+Afl.print(`PID: ${Process.id}`);
+
+const name = Process.enumerateModules()[0].name;
+Afl.print(`Name: ${name}`);
+
+new ModuleMap().values().forEach(m => {
+    Afl.print(`${m.base}-${m.base.add(m.size)} ${m.name}`);
+});
+
+const main = DebugSymbol.fromName('main').address;
+Afl.print(`main: ${main}`);
+
+const LLVMFuzzerTestOneInput = DebugSymbol.fromName('LLVMFuzzerTestOneInput').address;
+Afl.print(`LLVMFuzzerTestOneInput: ${LLVMFuzzerTestOneInput}`);
+
+const cm = new CModule(`
+
+    extern unsigned char * __afl_fuzz_ptr;
+    extern unsigned int * __afl_fuzz_len;
+    extern void LLVMFuzzerTestOneInput(char *buf, int len);
+
+    int main(int argc, char **argv)  {
+
+      LLVMFuzzerTestOneInput(__afl_fuzz_ptr, *__afl_fuzz_len);
+
+    }
+    `,
+    {
+        LLVMFuzzerTestOneInput: LLVMFuzzerTestOneInput,
+        __afl_fuzz_ptr: Afl.getAflFuzzPtr(),
+        __afl_fuzz_len: Afl.getAflFuzzLen()
+    });
+
+Afl.setEntryPoint(cm.main);
+Afl.setPersistentAddress(cm.main);
+Afl.setInMemoryFuzzing();
+Afl.setJsMainHook(cm.main);
+Afl.print("done");
+Afl.done();
diff --git a/frida_mode/ts/lib/afl.ts b/frida_mode/ts/lib/afl.ts
index e20ad3ec..0473cbf6 100644
--- a/frida_mode/ts/lib/afl.ts
+++ b/frida_mode/ts/lib/afl.ts
@@ -179,6 +179,14 @@ class Afl {
     Afl.jsApiSetInstrumentUnstableCoverageFile(buf);
   }
 
+  /*
+   * Set a callback to be called in place of the usual `main` function. This see
+   * `Scripting.md` for details.
+   */
+  public static setJsMainHook(address: NativePointer): void {
+    Afl.jsApiSetJsMainHook(address);
+  }
+
   /**
    * This is equivalent to setting `AFL_FRIDA_PERSISTENT_ADDR`, again a
    * `NativePointer` should be provided as it's argument.
@@ -387,6 +395,11 @@ class Afl {
     "void",
     ["pointer"]);
 
+  private static readonly jsApiSetJsMainHook = Afl.jsApiGetFunction(
+    "js_api_set_js_main_hook",
+    "void",
+    ["pointer"]);
+
   private static readonly jsApiSetPersistentAddress = Afl.jsApiGetFunction(
     "js_api_set_persistent_address",
     "void",
@@ -427,16 +440,16 @@ class Afl {
     "void",
     ["pointer"]);
 
-  private static readonly jsApiSetStalkerCallback = Afl.jsApiGetFunction(
-    "js_api_set_stalker_callback",
-    "void",
-    ["pointer"]);
-
   private static readonly jsApiSetStalkerAdjacentBlocks = Afl.jsApiGetFunction(
     "js_api_set_stalker_adjacent_blocks",
     "void",
     ["uint32"]);
 
+  private static readonly jsApiSetStalkerCallback = Afl.jsApiGetFunction(
+    "js_api_set_stalker_callback",
+    "void",
+    ["pointer"]);
+
   private static readonly jsApiSetStalkerIcEntries = Afl.jsApiGetFunction(
     "js_api_set_stalker_ic_entries",
     "void",