Added framework for logging register contents at the end of each basic block

14 files changed, 169 insertions, 6 deletions

diff --git a/frida_mode/README.md b/frida_mode/README.md
index 1ed368d4..29f7968b 100644
--- a/frida_mode/README.md
+++ b/frida_mode/README.md
@@ -178,6 +178,8 @@ Default is 256Mb.
   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_REGS_FILE` - File to write raw register contents at the start
+  of each block.
   `AFL_FRIDA_INST_TRACE`.
 * `AFL_FRIDA_INST_NO_CACHE` - Don't use a look-up table to cache real to
 instrumented address block translations.
diff --git a/frida_mode/Scripting.md b/frida_mode/Scripting.md
index 8634860b..2b18e200 100644
--- a/frida_mode/Scripting.md
+++ b/frida_mode/Scripting.md
@@ -850,6 +850,14 @@ class Afl {
   static setInstrumentNoOptimize() {
       Afl.jsApiSetInstrumentNoOptimize();
   }
+  /**
+   * See `AFL_FRIDA_INST_REGS_FILE`. This function takes a single `string` as
+   * an argument.
+   */
+  public static setInstrumentRegsFile(file: string): void {
+    const buf = Memory.allocUtf8String(file);
+    Afl.jsApiSetInstrumentRegsFile(buf);
+  }
   /*
     * See `AFL_FRIDA_INST_SEED`
     */
diff --git a/frida_mode/frida.map b/frida_mode/frida.map
index 6726dbbd..8e956460 100644
--- a/frida_mode/frida.map
+++ b/frida_mode/frida.map
@@ -19,6 +19,7 @@
     js_api_set_instrument_libraries;
     js_api_set_instrument_instructions;
     js_api_set_instrument_no_optimize;
+    js_api_set_instrument_regs_file;
     js_api_set_instrument_seed;
     js_api_set_instrument_trace;
     js_api_set_instrument_trace_unique;
diff --git a/frida_mode/include/instrument.h b/frida_mode/include/instrument.h
index cd480202..b85aa571 100644
--- a/frida_mode/include/instrument.h
+++ b/frida_mode/include/instrument.h
@@ -13,6 +13,7 @@ extern gboolean instrument_unique;
 extern guint64  instrument_hash_zero;
 extern char    *instrument_coverage_unstable_filename;
 extern gboolean instrument_coverage_insn;
+extern char *   instrument_regs_filename;
 
 extern gboolean instrument_use_fixed_seed;
 extern guint64  instrument_fixed_seed;
@@ -66,6 +67,8 @@ void instrument_cache_config(void);
 void instrument_cache_init(void);
 void instrument_cache_insert(gpointer real_address, gpointer code_address);
 void instrument_cache(const cs_insn *instr, GumStalkerOutput *output);
+void instrument_write_regs(GumCpuContext *cpu_context, gpointer user_data);
+void instrument_regs_format(int fd, char *format, ...);
 
 #endif
 
diff --git a/frida_mode/src/instrument/instrument.c b/frida_mode/src/instrument/instrument.c
index 35e85ec8..93c498e8 100644
--- a/frida_mode/src/instrument/instrument.c
+++ b/frida_mode/src/instrument/instrument.c
@@ -1,3 +1,4 @@
+#include <fcntl.h>
 #include <unistd.h>
 #include <sys/shm.h>
 #include <sys/mman.h>
@@ -32,6 +33,7 @@ gboolean instrument_use_fixed_seed = FALSE;
 guint64  instrument_fixed_seed = 0;
 char    *instrument_coverage_unstable_filename = NULL;
 gboolean instrument_coverage_insn = FALSE;
+char *   instrument_regs_filename = NULL;
 
 static GumStalkerTransformer *transformer = NULL;
 
@@ -39,6 +41,8 @@ static GumAddress previous_rip = 0;
 static GumAddress previous_end = 0;
 static u8        *edges_notified = NULL;
 
+static int regs_fd = -1;
+
 __thread guint64  instrument_previous_pc;
 __thread guint64 *instrument_previous_pc_addr = NULL;
 
@@ -232,6 +236,10 @@ static void instrument_basic_block(GumStalkerIterator *iterator,
 
         }
 
+        if (unlikely(instrument_regs_filename != NULL)) {
+          gum_stalker_iterator_put_callout(iterator, instrument_write_regs,
+                                           (void *)(size_t)regs_fd, NULL);
+        }
       }
 
     }
@@ -242,8 +250,6 @@ static void instrument_basic_block(GumStalkerIterator *iterator,
 
     }
 
-    instrument_debug_instruction(instr->address, instr->size, output);
-
     if (likely(!excluded)) {
 
       asan_instrument(instr, iterator);
@@ -268,7 +274,6 @@ static void instrument_basic_block(GumStalkerIterator *iterator,
   instrument_flush(output);
   instrument_debug_end(output);
   instrument_coverage_end(instr->address + instr->size);
-
 }
 
 void instrument_config(void) {
@@ -281,6 +286,7 @@ void instrument_config(void) {
   instrument_coverage_unstable_filename =
       (getenv("AFL_FRIDA_INST_UNSTABLE_COVERAGE_FILE"));
   instrument_coverage_insn = (getenv("AFL_FRIDA_INST_INSN") != NULL);
+  instrument_regs_filename = getenv("AFL_FRIDA_INST_REGS_FILE");
 
   instrument_debug_config();
   instrument_coverage_config();
@@ -394,6 +400,23 @@ void instrument_init(void) {
        instrument_hash_seed);
   instrument_hash_zero = instrument_get_offset_hash(0);
 
+  FOKF(cBLU "Instrumentation" cRST " - " cGRN "regs:" cYEL " [%s]",
+       instrument_regs_filename == NULL ? " " : instrument_regs_filename);
+
+  if (instrument_regs_filename != NULL) {
+    char *path =
+        g_canonicalize_filename(instrument_regs_filename, g_get_current_dir());
+
+    FOKF(cBLU "Instrumentation" cRST " - " cGRN "path:" cYEL " [%s]", path);
+
+    regs_fd = open(path, O_RDWR | O_CREAT | O_TRUNC,
+                   S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP);
+
+    if (regs_fd < 0) { FFATAL("Failed to open regs file '%s'", path); }
+
+    g_free(path);
+  }
+
   asan_init();
   cmplog_init();
   instrument_coverage_init();
@@ -420,3 +443,19 @@ void instrument_on_fork() {
 
 }
 
+void instrument_regs_format(int fd, char *format, ...) {
+  va_list ap;
+  char    buffer[4096] = {0};
+  int     ret;
+  int     len;
+
+  va_start(ap, format);
+  ret = vsnprintf(buffer, sizeof(buffer) - 1, format, ap);
+  va_end(ap);
+
+  if (ret < 0) { return; }
+
+  len = strnlen(buffer, sizeof(buffer));
+
+  IGNORED_RETURN(write(fd, buffer, len));
+}
diff --git a/frida_mode/src/instrument/instrument_arm32.c b/frida_mode/src/instrument/instrument_arm32.c
index 572b706c..73923326 100644
--- a/frida_mode/src/instrument/instrument_arm32.c
+++ b/frida_mode/src/instrument/instrument_arm32.c
@@ -80,5 +80,24 @@ void instrument_cache(const cs_insn *instr, GumStalkerOutput *output) {
 
 }
 
+void instrument_write_regs(GumCpuContext *cpu_context, gpointer user_data) {
+  int fd = (int)user_data;
+  instrument_regs_format(fd,
+                         "r0 : 0x%08x, r1 : 0x%08x, r2 : 0x%08x, r3 : 0x%08x\n",
+                         cpu_context->r[0], cpu_context->r[2],
+                         cpu_context->r[1], cpu_context->r[3]);
+  instrument_regs_format(fd,
+                         "r4 : 0x%08x, r5 : 0x%08x, r6 : 0x%08x, r7 : 0x%08x\n",
+                         cpu_context->r[4], cpu_context->r[5],
+                         cpu_context->r[6], cpu_context->r[7]);
+  instrument_regs_format(
+      fd, "r8 : 0x%08x, r9 : 0x%08x, r10: 0x%08x, r11: 0x%08x\n",
+      cpu_context->r8, cpu_context->r9, cpu_context->r10, cpu_context->r11);
+  instrument_regs_format(
+      fd, "r12: 0x%08x, sp : 0x%08x, lr : 0x%08x, pc : 0x%08x\n",
+      cpu_context->r12, cpu_context->sp, cpu_context->lr, cpu_context->pc);
+  instrument_regs_format(fd, "cpsr: 0x%08x\n\n", cpu_context->cpsr);
+}
+
 #endif
 
diff --git a/frida_mode/src/instrument/instrument_arm64.c b/frida_mode/src/instrument/instrument_arm64.c
index 3c37ea5f..9157f8f5 100644
--- a/frida_mode/src/instrument/instrument_arm64.c
+++ b/frida_mode/src/instrument/instrument_arm64.c
@@ -406,5 +406,41 @@ void instrument_cache(const cs_insn *instr, GumStalkerOutput *output) {
 
 }
 
+void instrument_write_regs(GumCpuContext *cpu_context, gpointer user_data) {
+  int fd = (int)(size_t)user_data;
+  instrument_regs_format(
+      fd, "x0 : 0x%016x, x1 : 0x%016x, x2 : 0x%016x, x3 : 0x%016x\n",
+      cpu_context->x[0], cpu_context->x[1], cpu_context->x[2],
+      cpu_context->x[3]);
+  instrument_regs_format(
+      fd, "x4 : 0x%016x, x5 : 0x%016x, x6 : 0x%016x, x7 : 0x%016x\n",
+      cpu_context->x[4], cpu_context->x[5], cpu_context->x[6],
+      cpu_context->x[7]);
+  instrument_regs_format(
+      fd, "x8 : 0x%016x, x9 : 0x%016x, x10: 0x%016x, x11: 0x%016x\n",
+      cpu_context->x[8], cpu_context->x[9], cpu_context->x[10],
+      cpu_context->x[11]);
+  instrument_regs_format(
+      fd, "x12: 0x%016x, x13: 0x%016x, x14: 0x%016x, x15: 0x%016x\n",
+      cpu_context->x[12], cpu_context->x[13], cpu_context->x[14],
+      cpu_context->x[15]);
+  instrument_regs_format(
+      fd, "x16: 0x%016x, x17: 0x%016x, x18: 0x%016x, x19: 0x%016x\n",
+      cpu_context->x[16], cpu_context->x[17], cpu_context->x[18],
+      cpu_context->x[19]);
+  instrument_regs_format(
+      fd, "x20: 0x%016x, x21: 0x%016x, x22: 0x%016x, x23: 0x%016x\n",
+      cpu_context->x[20], cpu_context->x[21], cpu_context->x[22],
+      cpu_context->x[23]);
+  instrument_regs_format(
+      fd, "x24: 0x%016x, x25: 0x%016x, x26: 0x%016x, x27: 0x%016x\n",
+      cpu_context->x[24], cpu_context->x[25], cpu_context->x[26],
+      cpu_context->x[27]);
+  instrument_regs_format(
+      fd, "x28: 0x%016x, fp : 0x%016x, lr : 0x%016x, sp : 0x%016x\n",
+      cpu_context->x[28], cpu_context->fp, cpu_context->lr, cpu_context->sp);
+  instrument_regs_format(fd, "pc : 0x%016x\n\n", cpu_context->pc);
+}
+
 #endif
 
diff --git a/frida_mode/src/instrument/instrument_debug.c b/frida_mode/src/instrument/instrument_debug.c
index d26f9cec..17245d65 100644
--- a/frida_mode/src/instrument/instrument_debug.c
+++ b/frida_mode/src/instrument/instrument_debug.c
@@ -63,14 +63,12 @@ static void instrument_disasm(guint8 *start, guint8 *end,
 
     count = cs_disasm(capstone, curr, size, GPOINTER_TO_SIZE(curr), 0, &insn);
     if (insn == NULL) {
-
       instrument_debug("\t0x%" G_GINT64_MODIFIER "x\t* 0x%016" G_GSIZE_MODIFIER
                        "x\n",
-                       (uint64_t)curr, *(size_t *)curr);
+                       (uint64_t)(size_t)curr, *(size_t *)curr);
 
       len += sizeof(size_t);
       continue;
-
     }
 
     for (i = 0; i != count; i++) {
diff --git a/frida_mode/src/instrument/instrument_x64.c b/frida_mode/src/instrument/instrument_x64.c
index b44c8439..9d754082 100644
--- a/frida_mode/src/instrument/instrument_x64.c
+++ b/frida_mode/src/instrument/instrument_x64.c
@@ -468,5 +468,22 @@ gpointer instrument_cur(GumStalkerOutput *output) {
 
 }
 
+void instrument_write_regs(GumCpuContext *cpu_context, gpointer user_data) {
+  int fd = (int)(size_t)user_data;
+  instrument_regs_format(
+      fd, "rax: 0x%016x, rbx: 0x%016x, rcx: 0x%016x, rdx: 0x%016x\n",
+      cpu_context->rax, cpu_context->rbx, cpu_context->rcx, cpu_context->rdx);
+  instrument_regs_format(
+      fd, "rdi: 0x%016x, rsi: 0x%016x, rbp: 0x%016x, rsp: 0x%016x\n",
+      cpu_context->rdi, cpu_context->rsi, cpu_context->rbp, cpu_context->rsp);
+  instrument_regs_format(
+      fd, "r8 : 0x%016x, r9 : 0x%016x, r10: 0x%016x, r11: 0x%016x\n",
+      cpu_context->r8, cpu_context->r9, cpu_context->r10, cpu_context->r11);
+  instrument_regs_format(
+      fd, "r12: 0x%016x, r13: 0x%016x, r14: 0x%016x, r15: 0x%016x\n",
+      cpu_context->r12, cpu_context->r13, cpu_context->r14, cpu_context->r15);
+  instrument_regs_format(fd, "rip: 0x%016x\n\n", cpu_context->rip);
+}
+
 #endif
 
diff --git a/frida_mode/src/instrument/instrument_x86.c b/frida_mode/src/instrument/instrument_x86.c
index 8798cfcf..eb0c7184 100644
--- a/frida_mode/src/instrument/instrument_x86.c
+++ b/frida_mode/src/instrument/instrument_x86.c
@@ -270,5 +270,16 @@ void instrument_cache(const cs_insn *instr, GumStalkerOutput *output) {
 
 }
 
+void instrument_write_regs(GumCpuContext *cpu_context, gpointer user_data) {
+  int fd = (int)(size_t)user_data;
+  instrument_regs_format(
+      fd, "eax: 0x%08x, ebx: 0x%08x, ecx: 0x%08x, edx: 0x%08x\n",
+      cpu_context->eax, cpu_context->ebx, cpu_context->ecx, cpu_context->edx);
+  instrument_regs_format(
+      fd, "esi: 0x%08x, edi: 0x%08x, ebp: 0x%08x, esp: 0x%08x\n",
+      cpu_context->esi, cpu_context->edi, cpu_context->ebp, cpu_context->esp);
+  instrument_regs_format(fd, "eip: 0x%08x\n\n", cpu_context->eip);
+}
+
 #endif
 
diff --git a/frida_mode/src/js/api.js b/frida_mode/src/js/api.js
index c1c9d36f..721ef82c 100644
--- a/frida_mode/src/js/api.js
+++ b/frida_mode/src/js/api.js
@@ -150,6 +150,14 @@ class Afl {
     static setInstrumentNoOptimize() {
         Afl.jsApiSetInstrumentNoOptimize();
     }
+    /**
+     * See `AFL_FRIDA_INST_REGS_FILE`. This function takes a single `string` as
+     * an argument.
+     */
+    static setInstrumentRegsFile(file) {
+        const buf = Memory.allocUtf8String(file);
+        Afl.jsApiSetInstrumentRegsFile(buf);
+    }
     /*
      * See `AFL_FRIDA_INST_SEED`
      */
@@ -322,6 +330,7 @@ Afl.jsApiSetInstrumentInstructions = Afl.jsApiGetFunction("js_api_set_instrument
 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.jsApiSetInstrumentRegsFile = Afl.jsApiGetFunction("js_api_set_instrument_regs_file", "void", ["pointer"]);
 Afl.jsApiSetInstrumentSeed = Afl.jsApiGetFunction("js_api_set_instrument_seed", "void", ["uint64"]);
 Afl.jsApiSetInstrumentTrace = Afl.jsApiGetFunction("js_api_set_instrument_trace", "void", []);
 Afl.jsApiSetInstrumentTraceUnique = Afl.jsApiGetFunction("js_api_set_instrument_trace_unique", "void", []);
diff --git a/frida_mode/src/js/js_api.c b/frida_mode/src/js/js_api.c
index 7cc8ffc7..d0c0aa60 100644
--- a/frida_mode/src/js/js_api.c
+++ b/frida_mode/src/js/js_api.c
@@ -156,6 +156,11 @@ __attribute__((visibility("default"))) void js_api_set_instrument_no_optimize(
 
 }
 
+__attribute__((visibility("default"))) void js_api_set_instrument_regs_file(
+    char *path) {
+  instrument_regs_filename = g_strdup(path);
+}
+
 __attribute__((visibility("default"))) void js_api_set_instrument_seed(
     guint64 seed) {
 
diff --git a/frida_mode/ts/lib/afl.ts b/frida_mode/ts/lib/afl.ts
index a858f074..455d4305 100644
--- a/frida_mode/ts/lib/afl.ts
+++ b/frida_mode/ts/lib/afl.ts
@@ -178,6 +178,15 @@ class Afl {
     Afl.jsApiSetInstrumentNoOptimize();
   }
 
+  /**
+   * See `AFL_FRIDA_INST_REGS_FILE`. This function takes a single `string` as
+   * an argument.
+   */
+  public static setInstrumentRegsFile(file: string): void {
+    const buf = Memory.allocUtf8String(file);
+    Afl.jsApiSetInstrumentRegsFile(buf);
+  }
+
   /*
    * See `AFL_FRIDA_INST_SEED`
    */
@@ -419,6 +428,11 @@ class Afl {
     "void",
     []);
 
+  private static readonly jsApiSetInstrumentRegsFile = Afl.jsApiGetFunction(
+    "js_api_set_instrument_regs_file",
+    "void",
+    ["pointer"]);
+
   private static readonly jsApiSetInstrumentSeed = Afl.jsApiGetFunction(
     "js_api_set_instrument_seed",
     "void",
diff --git a/include/envs.h b/include/envs.h
index 9b8917f9..853edbd9 100644
--- a/include/envs.h
+++ b/include/envs.h
@@ -67,6 +67,7 @@ static char *afl_environment_variables[] = {
     "AFL_FRIDA_INST_NO_PREFETCH",
     "AFL_FRIDA_INST_NO_PREFETCH_BACKPATCH",
     "AFL_FRIDA_INST_RANGES",
+    "AFL_FRIDA_INST_REGS_FILE",
     "AFL_FRIDA_INST_SEED",
     "AFL_FRIDA_INST_TRACE",
     "AFL_FRIDA_INST_TRACE_UNIQUE",