Frida Refactor: Separate Frida other helper functions into each Runtime (#418)

* separate asan

* fmt

* move asan out of helper.rs

* fmt

* move cmplog out of helper.rs

* fmt

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* final fix & fmt

* Fix unused imports

* Fmt

* rename files

* fix Makefile

* fmt

* clippy

Co-authored-by: s1341 <github@shmarya.net>
Co-authored-by: Dominik Maier <domenukk@gmail.com>

fuzzers/frida_libpng/src/fuzzer.rs

@@ -41,8 +41,6 @@ use libafl::{
     Error,
 };
 
-#[cfg(unix)]
-use libafl_frida::asan_errors::{AsanErrorsFeedback, AsanErrorsObserver, ASAN_ERRORS};
 use libafl_frida::{
     coverage_rt::MAP_SIZE,
     executor::FridaInProcessExecutor,
@@ -51,6 +49,8 @@ use libafl_frida::{
 };
 use libafl_targets::cmplog::{CmpLogObserver, CMPLOG_MAP};
 
+#[cfg(unix)]
+use libafl_frida::asan::errors::{AsanErrorsFeedback, AsanErrorsObserver, ASAN_ERRORS};
 fn timeout_from_millis_str(time: &str) -> Result<Duration, Error> {
     Ok(Duration::from_millis(time.parse()?))
 }

libafl_frida/src/alloc.rs

@@ -13,11 +13,11 @@ use serde::{Deserialize, Serialize};
 use std::{ffi::c_void, io};
 
 use crate::{
-    asan_errors::{AsanError, AsanErrors},
+    asan::errors::{AsanError, AsanErrors},
     FridaOptions,
 };
 
-pub(crate) struct Allocator {
+pub struct Allocator {
     #[allow(dead_code)]
     options: FridaOptions,
     page_size: usize,
@@ -44,7 +44,7 @@ macro_rules! map_to_shadow {
 }
 
 #[derive(Clone, Debug, Default, Serialize, Deserialize)]
-pub(crate) struct AllocationMetadata {
+pub struct AllocationMetadata {
     pub address: usize,
     pub size: usize,
     pub actual_size: usize,
@@ -55,6 +55,7 @@ pub(crate) struct AllocationMetadata {
 }
 
 impl Allocator {
+    #[must_use]
     pub fn new(options: FridaOptions) -> Self {
         let ret = unsafe { sysconf(_SC_PAGESIZE) };
         assert!(
@@ -193,6 +194,7 @@ impl Allocator {
     }
 
     #[must_use]
+    #[allow(clippy::missing_safety_doc)]
     pub unsafe fn alloc(&mut self, size: usize, _alignment: usize) -> *mut c_void {
         let mut is_malloc_zero = false;
         let size = if size == 0 {
@@ -270,6 +272,7 @@ impl Allocator {
         address
     }
 
+    #[allow(clippy::missing_safety_doc)]
     pub unsafe fn release(&mut self, ptr: *mut c_void) {
         let mut metadata = if let Some(metadata) = self.allocations.get_mut(&(ptr as usize)) {
             metadata
@@ -443,6 +446,7 @@ impl Allocator {
         (shadow_mapping_start, (end - start) / 8)
     }
 
+    #[must_use]
     pub fn map_to_shadow(&self, start: usize) -> usize {
         map_to_shadow!(self, start)
     }

libafl_frida/src/asan/errors.rs

@@ -21,9 +21,9 @@ use std::io::Write;
 use termcolor::{Color, ColorSpec, WriteColor};
 
 #[cfg(target_arch = "x86_64")]
-use crate::asan_rt::ASAN_SAVE_REGISTER_NAMES;
+use crate::asan::asan_rt::ASAN_SAVE_REGISTER_NAMES;
 
-use crate::{alloc::AllocationMetadata, asan_rt::ASAN_SAVE_REGISTER_COUNT, FridaOptions};
+use crate::{alloc::AllocationMetadata, asan::asan_rt::ASAN_SAVE_REGISTER_COUNT, FridaOptions};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub(crate) struct AsanReadWriteError {

libafl_frida/src/asan/mod.rs

@@ -0,0 +1,3 @@
+pub mod asan_rt;
+pub mod errors;
+pub mod hook_funcs;

libafl_frida/src/cmplog_rt.rs

@@ -7,12 +7,40 @@ extern "C" {
     pub fn __libafl_targets_cmplog_instructions(k: u64, shape: u8, arg1: u64, arg2: u64);
 }
 
+#[cfg(target_arch = "aarch64")]
+use frida_gum::{
+    instruction_writer::{Aarch64Register, IndexMode, InstructionWriter},
+    stalker::StalkerOutput,
+};
+
+#[cfg(all(feature = "cmplog", target_arch = "aarch64"))]
+use crate::helper::FridaInstrumentationHelper;
+
+pub enum SpecialCmpLogCase {
+    Tbz,
+    Tbnz,
+}
+
+#[cfg(target_arch = "aarch64")]
+use capstone::{
+    arch::{arm64::Arm64OperandType, ArchOperand::Arm64Operand},
+    Capstone, Insn,
+};
+
 pub struct CmpLogRuntime {
     ops_save_register_and_blr_to_populate: Option<Box<[u8]>>,
     ops_handle_tbz_masking: Option<Box<[u8]>>,
     ops_handle_tbnz_masking: Option<Box<[u8]>>,
 }
 
+#[cfg(all(feature = "cmplog", target_arch = "aarch64"))]
+pub enum CmplogOperandType {
+    Regid(capstone::RegId),
+    Imm(u64),
+    Cimm(u64),
+    Mem(capstone::RegId, capstone::RegId, i32, u32),
+}
+
 impl CmpLogRuntime {
     #[must_use]
     pub fn new() -> CmpLogRuntime {
@@ -175,6 +203,426 @@ impl CmpLogRuntime {
     pub fn ops_handle_tbnz_masking(&self) -> &[u8] {
         self.ops_handle_tbnz_masking.as_ref().unwrap()
     }
+
+    #[cfg(all(feature = "cmplog", target_arch = "aarch64"))]
+    #[inline]
+    /// Emit the instrumentation code which is responsible for opernads value extraction and cmplog map population
+    pub fn emit_comparison_handling(
+        &self,
+        _address: u64,
+        output: &StalkerOutput,
+        op1: CmplogOperandType,
+        op2: CmplogOperandType,
+        special_case: Option<SpecialCmpLogCase>,
+    ) {
+        let writer = output.writer();
+
+        // Preserve x0, x1:
+        writer.put_stp_reg_reg_reg_offset(
+            Aarch64Register::X0,
+            Aarch64Register::X1,
+            Aarch64Register::Sp,
+            -(16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i32) as i64,
+            IndexMode::PreAdjust,
+        );
+
+        // make sure operand1 value is saved into x0
+        match op1 {
+            CmplogOperandType::Imm(value) | CmplogOperandType::Cimm(value) => {
+                writer.put_ldr_reg_u64(Aarch64Register::X0, value);
+            }
+            CmplogOperandType::Regid(reg) => {
+                let reg = FridaInstrumentationHelper::writer_register(reg);
+                match reg {
+                    Aarch64Register::X0 | Aarch64Register::W0 => {}
+                    Aarch64Register::X1 | Aarch64Register::W1 => {
+                        writer.put_mov_reg_reg(Aarch64Register::X0, Aarch64Register::X1);
+                    }
+                    _ => {
+                        if !writer.put_mov_reg_reg(Aarch64Register::X0, reg) {
+                            writer.put_mov_reg_reg(Aarch64Register::W0, reg);
+                        }
+                    }
+                }
+            }
+            CmplogOperandType::Mem(basereg, indexreg, displacement, _width) => {
+                let basereg = FridaInstrumentationHelper::writer_register(basereg);
+                let indexreg = if indexreg.0 != 0 {
+                    Some(FridaInstrumentationHelper::writer_register(indexreg))
+                } else {
+                    None
+                };
+
+                // calculate base+index+displacment into x0
+                let displacement = displacement
+                    + if basereg == Aarch64Register::Sp {
+                        16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i32
+                    } else {
+                        0
+                    };
+
+                if indexreg.is_some() {
+                    if let Some(indexreg) = indexreg {
+                        writer.put_add_reg_reg_reg(Aarch64Register::X0, basereg, indexreg);
+                    }
+                } else {
+                    match basereg {
+                        Aarch64Register::X0 | Aarch64Register::W0 => {}
+                        Aarch64Register::X1 | Aarch64Register::W1 => {
+                            writer.put_mov_reg_reg(Aarch64Register::X0, Aarch64Register::X1);
+                        }
+                        _ => {
+                            if !writer.put_mov_reg_reg(Aarch64Register::X0, basereg) {
+                                writer.put_mov_reg_reg(Aarch64Register::W0, basereg);
+                            }
+                        }
+                    }
+                }
+
+                //add displacement
+                writer.put_add_reg_reg_imm(
+                    Aarch64Register::X0,
+                    Aarch64Register::X0,
+                    displacement as u64,
+                );
+
+                //deref into x0 to get the real value
+                writer.put_ldr_reg_reg_offset(Aarch64Register::X0, Aarch64Register::X0, 0u64);
+            }
+        }
+
+        // make sure operand2 value is saved into x1
+        match op2 {
+            CmplogOperandType::Imm(value) | CmplogOperandType::Cimm(value) => {
+                writer.put_ldr_reg_u64(Aarch64Register::X1, value);
+                match special_case {
+                    Some(inst) => match inst {
+                        SpecialCmpLogCase::Tbz => {
+                            writer.put_bytes(&self.ops_handle_tbz_masking());
+                        }
+                        SpecialCmpLogCase::Tbnz => {
+                            writer.put_bytes(&self.ops_handle_tbnz_masking());
+                        }
+                    },
+                    None => (),
+                }
+            }
+            CmplogOperandType::Regid(reg) => {
+                let reg = FridaInstrumentationHelper::writer_register(reg);
+                match reg {
+                    Aarch64Register::X1 | Aarch64Register::W1 => {}
+                    Aarch64Register::X0 | Aarch64Register::W0 => {
+                        writer.put_ldr_reg_reg_offset(
+                            Aarch64Register::X1,
+                            Aarch64Register::Sp,
+                            0u64,
+                        );
+                    }
+                    _ => {
+                        if !writer.put_mov_reg_reg(Aarch64Register::X1, reg) {
+                            writer.put_mov_reg_reg(Aarch64Register::W1, reg);
+                        }
+                    }
+                }
+            }
+            CmplogOperandType::Mem(basereg, indexreg, displacement, _width) => {
+                let basereg = FridaInstrumentationHelper::writer_register(basereg);
+                let indexreg = if indexreg.0 != 0 {
+                    Some(FridaInstrumentationHelper::writer_register(indexreg))
+                } else {
+                    None
+                };
+
+                // calculate base+index+displacment into x1
+                let displacement = displacement
+                    + if basereg == Aarch64Register::Sp {
+                        16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i32
+                    } else {
+                        0
+                    };
+
+                if indexreg.is_some() {
+                    if let Some(indexreg) = indexreg {
+                        match indexreg {
+                            Aarch64Register::X0 | Aarch64Register::W0 => {
+                                match basereg {
+                                    Aarch64Register::X1 | Aarch64Register::W1 => {
+                                        // x0 is overwrittern indexreg by op1 value.
+                                        // x1 is basereg
+
+                                        // Preserve x2, x3:
+                                        writer.put_stp_reg_reg_reg_offset(
+                                            Aarch64Register::X2,
+                                            Aarch64Register::X3,
+                                            Aarch64Register::Sp,
+                                            -(16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i32) as i64,
+                                            IndexMode::PreAdjust,
+                                        );
+
+                                        //reload indexreg to x2
+                                        writer.put_ldr_reg_reg_offset(
+                                            Aarch64Register::X2,
+                                            Aarch64Register::Sp,
+                                            0u64,
+                                        );
+                                        //add them into basereg==x1
+                                        writer.put_add_reg_reg_reg(
+                                            basereg,
+                                            basereg,
+                                            Aarch64Register::X2,
+                                        );
+
+                                        // Restore x2, x3
+                                        assert!(writer.put_ldp_reg_reg_reg_offset(
+                                            Aarch64Register::X2,
+                                            Aarch64Register::X3,
+                                            Aarch64Register::Sp,
+                                            16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i64,
+                                            IndexMode::PostAdjust,
+                                        ));
+                                    }
+                                    _ => {
+                                        // x0 is overwrittern indexreg by op1 value.
+                                        // basereg is not x1 nor x0
+
+                                        //reload indexreg to x1
+                                        writer.put_ldr_reg_reg_offset(
+                                            Aarch64Register::X1,
+                                            Aarch64Register::Sp,
+                                            0u64,
+                                        );
+                                        //add basereg into indexreg==x1
+                                        writer.put_add_reg_reg_reg(
+                                            Aarch64Register::X1,
+                                            basereg,
+                                            Aarch64Register::X1,
+                                        );
+                                    }
+                                }
+                            }
+                            Aarch64Register::X1 | Aarch64Register::W1 => {
+                                match basereg {
+                                    Aarch64Register::X0 | Aarch64Register::W0 => {
+                                        // x0 is overwrittern basereg by op1 value.
+                                        // x1 is indexreg
+
+                                        // Preserve x2, x3:
+                                        writer.put_stp_reg_reg_reg_offset(
+                                            Aarch64Register::X2,
+                                            Aarch64Register::X3,
+                                            Aarch64Register::Sp,
+                                            -(16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i32) as i64,
+                                            IndexMode::PreAdjust,
+                                        );
+
+                                        //reload basereg to x2
+                                        writer.put_ldr_reg_reg_offset(
+                                            Aarch64Register::X2,
+                                            Aarch64Register::Sp,
+                                            0u64,
+                                        );
+                                        //add basereg into indexreg==x1
+                                        writer.put_add_reg_reg_reg(
+                                            indexreg,
+                                            Aarch64Register::X2,
+                                            indexreg,
+                                        );
+
+                                        // Restore x2, x3
+                                        assert!(writer.put_ldp_reg_reg_reg_offset(
+                                            Aarch64Register::X2,
+                                            Aarch64Register::X3,
+                                            Aarch64Register::Sp,
+                                            16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i64,
+                                            IndexMode::PostAdjust,
+                                        ));
+                                    }
+                                    _ => {
+                                        // indexreg is x1
+                                        // basereg is not x0 and not x1
+
+                                        //add them into x1
+                                        writer.put_add_reg_reg_reg(indexreg, basereg, indexreg);
+                                    }
+                                }
+                            }
+                            _ => {
+                                match basereg {
+                                    Aarch64Register::X0 | Aarch64Register::W0 => {
+                                        //basereg is overwrittern by op1 value
+                                        //index reg is not x0 nor x1
+
+                                        //reload basereg to x1
+                                        writer.put_ldr_reg_reg_offset(
+                                            Aarch64Register::X1,
+                                            Aarch64Register::Sp,
+                                            0u64,
+                                        );
+                                        //add indexreg to basereg==x1
+                                        writer.put_add_reg_reg_reg(
+                                            Aarch64Register::X1,
+                                            Aarch64Register::X1,
+                                            indexreg,
+                                        );
+                                    }
+                                    _ => {
+                                        //basereg is not x0, can be x1
+                                        //index reg is not x0 nor x1
+
+                                        //add them into x1
+                                        writer.put_add_reg_reg_reg(
+                                            Aarch64Register::X1,
+                                            basereg,
+                                            indexreg,
+                                        );
+                                    }
+                                }
+                            }
+                        }
+                    }
+                } else {
+                    match basereg {
+                        Aarch64Register::X1 | Aarch64Register::W1 => {}
+                        Aarch64Register::X0 | Aarch64Register::W0 => {
+                            // x0 is overwrittern basereg by op1 value.
+                            //reload basereg to x1
+                            writer.put_ldr_reg_reg_offset(
+                                Aarch64Register::X1,
+                                Aarch64Register::Sp,
+                                0u64,
+                            );
+                        }
+                        _ => {
+                            writer.put_mov_reg_reg(Aarch64Register::W1, basereg);
+                        }
+                    }
+                }
+
+                // add displacement
+                writer.put_add_reg_reg_imm(
+                    Aarch64Register::X1,
+                    Aarch64Register::X1,
+                    displacement as u64,
+                );
+                //deref into x1 to get the real value
+                writer.put_ldr_reg_reg_offset(Aarch64Register::X1, Aarch64Register::X1, 0u64);
+            }
+        }
+
+        //call cmplog runtime to populate the values map
+        writer.put_bytes(&self.ops_save_register_and_blr_to_populate());
+
+        // Restore x0, x1
+        assert!(writer.put_ldp_reg_reg_reg_offset(
+            Aarch64Register::X0,
+            Aarch64Register::X1,
+            Aarch64Register::Sp,
+            16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i64,
+            IndexMode::PostAdjust,
+        ));
+    }
+
+    #[cfg(all(feature = "cmplog", target_arch = "aarch64"))]
+    #[inline]
+    /// Check if the current instruction is cmplog relevant one(any opcode which sets the flags)
+    pub fn cmplog_is_interesting_instruction(
+        &self,
+        capstone: &Capstone,
+        _address: u64,
+        instr: &Insn,
+    ) -> Result<
+        (
+            CmplogOperandType,
+            CmplogOperandType,
+            Option<SpecialCmpLogCase>,
+        ),
+        (),
+    > {
+        // We only care for compare instrunctions - aka instructions which set the flags
+        match instr.mnemonic().unwrap() {
+            "cmp" | "ands" | "subs" | "adds" | "negs" | "ngcs" | "sbcs" | "bics" | "cbz"
+            | "cbnz" | "tbz" | "tbnz" | "adcs" => (),
+            _ => return Err(()),
+        }
+        let mut operands = capstone
+            .insn_detail(instr)
+            .unwrap()
+            .arch_detail()
+            .operands();
+
+        // cbz - 1 operand, tbz - 3 operands
+        let special_case = [
+            "cbz", "cbnz", "tbz", "tbnz", "subs", "adds", "ands", "sbcs", "bics", "adcs",
+        ]
+        .contains(&instr.mnemonic().unwrap());
+        if operands.len() != 2 && !special_case {
+            return Err(());
+        }
+
+        // handle special opcodes case which have 3 operands, but the 1st(dest) is not important to us
+        if ["subs", "adds", "ands", "sbcs", "bics", "adcs"].contains(&instr.mnemonic().unwrap()) {
+            //remove the dest operand from the list
+            operands.remove(0);
+        }
+
+        // cbz marked as special since there is only 1 operand
+        let special_case = match instr.mnemonic().unwrap() {
+            "cbz" | "cbnz" => true,
+            _ => false,
+        };
+
+        let operand1 = if let Arm64Operand(arm64operand) = operands.first().unwrap() {
+            match arm64operand.op_type {
+                Arm64OperandType::Reg(regid) => Some(CmplogOperandType::Regid(regid)),
+                Arm64OperandType::Imm(val) => Some(CmplogOperandType::Imm(val as u64)),
+                Arm64OperandType::Mem(opmem) => Some(CmplogOperandType::Mem(
+                    opmem.base(),
+                    opmem.index(),
+                    opmem.disp(),
+                    FridaInstrumentationHelper::instruction_width(instr, &operands),
+                )),
+                Arm64OperandType::Cimm(val) => Some(CmplogOperandType::Cimm(val as u64)),
+                _ => return Err(()),
+            }
+        } else {
+            None
+        };
+
+        let operand2 = match special_case {
+            true => Some(CmplogOperandType::Imm(0)),
+            false => {
+                if let Arm64Operand(arm64operand2) = &operands[1] {
+                    match arm64operand2.op_type {
+                        Arm64OperandType::Reg(regid) => Some(CmplogOperandType::Regid(regid)),
+                        Arm64OperandType::Imm(val) => Some(CmplogOperandType::Imm(val as u64)),
+                        Arm64OperandType::Mem(opmem) => Some(CmplogOperandType::Mem(
+                            opmem.base(),
+                            opmem.index(),
+                            opmem.disp(),
+                            FridaInstrumentationHelper::instruction_width(instr, &operands),
+                        )),
+                        Arm64OperandType::Cimm(val) => Some(CmplogOperandType::Cimm(val as u64)),
+                        _ => return Err(()),
+                    }
+                } else {
+                    None
+                }
+            }
+        };
+
+        // tbz will need to have special handling at emit time(masking operand1 value with operand2)
+        let special_case = match instr.mnemonic().unwrap() {
+            "tbz" => Some(SpecialCmpLogCase::Tbz),
+            "tbnz" => Some(SpecialCmpLogCase::Tbnz),
+            _ => None,
+        };
+
+        if operand1.is_some() && operand2.is_some() {
+            Ok((operand1.unwrap(), operand2.unwrap(), special_case))
+        } else {
+            Err(())
+        }
+    }
 }
 
 impl Default for CmpLogRuntime {

libafl_frida/src/executor.rs

@@ -18,7 +18,7 @@ use libafl::{
 };
 
 #[cfg(unix)]
-use crate::asan_errors::ASAN_ERRORS;
+use crate::asan::errors::ASAN_ERRORS;
 
 #[cfg(windows)]
 use libafl::executors::inprocess::{HasInProcessHandlers, InProcessHandlers};
@@ -116,7 +116,7 @@ where
         #[cfg(all(not(debug_assertions), target_arch = "x86_64"))]
         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(

libafl_frida/src/lib.rs

@@ -6,12 +6,9 @@ It can report coverage and, on supported architecutres, even reports memory acce
 /// The frida-asan allocator
 #[cfg(unix)]
 pub mod alloc;
-/// Handling of ASAN errors
+
 #[cfg(unix)]
-pub mod asan_errors;
+pub mod asan;
-/// The frida address sanitizer runtime
-#[cfg(unix)]
-pub mod asan_rt;
 
 pub mod coverage_rt;