1 files changed, 46 insertions, 3 deletions

diff --git a/frida_mode/src/instrument/instrument.c b/frida_mode/src/instrument/instrument.c
index 9ee7db2d..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>
@@ -20,6 +21,8 @@
 #include "stats.h"
 #include "util.h"
 
+#define FRIDA_DEFAULT_MAP_SIZE (64UL << 10)
+
 gboolean instrument_tracing = false;
 gboolean instrument_optimize = false;
 gboolean instrument_unique = false;
@@ -30,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;
 
@@ -37,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;
 
@@ -230,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);
+        }
       }
 
     }
@@ -240,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);
@@ -266,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) {
@@ -279,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();
@@ -290,6 +298,8 @@ void instrument_config(void) {
 
 void instrument_init(void) {
 
+  if (__afl_map_size == MAP_SIZE) __afl_map_size = FRIDA_DEFAULT_MAP_SIZE;
+
   if (!instrument_is_coverage_optimize_supported()) instrument_optimize = false;
 
   FOKF(cBLU "Instrumentation" cRST " - " cGRN "optimize:" cYEL " [%c]",
@@ -390,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();
@@ -416,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));
+}