From 5fd44ce98ebcd1ec912d0bb8c40cceff9879b3a5 Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Sun, 22 Nov 2020 15:34:25 +0100 Subject: [PATCH] wip --- src/corpus/mod.rs | 106 +++++++------------- src/engines/mod.rs | 150 +++++++++++++++++++++------- src/events/mod.rs | 6 ++ src/generators/mod.rs | 77 ++++++++------- src/mutators/mod.rs | 13 ++- src/mutators/scheduled.rs | 202 +++++++++++++++++++------------------- src/stages/mod.rs | 10 +- src/stages/mutational.rs | 97 +++++++++--------- src/utils.rs | 33 ++++++- 9 files changed, 391 insertions(+), 303 deletions(-) diff --git a/src/corpus/mod.rs b/src/corpus/mod.rs index abf1a844a2..8d9cfc2c7b 100644 --- a/src/corpus/mod.rs +++ b/src/corpus/mod.rs @@ -10,7 +10,7 @@ use core::marker::PhantomData; use std::path::PathBuf; use crate::inputs::Input; -use crate::utils::{HasRand, Rand}; +use crate::utils::Rand; use crate::AflError; pub trait HasEntriesVec @@ -25,9 +25,10 @@ where } /// Corpus with all current testcases -pub trait Corpus: HasEntriesVec + HasRand +pub trait Corpus: HasEntriesVec where I: Input, + R: Rand, { /// Returns the number of elements fn count(&self) -> usize { @@ -58,20 +59,20 @@ where } /// Gets a random entry - fn random_entry(&self) -> Result<(Rc>>, usize), AflError> { + fn random_entry(&self, rand: &mut R) -> Result<(Rc>>, usize), AflError> { if self.count() == 0 { Err(AflError::Empty("No entries in corpus".to_owned())) } else { let len = { self.entries().len() }; - let id = self.rand_below(len as u64) as usize; + let id = rand.below(len as u64) as usize; Ok((self.entries()[id].clone(), id)) } } // TODO: IntoIter /// Gets the next entry - fn next(&mut self) -> Result<(Rc>>, usize), AflError> { - self.random_entry() + fn next(&mut self, rand: &mut R) -> Result<(Rc>>, usize), AflError> { + self.random_entry(rand) } } @@ -80,8 +81,8 @@ where I: Input, R: Rand, { - rand: Rc>, entries: Vec>>>, + phantom: PhantomData, } impl HasEntriesVec for InMemoryCorpus @@ -97,19 +98,7 @@ where } } -impl HasRand for InMemoryCorpus -where - I: Input, - R: Rand, -{ - type R = R; - - fn rand(&self) -> &Rc> { - &self.rand - } -} - -impl Corpus for InMemoryCorpus +impl Corpus for InMemoryCorpus where I: Input, R: Rand, @@ -122,10 +111,10 @@ where I: Input, R: Rand, { - pub fn new(rand: &Rc>) -> Self { + pub fn new() -> Self { InMemoryCorpus { - rand: Rc::clone(rand), entries: vec![], + phantom: PhantomData, } } } @@ -136,9 +125,9 @@ where I: Input, R: Rand, { - rand: Rc>, entries: Vec>>>, dir_path: PathBuf, + phantom: PhantomData, } #[cfg(feature = "std")] @@ -156,20 +145,7 @@ where } #[cfg(feature = "std")] -impl HasRand for OnDiskCorpus -where - I: Input, - R: Rand, -{ - type R = R; - - fn rand(&self) -> &Rc> { - &self.rand - } -} - -#[cfg(feature = "std")] -impl Corpus for OnDiskCorpus +impl Corpus for OnDiskCorpus where I: Input, R: Rand, @@ -194,31 +170,33 @@ where I: Input, R: Rand, { - pub fn new(rand: &Rc>, dir_path: PathBuf) -> Self { + pub fn new(dir_path: PathBuf) -> Self { OnDiskCorpus { - rand: Rc::clone(rand), dir_path: dir_path, entries: vec![], + phantom: PhantomData, } } } /// A Queue-like corpus, wrapping an existing Corpus instance -pub struct QueueCorpus +pub struct QueueCorpus where + C: Corpus, I: Input, - C: Corpus, + R: Rand, { corpus: C, - phantom: PhantomData, + phantom: PhantomData<(I, R)>, pos: usize, cycles: u64, } -impl<'a, I, C> HasEntriesVec for QueueCorpus +impl HasEntriesVec for QueueCorpus where + C: Corpus, I: Input, - C: Corpus, + R: Rand, { fn entries(&self) -> &[Rc>>] { self.corpus.entries() @@ -228,22 +206,11 @@ where } } -impl<'a, I, C> HasRand for QueueCorpus +impl Corpus for QueueCorpus where + C: Corpus, I: Input, - C: Corpus, -{ - type R = C::R; - - fn rand(&self) -> &Rc> { - self.corpus.rand() - } -} - -impl<'a, I, C> Corpus for QueueCorpus -where - I: Input, - C: Corpus, + R: Rand, { /// Returns the number of elements fn count(&self) -> usize { @@ -260,12 +227,12 @@ where } /// Gets a random entry - fn random_entry(&self) -> Result<(Rc>>, usize), AflError> { - self.corpus.random_entry() + fn random_entry(&self, rand: &mut R) -> Result<(Rc>>, usize), AflError> { + self.corpus.random_entry(rand) } /// Gets the next entry - fn next(&mut self) -> Result<(Rc>>, usize), AflError> { + fn next(&mut self, _rand: &mut R) -> Result<(Rc>>, usize), AflError> { self.pos += 1; if self.corpus.count() == 0 { return Err(AflError::Empty("Corpus".to_owned())); @@ -279,13 +246,14 @@ where } } -impl<'a, I, C> QueueCorpus +impl QueueCorpus where + C: Corpus, I: Input, - C: Corpus, + R: Rand, { pub fn new(corpus: C) -> Self { - QueueCorpus:: { + QueueCorpus { corpus: corpus, phantom: PhantomData, cycles: 0, @@ -365,13 +333,15 @@ mod tests { #[test] fn test_queuecorpus() { - let rand: Rc<_> = DefaultRand::new(0).into(); - let mut q = QueueCorpus::new(OnDiskCorpus::new(&rand, PathBuf::from("fancy/path"))); + let mut rand = DefaultRand::new(0); + let mut q = QueueCorpus::new(OnDiskCorpus::::new(PathBuf::from( + "fancy/path", + ))); let t: Rc<_> = Testcase::with_filename(BytesInput::new(vec![0 as u8; 4]), "fancyfile".into()).into(); q.add(t); let filename = q - .next() + .next(&mut rand) .unwrap() .0 .borrow() @@ -381,7 +351,7 @@ mod tests { .to_owned(); assert_eq!( filename, - q.next() + q.next(&mut rand) .unwrap() .0 .borrow() diff --git a/src/engines/mod.rs b/src/engines/mod.rs index 68a8879814..e018bd9688 100644 --- a/src/engines/mod.rs +++ b/src/engines/mod.rs @@ -6,21 +6,30 @@ use alloc::vec::Vec; use core::cell::RefCell; use core::fmt::Debug; +use hashbrown::HashMap; + use crate::corpus::{Corpus, Testcase}; +use crate::events::EventManager; use crate::executors::Executor; use crate::feedbacks::Feedback; use crate::inputs::Input; use crate::observers::Observer; use crate::stages::Stage; +use crate::utils::{HasRand, Rand}; use crate::AflError; -pub trait StateMetadata: Debug {} +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 +pub trait State: HasRand where - C: Corpus, + C: Corpus, E: Executor, + EM: EventManager, I: Input, + R: Rand, { /// Get executions fn executions(&self) -> usize; @@ -28,6 +37,21 @@ 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>; + + /// Get all the metadatas into an HashMap (mutable) + fn metadatas_mut(&mut self) -> &mut HashMap<&'static str, Box>; + + /// Add a metadata + fn add_metadata(&mut self, meta: Box) { + self.metadatas_mut().insert(meta.name(), meta); + } + /// Get the linked observers fn observers(&self) -> &[Rc>]; @@ -105,24 +129,49 @@ where } } -pub struct DefaultState +pub struct DefaultState where - C: Corpus, + C: Corpus, E: Executor, + EM: EventManager, I: Input, + R: Rand, { + rand: R, executions: usize, + events_manager: EM, + metadatas: HashMap<&'static str, Box>, observers: Vec>>, feedbacks: Vec>>, corpus: C, executor: E, } -impl State for DefaultState +impl HasRand for DefaultState where - C: Corpus, + C: Corpus, E: Executor, + EM: EventManager, I: Input, + R: Rand, +{ + type R = R; + + fn rand(&self) -> &Self::R { + &self.rand + } + fn rand_mut(&mut self) -> &mut Self::R { + &mut self.rand + } +} + +impl State for DefaultState +where + C: Corpus, + E: Executor, + EM: EventManager, + I: Input, + R: Rand, { fn executions(&self) -> usize { self.executions @@ -132,6 +181,21 @@ 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 + } + + fn metadatas_mut(&mut self) -> &mut HashMap<&'static str, Box> { + &mut self.metadatas + } + fn observers(&self) -> &[Rc>] { &self.observers } @@ -165,15 +229,20 @@ where } } -impl DefaultState +impl DefaultState where - C: Corpus, + C: Corpus, E: Executor, + EM: EventManager, I: Input, + R: Rand, { - pub fn new(corpus: C, executor: E) -> Self { + pub fn new(corpus: C, executor: E, events_manager: EM, rand: R) -> Self { DefaultState { + rand: rand, executions: 0, + events_manager: events_manager, + metadatas: HashMap::default(), observers: vec![], feedbacks: vec![], corpus: corpus, @@ -182,24 +251,25 @@ where } } -pub trait Engine +pub trait Engine where - S: State, - C: Corpus, + 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()?; - #[cfg(feature = "std")] + let (testcase, idx) = state.corpus_mut().next(state.rand_mut())?; println!("Cur entry: {}\tExecutions: {}", idx, state.executions()); for stage in self.stages_mut() { stage.perform(testcase.clone(), state)?; @@ -208,38 +278,44 @@ where } } -pub struct DefaultEngine +pub struct DefaultEngine where - S: State, - C: Corpus, + S: State, + C: Corpus, E: Executor, + EM: EventManager, I: Input, + R: Rand, { - stages: Vec>>, + stages: Vec>>, } -impl Engine for DefaultEngine +impl Engine for DefaultEngine where - S: State, - C: Corpus, + 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 +impl DefaultEngine where - S: State, - C: Corpus, + S: State, + C: Corpus, E: Executor, + EM: EventManager, I: Input, + R: Rand, { pub fn new() -> Self { DefaultEngine { stages: vec![] } @@ -253,6 +329,7 @@ mod tests { use crate::corpus::{Corpus, InMemoryCorpus, Testcase}; use crate::engines::{DefaultEngine, DefaultState, Engine}; + use crate::events::LoggerEventManager; use crate::executors::inmemory::InMemoryExecutor; use crate::executors::{Executor, ExitKind}; use crate::inputs::bytes::BytesInput; @@ -268,19 +345,20 @@ mod tests { #[test] fn test_engine() { - let rand = DefaultRand::new(0).into(); - - let mut corpus = InMemoryCorpus::::new(&rand); + let mut corpus = InMemoryCorpus::::new(); let testcase = Testcase::new(vec![0; 4]).into(); corpus.add(testcase); let executor = InMemoryExecutor::::new(harness); - let mut state = DefaultState::new(corpus, executor); + let events = LoggerEventManager::new(); + let rand = DefaultRand::new(0); + + let mut state = DefaultState::new(corpus, executor, events, rand); let mut engine = DefaultEngine::new(); - let mut mutator = DefaultScheduledMutator::new(&rand); + let mut mutator = DefaultScheduledMutator::new(); mutator.add_mutation(mutation_bitflip); - let stage = DefaultMutationalStage::new(&rand, mutator); + let stage = DefaultMutationalStage::new(mutator); engine.add_stage(Box::new(stage)); // diff --git a/src/events/mod.rs b/src/events/mod.rs index f0da136f92..2b32eb7df6 100644 --- a/src/events/mod.rs +++ b/src/events/mod.rs @@ -94,3 +94,9 @@ impl EventManager for LoggerEventManager { Ok(0) } } + +impl LoggerEventManager { + pub fn new() -> Self { + LoggerEventManager {} + } +} diff --git a/src/generators/mod.rs b/src/generators/mod.rs index a8d30f1d54..a5aac2e5a0 100644 --- a/src/generators/mod.rs +++ b/src/generators/mod.rs @@ -1,19 +1,19 @@ -use alloc::rc::Rc; use alloc::vec::Vec; -use core::cell::RefCell; use core::cmp::min; +use core::marker::PhantomData; use crate::inputs::bytes::BytesInput; use crate::inputs::Input; -use crate::utils::{HasRand, Rand}; +use crate::utils::Rand; use crate::AflError; -pub trait Generator: HasRand +pub trait Generator where I: Input, + R: Rand, { /// Generate a new input - fn generate(&mut self) -> Result; + fn generate(&mut self, rand: &mut R) -> Result; /// Generate a new dummy input fn generate_dummy(&self) -> I; @@ -21,29 +21,21 @@ where const DUMMY_BYTES_SIZE: usize = 64; -pub struct RandBytesGenerator { - rand: Rc>, +pub struct RandBytesGenerator +where + R: Rand, +{ max_size: usize, + phantom: PhantomData, } -impl HasRand for RandBytesGenerator +impl Generator for RandBytesGenerator where R: Rand, { - type R = R; - - fn rand(&self) -> &Rc> { - &self.rand - } -} - -impl Generator for RandBytesGenerator -where - R: Rand, -{ - fn generate(&mut self) -> Result { - let size = self.rand_below(self.max_size as u64); - let random_bytes: Vec = (0..size).map(|_| self.rand_below(256) as u8).collect(); + fn generate(&mut self, rand: &mut R) -> Result { + let size = rand.below(self.max_size as u64); + let random_bytes: Vec = (0..size).map(|_| rand.below(256) as u8).collect(); Ok(BytesInput::new(random_bytes)) } @@ -53,31 +45,32 @@ where } } -pub struct RandPrintablesGenerator { - rand: Rc>, - max_size: usize, -} - -impl HasRand for RandPrintablesGenerator +impl RandBytesGenerator where R: Rand, { - type R = R; - - fn rand(&self) -> &Rc> { - &self.rand + pub fn new(max_size: usize) -> Self { + RandBytesGenerator { + max_size: max_size, + phantom: PhantomData, + } } } -impl Generator for RandPrintablesGenerator +pub struct RandPrintablesGenerator { + max_size: usize, + phantom: PhantomData, +} + +impl Generator for RandPrintablesGenerator where R: Rand, { - fn generate(&mut self) -> Result { - let size = self.rand_below(self.max_size as u64); + fn generate(&mut self, rand: &mut R) -> Result { + let size = rand.below(self.max_size as u64); let printables = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz \t\n!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~".as_bytes(); let random_bytes: Vec = (0..size) - .map(|_| printables[self.rand_below(printables.len() as u64) as usize]) + .map(|_| printables[rand.below(printables.len() as u64) as usize]) .collect(); Ok(BytesInput::new(random_bytes)) } @@ -87,3 +80,15 @@ where BytesInput::new(vec!['0' as u8; size]) } } + +impl RandPrintablesGenerator +where + R: Rand, +{ + pub fn new(max_size: usize) -> Self { + RandPrintablesGenerator { + max_size: max_size, + phantom: PhantomData, + } + } +} diff --git a/src/mutators/mod.rs b/src/mutators/mod.rs index dde4b6169b..0c0ecdb9ec 100644 --- a/src/mutators/mod.rs +++ b/src/mutators/mod.rs @@ -4,18 +4,21 @@ pub use scheduled::DefaultScheduledMutator; pub use scheduled::HavocBytesMutator; pub use scheduled::ScheduledMutator; -use crate::corpus::Corpus; +use crate::corpus::{Corpus, HasCorpus}; +//use crate::engines::State; use crate::inputs::Input; -use crate::utils::HasRand; +use crate::utils::{HasRand, Rand}; use crate::AflError; -pub trait Mutator: HasRand +pub trait Mutator where - C: Corpus, + S: HasRand + HasCorpus, + C: Corpus, I: Input, + R: Rand, { /// Mutate a given input - fn mutate(&mut self, corpus: &mut C, input: &mut I, stage_idx: i32) -> Result<(), AflError>; + fn mutate(&mut self, state: &mut S, input: &mut I, stage_idx: i32) -> Result<(), AflError>; /// Post-process given the outcome of the execution fn post_exec(&mut self, _is_interesting: bool, _stage_idx: i32) -> Result<(), AflError> { diff --git a/src/mutators/scheduled.rs b/src/mutators/scheduled.rs index 91b3122bc6..1c5460aa17 100644 --- a/src/mutators/scheduled.rs +++ b/src/mutators/scheduled.rs @@ -1,12 +1,10 @@ use crate::inputs::{HasBytesVec, Input}; -use crate::mutators::Corpus; use crate::mutators::Mutator; +use crate::mutators::{Corpus, HasCorpus}; use crate::utils::{HasRand, Rand}; use crate::AflError; -use alloc::rc::Rc; use alloc::vec::Vec; -use core::cell::RefCell; use core::marker::PhantomData; pub enum MutationResult { @@ -14,43 +12,53 @@ pub enum MutationResult { Skipped, } +// TODO maybe the mutator arg is not needed /// The generic function type that identifies mutations -type MutationFunction = fn(&mut M, &mut C, &mut I) -> Result; +type MutationFunction = fn(&mut M, &mut S, &mut I) -> Result; -pub trait ComposedByMutations +pub trait ComposedByMutations where - C: Corpus, + S: HasRand + HasCorpus, + C: Corpus, I: Input, + R: Rand, { /// Get a mutation by index - fn mutation_by_idx(&self, index: usize) -> Result, AflError>; + fn mutation_by_idx(&self, index: usize) -> Result, AflError>; /// Get the number of mutations fn mutations_count(&self) -> usize; /// Add a mutation - fn add_mutation(&mut self, mutation: MutationFunction); + fn add_mutation(&mut self, mutation: MutationFunction); } -pub trait ScheduledMutator: Mutator + ComposedByMutations +pub trait ScheduledMutator: + Mutator + ComposedByMutations where - C: Corpus, + S: HasRand + HasCorpus, + C: Corpus, I: Input, + R: Rand, { /// Compute the number of iterations used to apply stacked mutations - fn iterations(&mut self, _input: &I) -> u64 { - 1 << (1 + self.rand_below(7)) + fn iterations(&mut self, rand: &mut R, _input: &I) -> u64 { + 1 << (1 + rand.below(7)) } /// Get the next mutation to apply - fn schedule(&mut self, _input: &I) -> Result, AflError> { + fn schedule( + &mut self, + rand: &mut R, + _input: &I, + ) -> Result, AflError> { let count = self.mutations_count() as u64; if count == 0 { return Err(AflError::Empty("no mutations".into())); } let idx; { - idx = self.rand_below(count) as usize; + idx = rand.below(count) as usize; } self.mutation_by_idx(idx) } @@ -59,59 +67,51 @@ where /// Implementations must forward mutate() to this method fn scheduled_mutate( &mut self, - corpus: &mut C, + state: &mut S, input: &mut I, _stage_idx: i32, ) -> Result<(), AflError> { - let num = self.iterations(input); + let num = self.iterations(state.rand_mut(), input); for _ in 0..num { - self.schedule(input)?(self, corpus, input)?; + self.schedule(state.rand_mut(), input)?(self, state, input)?; } Ok(()) } } -pub struct DefaultScheduledMutator +pub struct DefaultScheduledMutator where - C: Corpus, + C: Corpus, I: Input, R: Rand, { - rand: Rc>, - mutations: Vec>, + mutations: Vec>, } -impl HasRand for DefaultScheduledMutator +impl Mutator for DefaultScheduledMutator where - C: Corpus, + C: Corpus, I: Input, R: Rand, { - type R = R; - - fn rand(&self) -> &Rc> { - &self.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) } } -impl Mutator for DefaultScheduledMutator +impl ComposedByMutations for DefaultScheduledMutator where - C: Corpus, + C: Corpus, I: Input, R: Rand, { - fn mutate(&mut self, corpus: &mut C, input: &mut I, _stage_idx: i32) -> Result<(), AflError> { - self.scheduled_mutate(corpus, input, _stage_idx) - } -} - -impl ComposedByMutations for DefaultScheduledMutator -where - C: Corpus, - I: Input, - R: Rand, -{ - fn mutation_by_idx(&self, index: usize) -> Result, AflError> { + fn mutation_by_idx(&self, index: usize) -> Result, AflError> { if index >= self.mutations.len() { return Err(AflError::Unknown("oob".into())); } @@ -122,14 +122,14 @@ where self.mutations.len() } - fn add_mutation(&mut self, mutation: MutationFunction) { + fn add_mutation(&mut self, mutation: MutationFunction) { self.mutations.push(mutation) } } -impl ScheduledMutator for DefaultScheduledMutator +impl ScheduledMutator for DefaultScheduledMutator where - C: Corpus, + C: Corpus, I: Input, R: Rand, { @@ -138,39 +138,37 @@ where impl DefaultScheduledMutator where - C: Corpus, + C: Corpus, I: Input, R: Rand, { /// Create a new DefaultScheduledMutator instance without mutations and corpus - pub fn new(rand: &Rc>) -> Self { - DefaultScheduledMutator { - rand: Rc::clone(rand), - mutations: vec![], - } + pub fn new() -> Self { + DefaultScheduledMutator { mutations: vec![] } } /// Create a new DefaultScheduledMutator instance specifying mutations and corpus too - pub fn new_all(rand: &Rc>, mutations: Vec>) -> Self { + pub fn new_all(mutations: Vec>) -> Self { DefaultScheduledMutator { - rand: Rc::clone(rand), mutations: mutations, } } } /// 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, input: &mut I, ) -> Result where - C: Corpus, - M: HasRand, + M: Mutator, + C: Corpus, I: Input + HasBytesVec, + R: Rand, { - let bit = mutator.rand_below((input.bytes().len() * 8) as u64) as usize; + let bit = rand.below((input.bytes().len() * 8) as u64) as usize; input.bytes_mut()[bit >> 3] ^= (128 >> (bit & 7)) as u8; Ok(MutationResult::Mutated) } @@ -192,21 +190,23 @@ 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, input: &mut I, ) -> Result where - C: Corpus, - M: HasRand, + M: Mutator, + C: Corpus, I: Input + HasBytesVec, + R: Rand, { let mut retry_count = 0; // 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()?.clone(); + let (other_rr, _) = corpus.random_entry(rand)?.clone(); match other_rr.try_borrow_mut() { Ok(_) => found = true, Err(_) => { @@ -238,7 +238,7 @@ where counter += 1; }; - let split_at = mutator.rand_between(first_diff as u64, last_diff as u64) as usize; + let split_at = rand.between(first_diff as u64, last_diff as u64) as usize; // println!("Splicing at {}", split_at); @@ -253,51 +253,48 @@ where } /// Schedule some selected byte level mutations given a ScheduledMutator type -pub struct HavocBytesMutator +pub struct HavocBytesMutator where - C: Corpus, + C: Corpus, I: Input + HasBytesVec, - S: ScheduledMutator, + SM: ScheduledMutator, + R: Rand, { - scheduled: S, - phantom: PhantomData, + scheduled: SM, + phantom: PhantomData<(I, R)>, _phantom_corpus: PhantomData, } -impl HasRand for HavocBytesMutator +impl Mutator for HavocBytesMutator where - C: Corpus, + C: Corpus, I: Input + HasBytesVec, - S: ScheduledMutator, -{ - type R = S::R; - - fn rand(&self) -> &Rc> { - &self.scheduled.rand() - } -} - -impl Mutator for HavocBytesMutator -where - C: Corpus, - I: Input + HasBytesVec, - S: ScheduledMutator, + SM: ScheduledMutator, + R: Rand, { /// Mutate bytes - fn mutate(&mut self, corpus: &mut C, input: &mut I, stage_idx: i32) -> Result<(), AflError> { - self.scheduled.mutate(corpus, input, stage_idx) + 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) } } -impl HavocBytesMutator +impl HavocBytesMutator where - C: Corpus, + C: Corpus, I: Input + HasBytesVec, - S: ScheduledMutator, + SM: ScheduledMutator, + R: Rand, { /// Create a new HavocBytesMutator instance given a ScheduledMutator to wrap - pub fn new(mut scheduled: S) -> Self { + pub fn new(mut scheduled: SM) -> Self { scheduled.add_mutation(mutation_bitflip); + scheduled.add_mutation(mutation_splice); HavocBytesMutator { scheduled: scheduled, phantom: PhantomData, @@ -306,15 +303,15 @@ where } } -impl HavocBytesMutator> +impl HavocBytesMutator, R> where - C: Corpus, + C: Corpus, I: Input + HasBytesVec, R: Rand, { /// Create a new HavocBytesMutator instance wrapping DefaultScheduledMutator - pub fn new_default(rand: &Rc>) -> Self { - let mut scheduled = DefaultScheduledMutator::::new(rand); + pub fn new_default() -> Self { + let mut scheduled = DefaultScheduledMutator::::new(); scheduled.add_mutation(mutation_bitflip); scheduled.add_mutation(mutation_splice); HavocBytesMutator { @@ -328,29 +325,32 @@ where #[cfg(test)] mod tests { use crate::inputs::BytesInput; - use crate::mutators::scheduled::mutation_splice; - use crate::utils::{DefaultHasRand, Rand, XKCDRand}; + use crate::mutators::scheduled::{mutation_splice, DefaultScheduledMutator}; + use crate::utils::{Rand, XKCDRand}; use crate::{ corpus::{Corpus, InMemoryCorpus, Testcase}, inputs::HasBytesVec, }; - use alloc::rc::Rc; #[test] fn test_mut_splice() { // With the current impl, seed of 1 will result in a split at pos 2. - let rand: Rc<_> = XKCDRand::new().into(); - let mut has_rand = DefaultHasRand::new(&rand); - let mut corpus: InMemoryCorpus = InMemoryCorpus::new(&rand); + let mut rand = XKCDRand::new(); + let mut corpus: InMemoryCorpus = InMemoryCorpus::new(); corpus.add(Testcase::new(vec!['a' as u8, 'b' as u8, 'c' as u8]).into()); corpus.add(Testcase::new(vec!['d' as u8, 'e' as u8, 'f' as u8]).into()); - let (testcase_rr, _) = corpus.next().expect("Corpus did not contain entries"); + let (testcase_rr, _) = corpus + .next(&mut rand) + .expect("Corpus did not contain entries"); let mut testcase = testcase_rr.borrow_mut(); let mut input = testcase.load_input().expect("No input in testcase").clone(); - rand.borrow_mut().set_seed(5); - mutation_splice(&mut has_rand, &mut corpus, &mut input).unwrap(); + rand.set_seed(5); + let mut mutator = + DefaultScheduledMutator::, BytesInput, XKCDRand>::new(); + + mutation_splice(&mut mutator, &mut corpus, &mut rand, &mut input).unwrap(); #[cfg(feature = "std")] println!("{:?}", input.bytes()); diff --git a/src/stages/mod.rs b/src/stages/mod.rs index 2df141214e..8d06aa2fb8 100644 --- a/src/stages/mod.rs +++ b/src/stages/mod.rs @@ -4,18 +4,22 @@ 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; use crate::AflError; use alloc::rc::Rc; use core::cell::RefCell; -pub trait Stage +pub trait Stage where - S: State, - C: Corpus, + S: State, + C: Corpus, E: Executor, + EM: EventManager, I: Input, + R: Rand, { /// Run the stage fn perform( diff --git a/src/stages/mutational.rs b/src/stages/mutational.rs index c024d4e36a..f700781396 100644 --- a/src/stages/mutational.rs +++ b/src/stages/mutational.rs @@ -4,23 +4,26 @@ use core::marker::PhantomData; use crate::corpus::testcase::Testcase; use crate::engines::State; +use crate::events::EventManager; use crate::executors::Executor; use crate::inputs::Input; use crate::mutators::Mutator; use crate::stages::Corpus; use crate::stages::Stage; -use crate::utils::{HasRand, Rand}; +use crate::utils::Rand; use crate::AflError; // TODO multi mutators stage -pub trait MutationalStage: Stage + HasRand +pub trait MutationalStage: Stage where - M: Mutator, - S: State, + M: Mutator, + S: State, + C: Corpus, E: Executor, - C: Corpus, + EM: EventManager, I: Input, + R: Rand, { /// The mutator registered for this stage fn mutator(&self) -> &M; @@ -30,8 +33,8 @@ where /// Gets the number of iterations this mutator should run for. /// This call uses internal mutability, so it may change for each call - fn iterations(&mut self) -> usize { - 1 + self.rand_below(128) as usize + fn iterations(&mut self, rand: &mut R) -> usize { + 1 + rand.below(128) as usize } /// Runs this (mutational) stage for the given testcase @@ -40,13 +43,12 @@ where testcase: Rc>>, state: &mut S, ) -> Result<(), AflError> { - let num = self.iterations(); + let num = self.iterations(state.rand_mut()); let input = testcase.borrow_mut().load_input()?.clone(); for i in 0..num { let mut input_tmp = input.clone(); - self.mutator_mut() - .mutate(state.corpus_mut(), &mut input_tmp, i as i32)?; + self.mutator_mut().mutate(state, &mut input_tmp, i as i32)?; let interesting = state.evaluate_input(&input_tmp)?; @@ -61,39 +63,28 @@ where } /// The default mutational stage -pub struct DefaultMutationalStage +pub struct DefaultMutationalStage where - M: Mutator, - C: Corpus, - I: Input, - R: Rand, -{ - rand: Rc>, - mutator: M, - _phantom_corpus: PhantomData, - _phantom_input: PhantomData, -} - -impl HasRand for DefaultMutationalStage -where - M: Mutator, - C: Corpus, - I: Input, - R: Rand, -{ - type R = R; - - fn rand(&self) -> &Rc> { - &self.rand - } -} - -impl MutationalStage for DefaultMutationalStage -where - M: Mutator, - S: State, - C: Corpus, + M: Mutator, + S: State, + C: Corpus, E: Executor, + EM: EventManager, + I: Input, + R: Rand, +{ + mutator: M, + phantom: PhantomData<(S, C, E, EM, I, R)>, +} + +impl MutationalStage + for DefaultMutationalStage +where + M: Mutator, + S: State, + C: Corpus, + E: Executor, + EM: EventManager, I: Input, R: Rand, { @@ -108,12 +99,13 @@ where } } -impl Stage for DefaultMutationalStage +impl Stage for DefaultMutationalStage where - M: Mutator, - S: State, - C: Corpus, + M: Mutator, + S: State, + C: Corpus, E: Executor, + EM: EventManager, I: Input, R: Rand, { @@ -126,20 +118,21 @@ where } } -impl DefaultMutationalStage +impl DefaultMutationalStage where - M: Mutator, - C: Corpus, + M: Mutator, + S: State, + C: Corpus, + E: Executor, + EM: EventManager, I: Input, R: Rand, { /// Creates a new default mutational stage - pub fn new(rand: &Rc>, mutator: M) -> Self { + pub fn new(mutator: M) -> Self { DefaultMutationalStage { - rand: Rc::clone(rand), mutator: mutator, - _phantom_corpus: PhantomData, - _phantom_input: PhantomData, + phantom: PhantomData, } } } diff --git a/src/utils.rs b/src/utils.rs index 39bab1b8b6..e7db765b8f 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -46,10 +46,35 @@ pub trait Rand: Debug { } } -/// Has a Rand Rc RefCell field (internal mutability), that can be used to get random values +/// 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; + + // Gets the next 64 bit value + fn rand_next(&mut self) -> u64 { + self.rand_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) + } + + // 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) + } +} + +/// Has a Rand Rc RefCell field (internal mutability), that can be used to get random values +pub trait HasRandRR { + type R: Rand; + /// Get the hold Rand instance fn rand(&self) -> &Rc>; @@ -172,6 +197,7 @@ impl XKCDRand { } } +/* /// A very basic HasRand pub struct DefaultHasRand where @@ -205,6 +231,7 @@ where } } } +*/ /// Get the next higher power of two pub const fn next_pow2(val: u64) -> u64 { @@ -219,7 +246,7 @@ pub const fn next_pow2(val: u64) -> u64 { #[cfg(test)] mod tests { - use crate::utils::{next_pow2, DefaultHasRand, DefaultRand, HasRand, Rand}; + use crate::utils::{next_pow2, DefaultRand, HasRand, Rand}; #[test] fn test_rand() { @@ -244,6 +271,7 @@ mod tests { assert!(rand.between(11, 20) > 10); } + /* #[test] fn test_has_rand() { let rand = DefaultRand::new(0).into(); @@ -252,6 +280,7 @@ mod tests { assert!(has_rand.rand_below(100) < 100); assert_eq!(has_rand.rand_below(1), 0); } + */ #[test] fn test_next_pow2() {