13 files changed, 763 insertions, 50 deletions

diff --git a/frida_mode/src/asan/asan_arm.c b/frida_mode/src/asan/asan_arm.c
index 526017be..79475ced 100644
--- a/frida_mode/src/asan/asan_arm.c
+++ b/frida_mode/src/asan/asan_arm.c
@@ -18,5 +18,11 @@ void asan_instrument(const cs_insn *instr, GumStalkerIterator *iterator) {
 
 }
 
+void asan_arch_init(void) {
+
+  FATAL("ASAN mode not supported on this architecture");
+
+}
+
 #endif
 
diff --git a/frida_mode/src/asan/asan_arm64.c b/frida_mode/src/asan/asan_arm64.c
index 4e3fbafd..6262ee18 100644
--- a/frida_mode/src/asan/asan_arm64.c
+++ b/frida_mode/src/asan/asan_arm64.c
@@ -18,5 +18,11 @@ void asan_instrument(const cs_insn *instr, GumStalkerIterator *iterator) {
 
 }
 
+void asan_arch_init(void) {
+
+  FATAL("ASAN mode not supported on this architecture");
+
+}
+
 #endif
 
diff --git a/frida_mode/src/asan/asan_x64.c b/frida_mode/src/asan/asan_x64.c
index bdf4ac30..a2eabe3c 100644
--- a/frida_mode/src/asan/asan_x64.c
+++ b/frida_mode/src/asan/asan_x64.c
@@ -7,23 +7,23 @@
 #include "ctx.h"
 #include "util.h"
 
+#if defined(__x86_64__)
+
 typedef void (*asan_loadN_t)(uint64_t address, uint8_t size);
 typedef void (*asan_storeN_t)(uint64_t address, uint8_t size);
 
 asan_loadN_t  asan_loadN = NULL;
 asan_storeN_t asan_storeN = NULL;
 
-#if defined(__x86_64__)
-
 static void asan_callout(GumCpuContext *ctx, gpointer user_data) {
 
   UNUSED_PARAMETER(user_data);
 
   cs_x86_op * operand = (cs_x86_op *)user_data;
   x86_op_mem *mem = &operand->mem;
-  uint64_t    base = 0;
-  uint64_t    index = 0;
-  uint64_t    address;
+  gsize       base = 0;
+  gsize       index = 0;
+  gsize       address;
   uint8_t     size;
 
   if (mem->base != X86_REG_INVALID) { base = ctx_read_reg(ctx, mem->base); }
diff --git a/frida_mode/src/asan/asan_x86.c b/frida_mode/src/asan/asan_x86.c
index b946b3bf..8490b490 100644
--- a/frida_mode/src/asan/asan_x86.c
+++ b/frida_mode/src/asan/asan_x86.c
@@ -1,18 +1,89 @@
+#include <dlfcn.h>
 #include "frida-gum.h"
 
 #include "debug.h"
 
 #include "asan.h"
+#include "ctx.h"
 #include "util.h"
 
 #if defined(__i386__)
+
+typedef void (*asan_loadN_t)(gsize address, uint8_t size);
+typedef void (*asan_storeN_t)(gsize address, uint8_t size);
+
+asan_loadN_t  asan_loadN = NULL;
+asan_storeN_t asan_storeN = NULL;
+
+static void asan_callout(GumCpuContext *ctx, gpointer user_data) {
+
+  UNUSED_PARAMETER(user_data);
+
+  cs_x86_op * operand = (cs_x86_op *)user_data;
+  x86_op_mem *mem = &operand->mem;
+  gsize       base = 0;
+  gsize       index = 0;
+  gsize       address;
+  uint8_t     size;
+
+  if (mem->base != X86_REG_INVALID) { base = ctx_read_reg(ctx, mem->base); }
+
+  if (mem->index != X86_REG_INVALID) { index = ctx_read_reg(ctx, mem->index); }
+
+  address = base + (mem->scale * index) + mem->disp;
+  size = operand->size;
+
+  if (operand->access == CS_AC_READ) {
+
+    asan_loadN(address, size);
+
+  } else if (operand->access == CS_AC_WRITE) {
+
+    asan_storeN(address, size);
+
+  }
+
+}
+
 void asan_instrument(const cs_insn *instr, GumStalkerIterator *iterator) {
 
-  UNUSED_PARAMETER(instr);
   UNUSED_PARAMETER(iterator);
-  if (asan_initialized) {
 
-    FATAL("ASAN mode not supported on this architecture");
+  cs_x86      x86 = instr->detail->x86;
+  cs_x86_op * operand;
+  x86_op_mem *mem;
+  cs_x86_op * ctx;
+
+  if (!asan_initialized) return;
+
+  if (instr->id == X86_INS_LEA) return;
+
+  if (instr->id == X86_INS_NOP) return;
+
+  for (uint8_t i = 0; i < x86.op_count; i++) {
+
+    operand = &x86.operands[i];
+
+    if (operand->type != X86_OP_MEM) { continue; }
+
+    mem = &operand->mem;
+    if (mem->segment != X86_REG_INVALID) { continue; }
+
+    ctx = g_malloc0(sizeof(cs_x86_op));
+    memcpy(ctx, operand, sizeof(cs_x86_op));
+    gum_stalker_iterator_put_callout(iterator, asan_callout, ctx, g_free);
+
+  }
+
+}
+
+void asan_arch_init(void) {
+
+  asan_loadN = (asan_loadN_t)dlsym(RTLD_DEFAULT, "__asan_loadN");
+  asan_storeN = (asan_loadN_t)dlsym(RTLD_DEFAULT, "__asan_storeN");
+  if (asan_loadN == NULL || asan_storeN == NULL) {
+
+    FATAL("Frida ASAN failed to find '__asan_loadN' or '__asan_storeN'");
 
   }
 
diff --git a/frida_mode/src/cmplog/cmplog_x64.c b/frida_mode/src/cmplog/cmplog_x64.c
index c3621a29..9f56c32a 100644
--- a/frida_mode/src/cmplog/cmplog_x64.c
+++ b/frida_mode/src/cmplog/cmplog_x64.c
@@ -31,12 +31,12 @@ typedef struct {
 
 } cmplog_pair_ctx_t;
 
-static gboolean cmplog_read_mem(GumX64CpuContext *ctx, uint8_t size,
-                                x86_op_mem *mem, guint64 *val) {
+static gboolean cmplog_read_mem(GumCpuContext *ctx, uint8_t size,
+                                x86_op_mem *mem, gsize *val) {
 
-  guint64 base = 0;
-  guint64 index = 0;
-  guint64 address;
+  gsize base = 0;
+  gsize index = 0;
+  gsize address;
 
   if (mem->base != X86_REG_INVALID) base = ctx_read_reg(ctx, mem->base);
 
@@ -49,16 +49,16 @@ static gboolean cmplog_read_mem(GumX64CpuContext *ctx, uint8_t size,
   switch (size) {
 
     case 1:
-      *val = *((guint8 *)address);
+      *val = *((guint8 *)GSIZE_TO_POINTER(address));
       return TRUE;
     case 2:
-      *val = *((guint16 *)address);
+      *val = *((guint16 *)GSIZE_TO_POINTER(address));
       return TRUE;
     case 4:
-      *val = *((guint32 *)address);
+      *val = *((guint32 *)GSIZE_TO_POINTER(address));
       return TRUE;
     case 8:
-      *val = *((guint64 *)address);
+      *val = *((guint64 *)GSIZE_TO_POINTER(address));
       return TRUE;
     default:
       FATAL("Invalid operand size: %d\n", size);
@@ -70,7 +70,7 @@ static gboolean cmplog_read_mem(GumX64CpuContext *ctx, uint8_t size,
 }
 
 static gboolean cmplog_get_operand_value(GumCpuContext *context,
-                                         cmplog_ctx_t *ctx, guint64 *val) {
+                                         cmplog_ctx_t *ctx, gsize *val) {
 
   switch (ctx->type) {
 
@@ -95,9 +95,9 @@ static void cmplog_call_callout(GumCpuContext *context, gpointer user_data) {
 
   UNUSED_PARAMETER(user_data);
 
-  guint64 address = ctx_read_reg(context, X86_REG_RIP);
-  guint64 rdi = ctx_read_reg(context, X86_REG_RDI);
-  guint64 rsi = ctx_read_reg(context, X86_REG_RSI);
+  gsize address = ctx_read_reg(context, X86_REG_RIP);
+  gsize rdi = ctx_read_reg(context, X86_REG_RDI);
+  gsize rsi = ctx_read_reg(context, X86_REG_RSI);
 
   if (((G_MAXULONG - rdi) < 32) || ((G_MAXULONG - rsi) < 32)) return;
 
@@ -169,10 +169,10 @@ static void cmplog_instrument_call(const cs_insn *     instr,
 
 }
 
-static void cmplog_handle_cmp_sub(GumCpuContext *context, guint64 operand1,
-                                  guint64 operand2, uint8_t size) {
+static void cmplog_handle_cmp_sub(GumCpuContext *context, gsize operand1,
+                                  gsize operand2, uint8_t size) {
 
-  guint64 address = ctx_read_reg(context, X86_REG_RIP);
+  gsize address = ctx_read_reg(context, X86_REG_RIP);
 
   register uintptr_t k = (uintptr_t)address;
 
@@ -195,8 +195,8 @@ static void cmplog_handle_cmp_sub(GumCpuContext *context, guint64 operand1,
 static void cmplog_cmp_sub_callout(GumCpuContext *context, gpointer user_data) {
 
   cmplog_pair_ctx_t *ctx = (cmplog_pair_ctx_t *)user_data;
-  guint64            operand1;
-  guint64            operand2;
+  gsize              operand1;
+  gsize              operand2;
 
   if (ctx->operand1.size != ctx->operand2.size) FATAL("Operand size mismatch");
 
diff --git a/frida_mode/src/cmplog/cmplog_x86.c b/frida_mode/src/cmplog/cmplog_x86.c
index 2401180c..a27df0af 100644
--- a/frida_mode/src/cmplog/cmplog_x86.c
+++ b/frida_mode/src/cmplog/cmplog_x86.c
@@ -1,17 +1,275 @@
 #include "frida-gum.h"
 
 #include "debug.h"
+#include "cmplog.h"
 
+#include "ctx.h"
 #include "frida_cmplog.h"
 #include "util.h"
 
 #if defined(__i386__)
+
+typedef struct {
+
+  x86_op_type type;
+  uint8_t     size;
+
+  union {
+
+    x86_op_mem mem;
+    x86_reg    reg;
+    int64_t    imm;
+
+  };
+
+} cmplog_ctx_t;
+
+typedef struct {
+
+  cmplog_ctx_t operand1;
+  cmplog_ctx_t operand2;
+
+} cmplog_pair_ctx_t;
+
+static gboolean cmplog_read_mem(GumCpuContext *ctx, uint8_t size,
+                                x86_op_mem *mem, gsize *val) {
+
+  gsize base = 0;
+  gsize index = 0;
+  gsize address;
+
+  if (mem->base != X86_REG_INVALID) base = ctx_read_reg(ctx, mem->base);
+
+  if (mem->index != X86_REG_INVALID) index = ctx_read_reg(ctx, mem->index);
+
+  address = base + (index * mem->scale) + mem->disp;
+
+  if (!cmplog_is_readable(address, size)) { return FALSE; }
+
+  switch (size) {
+
+    case 1:
+      *val = *((guint8 *)GSIZE_TO_POINTER(address));
+      return TRUE;
+    case 2:
+      *val = *((guint16 *)GSIZE_TO_POINTER(address));
+      return TRUE;
+    case 4:
+      *val = *((guint32 *)GSIZE_TO_POINTER(address));
+      return TRUE;
+    default:
+      FATAL("Invalid operand size: %d\n", size);
+
+  }
+
+  return FALSE;
+
+}
+
+static gboolean cmplog_get_operand_value(GumCpuContext *context,
+                                         cmplog_ctx_t *ctx, gsize *val) {
+
+  switch (ctx->type) {
+
+    case X86_OP_REG:
+      *val = ctx_read_reg(context, ctx->reg);
+      return TRUE;
+    case X86_OP_IMM:
+      *val = ctx->imm;
+      return TRUE;
+    case X86_OP_MEM:
+      return cmplog_read_mem(context, ctx->size, &ctx->mem, val);
+    default:
+      FATAL("Invalid operand type: %d\n", ctx->type);
+
+  }
+
+  return FALSE;
+
+}
+
+static void cmplog_call_callout(GumCpuContext *context, gpointer user_data) {
+
+  UNUSED_PARAMETER(user_data);
+
+  gsize  address = ctx_read_reg(context, X86_REG_EIP);
+  gsize *esp = (gsize *)ctx_read_reg(context, X86_REG_ESP);
+
+  if (!cmplog_is_readable(GPOINTER_TO_SIZE(esp), 12)) return;
+
+  /*
+   * This callout is place immediately before the call instruction, and hence
+   * the return address is not yet pushed on the top of the stack.
+   */
+  gsize arg1 = esp[0];
+  gsize arg2 = esp[1];
+
+  if (((G_MAXULONG - arg1) < 32) || ((G_MAXULONG - arg2) < 32)) return;
+
+  if (!cmplog_is_readable(arg1, 32) || !cmplog_is_readable(arg2, 32)) return;
+
+  void *ptr1 = GSIZE_TO_POINTER(arg1);
+  void *ptr2 = GSIZE_TO_POINTER(arg2);
+
+  uintptr_t k = address;
+
+  k = (k >> 4) ^ (k << 8);
+  k &= CMP_MAP_W - 1;
+
+  __afl_cmp_map->headers[k].type = CMP_TYPE_RTN;
+
+  u32 hits = __afl_cmp_map->headers[k].hits;
+  __afl_cmp_map->headers[k].hits = hits + 1;
+
+  __afl_cmp_map->headers[k].shape = 31;
+
+  hits &= CMP_MAP_RTN_H - 1;
+  gum_memcpy(((struct cmpfn_operands *)__afl_cmp_map->log[k])[hits].v0, ptr1,
+             32);
+  gum_memcpy(((struct cmpfn_operands *)__afl_cmp_map->log[k])[hits].v1, ptr2,
+             32);
+
+}
+
+static void cmplog_instrument_put_operand(cmplog_ctx_t *ctx,
+                                          cs_x86_op *   operand) {
+
+  ctx->type = operand->type;
+  ctx->size = operand->size;
+  switch (operand->type) {
+
+    case X86_OP_REG:
+      gum_memcpy(&ctx->reg, &operand->reg, sizeof(x86_reg));
+      break;
+    case X86_OP_IMM:
+      gum_memcpy(&ctx->imm, &operand->imm, sizeof(int64_t));
+      break;
+    case X86_OP_MEM:
+      gum_memcpy(&ctx->mem, &operand->mem, sizeof(x86_op_mem));
+      break;
+    default:
+      FATAL("Invalid operand type: %d\n", operand->type);
+
+  }
+
+}
+
+static void cmplog_instrument_call(const cs_insn *     instr,
+                                   GumStalkerIterator *iterator) {
+
+  cs_x86     x86 = instr->detail->x86;
+  cs_x86_op *operand;
+
+  if (instr->id != X86_INS_CALL) return;
+
+  if (x86.op_count != 1) return;
+
+  operand = &x86.operands[0];
+
+  if (operand->type == X86_OP_INVALID) return;
+  if (operand->type == X86_OP_MEM && operand->mem.segment != X86_REG_INVALID)
+    return;
+
+  gum_stalker_iterator_put_callout(iterator, cmplog_call_callout, NULL, NULL);
+
+}
+
+static void cmplog_handle_cmp_sub(GumCpuContext *context, gsize operand1,
+                                  gsize operand2, uint8_t size) {
+
+  gsize address = ctx_read_reg(context, X86_REG_EIP);
+
+  register uintptr_t k = (uintptr_t)address;
+
+  k = (k >> 4) ^ (k << 8);
+  k &= CMP_MAP_W - 1;
+
+  __afl_cmp_map->headers[k].type = CMP_TYPE_INS;
+
+  u32 hits = __afl_cmp_map->headers[k].hits;
+  __afl_cmp_map->headers[k].hits = hits + 1;
+
+  __afl_cmp_map->headers[k].shape = (size - 1);
+
+  hits &= CMP_MAP_H - 1;
+  __afl_cmp_map->log[k][hits].v0 = operand1;
+  __afl_cmp_map->log[k][hits].v1 = operand2;
+
+}
+
+static void cmplog_cmp_sub_callout(GumCpuContext *context, gpointer user_data) {
+
+  cmplog_pair_ctx_t *ctx = (cmplog_pair_ctx_t *)user_data;
+  gsize              operand1;
+  gsize              operand2;
+
+  if (ctx->operand1.size != ctx->operand2.size) FATAL("Operand size mismatch");
+
+  if (!cmplog_get_operand_value(context, &ctx->operand1, &operand1)) { return; }
+  if (!cmplog_get_operand_value(context, &ctx->operand2, &operand2)) { return; }
+
+  cmplog_handle_cmp_sub(context, operand1, operand2, ctx->operand1.size);
+
+}
+
+static void cmplog_instrument_cmp_sub_put_callout(GumStalkerIterator *iterator,
+                                                  cs_x86_op *         operand1,
+                                                  cs_x86_op *operand2) {
+
+  cmplog_pair_ctx_t *ctx = g_malloc(sizeof(cmplog_pair_ctx_t));
+  if (ctx == NULL) return;
+
+  cmplog_instrument_put_operand(&ctx->operand1, operand1);
+  cmplog_instrument_put_operand(&ctx->operand2, operand2);
+
+  gum_stalker_iterator_put_callout(iterator, cmplog_cmp_sub_callout, ctx,
+                                   g_free);
+
+}
+
+static void cmplog_instrument_cmp_sub(const cs_insn *     instr,
+                                      GumStalkerIterator *iterator) {
+
+  cs_x86     x86 = instr->detail->x86;
+  cs_x86_op *operand1;
+  cs_x86_op *operand2;
+
+  switch (instr->id) {
+
+    case X86_INS_CMP:
+    case X86_INS_SUB:
+      break;
+    default:
+      return;
+
+  }
+
+  if (x86.op_count != 2) return;
+
+  operand1 = &x86.operands[0];
+  operand2 = &x86.operands[1];
+
+  if (operand1->type == X86_OP_INVALID) return;
+  if (operand2->type == X86_OP_INVALID) return;
+
+  if ((operand1->type == X86_OP_MEM) &&
+      (operand1->mem.segment != X86_REG_INVALID))
+    return;
+
+  if ((operand2->type == X86_OP_MEM) &&
+      (operand2->mem.segment != X86_REG_INVALID))
+    return;
+
+  cmplog_instrument_cmp_sub_put_callout(iterator, operand1, operand2);
+
+}
+
 void cmplog_instrument(const cs_insn *instr, GumStalkerIterator *iterator) {
 
-  UNUSED_PARAMETER(instr);
-  UNUSED_PARAMETER(iterator);
-  if (__afl_cmp_map == NULL) { return; }
-  FATAL("CMPLOG mode not supported on this architecture");
+  if (__afl_cmp_map == NULL) return;
+
+  cmplog_instrument_call(instr, iterator);
+  cmplog_instrument_cmp_sub(instr, iterator);
 
 }
 
diff --git a/frida_mode/src/ctx/ctx_x64.c b/frida_mode/src/ctx/ctx_x64.c
index dec759f4..c5900533 100644
--- a/frida_mode/src/ctx/ctx_x64.c
+++ b/frida_mode/src/ctx/ctx_x64.c
@@ -41,7 +41,7 @@
                                  \
     }
 
-guint64 ctx_read_reg(GumX64CpuContext *ctx, x86_reg reg) {
+gsize ctx_read_reg(GumX64CpuContext *ctx, x86_reg reg) {
 
   switch (reg) {
 
diff --git a/frida_mode/src/ctx/ctx_x86.c b/frida_mode/src/ctx/ctx_x86.c
new file mode 100644
index 00000000..45308272
--- /dev/null
+++ b/frida_mode/src/ctx/ctx_x86.c
@@ -0,0 +1,81 @@
+#include "frida-gum.h"
+
+#include "debug.h"
+
+#include "ctx.h"
+
+#if defined(__i386__)
+
+  #define X86_REG_8L(LABEL, REG)  \
+    case LABEL: {                 \
+                                  \
+      return REG & GUM_INT8_MASK; \
+                                  \
+    }
+
+  #define X86_REG_8H(LABEL, REG)          \
+    case LABEL: {                         \
+                                          \
+      return (REG & GUM_INT16_MASK) >> 8; \
+                                          \
+    }
+
+  #define X86_REG_16(LABEL, REG)     \
+    case LABEL: {                    \
+                                     \
+      return (REG & GUM_INT16_MASK); \
+                                     \
+    }
+
+  #define X86_REG_32(LABEL, REG)     \
+    case LABEL: {                    \
+                                     \
+      return (REG & GUM_INT32_MASK); \
+                                     \
+    }
+
+gsize ctx_read_reg(GumIA32CpuContext *ctx, x86_reg reg) {
+
+  switch (reg) {
+
+    X86_REG_8L(X86_REG_AL, ctx->eax)
+    X86_REG_8L(X86_REG_BL, ctx->ebx)
+    X86_REG_8L(X86_REG_CL, ctx->ecx)
+    X86_REG_8L(X86_REG_DL, ctx->edx)
+    X86_REG_8L(X86_REG_BPL, ctx->ebp)
+    X86_REG_8L(X86_REG_SIL, ctx->esi)
+    X86_REG_8L(X86_REG_DIL, ctx->edi)
+
+    X86_REG_8H(X86_REG_AH, ctx->eax)
+    X86_REG_8H(X86_REG_BH, ctx->ebx)
+    X86_REG_8H(X86_REG_CH, ctx->ecx)
+    X86_REG_8H(X86_REG_DH, ctx->edx)
+
+    X86_REG_16(X86_REG_AX, ctx->eax)
+    X86_REG_16(X86_REG_BX, ctx->ebx)
+    X86_REG_16(X86_REG_CX, ctx->ecx)
+    X86_REG_16(X86_REG_DX, ctx->edx)
+    X86_REG_16(X86_REG_DI, ctx->edi)
+    X86_REG_16(X86_REG_SI, ctx->esi)
+    X86_REG_16(X86_REG_BP, ctx->ebp)
+
+    X86_REG_32(X86_REG_EAX, ctx->eax)
+    X86_REG_32(X86_REG_ECX, ctx->ecx)
+    X86_REG_32(X86_REG_EDX, ctx->edx)
+    X86_REG_32(X86_REG_EBX, ctx->ebx)
+    X86_REG_32(X86_REG_ESP, ctx->esp)
+    X86_REG_32(X86_REG_EBP, ctx->ebp)
+    X86_REG_32(X86_REG_ESI, ctx->esi)
+    X86_REG_32(X86_REG_EDI, ctx->edi)
+    X86_REG_32(X86_REG_EIP, ctx->eip)
+
+    default:
+      FATAL("Failed to read register: %d", reg);
+      return 0;
+
+  }
+
+}
+
+#endif
+
diff --git a/frida_mode/src/instrument/instrument.c b/frida_mode/src/instrument/instrument.c
index 5c77ade6..67eadc3f 100644
--- a/frida_mode/src/instrument/instrument.c
+++ b/frida_mode/src/instrument/instrument.c
@@ -34,7 +34,7 @@ __attribute__((hot)) static void on_basic_block(GumCpuContext *context,
    */
   static char buffer[200];
   int         len;
-  guint64     current_pc = (guint64)user_data;
+  GumAddress  current_pc = GUM_ADDRESS(user_data);
   uint8_t *   cursor;
   uint64_t    value;
   if (unlikely(tracing)) {
@@ -86,8 +86,8 @@ static void instr_basic_block(GumStalkerIterator *iterator,
 
     if (begin) {
 
-      prefetch_write((void *)instr->address);
-      if (!range_is_excluded((void *)instr->address)) {
+      prefetch_write(GSIZE_TO_POINTER(instr->address));
+      if (!range_is_excluded(GSIZE_TO_POINTER(instr->address))) {
 
         if (optimize) {
 
@@ -95,8 +95,8 @@ static void instr_basic_block(GumStalkerIterator *iterator,
 
         } else {
 
-          gum_stalker_iterator_put_callout(iterator, on_basic_block,
-                                           (gpointer)instr->address, NULL);
+          gum_stalker_iterator_put_callout(
+              iterator, on_basic_block, GSIZE_TO_POINTER(instr->address), NULL);
 
         }
 
@@ -106,7 +106,7 @@ static void instr_basic_block(GumStalkerIterator *iterator,
 
     }
 
-    if (!range_is_excluded((void *)instr->address)) {
+    if (!range_is_excluded(GSIZE_TO_POINTER(instr->address))) {
 
       asan_instrument(instr, iterator);
       cmplog_instrument(instr, iterator);
diff --git a/frida_mode/src/instrument/instrument_arm32.c b/frida_mode/src/instrument/instrument_arm32.c
index c2d720a7..1a3c40bb 100644
--- a/frida_mode/src/instrument/instrument_arm32.c
+++ b/frida_mode/src/instrument/instrument_arm32.c
@@ -3,6 +3,7 @@
 #include "debug.h"
 
 #include "instrument.h"
+#include "util.h"
 
 #if defined(__arm__)
 
@@ -15,6 +16,8 @@ gboolean instrument_is_coverage_optimize_supported(void) {
 void instrument_coverage_optimize(const cs_insn *   instr,
                                   GumStalkerOutput *output) {
 
+  UNUSED_PARAMETER(instr);
+  UNUSED_PARAMETER(output);
   FATAL("Optimized coverage not supported on this architecture");
 
 }
diff --git a/frida_mode/src/instrument/instrument_x86.c b/frida_mode/src/instrument/instrument_x86.c
index 5b8cbbba..585bb5b8 100644
--- a/frida_mode/src/instrument/instrument_x86.c
+++ b/frida_mode/src/instrument/instrument_x86.c
@@ -3,19 +3,81 @@
 #include "debug.h"
 
 #include "instrument.h"
+#include "util.h"
 
 #if defined(__i386__)
 
+static GumAddress current_log_impl = GUM_ADDRESS(0);
+
+static void instrument_coverage_function(GumX86Writer *cw) {
+
+  gum_x86_writer_put_pushfx(cw);
+  gum_x86_writer_put_push_reg(cw, GUM_REG_ECX);
+  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_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);
+
+  gum_x86_writer_put_add_reg_imm(cw, GUM_REG_EDX, GUM_ADDRESS(__afl_area_ptr));
+
+  /* add byte ptr [edx], 1 */
+  uint8_t add_byte_ptr_edx_1[] = {0x80, 0x02, 0x01};
+  gum_x86_writer_put_bytes(cw, add_byte_ptr_edx_1, sizeof(add_byte_ptr_edx_1));
+
+  /* adc byte ptr [edx], 0 */
+  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);
+  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);
+  gum_x86_writer_put_pop_reg(cw, GUM_REG_ECX);
+  gum_x86_writer_put_popfx(cw);
+  gum_x86_writer_put_ret(cw);
+
+}
+
 gboolean instrument_is_coverage_optimize_supported(void) {
 
-  return false;
+  return true;
 
 }
 
 void instrument_coverage_optimize(const cs_insn *   instr,
                                   GumStalkerOutput *output) {
 
-  FATAL("Optimized coverage not supported on this architecture");
+  UNUSED_PARAMETER(instr);
+  UNUSED_PARAMETER(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 ||
+      !gum_x86_writer_can_branch_directly_between(cw->pc, current_log_impl) ||
+      !gum_x86_writer_can_branch_directly_between(cw->pc + 128,
+                                                  current_log_impl)) {
+
+    gconstpointer after_log_impl = cw->code + 1;
+
+    gum_x86_writer_put_jmp_near_label(cw, after_log_impl);
+
+    current_log_impl = cw->pc;
+    instrument_coverage_function(cw);
+
+    gum_x86_writer_put_label(cw, after_log_impl);
+
+  }
+
+  // gum_x86_writer_put_breakpoint(cw);
+  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);
+  gum_x86_writer_put_pop_reg(cw, GUM_REG_EDI);
 
 }
 
diff --git a/frida_mode/src/lib/lib.c b/frida_mode/src/lib/lib.c
index c5045533..13a7d1e7 100644
--- a/frida_mode/src/lib/lib.c
+++ b/frida_mode/src/lib/lib.c
@@ -90,7 +90,8 @@ static void lib_read_text_section(lib_details_t *lib_details, Elf_Ehdr *hdr) {
 
   if (!found_preferred_base) { FATAL("Failed to find preferred load address"); }
 
-  OKF("Image preferred load address 0x%016lx", preferred_base);
+  OKF("Image preferred load address 0x%016" G_GSIZE_MODIFIER "x",
+      preferred_base);
 
   shdr = (Elf_Shdr *)((char *)hdr + hdr->e_shoff);
   shstrtab = &shdr[hdr->e_shstrndx];
@@ -107,15 +108,16 @@ static void lib_read_text_section(lib_details_t *lib_details, Elf_Ehdr *hdr) {
     if (curr->sh_name == 0) continue;
 
     section_name = &shstr[curr->sh_name];
-    OKF("Section: %2lu - base: 0x%016lX size: 0x%016lX %s", i, curr->sh_addr,
-        curr->sh_size, section_name);
+    OKF("Section: %2" G_GSIZE_MODIFIER "u - base: 0x%016" G_GSIZE_MODIFIER
+        "X size: 0x%016" G_GSIZE_MODIFIER "X %s",
+        i, curr->sh_addr, curr->sh_size, section_name);
     if (memcmp(section_name, text_name, sizeof(text_name)) == 0 &&
         text_base == 0) {
 
       text_base = lib_details->base_address + curr->sh_addr - preferred_base;
       text_limit = text_base + curr->sh_size;
-      OKF("> text_addr: 0x%016lX", text_base);
-      OKF("> text_limit: 0x%016lX", text_limit);
+      OKF("> text_addr: 0x%016" G_GINT64_MODIFIER "X", text_base);
+      OKF("> text_limit: 0x%016" G_GINT64_MODIFIER "X", text_limit);
 
     }
 
@@ -153,7 +155,8 @@ void lib_init(void) {
 
   lib_details_t lib_details;
   gum_process_enumerate_modules(lib_find_exe, &lib_details);
-  OKF("Executable: 0x%016lx - %s", lib_details.base_address, lib_details.path);
+  OKF("Executable: 0x%016" G_GINT64_MODIFIER "x - %s", lib_details.base_address,
+      lib_details.path);
   lib_get_text_section(&lib_details);
 
 }
diff --git a/frida_mode/src/persistent/persistent_x86.c b/frida_mode/src/persistent/persistent_x86.c
index 9d39c4e9..bd7171b9 100644
--- a/frida_mode/src/persistent/persistent_x86.c
+++ b/frida_mode/src/persistent/persistent_x86.c
@@ -1,9 +1,9 @@
 #include "frida-gum.h"
 
-#include "debug.h"
+#include "config.h"
 
+#include "instrument.h"
 #include "persistent.h"
-#include "util.h"
 
 #if defined(__i386__)
 
@@ -38,16 +38,239 @@ struct x86_regs {
 
 typedef struct x86_regs arch_api_regs;
 
+static arch_api_regs saved_regs = {0};
+static void *        saved_return = NULL;
+
 gboolean persistent_is_supported(void) {
 
-  return false;
+  return true;
+
+}
+
+static void instrument_persitent_save_regs(GumX86Writer *   cw,
+                                           struct x86_regs *regs) {
+
+  GumAddress regs_address = GUM_ADDRESS(regs);
+
+  /* Should be pushing FPU here, but meh */
+  gum_x86_writer_put_pushfx(cw);
+  gum_x86_writer_put_push_reg(cw, GUM_REG_EAX);
+
+  gum_x86_writer_put_mov_reg_address(cw, GUM_REG_EAX, regs_address);
+
+  gum_x86_writer_put_mov_reg_offset_ptr_reg(cw, GUM_REG_EAX, (0x4 * 1),
+                                            GUM_REG_EBX);
+  gum_x86_writer_put_mov_reg_offset_ptr_reg(cw, GUM_REG_EAX, (0x4 * 2),
+                                            GUM_REG_ECX);
+  gum_x86_writer_put_mov_reg_offset_ptr_reg(cw, GUM_REG_EAX, (0x4 * 3),
+                                            GUM_REG_EDX);
+  gum_x86_writer_put_mov_reg_offset_ptr_reg(cw, GUM_REG_EAX, (0x4 * 4),
+                                            GUM_REG_EDI);
+  gum_x86_writer_put_mov_reg_offset_ptr_reg(cw, GUM_REG_EAX, (0x4 * 5),
+                                            GUM_REG_ESI);
+  gum_x86_writer_put_mov_reg_offset_ptr_reg(cw, GUM_REG_EAX, (0x4 * 6),
+                                            GUM_REG_EBP);
+
+  /* Store RIP */
+  gum_x86_writer_put_mov_reg_address(cw, GUM_REG_EBX,
+                                     GUM_ADDRESS(persistent_start));
+
+  gum_x86_writer_put_mov_reg_offset_ptr_reg(cw, GUM_REG_EAX, (0x4 * 7),
+                                            GUM_REG_EBX);
+
+  /* Store adjusted RSP */
+  gum_x86_writer_put_mov_reg_reg(cw, GUM_REG_EBX, GUM_REG_ESP);
+
+  /* RED_ZONE + Saved flags, RAX */
+  gum_x86_writer_put_add_reg_imm(cw, GUM_REG_EBX, (0x4 * 2));
+  gum_x86_writer_put_mov_reg_offset_ptr_reg(cw, GUM_REG_EAX, (0x4 * 8),
+                                            GUM_REG_EBX);
+
+  /* Save the flags */
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_EBX, GUM_REG_ESP, 0x4);
+  gum_x86_writer_put_mov_reg_offset_ptr_reg(cw, GUM_REG_EAX, (0x4 * 9),
+                                            GUM_REG_EBX);
+
+  /* Save the RAX */
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_EBX, GUM_REG_ESP, 0x0);
+  gum_x86_writer_put_mov_reg_offset_ptr_reg(cw, GUM_REG_EAX, (0x4 * 0),
+                                            GUM_REG_EBX);
+
+  /* Pop the saved values */
+  gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_REG_ESP, GUM_REG_ESP, 0x8);
+
+}
+
+static void instrument_persitent_restore_regs(GumX86Writer *   cw,
+                                              struct x86_regs *regs) {
+
+  GumAddress regs_address = GUM_ADDRESS(regs);
+  gum_x86_writer_put_mov_reg_address(cw, GUM_REG_EAX, regs_address);
+
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_ECX, GUM_REG_EAX,
+                                            (0x4 * 2));
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_EDX, GUM_REG_EAX,
+                                            (0x4 * 3));
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_EDI, GUM_REG_EAX,
+                                            (0x4 * 4));
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_ESI, GUM_REG_EAX,
+                                            (0x4 * 5));
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_EBP, GUM_REG_EAX,
+                                            (0x4 * 6));
+
+  /* Don't restore RIP or RSP */
+
+  /* Restore RBX, RAX & Flags */
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_EBX, GUM_REG_EAX,
+                                            (0x4 * 1));
+  gum_x86_writer_put_push_reg(cw, GUM_REG_EBX);
+
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_EBX, GUM_REG_EAX,
+                                            (0x4 * 0));
+  gum_x86_writer_put_push_reg(cw, GUM_REG_EBX);
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_EBX, GUM_REG_EAX,
+                                            (0x4 * 9));
+  gum_x86_writer_put_push_reg(cw, GUM_REG_EBX);
+
+  gum_x86_writer_put_popfx(cw);
+  gum_x86_writer_put_pop_reg(cw, GUM_REG_EAX);
+  gum_x86_writer_put_pop_reg(cw, GUM_REG_EBX);
+
+}
+
+static void instrument_save_ret(GumX86Writer *cw, void **saved_return_ptr) {
+
+  GumAddress saved_return_address = GUM_ADDRESS(saved_return_ptr);
+
+  gum_x86_writer_put_push_reg(cw, GUM_REG_EAX);
+  gum_x86_writer_put_push_reg(cw, GUM_REG_EBX);
+
+  gum_x86_writer_put_mov_reg_address(cw, GUM_REG_EAX, saved_return_address);
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_EBX, GUM_REG_ESP, 0x8);
+  gum_x86_writer_put_mov_reg_offset_ptr_reg(cw, GUM_REG_EAX, 0, GUM_REG_EBX);
+
+  gum_x86_writer_put_pop_reg(cw, GUM_REG_EBX);
+  gum_x86_writer_put_pop_reg(cw, GUM_REG_EAX);
+
+}
+
+static void instrument_jump_ret(GumX86Writer *cw, void **saved_return_ptr) {
+
+  GumAddress saved_return_address = GUM_ADDRESS(saved_return_ptr);
+
+  /* Place holder for ret */
+  gum_x86_writer_put_push_reg(cw, GUM_REG_EAX);
+  gum_x86_writer_put_push_reg(cw, GUM_REG_EAX);
+
+  gum_x86_writer_put_mov_reg_address(cw, GUM_REG_EAX, saved_return_address);
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_EAX, GUM_REG_EAX, 0);
+
+  gum_x86_writer_put_mov_reg_offset_ptr_reg(cw, GUM_REG_ESP, 0x4, GUM_REG_EAX);
+  gum_x86_writer_put_pop_reg(cw, GUM_REG_EAX);
+  gum_x86_writer_put_ret(cw);
+
+}
+
+static int instrument_afl_persistent_loop_func(void) {
+
+  int ret = __afl_persistent_loop(persistent_count);
+  previous_pc = 0;
+  return ret;
+
+}
+
+static void instrument_afl_persistent_loop(GumX86Writer *cw) {
+
+  gum_x86_writer_put_call_address_with_arguments(
+      cw, GUM_CALL_CAPI, GUM_ADDRESS(instrument_afl_persistent_loop_func), 0);
+  gum_x86_writer_put_test_reg_reg(cw, GUM_REG_EAX, GUM_REG_EAX);
+
+}
+
+static void persistent_prologue_hook(GumX86Writer *cw, struct x86_regs *regs) {
+
+  if (hook == NULL) return;
+
+  gum_x86_writer_put_mov_reg_address(cw, GUM_REG_ECX,
+                                     GUM_ADDRESS(&__afl_fuzz_len));
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_ECX, GUM_REG_ECX, 0);
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_ECX, GUM_REG_ECX, 0);
+
+  gum_x86_writer_put_mov_reg_address(cw, GUM_REG_EDX,
+                                     GUM_ADDRESS(&__afl_fuzz_ptr));
+  gum_x86_writer_put_mov_reg_reg_offset_ptr(cw, GUM_REG_EDX, GUM_REG_EDX, 0);
+
+  /* Base address is 64-bits (hence two zero arguments) */
+  gum_x86_writer_put_call_address_with_arguments(
+      cw, GUM_CALL_CAPI, GUM_ADDRESS(hook), 5, GUM_ARG_ADDRESS,
+      GUM_ADDRESS(regs), GUM_ARG_ADDRESS, GUM_ADDRESS(0), GUM_ARG_ADDRESS,
+      GUM_ADDRESS(0), GUM_ARG_REGISTER, GUM_REG_EDX, GUM_ARG_REGISTER,
+      GUM_REG_ECX);
 
 }
 
 void persistent_prologue(GumStalkerOutput *output) {
 
-  UNUSED_PARAMETER(output);
-  FATAL("Persistent mode not supported on this architecture");
+  /*
+   *  SAVE REGS
+   *  SAVE RET
+   *  POP RET
+   * loop:
+   *  CALL instrument_afl_persistent_loop
+   *  TEST EAX, EAX
+   *  JZ end:
+   *  call hook (optionally)
+   *  RESTORE REGS
+   *  call original
+   *  jmp loop:
+   *
+   * end:
+   *  JMP SAVED RET
+   *
+   * original:
+   *  INSTRUMENTED PERSISTENT FUNC
+   */
+
+  GumX86Writer *cw = output->writer.x86;
+
+  gconstpointer loop = cw->code + 1;
+
+  /* Stack must be 16-byte aligned per ABI */
+  instrument_persitent_save_regs(cw, &saved_regs);
+
+  /* Stash and pop the return value */
+  instrument_save_ret(cw, &saved_return);
+  gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_REG_ESP, GUM_REG_ESP, (4));
+
+  /* loop: */
+  gum_x86_writer_put_label(cw, loop);
+
+  /* call instrument_prologue_func */
+  instrument_afl_persistent_loop(cw);
+
+  /* jz done */
+  gconstpointer done = cw->code + 1;
+  gum_x86_writer_put_jcc_near_label(cw, X86_INS_JE, done, GUM_UNLIKELY);
+
+  /* Optionally call the persistent hook */
+  persistent_prologue_hook(cw, &saved_regs);
+
+  instrument_persitent_restore_regs(cw, &saved_regs);
+  gconstpointer original = cw->code + 1;
+  /* call original */
+  gum_x86_writer_put_call_near_label(cw, original);
+  /* jmp loop */
+  gum_x86_writer_put_jmp_near_label(cw, loop);
+
+  /* done: */
+  gum_x86_writer_put_label(cw, done);
+
+  instrument_jump_ret(cw, &saved_return);
+
+  /* original: */
+  gum_x86_writer_put_label(cw, original);
+
+  gum_x86_writer_flush(cw);
 
 }