about summary refs log tree commit diff
path: root/frida_mode/src/instrument
diff options
context:
space:
mode:
Diffstat (limited to 'frida_mode/src/instrument')
-rw-r--r--frida_mode/src/instrument/instrument.c242
-rw-r--r--frida_mode/src/instrument/instrument_arm32.c14
-rw-r--r--frida_mode/src/instrument/instrument_arm64.c32
-rw-r--r--frida_mode/src/instrument/instrument_debug.c73
-rw-r--r--frida_mode/src/instrument/instrument_x64.c52
-rw-r--r--frida_mode/src/instrument/instrument_x86.c38
6 files changed, 335 insertions, 116 deletions
diff --git a/frida_mode/src/instrument/instrument.c b/frida_mode/src/instrument/instrument.c
index cd1ac0be..67aafa5a 100644
--- a/frida_mode/src/instrument/instrument.c
+++ b/frida_mode/src/instrument/instrument.c
@@ -1,14 +1,19 @@
 #include <unistd.h>
+#include <sys/shm.h>
+#include <sys/mman.h>
+#include <sys/syscall.h>
 
-#include "frida-gum.h"
+#include "frida-gumjs.h"
 
 #include "config.h"
 #include "debug.h"
+#include "hash.h"
 
 #include "asan.h"
 #include "entry.h"
 #include "frida_cmplog.h"
 #include "instrument.h"
+#include "js.h"
 #include "persistent.h"
 #include "prefetch.h"
 #include "ranges.h"
@@ -16,46 +21,55 @@
 #include "stats.h"
 #include "util.h"
 
-static gboolean               tracing = false;
-static gboolean               optimize = false;
+gboolean instrument_tracing = false;
+gboolean instrument_optimize = false;
+gboolean instrument_unique = false;
+guint64  instrument_hash_zero = 0;
+guint64  instrument_hash_seed = 0;
+
+gboolean instrument_use_fixed_seed = FALSE;
+guint64  instrument_fixed_seed = 0;
+
 static GumStalkerTransformer *transformer = NULL;
 
-__thread uint64_t previous_pc = 0;
+__thread guint64 instrument_previous_pc = 0;
 
-__attribute__((hot)) static void on_basic_block(GumCpuContext *context,
-                                                gpointer       user_data) {
+static GumAddress previous_rip = 0;
+static u8 *       edges_notified = NULL;
 
-  UNUSED_PARAMETER(context);
-  /*
-   * This function is performance critical as it is called to instrument every
-   * basic block. By moving our print buffer to a global, we avoid it affecting
-   * the critical path with additional stack adjustments if tracing is not
-   * enabled. If tracing is enabled, then we're printing a load of diagnostic
-   * information so this overhead is unlikely to be noticeable.
-   */
-  static char buffer[200];
-  int         len;
-  GumAddress  current_pc = GUM_ADDRESS(user_data);
-  uint8_t *   cursor;
-  uint64_t    value;
-  if (unlikely(tracing)) {
-
-    /* Avoid any functions which may cause an allocation since the target app
-     * may already be running inside malloc and it isn't designed to be
-     * re-entrant on a single thread */
-    len = snprintf(buffer, sizeof(buffer),
-                   "current_pc: 0x%016" G_GINT64_MODIFIER
-                   "x, previous_pc: 0x%016" G_GINT64_MODIFIER "x\n",
-                   current_pc, previous_pc);
-
-    IGNORED_RETURN(write(STDOUT_FILENO, buffer, len + 1));
+static void trace_debug(char *format, ...) {
 
-  }
+  va_list ap;
+  char    buffer[4096] = {0};
+  int     ret;
+  int     len;
 
-  current_pc = (current_pc >> 4) ^ (current_pc << 8);
-  current_pc &= MAP_SIZE - 1;
+  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(STDOUT_FILENO, buffer, len));
+
+}
 
-  cursor = &__afl_area_ptr[current_pc ^ previous_pc];
+guint64 instrument_get_offset_hash(GumAddress current_rip) {
+
+  guint64 area_offset = hash64((unsigned char *)&current_rip,
+                               sizeof(GumAddress), instrument_hash_seed);
+  return area_offset &= MAP_SIZE - 1;
+
+}
+
+__attribute__((hot)) static void instrument_increment_map(GumAddress edge) {
+
+  uint8_t *cursor;
+  uint64_t value;
+
+  cursor = &__afl_area_ptr[edge];
   value = *cursor;
 
   if (value == 0xff) {
@@ -69,12 +83,47 @@ __attribute__((hot)) static void on_basic_block(GumCpuContext *context,
   }
 
   *cursor = value;
-  previous_pc = current_pc >> 1;
 
 }
 
-static void instr_basic_block(GumStalkerIterator *iterator,
-                              GumStalkerOutput *output, gpointer user_data) {
+__attribute__((hot)) static void on_basic_block(GumCpuContext *context,
+                                                gpointer       user_data) {
+
+  UNUSED_PARAMETER(context);
+
+  GumAddress current_rip = GUM_ADDRESS(user_data);
+  guint64    current_pc = instrument_get_offset_hash(current_rip);
+  guint64    edge;
+
+  edge = current_pc ^ instrument_previous_pc;
+
+  instrument_increment_map(edge);
+
+  if (unlikely(instrument_tracing)) {
+
+    if (!instrument_unique || edges_notified[edge] == 0) {
+
+      trace_debug("TRACE: edge: %10" G_GINT64_MODIFIER
+                  "d, current_rip: 0x%016" G_GINT64_MODIFIER
+                  "x, previous_rip: 0x%016" G_GINT64_MODIFIER "x\n",
+                  edge, current_rip, previous_rip);
+
+    }
+
+    if (instrument_unique) { edges_notified[edge] = 1; }
+
+    previous_rip = current_rip;
+
+  }
+
+  instrument_previous_pc =
+      ((current_pc & (MAP_SIZE - 1) >> 1)) | ((current_pc & 0x1) << 15);
+
+}
+
+static void instrument_basic_block(GumStalkerIterator *iterator,
+                                   GumStalkerOutput *  output,
+                                   gpointer            user_data) {
 
   UNUSED_PARAMETER(user_data);
 
@@ -84,7 +133,9 @@ static void instr_basic_block(GumStalkerIterator *iterator,
 
   while (gum_stalker_iterator_next(iterator, &instr)) {
 
-    if (instr->address == entry_start) { entry_prologue(iterator, output); }
+    if (unlikely(begin)) { instrument_debug_start(instr->address, output); }
+
+    if (instr->address == entry_point) { entry_prologue(iterator, output); }
     if (instr->address == persistent_start) { persistent_prologue(output); }
     if (instr->address == persistent_ret) { persistent_epilogue(output); }
 
@@ -121,11 +172,15 @@ static void instr_basic_block(GumStalkerIterator *iterator,
 
       instrument_debug_start(instr->address, output);
 
-      prefetch_write(GSIZE_TO_POINTER(instr->address));
+      if (likely(entry_reached)) {
+
+        prefetch_write(GSIZE_TO_POINTER(instr->address));
+
+      }
 
       if (likely(!excluded)) {
 
-        if (likely(optimize)) {
+        if (likely(instrument_optimize)) {
 
           instrument_coverage_optimize(instr, output);
 
@@ -138,8 +193,6 @@ static void instr_basic_block(GumStalkerIterator *iterator,
 
       }
 
-      begin = FALSE;
-
     }
 
     instrument_debug_instruction(instr->address, instr->size);
@@ -151,38 +204,117 @@ static void instr_basic_block(GumStalkerIterator *iterator,
 
     }
 
-    gum_stalker_iterator_keep(iterator);
+    if (js_stalker_callback(instr, begin, excluded, output)) {
+
+      gum_stalker_iterator_keep(iterator);
+
+    }
+
+    begin = FALSE;
 
   }
 
+  instrument_flush(output);
   instrument_debug_end(output);
 
 }
 
+void instrument_config(void) {
+
+  instrument_optimize = (getenv("AFL_FRIDA_INST_NO_OPTIMIZE") == NULL);
+  instrument_tracing = (getenv("AFL_FRIDA_INST_TRACE") != NULL);
+  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_debug_config();
+  asan_config();
+  cmplog_config();
+
+}
+
 void instrument_init(void) {
 
-  optimize = (getenv("AFL_FRIDA_INST_NO_OPTIMIZE") == NULL);
-  tracing = (getenv("AFL_FRIDA_INST_TRACE") != NULL);
+  if (!instrument_is_coverage_optimize_supported()) instrument_optimize = false;
 
-  if (!instrument_is_coverage_optimize_supported()) optimize = false;
+  OKF("Instrumentation - optimize [%c]", instrument_optimize ? 'X' : ' ');
+  OKF("Instrumentation - tracing [%c]", instrument_tracing ? 'X' : ' ');
+  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 - optimize [%c]", optimize ? 'X' : ' ');
-  OKF("Instrumentation - tracing [%c]", tracing ? 'X' : ' ');
+  if (instrument_tracing && instrument_optimize) {
 
-  if (tracing && optimize) {
+    WARNF("AFL_FRIDA_INST_TRACE implies AFL_FRIDA_INST_NO_OPTIMIZE");
+    instrument_optimize = FALSE;
+
+  }
 
-    FATAL("AFL_FRIDA_INST_OPTIMIZE and AFL_FRIDA_INST_TRACE are incompatible");
+  if (instrument_unique && instrument_optimize) {
+
+    WARNF("AFL_FRIDA_INST_TRACE_UNIQUE implies AFL_FRIDA_INST_NO_OPTIMIZE");
+    instrument_optimize = FALSE;
 
   }
 
+  if (instrument_unique) { instrument_tracing = TRUE; }
+
   if (__afl_map_size != 0x10000) {
 
     FATAL("Bad map size: 0x%08x", __afl_map_size);
 
   }
 
-  transformer =
-      gum_stalker_transformer_make_from_callback(instr_basic_block, NULL, NULL);
+  transformer = gum_stalker_transformer_make_from_callback(
+      instrument_basic_block, NULL, NULL);
+
+  if (instrument_unique) {
+
+    int shm_id = shmget(IPC_PRIVATE, MAP_SIZE, IPC_CREAT | IPC_EXCL | 0600);
+    if (shm_id < 0) { FATAL("shm_id < 0 - errno: %d\n", errno); }
+
+    edges_notified = shmat(shm_id, NULL, 0);
+    g_assert(edges_notified != MAP_FAILED);
+
+    /*
+     * Configure the shared memory region to be removed once the process
+     * dies.
+     */
+    if (shmctl(shm_id, IPC_RMID, NULL) < 0) {
+
+      FATAL("shmctl (IPC_RMID) < 0 - errno: %d\n", errno);
+
+    }
+
+    /* Clear it, not sure it's necessary, just seems like good practice */
+    memset(edges_notified, '\0', MAP_SIZE);
+
+  }
+
+  if (instrument_use_fixed_seed) {
+
+    /*
+     * This configuration option may be useful for diagnostics or
+     * debugging.
+     */
+    instrument_hash_seed = instrument_fixed_seed;
+
+  } else {
+
+    /*
+     * By using a different seed value for the hash, we can make different
+     * instances have edge collisions in different places when carrying out
+     * parallel fuzzing. The seed itself, doesn't have to be random, it
+     * just needs to be different for each instance.
+     */
+    instrument_hash_seed = g_get_monotonic_time() ^
+                           (((guint64)getpid()) << 32) ^ syscall(SYS_gettid);
+
+  }
+
+  OKF("Instrumentation - seed [0x%016" G_GINT64_MODIFIER "x]",
+      instrument_hash_seed);
+  instrument_hash_zero = instrument_get_offset_hash(0);
 
   instrument_debug_init();
   asan_init();
@@ -197,3 +329,9 @@ GumStalkerTransformer *instrument_get_transformer(void) {
 
 }
 
+void instrument_on_fork() {
+
+  instrument_previous_pc = instrument_hash_zero;
+
+}
+
diff --git a/frida_mode/src/instrument/instrument_arm32.c b/frida_mode/src/instrument/instrument_arm32.c
index 1a3c40bb..0e15940a 100644
--- a/frida_mode/src/instrument/instrument_arm32.c
+++ b/frida_mode/src/instrument/instrument_arm32.c
@@ -1,4 +1,4 @@
-#include "frida-gum.h"
+#include "frida-gumjs.h"
 
 #include "debug.h"
 
@@ -22,5 +22,17 @@ void instrument_coverage_optimize(const cs_insn *   instr,
 
 }
 
+void instrument_flush(GumStalkerOutput *output) {
+
+  gum_arm_writer_flush(output->writer.arm);
+
+}
+
+gpointer instrument_cur(GumStalkerOutput *output) {
+
+  return gum_arm_writer_cur(output->writer.arm);
+
+}
+
 #endif
 
diff --git a/frida_mode/src/instrument/instrument_arm64.c b/frida_mode/src/instrument/instrument_arm64.c
index fa3afb48..cf37e048 100644
--- a/frida_mode/src/instrument/instrument_arm64.c
+++ b/frida_mode/src/instrument/instrument_arm64.c
@@ -1,4 +1,4 @@
-#include "frida-gum.h"
+#include "frida-gumjs.h"
 
 #include "config.h"
 #include "debug.h"
@@ -12,15 +12,15 @@ static GumAddress current_log_impl = GUM_ADDRESS(0);
 static const guint8 afl_log_code[] = {
 
     // __afl_area_ptr[current_pc ^ previous_pc]++;
-    // previous_pc = current_pc >> 1;
+    // previous_pc = current_pc ROR 1;
     0xE1, 0x0B, 0xBF, 0xA9,  // stp x1, x2, [sp, -0x10]!
     0xE3, 0x13, 0xBF, 0xA9,  // stp x3, x4, [sp, -0x10]!
 
     // x0 = current_pc
-    0xe1, 0x01, 0x00, 0x58,  // ldr x1, #0x3c, =&__afl_area_ptr
+    0x21, 0x02, 0x00, 0x58,  // ldr x1, #0x44, =&__afl_area_ptr
     0x21, 0x00, 0x40, 0xf9,  // ldr x1, [x1] (=__afl_area_ptr)
 
-    0xe2, 0x01, 0x00, 0x58,  // ldr x2, #0x3c, =&previous_pc
+    0x22, 0x02, 0x00, 0x58,  // ldr x2, #0x44, =&previous_pc
     0x42, 0x00, 0x40, 0xf9,  // ldr x2, [x2] (=previous_pc)
 
     // __afl_area_ptr[current_pc ^ previous_pc]++;
@@ -30,8 +30,11 @@ static const guint8 afl_log_code[] = {
     0x63, 0x00, 0x1f, 0x9a,  // adc x3, x3, xzr
     0x23, 0x68, 0x22, 0xf8,  // str x3, [x1, x2]
 
-    // previous_pc = current_pc >> 1;
-    0xe0, 0x07, 0x40, 0x8b,  // add x0, xzr, x0, LSR #1
+    // previous_pc = current_pc ROR 1;
+    0xe4, 0x07, 0x40, 0x8b,  // add x4, xzr, x0, LSR #1
+    0xe0, 0xff, 0x00, 0x8b,  // add x0, xzr, x0, LSL #63
+    0x80, 0xc0, 0x40, 0x8b,  // add x0, x4, x0, LSR #48
+
     0xe2, 0x00, 0x00, 0x58,  // ldr x2, #0x1c, =&previous_pc
     0x40, 0x00, 0x00, 0xf9,  // str x0, [x2]
 
@@ -54,8 +57,7 @@ void instrument_coverage_optimize(const cs_insn *   instr,
                                   GumStalkerOutput *output) {
 
   guint64 current_pc = instr->address;
-  guint64 area_offset = (current_pc >> 4) ^ (current_pc << 8);
-  area_offset &= MAP_SIZE - 1;
+  guint64 area_offset = instrument_get_offset_hash(GUM_ADDRESS(instr->address));
   GumArm64Writer *cw = output->writer.arm64;
 
   if (current_log_impl == 0 ||
@@ -72,7 +74,7 @@ void instrument_coverage_optimize(const cs_insn *   instr,
     gum_arm64_writer_put_bytes(cw, afl_log_code, sizeof(afl_log_code));
 
     uint8_t **afl_area_ptr_ptr = &__afl_area_ptr;
-    uint64_t *afl_prev_loc_ptr = &previous_pc;
+    uint64_t *afl_prev_loc_ptr = &instrument_previous_pc;
     gum_arm64_writer_put_bytes(cw, (const guint8 *)&afl_area_ptr_ptr,
                                sizeof(afl_area_ptr_ptr));
     gum_arm64_writer_put_bytes(cw, (const guint8 *)&afl_prev_loc_ptr,
@@ -93,5 +95,17 @@ void instrument_coverage_optimize(const cs_insn *   instr,
 
 }
 
+void instrument_flush(GumStalkerOutput *output) {
+
+  gum_arm64_writer_flush(output->writer.arm64);
+
+}
+
+gpointer instrument_cur(GumStalkerOutput *output) {
+
+  return gum_arm64_writer_cur(output->writer.arm64);
+
+}
+
 #endif
 
diff --git a/frida_mode/src/instrument/instrument_debug.c b/frida_mode/src/instrument/instrument_debug.c
index f8c1df77..b8cca634 100644
--- a/frida_mode/src/instrument/instrument_debug.c
+++ b/frida_mode/src/instrument/instrument_debug.c
@@ -3,15 +3,18 @@
 #include <stdio.h>
 #include <unistd.h>
 
-#include "frida-gum.h"
+#include "frida-gumjs.h"
 
 #include "debug.h"
 
+#include "instrument.h"
 #include "util.h"
 
 static int      debugging_fd = -1;
 static gpointer instrument_gen_start = NULL;
 
+char *instrument_debug_filename = NULL;
+
 static void instrument_debug(char *format, ...) {
 
   va_list ap;
@@ -31,24 +34,44 @@ static void instrument_debug(char *format, ...) {
 
 }
 
-static void instrument_disasm(guint8 *code, guint size) {
+static void instrument_disasm(guint8 *start, guint8 *end) {
 
   csh      capstone;
   cs_err   err;
+  uint16_t size;
   cs_insn *insn;
-  size_t   count, i;
+  size_t   count = 0;
+  size_t   i;
+  uint16_t len;
 
   err = cs_open(GUM_DEFAULT_CS_ARCH,
                 GUM_DEFAULT_CS_MODE | GUM_DEFAULT_CS_ENDIAN, &capstone);
   g_assert(err == CS_ERR_OK);
 
-  count = cs_disasm(capstone, code, size, GPOINTER_TO_SIZE(code), 0, &insn);
-  g_assert(insn != NULL);
+  size = GPOINTER_TO_SIZE(end) - GPOINTER_TO_SIZE(start);
+
+  for (guint8 *curr = start; curr < end; curr += len, size -= len, len = 0) {
+
+    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",
+                       curr, *(size_t *)curr);
+
+      len += sizeof(size_t);
+      continue;
+
+    }
+
+    for (i = 0; i != count; i++) {
+
+      instrument_debug("\t0x%" G_GINT64_MODIFIER "x\t%s %s\n", insn[i].address,
+                       insn[i].mnemonic, insn[i].op_str);
 
-  for (i = 0; i != count; i++) {
+      len += insn[i].size;
 
-    instrument_debug("\t0x%" G_GINT64_MODIFIER "x\t%s %s\n", insn[i].address,
-                     insn[i].mnemonic, insn[i].op_str);
+    }
 
   }
 
@@ -58,32 +81,25 @@ static void instrument_disasm(guint8 *code, guint size) {
 
 }
 
-static gpointer instrument_cur(GumStalkerOutput *output) {
+void instrument_debug_config(void) {
 
-#if defined(__i386__) || defined(__x86_64__)
-  return gum_x86_writer_cur(output->writer.x86);
-#elif defined(__aarch64__)
-  return gum_arm64_writer_cur(output->writer.arm64);
-#elif defined(__arm__)
-  return gum_arm_writer_cur(output->writer.arm);
-#else
-  #error "Unsupported architecture"
-#endif
+  instrument_debug_filename = getenv("AFL_FRIDA_INST_DEBUG_FILE");
 
 }
 
 void instrument_debug_init(void) {
 
-  char *filename = getenv("AFL_FRIDA_INST_DEBUG_FILE");
-  OKF("Instrumentation debugging - enabled [%c]", filename == NULL ? ' ' : 'X');
+  OKF("Instrumentation debugging - enabled [%c]",
+      instrument_debug_filename == NULL ? ' ' : 'X');
 
-  if (filename == NULL) { return; }
+  if (instrument_debug_filename == NULL) { return; }
 
-  OKF("Instrumentation debugging - file [%s]", filename);
+  OKF("Instrumentation debugging - file [%s]", instrument_debug_filename);
 
-  if (filename == NULL) { return; }
+  if (instrument_debug_filename == NULL) { return; }
 
-  char *path = g_canonicalize_filename(filename, g_get_current_dir());
+  char *path =
+      g_canonicalize_filename(instrument_debug_filename, g_get_current_dir());
 
   OKF("Instrumentation debugging - path [%s]", path);
 
@@ -111,7 +127,7 @@ void instrument_debug_instruction(uint64_t address, uint16_t size) {
 
   if (likely(debugging_fd < 0)) { return; }
   uint8_t *start = (uint8_t *)GSIZE_TO_POINTER(address);
-  instrument_disasm(start, size);
+  instrument_disasm(start, start + size);
 
 }
 
@@ -119,11 +135,10 @@ void instrument_debug_end(GumStalkerOutput *output) {
 
   if (likely(debugging_fd < 0)) { return; }
   gpointer instrument_gen_end = instrument_cur(output);
-  uint16_t size = GPOINTER_TO_SIZE(instrument_gen_end) -
-                  GPOINTER_TO_SIZE(instrument_gen_start);
 
-  instrument_debug("\nGenerated block %p\n", instrument_gen_start);
-  instrument_disasm(instrument_gen_start, size);
+  instrument_debug("\nGenerated block %p-%p\n", instrument_gen_start,
+                   instrument_gen_end);
+  instrument_disasm(instrument_gen_start, instrument_gen_end);
 
 }
 
diff --git a/frida_mode/src/instrument/instrument_x64.c b/frida_mode/src/instrument/instrument_x64.c
index 901f3bd0..fec8afbb 100644
--- a/frida_mode/src/instrument/instrument_x64.c
+++ b/frida_mode/src/instrument/instrument_x64.c
@@ -1,4 +1,4 @@
-#include "frida-gum.h"
+#include "frida-gumjs.h"
 
 #include "config.h"
 
@@ -10,23 +10,21 @@ static GumAddress current_log_impl = GUM_ADDRESS(0);
 
 static const guint8 afl_log_code[] = {
 
-    // 0xcc,
-
     0x9c,                                                         /* pushfq */
     0x51,                                                       /* push rcx */
     0x52,                                                       /* push rdx */
 
-    0x48, 0x8b, 0x0d, 0x28,
+    0x48, 0x8b, 0x0d, 0x26,
     0x00, 0x00, 0x00,                          /* mov rcx, sym.&previous_pc */
     0x48, 0x8b, 0x11,                               /* mov rdx, qword [rcx] */
     0x48, 0x31, 0xfa,                                       /* xor rdx, rdi */
 
-    0x48, 0x03, 0x15, 0x13,
+    0x48, 0x03, 0x15, 0x11,
     0x00, 0x00, 0x00,                     /* add rdx, sym._afl_area_ptr_ptr */
 
     0x80, 0x02, 0x01,                              /* add byte ptr [rdx], 1 */
     0x80, 0x12, 0x00,                              /* adc byte ptr [rdx], 0 */
-    0x48, 0xd1, 0xef,                                         /* shr rdi, 1 */
+    0x66, 0xd1, 0xcf,                                          /* ror di, 1 */
     0x48, 0x89, 0x39,                               /* mov qword [rcx], rdi */
 
     0x5a,                                                        /* pop rdx */
@@ -34,7 +32,8 @@ static const guint8 afl_log_code[] = {
     0x9d,                                                          /* popfq */
 
     0xc3,                                                            /* ret */
-    0x90, 0x90, 0x90                                             /* nop pad */
+
+    0x90
 
     /* Read-only data goes here: */
     /* uint8_t* __afl_area_ptr */
@@ -48,12 +47,11 @@ gboolean instrument_is_coverage_optimize_supported(void) {
 
 }
 
-void instrument_coverage_optimize(const cs_insn *   instr,
-                                  GumStalkerOutput *output) {
+static guint8 align_pad[] = {0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90};
+
+static void instrument_coverate_write_function(GumStalkerOutput *output) {
 
-  guint64 current_pc = instr->address;
-  guint64 area_offset = (current_pc >> 4) ^ (current_pc << 8);
-  area_offset &= MAP_SIZE - 1;
+  guint64       misalign = 0;
   GumX86Writer *cw = output->writer.x86;
 
   if (current_log_impl == 0 ||
@@ -65,10 +63,17 @@ void instrument_coverage_optimize(const cs_insn *   instr,
 
     gum_x86_writer_put_jmp_near_label(cw, after_log_impl);
 
+    misalign = (cw->pc & 0x7);
+    if (misalign != 0) {
+
+      gum_x86_writer_put_bytes(cw, align_pad, 8 - misalign);
+
+    }
+
     current_log_impl = cw->pc;
     gum_x86_writer_put_bytes(cw, afl_log_code, sizeof(afl_log_code));
 
-    uint64_t *afl_prev_loc_ptr = &previous_pc;
+    uint64_t *afl_prev_loc_ptr = &instrument_previous_pc;
     gum_x86_writer_put_bytes(cw, (const guint8 *)&__afl_area_ptr,
                              sizeof(__afl_area_ptr));
     gum_x86_writer_put_bytes(cw, (const guint8 *)&afl_prev_loc_ptr,
@@ -78,6 +83,15 @@ void instrument_coverage_optimize(const cs_insn *   instr,
 
   }
 
+}
+
+void instrument_coverage_optimize(const cs_insn *   instr,
+                                  GumStalkerOutput *output) {
+
+  GumX86Writer *cw = output->writer.x86;
+  guint64 area_offset = instrument_get_offset_hash(GUM_ADDRESS(instr->address));
+  instrument_coverate_write_function(output);
+
   gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_REG_RSP, GUM_REG_RSP,
                                         -GUM_RED_ZONE_SIZE);
   gum_x86_writer_put_push_reg(cw, GUM_REG_RDI);
@@ -89,5 +103,17 @@ void instrument_coverage_optimize(const cs_insn *   instr,
 
 }
 
+void instrument_flush(GumStalkerOutput *output) {
+
+  gum_x86_writer_flush(output->writer.x86);
+
+}
+
+gpointer instrument_cur(GumStalkerOutput *output) {
+
+  return gum_x86_writer_cur(output->writer.x86);
+
+}
+
 #endif
 
diff --git a/frida_mode/src/instrument/instrument_x86.c b/frida_mode/src/instrument/instrument_x86.c
index 585bb5b8..7bf48f96 100644
--- a/frida_mode/src/instrument/instrument_x86.c
+++ b/frida_mode/src/instrument/instrument_x86.c
@@ -1,4 +1,4 @@
-#include "frida-gum.h"
+#include "frida-gumjs.h"
 
 #include "debug.h"
 
@@ -16,7 +16,7 @@ static void instrument_coverage_function(GumX86Writer *cw) {
   gum_x86_writer_put_push_reg(cw, GUM_REG_EDX);
 
   gum_x86_writer_put_mov_reg_address(cw, GUM_REG_ECX,
-                                     GUM_ADDRESS(&previous_pc));
+                                     GUM_ADDRESS(&instrument_previous_pc));
   gum_x86_writer_put_mov_reg_reg_ptr(cw, GUM_REG_EDX, GUM_REG_ECX);
   gum_x86_writer_put_xor_reg_reg(cw, GUM_REG_EDX, GUM_REG_EDI);
 
@@ -30,7 +30,8 @@ static void instrument_coverage_function(GumX86Writer *cw) {
   uint8_t adc_byte_ptr_edx_0[] = {0x80, 0x12, 0x00};
   gum_x86_writer_put_bytes(cw, adc_byte_ptr_edx_0, sizeof(adc_byte_ptr_edx_0));
 
-  gum_x86_writer_put_shr_reg_u8(cw, GUM_REG_EDI, 1);
+  uint8_t ror_di_1[] = {0x66, 0xd1, 0xcf};
+  gum_x86_writer_put_bytes(cw, ror_di_1, sizeof(ror_di_1));
   gum_x86_writer_put_mov_reg_ptr_reg(cw, GUM_REG_ECX, GUM_REG_EDI);
 
   gum_x86_writer_put_pop_reg(cw, GUM_REG_EDX);
@@ -46,15 +47,8 @@ gboolean instrument_is_coverage_optimize_supported(void) {
 
 }
 
-void instrument_coverage_optimize(const cs_insn *   instr,
-                                  GumStalkerOutput *output) {
-
-  UNUSED_PARAMETER(instr);
-  UNUSED_PARAMETER(output);
+static void instrument_coverate_write_function(GumStalkerOutput *output) {
 
-  guint64 current_pc = instr->address;
-  guint64 area_offset = (current_pc >> 4) ^ (current_pc << 8);
-  area_offset &= MAP_SIZE - 1;
   GumX86Writer *cw = output->writer.x86;
 
   if (current_log_impl == 0 ||
@@ -73,7 +67,15 @@ void instrument_coverage_optimize(const cs_insn *   instr,
 
   }
 
-  // gum_x86_writer_put_breakpoint(cw);
+}
+
+void instrument_coverage_optimize(const cs_insn *   instr,
+                                  GumStalkerOutput *output) {
+
+  GumX86Writer *cw = output->writer.x86;
+  guint64 area_offset = instrument_get_offset_hash(GUM_ADDRESS(instr->address));
+  instrument_coverate_write_function(output);
+
   gum_x86_writer_put_push_reg(cw, GUM_REG_EDI);
   gum_x86_writer_put_mov_reg_address(cw, GUM_REG_EDI, area_offset);
   gum_x86_writer_put_call_address(cw, current_log_impl);
@@ -81,5 +83,17 @@ void instrument_coverage_optimize(const cs_insn *   instr,
 
 }
 
+void instrument_flush(GumStalkerOutput *output) {
+
+  gum_x86_writer_flush(output->writer.x86);
+
+}
+
+gpointer instrument_cur(GumStalkerOutput *output) {
+
+  return gum_x86_writer_cur(output->writer.x86);
+
+}
+
 #endif