3 files changed, 72 insertions, 9 deletions

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