From 9cd0d2228cf382df7c6d2134b69b2483f9fc3468 Mon Sep 17 00:00:00 2001 From: tokatoka Date: Fri, 24 Dec 2021 15:45:08 +0900 Subject: [PATCH 1/3] drcov runtime --- libafl_frida/src/drcov_rt.rs | 38 ++++++++++++++++++++++++++++++++++++ libafl_frida/src/helper.rs | 29 ++++++++++----------------- libafl_frida/src/lib.rs | 2 ++ 3 files changed, 50 insertions(+), 19 deletions(-) create mode 100644 libafl_frida/src/drcov_rt.rs diff --git a/libafl_frida/src/drcov_rt.rs b/libafl_frida/src/drcov_rt.rs new file mode 100644 index 0000000000..c08c18d36c --- /dev/null +++ b/libafl_frida/src/drcov_rt.rs @@ -0,0 +1,38 @@ +use std::hash::Hasher; +use ahash::AHasher; +use libafl::inputs::{HasTargetBytes, Input}; +use libafl_targets::drcov::{DrCovBasicBlock, DrCovWriter}; +use libafl::Error; +use rangemap::RangeMap; + + +pub struct DrCovRuntime { + pub drcov_basic_blocks: Vec, + ranges: RangeMap, +} + +impl DrCovRuntime { + pub fn new() -> Self { + Self { + drcov_basic_blocks: vec![], + ranges: RangeMap::new(), + } + } + + pub fn pre_exec(&mut self, _input: &I) -> Result<(), Error> { + Ok(()) + } + + pub 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()); + + let filename = format!("./coverage/{:016x}.drcov", hasher.finish(),); + DrCovWriter::new(&self.ranges).write(&filename, &self.drcov_basic_blocks)?; + self.drcov_basic_blocks.clear(); + + Ok(()) + } + + +} \ No newline at end of file diff --git a/libafl_frida/src/helper.rs b/libafl_frida/src/helper.rs index 83c222695d..1d4f583a85 100644 --- a/libafl_frida/src/helper.rs +++ b/libafl_frida/src/helper.rs @@ -1,9 +1,6 @@ -use ahash::AHasher; -use std::hash::Hasher; - use libafl::inputs::{HasTargetBytes, Input}; use libafl::Error; -use libafl_targets::drcov::{DrCovBasicBlock, DrCovWriter}; +use libafl_targets::drcov::{DrCovBasicBlock}; #[cfg(target_arch = "aarch64")] use capstone::{ @@ -46,6 +43,8 @@ use crate::{asan::asan_rt::AsanRuntime, FridaOptions}; #[cfg(windows)] use crate::FridaOptions; +use crate::drcov_rt::DrCovRuntime; + use crate::coverage_rt::CoverageRuntime; #[cfg(feature = "cmplog")] @@ -66,7 +65,7 @@ pub trait FridaHelper<'a> { fn register_thread(&mut self); /// Called prior to execution of an input - fn pre_exec(&mut self, input: &I); + fn pre_exec(&mut self, input: &I) -> Result<(), Error>; /// Called after execution of an input fn post_exec(&mut self, input: &I) -> Result<(), Error>; @@ -93,10 +92,10 @@ pub struct FridaInstrumentationHelper<'a> { asan_runtime: AsanRuntime, #[cfg(feature = "cmplog")] cmplog_runtime: CmpLogRuntime, + drcov_runtime: DrCovRuntime, ranges: RangeMap, module_map: ModuleMap, options: &'a FridaOptions, - drcov_basic_blocks: Vec, } impl<'a> FridaHelper<'a> for FridaInstrumentationHelper<'a> { @@ -114,7 +113,7 @@ impl<'a> FridaHelper<'a> for FridaInstrumentationHelper<'a> { fn pre_exec(&mut self, _input: &I) {} #[cfg(unix)] - fn pre_exec(&mut self, input: &I) { + 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); @@ -122,18 +121,11 @@ impl<'a> FridaHelper<'a> for FridaInstrumentationHelper<'a> { self.asan_runtime .unpoison(slice.as_ptr() as usize, slice.len()); } + Ok(()) } fn post_exec(&mut self, input: &I) -> Result<(), Error> { - if self.options.drcov_enabled() { - let mut hasher = AHasher::new_with_keys(0, 0); - hasher.write(input.target_bytes().as_slice()); - - let filename = format!("./coverage/{:016x}.drcov", hasher.finish(),); - DrCovWriter::new(&self.ranges).write(&filename, &self.drcov_basic_blocks)?; - self.drcov_basic_blocks.clear(); - } - + self.drcov_runtime.post_exec(input)?; #[cfg(unix)] if self.options.asan_enabled() { if self.options.asan_detect_leaks() { @@ -246,10 +238,10 @@ impl<'a> FridaInstrumentationHelper<'a> { 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, - drcov_basic_blocks: vec![], }; if helper.options().stalker_enabled() { @@ -309,8 +301,7 @@ impl<'a> FridaInstrumentationHelper<'a> { helper.asan_runtime.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_basic_blocks + helper.drcov_runtime.drcov_basic_blocks .push(DrCovBasicBlock::new(real_address, real_address + 4)); }); } diff --git a/libafl_frida/src/lib.rs b/libafl_frida/src/lib.rs index 4c1d8f805c..70915cbc95 100644 --- a/libafl_frida/src/lib.rs +++ b/libafl_frida/src/lib.rs @@ -19,6 +19,8 @@ pub mod cmplog_rt; /// The `LibAFL` firda helper pub mod helper; +pub mod drcov_rt; + /// The frida executor pub mod executor; From e6434d2ec2b8383a254e90b5cd86a26d467bc995 Mon Sep 17 00:00:00 2001 From: tokatoka Date: Fri, 24 Dec 2021 15:46:27 +0900 Subject: [PATCH 2/3] fmt --- libafl_frida/src/drcov_rt.rs | 9 +++------ libafl_frida/src/helper.rs | 8 +++++--- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/libafl_frida/src/drcov_rt.rs b/libafl_frida/src/drcov_rt.rs index c08c18d36c..d2ef88b2f3 100644 --- a/libafl_frida/src/drcov_rt.rs +++ b/libafl_frida/src/drcov_rt.rs @@ -1,10 +1,9 @@ -use std::hash::Hasher; use ahash::AHasher; use libafl::inputs::{HasTargetBytes, Input}; -use libafl_targets::drcov::{DrCovBasicBlock, DrCovWriter}; use libafl::Error; +use libafl_targets::drcov::{DrCovBasicBlock, DrCovWriter}; use rangemap::RangeMap; - +use std::hash::Hasher; pub struct DrCovRuntime { pub drcov_basic_blocks: Vec, @@ -33,6 +32,4 @@ impl DrCovRuntime { Ok(()) } - - -} \ No newline at end of file +} diff --git a/libafl_frida/src/helper.rs b/libafl_frida/src/helper.rs index 1d4f583a85..2b035adc6f 100644 --- a/libafl_frida/src/helper.rs +++ b/libafl_frida/src/helper.rs @@ -1,6 +1,6 @@ use libafl::inputs::{HasTargetBytes, Input}; use libafl::Error; -use libafl_targets::drcov::{DrCovBasicBlock}; +use libafl_targets::drcov::DrCovBasicBlock; #[cfg(target_arch = "aarch64")] use capstone::{ @@ -113,7 +113,7 @@ impl<'a> FridaHelper<'a> for FridaInstrumentationHelper<'a> { fn pre_exec(&mut self, _input: &I) {} #[cfg(unix)] - fn pre_exec(&mut self, input: &I) -> Result<(), Error>{ + 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); @@ -301,7 +301,9 @@ impl<'a> FridaInstrumentationHelper<'a> { helper.asan_runtime.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 + helper + .drcov_runtime + .drcov_basic_blocks .push(DrCovBasicBlock::new(real_address, real_address + 4)); }); } From 97c169fe63142593dc974f6683e693a5ad51a889 Mon Sep 17 00:00:00 2001 From: tokatoka Date: Fri, 24 Dec 2021 16:34:53 +0900 Subject: [PATCH 3/3] init ranges later --- libafl_frida/src/drcov_rt.rs | 6 ++++++ libafl_frida/src/helper.rs | 2 ++ 2 files changed, 8 insertions(+) diff --git a/libafl_frida/src/drcov_rt.rs b/libafl_frida/src/drcov_rt.rs index d2ef88b2f3..66bdf2271c 100644 --- a/libafl_frida/src/drcov_rt.rs +++ b/libafl_frida/src/drcov_rt.rs @@ -11,6 +11,7 @@ pub struct DrCovRuntime { } impl DrCovRuntime { + #[must_use] pub fn new() -> Self { Self { drcov_basic_blocks: vec![], @@ -18,6 +19,11 @@ impl DrCovRuntime { } } + pub fn init(&mut self, ranges: &RangeMap) { + self.ranges = ranges.clone(); + } + + #[allow(clippy::unused_self)] pub fn pre_exec(&mut self, _input: &I) -> Result<(), Error> { Ok(()) } diff --git a/libafl_frida/src/helper.rs b/libafl_frida/src/helper.rs index 2b035adc6f..bd45710d65 100644 --- a/libafl_frida/src/helper.rs +++ b/libafl_frida/src/helper.rs @@ -377,6 +377,8 @@ impl<'a> FridaInstrumentationHelper<'a> { 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();