Frida Refactor: Split FridaHelper into each Runtime (#368)

* dynasm maybe_log

* create coverage_rt, trim helper

* add

* amd64 working

* aarch64 instrumentation, untested

* asan dir

* Revert "asan dir"

This reverts commit c7afc784819072d9fa7b8ce23adb7c9f07a21b10.

* non x86_64 fix

* clippy

* change

* change

* fix

* Fix aarch64-linux-android build

* Fix aarch64 execution

* Fix fmt

Co-authored-by: s1341 <github@shmarya.net>

fuzzers/frida_libpng/src/fuzzer.rs

@@ -52,7 +52,8 @@ use std::{
 };
 
 use libafl_frida::{
-    helper::{FridaHelper, FridaInstrumentationHelper, MAP_SIZE},
+    coverage_rt::MAP_SIZE,
+    helper::{FridaHelper, FridaInstrumentationHelper},
     FridaOptions,
 };
 
@@ -156,7 +157,7 @@ where
     ) -> Self {
         let mut stalker = Stalker::new(gum);
 
-        #[cfg(all(not(debug_assertions), target_arch = "x86_64"))]
+        #[cfg(not(all(debug_assertions, target_arch = "x86_64")))]
         for range in helper.ranges().gaps(&(0..usize::MAX)) {
             println!("excluding range: {:x}-{:x}", range.start, range.end);
             stalker.exclude(&MemoryRange::new(
@@ -312,7 +313,7 @@ unsafe fn fuzz(
         // Create an observation channel using the coverage map
         let edges_observer = HitcountsMapObserver::new(StdMapObserver::new_from_ptr(
             "edges",
-            frida_helper.map_ptr(),
+            frida_helper.map_ptr_mut(),
             MAP_SIZE,
         ));
 

libafl_frida/src/coverage_rt.rs

@@ -0,0 +1,178 @@
+use dynasmrt::{dynasm, DynasmApi, DynasmLabelApi};
+use std::ffi::c_void;
+
+#[cfg(target_arch = "x86_64")]
+use frida_gum::instruction_writer::X86Register;
+#[cfg(target_arch = "aarch64")]
+use frida_gum::instruction_writer::{Aarch64Register, IndexMode};
+
+use frida_gum::{instruction_writer::InstructionWriter, stalker::StalkerOutput};
+
+/// (Default) map size for frida coverage reporting
+pub const MAP_SIZE: usize = 64 * 1024;
+
+pub struct CoverageRuntime {
+    map: [u8; MAP_SIZE],
+    previous_pc: u64,
+    current_log_impl: u64,
+    blob_maybe_log: Option<Box<[u8]>>,
+}
+
+impl Default for CoverageRuntime {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl CoverageRuntime {
+    #[must_use]
+    pub fn new() -> Self {
+        Self {
+            map: [0u8; MAP_SIZE],
+            previous_pc: 0,
+            current_log_impl: 0,
+            blob_maybe_log: None,
+        }
+    }
+
+    pub fn init(&mut self) {
+        self.generate_maybe_log_blob();
+    }
+
+    pub fn map_ptr_mut(&mut self) -> *mut u8 {
+        self.map.as_mut_ptr()
+    }
+    #[must_use]
+    pub fn blob_maybe_log(&self) -> &[u8] {
+        self.blob_maybe_log.as_ref().unwrap()
+    }
+
+    /// A minimal `maybe_log` implementation. We insert this into the transformed instruction stream
+    /// every time we need a copy that is within a direct branch of the start of the transformed basic
+    /// block.
+    #[cfg(target_arch = "aarch64")]
+    pub fn generate_maybe_log_blob(&mut self) {
+        let mut ops = dynasmrt::VecAssembler::<dynasmrt::aarch64::Aarch64Relocation>::new(0);
+        dynasm!(ops
+            ;   .arch aarch64
+            ;   stp x1, x2, [sp, -0x10]!
+            ;   stp x3, x4, [sp, -0x10]!
+            ;   ldr x1, >map_addr
+            ;   ldr x2, >previous_loc
+            ;   ldr x4, [x2]
+            ;   eor x4, x4, x0
+            ;   mov x3, ((MAP_SIZE - 1) as u32) as u64
+            ;   and x4, x4, x3
+            ;   ldr x3, [x1, x4]
+            ;   add x3, x3, #1
+            ;   str x3, [x1, x4]
+            ;   add x0, xzr, x0, LSR #1
+            ;   str x0, [x2]
+            ;   ldp x3, x4, [sp], #0x10
+            ;   ldp x1, x2, [sp], #0x10
+            ;   ret
+            ;map_addr:
+            ;.qword &mut self.map as *mut _ as *mut c_void as i64
+            ;previous_loc:
+            ;.qword 0
+        );
+        let ops_vec = ops.finalize().unwrap();
+        self.blob_maybe_log = Some(ops_vec[..ops_vec.len() - 8].to_vec().into_boxed_slice())
+    }
+
+    /// A minimal `maybe_log` implementation. We insert this into the transformed instruction stream
+    /// every time we need a copy that is within a direct branch of the start of the transformed basic
+    /// block.
+    #[cfg(target_arch = "x86_64")]
+    pub fn generate_maybe_log_blob(&mut self) {
+        let mut ops = dynasmrt::VecAssembler::<dynasmrt::x64::X64Relocation>::new(0);
+        dynasm!(ops
+            ;   .arch x64
+            ;   pushfq
+            ;   push rax
+            ;   push rcx
+            ;   push rdx
+            ;   lea rax, [>map_addr]
+            ;   mov rax, QWORD [rax]
+            ;   lea rcx, [>previous_loc]
+            ;   mov rdx, QWORD [rcx]
+            ;   mov rdx, QWORD [rdx]
+            ;   xor rdx, rdi
+            ;   inc BYTE [rax + rdx]
+            ;   shr rdi, 1
+            ;   mov rax, QWORD [rcx]
+            ;   mov QWORD [rax], rdi
+            ;   pop rdx
+            ;   pop rcx
+            ;   pop rax
+            ;   popfq
+            ;   ret
+            ;map_addr:
+            ;.qword &mut self.map as *mut _ as *mut c_void as i64
+            ;previous_loc:
+            ;.qword 0
+        );
+        let ops_vec = ops.finalize().unwrap();
+        self.blob_maybe_log = Some(ops_vec[..ops_vec.len() - 8].to_vec().into_boxed_slice());
+    }
+
+    #[inline]
+    pub fn emit_coverage_mapping(&mut self, address: u64, output: &StalkerOutput) {
+        let writer = output.writer();
+        #[allow(clippy::cast_possible_wrap)] // gum redzone size is u32, we need an offset as i32.
+        let redzone_size = i64::from(frida_gum_sys::GUM_RED_ZONE_SIZE);
+        if self.current_log_impl == 0
+            || !writer.can_branch_directly_to(self.current_log_impl)
+            || !writer.can_branch_directly_between(writer.pc() + 128, self.current_log_impl)
+        {
+            let after_log_impl = writer.code_offset() + 1;
+
+            #[cfg(target_arch = "x86_64")]
+            writer.put_jmp_near_label(after_log_impl);
+            #[cfg(target_arch = "aarch64")]
+            writer.put_b_label(after_log_impl);
+
+            self.current_log_impl = writer.pc();
+            writer.put_bytes(self.blob_maybe_log());
+            let prev_loc_pointer = &mut self.previous_pc as *mut _ as u64; // Get the pointer to self.previous_pc
+
+            writer.put_bytes(&prev_loc_pointer.to_ne_bytes());
+
+            writer.put_label(after_log_impl);
+        }
+        #[cfg(target_arch = "x86_64")]
+        {
+            writer.put_lea_reg_reg_offset(X86Register::Rsp, X86Register::Rsp, -(redzone_size));
+            writer.put_push_reg(X86Register::Rdi);
+            writer.put_mov_reg_address(
+                X86Register::Rdi,
+                ((address >> 4) ^ (address << 8)) & (MAP_SIZE - 1) as u64,
+            );
+            writer.put_call_address(self.current_log_impl);
+            writer.put_pop_reg(X86Register::Rdi);
+            writer.put_lea_reg_reg_offset(X86Register::Rsp, X86Register::Rsp, redzone_size);
+        }
+        #[cfg(target_arch = "aarch64")]
+        {
+            writer.put_stp_reg_reg_reg_offset(
+                Aarch64Register::Lr,
+                Aarch64Register::X0,
+                Aarch64Register::Sp,
+                -(16 + redzone_size) as i64,
+                IndexMode::PreAdjust,
+            );
+            writer.put_ldr_reg_u64(
+                Aarch64Register::X0,
+                ((address >> 4) ^ (address << 8)) & (MAP_SIZE - 1) as u64,
+            );
+            writer.put_bl_imm(self.current_log_impl);
+            writer.put_ldp_reg_reg_reg_offset(
+                Aarch64Register::Lr,
+                Aarch64Register::X0,
+                Aarch64Register::Sp,
+                16 + redzone_size as i64,
+                IndexMode::PostAdjust,
+            );
+        }
+    }
+}

libafl_frida/src/helper.rs

@@ -51,6 +51,8 @@ use crate::{asan_rt::AsanRuntime, FridaOptions};
 #[cfg(windows)]
 use crate::FridaOptions;
 
+use crate::coverage_rt::CoverageRuntime;
+
 #[cfg(feature = "cmplog")]
 use crate::cmplog_rt::CmpLogRuntime;
 
@@ -92,19 +94,14 @@ pub trait FridaHelper<'a> {
     fn stalker_enabled(&self) -> bool;
 
     /// pointer to the frida coverage map
-    fn map_ptr(&mut self) -> *mut u8;
+    fn map_ptr_mut(&mut self) -> *mut u8;
 
     fn ranges(&self) -> &RangeMap<usize, (u16, String)>;
 }
 
-/// (Default) map size for frida coverage reporting
-pub const MAP_SIZE: usize = 64 * 1024;
-
 /// An helper that feeds [`FridaInProcessExecutor`] with edge-coverage instrumentation
 pub struct FridaInstrumentationHelper<'a> {
-    map: [u8; MAP_SIZE],
+    coverage_rt: CoverageRuntime,
-    previous_pc: [u64; 1],
-    current_log_impl: u64,
     #[cfg(unix)]
     current_report_impl: u64,
     /// Transformer that has to be passed to FridaInProcessExecutor
@@ -173,8 +170,8 @@ impl<'a> FridaHelper<'a> for FridaInstrumentationHelper<'a> {
         self.options.stalker_enabled()
     }
 
-    fn map_ptr(&mut self) -> *mut u8 {
+    fn map_ptr_mut(&mut self) -> *mut u8 {
-        self.map.as_mut_ptr()
+        self.coverage_rt.map_ptr_mut()
     }
 
     fn ranges(&self) -> &RangeMap<usize, (u16, String)> {
@@ -195,64 +192,6 @@ pub fn get_module_size(module_name: &str) -> usize {
     code_size
 }
 
-/// A minimal `maybe_log` implementation. We insert this into the transformed instruction stream
-/// every time we need a copy that is within a direct branch of the start of the transformed basic
-/// block.
-#[cfg(target_arch = "x86_64")]
-const MAYBE_LOG_CODE: [u8; 47] = [
-    0x9c, /* pushfq */
-    0x50, /* push rax */
-    0x51, /* push rcx */
-    0x52, /* push rdx */
-    0x48, 0x8d, 0x05, 0x24, 0x00, 0x00, 0x00, /* lea rax, sym._afl_area_ptr_ptr */
-    0x48, 0x8b, 0x00, /* mov rax, qword [rax] */
-    0x48, 0x8d, 0x0d, 0x22, 0x00, 0x00, 0x00, /* lea rcx, sym.previous_pc     */
-    0x48, 0x8b, 0x11, /* mov rdx, qword [rcx] */
-    0x48, 0x8b, 0x12, /* mov rdx, qword [rdx] */
-    0x48, 0x31, 0xfa, /* xor rdx, rdi */
-    0xfe, 0x04, 0x10, /* inc byte [rax + rdx] */
-    0x48, 0xd1, 0xef, /* shr rdi, 1 */
-    0x48, 0x8b, 0x01, /* mov rax, qword [rcx] */
-    0x48, 0x89, 0x38, /* mov qword [rax], rdi */
-    0x5a, /* pop rdx */
-    0x59, /* pop rcx */
-    0x58, /* pop rax */
-    0x9d, /* popfq */
-    0xc3, /* ret */
-
-          /* Read-only data goes here: */
-          /* uint8_t* afl_area_ptr */
-          /* uint64_t* afl_prev_loc_ptr */
-];
-
-#[cfg(target_arch = "aarch64")]
-const MAYBE_LOG_CODE: [u8; 60] = [
-    // __afl_area_ptr[current_pc ^ previous_pc]++;
-    // previous_pc = current_pc >> 1;
-    0xE1, 0x0B, 0xBF, 0xA9, // stp x1, x2, [sp, -0x10]!
-    0xE3, 0x13, 0xBF, 0xA9, // stp x3, x4, [sp, -0x10]!
-    // x0 = current_pc
-    0xa1, 0x01, 0x00, 0x58, // ldr x1, #0x30, =__afl_area_ptr
-    0x82, 0x01, 0x00, 0x58, // ldr x2, #0x38, =&previous_pc
-    0x44, 0x00, 0x40, 0xf9, // ldr x4, [x2] (=previous_pc)
-    // __afl_area_ptr[current_pc ^ previous_pc]++;
-    0x84, 0x00, 0x00, 0xca, // eor x4, x4, x0
-    0x84, 0x3c, 0x40, 0x92, // and x4, x4, 0xffff (=MAP_SIZE - 1)
-    //0x20, 0x13, 0x20, 0xd4,
-    0x23, 0x68, 0x64, 0xf8, // ldr x3, [x1, x4]
-    0x63, 0x04, 0x00, 0x91, // add x3, x3, #1
-    0x23, 0x68, 0x24, 0xf8, // str x3, [x1, x4]
-    // previous_pc = current_pc >> 1;
-    0xe0, 0x07, 0x40, 0x8b, // add x0, xzr, x0, LSR #1
-    0x40, 0x00, 0x00, 0xf9, // str x0, [x2]
-    0xE3, 0x13, 0xc1, 0xA8, // ldp x3, x4, [sp], #0x10
-    0xE1, 0x0B, 0xc1, 0xA8, // ldp x1, x2, [sp], #0x10
-    0xC0, 0x03, 0x5F, 0xD6, // ret
-
-          // &afl_area_ptr
-          // &afl_prev_loc_ptr
-];
-
 #[cfg(target_arch = "aarch64")]
 fn pc(context: &CpuContext) -> usize {
     context.pc() as usize
@@ -300,9 +239,7 @@ impl<'a> FridaInstrumentationHelper<'a> {
         }
 
         let mut helper = Self {
-            map: [0u8; MAP_SIZE],
+            coverage_rt: CoverageRuntime::new(),
-            previous_pc: [0u64; 1],
-            current_log_impl: 0,
             #[cfg(unix)]
             current_report_impl: 0,
             transformer: None,
@@ -355,6 +292,10 @@ impl<'a> FridaInstrumentationHelper<'a> {
                     .expect("failed to create directory for coverage files");
             }
 
+            if helper.options().coverage_enabled() {
+                helper.coverage_rt.init();
+            }
+
             let transformer = Transformer::from_callback(gum, |basic_block, output| {
                 let mut first = true;
                 for instruction in basic_block {
@@ -374,7 +315,7 @@ impl<'a> FridaInstrumentationHelper<'a> {
                             first = false;
                             // println!("block @ {:x} transformed to {:x}", address, output.writer().pc());
                             if helper.options().coverage_enabled() {
-                                helper.emit_coverage_mapping(address, &output);
+                                helper.coverage_rt.emit_coverage_mapping(address, &output);
                             }
                             #[cfg(unix)]
                             if helper.options().drcov_enabled() {
@@ -1484,66 +1425,4 @@ impl<'a> FridaInstrumentationHelper<'a> {
             Err(())
         }
     }
-
-    #[inline]
-    fn emit_coverage_mapping(&mut self, address: u64, output: &StalkerOutput) {
-        let writer = output.writer();
-        #[allow(clippy::cast_possible_wrap)] // gum redzone size is u32, we need an offset as i32.
-        let redzone_size = i64::from(frida_gum_sys::GUM_RED_ZONE_SIZE);
-        if self.current_log_impl == 0
-            || !writer.can_branch_directly_to(self.current_log_impl)
-            || !writer.can_branch_directly_between(writer.pc() + 128, self.current_log_impl)
-        {
-            let after_log_impl = writer.code_offset() + 1;
-
-            #[cfg(target_arch = "x86_64")]
-            writer.put_jmp_near_label(after_log_impl);
-            #[cfg(target_arch = "aarch64")]
-            writer.put_b_label(after_log_impl);
-
-            self.current_log_impl = writer.pc();
-            writer.put_bytes(&MAYBE_LOG_CODE);
-            let prev_loc_pointer = self.previous_pc.as_ptr() as usize;
-            let map_pointer = self.map.as_ptr() as usize;
-
-            writer.put_bytes(&map_pointer.to_ne_bytes());
-            writer.put_bytes(&prev_loc_pointer.to_ne_bytes());
-
-            writer.put_label(after_log_impl);
-        }
-        #[cfg(target_arch = "x86_64")]
-        {
-            writer.put_lea_reg_reg_offset(X86Register::Rsp, X86Register::Rsp, -(redzone_size));
-            writer.put_push_reg(X86Register::Rdi);
-            writer.put_mov_reg_address(
-                X86Register::Rdi,
-                ((address >> 4) ^ (address << 8)) & (MAP_SIZE - 1) as u64,
-            );
-            writer.put_call_address(self.current_log_impl);
-            writer.put_pop_reg(X86Register::Rdi);
-            writer.put_lea_reg_reg_offset(X86Register::Rsp, X86Register::Rsp, redzone_size);
-        }
-        #[cfg(target_arch = "aarch64")]
-        {
-            writer.put_stp_reg_reg_reg_offset(
-                Aarch64Register::Lr,
-                Aarch64Register::X0,
-                Aarch64Register::Sp,
-                -(16 + redzone_size) as i64,
-                IndexMode::PreAdjust,
-            );
-            writer.put_ldr_reg_u64(
-                Aarch64Register::X0,
-                ((address >> 4) ^ (address << 8)) & (MAP_SIZE - 1) as u64,
-            );
-            writer.put_bl_imm(self.current_log_impl);
-            writer.put_ldp_reg_reg_reg_offset(
-                Aarch64Register::Lr,
-                Aarch64Register::X0,
-                Aarch64Register::Sp,
-                16 + redzone_size as i64,
-                IndexMode::PostAdjust,
-            );
-        }
-    }
 }

libafl_frida/src/lib.rs

@@ -13,6 +13,8 @@ pub mod asan_errors;
 #[cfg(unix)]
 pub mod asan_rt;
 
+pub mod coverage_rt;
+
 #[cfg(feature = "cmplog")]
 /// The frida cmplog runtime
 pub mod cmplog_rt;