From 849ff1fa049b70d4d391a762baa0760063e3a8bf Mon Sep 17 00:00:00 2001 From: Toka Date: Mon, 5 Jul 2021 20:54:15 +0900 Subject: [PATCH] MOpt scheduler (#161) * add the struct for MOpt globals * constants * RAND_C * more comments & reorder class members * select_algorithm * no_std fixes * clippy fixes * MOptMutator * MutatorsTuple has HasLen * MOptStage * pso_update * HasMOpt trait * ScheduledMutator, core_fuzzing * clippy fix * fmt * core_fuzzing * core_fuzzing done * fix * pilot_mutate * pilot_fuzzing * pilot_fuzzing done * MOpt metadata * Make MOptMutator into a trait * initialize_mopt * No getter/setters * fmt * fixed compiler warnings & clippy warnings * Comments * fix type paramter, integrate into libpng * fmt * fmt * No HasMOpt * fmt * improve * pso_initialize, various fixes * clippy * fmt * always pacemaker mode * fmt * fix * less noisy fmt::Debug Co-authored-by: Dominik Maier Co-authored-by: Andrea Fioraldi --- fuzzers/libfuzzer_libpng/src/lib.rs | 11 +- libafl/src/lib.rs | 3 + libafl/src/mutators/mod.rs | 2 + libafl/src/mutators/mopt_mutator.rs | 620 ++++++++++++++++++++++++++++ libafl/src/stages/mod.rs | 3 + libafl/src/stages/mopt.rs | 237 +++++++++++ 6 files changed, 874 insertions(+), 2 deletions(-) create mode 100644 libafl/src/mutators/mopt_mutator.rs create mode 100644 libafl/src/stages/mopt.rs diff --git a/fuzzers/libfuzzer_libpng/src/lib.rs b/fuzzers/libfuzzer_libpng/src/lib.rs index 4a192a4b5c..b36eaf09aa 100644 --- a/fuzzers/libfuzzer_libpng/src/lib.rs +++ b/fuzzers/libfuzzer_libpng/src/lib.rs @@ -17,10 +17,11 @@ use libafl::{ feedbacks::{CrashFeedback, MapFeedbackState, MaxMapFeedback, TimeFeedback, TimeoutFeedback}, fuzzer::{Fuzzer, StdFuzzer}, inputs::{BytesInput, HasTargetBytes}, - mutators::scheduled::{havoc_mutations, tokens_mutations, StdScheduledMutator}, + mutators::mopt_mutator::StdMOptMutator, + mutators::scheduled::{havoc_mutations, tokens_mutations}, mutators::token_mutations::Tokens, observers::{HitcountsMapObserver, StdMapObserver, TimeObserver}, - stages::mutational::StdMutationalStage, + stages::mopt::MOptStage, state::{HasCorpus, HasMetadata, StdState}, stats::MultiStats, Error, @@ -119,8 +120,14 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re } // Setup a basic mutator with a mutational stage + + let mutator = StdMOptMutator::new(havoc_mutations().merge(tokens_mutations())); + let mut stages = tuple_list!(MOptStage::new(mutator, &mut state, 5)?); + + /* let mutator = StdScheduledMutator::new(havoc_mutations().merge(tokens_mutations())); let mut stages = tuple_list!(StdMutationalStage::new(mutator)); + */ // A minimization+queue policy to get testcasess from the corpus let scheduler = IndexesLenTimeMinimizerCorpusScheduler::new(QueueCorpusScheduler::new()); diff --git a/libafl/src/lib.rs b/libafl/src/lib.rs index 63627511c0..0558648a14 100644 --- a/libafl/src/lib.rs +++ b/libafl/src/lib.rs @@ -72,6 +72,8 @@ pub enum Error { IllegalArgument(String), /// Forkserver related Error Forkserver(String), + /// MOpt related Error + MOpt(String), /// Shutting down, not really an error. ShuttingDown, /// Something else happened @@ -96,6 +98,7 @@ impl fmt::Display for Error { Self::IllegalState(s) => write!(f, "Illegal state: {0}", &s), Self::IllegalArgument(s) => write!(f, "Illegal argument: {0}", &s), Self::Forkserver(s) => write!(f, "Forkserver : {0}", &s), + Self::MOpt(s) => write!(f, "MOpt: {0}", &s), Self::ShuttingDown => write!(f, "Shutting down!"), Self::Unknown(s) => write!(f, "Unknown error: {0}", &s), } diff --git a/libafl/src/mutators/mod.rs b/libafl/src/mutators/mod.rs index b2e7911d15..ab4963c9d9 100644 --- a/libafl/src/mutators/mod.rs +++ b/libafl/src/mutators/mod.rs @@ -6,6 +6,8 @@ pub mod mutations; pub use mutations::*; pub mod token_mutations; pub use token_mutations::*; +pub mod mopt_mutator; +pub use mopt_mutator::*; use crate::{ bolts::tuples::{HasLen, Named}, diff --git a/libafl/src/mutators/mopt_mutator.rs b/libafl/src/mutators/mopt_mutator.rs new file mode 100644 index 0000000000..4ab5dfd98a --- /dev/null +++ b/libafl/src/mutators/mopt_mutator.rs @@ -0,0 +1,620 @@ +//! The `MOpt` mutator scheduler, see and +use alloc::{string::ToString, vec::Vec}; + +use crate::{ + bolts::{rands::Rand, rands::StdRand}, + inputs::Input, + mutators::{ComposedByMutations, MutationResult, Mutator, MutatorsTuple, ScheduledMutator}, + state::{HasMetadata, HasRand}, + Error, +}; +use core::{ + fmt::{self, Debug}, + marker::PhantomData, +}; +use serde::{Deserialize, Serialize}; + +/// A Struct for managing MOpt-mutator parameters +/// There are 2 modes for `MOpt` scheduler, the core fuzzing mode and the pilot fuzzing mode +/// In short, in the pilot fuzzing mode, the fuzzer employs several `swarms` to compute the probability to choose the mutation operator +/// On the other hand, in the core fuzzing mode, the fuzzer chooses the best `swarms`, which was determined during the pilot fuzzing mode, to compute the probability to choose the operation operator +/// With the current implementation we are always in the pacemaker fuzzing mode. +#[derive(Serialize, Deserialize, Clone)] +pub struct MOpt { + /// Random number generator + pub rand: StdRand, + /// The number of total findings (unique crashes and unique interesting paths). This is equivalent to `state.corpus().count() + state.solutions().count()`; + pub total_finds: usize, + /// The number of finds before switching to this mode. + pub finds_before_switch: usize, + /// The MOpt mode that we are currently using the pilot fuzzing mode or the core_fuzzing mode + pub key_module: MOptMode, + /// These w_* and g_* values are the coefficients for updating variables according to the PSO algorithms + pub w_init: f64, + pub w_end: f64, + pub w_now: f64, + pub g_now: i32, + pub g_max: i32, + /// The number of mutation operators + pub operator_num: usize, + /// The number of swarms that we want to employ during the pilot fuzzing mode + pub swarm_num: usize, + /// We'll generate testcases for `period_pilot` times before we call pso_update in core fuzzing module + pub period_pilot: usize, + /// We'll generate testcases for `period_core` times before we call pso_update in core fuzzing module + pub period_core: usize, + /// The number of testcases generated during this pilot fuzzing mode + pub pilot_time: usize, + /// The number of testcases generated during this core fuzzing mode + pub core_time: usize, + /// The swarm identifier that we are currently using in the pilot fuzzing mode + pub swarm_now: usize, + /// These are the parameters for the PSO algorithm + x_now: Vec>, + l_best: Vec>, + eff_best: Vec>, + g_best: Vec, + v_now: Vec>, + /// The probability that we want to use to choose the mutation operator. + probability_now: Vec>, + /// The fitness for each swarm, we'll calculate the fitness in the pilot fuzzing mode and use the best one in the core fuzzing mode + pub swarm_fitness: Vec, + /// (Pilot Mode) Finds by each operators. This vector is used in pso_update + pub pilot_operator_finds_pso: Vec>, + /// (Pilot Mode) Finds by each operator till now. + pub pilot_operator_finds_this: Vec>, + /// (Pilot Mode) The number of mutation operator used. This vector is used in pso_update + pub pilot_operator_ctr_pso: Vec>, + /// (Pilot Mode) The number of mutation operator used till now + pub pilot_operator_ctr_this: Vec>, + /// (Pilot Mode) The number of mutation operator used till last execution + pub pilot_operator_ctr_last: Vec>, + /// Vector used in pso_update + pub operator_finds_puppet: Vec, + /// (Core Mode) Finds by each operators. This vector is used in pso_update + pub core_operator_finds_pso: Vec, + /// (Core Mode) Finds by each operator till now. + pub core_operator_finds_this: Vec, + /// (Core Mode) The number of mutation operator used. This vector is used in pso_update + pub core_operator_ctr_pso: Vec, + /// (Core Mode) The number of mutation operator used till now + pub core_operator_ctr_this: Vec, + /// (Core Mode) The number of mutation operator used till last execution + pub core_operator_ctr_last: Vec, +} + +crate::impl_serdeany!(MOpt); + +impl fmt::Debug for MOpt { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MOpt") + .field("\ntotal_finds", &self.total_finds) + .field("\nfinds_before_switch", &self.finds_before_switch) + .field("\nkey_module", &self.key_module) + .field("\nw_init", &self.w_init) + .field("\nw_end", &self.w_end) + .field("\nw_now", &self.g_now) + .field("\ng_now", &self.g_max) + .field("\npilot_time", &self.pilot_time) + .field("\ncore_time", &self.core_time) + .field("\n\nx_now", &self.x_now) + .field("\n\nl_best", &self.l_best) + .field("\n\neff_best", &self.eff_best) + .field("\n\ng_best", &self.g_best) + .field("\n\nv_now", &self.v_now) + .field("\n\nprobability_now", &self.probability_now) + .field("\n\nswarm_fitness", &self.swarm_fitness) + .field( + "\n\npilot_operator_finds_pso", + &self.pilot_operator_finds_pso, + ) + .field( + "\n\npilot_operator_finds_this", + &self.pilot_operator_finds_this, + ) + .field("\n\npilot_operator_ctr_pso", &self.pilot_operator_ctr_pso) + .field("\n\npilot_operator_ctr_this", &self.pilot_operator_ctr_this) + .field("\n\npilot_operator_ctr_last", &self.pilot_operator_ctr_last) + .field("\n\noperator_finds_puppuet", &self.operator_finds_puppet) + .field("\n\ncore_operator_finds_pso", &self.core_operator_finds_pso) + .field( + "\n\ncore_operator_finds_this", + &self.core_operator_finds_this, + ) + .field("\n\ncore_operator_ctr_pso", &self.core_operator_ctr_pso) + .field("\n\ncore_operator_ctr_this", &self.core_operator_ctr_this) + .field("\n\ncore_operator_ctr_last", &self.core_operator_ctr_last) + .finish() + } +} + +impl MOpt { + pub fn new(operator_num: usize, swarm_num: usize) -> Result { + let mut mopt = Self { + rand: StdRand::with_seed(0), + total_finds: 0, + finds_before_switch: 0, + key_module: MOptMode::Pilotfuzzing, + w_init: 0.9, + w_end: 0.3, + w_now: 0.0, + g_now: 0, + g_max: 5000, + operator_num, + swarm_num, + period_pilot: 50000, + period_core: 500000, + pilot_time: 0, + core_time: 0, + swarm_now: 0, + x_now: vec![vec![0.0; operator_num]; swarm_num], + l_best: vec![vec![0.0; operator_num]; swarm_num], + eff_best: vec![vec![0.0; operator_num]; swarm_num], + g_best: vec![0.0; operator_num], + v_now: vec![vec![0.0; operator_num]; swarm_num], + probability_now: vec![vec![0.0; operator_num]; swarm_num], + swarm_fitness: vec![0.0; swarm_num], + pilot_operator_finds_pso: vec![vec![0; operator_num]; swarm_num], + pilot_operator_finds_this: vec![vec![0; operator_num]; swarm_num], + pilot_operator_ctr_pso: vec![vec![0; operator_num]; swarm_num], + pilot_operator_ctr_this: vec![vec![0; operator_num]; swarm_num], + pilot_operator_ctr_last: vec![vec![0; operator_num]; swarm_num], + operator_finds_puppet: vec![0; operator_num], + core_operator_finds_pso: vec![0; operator_num], + core_operator_finds_this: vec![0; operator_num], + core_operator_ctr_pso: vec![0; operator_num], + core_operator_ctr_this: vec![0; operator_num], + core_operator_ctr_last: vec![0; operator_num], + }; + mopt.pso_initialize()?; + Ok(mopt) + } + + /// Initialize `core_operator_*` values + pub fn init_core_module(&mut self) -> Result<(), Error> { + for i in 0..self.operator_num { + self.core_operator_ctr_this[i] = self.core_operator_ctr_pso[i]; + self.core_operator_ctr_last[i] = self.core_operator_ctr_pso[i]; + self.core_operator_finds_this[i] = self.core_operator_finds_pso[i] + } + + let mut swarm_eff = 0.0; + let mut best_swarm = 0; + for i in 0..self.swarm_num { + if self.swarm_fitness[i] > swarm_eff { + swarm_eff = self.swarm_fitness[i]; + best_swarm = i; + } + } + + self.swarm_now = best_swarm; + Ok(()) + } + + #[inline] + pub fn update_pilot_operator_ctr_last(&mut self, swarm_now: usize) { + for i in 0..self.operator_num { + self.pilot_operator_ctr_last[swarm_now][i] = self.pilot_operator_ctr_this[swarm_now][i] + } + } + + #[inline] + pub fn update_core_operator_ctr_last(&mut self) { + for i in 0..self.operator_num { + self.core_operator_ctr_last[i] = self.core_operator_ctr_this[i]; + } + } + + /// Finds the local optimum for each operator + /// See + + #[allow(clippy::cast_precision_loss)] + pub fn update_pilot_operator_ctr_pso(&mut self, swarm_now: usize) { + let mut eff = 0.0; + for i in 0..self.operator_num { + if self.pilot_operator_ctr_this[swarm_now][i] + > self.pilot_operator_ctr_pso[swarm_now][i] + { + eff = ((self.pilot_operator_finds_this[swarm_now][i] + - self.pilot_operator_finds_pso[swarm_now][i]) as f64) + / ((self.pilot_operator_ctr_this[swarm_now][i] + - self.pilot_operator_ctr_pso[swarm_now][i]) as f64) + } + + if self.eff_best[swarm_now][i] < eff { + self.eff_best[swarm_now][i] = eff; + self.l_best[swarm_now][i] = self.x_now[swarm_now][i]; + } + + self.pilot_operator_finds_pso[swarm_now][i] = + self.pilot_operator_finds_this[swarm_now][i]; + self.pilot_operator_ctr_pso[swarm_now][i] = self.pilot_operator_ctr_this[swarm_now][i]; + } + } + + #[inline] + pub fn update_core_operator_ctr_pso(&mut self) { + for i in 0..self.operator_num { + self.core_operator_finds_pso[i] = self.core_operator_finds_this[i]; + self.core_operator_ctr_pso[i] = self.core_operator_ctr_this[i]; + } + } + + #[allow(clippy::cast_precision_loss)] + pub fn pso_initialize(&mut self) -> Result<(), Error> { + if self.g_now > self.g_max { + self.g_now = 0; + } + self.w_now = (self.w_init - self.w_end) * f64::from(self.g_max - self.g_now) + / f64::from(self.g_max) + + self.w_end; + + for swarm in 0..self.swarm_num { + let mut total_x_now = 0.0; + let mut x_sum = 0.0; + for i in 0..self.operator_num { + self.x_now[swarm][i] = (self.rand.below(7000) as f64) * 0.0001 + 0.1; + total_x_now += self.x_now[swarm][i]; + self.v_now[swarm][i] = 0.1; + self.l_best[swarm][i] = 0.5; + self.g_best[i] = 0.5; + } + + for i in 0..self.operator_num { + self.x_now[swarm][i] /= total_x_now + } + + for i in 0..self.operator_num { + self.v_now[swarm][i] = self.w_now * self.v_now[swarm][i] + + (self.rand.below(1000) as f64) + * 0.001 + * (self.l_best[swarm][i] - self.x_now[swarm][i]) + + (self.rand.below(1000) as f64) + * 0.001 + * (self.g_best[i] - self.x_now[swarm][i]); + self.x_now[swarm][i] += self.v_now[swarm][i]; + + if self.x_now[swarm][i] > V_MAX { + self.x_now[swarm][i] = V_MAX; + } else if self.x_now[swarm][i] < V_MIN { + self.x_now[swarm][i] = V_MIN; + } + + x_sum += self.x_now[swarm][i] + } + + for i in 0..self.operator_num { + self.x_now[swarm][i] /= x_sum; + if i == 0 { + self.probability_now[swarm][i] = self.x_now[swarm][i]; + } else { + self.probability_now[swarm][i] = + self.probability_now[swarm][i - 1] + self.x_now[swarm][i]; + } + } + if self.probability_now[swarm][self.operator_num - 1] < 0.99 + || self.probability_now[swarm][self.operator_num - 1] > 1.01 + { + return Err(Error::MOpt("Error in pso_update".to_string())); + } + } + Ok(()) + } + + /// Update the PSO algorithm parameters + /// See + #[allow(clippy::cast_precision_loss)] + pub fn pso_update(&mut self) -> Result<(), Error> { + self.g_now += 1; + if self.g_now > self.g_max { + self.g_now = 0; + } + self.w_now = (self.w_init - self.w_end) * f64::from(self.g_max - self.g_now) + / f64::from(self.g_max) + + self.w_end; + + let mut operator_find_sum = 0; + + for i in 0..self.operator_num { + self.operator_finds_puppet[i] = self.core_operator_ctr_pso[i]; + + for j in 0..self.swarm_num { + self.operator_finds_puppet[i] += self.pilot_operator_finds_pso[j][i]; + } + operator_find_sum += self.operator_finds_puppet[i]; + } + + for i in 0..self.operator_num { + if self.operator_finds_puppet[i] > 0 { + self.g_best[i] = + (self.operator_finds_puppet[i] as f64) / (operator_find_sum as f64); + } + } + + for swarm in 0..self.swarm_num { + let mut x_sum = 0.0; + for i in 0..self.operator_num { + self.probability_now[swarm][i] = 0.0; + self.v_now[swarm][i] = self.w_now * self.v_now[swarm][i] + + (self.rand.below(1000) as f64) + * 0.001 + * (self.l_best[swarm][i] - self.x_now[swarm][i]) + + (self.rand.below(1000) as f64) + * 0.001 + * (self.g_best[i] - self.x_now[swarm][i]); + self.x_now[swarm][i] += self.v_now[swarm][i]; + + if self.x_now[swarm][i] > V_MAX { + self.x_now[swarm][i] = V_MAX; + } else if self.x_now[swarm][i] < V_MIN { + self.x_now[swarm][i] = V_MIN; + } + x_sum += self.x_now[swarm][i]; + } + + for i in 0..self.operator_num { + self.x_now[swarm][i] /= x_sum; + if i == 0 { + self.probability_now[swarm][i] = self.x_now[swarm][i]; + } else { + self.probability_now[swarm][i] = + self.probability_now[swarm][i - 1] + self.x_now[swarm][i]; + } + } + if self.probability_now[swarm][self.operator_num - 1] < 0.99 + || self.probability_now[swarm][self.operator_num - 1] > 1.01 + { + return Err(Error::MOpt("Error in pso_update".to_string())); + } + } + self.swarm_now = 0; + + self.key_module = MOptMode::Pilotfuzzing; + //println!("Mopt struct:\n{:?}", self); + Ok(()) + } + + /// This function is used to decide the operator that we want to apply next + /// see + #[allow(clippy::cast_precision_loss)] + pub fn select_algorithm(&mut self) -> Result { + let mut res = 0; + let mut sentry = 0; + + let operator_num = self.operator_num; + + // Fetch a random sele value + let select_prob: f64 = self.probability_now[self.swarm_now][operator_num - 1] + * ((self.rand.below(10000) as f64) * 0.0001); + + for i in 0..operator_num { + if i == 0 { + if select_prob < self.probability_now[self.swarm_now][i] { + res = i; + break; + } + } else if select_prob < self.probability_now[self.swarm_now][i] { + res = i; + sentry = 1; + break; + } + } + + if (sentry == 1 && select_prob < self.probability_now[self.swarm_now][res - 1]) + || (res + 1 < operator_num + && select_prob > self.probability_now[self.swarm_now][res + 1]) + { + return Err(Error::MOpt("Error in select_algorithm".to_string())); + } + Ok(res) + } +} + +const V_MAX: f64 = 1.0; +const V_MIN: f64 = 0.05; + +#[derive(Serialize, Deserialize, Clone, Copy, Debug)] +pub enum MOptMode { + Pilotfuzzing, + Corefuzzing, +} + +pub struct StdMOptMutator +where + I: Input, + MT: MutatorsTuple, + R: Rand, + S: HasRand + HasMetadata, +{ + mutations: MT, + phantom: PhantomData<(I, R, S)>, +} + +impl Debug for StdMOptMutator +where + I: Input, + MT: MutatorsTuple, + R: Rand, + S: HasRand + HasMetadata, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "StdMOptMutator with {} mutations for Input type {}", + self.mutations.len(), + core::any::type_name::() + ) + } +} + +impl Mutator for StdMOptMutator +where + I: Input, + MT: MutatorsTuple, + R: Rand, + S: HasRand + HasMetadata, +{ + #[inline] + fn mutate( + &mut self, + state: &mut S, + input: &mut I, + stage_idx: i32, + ) -> Result { + self.scheduled_mutate(state, input, stage_idx) + } +} + +impl StdMOptMutator +where + I: Input, + MT: MutatorsTuple, + R: Rand, + S: HasRand + HasMetadata, +{ + pub fn new(mutations: MT) -> Self { + Self { + mutations, + phantom: PhantomData, + } + } + fn core_mutate( + &mut self, + state: &mut S, + input: &mut I, + stage_idx: i32, + ) -> Result { + // TODO + let mut r = MutationResult::Skipped; + state + .metadata_mut() + .get_mut::() + .unwrap() + .update_core_operator_ctr_last(); + + for _i in 0..self.iterations(state, input) { + let idx = self.schedule(state, input); + let outcome = self + .mutations_mut() + .get_and_mutate(idx, state, input, stage_idx)?; + if outcome == MutationResult::Mutated { + r = MutationResult::Mutated; + } + + state + .metadata_mut() + .get_mut::() + .unwrap() + .core_operator_ctr_this[idx] += 1; + } + + Ok(r) + } + + fn pilot_mutate( + &mut self, + state: &mut S, + input: &mut I, + stage_idx: i32, + ) -> Result { + let mut r = MutationResult::Skipped; + let swarm_now; + { + let mopt = state.metadata_mut().get_mut::().unwrap(); + swarm_now = mopt.swarm_now; + mopt.update_pilot_operator_ctr_last(swarm_now); + } + + for _i in 0..self.iterations(state, input) { + let idx = self.schedule(state, input); + let outcome = self + .mutations_mut() + .get_and_mutate(idx, state, input, stage_idx)?; + if outcome == MutationResult::Mutated { + r = MutationResult::Mutated; + } + + state + .metadata_mut() + .get_mut::() + .unwrap() + .pilot_operator_ctr_this[swarm_now][idx] += 1; + } + + Ok(r) + } +} + +impl ComposedByMutations for StdMOptMutator +where + I: Input, + MT: MutatorsTuple, + R: Rand, + S: HasRand + HasMetadata, +{ + /// Get the mutations + #[inline] + fn mutations(&self) -> &MT { + &self.mutations + } + + // Get the mutations (mut) + #[inline] + fn mutations_mut(&mut self) -> &mut MT { + &mut self.mutations + } +} + +impl ScheduledMutator for StdMOptMutator +where + I: Input, + MT: MutatorsTuple, + R: Rand, + S: HasRand + HasMetadata, +{ + /// Compute the number of iterations used to apply stacked mutations + fn iterations(&self, state: &mut S, _: &I) -> u64 { + 1 << (1 + state.rand_mut().below(6)) + } + + /// Get the next mutation to apply + fn schedule(&self, state: &mut S, _: &I) -> usize { + state + .metadata_mut() + .get_mut::() + .unwrap() + .select_algorithm() + .unwrap() + } + + fn scheduled_mutate( + &mut self, + state: &mut S, + input: &mut I, + stage_idx: i32, + ) -> Result { + let mode = state.metadata().get::().unwrap().key_module; + match mode { + MOptMode::Corefuzzing => self.core_mutate(state, input, stage_idx), + MOptMode::Pilotfuzzing => self.pilot_mutate(state, input, stage_idx), + } + } +} + +pub trait MOptMutator: ScheduledMutator +where + I: Input, + MT: MutatorsTuple, + R: Rand, + S: HasRand + HasMetadata, +{ +} + +impl MOptMutator for StdMOptMutator +where + I: Input, + MT: MutatorsTuple, + R: Rand, + S: HasRand + HasMetadata, +{ +} diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs index 5042d8e680..d1cc962f73 100644 --- a/libafl/src/stages/mod.rs +++ b/libafl/src/stages/mod.rs @@ -11,6 +11,9 @@ pub use mutational::{MutationalStage, StdMutationalStage}; pub mod tracing; pub use tracing::{ShadowTracingStage, TracingStage}; +pub mod mopt; +pub use mopt::*; + //pub mod power; //pub use power::PowerMutationalStage; use crate::Error; diff --git a/libafl/src/stages/mopt.rs b/libafl/src/stages/mopt.rs new file mode 100644 index 0000000000..3abde341e5 --- /dev/null +++ b/libafl/src/stages/mopt.rs @@ -0,0 +1,237 @@ +use core::marker::PhantomData; + +use crate::{ + bolts::rands::Rand, + corpus::Corpus, + fuzzer::Evaluator, + inputs::Input, + mutators::{MOpt, MOptMode, MOptMutator, MutatorsTuple}, + stages::{MutationalStage, Stage}, + state::{HasClientPerfStats, HasCorpus, HasMetadata, HasRand, HasSolutions}, + Error, +}; + +const PERIOD_PILOT_COEF: f64 = 5000.0; + +#[derive(Clone, Debug)] +pub struct MOptStage +where + C: Corpus, + M: MOptMutator, + MT: MutatorsTuple, + I: Input, + R: Rand, + S: HasClientPerfStats + HasCorpus + HasSolutions + HasRand + HasMetadata, + SC: Corpus, + Z: Evaluator, +{ + mutator: M, + #[allow(clippy::type_complexity)] + phantom: PhantomData<(C, E, EM, I, MT, R, S, SC, Z)>, +} + +impl MutationalStage + for MOptStage +where + C: Corpus, + M: MOptMutator, + MT: MutatorsTuple, + I: Input, + R: Rand, + S: HasClientPerfStats + HasCorpus + HasSolutions + HasRand + HasMetadata, + SC: Corpus, + Z: Evaluator, +{ + /// The mutator, added to this stage + #[inline] + fn mutator(&self) -> &M { + &self.mutator + } + + /// The list of mutators, added to this stage (as mutable ref) + #[inline] + fn mutator_mut(&mut self) -> &mut M { + &mut self.mutator + } + + /// Gets the number of iterations as a random number + fn iterations(&self, state: &mut S) -> usize { + // TODO: we want to use calculate_score here + + 1 + state.rand_mut().below(128) as usize + } + + #[allow( + clippy::cast_possible_wrap, + clippy::cast_precision_loss, + clippy::too_many_lines + )] + fn perform_mutational( + &mut self, + fuzzer: &mut Z, + executor: &mut E, + state: &mut S, + manager: &mut EM, + corpus_idx: usize, + ) -> Result<(), Error> { + let key_module = state.metadata().get::().unwrap().key_module; + + match key_module { + MOptMode::Corefuzzing => { + let num = self.iterations(state); + + for stage_id in 0..num { + let mut input = state + .corpus() + .get(corpus_idx)? + .borrow_mut() + .load_input()? + .clone(); + + self.mutator_mut() + .mutate(state, &mut input, stage_id as i32)?; + + let finds_before = state.corpus().count() + state.solutions().count(); + + let (_, corpus_idx) = fuzzer.evaluate_input(state, executor, manager, input)?; + + self.mutator_mut() + .post_exec(state, stage_id as i32, corpus_idx)?; + + let finds_after = state.corpus().count() + state.solutions().count(); + + let mopt = state.metadata_mut().get_mut::().unwrap(); + + mopt.core_time += 1; + + if finds_after > finds_before { + let diff = finds_after - finds_before; + mopt.total_finds += diff; + for i in 0..mopt.operator_num { + if mopt.core_operator_ctr_this[i] > mopt.core_operator_ctr_last[i] { + mopt.core_operator_finds_this[i] += diff; + } + } + } + + if mopt.core_time > mopt.period_core { + // Make a call to pso_update() + mopt.core_time = 0; + let total_finds = mopt.total_finds; + mopt.finds_before_switch = total_finds; + mopt.update_core_operator_ctr_pso(); + mopt.pso_update()?; + } + } + } + MOptMode::Pilotfuzzing => { + let num = self.iterations(state); + for stage_id in 0..num { + let mut input = state + .corpus() + .get(corpus_idx)? + .borrow_mut() + .load_input()? + .clone(); + + self.mutator_mut() + .mutate(state, &mut input, stage_id as i32)?; + + let finds_before = state.corpus().count() + state.solutions().count(); + + let (_, corpus_idx) = fuzzer.evaluate_input(state, executor, manager, input)?; + + self.mutator_mut() + .post_exec(state, stage_id as i32, corpus_idx)?; + + let finds_after = state.corpus().count() + state.solutions().count(); + + let mopt = state.metadata_mut().get_mut::().unwrap(); + + mopt.pilot_time += 1; + let swarm_now = mopt.swarm_now; + + if finds_after > finds_before { + let diff = finds_after - finds_before; + mopt.total_finds += diff; + for i in 0..mopt.operator_num { + if mopt.pilot_operator_ctr_this[swarm_now][i] + > mopt.pilot_operator_ctr_last[swarm_now][i] + { + mopt.pilot_operator_finds_this[swarm_now][i] += diff; + } + } + } + + if mopt.pilot_time > mopt.period_pilot { + let new_finds = mopt.total_finds - mopt.finds_before_switch; + let f = + (new_finds as f64) / ((mopt.pilot_time as f64) / (PERIOD_PILOT_COEF)); + mopt.swarm_fitness[swarm_now] = f; + mopt.pilot_time = 0; + let total_finds = mopt.total_finds; + mopt.finds_before_switch = total_finds; + mopt.update_pilot_operator_ctr_pso(swarm_now); + + mopt.swarm_now += 1; + + if mopt.swarm_now == mopt.swarm_num { + // Move to CORE_FUZING mode + mopt.key_module = MOptMode::Corefuzzing; + + mopt.init_core_module()?; + } + } + } + } + } + + Ok(()) + } +} + +impl Stage + for MOptStage +where + C: Corpus, + M: MOptMutator, + MT: MutatorsTuple, + I: Input, + R: Rand, + S: HasClientPerfStats + HasCorpus + HasSolutions + HasRand + HasMetadata, + SC: Corpus, + Z: Evaluator, +{ + #[inline] + fn perform( + &mut self, + fuzzer: &mut Z, + executor: &mut E, + state: &mut S, + manager: &mut EM, + corpus_idx: usize, + ) -> Result<(), Error> { + self.perform_mutational(fuzzer, executor, state, manager, corpus_idx) + } +} + +impl MOptStage +where + C: Corpus, + M: MOptMutator, + MT: MutatorsTuple, + I: Input, + R: Rand, + S: HasClientPerfStats + HasCorpus + HasSolutions + HasRand + HasMetadata, + SC: Corpus, + Z: Evaluator, +{ + /// Creates a new default mutational stage + pub fn new(mutator: M, state: &mut S, swarm_num: usize) -> Result { + state.add_metadata::(MOpt::new(mutator.mutations().len(), swarm_num)?); + Ok(Self { + mutator, + phantom: PhantomData, + }) + } +}