From 2a79ee5b4fe5d73c03803c10beb336925e9a920f Mon Sep 17 00:00:00 2001 From: s1341 Date: Wed, 25 Dec 2024 14:42:54 +0200 Subject: [PATCH] frida: Deduplicate with IfElseRuntime (#2792) * frida: Deduplicate with IfElseRuntime * clippy' * get rid of cfg * fmt * documentation * fix lint * fix lint * debug: add tmate * debug: add tmate * frida_windows_gdiplus: move to mimalloc on windows * remove tmate --- .../binary_only/frida_libpng/src/fuzzer.rs | 559 ++++++------------ .../frida_windows_gdiplus/Cargo.toml | 2 +- .../frida_windows_gdiplus/src/fuzzer.rs | 7 - libafl_frida/src/asan/asan_rt.rs | 4 +- libafl_frida/src/helper.rs | 72 +++ 5 files changed, 243 insertions(+), 401 deletions(-) diff --git a/fuzzers/binary_only/frida_libpng/src/fuzzer.rs b/fuzzers/binary_only/frida_libpng/src/fuzzer.rs index 6f49f06a72..9a94d2459a 100644 --- a/fuzzers/binary_only/frida_libpng/src/fuzzer.rs +++ b/fuzzers/binary_only/frida_libpng/src/fuzzer.rs @@ -21,7 +21,7 @@ use libafl::{ }, observers::{CanTrack, HitcountsMapObserver, StdMapObserver, TimeObserver}, schedulers::{IndexesLenTimeMinimizerScheduler, QueueScheduler}, - stages::{ShadowTracingStage, StdMutationalStage}, + stages::{IfElseStage, ShadowTracingStage, StdMutationalStage}, state::{HasCorpus, StdState}, Error, HasMetadata, }; @@ -34,16 +34,15 @@ use libafl_bolts::{ tuples::{tuple_list, Merge}, AsSlice, }; -#[cfg(unix)] -use libafl_frida::asan::{ - asan_rt::AsanRuntime, - errors::{AsanErrorsFeedback, AsanErrorsObserver}, -}; use libafl_frida::{ + asan::{ + asan_rt::AsanRuntime, + errors::{AsanErrorsFeedback, AsanErrorsObserver}, + }, cmplog_rt::CmpLogRuntime, coverage_rt::{CoverageRuntime, MAP_SIZE}, executor::FridaInProcessExecutor, - helper::FridaInstrumentationHelper, + helper::{FridaInstrumentationHelper, IfElseRuntime}, }; use libafl_targets::cmplog::CmpLogObserver; use mimalloc::MiMalloc; @@ -74,6 +73,12 @@ unsafe fn fuzz(options: &FuzzerOptions) -> Result<(), Error> { let monitor = MultiMonitor::new(|s| println!("{s}")); let shmem_provider = StdShMemProvider::new()?; + let is_asan = |options: &FuzzerOptions, client_description: &ClientDescription| { + options.asan && options.asan_cores.contains(client_description.core_id()) + }; + let is_cmplog = |options: &FuzzerOptions, client_description: &ClientDescription| { + options.cmplog && options.cmplog_cores.contains(client_description.core_id()) + }; let mut run_client = |state: Option<_>, mgr: LlmpRestartingEventManager<_, _, _>, @@ -94,390 +99,162 @@ unsafe fn fuzz(options: &FuzzerOptions) -> Result<(), Error> { ExitKind::Ok }; - if options.asan && options.asan_cores.contains(client_description.core_id()) { - (|state: Option<_>, - mut mgr: LlmpRestartingEventManager<_, _, _>, - _client_description| { - let gum = Gum::obtain(); - - let coverage = CoverageRuntime::new(); - #[cfg(unix)] - let asan = AsanRuntime::new(options); - - #[cfg(unix)] - let mut frida_helper = - FridaInstrumentationHelper::new(&gum, options, tuple_list!(asan, coverage)); - #[cfg(windows)] - let mut frida_helper = - FridaInstrumentationHelper::new(&gum, &options, tuple_list!(coverage)); - - // Create an observation channel using the coverage map - let edges_observer = HitcountsMapObserver::new(StdMapObserver::from_mut_ptr( - "edges", - frida_helper.map_mut_ptr().unwrap(), - MAP_SIZE, - )) - .track_indices(); - - // Create an observation channel to keep track of the execution time - let time_observer = TimeObserver::new("time"); - #[cfg(unix)] - let asan_observer = AsanErrorsObserver::from_static_asan_errors(); - - // Feedback to rate the interestingness of an input - // This one is composed by two Feedbacks in OR - let mut feedback = feedback_or!( - // New maximization map feedback linked to the edges observer and the feedback state - MaxMapFeedback::new(&edges_observer), - // Time feedback, this one does not need a feedback state - TimeFeedback::new(&time_observer) - ); - - // Feedbacks to recognize an input as solution - #[cfg(unix)] - let mut objective = feedback_or_fast!( - CrashFeedback::new(), - TimeoutFeedback::new(), - // true enables the AsanErrorFeedback - feedback_and_fast!( - ConstFeedback::from(true), - AsanErrorsFeedback::new(&asan_observer) - ) - ); - #[cfg(windows)] - let mut objective = feedback_or_fast!(CrashFeedback::new(), TimeoutFeedback::new()); - - // If not restarting, create a State from scratch - let mut state = state.unwrap_or_else(|| { - StdState::new( - // RNG - StdRand::new(), - // Corpus that will be evolved, we keep it in memory for performance - CachedOnDiskCorpus::no_meta(PathBuf::from("./corpus_discovered"), 64) - .unwrap(), - // Corpus in which we store solutions (crashes in this example), - // on disk so the user can get them after stopping the fuzzer - OnDiskCorpus::new(options.output.clone()).unwrap(), - &mut feedback, - &mut objective, - ) - .unwrap() - }); - - println!("We're a client, let's fuzz :)"); - - // Create a PNG dictionary if not existing - if state.metadata_map().get::().is_none() { - state.add_metadata(Tokens::from([ - vec![137, 80, 78, 71, 13, 10, 26, 10], // PNG header - b"IHDR".to_vec(), - b"IDAT".to_vec(), - b"PLTE".to_vec(), - b"IEND".to_vec(), - ])); - } - - // Setup a basic mutator with a mutational stage - let mutator = StdScheduledMutator::new(havoc_mutations().merge(tokens_mutations())); - - // A minimization+queue policy to get testcasess from the corpus - let scheduler = - IndexesLenTimeMinimizerScheduler::new(&edges_observer, QueueScheduler::new()); - - // A fuzzer with feedbacks and a corpus scheduler - let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); - - #[cfg(unix)] - let observers = tuple_list!(edges_observer, time_observer, asan_observer); - #[cfg(windows)] - let observers = tuple_list!(edges_observer, time_observer); - - // Create the executor for an in-process function with just one observer for edge coverage - let mut executor = FridaInProcessExecutor::new( - &gum, - InProcessExecutor::new( - &mut frida_harness, - observers, - &mut fuzzer, - &mut state, - &mut mgr, - )?, - &mut frida_helper, - ); - - // In case the corpus is empty (on first run), reset - if state.must_load_initial_inputs() { - state - .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &options.input) - .unwrap_or_else(|_| { - panic!("Failed to load initial corpus at {:?}", &options.input) - }); - println!("We imported {} inputs from disk.", state.corpus().count()); - } - - let mut stages = tuple_list!(StdMutationalStage::new(mutator)); - - fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?; - - Ok(()) - })(state, mgr, client_description) - } else if options.cmplog && options.cmplog_cores.contains(client_description.core_id()) { - (|state: Option<_>, - mut mgr: LlmpRestartingEventManager<_, _, _>, - _client_description| { - let gum = Gum::obtain(); - - let coverage = CoverageRuntime::new(); - let cmplog = CmpLogRuntime::new(); - println!("cmplog runtime created"); - - let mut frida_helper = - FridaInstrumentationHelper::new(&gum, options, tuple_list!(coverage, cmplog)); - - // Create an observation channel using the coverage map - let edges_observer = HitcountsMapObserver::new(StdMapObserver::from_mut_ptr( - "edges", - frida_helper.map_mut_ptr().unwrap(), - MAP_SIZE, - )) - .track_indices(); - - // Create an observation channel to keep track of the execution time - let time_observer = TimeObserver::new("time"); - #[cfg(unix)] - let asan_observer = AsanErrorsObserver::from_static_asan_errors(); - - // Feedback to rate the interestingness of an input - // This one is composed by two Feedbacks in OR - let mut feedback = feedback_or!( - // New maximization map feedback linked to the edges observer and the feedback state - MaxMapFeedback::new(&edges_observer), - // Time feedback, this one does not need a feedback state - TimeFeedback::new(&time_observer) - ); - - #[cfg(unix)] - let mut objective = feedback_or_fast!( - CrashFeedback::new(), - TimeoutFeedback::new(), - feedback_and_fast!( - ConstFeedback::from(false), - AsanErrorsFeedback::new(&asan_observer) - ) - ); - #[cfg(windows)] - let mut objective = feedback_or_fast!(CrashFeedback::new(), TimeoutFeedback::new()); - - // If not restarting, create a State from scratch - let mut state = state.unwrap_or_else(|| { - StdState::new( - // RNG - StdRand::new(), - // Corpus that will be evolved, we keep it in memory for performance - CachedOnDiskCorpus::no_meta(PathBuf::from("./corpus_discovered"), 64) - .unwrap(), - // Corpus in which we store solutions (crashes in this example), - // on disk so the user can get them after stopping the fuzzer - OnDiskCorpus::new(options.output.clone()).unwrap(), - &mut feedback, - &mut objective, - ) - .unwrap() - }); - - println!("We're a client, let's fuzz :)"); - - // Create a PNG dictionary if not existing - if state.metadata_map().get::().is_none() { - state.add_metadata(Tokens::from([ - vec![137, 80, 78, 71, 13, 10, 26, 10], // PNG header - b"IHDR".to_vec(), - b"IDAT".to_vec(), - b"PLTE".to_vec(), - b"IEND".to_vec(), - ])); - } - - // Setup a basic mutator with a mutational stage - let mutator = StdScheduledMutator::new(havoc_mutations().merge(tokens_mutations())); - - // A minimization+queue policy to get testcasess from the corpus - let scheduler = - IndexesLenTimeMinimizerScheduler::new(&edges_observer, QueueScheduler::new()); - - // A fuzzer with feedbacks and a corpus scheduler - let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); - - #[cfg(unix)] - let observers = tuple_list!(edges_observer, time_observer, asan_observer); - #[cfg(windows)] - let observers = tuple_list!(edges_observer, time_observer); - - // Create the executor for an in-process function with just one observer for edge coverage - let mut executor = FridaInProcessExecutor::new( - &gum, - InProcessExecutor::new( - &mut frida_harness, - observers, - &mut fuzzer, - &mut state, - &mut mgr, - )?, - &mut frida_helper, - ); - - // In case the corpus is empty (on first run), reset - if state.must_load_initial_inputs() { - state - .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &options.input) - .unwrap_or_else(|_| { - panic!("Failed to load initial corpus at {:?}", &options.input) - }); - println!("We imported {} inputs from disk.", state.corpus().count()); - } - - // Create an observation channel using cmplog map - let cmplog_observer = CmpLogObserver::new("cmplog", true); - - let mut executor = ShadowExecutor::new(executor, tuple_list!(cmplog_observer)); - - let tracing = ShadowTracingStage::new(&mut executor); - - // Setup a randomic Input2State stage - let i2s = StdMutationalStage::new(StdScheduledMutator::new(tuple_list!( - I2SRandReplace::new() - ))); - - // Setup a basic mutator - let mutational = StdMutationalStage::new(mutator); - - // The order of the stages matter! - let mut stages = tuple_list!(tracing, i2s, mutational); - - fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?; - - Ok(()) - })(state, mgr, client_description) - } else { - (|state: Option<_>, - mut mgr: LlmpRestartingEventManager<_, _, _>, - _client_description| { - let gum = Gum::obtain(); - - let coverage = CoverageRuntime::new(); - - let mut frida_helper = - FridaInstrumentationHelper::new(&gum, options, tuple_list!(coverage)); - - // Create an observation channel using the coverage map - let edges_observer = HitcountsMapObserver::new(StdMapObserver::from_mut_ptr( - "edges", - frida_helper.map_mut_ptr().unwrap(), - MAP_SIZE, - )) - .track_indices(); - - // Create an observation channel to keep track of the execution time - let time_observer = TimeObserver::new("time"); - #[cfg(unix)] - let asan_observer = AsanErrorsObserver::from_static_asan_errors(); - - // Feedback to rate the interestingness of an input - // This one is composed by two Feedbacks in OR - let mut feedback = feedback_or!( - // New maximization map feedback linked to the edges observer and the feedback state - MaxMapFeedback::new(&edges_observer), - // Time feedback, this one does not need a feedback state - TimeFeedback::new(&time_observer) - ); - - #[cfg(unix)] - let mut objective = feedback_or_fast!( - CrashFeedback::new(), - TimeoutFeedback::new(), - feedback_and_fast!( - ConstFeedback::from(false), - AsanErrorsFeedback::new(&asan_observer) - ) - ); - #[cfg(windows)] - let mut objective = feedback_or_fast!(CrashFeedback::new(), TimeoutFeedback::new()); - - // If not restarting, create a State from scratch - let mut state = state.unwrap_or_else(|| { - StdState::new( - // RNG - StdRand::new(), - // Corpus that will be evolved, we keep it in memory for performance - CachedOnDiskCorpus::no_meta(PathBuf::from("./corpus_discovered"), 64) - .unwrap(), - // Corpus in which we store solutions (crashes in this example), - // on disk so the user can get them after stopping the fuzzer - OnDiskCorpus::new(options.output.clone()).unwrap(), - &mut feedback, - &mut objective, - ) - .unwrap() - }); - - println!("We're a client, let's fuzz :)"); - - // Create a PNG dictionary if not existing - if state.metadata_map().get::().is_none() { - state.add_metadata(Tokens::from([ - vec![137, 80, 78, 71, 13, 10, 26, 10], // PNG header - b"IHDR".to_vec(), - b"IDAT".to_vec(), - b"PLTE".to_vec(), - b"IEND".to_vec(), - ])); - } - - // Setup a basic mutator with a mutational stage - let mutator = StdScheduledMutator::new(havoc_mutations().merge(tokens_mutations())); - - // A minimization+queue policy to get testcasess from the corpus - let scheduler = - IndexesLenTimeMinimizerScheduler::new(&edges_observer, QueueScheduler::new()); - - // A fuzzer with feedbacks and a corpus scheduler - let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); - - #[cfg(unix)] - let observers = tuple_list!(edges_observer, time_observer, asan_observer); - #[cfg(windows)] - let observers = tuple_list!(edges_observer, time_observer); - - // Create the executor for an in-process function with just one observer for edge coverage - let mut executor = FridaInProcessExecutor::new( - &gum, - InProcessExecutor::new( - &mut frida_harness, - observers, - &mut fuzzer, - &mut state, - &mut mgr, - )?, - &mut frida_helper, - ); - - // In case the corpus is empty (on first run), reset - if state.must_load_initial_inputs() { - state - .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &options.input) - .unwrap_or_else(|_| { - panic!("Failed to load initial corpus at {:?}", &options.input) - }); - println!("We imported {} inputs from disk.", state.corpus().count()); - } - - let mut stages = tuple_list!(StdMutationalStage::new(mutator)); - - fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?; - - Ok(()) - })(state, mgr, client_description) - } + // if options.asan && options.asan_cores.contains(client_description.core_id()) { + (|state: Option<_>, mut mgr: LlmpRestartingEventManager<_, _, _>, _client_description| { + let gum = Gum::obtain(); + + let coverage = CoverageRuntime::new(); + let asan = AsanRuntime::new(options); + let cmplog = CmpLogRuntime::new(); + + let client_description_clone = client_description.clone(); + let options_clone = options.clone(); + let client_description_clone2 = client_description.clone(); + let options_clone2 = options.clone(); + let mut frida_helper = FridaInstrumentationHelper::new( + &gum, + options, + tuple_list!( + IfElseRuntime::new( + move || Ok(is_asan(&options_clone, &client_description_clone)), + tuple_list!(asan), + tuple_list!() + ), + IfElseRuntime::new( + move || Ok(is_cmplog(&options_clone2, &client_description_clone2)), + tuple_list!(cmplog), + tuple_list!() + ), + coverage + ), + ); + + // Create an observation channel using the coverage map + let edges_observer = HitcountsMapObserver::new(StdMapObserver::from_mut_ptr( + "edges", + frida_helper.map_mut_ptr().unwrap(), + MAP_SIZE, + )) + .track_indices(); + + // Create an observation channel to keep track of the execution time + let time_observer = TimeObserver::new("time"); + let asan_observer = AsanErrorsObserver::from_static_asan_errors(); + + // Feedback to rate the interestingness of an input + // This one is composed by two Feedbacks in OR + let mut feedback = feedback_or!( + // New maximization map feedback linked to the edges observer and the feedback state + MaxMapFeedback::new(&edges_observer), + // Time feedback, this one does not need a feedback state + TimeFeedback::new(&time_observer) + ); + + // Feedbacks to recognize an input as solution + #[cfg(unix)] + let mut objective = feedback_or_fast!( + CrashFeedback::new(), + TimeoutFeedback::new(), + // true enables the AsanErrorFeedback + feedback_and_fast!( + ConstFeedback::from(true), + AsanErrorsFeedback::new(&asan_observer) + ) + ); + #[cfg(windows)] + let mut objective = feedback_or_fast!(CrashFeedback::new(), TimeoutFeedback::new()); + + // If not restarting, create a State from scratch + let mut state = state.unwrap_or_else(|| { + StdState::new( + // RNG + StdRand::new(), + // Corpus that will be evolved, we keep it in memory for performance + CachedOnDiskCorpus::no_meta(PathBuf::from("./corpus_discovered"), 64).unwrap(), + // Corpus in which we store solutions (crashes in this example), + // on disk so the user can get them after stopping the fuzzer + OnDiskCorpus::new(options.output.clone()).unwrap(), + &mut feedback, + &mut objective, + ) + .unwrap() + }); + + println!("We're a client, let's fuzz :)"); + + // Create a PNG dictionary if not existing + if state.metadata_map().get::().is_none() { + state.add_metadata(Tokens::from([ + vec![137, 80, 78, 71, 13, 10, 26, 10], // PNG header + b"IHDR".to_vec(), + b"IDAT".to_vec(), + b"PLTE".to_vec(), + b"IEND".to_vec(), + ])); + } + + // Setup a basic mutator with a mutational stage + let mutator = StdScheduledMutator::new(havoc_mutations().merge(tokens_mutations())); + + // A minimization+queue policy to get testcasess from the corpus + let scheduler = + IndexesLenTimeMinimizerScheduler::new(&edges_observer, QueueScheduler::new()); + + // A fuzzer with feedbacks and a corpus scheduler + let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); + + #[cfg(unix)] + let observers = tuple_list!(edges_observer, time_observer, asan_observer); + #[cfg(windows)] + let observers = tuple_list!(edges_observer, time_observer); + + // Create the executor for an in-process function with just one observer for edge coverage + let mut executor = FridaInProcessExecutor::new( + &gum, + InProcessExecutor::new( + &mut frida_harness, + observers, + &mut fuzzer, + &mut state, + &mut mgr, + )?, + &mut frida_helper, + ); + // Create an observation channel using cmplog map + let cmplog_observer = CmpLogObserver::new("cmplog", true); + + let mut executor = ShadowExecutor::new(executor, tuple_list!(cmplog_observer)); + + let tracing = ShadowTracingStage::new(&mut executor); + + // Setup a randomic Input2State stage + let i2s = StdMutationalStage::new(StdScheduledMutator::new(tuple_list!( + I2SRandReplace::new() + ))); + + // In case the corpus is empty (on first run), reset + if state.must_load_initial_inputs() { + state + .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &options.input) + .unwrap_or_else(|_| { + panic!("Failed to load initial corpus at {:?}", &options.input) + }); + println!("We imported {} inputs from disk.", state.corpus().count()); + } + + let mut stages = tuple_list!( + IfElseStage::new( + |_, _, _, _| Ok(is_cmplog(&options, &client_description)), + tuple_list!(tracing, i2s), + tuple_list!() + ), + StdMutationalStage::new(mutator) + ); + + fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?; + + Ok(()) + })(state, mgr, client_description.clone()) }; Launcher::builder() diff --git a/fuzzers/binary_only/frida_windows_gdiplus/Cargo.toml b/fuzzers/binary_only/frida_windows_gdiplus/Cargo.toml index 43b2252748..92ae48c01e 100644 --- a/fuzzers/binary_only/frida_windows_gdiplus/Cargo.toml +++ b/fuzzers/binary_only/frida_windows_gdiplus/Cargo.toml @@ -36,7 +36,7 @@ libafl_targets = { path = "../../../libafl_targets", features = [ libloading = "0.8.5" log = { version = "0.4.22", features = ["release_max_level_info"] } mimalloc = { version = "0.1.43", default-features = false } -dlmalloc = { version = "0.2.6", features = ["global"] } +#dlmalloc = { version = "0.2.6", features = ["global"] } color-backtrace = "0.6.1" env_logger = "0.11.5" iced-x86 = { version = "1.21.0", features = ["code_asm"] } diff --git a/fuzzers/binary_only/frida_windows_gdiplus/src/fuzzer.rs b/fuzzers/binary_only/frida_windows_gdiplus/src/fuzzer.rs index 41afe956b5..42e9d5de25 100644 --- a/fuzzers/binary_only/frida_windows_gdiplus/src/fuzzer.rs +++ b/fuzzers/binary_only/frida_windows_gdiplus/src/fuzzer.rs @@ -6,16 +6,9 @@ //! going to make it compilable only for Windows, don't forget to modify the //! `scripts/test_fuzzer.sh` to opt-out this fuzzer from that test. -#[cfg(unix)] use mimalloc::MiMalloc; -#[cfg(unix)] #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; -#[cfg(windows)] -use dlmalloc::GlobalDlmalloc; -#[cfg(windows)] -#[global_allocator] -static GLOBAL: GlobalDlmalloc = GlobalDlmalloc; use std::path::PathBuf; diff --git a/libafl_frida/src/asan/asan_rt.rs b/libafl_frida/src/asan/asan_rt.rs index 0fd25d79d2..727b248249 100644 --- a/libafl_frida/src/asan/asan_rt.rs +++ b/libafl_frida/src/asan/asan_rt.rs @@ -509,7 +509,7 @@ impl AsanRuntime { let _ = [<$lib_ident:snake:upper _ $name:snake:upper _PTR>].set(unsafe {std::mem::transmute::<*const c_void, extern "C" fn($($param: $param_type),*) -> $return_type>(target_function.0)}).unwrap(); - #[expect(non_snake_case)] + #[allow(non_snake_case)] unsafe extern "C" fn []($($param: $param_type),*) -> $return_type { let mut invocation = Interceptor::current_invocation(); let this = &mut *(invocation.replacement_data().unwrap().0 as *mut AsanRuntime); @@ -593,7 +593,7 @@ impl AsanRuntime { let _ = [<$lib_ident:snake:upper _ $name:snake:upper _PTR>].set(unsafe {std::mem::transmute::<*const c_void, extern "C" fn($($param: $param_type),*) -> $return_type>(target_function.0)}).unwrap_or_else(|e| println!("{:?}", e)); - #[expect(non_snake_case)] + #[allow(non_snake_case)] unsafe extern "C" fn []($($param: $param_type),*) -> $return_type { let mut invocation = Interceptor::current_invocation(); let this = &mut *(invocation.replacement_data().unwrap().0 as *mut AsanRuntime); diff --git a/libafl_frida/src/helper.rs b/libafl_frida/src/helper.rs index d7148a6a7b..278ba6f553 100644 --- a/libafl_frida/src/helper.rs +++ b/libafl_frida/src/helper.rs @@ -52,6 +52,78 @@ pub trait FridaRuntime: 'static + Debug { fn post_exec(&mut self, input_bytes: &[u8]) -> Result<(), Error>; } +/// Use the runtime if closure evaluates to true +pub struct IfElseRuntime { + closure: CB, + if_runtimes: FR1, + else_runtimes: FR2, +} + +impl Debug for IfElseRuntime +where + FR1: Debug, + FR2: Debug, +{ + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Debug::fmt(&self.if_runtimes, f)?; + Debug::fmt(&self.else_runtimes, f)?; + Ok(()) + } +} +impl IfElseRuntime { + /// Constructor for this conditionally enabled runtime + pub fn new(closure: CB, if_runtimes: FR1, else_runtimes: FR2) -> Self { + Self { + closure, + if_runtimes, + else_runtimes, + } + } +} + +impl FridaRuntime for IfElseRuntime +where + CB: FnMut() -> Result + 'static, + FR1: FridaRuntimeTuple + 'static, + FR2: FridaRuntimeTuple + 'static, +{ + fn init( + &mut self, + gum: &Gum, + ranges: &RangeMap, + module_map: &Rc, + ) { + if (self.closure)().unwrap() { + self.if_runtimes.init_all(gum, ranges, module_map); + } else { + self.else_runtimes.init_all(gum, ranges, module_map); + } + } + + fn deinit(&mut self, gum: &Gum) { + if (self.closure)().unwrap() { + self.if_runtimes.deinit_all(gum); + } else { + self.else_runtimes.deinit_all(gum); + } + } + + fn pre_exec(&mut self, input_bytes: &[u8]) -> Result<(), Error> { + if (self.closure)()? { + self.if_runtimes.pre_exec_all(input_bytes) + } else { + self.else_runtimes.pre_exec_all(input_bytes) + } + } + + fn post_exec(&mut self, input_bytes: &[u8]) -> Result<(), Error> { + if (self.closure)()? { + self.if_runtimes.post_exec_all(input_bytes) + } else { + self.else_runtimes.post_exec_all(input_bytes) + } + } +} /// The tuple for Frida Runtime pub trait FridaRuntimeTuple: MatchFirstType + Debug { /// Initialization