1 files changed, 47 insertions, 12 deletions

diff --git a/frida_mode/src/instrument/instrument.c b/frida_mode/src/instrument/instrument.c
index fddff19a..fd0982f8 100644
--- a/frida_mode/src/instrument/instrument.c
+++ b/frida_mode/src/instrument/instrument.c
@@ -29,14 +29,23 @@ guint64  instrument_hash_seed = 0;
 
 gboolean instrument_use_fixed_seed = FALSE;
 guint64  instrument_fixed_seed = 0;
+char *   instrument_coverage_unstable_filename = NULL;
 
 static GumStalkerTransformer *transformer = NULL;
 
 __thread guint64 instrument_previous_pc = 0;
 
 static GumAddress previous_rip = 0;
+static GumAddress previous_end = 0;
 static u8 *       edges_notified = NULL;
 
+typedef struct {
+
+  GumAddress address;
+  GumAddress end;
+
+} block_ctx_t;
+
 static void trace_debug(char *format, ...) {
 
   va_list ap;
@@ -91,9 +100,11 @@ __attribute__((hot)) static void on_basic_block(GumCpuContext *context,
 
   UNUSED_PARAMETER(context);
 
-  GumAddress current_rip = GUM_ADDRESS(user_data);
-  guint64    current_pc = instrument_get_offset_hash(current_rip);
-  guint64    edge;
+  block_ctx_t *ctx = (block_ctx_t *)user_data;
+  GumAddress   current_rip = ctx->address;
+  guint16      current_end = ctx->end;
+  guint64      current_pc = instrument_get_offset_hash(current_rip);
+  guint64      edge;
 
   edge = current_pc ^ instrument_previous_pc;
 
@@ -112,10 +123,18 @@ __attribute__((hot)) static void on_basic_block(GumCpuContext *context,
 
     if (instrument_unique) { edges_notified[edge] = 1; }
 
-    previous_rip = current_rip;
+  }
+
+  if (unlikely(instrument_coverage_unstable_filename != NULL)) {
+
+    instrument_coverage_unstable(edge, previous_rip, previous_end, current_rip,
+                                 current_end);
 
   }
 
+  previous_rip = current_rip;
+  previous_end = current_end;
+
   instrument_previous_pc = ((current_pc & (MAP_SIZE - 1) >> 1)) |
                            ((current_pc & 0x1) << (MAP_SIZE_POW2 - 1));
 
@@ -130,6 +149,7 @@ static void instrument_basic_block(GumStalkerIterator *iterator,
   const cs_insn *instr;
   gboolean       begin = TRUE;
   gboolean       excluded;
+  block_ctx_t *  ctx = NULL;
 
   while (gum_stalker_iterator_next(iterator, &instr)) {
 
@@ -164,19 +184,16 @@ static void instrument_basic_block(GumStalkerIterator *iterator,
      * our AFL_ENTRYPOINT, since it is not until then that we start the
      * fork-server and thus start executing in the child.
      */
-    excluded = range_is_excluded(GSIZE_TO_POINTER(instr->address));
+    excluded = range_is_excluded(GUM_ADDRESS(instr->address));
 
     stats_collect(instr, begin);
 
     if (unlikely(begin)) {
 
       instrument_debug_start(instr->address, output);
+      instrument_coverage_start(instr->address);
 
-      if (likely(entry_reached)) {
-
-        prefetch_write(GSIZE_TO_POINTER(instr->address));
-
-      }
+      prefetch_write(GSIZE_TO_POINTER(instr->address));
 
       if (likely(!excluded)) {
 
@@ -186,8 +203,9 @@ static void instrument_basic_block(GumStalkerIterator *iterator,
 
         } else {
 
-          gum_stalker_iterator_put_callout(
-              iterator, on_basic_block, GSIZE_TO_POINTER(instr->address), NULL);
+          ctx = gum_malloc0(sizeof(block_ctx_t));
+          ctx->address = GUM_ADDRESS(instr->address);
+          gum_stalker_iterator_put_callout(iterator, on_basic_block, ctx, NULL);
 
         }
 
@@ -214,8 +232,11 @@ static void instrument_basic_block(GumStalkerIterator *iterator,
 
   }
 
+  if (ctx != NULL) { ctx->end = (instr->address + instr->size); }
+
   instrument_flush(output);
   instrument_debug_end(output);
+  instrument_coverage_end(instr->address + instr->size);
 
 }
 
@@ -226,8 +247,11 @@ void instrument_config(void) {
   instrument_unique = (getenv("AFL_FRIDA_INST_TRACE_UNIQUE") != NULL);
   instrument_use_fixed_seed = (getenv("AFL_FRIDA_INST_SEED") != NULL);
   instrument_fixed_seed = util_read_num("AFL_FRIDA_INST_SEED");
+  instrument_coverage_unstable_filename =
+      (getenv("AFL_FRIDA_INST_UNSTABLE_COVERAGE_FILE"));
 
   instrument_debug_config();
+  instrument_coverage_config();
   asan_config();
   cmplog_config();
 
@@ -242,6 +266,9 @@ void instrument_init(void) {
   OKF("Instrumentation - unique [%c]", instrument_unique ? 'X' : ' ');
   OKF("Instrumentation - fixed seed [%c] [0x%016" G_GINT64_MODIFIER "x]",
       instrument_use_fixed_seed ? 'X' : ' ', instrument_fixed_seed);
+  OKF("Instrumentation - unstable coverage [%c] [%s]",
+      instrument_coverage_unstable_filename == NULL ? ' ' : 'X',
+      instrument_coverage_unstable_filename);
 
   if (instrument_tracing && instrument_optimize) {
 
@@ -250,6 +277,13 @@ void instrument_init(void) {
 
   }
 
+  if (instrument_coverage_unstable_filename && instrument_optimize) {
+
+    WARNF("AFL_FRIDA_INST_COVERAGE_FILE implies AFL_FRIDA_INST_NO_OPTIMIZE");
+    instrument_optimize = FALSE;
+
+  }
+
   if (instrument_unique && instrument_optimize) {
 
     WARNF("AFL_FRIDA_INST_TRACE_UNIQUE implies AFL_FRIDA_INST_NO_OPTIMIZE");
@@ -317,6 +351,7 @@ void instrument_init(void) {
   instrument_hash_zero = instrument_get_offset_hash(0);
 
   instrument_debug_init();
+  instrument_coverage_init();
   asan_init();
   cmplog_init();