From f14b3bd9de4570a30c0f89bac2879ee3a7cfbcae Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Tue, 26 Oct 2021 18:48:13 +0100
Subject: Fix issues with Yama restrictions in FRIDA mode

---
 frida_mode/src/js/api.js | 7 +++++++
 1 file changed, 7 insertions(+)

(limited to 'frida_mode/src/js/api.js')

diff --git a/frida_mode/src/js/api.js b/frida_mode/src/js/api.js
index 40bb4a16..6f9f05d8 100644
--- a/frida_mode/src/js/api.js
+++ b/frida_mode/src/js/api.js
@@ -243,6 +243,12 @@ class Afl {
         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);
@@ -286,6 +292,7 @@ Afl.jsApiSetStatsFile = Afl.jsApiGetFunction("js_api_set_stats_file", "void", ["
 Afl.jsApiSetStatsInterval = Afl.jsApiGetFunction("js_api_set_stats_interval", "void", ["uint64"]);
 Afl.jsApiSetStdErr = Afl.jsApiGetFunction("js_api_set_stderr", "void", ["pointer"]);
 Afl.jsApiSetStdOut = Afl.jsApiGetFunction("js_api_set_stdout", "void", ["pointer"]);
+Afl.jsApiSetTraceable = Afl.jsApiGetFunction("js_api_set_traceable", "void", []);
 Afl.jsApiWrite = new NativeFunction(
 /* tslint:disable-next-line:no-null-keyword */
 Module.getExportByName(null, "write"), "int", ["int", "pointer", "int"]);
-- 
cgit 1.4.1


From 93b8f17242f6f354da189fee68e69c55288e2488 Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Mon, 8 Nov 2021 18:17:09 +0000
Subject: Add support for option to disable backpatch

---
 frida_mode/README.md         |  2 ++
 frida_mode/frida.map         |  1 +
 frida_mode/include/stalker.h |  3 ++-
 frida_mode/src/js/api.js     |  7 +++++++
 frida_mode/src/js/js_api.c   |  6 ++++++
 frida_mode/src/stalker.c     |  9 +++++++--
 frida_mode/ts/lib/afl.ts     | 12 ++++++++++++
 include/envs.h               |  1 +
 8 files changed, 38 insertions(+), 3 deletions(-)

(limited to 'frida_mode/src/js/api.js')

diff --git a/frida_mode/README.md b/frida_mode/README.md
index bb194080..8211224d 100644
--- a/frida_mode/README.md
+++ b/frida_mode/README.md
@@ -174,6 +174,8 @@ file.
 * `AFL_FRIDA_INST_NO_OPTIMIZE` - Don't use optimized inline assembly coverage
 instrumentation (the default where available). Required to use
 `AFL_FRIDA_INST_TRACE`.
+* `AFL_FRIDA_INST_NO_BACKPATCH` - Disable backpatching. At the end of executing
+each block, control will return to FRIDA to identify the next block to execute.
 * `AFL_FRIDA_INST_NO_PREFETCH` - Disable prefetching. By default the child will
 report instrumented blocks back to the parent so that it can also instrument
 them and they be inherited by the next child on fork, implies
diff --git a/frida_mode/frida.map b/frida_mode/frida.map
index e2ae87a7..7be41aa0 100644
--- a/frida_mode/frida.map
+++ b/frida_mode/frida.map
@@ -8,6 +8,7 @@
     js_api_add_include_range;
     js_api_done;
     js_api_error;
+    js_api_set_backpatch_disable;
     js_api_set_debug_maps;
     js_api_set_entrypoint;
     js_api_set_instrument_coverage_file;
diff --git a/frida_mode/include/stalker.h b/frida_mode/include/stalker.h
index 955f3913..8a111b90 100644
--- a/frida_mode/include/stalker.h
+++ b/frida_mode/include/stalker.h
@@ -3,7 +3,8 @@
 
 #include "frida-gumjs.h"
 
-extern guint stalker_ic_entries;
+extern guint    stalker_ic_entries;
+extern gboolean backpatch_enable;
 
 void        stalker_config(void);
 void        stalker_init(void);
diff --git a/frida_mode/src/js/api.js b/frida_mode/src/js/api.js
index 6f9f05d8..5db62389 100644
--- a/frida_mode/src/js/api.js
+++ b/frida_mode/src/js/api.js
@@ -62,6 +62,12 @@ class Afl {
         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`.
      */
@@ -267,6 +273,7 @@ Afl.jsApiAddIncludeRange = Afl.jsApiGetFunction("js_api_add_include_range", "voi
 Afl.jsApiAflSharedMemFuzzing = Afl.jsApiGetSymbol("__afl_sharedmem_fuzzing");
 Afl.jsApiDone = Afl.jsApiGetFunction("js_api_done", "void", []);
 Afl.jsApiError = Afl.jsApiGetFunction("js_api_error", "void", ["pointer"]);
+Afl.jsApiSetBackpatchDisable = Afl.jsApiGetFunction("js_api_set_backpatch_disable", "void", []);
 Afl.jsApiSetDebugMaps = Afl.jsApiGetFunction("js_api_set_debug_maps", "void", []);
 Afl.jsApiSetEntryPoint = Afl.jsApiGetFunction("js_api_set_entrypoint", "void", ["pointer"]);
 Afl.jsApiSetInstrumentCoverageFile = Afl.jsApiGetFunction("js_api_set_instrument_coverage_file", "void", ["pointer"]);
diff --git a/frida_mode/src/js/js_api.c b/frida_mode/src/js/js_api.c
index f3d81a32..abc0ac30 100644
--- a/frida_mode/src/js/js_api.c
+++ b/frida_mode/src/js/js_api.c
@@ -237,3 +237,9 @@ __attribute__((visibility("default"))) void js_api_set_traceable(void) {
 
 }
 
+__attribute__((visibility("default"))) void js_api_set_backpatch_disable(void) {
+
+  backpatch_enable = FALSE;
+
+}
+
diff --git a/frida_mode/src/stalker.c b/frida_mode/src/stalker.c
index 814aaeb3..6ba41bc6 100644
--- a/frida_mode/src/stalker.c
+++ b/frida_mode/src/stalker.c
@@ -6,7 +6,8 @@
 #include "stats.h"
 #include "util.h"
 
-guint stalker_ic_entries = 0;
+guint    stalker_ic_entries = 0;
+gboolean backpatch_enable = TRUE;
 
 static GumStalker *stalker = NULL;
 
@@ -58,6 +59,8 @@ void stalker_config(void) {
 
   if (!gum_stalker_is_supported()) { FATAL("Failed to initialize embedded"); }
 
+  backpatch_enable = (getenv("AFL_FRIDA_INST_NO_BACKPATCH") == NULL);
+
   stalker_ic_entries = util_read_num("AFL_FRIDA_STALKER_IC_ENTRIES");
 
   observer = g_object_new(GUM_TYPE_AFL_STALKER_OBSERVER, NULL);
@@ -87,6 +90,8 @@ static gboolean stalker_exclude_self(const GumRangeDetails *details,
 
 void stalker_init(void) {
 
+  OKF("Instrumentation - backpatch [%c]", backpatch_enable ? 'X' : ' ');
+
   OKF("Stalker - ic_entries [%u]", stalker_ic_entries);
 
 #if !(defined(__x86_64__) || defined(__i386__))
@@ -134,7 +139,7 @@ void stalker_start(void) {
 
 void stalker_trust(void) {
 
-  gum_stalker_set_trust_threshold(stalker, 0);
+  if (backpatch_enable) { gum_stalker_set_trust_threshold(stalker, 0); }
 
 }
 
diff --git a/frida_mode/ts/lib/afl.ts b/frida_mode/ts/lib/afl.ts
index 538d9b70..3639d670 100644
--- a/frida_mode/ts/lib/afl.ts
+++ b/frida_mode/ts/lib/afl.ts
@@ -77,6 +77,13 @@ class Afl {
     Afl.jsApiWrite(STDOUT_FILENO, buf, log.length);
   }
 
+  /**
+   * See `AFL_FRIDA_INST_NO_BACKPATCH`.
+   */
+  public static setBackpatchDisable(): void {
+    Afl.jsApiSetBackpatchDisable();
+  }
+
   /**
    * See `AFL_FRIDA_DEBUG_MAPS`.
    */
@@ -313,6 +320,11 @@ class Afl {
     "void",
     ["pointer"]);
 
+  private static readonly jsApiSetBackpatchDisable = Afl.jsApiGetFunction(
+    "js_api_set_backpatch_disable",
+    "void",
+    []);
+
   private static readonly jsApiSetDebugMaps = Afl.jsApiGetFunction(
     "js_api_set_debug_maps",
     "void",
diff --git a/include/envs.h b/include/envs.h
index fd459f14..0ba79092 100644
--- a/include/envs.h
+++ b/include/envs.h
@@ -59,6 +59,7 @@ static char *afl_environment_variables[] = {
     "AFL_FRIDA_INST_COVERAGE_FILE",
     "AFL_FRIDA_INST_DEBUG_FILE",
     "AFL_FRIDA_INST_JIT",
+    "AFL_FRIDA_INST_NO_BACKPATCH",
     "AFL_FRIDA_INST_NO_OPTIMIZE",
     "AFL_FRIDA_INST_NO_PREFETCH",
     "AFL_FRIDA_INST_NO_PREFETCH_BACKPATCH",
-- 
cgit 1.4.1


From e1d3fe30dec150aa7111cb623a7362a8026963a8 Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Mon, 15 Nov 2021 17:14:04 +0000
Subject: Changes to allow configuration of stalker adjacent blocks

---
 frida_mode/README.md         |  5 +++++
 frida_mode/frida.map         |  1 +
 frida_mode/include/stalker.h |  1 +
 frida_mode/src/js/api.js     |  7 +++++++
 frida_mode/src/js/js_api.c   |  8 ++++++++
 frida_mode/src/stalker.c     | 18 +++++++++++++++---
 frida_mode/ts/lib/afl.ts     | 12 ++++++++++++
 include/envs.h               |  2 ++
 8 files changed, 51 insertions(+), 3 deletions(-)

(limited to 'frida_mode/src/js/api.js')

diff --git a/frida_mode/README.md b/frida_mode/README.md
index 8211224d..a75324d5 100644
--- a/frida_mode/README.md
+++ b/frida_mode/README.md
@@ -215,6 +215,11 @@ gdb \
 ```
 * `AFL_FRIDA_SECCOMP_FILE` - Write a log of any syscalls made by the target to
 the specified file.
+* `AFL_FRIDA_STALKER_ADJACENT_BLOCKS` - Configure the number of adjacent blocks
+ to fetch when generating instrumented code. By fetching blocks in the same
+ order they appear in the original program, rather than the order of execution
+ should help reduce locallity and adjacency. This includes allowing us to vector
+ between adjancent blocks using a NOP slide rather than an immediate branch.
 * `AFL_FRIDA_STALKER_IC_ENTRIES` - Configure the number of inline cache entries
 stored along-side branch instructions which provide a cache to avoid having to
 call back into FRIDA to find the next block. Default is 32.
diff --git a/frida_mode/frida.map b/frida_mode/frida.map
index 7be41aa0..61eb19ee 100644
--- a/frida_mode/frida.map
+++ b/frida_mode/frida.map
@@ -29,6 +29,7 @@
     js_api_set_prefetch_disable;
     js_api_set_seccomp_file;
     js_api_set_stalker_callback;
+    js_api_set_stalker_adjacent_blocks;
     js_api_set_stalker_ic_entries;
     js_api_set_stats_file;
     js_api_set_stats_interval;
diff --git a/frida_mode/include/stalker.h b/frida_mode/include/stalker.h
index 8a111b90..666787e9 100644
--- a/frida_mode/include/stalker.h
+++ b/frida_mode/include/stalker.h
@@ -5,6 +5,7 @@
 
 extern guint    stalker_ic_entries;
 extern gboolean backpatch_enable;
+extern guint    stalker_adjacent_blocks;
 
 void        stalker_config(void);
 void        stalker_init(void);
diff --git a/frida_mode/src/js/api.js b/frida_mode/src/js/api.js
index 5db62389..8e810d09 100644
--- a/frida_mode/src/js/api.js
+++ b/frida_mode/src/js/api.js
@@ -205,6 +205,12 @@ class Afl {
         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.
@@ -294,6 +300,7 @@ Afl.jsApiSetPrefetchBackpatchDisable = Afl.jsApiGetFunction("js_api_set_prefetch
 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.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_api.c b/frida_mode/src/js/js_api.c
index c1f092c9..8e0a549c 100644
--- a/frida_mode/src/js/js_api.c
+++ b/frida_mode/src/js/js_api.c
@@ -250,3 +250,11 @@ __attribute__((visibility("default"))) void js_api_set_backpatch_disable(void) {
 
 }
 
+
+__attribute__((visibility("default"))) void js_api_set_stalker_adjacent_blocks(
+    guint val) {
+
+  stalker_adjacent_blocks = val;
+
+}
+
diff --git a/frida_mode/src/stalker.c b/frida_mode/src/stalker.c
index 35a9d856..65ed5d50 100644
--- a/frida_mode/src/stalker.c
+++ b/frida_mode/src/stalker.c
@@ -7,6 +7,7 @@
 
 guint    stalker_ic_entries = 0;
 gboolean backpatch_enable = TRUE;
+guint    stalker_adjacent_blocks = 0;
 
 static GumStalker *stalker = NULL;
 
@@ -60,7 +61,9 @@ void stalker_config(void) {
 
   backpatch_enable = (getenv("AFL_FRIDA_INST_NO_BACKPATCH") == NULL);
 
-  stalker_ic_entries = util_read_num("AFL_FRIDA_STALKER_IC_ENTRIES");
+  stalker_ic_entries = util_read_num("AFL_FRIDA_STALKER_ADJACENT_BLOCKS");
+
+  stalker_adjacent_blocks = util_read_num("AFL_FRIDA_STALKER_IC_ENTRIES");
 
   observer = g_object_new(GUM_TYPE_AFL_STALKER_OBSERVER, NULL);
 
@@ -92,6 +95,7 @@ void stalker_init(void) {
   FOKF("Instrumentation - backpatch [%c]", backpatch_enable ? 'X' : ' ');
 
   FOKF("Stalker - ic_entries [%u]", stalker_ic_entries);
+  FOKF("Stalker - adjacent_blocks [%u]", stalker_adjacent_blocks);
 
 #if !(defined(__x86_64__) || defined(__i386__))
   if (stalker_ic_entries != 0) {
@@ -100,13 +104,21 @@ void stalker_init(void) {
 
   }
 
+  if (stalker_adjacent_blocks != 0) {
+
+    FFATAL("AFL_FRIDA_STALKER_ADJACENT_BLOCKS not supported");
+
+  }
+
 #endif
 
   if (stalker_ic_entries == 0) { stalker_ic_entries = 32; }
 
+  if (stalker_adjacent_blocks == 0) { stalker_adjacent_blocks = 32; }
+
 #if defined(__x86_64__) || defined(__i386__)
-  stalker =
-      g_object_new(GUM_TYPE_STALKER, "ic-entries", stalker_ic_entries, NULL);
+  stalker = g_object_new(GUM_TYPE_STALKER, "ic-entries", stalker_ic_entries,
+                         "adjacent-blocks", stalker_adjacent_blocks, NULL);
 #else
   stalker = gum_stalker_new();
 #endif
diff --git a/frida_mode/ts/lib/afl.ts b/frida_mode/ts/lib/afl.ts
index 3639d670..e20ad3ec 100644
--- a/frida_mode/ts/lib/afl.ts
+++ b/frida_mode/ts/lib/afl.ts
@@ -241,6 +241,13 @@ class Afl {
     Afl.jsApiSetSeccompFile(buf);
   }
 
+  /**
+   * See `AFL_FRIDA_STALKER_ADJACENT_BLOCKS`.
+   */
+  public static setStalkerAdjacentBlocks(val: number): void {
+    Afl.jsApiSetStalkerAdjacentBlocks(val);
+  }
+
   /*
    * Set a function to be called for each instruction which is instrumented
    * by AFL FRIDA mode.
@@ -425,6 +432,11 @@ class Afl {
     "void",
     ["pointer"]);
 
+  private static readonly jsApiSetStalkerAdjacentBlocks = Afl.jsApiGetFunction(
+    "js_api_set_stalker_adjacent_blocks",
+    "void",
+    ["uint32"]);
+
   private static readonly jsApiSetStalkerIcEntries = Afl.jsApiGetFunction(
     "js_api_set_stalker_ic_entries",
     "void",
diff --git a/include/envs.h b/include/envs.h
index 0ba79092..a3ba5e88 100644
--- a/include/envs.h
+++ b/include/envs.h
@@ -76,6 +76,8 @@ static char *afl_environment_variables[] = {
     "AFL_FRIDA_PERSISTENT_DEBUG",
     "AFL_FRIDA_PERSISTENT_HOOK",
     "AFL_FRIDA_PERSISTENT_RET",
+    "AFL_FRIDA_STALKER_IC_ENTRIES",
+    "AFL_FRIDA_STALKER_ADJACENT_BLOCKS",
     "AFL_FRIDA_STATS_FILE",
     "AFL_FRIDA_STATS_INTERVAL",
     "AFL_FRIDA_TRACEABLE",
-- 
cgit 1.4.1


From af02fa1670db6d19feaf0a3e54d9d8013ad3312f Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Thu, 18 Nov 2021 17:08:39 +0000
Subject: Improve JS bindings for hooking functions

---
 frida_mode/Scripting.md        | 333 ++++++++++++++++++++++++++++-------------
 frida_mode/frida.map           |   1 +
 frida_mode/include/js.h        |   3 +
 frida_mode/src/js/api.js       |  10 +-
 frida_mode/src/js/js.c         |   3 +-
 frida_mode/src/js/js_api.c     |   7 +
 frida_mode/src/main.c          |  20 ++-
 frida_mode/test/js/GNUmakefile |  12 ++
 frida_mode/test/js/main.js     |  44 ++++++
 frida_mode/ts/lib/afl.ts       |  23 ++-
 10 files changed, 337 insertions(+), 119 deletions(-)
 create mode 100644 frida_mode/test/js/main.js

(limited to 'frida_mode/src/js/api.js')

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",
-- 
cgit 1.4.1