1 files changed, 86 insertions, 1 deletions

diff --git a/utils/afl_frida/afl-frida.c b/utils/afl_frida/afl-frida.c
index 85cf2e9d..711d8f33 100644
--- a/utils/afl_frida/afl-frida.c
+++ b/utils/afl_frida/afl-frida.c
@@ -94,6 +94,8 @@ typedef struct {
 
   GumAddress base_address;
   guint64    code_start, code_end;
+  GumAddress current_log_impl;
+  uint64_t   afl_prev_loc;
 
 } range_t;
 
@@ -109,12 +111,58 @@ inline static void afl_maybe_log(guint64 current_pc) {
 
 }
 
+#if GUM_NATIVE_CPU == GUM_CPU_AMD64
+
+static const guint8 afl_maybe_log_code[] = {
+
+    0x9c,  // pushfq
+    0x50,  // push rax
+    0x51,  // push rcx
+    0x52,  // push rdx
+    0x56,  // push rsi
+
+    0x89, 0xf8,                                // mov eax, edi
+    0xc1, 0xe0, 0x08,                          // shl eax, 8
+    0xc1, 0xef, 0x04,                          // shr edi, 4
+    0x31, 0xc7,                                // xor edi, eax
+    0x0f, 0xb7, 0xc7,                          // movzx eax, di
+    0x48, 0x8d, 0x0d, 0x30, 0x00, 0x00, 0x00,  // lea rcx, sym._afl_area_ptr_ptr
+    0x48, 0x8b, 0x09,                          // mov rcx, qword [rcx]
+    0x48, 0x8b, 0x09,                          // mov rcx, qword [rcx]
+    0x48, 0x8d, 0x15, 0x1b, 0x00, 0x00, 0x00,  // lea rdx, sym._afl_prev_loc_ptr
+    0x48, 0x8b, 0x32,                          // mov rsi, qword [rdx]
+    0x48, 0x8b, 0x36,                          // mov rsi, qword [rsi]
+    0x48, 0x31, 0xc6,                          // xor rsi, rax
+    0xfe, 0x04, 0x31,                          // inc byte [rcx + rsi]
+
+    0x48, 0xd1, 0xe8,  // shr rax, 1
+    0x48, 0x8b, 0x0a,  // mov rcx, qword [rdx]
+    0x48, 0x89, 0x01,  // mov qword [rcx], rax
+
+    0x5e,  // pop rsi
+    0x5a,  // pop rdx
+    0x59,  // pop rcx
+    0x58,  // pop rax
+    0x9d,  // popfq
+
+    0xc3,  // ret
+           // Read-only data goes here:
+           // uint64_t* afl_prev_loc_ptr
+           // uint8_t** afl_area_ptr_ptr
+           // unsigned int afl_instr_rms
+
+};
+
+#else
+
 static void on_basic_block(GumCpuContext *context, gpointer user_data) {
 
   afl_maybe_log((guint64)user_data);
 
 }
 
+#endif
+
 void instr_basic_block(GumStalkerIterator *iterator, GumStalkerOutput *output,
                        gpointer user_data) {
 
@@ -129,8 +177,45 @@ void instr_basic_block(GumStalkerIterator *iterator, GumStalkerOutput *output,
       if (instr->address >= range->code_start &&
           instr->address <= range->code_end) {
 
+#if GUM_NATIVE_CPU == GUM_CPU_AMD64
+        GumX86Writer *cw = output->writer.x86;
+        if (range->current_log_impl == 0 ||
+            !gum_x86_writer_can_branch_directly_between(
+                cw->pc, range->current_log_impl) ||
+            !gum_x86_writer_can_branch_directly_between(
+                cw->pc + 128, range->current_log_impl)) {
+
+          gconstpointer after_log_impl = cw->code + 1;
+
+          gum_x86_writer_put_jmp_near_label(cw, after_log_impl);
+
+          range->current_log_impl = cw->pc;
+          gum_x86_writer_put_bytes(cw, afl_maybe_log_code,
+                                   sizeof(afl_maybe_log_code));
+
+          uint64_t *afl_prev_loc_ptr = &range->afl_prev_loc;
+          uint8_t **afl_area_ptr_ptr = &__afl_area_ptr;
+          gum_x86_writer_put_bytes(cw, (const guint8 *)&afl_prev_loc_ptr,
+                                   sizeof(afl_prev_loc_ptr));
+          gum_x86_writer_put_bytes(cw, (const guint8 *)&afl_area_ptr_ptr,
+                                   sizeof(afl_area_ptr_ptr));
+          gum_x86_writer_put_label(cw, after_log_impl);
+
+        }
+
+        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);
+        gum_x86_writer_put_mov_reg_address(cw, GUM_REG_RDI,
+                                           GUM_ADDRESS(instr->address));
+        gum_x86_writer_put_call_address(cw, range->current_log_impl);
+        gum_x86_writer_put_pop_reg(cw, GUM_REG_RDI);
+        gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_REG_RSP, GUM_REG_RSP,
+                                              GUM_RED_ZONE_SIZE);
+#else
         gum_stalker_iterator_put_callout(iterator, on_basic_block,
                                          (gpointer)instr->address, NULL);
+#endif
         begin = FALSE;
 
       }
@@ -228,7 +313,7 @@ int main(int argc, char **argv) {
 
     guint64 code_start = code_range.base_address;
     guint64 code_end = code_range.base_address + code_range.size;
-    range_t instr_range = {0, code_start, code_end};
+    range_t instr_range = {0, code_start, code_end, 0, 0};
 
     printf("Frida instrumentation: base=0x%lx instrumenting=0x%lx-%lx\n",
            base_address, code_start, code_end);