diff --git a/src/corpus/mod.rs b/src/corpus/mod.rs index 8d9cfc2c7b..92bcddd3b0 100644 --- a/src/corpus/mod.rs +++ b/src/corpus/mod.rs @@ -13,7 +13,7 @@ use crate::inputs::Input; use crate::utils::Rand; use crate::AflError; -pub trait HasEntriesVec +pub trait HasTestcaseVec where I: Input, { @@ -25,7 +25,7 @@ where } /// Corpus with all current testcases -pub trait Corpus: HasEntriesVec +pub trait Corpus: HasTestcaseVec where I: Input, R: Rand, @@ -76,6 +76,19 @@ where } } +pub trait HasCorpus +where + C: Corpus, + I: Input, + R: Rand, +{ + /// Get the corpus field + fn corpus(&self) -> &C; + + /// Get thecorpus field (mutable) + fn corpus_mut(&mut self) -> &mut C; +} + pub struct InMemoryCorpus where I: Input, @@ -85,7 +98,7 @@ where phantom: PhantomData, } -impl HasEntriesVec for InMemoryCorpus +impl HasTestcaseVec for InMemoryCorpus where I: Input, R: Rand, @@ -131,7 +144,7 @@ where } #[cfg(feature = "std")] -impl HasEntriesVec for OnDiskCorpus +impl HasTestcaseVec for OnDiskCorpus where I: Input, R: Rand, @@ -192,7 +205,7 @@ where cycles: u64, } -impl HasEntriesVec for QueueCorpus +impl HasTestcaseVec for QueueCorpus where C: Corpus, I: Input, diff --git a/src/engines/mod.rs b/src/engines/mod.rs index e018bd9688..a861a15fb0 100644 --- a/src/engines/mod.rs +++ b/src/engines/mod.rs @@ -8,7 +8,7 @@ use core::fmt::Debug; use hashbrown::HashMap; -use crate::corpus::{Corpus, Testcase}; +use crate::corpus::{Corpus, HasCorpus, Testcase}; use crate::events::EventManager; use crate::executors::Executor; use crate::feedbacks::Feedback; @@ -18,16 +18,17 @@ use crate::stages::Stage; use crate::utils::{HasRand, Rand}; use crate::AflError; +// TODO FeedbackMetadata to store histroy_map + pub trait StateMetadata: Debug { /// The name of this metadata - used to find it in the list of avaliable metadatas fn name(&self) -> &'static str; } -pub trait State: HasRand +pub trait State: HasCorpus where C: Corpus, E: Executor, - EM: EventManager, I: Input, R: Rand, { @@ -37,10 +38,6 @@ where /// Set executions fn set_executions(&mut self, executions: usize); - fn events_manager(&self) -> &EM; - - fn events_manager_mut(&mut self) -> &mut EM; - /// Get all the metadatas into an HashMap fn metadatas(&self) -> &HashMap<&'static str, Box>; @@ -90,12 +87,6 @@ where self.feedbacks_mut().push(feedback); } - /// Return the corpus - fn corpus(&self) -> &C; - - /// Return the corpus (mutable) - fn corpus_mut(&mut self) -> &mut C; - /// Return the executor fn executor(&self) -> &E; @@ -129,47 +120,43 @@ where } } -pub struct DefaultState +pub struct DefaultState where C: Corpus, E: Executor, - EM: EventManager, I: Input, R: Rand, { - rand: R, executions: usize, - events_manager: EM, metadatas: HashMap<&'static str, Box>, + // additional_corpuses: HashMap<&'static str, Box>, observers: Vec>>, feedbacks: Vec>>, corpus: C, executor: E, } -impl HasRand for DefaultState +impl HasCorpus for DefaultState where C: Corpus, E: Executor, - EM: EventManager, I: Input, R: Rand, { - type R = R; - - fn rand(&self) -> &Self::R { - &self.rand + fn corpus(&self) -> &C { + &self.corpus } - fn rand_mut(&mut self) -> &mut Self::R { - &mut self.rand + + /// Get thecorpus field (mutable) + fn corpus_mut(&mut self) -> &mut C { + &mut self.corpus } } -impl State for DefaultState +impl State for DefaultState where C: Corpus, E: Executor, - EM: EventManager, I: Input, R: Rand, { @@ -181,13 +168,6 @@ where self.executions = executions } - fn events_manager(&self) -> &EM { - &self.events_manager - } - fn events_manager_mut(&mut self) -> &mut EM { - &mut self.events_manager - } - fn metadatas(&self) -> &HashMap<&'static str, Box> { &self.metadatas } @@ -212,14 +192,6 @@ where &mut self.feedbacks } - fn corpus(&self) -> &C { - &self.corpus - } - - fn corpus_mut(&mut self) -> &mut C { - &mut self.corpus - } - fn executor(&self) -> &E { &self.executor } @@ -229,19 +201,16 @@ where } } -impl DefaultState +impl DefaultState where C: Corpus, E: Executor, - EM: EventManager, I: Input, R: Rand, { - pub fn new(corpus: C, executor: E, events_manager: EM, rand: R) -> Self { + pub fn new(corpus: C, executor: E) -> Self { DefaultState { - rand: rand, executions: 0, - events_manager: events_manager, metadatas: HashMap::default(), observers: vec![], feedbacks: vec![], @@ -253,23 +222,28 @@ where pub trait Engine where - S: State, + S: State, C: Corpus, E: Executor, EM: EventManager, I: Input, R: Rand, { - fn stages(&self) -> &[Box>]; + fn stages(&self) -> &[Box>]; - fn stages_mut(&mut self) -> &mut Vec>>; + fn stages_mut(&mut self) -> &mut Vec>>; - fn add_stage(&mut self, stage: Box>) { + fn add_stage(&mut self, stage: Box>) { self.stages_mut().push(stage); } - fn fuzz_one(&mut self, state: &mut S) -> Result { - let (testcase, idx) = state.corpus_mut().next(state.rand_mut())?; + fn fuzz_one( + &mut self, + rand: &mut R, + state: &mut S, + events_manager: &mut EM, + ) -> Result { + let (testcase, idx) = state.corpus_mut().next(rand)?; println!("Cur entry: {}\tExecutions: {}", idx, state.executions()); for stage in self.stages_mut() { stage.perform(testcase.clone(), state)?; @@ -287,30 +261,30 @@ where I: Input, R: Rand, { - stages: Vec>>, + stages: Vec>>, } impl Engine for DefaultEngine where - S: State, + S: State, C: Corpus, E: Executor, EM: EventManager, I: Input, R: Rand, { - fn stages(&self) -> &[Box>] { + fn stages(&self) -> &[Box>] { &self.stages } - fn stages_mut(&mut self) -> &mut Vec>> { + fn stages_mut(&mut self) -> &mut Vec>> { &mut self.stages } } impl DefaultEngine where - S: State, + S: State, C: Corpus, E: Executor, EM: EventManager, diff --git a/src/mutators/mod.rs b/src/mutators/mod.rs index 0c0ecdb9ec..28ee775e27 100644 --- a/src/mutators/mod.rs +++ b/src/mutators/mod.rs @@ -12,7 +12,7 @@ use crate::AflError; pub trait Mutator where - S: HasRand + HasCorpus, + S: HasRand + HasCorpus, C: Corpus, I: Input, R: Rand, diff --git a/src/mutators/scheduled.rs b/src/mutators/scheduled.rs index 1c5460aa17..5046307eb0 100644 --- a/src/mutators/scheduled.rs +++ b/src/mutators/scheduled.rs @@ -18,7 +18,7 @@ type MutationFunction = fn(&mut M, &mut S, &mut I) -> Result where - S: HasRand + HasCorpus, + S: HasRand + HasCorpus, C: Corpus, I: Input, R: Rand, @@ -36,7 +36,7 @@ where pub trait ScheduledMutator: Mutator + ComposedByMutations where - S: HasRand + HasCorpus, + S: HasRand + HasCorpus, C: Corpus, I: Input, R: Rand, @@ -90,23 +90,19 @@ where impl Mutator for DefaultScheduledMutator where + S: HasRand + HasCorpus, C: Corpus, I: Input, R: Rand, { - fn mutate( - &mut self, - corpus: &mut C, - rand: &mut R, - input: &mut I, - _stage_idx: i32, - ) -> Result<(), AflError> { - self.scheduled_mutate(corpus, rand, input, _stage_idx) + fn mutate(&mut self, state: &mut S, input: &mut I, _stage_idx: i32) -> Result<(), AflError> { + self.scheduled_mutate(state, input, _stage_idx) } } impl ComposedByMutations for DefaultScheduledMutator where + S: HasRand + HasCorpus, C: Corpus, I: Input, R: Rand, @@ -127,8 +123,9 @@ where } } -impl ScheduledMutator for DefaultScheduledMutator +impl ScheduledMutator for DefaultScheduledMutator where + S: HasRand + HasCorpus, C: Corpus, I: Input, R: Rand, @@ -136,8 +133,9 @@ where // Just use the default methods } -impl DefaultScheduledMutator +impl DefaultScheduledMutator where + S: HasRand + HasCorpus, C: Corpus, I: Input, R: Rand, @@ -156,19 +154,19 @@ where } /// Bitflip mutation for inputs with a bytes vector -pub fn mutation_bitflip( +pub fn mutation_bitflip( mutator: &mut M, - _corpus: &mut C, - rand: &mut R, + state: &mut S, input: &mut I, ) -> Result where - M: Mutator, + M: Mutator, + S: HasRand + HasCorpus, C: Corpus, I: Input + HasBytesVec, R: Rand, { - let bit = rand.below((input.bytes().len() * 8) as u64) as usize; + let bit = state.rand_mut().below((input.bytes().len() * 8) as u64) as usize; input.bytes_mut()[bit >> 3] ^= (128 >> (bit & 7)) as u8; Ok(MutationResult::Mutated) } @@ -190,14 +188,14 @@ fn locate_diffs(this: &[u8], other: &[u8]) -> (i64, i64) { } /// Splicing mutator -pub fn mutation_splice( +pub fn mutation_splice( mutator: &mut M, - corpus: &mut C, - rand: &mut R, + state: &mut S, input: &mut I, ) -> Result where - M: Mutator, + M: Mutator, + S: HasRand + HasCorpus, C: Corpus, I: Input + HasBytesVec, R: Rand, @@ -206,7 +204,7 @@ where // We don't want to use the testcase we're already using for splicing let other_rr = loop { let mut found = false; - let (other_rr, _) = corpus.random_entry(rand)?.clone(); + let (other_rr, _) = state.corpus_mut().random_entry(state.rand_mut())?.clone(); match other_rr.try_borrow_mut() { Ok(_) => found = true, Err(_) => { diff --git a/src/stages/mod.rs b/src/stages/mod.rs index 8d06aa2fb8..b70318fb0f 100644 --- a/src/stages/mod.rs +++ b/src/stages/mod.rs @@ -4,7 +4,6 @@ pub use mutational::DefaultMutationalStage; use crate::corpus::testcase::Testcase; use crate::corpus::Corpus; use crate::engines::State; -use crate::events::EventManager; use crate::executors::Executor; use crate::inputs::Input; use crate::utils::Rand; @@ -12,19 +11,19 @@ use crate::AflError; use alloc::rc::Rc; use core::cell::RefCell; -pub trait Stage +pub trait Stage where - S: State, + S: State, C: Corpus, E: Executor, - EM: EventManager, I: Input, R: Rand, { /// Run the stage fn perform( &mut self, - testcase: Rc>>, + rand: &mut R, state: &mut S, + testcase: Rc>>, ) -> Result<(), AflError>; } diff --git a/src/stages/mutational.rs b/src/stages/mutational.rs index f700781396..79c1129ebe 100644 --- a/src/stages/mutational.rs +++ b/src/stages/mutational.rs @@ -15,13 +15,12 @@ use crate::AflError; // TODO multi mutators stage -pub trait MutationalStage: Stage +pub trait MutationalStage: Stage where M: Mutator, - S: State, + S: State, C: Corpus, E: Executor, - EM: EventManager, I: Input, R: Rand, { @@ -40,15 +39,17 @@ where /// Runs this (mutational) stage for the given testcase fn perform_mutational( &mut self, - testcase: Rc>>, + rand: &mut R, state: &mut S, + testcase: Rc>>, ) -> Result<(), AflError> { - let num = self.iterations(state.rand_mut()); + let num = self.iterations(rand); let input = testcase.borrow_mut().load_input()?.clone(); for i in 0..num { let mut input_tmp = input.clone(); - self.mutator_mut().mutate(state, &mut input_tmp, i as i32)?; + self.mutator_mut() + .mutate(rand, state, &mut input_tmp, i as i32)?; let interesting = state.evaluate_input(&input_tmp)?; diff --git a/src/utils.rs b/src/utils.rs index e7db765b8f..77261a2f25 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -47,27 +47,27 @@ pub trait Rand: Debug { } /// Has a Rand field, that can be used to get random values -pub trait HasRand { - type R: Rand; - - /// Get the hold Rand instance - fn rand(&self) -> &Self::R; - - /// Get the hold Rand instance (mutable) - fn rand_mut(&mut self) -> &mut Self::R; +pub trait HasRand +where + R: Rand, +{ + /// Get the hold RefCell Rand instance + fn rand(&self) -> &RefCell; // Gets the next 64 bit value fn rand_next(&mut self) -> u64 { - self.rand_mut().next() + self.rand().borrow_mut().next() } // Gets a value below the given 64 bit val (inclusive) fn rand_below(&mut self, upper_bound_excl: u64) -> u64 { - self.rand_mut().below(upper_bound_excl) + self.rand().borrow_mut().below(upper_bound_excl) } // Gets a value between the given lower bound (inclusive) and upper bound (inclusive) fn rand_between(&mut self, lower_bound_incl: u64, upper_bound_incl: u64) -> u64 { - self.rand_mut().between(lower_bound_incl, upper_bound_incl) + self.rand() + .borrow_mut() + .between(lower_bound_incl, upper_bound_incl) } }