Support for excluding JIT code (#1006)

Co-authored-by: Your Name <you@example.com>

8 files changed, 90 insertions, 9 deletions

diff --git a/frida_mode/README.md b/frida_mode/README.md
index c85cf3af..024fc140 100644
--- a/frida_mode/README.md
+++ b/frida_mode/README.md
@@ -153,6 +153,9 @@ Generated block 0x7ffff75e98e2
 
 ***
 ```
+* `AFL_FRIDA_INST_JIT` - Enable the instrumentation of Just-In-Time compiled
+code. Code is considered to be JIT if the executable segment is not backed by a
+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`.
diff --git a/frida_mode/frida.map b/frida_mode/frida.map
index cc072dd7..8fc0b174 100644
--- a/frida_mode/frida.map
+++ b/frida_mode/frida.map
@@ -11,6 +11,7 @@
     js_api_set_debug_maps;
     js_api_set_entrypoint;
     js_api_set_instrument_debug_file;
+    js_api_set_instrument_jit;
     js_api_set_instrument_libraries;
     js_api_set_instrument_no_optimize;
     js_api_set_instrument_trace;
diff --git a/frida_mode/include/ranges.h b/frida_mode/include/ranges.h
index a667fb76..2eb9b355 100644
--- a/frida_mode/include/ranges.h
+++ b/frida_mode/include/ranges.h
@@ -5,6 +5,7 @@
 
 extern gboolean ranges_debug_maps;
 extern gboolean ranges_inst_libs;
+extern gboolean ranges_inst_jit;
 
 void ranges_config(void);
 void ranges_init(void);
diff --git a/frida_mode/src/js/api.js b/frida_mode/src/js/api.js
index 4cb04704..1d843024 100644
--- a/frida_mode/src/js/api.js
+++ b/frida_mode/src/js/api.js
@@ -100,6 +100,12 @@ class Afl {
         Afl.jsApiSetInstrumentTrace();
     }
     /**
+     * See `AFL_FRIDA_INST_JIT`.
+     */
+    static setInstrumentJit() {
+        Afl.jsApiSetInstrumentJit();
+    }
+    /**
      * See `AFL_INST_LIBS`.
      */
     static setInstrumentLibraries() {
@@ -222,6 +228,7 @@ Afl.jsApiError = Afl.jsApiGetFunction("js_api_error", "void", ["pointer"]);
 Afl.jsApiSetDebugMaps = Afl.jsApiGetFunction("js_api_set_debug_maps", "void", []);
 Afl.jsApiSetEntryPoint = Afl.jsApiGetFunction("js_api_set_entrypoint", "void", ["pointer"]);
 Afl.jsApiSetInstrumentDebugFile = Afl.jsApiGetFunction("js_api_set_instrument_debug_file", "void", ["pointer"]);
+Afl.jsApiSetInstrumentJit = Afl.jsApiGetFunction("js_api_set_instrument_jit", "void", []);
 Afl.jsApiSetInstrumentLibraries = Afl.jsApiGetFunction("js_api_set_instrument_libraries", "void", []);
 Afl.jsApiSetInstrumentNoOptimize = Afl.jsApiGetFunction("js_api_set_instrument_no_optimize", "void", []);
 Afl.jsApiSetInstrumentTrace = Afl.jsApiGetFunction("js_api_set_instrument_trace", "void", []);
diff --git a/frida_mode/src/js/js_api.c b/frida_mode/src/js/js_api.c
index 58bf9ba3..36471387 100644
--- a/frida_mode/src/js/js_api.c
+++ b/frida_mode/src/js/js_api.c
@@ -77,6 +77,12 @@ __attribute__((visibility("default"))) void js_api_add_exclude_range(
 
 }
 
+__attribute__((visibility("default"))) void js_api_set_instrument_jit() {
+
+  ranges_inst_jit = TRUE;
+
+}
+
 __attribute__((visibility("default"))) void js_api_set_instrument_libraries() {
 
   ranges_inst_libs = TRUE;
diff --git a/frida_mode/src/ranges.c b/frida_mode/src/ranges.c
index 05e18156..5e78fa60 100644
--- a/frida_mode/src/ranges.c
+++ b/frida_mode/src/ranges.c
@@ -19,9 +19,11 @@ typedef struct {
 
 gboolean ranges_debug_maps = FALSE;
 gboolean ranges_inst_libs = FALSE;
+gboolean ranges_inst_jit = FALSE;
 
 static GArray *module_ranges = NULL;
 static GArray *libs_ranges = NULL;
+static GArray *jit_ranges = NULL;
 static GArray *include_ranges = NULL;
 static GArray *exclude_ranges = NULL;
 static GArray *ranges = NULL;
@@ -174,19 +176,27 @@ static gboolean print_ranges_callback(const GumRangeDetails *details,
                                       gpointer               user_data) {
 
   UNUSED_PARAMETER(user_data);
+
   if (details->file == NULL) {
 
-    OKF("MAP - 0x%016" G_GINT64_MODIFIER "x - 0x%016" G_GINT64_MODIFIER "X",
+    OKF("MAP - 0x%016" G_GINT64_MODIFIER "x - 0x%016" G_GINT64_MODIFIER
+        "X %c%c%c",
         details->range->base_address,
-        details->range->base_address + details->range->size);
+        details->range->base_address + details->range->size,
+        details->protection & GUM_PAGE_READ ? 'R' : '-',
+        details->protection & GUM_PAGE_WRITE ? 'W' : '-',
+        details->protection & GUM_PAGE_EXECUTE ? 'X' : '-');
 
   } else {
 
     OKF("MAP - 0x%016" G_GINT64_MODIFIER "x - 0x%016" G_GINT64_MODIFIER
-        "X %s(0x%016" G_GINT64_MODIFIER "x)",
+        "X %c%c%c %s(0x%016" G_GINT64_MODIFIER "x)",
         details->range->base_address,
         details->range->base_address + details->range->size,
-        details->file->path, details->file->offset);
+        details->protection & GUM_PAGE_READ ? 'R' : '-',
+        details->protection & GUM_PAGE_WRITE ? 'W' : '-',
+        details->protection & GUM_PAGE_EXECUTE ? 'X' : '-', details->file->path,
+        details->file->offset);
 
   }
 
@@ -331,6 +341,39 @@ static GArray *collect_libs_ranges(void) {
 
 }
 
+static gboolean collect_jit_ranges_callback(const GumRangeDetails *details,
+                                            gpointer               user_data) {
+
+  GArray *ranges = (GArray *)user_data;
+
+  /* If the executable code isn't backed by a file, it's probably JIT */
+  if (details->file == NULL) {
+
+    GumMemoryRange range = *details->range;
+    g_array_append_val(ranges, range);
+
+  }
+
+  return TRUE;
+
+}
+
+static GArray *collect_jit_ranges(void) {
+
+  GArray *result;
+  result = g_array_new(false, false, sizeof(GumMemoryRange));
+  if (!ranges_inst_jit) {
+
+    gum_process_enumerate_ranges(GUM_PAGE_EXECUTE, collect_jit_ranges_callback,
+                                 result);
+
+  }
+
+  print_ranges("JIT", result);
+  return result;
+
+}
+
 static gboolean intersect_range(GumMemoryRange *rr, GumMemoryRange *ra,
                                 GumMemoryRange *rb) {
 
@@ -510,6 +553,7 @@ void ranges_config(void) {
 
   if (getenv("AFL_FRIDA_DEBUG_MAPS") != NULL) { ranges_debug_maps = TRUE; }
   if (getenv("AFL_INST_LIBS") != NULL) { ranges_inst_libs = TRUE; }
+  if (getenv("AFL_FRIDA_INST_JIT") != NULL) { ranges_inst_jit = TRUE; }
 
   if (ranges_debug_maps) {
 
@@ -530,7 +574,9 @@ void ranges_init(void) {
   GArray *       step2;
   GArray *       step3;
   GArray *       step4;
+  GArray *       step5;
 
+  OKF("Ranges - Instrument jit [%c]", ranges_inst_jit ? 'X' : ' ');
   OKF("Ranges - Instrument libraries [%c]", ranges_inst_libs ? 'X' : ' ');
 
   print_ranges("AFL_FRIDA_INST_RANGES", include_ranges);
@@ -538,6 +584,7 @@ void ranges_init(void) {
 
   module_ranges = collect_module_ranges();
   libs_ranges = collect_libs_ranges();
+  jit_ranges = collect_jit_ranges();
 
   /* If include ranges is empty, then assume everything is included */
   if (include_ranges->len == 0) {
@@ -560,17 +607,20 @@ void ranges_init(void) {
   step3 = subtract_ranges(step2, exclude_ranges);
   print_ranges("step3", step3);
 
+  step4 = subtract_ranges(step3, jit_ranges);
+  print_ranges("step4", step4);
+
   /*
-   * After step3, we have the total ranges to be instrumented, we now subtract
+   * After step4, we have the total ranges to be instrumented, we now subtract
    * that from the original ranges of the modules to configure stalker.
    */
+  step5 = subtract_ranges(module_ranges, step4);
+  print_ranges("step5", step5);
 
-  step4 = subtract_ranges(module_ranges, step3);
-  print_ranges("step4", step4);
-
-  ranges = merge_ranges(step4);
+  ranges = merge_ranges(step5);
   print_ranges("final", ranges);
 
+  g_array_free(step5, TRUE);
   g_array_free(step4, TRUE);
   g_array_free(step3, TRUE);
   g_array_free(step2, TRUE);
diff --git a/frida_mode/ts/lib/afl.ts b/frida_mode/ts/lib/afl.ts
index 93368dac..67e21beb 100644
--- a/frida_mode/ts/lib/afl.ts
+++ b/frida_mode/ts/lib/afl.ts
@@ -120,6 +120,13 @@ class Afl {
   }
 
   /**
+   * See `AFL_FRIDA_INST_JIT`.
+   */
+  public static setInstrumentJit(): void {
+    Afl.jsApiSetInstrumentJit();
+  }
+
+  /**
    * See `AFL_INST_LIBS`.
    */
   public static setInstrumentLibraries(): void {
@@ -273,6 +280,11 @@ class Afl {
     "void",
     ["pointer"]);
 
+  private static readonly jsApiSetInstrumentJit = Afl.jsApiGetFunction(
+    "js_api_set_instrument_jit",
+    "void",
+    []);
+
   private static readonly jsApiSetInstrumentLibraries = Afl.jsApiGetFunction(
     "js_api_set_instrument_libraries",
     "void",
diff --git a/include/envs.h b/include/envs.h
index f89e8e62..4bab54ce 100644
--- a/include/envs.h
+++ b/include/envs.h
@@ -56,6 +56,7 @@ static char *afl_environment_variables[] = {
     "AFL_FRIDA_DEBUG_MAPS",
     "AFL_FRIDA_EXCLUDE_RANGES",
     "AFL_FRIDA_INST_DEBUG_FILE",
+    "AFL_FRIDA_INST_JIT",
     "AFL_FRIDA_INST_NO_OPTIMIZE",
     "AFL_FRIDA_INST_NO_PREFETCH",
     "AFL_FRIDA_INST_RANGES",