Minimize inline assembly

6 files changed, 203 insertions, 163 deletions

diff --git a/frida_mode/GNUmakefile b/frida_mode/GNUmakefile
index 3e35e2f6..ab6efecf 100644
--- a/frida_mode/GNUmakefile
+++ b/frida_mode/GNUmakefile
@@ -31,7 +31,8 @@ AFL_CFLAGS:=-Wno-unused-parameter \
 LDFLAGS+=-shared \
 		 -lpthread \
 		 -lresolv \
-		 -ldl
+		 -ldl \
+		 -lrt
 
 ifdef DEBUG
 CFLAGS+=-Werror \
diff --git a/frida_mode/include/instrument.h b/frida_mode/include/instrument.h
index 909b2a2c..cac5ee93 100644
--- a/frida_mode/include/instrument.h
+++ b/frida_mode/include/instrument.h
@@ -29,6 +29,7 @@ GumStalkerTransformer *instrument_get_transformer(void);
 /* Functions to be implemented by the different architectures */
 gboolean instrument_is_coverage_optimize_supported(void);
 
+void instrument_coverage_optimize_init(void);
 void instrument_coverage_optimize(const cs_insn *   instr,
                                   GumStalkerOutput *output);
 
diff --git a/frida_mode/include/ranges.h b/frida_mode/include/ranges.h
index 0220a59d..3bd9eaa6 100644
--- a/frida_mode/include/ranges.h
+++ b/frida_mode/include/ranges.h
@@ -10,6 +10,8 @@ extern gboolean ranges_inst_jit;
 void ranges_config(void);
 void ranges_init(void);
 
+void ranges_print_debug_maps(void);
+
 gboolean range_is_excluded(GumAddress address);
 
 void ranges_exclude();
diff --git a/frida_mode/src/instrument/instrument.c b/frida_mode/src/instrument/instrument.c
index 71d9bdf6..81d85aa1 100644
--- a/frida_mode/src/instrument/instrument.c
+++ b/frida_mode/src/instrument/instrument.c
@@ -356,6 +356,7 @@ void instrument_init(void) {
       instrument_hash_seed);
   instrument_hash_zero = instrument_get_offset_hash(0);
 
+  instrument_coverage_optimize_init();
   instrument_debug_init();
   instrument_coverage_init();
   asan_init();
diff --git a/frida_mode/src/instrument/instrument_x64.c b/frida_mode/src/instrument/instrument_x64.c
index 1c2cf113..dc040a20 100644
--- a/frida_mode/src/instrument/instrument_x64.c
+++ b/frida_mode/src/instrument/instrument_x64.c
@@ -1,4 +1,9 @@
+#include <fcntl.h>
 #include <stddef.h>
+#include <asm/prctl.h>
+#include <sys/mman.h>
+#include <sys/syscall.h>
+#include <sys/shm.h>
 
 #include "frida-gumjs.h"
 
@@ -6,245 +11,274 @@
 #include "debug.h"
 
 #include "instrument.h"
+#include "ranges.h"
 
 #if defined(__x86_64__)
 
-static GumAddress current_log_impl = GUM_ADDRESS(0);
+gboolean instrument_is_coverage_optimize_supported(void) {
 
-  #pragma pack(push, 1)
+  return true;
+
+}
+
+static gboolean instrument_coverage_in_range(gssize offset) {
+
+  return (offset >= G_MININT32 && offset <= G_MAXINT32);
 
+}
+
+  #pragma pack(push, 1)
 typedef struct {
 
-  /*
-   * pushfq
-   * push rdx
-   * mov rdx, [&previouspc] (rip relative addr)
-   * xor rdx, rdi (current_pc)
-   * shr rdi. 1
-   * mov [&previouspc], rdi
-   * lea rsi, [&_afl_area_ptr] (rip relative)
-   * add rdx, rsi
-   * add byte ptr [rdx], 1
-   * adc byte ptr [rdx], 0
-
-   * pop rdx
-   * popfq
-   */
+  // cur_location = (block_address >> 4) ^ (block_address << 8);
+  // shared_mem[cur_location ^ prev_location]++;
+  // prev_location = cur_location >> 1;
+
+  // => 0x7ffff6cfb086:      lea    rsp,[rsp-0x80]
+  //    0x7ffff6cfb08b:      pushf
+  //    0x7ffff6cfb08c:      push   rsi
+  //    0x7ffff6cfb08d:      mov    rsi,0x228
+  //    0x7ffff6cfb094:      xchg   QWORD PTR [rip+0x3136a5],rsi        # 0x7ffff700e740
+  //    0x7ffff6cfb09b:      xor    rsi,0x451
+  //    0x7ffff6cfb0a2:      add    BYTE PTR [rsi+0x10000],0x1
+  //    0x7ffff6cfb0a9:      adc    BYTE PTR [rsi+0x10000],0x0
+  //    0x7ffff6cfb0b0:      pop    rsi
+  //    0x7ffff6cfb0b1:      popf
+  //    0x7ffff6cfb0b2:      lea    rsp,[rsp+0x80]
+
+
+  uint8_t lea_rsp_rsp_sub_rz[5];
   uint8_t push_fq;
-  uint8_t push_rdx;
-  uint8_t mov_rdx_rip_off[7];
-  uint8_t xor_rdx_rdi[3];
-  uint8_t shr_rdi[3];
-  uint8_t mov_rip_off_rdi[7];
-
-  uint8_t lea_rdi_rip_off[7];
-  uint8_t add_rdx_rdi[3];
-  uint8_t add_byte_ptr_rdx[3];
-  uint8_t adc_byte_ptr_rdx[3];
-
-  uint8_t pop_rdx;
+  uint8_t push_rsi;
+
+  uint8_t mov_rsi_curr_loc_shr_1[7];
+  uint8_t xchg_rsi_prev_loc_curr_loc[7];
+  uint8_t xor_rsi_curr_loc[7];
+
+  uint8_t add_rsi_1[7];
+  uint8_t adc_rsi_0[7];
+
+  uint8_t pop_rsi;
   uint8_t pop_fq;
-  uint8_t ret;
+  uint8_t lsa_rsp_rsp_add_rz[8];
 
 } afl_log_code_asm_t;
 
   #pragma pack(pop)
 
-  #pragma pack(push, 8)
-typedef struct {
+typedef union {
 
-  afl_log_code_asm_t assembly;
-  uint64_t           current_pc;
+  afl_log_code_asm_t code;
+  uint8_t            bytes[0];
 
-} afl_log_code_t;
+} afl_log_code;
 
-  #pragma pack(pop)
+static const afl_log_code_asm_t template =
+    {
 
-typedef union {
+        .lea_rsp_rsp_sub_rz = {0x48, 0x8D, 0x64, 0x24, 0x80},
+        .push_fq = 0x9c,
+        .push_rsi = 0x56,
 
-  afl_log_code_t data;
-  uint8_t        bytes[0];
+        .mov_rsi_curr_loc_shr_1 = {0x48, 0xC7, 0xC6},
+        .xchg_rsi_prev_loc_curr_loc = {0x48, 0x87, 0x35},
+        .xor_rsi_curr_loc = {0x48, 0x81, 0xF6},
 
-} afl_log_code;
+        .add_rsi_1 = {0x80, 0x86, 0x00, 0x00, 0x00, 0x00, 0x01},
+        .adc_rsi_0 = {0x80, 0x96, 0x00, 0x00, 0x00, 0x00, 0x00},
 
-static const afl_log_code_asm_t template = {
+        .pop_rsi = 0x5E,
+        .pop_fq = 0x9D,
+        .lsa_rsp_rsp_add_rz = {0x48, 0x8D, 0xA4, 0x24, 0x80, 0x00, 0x00, 0x00},
 
-    .push_fq = 0x9c,
-    .push_rdx = 0x52,
-    .mov_rdx_rip_off =
-        {
+}
 
-            0x48, 0x8b, 0x15,
-            /* TBC */
+;
 
-        },
+static gboolean instrument_coverage_find_low(const GumRangeDetails *details,
+                                             gpointer               user_data) {
 
-    .xor_rdx_rdi =
-        {
+  static GumAddress last_limit = (64ULL << 10);
+  gpointer *        address = (gpointer *)user_data;
 
-            0x48,
-            0x31,
-            0xfa,
+  if ((details->range->base_address - last_limit) > __afl_map_size) {
 
-        },
+    *address = GSIZE_TO_POINTER(last_limit);
+    return FALSE;
 
-    .shr_rdi = {0x48, 0xd1, 0xef},
-    .mov_rip_off_rdi = {0x48, 0x89, 0x3d},
+  }
 
-    .lea_rdi_rip_off =
-        {
+  if (details->range->base_address > ((2ULL << 20) - __afl_map_size)) {
 
-            0x48,
-            0x8d,
-            0x3d,
+    return FALSE;
 
-        },
+  }
 
-    .add_rdx_rdi = {0x48, 0x01, 0xfA},
+  last_limit = details->range->base_address + details->range->size;
+  return TRUE;
 
-    .add_byte_ptr_rdx =
-        {
+}
 
-            0x80,
-            0x02,
-            0x01,
+static void instrument_coverage_optimize_map_mmap_anon(gpointer address) {
 
-        },
+  __afl_area_ptr =
+      mmap(address, __afl_map_size, PROT_READ | PROT_WRITE,
+           MAP_FIXED_NOREPLACE | MAP_SHARED | MAP_ANONYMOUS, -1, 0);
+  if (__afl_area_ptr != address) {
 
-    .adc_byte_ptr_rdx =
-        {
+    FATAL("Failed to map mmap __afl_area_ptr: %d", errno);
 
-            0x80,
-            0x12,
-            0x00,
+  }
 
-        },
+}
 
-    .pop_rdx = 0x5a,
-    .pop_fq = 0x9d,
-    .ret = 0xc3};
+static void instrument_coverage_optimize_map_mmap(char *   shm_file_path,
+                                                  gpointer address) {
 
-static guint8 align_pad[] = {0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90};
+  int shm_fd = -1;
 
-gboolean instrument_is_coverage_optimize_supported(void) {
+  if (munmap(__afl_area_ptr, __afl_map_size) != 0) {
 
-  return true;
+    FATAL("Failed to unmap previous __afl_area_ptr");
 
-}
+  }
 
-static gboolean instrument_coverage_in_range(gssize offset) {
+  __afl_area_ptr = NULL;
 
-  return (offset >= G_MININT32 && offset <= G_MAXINT32);
+  shm_fd = shm_open(shm_file_path, O_RDWR, DEFAULT_PERMISSION);
+  if (shm_fd == -1) { FATAL("shm_open() failed\n"); }
+
+  __afl_area_ptr = mmap(address, __afl_map_size, PROT_READ | PROT_WRITE,
+                        MAP_FIXED_NOREPLACE | MAP_SHARED, shm_fd, 0);
+  if (__afl_area_ptr != address) {
+
+    FATAL("Failed to map mmap __afl_area_ptr: %d", errno);
+
+  }
+
+  if (close(shm_fd) != 0) { FATAL("Failed to close shm_fd"); }
 
 }
 
-static void instrument_coverate_write_function(GumStalkerOutput *output) {
+static void instrument_coverage_optimize_map_shm(guint64  shm_env_val,
+                                                 gpointer address) {
 
-  guint64       misalign = 0;
-  GumX86Writer *cw = output->writer.x86;
-  GumAddress    code_addr = 0;
-  afl_log_code  code = {0};
-  /*guint64       instrument_hash_zero = 0;*/
+  if (shmdt(__afl_area_ptr) != 0) {
+
+    FATAL("Failed to detach previous __afl_area_ptr");
 
-  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;
+  __afl_area_ptr = shmat(shm_env_val, address, 0);
+  if (__afl_area_ptr != address) {
 
-    gum_x86_writer_put_jmp_near_label(cw, after_log_impl);
+    FATAL("Failed to map shm __afl_area_ptr: %d", errno);
 
-    misalign = (cw->pc & 0x7);
-    if (misalign != 0) {
+  }
 
-      gum_x86_writer_put_bytes(cw, align_pad, 8 - misalign);
+}
 
-    }
+void instrument_coverage_optimize_init(void) {
 
-    current_log_impl = cw->pc;
-    // gum_x86_writer_put_breakpoint(cw);
-    code_addr = cw->pc;
+  gpointer low_address = NULL;
 
-    code.data.assembly = template;
-    code.data.current_pc = instrument_get_offset_hash(0);
+  gum_process_enumerate_ranges(GUM_PAGE_NO_ACCESS, instrument_coverage_find_low,
+                               &low_address);
 
-    gssize current_pc_value1 =
-        GPOINTER_TO_SIZE(&instrument_previous_pc) -
-        (code_addr + offsetof(afl_log_code, data.assembly.mov_rdx_rip_off) +
-         sizeof(code.data.assembly.mov_rdx_rip_off));
-    gssize patch_offset1 =
-        offsetof(afl_log_code, data.assembly.mov_rdx_rip_off) +
-        sizeof(code.data.assembly.mov_rdx_rip_off) - sizeof(gint);
-    if (!instrument_coverage_in_range(current_pc_value1)) {
+  OKF("Low address: %p", low_address);
 
-      FATAL("Patch out of range (current_pc_value1): 0x%016lX",
-            current_pc_value1);
+  if (low_address == 0 ||
+      GPOINTER_TO_SIZE(low_address) > ((2UL << 20) - __afl_map_size)) {
 
-    }
+    FATAL("Invalid low_address: %p", low_address);
 
-    gint *dst_pc_value = (gint *)&code.bytes[patch_offset1];
-    *dst_pc_value = (gint)current_pc_value1;
+  }
 
-    gssize current_pc_value2 =
-        GPOINTER_TO_SIZE(&instrument_previous_pc) -
-        (code_addr + offsetof(afl_log_code, data.assembly.mov_rip_off_rdi) +
-         sizeof(code.data.assembly.mov_rip_off_rdi));
-    gssize patch_offset2 =
-        offsetof(afl_log_code, data.assembly.mov_rip_off_rdi) +
-        sizeof(code.data.assembly.mov_rip_off_rdi) - sizeof(gint);
+  ranges_print_debug_maps();
 
-    if (!instrument_coverage_in_range(current_pc_value2)) {
+  char *shm_env = getenv(SHM_ENV_VAR);
+  OKF("SHM_ENV_VAR: %s", shm_env);
 
-      FATAL("Patch out of range (current_pc_value2): 0x%016lX",
-            current_pc_value2);
+  if (shm_env == NULL) {
 
-    }
+    WARNF("SHM_ENV_VAR not set, using anonymous map for debugging purposes");
 
-    dst_pc_value = (gint *)&code.bytes[patch_offset2];
-    *dst_pc_value = (gint)current_pc_value2;
+    instrument_coverage_optimize_map_mmap_anon(low_address);
 
-    gsize afl_area_ptr_value =
-        GPOINTER_TO_SIZE(__afl_area_ptr) -
-        (code_addr + offsetof(afl_log_code, data.assembly.lea_rdi_rip_off) +
-         sizeof(code.data.assembly.lea_rdi_rip_off));
-    gssize afl_area_ptr_offset =
-        offsetof(afl_log_code, data.assembly.lea_rdi_rip_off) +
-        sizeof(code.data.assembly.lea_rdi_rip_off) - sizeof(gint);
+  } else {
 
-    if (!instrument_coverage_in_range(afl_area_ptr_value)) {
+    guint64 shm_env_val = g_ascii_strtoull(shm_env, NULL, 10);
 
-      FATAL("Patch out of range (afl_area_ptr_value): 0x%016lX",
-            afl_area_ptr_value);
+    if (shm_env_val == 0) {
 
-    }
+      instrument_coverage_optimize_map_mmap(shm_env, low_address);
 
-    gint *dst_afl_area_ptr_value = (gint *)&code.bytes[afl_area_ptr_offset];
-    *dst_afl_area_ptr_value = (gint)afl_area_ptr_value;
+    } else {
 
-    gum_x86_writer_put_bytes(cw, code.bytes, sizeof(afl_log_code));
+      instrument_coverage_optimize_map_shm(shm_env_val, low_address);
 
-    gum_x86_writer_put_label(cw, after_log_impl);
+    }
 
   }
 
+  OKF("__afl_area_ptr: %p", __afl_area_ptr);
+  OKF("instrument_previous_pc: %p", &instrument_previous_pc);
+
 }
 
 void instrument_coverage_optimize(const cs_insn *   instr,
                                   GumStalkerOutput *output) {
 
+  afl_log_code  code = {0};
   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);
-  gum_x86_writer_put_mov_reg_address(cw, GUM_REG_RDI, area_offset);
-  gum_x86_writer_put_call_address(cw, 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);
+  GumAddress code_addr = 0;
+
+  // gum_x86_writer_put_breakpoint(cw);
+  code_addr = cw->pc;
+  code.code = template;
+
+  gssize curr_loc_shr_1_offset =
+      offsetof(afl_log_code, code.mov_rsi_curr_loc_shr_1) +
+      sizeof(code.code.mov_rsi_curr_loc_shr_1) - sizeof(guint32);
+
+  *((guint32 *)&code.bytes[curr_loc_shr_1_offset]) =
+      (guint32)(area_offset >> 1);
+
+  gssize prev_loc_value =
+      GPOINTER_TO_SIZE(&instrument_previous_pc) -
+      (code_addr + offsetof(afl_log_code, code.xchg_rsi_prev_loc_curr_loc) +
+       sizeof(code.code.xchg_rsi_prev_loc_curr_loc));
+  gssize prev_loc_value_offset =
+      offsetof(afl_log_code, code.xchg_rsi_prev_loc_curr_loc) +
+      sizeof(code.code.xchg_rsi_prev_loc_curr_loc) - sizeof(gint);
+  if (!instrument_coverage_in_range(prev_loc_value)) {
+
+    FATAL("Patch out of range (current_pc_value1): 0x%016lX", prev_loc_value);
+
+  }
+
+  *((gint *)&code.bytes[prev_loc_value_offset]) = (gint)prev_loc_value;
+
+  gssize xor_curr_loc_offset = offsetof(afl_log_code, code.xor_rsi_curr_loc) +
+                               sizeof(code.code.xor_rsi_curr_loc) -
+                               sizeof(guint32);
+
+  *((guint32 *)&code.bytes[xor_curr_loc_offset]) = (guint32)(area_offset);
+
+  gssize add_rsi_1_offset = offsetof(afl_log_code, code.add_rsi_1) +
+                            sizeof(code.code.add_rsi_1) - sizeof(guint32) - 1;
+
+  *((guint32 *)&code.bytes[add_rsi_1_offset]) =
+      (guint32)GPOINTER_TO_SIZE(__afl_area_ptr);
+
+  gssize adc_rsi_0_ffset = offsetof(afl_log_code, code.adc_rsi_0) +
+                           sizeof(code.code.adc_rsi_0) - sizeof(guint32) - 1;
+
+  *((guint32 *)&code.bytes[adc_rsi_0_ffset]) =
+      (guint32)GPOINTER_TO_SIZE(__afl_area_ptr);
+
+  gum_x86_writer_put_bytes(cw, code.bytes, sizeof(afl_log_code));
 
 }
 
diff --git a/frida_mode/src/ranges.c b/frida_mode/src/ranges.c
index 5b6eb462..1b666fce 100644
--- a/frida_mode/src/ranges.c
+++ b/frida_mode/src/ranges.c
@@ -549,18 +549,19 @@ static GArray *merge_ranges(GArray *a) {
 
 }
 
+void ranges_print_debug_maps(void) {
+
+  gum_process_enumerate_ranges(GUM_PAGE_NO_ACCESS, print_ranges_callback, NULL);
+
+}
+
 void ranges_config(void) {
 
   if (getenv("AFL_FRIDA_DEBUG_MAPS") != NULL) { ranges_debug_maps = TRUE; }
   if (getenv("AFL_INST_LIBS") != NULL) { ranges_inst_libs = TRUE; }
   if (getenv("AFL_FRIDA_INST_JIT") != NULL) { ranges_inst_jit = TRUE; }
 
-  if (ranges_debug_maps) {
-
-    gum_process_enumerate_ranges(GUM_PAGE_NO_ACCESS, print_ranges_callback,
-                                 NULL);
-
-  }
+  if (ranges_debug_maps) { ranges_print_debug_maps(); }
 
   include_ranges = collect_ranges("AFL_FRIDA_INST_RANGES");
   exclude_ranges = collect_ranges("AFL_FRIDA_EXCLUDE_RANGES");