From fb21c4ff82ba9dd2396e8f50a064de9e6b01796c Mon Sep 17 00:00:00 2001 From: Dongjia Zhang Date: Tue, 1 Feb 2022 22:34:53 +0900 Subject: [PATCH] Frida Runtime Tuples (#457) * an attempt to make runtimes into tuples * wip * wip * wipp * getter * refactor * fmt * fix * compiles * fuzzer change * coverage working * asan & less unwrap() & fixes * inst size, fmt * build & coverage works on asan * amd64 fix --- fuzzers/frida_libpng/src/fuzzer.rs | 14 +- libafl_frida/src/asan/asan_rt.rs | 155 +++++---- libafl_frida/src/cmplog_rt.rs | 51 ++- libafl_frida/src/coverage_rt.rs | 37 ++- libafl_frida/src/drcov_rt.rs | 55 +++- libafl_frida/src/executor.rs | 31 +- libafl_frida/src/helper.rs | 492 ++++++++++++----------------- libafl_frida/src/lib.rs | 3 + libafl_frida/src/utils.rs | 136 ++++++++ 9 files changed, 562 insertions(+), 412 deletions(-) create mode 100644 libafl_frida/src/utils.rs diff --git a/fuzzers/frida_libpng/src/fuzzer.rs b/fuzzers/frida_libpng/src/fuzzer.rs index 5c986a3a1b..410e3a430b 100644 --- a/fuzzers/frida_libpng/src/fuzzer.rs +++ b/fuzzers/frida_libpng/src/fuzzer.rs @@ -45,10 +45,8 @@ use libafl::{ }; use libafl_frida::{ - coverage_rt::MAP_SIZE, - executor::FridaInProcessExecutor, - helper::{FridaHelper, FridaInstrumentationHelper}, - FridaOptions, + asan::asan_rt::AsanRuntime, coverage_rt::CoverageRuntime, coverage_rt::MAP_SIZE, + executor::FridaInProcessExecutor, helper::FridaInstrumentationHelper, FridaOptions, }; use libafl_targets::cmplog::{CmpLogObserver, CMPLOG_MAP}; @@ -208,17 +206,20 @@ unsafe fn fuzz( let gum = Gum::obtain(); let frida_options = FridaOptions::parse_env_options(); + let coverage = CoverageRuntime::new(); + // let asan = AsanRuntime::new(frida_options.clone()); let mut frida_helper = FridaInstrumentationHelper::new( &gum, &frida_options, module_name, modules_to_instrument, + tuple_list!(coverage), ); // Create an observation channel using the coverage map let edges_observer = HitcountsMapObserver::new(StdMapObserver::new_from_ptr( "edges", - frida_helper.map_ptr_mut(), + frida_helper.map_ptr_mut().unwrap(), MAP_SIZE, )); @@ -289,10 +290,7 @@ unsafe fn fuzz( // A fuzzer with feedbacks and a corpus scheduler let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); - #[cfg(unix)] - frida_helper.register_thread(); // Create the executor for an in-process function with just one observer for edge coverage - #[cfg(unix)] let mut executor = FridaInProcessExecutor::new( &gum, diff --git a/libafl_frida/src/asan/asan_rt.rs b/libafl_frida/src/asan/asan_rt.rs index f3a8a891ca..cc230d16f8 100644 --- a/libafl_frida/src/asan/asan_rt.rs +++ b/libafl_frida/src/asan/asan_rt.rs @@ -13,9 +13,9 @@ use core::{ }; use frida_gum::{ModuleDetails, NativePointer, RangeDetails}; use hashbrown::HashMap; +use libafl::bolts::AsSlice; use nix::sys::mman::{mmap, MapFlags, ProtFlags}; - -use crate::helper::FridaInstrumentationHelper; +use rangemap::RangeMap; #[cfg(target_arch = "aarch64")] use capstone::{ @@ -56,9 +56,14 @@ use std::{ffi::c_void, ptr::write_volatile}; use crate::{ alloc::Allocator, asan::errors::{AsanError, AsanErrors, AsanReadWriteError, ASAN_ERRORS}, + helper::FridaRuntime, + utils::writer_register, FridaOptions, }; +#[cfg(target_arch = "aarch64")] +use crate::utils::instruction_width; + extern "C" { fn __register_frame(begin: *mut c_void); } @@ -112,6 +117,7 @@ pub const ASAN_SAVE_REGISTER_COUNT: usize = 32; /// even if the target would not have crashed under normal conditions. /// this helps finding mem errors early. pub struct AsanRuntime { + check_for_leaks_enabled: bool, current_report_impl: u64, allocator: Allocator, regs: [usize; ASAN_SAVE_REGISTER_COUNT], @@ -146,38 +152,16 @@ impl Debug for AsanRuntime { } } -impl AsanRuntime { - /// Create a new `AsanRuntime` - #[must_use] - pub fn new(options: FridaOptions) -> AsanRuntime { - Self { - current_report_impl: 0, - allocator: Allocator::new(options.clone()), - regs: [0; ASAN_SAVE_REGISTER_COUNT], - blob_report: None, - blob_check_mem_byte: None, - blob_check_mem_halfword: None, - blob_check_mem_dword: None, - blob_check_mem_qword: None, - blob_check_mem_16bytes: None, - blob_check_mem_3bytes: None, - blob_check_mem_6bytes: None, - blob_check_mem_12bytes: None, - blob_check_mem_24bytes: None, - blob_check_mem_32bytes: None, - blob_check_mem_48bytes: None, - blob_check_mem_64bytes: None, - stalked_addresses: HashMap::new(), - options, - module_map: None, - suppressed_addresses: Vec::new(), - shadow_check_func: None, - } - } +impl FridaRuntime for AsanRuntime { /// Initialize the runtime so that it is read for action. Take care not to move the runtime /// instance after this function has been called, as the generated blobs would become /// invalid! - pub fn init(&mut self, _gum: &Gum, modules_to_instrument: &[&str]) { + fn init( + &mut self, + gum: &Gum, + _ranges: &RangeMap, + modules_to_instrument: &[&str], + ) { unsafe { ASAN_ERRORS = Some(AsanErrors::new(self.options.clone())); } @@ -196,7 +180,7 @@ impl AsanRuntime { } } - self.hook_functions(_gum); + self.hook_functions(gum); /* unsafe { let mem = self.allocator.alloc(0xac + 2, 8); @@ -271,6 +255,64 @@ impl AsanRuntime { // assert!((self.shadow_check_func.unwrap())(((mem2 as usize) + 8875) as *const c_void, 4)); } */ + self.register_thread(); + } + fn pre_exec( + &mut self, + input: &I, + ) -> Result<(), libafl::Error> { + let target_bytes = input.target_bytes(); + let slice = target_bytes.as_slice(); + + self.unpoison(slice.as_ptr() as usize, slice.len()); + Ok(()) + } + + fn post_exec( + &mut self, + input: &I, + ) -> Result<(), libafl::Error> { + if self.check_for_leaks_enabled { + self.check_for_leaks(); + } + + let target_bytes = input.target_bytes(); + let slice = target_bytes.as_slice(); + self.poison(slice.as_ptr() as usize, slice.len()); + self.reset_allocations(); + + Ok(()) + } +} + +impl AsanRuntime { + /// Create a new `AsanRuntime` + #[must_use] + pub fn new(options: FridaOptions) -> AsanRuntime { + Self { + check_for_leaks_enabled: options.asan_detect_leaks(), + current_report_impl: 0, + allocator: Allocator::new(options.clone()), + regs: [0; ASAN_SAVE_REGISTER_COUNT], + blob_report: None, + blob_check_mem_byte: None, + blob_check_mem_halfword: None, + blob_check_mem_dword: None, + blob_check_mem_qword: None, + blob_check_mem_16bytes: None, + blob_check_mem_3bytes: None, + blob_check_mem_6bytes: None, + blob_check_mem_12bytes: None, + blob_check_mem_24bytes: None, + blob_check_mem_32bytes: None, + blob_check_mem_48bytes: None, + blob_check_mem_64bytes: None, + stalked_addresses: HashMap::new(), + options, + module_map: None, + suppressed_addresses: Vec::new(), + shadow_check_func: None, + } } /// Reset all allocations so that they can be reused for new allocation requests. @@ -2091,29 +2133,27 @@ impl AsanRuntime { /// Determine if the instruction is 'interesting' for the purposes of ASAN #[cfg(target_arch = "aarch64")] + #[must_use] #[inline] pub fn asan_is_interesting_instruction( &self, capstone: &Capstone, _address: u64, instr: &Insn, - ) -> Result< - ( - capstone::RegId, - capstone::RegId, - i32, - u32, - Arm64Shift, - Arm64Extender, - ), - (), - > { + ) -> Option<( + capstone::RegId, + capstone::RegId, + i32, + u32, + Arm64Shift, + Arm64Extender, + )> { // We have to ignore these instructions. Simulating them with their side effects is // complex, to say the least. match instr.mnemonic().unwrap() { "ldaxr" | "stlxr" | "ldxr" | "stxr" | "ldar" | "stlr" | "ldarb" | "ldarh" | "ldaxp" | "ldaxrb" | "ldaxrh" | "stlrb" | "stlrh" | "stlxp" | "stlxrb" | "stlxrh" | "ldxrb" - | "ldxrh" | "stxrb" | "stxrh" => return Err(()), + | "ldxrh" | "stxrb" | "stxrh" => return None, _ => (), } @@ -2123,28 +2163,29 @@ impl AsanRuntime { .arch_detail() .operands(); if operands.len() < 2 { - return Err(()); + return None; } if let Arm64Operand(arm64operand) = operands.last().unwrap() { if let Arm64OperandType::Mem(opmem) = arm64operand.op_type { - return Ok(( + return Some(( opmem.base(), opmem.index(), opmem.disp(), - FridaInstrumentationHelper::instruction_width(instr, &operands), + instruction_width(instr, &operands), arm64operand.shift, arm64operand.ext, )); } } - Err(()) + None } /// Checks if the current instruction is interesting for address sanitization. #[cfg(all(target_arch = "x86_64", unix))] #[inline] + #[must_use] #[allow(clippy::unused_self)] #[allow(clippy::result_unit_err)] pub fn asan_is_interesting_instruction( @@ -2152,7 +2193,7 @@ impl AsanRuntime { capstone: &Capstone, _address: u64, instr: &Insn, - ) -> Result<(RegId, u8, RegId, RegId, i32, i64), ()> { + ) -> Option<(RegId, u8, RegId, RegId, i32, i64)> { let operands = capstone .insn_detail(instr) .unwrap() @@ -2162,7 +2203,7 @@ impl AsanRuntime { // put nop into the white-list so that instructions like // like `nop dword [rax + rax]` does not get caught. match instr.mnemonic().unwrap() { - "lea" | "nop" => return Err(()), + "lea" | "nop" => return None, _ => (), } @@ -2170,7 +2211,7 @@ impl AsanRuntime { // This is a TODO! In this case, both the src and the dst are mem operand // so we would need to return two operadns? if instr.mnemonic().unwrap().starts_with("rep") { - return Err(()); + return None; } for operand in operands { @@ -2190,7 +2231,7 @@ impl AsanRuntime { ); */ if opmem.segment() == RegId(0) { - return Ok(( + return Some(( opmem.segment(), x86operand.size, opmem.base(), @@ -2203,7 +2244,7 @@ impl AsanRuntime { } } - Err(()) + None } /// Emits a asan shadow byte check. @@ -2229,14 +2270,14 @@ impl AsanRuntime { let basereg = if basereg.0 == 0 { None } else { - let reg = FridaInstrumentationHelper::writer_register(basereg); + let reg = writer_register(basereg); Some(reg) }; let indexreg = if indexreg.0 == 0 { None } else { - let reg = FridaInstrumentationHelper::writer_register(indexreg); + let reg = writer_register(indexreg); Some(reg) }; @@ -2399,9 +2440,9 @@ impl AsanRuntime { let redzone_size = frida_gum_sys::GUM_RED_ZONE_SIZE as i32; let writer = output.writer(); - let basereg = FridaInstrumentationHelper::writer_register(basereg); + let basereg = writer_register(basereg); let indexreg = if indexreg.0 != 0 { - Some(FridaInstrumentationHelper::writer_register(indexreg)) + Some(writer_register(indexreg)) } else { None }; diff --git a/libafl_frida/src/cmplog_rt.rs b/libafl_frida/src/cmplog_rt.rs index 835dc40e83..59adb8a934 100644 --- a/libafl_frida/src/cmplog_rt.rs +++ b/libafl_frida/src/cmplog_rt.rs @@ -4,10 +4,16 @@ //! related to the input. //! Read the [`RedQueen`](https://www.ndss-symposium.org/ndss-paper/redqueen-fuzzing-with-input-to-state-correspondence/) paper for the general concepts. use dynasmrt::{dynasm, DynasmApi, DynasmLabelApi}; +use libafl::{ + inputs::{HasTargetBytes, Input}, + Error, +}; use libafl_targets; use libafl_targets::CMPLOG_MAP_W; +use rangemap::RangeMap; use std::ffi::c_void; +use crate::helper::FridaRuntime; extern "C" { /// Tracks cmplog instructions pub fn __libafl_targets_cmplog_instructions(k: u64, shape: u8, arg1: u64, arg2: u64); @@ -20,7 +26,7 @@ use frida_gum::{ }; #[cfg(all(feature = "cmplog", target_arch = "aarch64"))] -use crate::helper::FridaInstrumentationHelper; +use crate::utils::{instruction_width, writer_register}; #[cfg(all(feature = "cmplog", target_arch = "aarch64"))] /// Speciial CmpLog Cases for `aarch64` @@ -61,6 +67,27 @@ pub struct CmpLogRuntime { ops_handle_tbnz_masking: Option>, } +impl FridaRuntime for CmpLogRuntime { + /// Initialize this `CmpLog` runtime. + /// This will generate the instrumentation blobs for the current arch. + fn init( + &mut self, + _gum: &frida_gum::Gum, + _ranges: &RangeMap, + _modules_to_instrument: &[&str], + ) { + self.generate_instrumentation_blobs(); + } + + fn pre_exec(&mut self, _input: &I) -> Result<(), Error> { + Ok(()) + } + + fn post_exec(&mut self, _input: &I) -> Result<(), Error> { + Ok(()) + } +} + impl CmpLogRuntime { /// Create a new [`CmpLogRuntime`] #[must_use] @@ -201,12 +228,6 @@ impl CmpLogRuntime { ); } - /// Initialize this `CmpLog` runtime. - /// This will generate the instrumentation blobs for the current arch. - pub fn init(&mut self) { - self.generate_instrumentation_blobs(); - } - /// Get the blob which saves the context, jumps to the populate function and restores the context #[inline] #[must_use] @@ -256,7 +277,7 @@ impl CmpLogRuntime { writer.put_ldr_reg_u64(Aarch64Register::X0, value); } CmplogOperandType::Regid(reg) => { - let reg = FridaInstrumentationHelper::writer_register(reg); + let reg = writer_register(reg); match reg { Aarch64Register::X0 | Aarch64Register::W0 => {} Aarch64Register::X1 | Aarch64Register::W1 => { @@ -270,9 +291,9 @@ impl CmpLogRuntime { } } CmplogOperandType::Mem(basereg, indexreg, displacement, _width) => { - let basereg = FridaInstrumentationHelper::writer_register(basereg); + let basereg = writer_register(basereg); let indexreg = if indexreg.0 != 0 { - Some(FridaInstrumentationHelper::writer_register(indexreg)) + Some(writer_register(indexreg)) } else { None }; @@ -332,7 +353,7 @@ impl CmpLogRuntime { } } CmplogOperandType::Regid(reg) => { - let reg = FridaInstrumentationHelper::writer_register(reg); + let reg = writer_register(reg); match reg { Aarch64Register::X1 | Aarch64Register::W1 => {} Aarch64Register::X0 | Aarch64Register::W0 => { @@ -350,9 +371,9 @@ impl CmpLogRuntime { } } CmplogOperandType::Mem(basereg, indexreg, displacement, _width) => { - let basereg = FridaInstrumentationHelper::writer_register(basereg); + let basereg = writer_register(basereg); let indexreg = if indexreg.0 != 0 { - Some(FridaInstrumentationHelper::writer_register(indexreg)) + Some(writer_register(indexreg)) } else { None }; @@ -603,7 +624,7 @@ impl CmpLogRuntime { opmem.base(), opmem.index(), opmem.disp(), - FridaInstrumentationHelper::instruction_width(instr, &operands), + instruction_width(instr, &operands), )), Arm64OperandType::Cimm(val) => Some(CmplogOperandType::Cimm(val as u64)), _ => return Err(()), @@ -623,7 +644,7 @@ impl CmpLogRuntime { opmem.base(), opmem.index(), opmem.disp(), - FridaInstrumentationHelper::instruction_width(instr, &operands), + instruction_width(instr, &operands), )), Arm64OperandType::Cimm(val) => Some(CmplogOperandType::Cimm(val as u64)), _ => return Err(()), diff --git a/libafl_frida/src/coverage_rt.rs b/libafl_frida/src/coverage_rt.rs index c4f2c4d4a9..dc9ae2d932 100644 --- a/libafl_frida/src/coverage_rt.rs +++ b/libafl_frida/src/coverage_rt.rs @@ -1,6 +1,10 @@ //! Functionality regarding binary-only coverage collection. use core::ptr::addr_of_mut; use dynasmrt::{dynasm, DynasmApi, DynasmLabelApi}; +use rangemap::RangeMap; + +#[cfg(target_arch = "aarch64")] +use std::ffi::c_void; #[cfg(target_arch = "x86_64")] use frida_gum::instruction_writer::X86Register; @@ -9,6 +13,8 @@ use frida_gum::instruction_writer::{Aarch64Register, IndexMode}; use frida_gum::{instruction_writer::InstructionWriter, stalker::StalkerOutput}; +use crate::helper::FridaRuntime; + /// (Default) map size for frida coverage reporting pub const MAP_SIZE: usize = 64 * 1024; @@ -27,6 +33,32 @@ impl Default for CoverageRuntime { } } +impl FridaRuntime for CoverageRuntime { + /// Initialize the coverage runtime + fn init( + &mut self, + _gum: &frida_gum::Gum, + _ranges: &RangeMap, + _modules_to_instrument: &[&str], + ) { + self.generate_maybe_log_blob(); + } + + fn pre_exec( + &mut self, + _input: &I, + ) -> Result<(), libafl::Error> { + Ok(()) + } + + fn post_exec( + &mut self, + _input: &I, + ) -> Result<(), libafl::Error> { + Ok(()) + } +} + impl CoverageRuntime { /// Create a new coverage runtime #[must_use] @@ -39,11 +71,6 @@ impl CoverageRuntime { } } - /// Initialize the coverage runtime - pub fn init(&mut self) { - self.generate_maybe_log_blob(); - } - /// Retrieve the coverage map pointer pub fn map_ptr_mut(&mut self) -> *mut u8 { self.map.as_mut_ptr() diff --git a/libafl_frida/src/drcov_rt.rs b/libafl_frida/src/drcov_rt.rs index e4152e276a..1585fdb43d 100644 --- a/libafl_frida/src/drcov_rt.rs +++ b/libafl_frida/src/drcov_rt.rs @@ -1,4 +1,5 @@ //! Generates `DrCov` traces +use crate::helper::FridaRuntime; use ahash::AHasher; use libafl::{ bolts::AsSlice, @@ -7,6 +8,7 @@ use libafl::{ }; use libafl_targets::drcov::{DrCovBasicBlock, DrCovWriter}; use rangemap::RangeMap; +use std::collections::HashMap; use std::hash::Hasher; /// Generates `DrCov` traces @@ -16,32 +18,30 @@ pub struct DrCovRuntime { pub drcov_basic_blocks: Vec, /// The memory ragnes of this target ranges: RangeMap, + stalked_addresses: HashMap, } -impl DrCovRuntime { - /// Creates a new [`DrCovRuntime`] - #[must_use] - pub fn new() -> Self { - Self { - drcov_basic_blocks: vec![], - ranges: RangeMap::new(), - } - } - +impl FridaRuntime for DrCovRuntime { /// initializes this runtime wiith the given `ranges` - pub fn init(&mut self, ranges: &RangeMap) { + fn init( + &mut self, + _gum: &frida_gum::Gum, + ranges: &RangeMap, + _modules_to_instrument: &[&str], + ) { self.ranges = ranges.clone(); + std::fs::create_dir_all("./coverage") + .expect("failed to create directory for coverage files"); } /// Called before execution, does nothing - #[allow(clippy::unused_self)] - pub fn pre_exec(&mut self, _input: &I) -> Result<(), Error> { + fn pre_exec(&mut self, _input: &I) -> Result<(), Error> { Ok(()) } /// Called after execution, writes the trace to a unique `DrCov` file for this trace /// into `./coverage/.drcov` - pub fn post_exec(&mut self, input: &I) -> Result<(), Error> { + fn post_exec(&mut self, input: &I) -> Result<(), Error> { let mut hasher = AHasher::new_with_keys(0, 0); hasher.write(input.target_bytes().as_slice()); @@ -53,6 +53,33 @@ impl DrCovRuntime { } } +impl DrCovRuntime { + /// Creates a new [`DrCovRuntime`] + #[must_use] + pub fn new() -> Self { + Self { + drcov_basic_blocks: vec![], + ranges: RangeMap::new(), + stalked_addresses: HashMap::new(), + } + } + + /// Add a stalked address to real address mapping. + #[inline] + pub fn add_stalked_address(&mut self, stalked: usize, real: usize) { + self.stalked_addresses.insert(stalked, real); + } + + /// Resolves the real address from a stalker stalked address if possible, if there is no + /// real address, the stalked address is returned. + #[must_use] + pub fn real_address_for_stalked(&self, stalked: usize) -> usize { + self.stalked_addresses + .get(&stalked) + .map_or(stalked, |addr| *addr) + } +} + impl Default for DrCovRuntime { fn default() -> Self { Self::new() diff --git a/libafl_frida/src/executor.rs b/libafl_frida/src/executor.rs index 41e04d63f7..a8d69ae87a 100644 --- a/libafl_frida/src/executor.rs +++ b/libafl_frida/src/executor.rs @@ -1,4 +1,4 @@ -use crate::helper::FridaHelper; +use crate::helper::{FridaInstrumentationHelper, FridaRuntimeTuple}; use core::fmt::{self, Debug, Formatter}; use frida_gum::{ @@ -21,9 +21,8 @@ use crate::asan::errors::ASAN_ERRORS; use libafl::executors::inprocess::{HasInProcessHandlers, InProcessHandlers}; /// The [`FridaInProcessExecutor`] is an [`Executor`] that executes the target in the same process, usinig [`frida`](https://frida.re/) for binary-only instrumentation. -pub struct FridaInProcessExecutor<'a, 'b, 'c, FH, H, I, OT, S> +pub struct FridaInProcessExecutor<'a, 'b, 'c, H, I, OT, RT, S> where - FH: FridaHelper<'b>, H: FnMut(&I) -> ExitKind, I: Input + HasTargetBytes, OT: ObserversTuple, @@ -32,14 +31,13 @@ where /// Frida's dynamic rewriting engine stalker: Stalker<'a>, /// User provided callback for instrumentation - helper: &'c mut FH, + helper: &'c mut FridaInstrumentationHelper<'b, RT>, followed: bool, _phantom: PhantomData<&'b u8>, } -impl<'a, 'b, 'c, FH, H, I, OT, S> Debug for FridaInProcessExecutor<'a, 'b, 'c, FH, H, I, OT, S> +impl<'a, 'b, 'c, H, I, OT, RT, S> Debug for FridaInProcessExecutor<'a, 'b, 'c, H, I, OT, RT, S> where - FH: FridaHelper<'b>, H: FnMut(&I) -> ExitKind, I: Input + HasTargetBytes, OT: ObserversTuple, @@ -53,13 +51,13 @@ where } } -impl<'a, 'b, 'c, EM, FH, H, I, OT, S, Z> Executor - for FridaInProcessExecutor<'a, 'b, 'c, FH, H, I, OT, S> +impl<'a, 'b, 'c, EM, H, I, OT, RT, S, Z> Executor + for FridaInProcessExecutor<'a, 'b, 'c, H, I, OT, RT, S> where - FH: FridaHelper<'b>, H: FnMut(&I) -> ExitKind, I: Input + HasTargetBytes, OT: ObserversTuple, + RT: FridaRuntimeTuple, { /// Instruct the target about the input and run #[inline] @@ -96,10 +94,9 @@ where } } -impl<'a, 'b, 'c, FH, H, I, OT, S> HasObservers - for FridaInProcessExecutor<'a, 'b, 'c, FH, H, I, OT, S> +impl<'a, 'b, 'c, H, I, OT, RT, S> HasObservers + for FridaInProcessExecutor<'a, 'b, 'c, H, I, OT, RT, S> where - FH: FridaHelper<'b>, H: FnMut(&I) -> ExitKind, I: Input + HasTargetBytes, OT: ObserversTuple, @@ -115,15 +112,19 @@ where } } -impl<'a, 'b, 'c, FH, H, I, OT, S> FridaInProcessExecutor<'a, 'b, 'c, FH, H, I, OT, S> +impl<'a, 'b, 'c, H, I, OT, S, RT> FridaInProcessExecutor<'a, 'b, 'c, H, I, OT, RT, S> where - FH: FridaHelper<'b>, H: FnMut(&I) -> ExitKind, I: Input + HasTargetBytes, OT: ObserversTuple, + RT: FridaRuntimeTuple, { /// Creates a new [`FridaInProcessExecutor`] - pub fn new(gum: &'a Gum, base: InProcessExecutor<'a, H, I, OT, S>, helper: &'c mut FH) -> Self { + pub fn new( + gum: &'a Gum, + base: InProcessExecutor<'a, H, I, OT, S>, + helper: &'c mut FridaInstrumentationHelper<'b, RT>, + ) -> Self { let mut stalker = Stalker::new(gum); // Include the current module (the fuzzer) in stalked ranges. We clone the ranges so that // we don't add it to the INSTRUMENTED ranges. diff --git a/libafl_frida/src/helper.rs b/libafl_frida/src/helper.rs index 172c91708b..e44ccf731f 100644 --- a/libafl_frida/src/helper.rs +++ b/libafl_frida/src/helper.rs @@ -1,12 +1,11 @@ use libafl::{ - bolts::AsSlice, + bolts::tuples::MatchFirstType, inputs::{HasTargetBytes, Input}, Error, }; -#[cfg(unix)] use libafl_targets::drcov::DrCovBasicBlock; -#[cfg(feature = "cmplog")] +#[cfg(all(feature = "cmplog", target_arch = "aarch64"))] use crate::cmplog_rt::CmpLogRuntime; #[cfg(windows)] use crate::FridaOptions; @@ -15,19 +14,15 @@ use crate::{asan::asan_rt::AsanRuntime, FridaOptions}; use crate::{coverage_rt::CoverageRuntime, drcov_rt::DrCovRuntime}; #[cfg(target_arch = "aarch64")] use capstone::{ - arch::{self, arm64::Arm64OperandType, ArchOperand::Arm64Operand, BuildsCapstone}, - Capstone, Insn, + arch::{self, BuildsCapstone}, + Capstone, }; #[cfg(all(target_arch = "x86_64", unix))] use capstone::{ arch::{self, BuildsCapstone}, - Capstone, RegId, + Capstone, }; use core::fmt::{self, Debug, Formatter}; -#[cfg(target_arch = "aarch64")] -use frida_gum::instruction_writer::Aarch64Register; -#[cfg(all(target_arch = "x86_64", unix))] -use frida_gum::instruction_writer::X86Register; #[cfg(unix)] use frida_gum::CpuContext; @@ -36,8 +31,6 @@ use frida_gum::instruction_writer::InstructionWriter; use frida_gum::{stalker::Transformer, Gum, Module, ModuleDetails, ModuleMap, PageProtection}; #[cfg(unix)] use nix::sys::mman::{mmap, MapFlags, ProtFlags}; -#[cfg(target_arch = "aarch64")] -use num_traits::cast::FromPrimitive; use rangemap::RangeMap; #[cfg(any(target_vendor = "apple"))] @@ -45,134 +38,106 @@ const ANONYMOUS_FLAG: MapFlags = MapFlags::MAP_ANON; #[cfg(not(any(target_vendor = "apple", target_os = "windows")))] const ANONYMOUS_FLAG: MapFlags = MapFlags::MAP_ANONYMOUS; -/// An helper that feeds `FridaInProcessExecutor` with user-supplied instrumentation -pub trait FridaHelper<'a>: Debug { - /// Access to the stalker `Transformer` - fn transformer(&self) -> &Transformer<'a>; +/// The Runtime trait +pub trait FridaRuntime: 'static + Debug { + /// Initialization + fn init( + &mut self, + gum: &Gum, + ranges: &RangeMap, + modules_to_instrument: &[&str], + ); - /// Register a new thread with this `FridaHelper` - #[cfg(unix)] - fn register_thread(&mut self); - - /// Called prior to execution of an input + /// Method called before execution fn pre_exec(&mut self, input: &I) -> Result<(), Error>; - /// Called after execution of an input + /// Method called after execution fn post_exec(&mut self, input: &I) -> Result<(), Error>; +} - /// Returns `true` if stalker is enabled - fn stalker_enabled(&self) -> bool; +/// The tuple for Frida Runtime +pub trait FridaRuntimeTuple: MatchFirstType + Debug { + /// Initialization + fn init_all( + &mut self, + gum: &Gum, + ranges: &RangeMap, + modules_to_instrument: &[&str], + ); - /// pointer to the frida coverage map - fn map_ptr_mut(&mut self) -> *mut u8; + /// Method called before execution + fn pre_exec_all(&mut self, input: &I) -> Result<(), Error>; - /// Returns the mapped ranges of the target - fn ranges(&self) -> &RangeMap; + /// Method called after execution + fn post_exec_all(&mut self, input: &I) -> Result<(), Error>; +} - /// Returns the mapped ranges of the target, mutable - fn ranges_mut(&mut self) -> &mut RangeMap; +impl FridaRuntimeTuple for () { + fn init_all( + &mut self, + _gum: &Gum, + _ranges: &RangeMap, + _modules_to_instrument: &[&str], + ) { + } + fn pre_exec_all(&mut self, _input: &I) -> Result<(), Error> { + Ok(()) + } + fn post_exec_all(&mut self, _input: &I) -> Result<(), Error> { + Ok(()) + } +} + +impl FridaRuntimeTuple for (Head, Tail) +where + Head: FridaRuntime, + Tail: FridaRuntimeTuple, +{ + fn init_all( + &mut self, + gum: &Gum, + ranges: &RangeMap, + modules_to_instrument: &[&str], + ) { + self.0.init(gum, ranges, modules_to_instrument); + self.1.init_all(gum, ranges, modules_to_instrument); + } + + fn pre_exec_all(&mut self, input: &I) -> Result<(), Error> { + self.0.pre_exec(input)?; + self.1.pre_exec_all(input) + } + + fn post_exec_all(&mut self, input: &I) -> Result<(), Error> { + self.0.post_exec(input)?; + self.1.post_exec_all(input) + } } /// An helper that feeds `FridaInProcessExecutor` with edge-coverage instrumentation -pub struct FridaInstrumentationHelper<'a> { - coverage_rt: CoverageRuntime, +pub struct FridaInstrumentationHelper<'a, RT> { /// Transformer that has to be passed to FridaInProcessExecutor transformer: Option>, #[cfg(unix)] capstone: Capstone, - #[cfg(unix)] - asan_runtime: AsanRuntime, - #[cfg(feature = "cmplog")] - cmplog_runtime: CmpLogRuntime, - drcov_runtime: DrCovRuntime, ranges: RangeMap, module_map: ModuleMap, options: &'a FridaOptions, + runtimes: RT, } -impl Debug for FridaInstrumentationHelper<'_> { +impl Debug for FridaInstrumentationHelper<'_, RT> { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { let mut dbg_me = f.debug_struct("FridaInstrumentationHelper"); dbg_me - .field("coverage_rt", &self.coverage_rt) - .field("drcov_runtime", &self.drcov_runtime) + .field("capstone", &self.capstone) .field("ranges", &self.ranges) .field("module_map", &"") .field("options", &self.options); - #[cfg(feature = "cmplog")] - dbg_me.field("cmplog_runtime", &self.cmplog_runtime); - #[cfg(unix)] - dbg_me.field("capstone", &self.capstone); - #[cfg(unix)] - dbg_me.field("asan_runtime", &self.asan_runtime); dbg_me.finish() } } -impl<'a> FridaHelper<'a> for FridaInstrumentationHelper<'a> { - fn transformer(&self) -> &Transformer<'a> { - self.transformer.as_ref().unwrap() - } - - /// Register the current thread with the [`FridaInstrumentationHelper`] - #[cfg(unix)] - fn register_thread(&mut self) { - self.asan_runtime.register_thread(); - } - - #[cfg(not(unix))] - fn pre_exec(&mut self, _input: &I) -> Result<(), Error> { - Ok(()) - } - - #[cfg(unix)] - fn pre_exec(&mut self, input: &I) -> Result<(), Error> { - let target_bytes = input.target_bytes(); - let slice = target_bytes.as_slice(); - //println!("target_bytes: {:#x}: {:02x?}", slice.as_ptr() as usize, slice); - if self.options.asan_enabled() { - self.asan_runtime - .unpoison(slice.as_ptr() as usize, slice.len()); - } - Ok(()) - } - - fn post_exec(&mut self, input: &I) -> Result<(), Error> { - if self.options().enable_drcov { - self.drcov_runtime.post_exec(input)?; - } - #[cfg(unix)] - if self.options.asan_enabled() { - if self.options.asan_detect_leaks() { - self.asan_runtime.check_for_leaks(); - } - - let target_bytes = input.target_bytes(); - let slice = target_bytes.as_slice(); - self.asan_runtime - .poison(slice.as_ptr() as usize, slice.len()); - self.asan_runtime.reset_allocations(); - } - Ok(()) - } - - fn stalker_enabled(&self) -> bool { - self.options.stalker_enabled() - } - - fn map_ptr_mut(&mut self) -> *mut u8 { - self.coverage_rt.map_ptr_mut() - } - - fn ranges(&self) -> &RangeMap { - &self.ranges - } - - fn ranges_mut(&mut self) -> &mut RangeMap { - &mut self.ranges - } -} - /// Helper function to get the size of a module's CODE section from frida #[must_use] pub fn get_module_size(module_name: &str) -> usize { @@ -197,7 +162,10 @@ fn pc(context: &CpuContext) -> usize { } /// The implementation of the [`FridaInstrumentationHelper`] -impl<'a> FridaInstrumentationHelper<'a> { +impl<'a, RT> FridaInstrumentationHelper<'a, RT> +where + RT: FridaRuntimeTuple, +{ /// Constructor function to create a new [`FridaInstrumentationHelper`], given a `module_name`. #[allow(clippy::too_many_lines)] #[must_use] @@ -206,6 +174,7 @@ impl<'a> FridaInstrumentationHelper<'a> { options: &'a FridaOptions, _harness_module_name: &str, modules_to_instrument: &'a [&str], + runtimes: RT, ) -> Self { // workaround frida's frida-gum-allocate-near bug: #[cfg(unix)] @@ -233,7 +202,6 @@ impl<'a> FridaInstrumentationHelper<'a> { } let mut helper = Self { - coverage_rt: CoverageRuntime::new(), transformer: None, #[cfg(target_arch = "aarch64")] capstone: Capstone::new() @@ -249,14 +217,10 @@ impl<'a> FridaInstrumentationHelper<'a> { .detail(true) .build() .expect("Failed to create Capstone object"), - #[cfg(not(windows))] - asan_runtime: AsanRuntime::new(options.clone()), - #[cfg(feature = "cmplog")] - cmplog_runtime: CmpLogRuntime::new(), - drcov_runtime: DrCovRuntime::new(), ranges: RangeMap::new(), module_map: ModuleMap::new_from_names(modules_to_instrument), options, + runtimes, }; if helper.options().stalker_enabled() { @@ -279,19 +243,11 @@ impl<'a> FridaInstrumentationHelper<'a> { } } - if helper.options().drcov_enabled() { - std::fs::create_dir_all("./coverage") - .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 { let instr = instruction.instr(); + let instr_size = instr.bytes().len(); let address = instr.address(); //println!("block @ {:x} transformed to {:x}", address, output.writer().pc()); @@ -306,44 +262,44 @@ impl<'a> FridaInstrumentationHelper<'a> { if first { first = false; //println!("block @ {:x} transformed to {:x}", address, output.writer().pc()); - if helper.options().coverage_enabled() { - helper.coverage_rt.emit_coverage_mapping(address, &output); + if let Some(rt) = helper.runtime_mut::() { + rt.emit_coverage_mapping(address, &output); } - #[cfg(unix)] - if helper.options().drcov_enabled() { + + if let Some(rt) = helper.runtime_mut::() { instruction.put_callout(|context| { - let real_address = - helper.asan_runtime.real_address_for_stalked(pc(&context)); + let real_address = rt.real_address_for_stalked(pc(&context)); //let (range, (id, name)) = helper.ranges.get_key_value(&real_address).unwrap(); //println!("{}:0x{:016x}", name, real_address - range.start); - helper - .drcov_runtime - .drcov_basic_blocks - .push(DrCovBasicBlock::new(real_address, real_address + 4)); + rt.drcov_basic_blocks.push(DrCovBasicBlock::new( + real_address, + real_address + instr_size, + )); }); } } - if helper.options().asan_enabled() { - #[cfg(all(target_arch = "x86_64", unix))] - if let Ok((segment, width, basereg, indexreg, scale, disp)) = helper - .asan_runtime - .asan_is_interesting_instruction(&helper.capstone, address, instr) - { - helper.asan_runtime.emit_shadow_check( + let res = if let Some(rt) = helper.runtime::() { + rt.asan_is_interesting_instruction(&helper.capstone, address, instr) + } else { + None + }; + + #[cfg(all(target_arch = "x86_64", unix))] + if let Some((segment, width, basereg, indexreg, scale, disp)) = res { + if let Some(rt) = helper.runtime_mut::() { + rt.emit_shadow_check( address, &output, segment, width, basereg, indexreg, scale, disp, ); } - #[cfg(target_arch = "aarch64")] - if let Ok((basereg, indexreg, displacement, width, shift, extender)) = - helper.asan_runtime.asan_is_interesting_instruction( - &helper.capstone, - address, - instr, - ) - { - helper.asan_runtime.emit_shadow_check( + } + + #[cfg(target_arch = "aarch64")] + if let Some((basereg, indexreg, displacement, width, shift, extender)) = res + { + if let Some(rt) = helper.runtime_mut::() { + rt.emit_shadow_check( address, &output, basereg, @@ -355,17 +311,14 @@ impl<'a> FridaInstrumentationHelper<'a> { ); } } - if helper.options().cmplog_enabled() { - #[cfg(not(target_arch = "aarch64"))] - todo!("Implement cmplog for non-aarch64 targets"); - #[cfg(all(feature = "cmplog", target_arch = "aarch64"))] - // check if this instruction is a compare instruction and if so save the registers values - if let Ok((op1, op2, special_case)) = helper - .cmplog_runtime + + #[cfg(all(feature = "cmplog", target_arch = "aarch64"))] + if let Some(rt) = helper.runtime::() { + if let Ok((op1, op2, special_case)) = rt .cmplog_is_interesting_instruction(&helper.capstone, address, instr) { //emit code that saves the relevant data in runtime(passes it to x0, x1) - helper.cmplog_runtime.emit_comparison_handling( + rt.emit_comparison_handling( address, &output, op1, @@ -376,9 +329,17 @@ impl<'a> FridaInstrumentationHelper<'a> { } #[cfg(unix)] - if helper.options().asan_enabled() || helper.options().drcov_enabled() { - helper.asan_runtime.add_stalked_address( - output.writer().pc() as usize - 4, + if let Some(rt) = helper.runtime_mut::() { + rt.add_stalked_address( + output.writer().pc() as usize - instr_size, + address as usize, + ); + } + + #[cfg(unix)] + if let Some(rt) = helper.runtime_mut::() { + rt.add_stalked_address( + output.writer().pc() as usize - instr_size, address as usize, ); } @@ -387,146 +348,81 @@ impl<'a> FridaInstrumentationHelper<'a> { } }); helper.transformer = Some(transformer); - - #[cfg(unix)] - if helper.options().asan_enabled() || helper.options().drcov_enabled() { - helper.asan_runtime.init(gum, modules_to_instrument); - } - - helper.drcov_runtime.init(&helper.ranges); - #[cfg(feature = "cmplog")] - if helper.options.cmplog_enabled() { - helper.cmplog_runtime.init(); - } + helper + .runtimes + .init_all(gum, &helper.ranges, modules_to_instrument); } helper } + /// Return the runtime + pub fn runtime(&self) -> Option<&R> + where + R: FridaRuntime, + { + self.runtimes.match_first_type::() + } + + /// Return the mutable runtime + pub fn runtime_mut(&mut self) -> Option<&mut R> + where + R: FridaRuntime, + { + self.runtimes.match_first_type_mut::() + } + + /// Returns ref to the Transformer + pub fn transformer(&self) -> &Transformer<'a> { + self.transformer.as_ref().unwrap() + } + + /// Initializa all + pub fn init( + &mut self, + gum: &'a Gum, + ranges: &RangeMap, + modules_to_instrument: &'a [&str], + ) { + self.runtimes.init_all(gum, ranges, modules_to_instrument); + } + + /// Method called before execution + pub fn pre_exec(&mut self, input: &I) -> Result<(), Error> { + self.runtimes.pre_exec_all(input) + } + + /// Method called after execution + pub fn post_exec(&mut self, input: &I) -> Result<(), Error> { + self.runtimes.post_exec_all(input) + } + + /// If stalker is enabled + pub fn stalker_enabled(&self) -> bool { + self.options.stalker_enabled() + } + + /// Pointer to coverage map + pub fn map_ptr_mut(&mut self) -> Option<*mut u8> { + if let Some(rt) = self.runtime_mut::() { + Some(rt.map_ptr_mut()) + } else { + None + } + } + + /// Ranges + pub fn ranges(&self) -> &RangeMap { + &self.ranges + } + + /// Mutable ranges + pub fn ranges_mut(&mut self) -> &mut RangeMap { + &mut self.ranges + } + + /// Return the ref to options #[inline] - fn options(&self) -> &FridaOptions { + pub fn options(&self) -> &FridaOptions { self.options } - - /// Determine the width of the specified instruction - #[cfg(target_arch = "aarch64")] - #[inline] - pub fn instruction_width(instr: &Insn, operands: &Vec) -> u32 { - use capstone::arch::arm64::Arm64Insn as I; - use capstone::arch::arm64::Arm64Reg as R; - use capstone::arch::arm64::Arm64Vas as V; - - let num_registers = match instr.id().0.into() { - I::ARM64_INS_STP - | I::ARM64_INS_STXP - | I::ARM64_INS_STNP - | I::ARM64_INS_STLXP - | I::ARM64_INS_LDP - | I::ARM64_INS_LDXP - | I::ARM64_INS_LDNP => 2, - _ => 1, - }; - - let mnemonic = instr.mnemonic().unwrap(); - match mnemonic.as_bytes().last().unwrap() { - b'b' => return 1, - b'h' => return 2, - b'w' => return 4 * num_registers, - _ => (), - } - - if let Arm64Operand(operand) = operands.first().unwrap() { - if operand.vas != V::ARM64_VAS_INVALID { - let count_byte: u32 = if mnemonic.starts_with("st") || mnemonic.starts_with("ld") { - mnemonic.chars().nth(2).unwrap().to_digit(10).unwrap() - } else { - 1 - }; - - return match operand.vas { - V::ARM64_VAS_1B => 1 * count_byte, - V::ARM64_VAS_1H => 2 * count_byte, - V::ARM64_VAS_4B | V::ARM64_VAS_1S | V::ARM64_VAS_1D | V::ARM64_VAS_2H => { - 4 * count_byte - } - V::ARM64_VAS_8B - | V::ARM64_VAS_4H - | V::ARM64_VAS_2S - | V::ARM64_VAS_2D - | V::ARM64_VAS_1Q => 8 * count_byte, - V::ARM64_VAS_8H | V::ARM64_VAS_4S | V::ARM64_VAS_16B => 16 * count_byte, - V::ARM64_VAS_INVALID => { - panic!("should not be reached"); - } - }; - } else if let Arm64OperandType::Reg(operand) = operand.op_type { - match operand.0 as u32 { - R::ARM64_REG_W0..=R::ARM64_REG_W30 - | R::ARM64_REG_WZR - | R::ARM64_REG_WSP - | R::ARM64_REG_S0..=R::ARM64_REG_S31 => return 4 * num_registers, - R::ARM64_REG_D0..=R::ARM64_REG_D31 => return 8 * num_registers, - R::ARM64_REG_Q0..=R::ARM64_REG_Q31 => return 16, - _ => (), - } - }; - }; - - 8 * num_registers - } - - /// Convert from a capstone register id to a frida InstructionWriter register index - #[cfg(target_arch = "aarch64")] - #[inline] - pub fn writer_register(reg: capstone::RegId) -> Aarch64Register { - let regint: u16 = reg.0; - Aarch64Register::from_u32(regint as u32).unwrap() - } - - /// The writer registers - /// frida registers: - /// capstone registers: - #[cfg(all(target_arch = "x86_64", unix))] - #[must_use] - #[inline] - #[allow(clippy::unused_self)] - pub fn writer_register(reg: RegId) -> X86Register { - let regint: u16 = reg.0; - match regint { - 19 => X86Register::Eax, - 22 => X86Register::Ecx, - 24 => X86Register::Edx, - 21 => X86Register::Ebx, - 30 => X86Register::Esp, - 20 => X86Register::Ebp, - 29 => X86Register::Esi, - 23 => X86Register::Edi, - 226 => X86Register::R8d, - 227 => X86Register::R9d, - 228 => X86Register::R10d, - 229 => X86Register::R11d, - 230 => X86Register::R12d, - 231 => X86Register::R13d, - 232 => X86Register::R14d, - 233 => X86Register::R15d, - 26 => X86Register::Eip, - 35 => X86Register::Rax, - 38 => X86Register::Rcx, - 40 => X86Register::Rdx, - 37 => X86Register::Rbx, - 44 => X86Register::Rsp, - 36 => X86Register::Rbp, - 43 => X86Register::Rsi, - 39 => X86Register::Rdi, - 106 => X86Register::R8, - 107 => X86Register::R9, - 108 => X86Register::R10, - 109 => X86Register::R11, - 110 => X86Register::R12, - 111 => X86Register::R13, - 112 => X86Register::R14, - 113 => X86Register::R15, - 41 => X86Register::Rip, - _ => X86Register::None, // Ignore Xax..Xip - } - } } diff --git a/libafl_frida/src/lib.rs b/libafl_frida/src/lib.rs index 7ee7e08a28..b774d78c2a 100644 --- a/libafl_frida/src/lib.rs +++ b/libafl_frida/src/lib.rs @@ -80,6 +80,9 @@ pub mod drcov_rt; /// The frida executor pub mod executor; +/// Utilities +pub mod utils; + // for parsing asan and cmplog cores use libafl::bolts::os::{CoreId, Cores}; diff --git a/libafl_frida/src/utils.rs b/libafl_frida/src/utils.rs new file mode 100644 index 0000000000..e805ee41cc --- /dev/null +++ b/libafl_frida/src/utils.rs @@ -0,0 +1,136 @@ +#[cfg(target_arch = "aarch64")] +use frida_gum::instruction_writer::Aarch64Register; +#[cfg(target_arch = "x86_64")] +use frida_gum::instruction_writer::X86Register; + +#[cfg(target_arch = "aarch64")] +use capstone::{ + arch::{self, arm64::Arm64OperandType, ArchOperand::Arm64Operand}, + Insn, +}; + +#[cfg(target_arch = "aarch64")] +use num_traits::cast::FromPrimitive; + +/// Determine the width of the specified instruction +#[cfg(target_arch = "aarch64")] +#[inline] +pub fn instruction_width(instr: &Insn, operands: &Vec) -> u32 { + use capstone::arch::arm64::Arm64Insn as I; + use capstone::arch::arm64::Arm64Reg as R; + use capstone::arch::arm64::Arm64Vas as V; + + let num_registers = match instr.id().0.into() { + I::ARM64_INS_STP + | I::ARM64_INS_STXP + | I::ARM64_INS_STNP + | I::ARM64_INS_STLXP + | I::ARM64_INS_LDP + | I::ARM64_INS_LDXP + | I::ARM64_INS_LDNP => 2, + _ => 1, + }; + + let mnemonic = instr.mnemonic().unwrap(); + match mnemonic.as_bytes().last().unwrap() { + b'b' => return 1, + b'h' => return 2, + b'w' => return 4 * num_registers, + _ => (), + } + + if let Arm64Operand(operand) = operands.first().unwrap() { + if operand.vas != V::ARM64_VAS_INVALID { + let count_byte: u32 = if mnemonic.starts_with("st") || mnemonic.starts_with("ld") { + mnemonic.chars().nth(2).unwrap().to_digit(10).unwrap() + } else { + 1 + }; + + return match operand.vas { + V::ARM64_VAS_1B => 1 * count_byte, + V::ARM64_VAS_1H => 2 * count_byte, + V::ARM64_VAS_4B | V::ARM64_VAS_1S | V::ARM64_VAS_1D | V::ARM64_VAS_2H => { + 4 * count_byte + } + V::ARM64_VAS_8B + | V::ARM64_VAS_4H + | V::ARM64_VAS_2S + | V::ARM64_VAS_2D + | V::ARM64_VAS_1Q => 8 * count_byte, + V::ARM64_VAS_8H | V::ARM64_VAS_4S | V::ARM64_VAS_16B => 16 * count_byte, + V::ARM64_VAS_INVALID => { + panic!("should not be reached"); + } + }; + } else if let Arm64OperandType::Reg(operand) = operand.op_type { + match operand.0 as u32 { + R::ARM64_REG_W0..=R::ARM64_REG_W30 + | R::ARM64_REG_WZR + | R::ARM64_REG_WSP + | R::ARM64_REG_S0..=R::ARM64_REG_S31 => return 4 * num_registers, + R::ARM64_REG_D0..=R::ARM64_REG_D31 => return 8 * num_registers, + R::ARM64_REG_Q0..=R::ARM64_REG_Q31 => return 16, + _ => (), + } + }; + }; + + 8 * num_registers +} + +/// Convert from a capstone register id to a frida InstructionWriter register index +#[cfg(target_arch = "aarch64")] +#[inline] +pub fn writer_register(reg: capstone::RegId) -> Aarch64Register { + let regint: u16 = reg.0; + Aarch64Register::from_u32(regint as u32).unwrap() +} + +/// The writer registers +/// frida registers: +/// capstone registers: +#[cfg(all(target_arch = "x86_64", unix))] +#[must_use] +#[inline] +#[allow(clippy::unused_self)] +pub fn writer_register(reg: capstone::RegId) -> X86Register { + let regint: u16 = reg.0; + match regint { + 19 => X86Register::Eax, + 22 => X86Register::Ecx, + 24 => X86Register::Edx, + 21 => X86Register::Ebx, + 30 => X86Register::Esp, + 20 => X86Register::Ebp, + 29 => X86Register::Esi, + 23 => X86Register::Edi, + 226 => X86Register::R8d, + 227 => X86Register::R9d, + 228 => X86Register::R10d, + 229 => X86Register::R11d, + 230 => X86Register::R12d, + 231 => X86Register::R13d, + 232 => X86Register::R14d, + 233 => X86Register::R15d, + 26 => X86Register::Eip, + 35 => X86Register::Rax, + 38 => X86Register::Rcx, + 40 => X86Register::Rdx, + 37 => X86Register::Rbx, + 44 => X86Register::Rsp, + 36 => X86Register::Rbp, + 43 => X86Register::Rsi, + 39 => X86Register::Rdi, + 106 => X86Register::R8, + 107 => X86Register::R9, + 108 => X86Register::R10, + 109 => X86Register::R11, + 110 => X86Register::R12, + 111 => X86Register::R13, + 112 => X86Register::R14, + 113 => X86Register::R15, + 41 => X86Register::Rip, + _ => X86Register::None, // Ignore Xax..Xip + } +}