From bfe0a97077fce40eda412e2f214786b1300b135f Mon Sep 17 00:00:00 2001 From: Toka Date: Thu, 29 Jul 2021 04:19:50 +0900 Subject: [PATCH] AFLFast (#208) * Defined PowerScheduleTestcase * calibrate stage * powerschedule * PowerQueueCorpusScheduler * Calstat * Integer for observer * update power.rs * calculate_score * no_std * no_std * calibrate_case * calculate_score done * update fuzz_level * update depth * rename, PowerScheduleQueueScheduler::next * increment hashmap * n_fuzz * no_std * fmt * clippy * check bitmap_size * hash * compile * rename, vec * fuzzer runs * fixes * rename, no_std log2 * fmt * clippy * fmt * unused imports * use exec_time * getter/setter * No unwrap * ToString * fixes * ahash * fmt * use favored from power.rs side * rename * IsFavoredMetadata * run_target * clippy fixes * doc & fix * doc * rename * forgot log2 * fix * fix * added comments explaining why the COE and FAST schedules are different from the thesis * saturated increment * saturating_add * put n_fuzz in PowerScheduleMetadata Co-authored-by: Andrea Fioraldi --- fuzzers/libfuzzer_libpng/src/lib.rs | 23 +- libafl/Cargo.toml | 1 + libafl/src/corpus/mod.rs | 5 +- libafl/src/corpus/powersched.rs | 105 +++++++ libafl/src/corpus/testcase.rs | 82 ++++++ libafl/src/observers/map.rs | 64 +++-- libafl/src/stages/calibrate.rs | 236 ++++++++++++++++ libafl/src/stages/mod.rs | 7 +- libafl/src/stages/mutational.rs | 8 +- libafl/src/stages/power.rs | 420 +++++++++++++++++++++++----- 10 files changed, 849 insertions(+), 102 deletions(-) create mode 100644 libafl/src/corpus/powersched.rs create mode 100644 libafl/src/stages/calibrate.rs diff --git a/fuzzers/libfuzzer_libpng/src/lib.rs b/fuzzers/libfuzzer_libpng/src/lib.rs index cc5e056399..5afe9a76f5 100644 --- a/fuzzers/libfuzzer_libpng/src/lib.rs +++ b/fuzzers/libfuzzer_libpng/src/lib.rs @@ -9,7 +9,7 @@ use libafl::{ bolts::{current_nanos, rands::StdRand}, corpus::{ Corpus, InMemoryCorpus, IndexesLenTimeMinimizerCorpusScheduler, OnDiskCorpus, - QueueCorpusScheduler, + PowerQueueCorpusScheduler, }, events::{setup_restarting_mgr_std, EventRestarter}, executors::{inprocess::InProcessExecutor, ExitKind, TimeoutExecutor}, @@ -17,11 +17,13 @@ use libafl::{ feedbacks::{CrashFeedback, MapFeedbackState, MaxMapFeedback, TimeFeedback, TimeoutFeedback}, fuzzer::{Fuzzer, StdFuzzer}, inputs::{BytesInput, HasTargetBytes}, - mutators::mopt_mutator::StdMOptMutator, - mutators::scheduled::{havoc_mutations, tokens_mutations}, + mutators::scheduled::{havoc_mutations, tokens_mutations, StdScheduledMutator}, mutators::token_mutations::Tokens, observers::{HitcountsMapObserver, StdMapObserver, TimeObserver}, - stages::mutational::StdMutationalStage, + stages::{ + calibrate::CalibrationStage, + power::{PowerMutationalStage, PowerSchedule}, + }, state::{HasCorpus, HasMetadata, StdState}, stats::MultiStats, Error, @@ -121,16 +123,15 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re // Setup a basic mutator with a mutational stage - let mutator = StdMOptMutator::new(&mut state, havoc_mutations().merge(tokens_mutations()), 5)?; - let mut stages = tuple_list!(StdMutationalStage::new(mutator)); - - /* let mutator = StdScheduledMutator::new(havoc_mutations().merge(tokens_mutations())); - let mut stages = tuple_list!(StdMutationalStage::new(mutator)); - */ + + let calibration = CalibrationStage::new(&mut state, &edges_observer); + let power = PowerMutationalStage::new(mutator, PowerSchedule::FAST, &edges_observer); + + let mut stages = tuple_list!(calibration, power); // A minimization+queue policy to get testcasess from the corpus - let scheduler = IndexesLenTimeMinimizerCorpusScheduler::new(QueueCorpusScheduler::new()); + let scheduler = IndexesLenTimeMinimizerCorpusScheduler::new(PowerQueueCorpusScheduler::new()); // A fuzzer with feedbacks and a corpus scheduler let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); diff --git a/libafl/Cargo.toml b/libafl/Cargo.toml index 7a774f5817..9cbfa44b6a 100644 --- a/libafl/Cargo.toml +++ b/libafl/Cargo.toml @@ -74,6 +74,7 @@ hostname = { version = "^0.3", optional = true } # Is there really no gethostnam rand = { version = "0.8.1", optional = true } # rand_core = { version = "0.6.2", optional = true } # This dependency allows us to export our RomuRand as rand::Rng. nix = { version = "0.20.0", optional = true } +libm = "0.2.1" [target.'cfg(target_os = "android")'.dependencies] backtrace = { version = "0.3", optional = true, default-features = false, features = ["std", "libbacktrace"] } # for llmp_debug diff --git a/libafl/src/corpus/mod.rs b/libafl/src/corpus/mod.rs index ac638c1de6..28f4c7b0a2 100644 --- a/libafl/src/corpus/mod.rs +++ b/libafl/src/corpus/mod.rs @@ -1,7 +1,7 @@ //! Corpuses contain the testcases, either in mem, on disk, or somewhere else. pub mod testcase; -pub use testcase::Testcase; +pub use testcase::{PowerScheduleTestcaseMetaData, Testcase}; pub mod inmemory; pub use inmemory::InMemoryCorpus; @@ -26,6 +26,9 @@ pub use minimizer::{ TopRatedsMetadata, }; +pub mod powersched; +pub use powersched::PowerQueueCorpusScheduler; + use alloc::borrow::ToOwned; use core::{cell::RefCell, marker::PhantomData}; diff --git a/libafl/src/corpus/powersched.rs b/libafl/src/corpus/powersched.rs new file mode 100644 index 0000000000..4d8044d606 --- /dev/null +++ b/libafl/src/corpus/powersched.rs @@ -0,0 +1,105 @@ +//! The queue corpus scheduler for power schedules. + +use alloc::string::{String, ToString}; +use core::marker::PhantomData; + +use crate::{ + corpus::{Corpus, CorpusScheduler, PowerScheduleTestcaseMetaData}, + inputs::Input, + stages::PowerScheduleMetadata, + state::{HasCorpus, HasMetadata}, + Error, +}; + +pub struct PowerQueueCorpusScheduler +where + S: HasCorpus + HasMetadata, + C: Corpus, + I: Input, +{ + phantom: PhantomData<(C, I, S)>, +} + +impl Default for PowerQueueCorpusScheduler +where + S: HasCorpus + HasMetadata, + C: Corpus, + I: Input, +{ + fn default() -> Self { + Self::new() + } +} + +impl CorpusScheduler for PowerQueueCorpusScheduler +where + S: HasCorpus + HasMetadata, + C: Corpus, + I: Input, +{ + /// Add an entry to the corpus and return its index + fn on_add(&self, state: &mut S, idx: usize) -> Result<(), Error> { + let current_idx = *state.corpus().current(); + + let mut depth = match current_idx { + Some(idx) => state + .corpus() + .get(idx)? + .borrow_mut() + .metadata_mut() + .get_mut::() + .ok_or_else(|| Error::KeyNotFound("PowerScheduleTestData not found".to_string()))? + .depth(), + None => 0, + }; + + // Attach a `PowerScheduleTestData` to the queue entry. + depth += 1; + state + .corpus() + .get(idx)? + .borrow_mut() + .add_metadata(PowerScheduleTestcaseMetaData::new(depth)); + Ok(()) + } + + fn next(&self, state: &mut S) -> Result { + if state.corpus().count() == 0 { + Err(Error::Empty(String::from("No entries in corpus"))) + } else { + let id = match state.corpus().current() { + Some(cur) => { + if *cur + 1 >= state.corpus().count() { + let psmeta = state + .metadata_mut() + .get_mut::() + .ok_or_else(|| { + Error::KeyNotFound("PowerScheduleMetadata not found".to_string()) + })?; + psmeta.set_queue_cycles(psmeta.queue_cycles() + 1); + 0 + } else { + *cur + 1 + } + } + None => 0, + }; + *state.corpus_mut().current_mut() = Some(id); + Ok(id) + } + } +} + +impl PowerQueueCorpusScheduler +where + S: HasCorpus + HasMetadata, + C: Corpus, + I: Input, +{ + #[must_use] + pub fn new() -> Self { + Self { + phantom: PhantomData, + } + } +} diff --git a/libafl/src/corpus/testcase.rs b/libafl/src/corpus/testcase.rs index 6906101a6b..1b539e8f57 100644 --- a/libafl/src/corpus/testcase.rs +++ b/libafl/src/corpus/testcase.rs @@ -120,15 +120,22 @@ where } /// Get the execution time of the testcase + #[inline] pub fn exec_time(&self) -> &Option { &self.exec_time } /// Get the execution time of the testcase (mut) + #[inline] pub fn exec_time_mut(&mut self) -> &mut Option { &mut self.exec_time } + #[inline] + pub fn set_exec_time(&mut self, time: Duration) { + self.exec_time = Some(time); + } + /// Create a new Testcase instace given an input #[inline] pub fn new(input: T) -> Self @@ -206,3 +213,78 @@ where Testcase::new(input) } } + +/// The Metadata for each testcase used in power schedules. +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct PowerScheduleTestcaseMetaData { + /// Number of bits set in bitmap, updated in calibrate_case + bitmap_size: u64, + /// Number of fuzzing iterations, updated in perform_mutational + fuzz_level: u64, + /// Number of queue cycles behind + handicap: u64, + /// Path depth, initialized in on_add + depth: u64, + /// Offset in n_fuzz + n_fuzz_entry: usize, +} + +impl PowerScheduleTestcaseMetaData { + #[must_use] + pub fn new(depth: u64) -> Self { + Self { + bitmap_size: 0, + fuzz_level: 0, + handicap: 0, + depth, + n_fuzz_entry: 0, + } + } + + #[must_use] + pub fn bitmap_size(&self) -> u64 { + self.bitmap_size + } + + pub fn set_bitmap_size(&mut self, val: u64) { + self.bitmap_size = val; + } + + #[must_use] + pub fn fuzz_level(&self) -> u64 { + self.fuzz_level + } + + pub fn set_fuzz_level(&mut self, val: u64) { + self.fuzz_level = val; + } + + #[must_use] + pub fn handicap(&self) -> u64 { + self.handicap + } + + pub fn set_handicap(&mut self, val: u64) { + self.handicap = val; + } + + #[must_use] + pub fn depth(&self) -> u64 { + self.depth + } + + pub fn set_depth(&mut self, val: u64) { + self.depth = val; + } + + #[must_use] + pub fn n_fuzz_entry(&self) -> usize { + self.n_fuzz_entry + } + + pub fn set_n_fuzz_entry(&mut self, val: usize) { + self.n_fuzz_entry = val; + } +} + +crate::impl_serdeany!(PowerScheduleTestcaseMetaData); diff --git a/libafl/src/observers/map.rs b/libafl/src/observers/map.rs index 8fb6aac66e..22416318b3 100644 --- a/libafl/src/observers/map.rs +++ b/libafl/src/observers/map.rs @@ -1,10 +1,15 @@ //! The `MapObserver` provides access a map, usually injected into the target +use ahash::AHasher; use alloc::{ string::{String, ToString}, vec::Vec, }; -use core::slice::from_raw_parts_mut; +use core::{ + hash::Hasher, + slice::{from_raw_parts, from_raw_parts_mut}, +}; +use num::Integer; use serde::{Deserialize, Serialize}; use crate::{ @@ -19,7 +24,7 @@ use crate::{ /// A [`MapObserver`] observes the static map, as oftentimes used for afl-like coverage information pub trait MapObserver: Named + serde::Serialize + serde::de::DeserializeOwned where - T: Default + Copy, + T: Integer + Default + Copy, { /// Get the map fn map(&self) -> &[T]; @@ -32,6 +37,31 @@ where self.map().len() } + /// Count the set bytes in the map + fn count_bytes(&self) -> u64 { + let initial = self.initial(); + let cnt = self.usable_count(); + let mut res = 0; + for x in self.map()[0..cnt].iter() { + if *x != initial { + res += 1; + } + } + res + } + + /// Compute the hash of the map + fn hash(&self) -> u64 { + let mut hasher = AHasher::new_with_keys(0, 0); + let ptr = self.map().as_ptr() as *const u8; + let map_size = self.map().len() / core::mem::size_of::(); + unsafe { + hasher.write(from_raw_parts(ptr, map_size)); + } + + hasher.finish() + } + /// Get the initial value for reset() fn initial(&self) -> T; @@ -62,7 +92,7 @@ where #[allow(clippy::unsafe_derive_deserialize)] pub struct StdMapObserver<'a, T> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, { map: OwnedSliceMut<'a, T>, initial: T, @@ -71,7 +101,7 @@ where impl<'a, I, S, T> Observer for StdMapObserver<'a, T> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, Self: MapObserver, { #[inline] @@ -82,7 +112,7 @@ where impl<'a, T> Named for StdMapObserver<'a, T> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, { #[inline] fn name(&self) -> &str { @@ -92,7 +122,7 @@ where impl<'a, T> MapObserver for StdMapObserver<'a, T> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, { #[inline] fn map(&self) -> &[T] { @@ -122,7 +152,7 @@ where impl<'a, T> StdMapObserver<'a, T> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, { /// Creates a new [`MapObserver`] #[must_use] @@ -167,7 +197,7 @@ where #[allow(clippy::unsafe_derive_deserialize)] pub struct ConstMapObserver<'a, T, const N: usize> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, { map: OwnedSliceMut<'a, T>, initial: T, @@ -176,7 +206,7 @@ where impl<'a, I, S, T, const N: usize> Observer for ConstMapObserver<'a, T, N> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, Self: MapObserver, { #[inline] @@ -187,7 +217,7 @@ where impl<'a, T, const N: usize> Named for ConstMapObserver<'a, T, N> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, { #[inline] fn name(&self) -> &str { @@ -197,7 +227,7 @@ where impl<'a, T, const N: usize> MapObserver for ConstMapObserver<'a, T, N> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, { #[inline] fn usable_count(&self) -> usize { @@ -232,7 +262,7 @@ where impl<'a, T, const N: usize> ConstMapObserver<'a, T, N> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, { /// Creates a new [`MapObserver`] #[must_use] @@ -278,7 +308,7 @@ where #[allow(clippy::unsafe_derive_deserialize)] pub struct VariableMapObserver<'a, T> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, { map: OwnedSliceMut<'a, T>, size: OwnedRefMut<'a, usize>, @@ -288,7 +318,7 @@ where impl<'a, I, S, T> Observer for VariableMapObserver<'a, T> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, Self: MapObserver, { #[inline] @@ -299,7 +329,7 @@ where impl<'a, T> Named for VariableMapObserver<'a, T> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, { #[inline] fn name(&self) -> &str { @@ -309,7 +339,7 @@ where impl<'a, T> MapObserver for VariableMapObserver<'a, T> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, { #[inline] fn map(&self) -> &[T] { @@ -344,7 +374,7 @@ where impl<'a, T> VariableMapObserver<'a, T> where - T: Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, { /// Creates a new [`MapObserver`] pub fn new(name: &'static str, map: &'a mut [T], size: &'a mut usize) -> Self { diff --git a/libafl/src/stages/calibrate.rs b/libafl/src/stages/calibrate.rs new file mode 100644 index 0000000000..7b116b3863 --- /dev/null +++ b/libafl/src/stages/calibrate.rs @@ -0,0 +1,236 @@ +//! The calibration stage. The fuzzer measures the average exec time and the bitmap size. + +use crate::{ + bolts::current_time, + corpus::{Corpus, PowerScheduleTestcaseMetaData}, + executors::{Executor, HasObservers}, + fuzzer::Evaluator, + inputs::Input, + observers::{MapObserver, ObserversTuple}, + stages::Stage, + state::{HasCorpus, HasMetadata}, + Error, +}; +use alloc::{ + string::{String, ToString}, + vec::Vec, +}; +use core::{marker::PhantomData, time::Duration}; +use num::Integer; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug)] +pub struct CalibrationStage +where + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + C: Corpus, + E: Executor + HasObservers, + I: Input, + O: MapObserver, + OT: ObserversTuple, + S: HasCorpus + HasMetadata, + Z: Evaluator, +{ + map_observer_name: String, + stage_max: usize, + #[allow(clippy::type_complexity)] + phantom: PhantomData<(C, E, EM, I, O, OT, S, T, Z)>, +} + +const CAL_STAGE_MAX: usize = 8; + +impl Stage + for CalibrationStage +where + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + C: Corpus, + E: Executor + HasObservers, + I: Input, + O: MapObserver, + OT: ObserversTuple, + S: HasCorpus + HasMetadata, + Z: Evaluator, +{ + #[inline] + #[allow(clippy::let_and_return)] + fn perform( + &mut self, + fuzzer: &mut Z, + executor: &mut E, + state: &mut S, + manager: &mut EM, + corpus_idx: usize, + ) -> Result<(), Error> { + let iter = self.stage_max; + let handicap = state + .metadata() + .get::() + .ok_or_else(|| Error::KeyNotFound("PowerScheduleMetadata not found".to_string()))? + .queue_cycles; + + let start = current_time(); + + for _i in 0..iter { + let input = state + .corpus() + .get(corpus_idx)? + .borrow_mut() + .load_input()? + .clone(); + let _ = executor.run_target(fuzzer, state, manager, &input)?; + } + + let end = current_time(); + + let map = executor + .observers() + .match_name::(&self.map_observer_name) + .ok_or_else(|| Error::KeyNotFound("MapObserver not found".to_string()))?; + + let bitmap_size = map.count_bytes(); + + let psmeta = state + .metadata_mut() + .get_mut::() + .ok_or_else(|| Error::KeyNotFound("PowerScheduleMetadata not found".to_string()))?; + + psmeta.set_exec_time(psmeta.exec_time() + (end - start)); + psmeta.set_cycles(psmeta.cycles() + (iter as u64)); + psmeta.set_bitmap_size(psmeta.bitmap_size() + bitmap_size); + psmeta.set_bitmap_entries(psmeta.bitmap_entries() + 1); + + // println!("psmeta: {:#?}", psmeta); + let mut testcase = state.corpus().get(corpus_idx)?.borrow_mut(); + + testcase.set_exec_time((end - start) / (iter as u32)); + // println!("time: {:#?}", testcase.exec_time()); + let data = testcase + .metadata_mut() + .get_mut::() + .ok_or_else(|| Error::KeyNotFound("PowerScheduleTestData not found".to_string()))?; + + data.set_bitmap_size(bitmap_size); + data.set_handicap(handicap); + data.set_fuzz_level(data.fuzz_level() + 1); + // println!("data: {:#?}", data); + + Ok(()) + } +} + +pub const N_FUZZ_SIZE: usize = 1 << 21; + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct PowerScheduleMetadata { + /// Measured exec time during calibration + exec_time: Duration, + /// Calibration cycles + cycles: u64, + /// Size of the observer map + bitmap_size: u64, + /// Number of filled map entries + bitmap_entries: u64, + /// Queue cycles + queue_cycles: u64, + /// The vector to contain the frequency of each execution path. + n_fuzz: Vec, +} + +/// The metadata for runs in the calibration stage. +impl PowerScheduleMetadata { + #[must_use] + pub fn new() -> Self { + Self { + exec_time: Duration::from_millis(0), + cycles: 0, + bitmap_size: 0, + bitmap_entries: 0, + queue_cycles: 0, + n_fuzz: vec![0; N_FUZZ_SIZE], + } + } + + #[must_use] + pub fn exec_time(&self) -> Duration { + self.exec_time + } + + pub fn set_exec_time(&mut self, time: Duration) { + self.exec_time = time; + } + + #[must_use] + pub fn cycles(&self) -> u64 { + self.cycles + } + + pub fn set_cycles(&mut self, val: u64) { + self.cycles = val; + } + + #[must_use] + pub fn bitmap_size(&self) -> u64 { + self.bitmap_size + } + + pub fn set_bitmap_size(&mut self, val: u64) { + self.bitmap_size = val; + } + + #[must_use] + pub fn bitmap_entries(&self) -> u64 { + self.bitmap_entries + } + + pub fn set_bitmap_entries(&mut self, val: u64) { + self.bitmap_entries = val; + } + + #[must_use] + pub fn queue_cycles(&self) -> u64 { + self.queue_cycles + } + + pub fn set_queue_cycles(&mut self, val: u64) { + self.queue_cycles = val; + } + + #[must_use] + pub fn n_fuzz(&self) -> &[u32] { + &self.n_fuzz + } + + #[must_use] + pub fn n_fuzz_mut(&mut self) -> &mut [u32] { + &mut self.n_fuzz + } +} + +crate::impl_serdeany!(PowerScheduleMetadata); + +impl CalibrationStage +where + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + C: Corpus, + E: Executor + HasObservers, + I: Input, + O: MapObserver, + OT: ObserversTuple, + S: HasCorpus + HasMetadata, + Z: Evaluator, +{ + pub fn new(state: &mut S, map_observer_name: &O) -> Self { + state.add_metadata::(PowerScheduleMetadata::new()); + Self { + map_observer_name: map_observer_name.name().to_string(), + stage_max: CAL_STAGE_MAX, + phantom: PhantomData, + } + } +} + +impl Default for PowerScheduleMetadata { + fn default() -> Self { + Self::new() + } +} diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs index 5042d8e680..f27279e25e 100644 --- a/libafl/src/stages/mod.rs +++ b/libafl/src/stages/mod.rs @@ -11,9 +11,12 @@ pub use mutational::{MutationalStage, StdMutationalStage}; pub mod tracing; pub use tracing::{ShadowTracingStage, TracingStage}; -//pub mod power; -//pub use power::PowerMutationalStage; +pub mod calibrate; +pub use calibrate::{CalibrationStage, PowerScheduleMetadata}; + +pub mod power; use crate::Error; +pub use power::PowerMutationalStage; /// A stage is one step in the fuzzing process. /// Multiple stages will be scheduled one by one for each input. diff --git a/libafl/src/stages/mutational.rs b/libafl/src/stages/mutational.rs index a5df0bda43..0d5f73806f 100644 --- a/libafl/src/stages/mutational.rs +++ b/libafl/src/stages/mutational.rs @@ -36,7 +36,7 @@ where fn mutator_mut(&mut self) -> &mut M; /// Gets the number of iterations this mutator should run for. - fn iterations(&self, state: &mut S) -> usize; + fn iterations(&self, state: &mut S, corpus_idx: usize) -> Result; /// Runs this (mutational) stage for the given testcase #[allow(clippy::cast_possible_wrap)] // more than i32 stages on 32 bit system - highly unlikely... @@ -48,7 +48,7 @@ where manager: &mut EM, corpus_idx: usize, ) -> Result<(), Error> { - let num = self.iterations(state); + let num = self.iterations(state, corpus_idx)?; for i in 0..num { start_timer!(state); @@ -118,8 +118,8 @@ where } /// Gets the number of iterations as a random number - fn iterations(&self, state: &mut S) -> usize { - 1 + state.rand_mut().below(DEFAULT_MUTATIONAL_MAX_ITERATIONS) as usize + fn iterations(&self, state: &mut S, _corpus_idx: usize) -> Result { + Ok(1 + state.rand_mut().below(DEFAULT_MUTATIONAL_MAX_ITERATIONS) as usize) } } diff --git a/libafl/src/stages/power.rs b/libafl/src/stages/power.rs index 61183ee462..af27d26937 100644 --- a/libafl/src/stages/power.rs +++ b/libafl/src/stages/power.rs @@ -1,48 +1,69 @@ +//! The power schedules. This stage should be invoked after the calibration stage. + +use alloc::string::{String, ToString}; use core::marker::PhantomData; +use num::Integer; use crate::{ - bolts::rands::Rand, - corpus::{Corpus, CorpusScheduler}, - events::EventManager, + corpus::{Corpus, IsFavoredMetadata, PowerScheduleTestcaseMetaData, Testcase}, executors::{Executor, HasObservers}, + fuzzer::Evaluator, inputs::Input, mutators::Mutator, - observers::ObserversTuple, - stages::{Stage, MutationalStage}, - state::{Evaluator, HasCorpus, HasRand}, + observers::{MapObserver, ObserversTuple}, + stages::{MutationalStage, PowerScheduleMetadata, Stage}, + state::{HasClientPerfStats, HasCorpus, HasMetadata}, Error, }; -/// The mutational stage using power schedules -#[derive(Clone, Debug)] -pub struct PowerMutationalStage -where - M: Mutator, - I: Input, - S: HasCorpus + Evaluator + HasRand, - C: Corpus, - EM: EventManager, - E: Executor + HasObservers, - OT: ObserversTuple, - CS: CorpusScheduler, - R: Rand, -{ - mutator: M, - phantom: PhantomData<(C, CS, E, EM, I, OT, R, S)>, +#[derive(Clone, Debug, PartialEq)] +pub enum PowerSchedule { + EXPLORE, + FAST, + COE, + LIN, + QUAD, + EXPLOIT, } -impl MutationalStage - for PowerMutationalStage +const POWER_BETA: f64 = 1.0; +const MAX_FACTOR: f64 = POWER_BETA * 32.0; +const HAVOC_MAX_MULT: f64 = 64.0; + +/// The mutational stage using power schedules +#[derive(Clone, Debug)] +pub struct PowerMutationalStage where - M: Mutator, - I: Input, - S: HasCorpus + Evaluator + HasRand, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, C: Corpus, - EM: EventManager, - E: Executor + HasObservers, - OT: ObserversTuple, - CS: CorpusScheduler, - R: Rand, + E: Executor + HasObservers, + I: Input, + M: Mutator, + O: MapObserver, + OT: ObserversTuple, + S: HasClientPerfStats + HasCorpus + HasMetadata, + Z: Evaluator, +{ + map_observer_name: String, + mutator: M, + /// The employed power schedule strategy + strat: PowerSchedule, + #[allow(clippy::type_complexity)] + phantom: PhantomData<(C, E, EM, I, O, OT, S, T, Z)>, +} + +impl MutationalStage + for PowerMutationalStage +where + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + C: Corpus, + E: Executor + HasObservers, + I: Input, + M: Mutator, + O: MapObserver, + OT: ObserversTuple, + S: HasClientPerfStats + HasCorpus + HasMetadata, + Z: Evaluator, { /// The mutator, added to this stage #[inline] @@ -57,54 +78,319 @@ where } /// Gets the number of iterations as a random number - fn iterations(&self, state: &mut S) -> usize { - 1 + state.rand_mut().below(DEFAULT_MUTATIONAL_MAX_ITERATIONS) as usize - } -} + fn iterations(&self, state: &mut S, corpus_idx: usize) -> Result { + let mut testcase = state.corpus().get(corpus_idx)?.borrow_mut(); + let psmeta = state + .metadata() + .get::() + .ok_or_else(|| Error::KeyNotFound("PowerScheduleMetadata not found".to_string()))?; -impl Stage - for PowerMutationalStage -where - M: Mutator, - I: Input, - S: HasCorpus + Evaluator + HasRand, - C: Corpus, - EM: EventManager, - E: Executor + HasObservers, - OT: ObserversTuple, - CS: CorpusScheduler, - R: Rand, -{ - #[inline] - fn perform( - &self, - state: &mut S, + let mut fuzz_mu = 0.0; + if self.strat == PowerSchedule::COE { + fuzz_mu = self.fuzz_mu(state, psmeta)?; + } + + // 1 + state.rand_mut().below(DEFAULT_MUTATIONAL_MAX_ITERATIONS) as usize + self.calculate_score(&mut testcase, psmeta, fuzz_mu) + } + + #[allow(clippy::cast_possible_wrap)] + fn perform_mutational( + &mut self, + fuzzer: &mut Z, executor: &mut E, + state: &mut S, manager: &mut EM, - scheduler: &CS, corpus_idx: usize, ) -> Result<(), Error> { - self.perform_mutational(state, executor, manager, scheduler, corpus_idx) + let num = self.iterations(state, corpus_idx)?; + + for i in 0..num { + let mut input = state + .corpus() + .get(corpus_idx)? + .borrow_mut() + .load_input()? + .clone(); + + self.mutator_mut().mutate(state, &mut input, i as i32)?; + + let (_, corpus_idx) = fuzzer.evaluate_input(state, executor, manager, input)?; + + let observer = executor + .observers() + .match_name::(&self.map_observer_name) + .ok_or_else(|| Error::KeyNotFound("MapObserver not found".to_string()))?; + + let mut hash = observer.hash() as usize; + + let psmeta = state + .metadata_mut() + .get_mut::() + .ok_or_else(|| Error::KeyNotFound("PowerScheduleMetadata not found".to_string()))?; + + hash %= psmeta.n_fuzz().len(); + // Update the path frequency + psmeta.n_fuzz_mut()[hash] = psmeta.n_fuzz()[hash].saturating_add(1); + + if let Some(idx) = corpus_idx { + state + .corpus() + .get(idx)? + .borrow_mut() + .metadata_mut() + .get_mut::() + .ok_or_else(|| { + Error::KeyNotFound("PowerScheduleTestData not found".to_string()) + })? + .set_n_fuzz_entry(hash); + } + + self.mutator_mut().post_exec(state, i as i32, corpus_idx)?; + } + + Ok(()) } } -impl PowerMutationalStage +impl Stage + for PowerMutationalStage where - M: Mutator, - I: Input, - S: HasCorpus + Evaluator + HasRand, + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, C: Corpus, - EM: EventManager, - E: Executor + HasObservers, - OT: ObserversTuple, - CS: CorpusScheduler, - R: Rand, + E: Executor + HasObservers, + I: Input, + M: Mutator, + O: MapObserver, + OT: ObserversTuple, + S: HasClientPerfStats + HasCorpus + HasMetadata, + Z: Evaluator, { - /// Creates a new default mutational stage - pub fn new(mutator: M) -> Self { + #[inline] + #[allow(clippy::let_and_return)] + fn perform( + &mut self, + fuzzer: &mut Z, + executor: &mut E, + state: &mut S, + manager: &mut EM, + corpus_idx: usize, + ) -> Result<(), Error> { + let ret = self.perform_mutational(fuzzer, executor, state, manager, corpus_idx); + ret + } +} + +impl PowerMutationalStage +where + T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned, + C: Corpus, + E: Executor + HasObservers, + I: Input, + M: Mutator, + O: MapObserver, + OT: ObserversTuple, + S: HasClientPerfStats + HasCorpus + HasMetadata, + Z: Evaluator, +{ + pub fn new(mutator: M, strat: PowerSchedule, map_observer_name: &O) -> Self { Self { - mutator: mutator, + map_observer_name: map_observer_name.name().to_string(), + mutator, + strat, phantom: PhantomData, } } + + /// Compute the parameter `μ` used in the COE schedule. + #[inline] + pub fn fuzz_mu(&self, state: &S, psmeta: &PowerScheduleMetadata) -> Result { + let corpus = state.corpus(); + let mut n_paths = 0; + let mut fuzz_mu = 0.0; + for idx in 0..corpus.count() { + let n_fuzz_entry = corpus + .get(idx)? + .borrow() + .metadata() + .get::() + .ok_or_else(|| Error::KeyNotFound("PowerScheduleTestData not found".to_string()))? + .n_fuzz_entry(); + fuzz_mu += libm::log2(f64::from(psmeta.n_fuzz()[n_fuzz_entry])); + n_paths += 1; + } + + if n_paths == 0 { + return Err(Error::Unknown(String::from("Queue state corrput"))); + } + + fuzz_mu /= f64::from(n_paths); + Ok(fuzz_mu) + } + + /// Compute the `power` we assign to each corpus entry + #[inline] + #[allow( + clippy::cast_precision_loss, + clippy::too_many_lines, + clippy::cast_sign_loss + )] + fn calculate_score( + &self, + testcase: &mut Testcase, + psmeta: &PowerScheduleMetadata, + fuzz_mu: f64, + ) -> Result { + let mut perf_score = 100.0; + let q_exec_us = testcase + .exec_time() + .ok_or_else(|| Error::KeyNotFound("exec_time not set".to_string()))? + .as_nanos() as f64; + + let avg_exec_us = psmeta.exec_time().as_nanos() as f64 / psmeta.cycles() as f64; + let avg_bitmap_size = psmeta.bitmap_size() / psmeta.bitmap_entries(); + + let favored = testcase.has_metadata::(); + let tcmeta = testcase + .metadata_mut() + .get_mut::() + .ok_or_else(|| Error::KeyNotFound("PowerScheduleTestData not found".to_string()))?; + + if q_exec_us * 0.1 > avg_exec_us { + perf_score = 10.0; + } else if q_exec_us * 0.2 > avg_exec_us { + perf_score = 25.0; + } else if q_exec_us * 0.5 > avg_exec_us { + perf_score = 50.0; + } else if q_exec_us * 0.75 > avg_exec_us { + perf_score = 75.0; + } else if q_exec_us * 4.0 < avg_exec_us { + perf_score = 300.0; + } else if q_exec_us * 3.0 < avg_exec_us { + perf_score = 200.0; + } else if q_exec_us * 2.0 < avg_exec_us { + perf_score = 150.0; + } + + let q_bitmap_size = tcmeta.bitmap_size() as f64; + if q_bitmap_size * 0.3 > avg_bitmap_size as f64 { + perf_score *= 3.0; + } else if q_bitmap_size * 0.5 > avg_bitmap_size as f64 { + perf_score *= 2.0; + } else if q_bitmap_size * 0.75 > avg_bitmap_size as f64 { + perf_score *= 1.5; + } else if q_bitmap_size * 3.0 < avg_bitmap_size as f64 { + perf_score *= 0.25; + } else if q_bitmap_size * 2.0 < avg_bitmap_size as f64 { + perf_score *= 0.5; + } else if q_bitmap_size * 1.5 < avg_bitmap_size as f64 { + perf_score *= 0.75; + } + + if tcmeta.handicap() >= 4 { + perf_score *= 4.0; + tcmeta.set_handicap(tcmeta.handicap() - 4); + } else if tcmeta.handicap() > 0 { + perf_score *= 2.0; + tcmeta.set_handicap(tcmeta.handicap() - 1); + } + + if tcmeta.depth() >= 4 && tcmeta.depth() < 8 { + perf_score *= 2.0; + } else if tcmeta.depth() >= 8 && tcmeta.depth() < 14 { + perf_score *= 3.0; + } else if tcmeta.depth() >= 14 && tcmeta.depth() < 25 { + perf_score *= 4.0; + } else if tcmeta.depth() >= 25 { + perf_score *= 5.0; + } + + let mut factor: f64 = 1.0; + + // COE and Fast schedule are fairly different from what are described in the original thesis, + // This implementation follows the changes made in this pull request https://github.com/AFLplusplus/AFLplusplus/pull/568 + match &self.strat { + PowerSchedule::EXPLORE => { + // Nothing happens in EXPLORE + } + PowerSchedule::EXPLOIT => { + factor = MAX_FACTOR; + } + PowerSchedule::COE => { + if libm::log2(f64::from(psmeta.n_fuzz()[tcmeta.n_fuzz_entry()])) > fuzz_mu + && !favored + { + // Never skip favorites. + factor = 0.0; + } + } + PowerSchedule::FAST => { + if tcmeta.fuzz_level() != 0 { + let lg = libm::log2(f64::from(psmeta.n_fuzz()[tcmeta.n_fuzz_entry()])); + + match lg { + f if f < 2.0 => { + factor = 4.0; + } + f if (2.0..4.0).contains(&f) => { + factor = 3.0; + } + f if (4.0..5.0).contains(&f) => { + factor = 2.0; + } + f if (6.0..7.0).contains(&f) => { + if !favored { + factor = 0.8; + } + } + f if (7.0..8.0).contains(&f) => { + if !favored { + factor = 0.6; + } + } + f if f >= 8.0 => { + if !favored { + factor = 0.4 + } + } + _ => { + factor = 1.0; + } + } + + if favored { + factor *= 1.15; + } + } + } + PowerSchedule::LIN => { + factor = (tcmeta.fuzz_level() as f64) + / f64::from(psmeta.n_fuzz()[tcmeta.n_fuzz_entry()] + 1); + } + PowerSchedule::QUAD => { + factor = ((tcmeta.fuzz_level() * tcmeta.fuzz_level()) as f64) + / f64::from(psmeta.n_fuzz()[tcmeta.n_fuzz_entry()] + 1); + } + } + + if self.strat != PowerSchedule::EXPLORE { + if factor > MAX_FACTOR { + factor = MAX_FACTOR; + } + + perf_score *= factor / POWER_BETA; + } + + // Lower bound if the strat is not COE. + if self.strat == PowerSchedule::COE && perf_score < 1.0 { + perf_score = 1.0; + } + + // Upper bound + if perf_score > HAVOC_MAX_MULT * 100.0 { + perf_score = HAVOC_MAX_MULT * 100.0; + } + + Ok(perf_score as usize) + } }