From c9cc66311af32f44d238a6bc99ad336eb6e36a19 Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Fri, 20 Nov 2020 12:53:33 +0100 Subject: [PATCH] state trait --- src/corpus/testcase.rs | 23 +++-- src/engines/mod.rs | 196 +++++++++++++++++++++++++++++++++------ src/executors/mod.rs | 18 ++-- src/feedbacks/mod.rs | 23 +++-- src/stages/mod.rs | 8 +- src/stages/mutational.rs | 60 ++++++------ 6 files changed, 245 insertions(+), 83 deletions(-) diff --git a/src/corpus/testcase.rs b/src/corpus/testcase.rs index 1e4135e7f2..a8207a074b 100644 --- a/src/corpus/testcase.rs +++ b/src/corpus/testcase.rs @@ -80,10 +80,7 @@ where match self { Self::Stored { filename } => { let input = I::from_file(&filename)?; - Ok(Self::Loaded { - filename, - input, - }) + Ok(Self::Loaded { filename, input }) } Self::Loaded { input: _, @@ -170,11 +167,13 @@ where { /// Make sure to return a valid input instance loading it from disk if not in memory pub fn load_input(&mut self) -> Result<&I, AflError> { - // TODO: Implement cache to disk - match self.input.as_ref() { - Some(i) => Ok(i), - None => Err(AflError::NotImplemented("load_input".into())), + if self.input.is_none() { + let input = I::from_file(self.filename.as_ref().ok_or(AflError::EmptyOptional( + "filename not specified".to_string(), + ))?)?; + self.input = Some(input); } + Ok(self.input.as_ref().unwrap()) } /// Get the input, if any @@ -226,4 +225,12 @@ where metadatas: HashMap::default(), } } + + pub fn default() -> Self { + Testcase { + input: None, + filename: None, + metadatas: HashMap::default(), + } + } } diff --git a/src/engines/mod.rs b/src/engines/mod.rs index dbd7e32663..da8528553a 100644 --- a/src/engines/mod.rs +++ b/src/engines/mod.rs @@ -1,58 +1,126 @@ //! The engine is the core piece of every good fuzzer use alloc::boxed::Box; +use alloc::rc::Rc; use alloc::vec::Vec; +use core::cell::RefCell; -use crate::corpus::Corpus; +use crate::corpus::{Corpus, Testcase}; +use crate::executors::Executor; use crate::feedbacks::Feedback; use crate::inputs::Input; +use crate::observers::Observer; use crate::stages::Stage; use crate::AflError; -pub trait Engine +pub trait State where C: Corpus, + E: Executor, I: Input, { + /// Get the linked observers + fn observers(&self) -> &[Box]; + + /// Get the linked observers + fn observers_mut(&mut self) -> &mut Vec>; + + /// Add a linked observer + fn add_observer(&mut self, observer: Box) { + self.observers_mut().push(observer); + } + + /// Reset the state of all the observes linked to this executor + fn reset_observers(&mut self) -> Result<(), AflError> { + for observer in self.observers_mut() { + observer.reset()?; + } + Ok(()) + } + + /// Run the post exec hook for all the observes linked to this executor + fn post_exec_observers(&mut self) -> Result<(), AflError> { + self.observers_mut() + .iter_mut() + .map(|x| x.post_exec()) + .fold(Ok(()), |acc, x| if x.is_err() { x } else { acc }) + } + + /// Returns vector of feebacks fn feedbacks(&self) -> &[Box>]; + /// Returns vector of feebacks (mutable) fn feedbacks_mut(&mut self) -> &mut Vec>>; + /// Adds a feedback fn add_feedback(&mut self, feedback: Box>) { self.feedbacks_mut().push(feedback); } - fn stages(&self) -> &[Box>]; + /// Return the corpus + fn corpus(&self) -> &C; - fn stages_mut(&mut self) -> &mut Vec>>; + /// Return the corpus (mutable) + fn corpus_mut(&mut self) -> &mut C; - fn add_stage(&mut self, stage: Box>) { - self.stages_mut().push(stage); - } + /// Return the executor + fn executor(&self) -> &E; - fn fuzz_one(&mut self, corpus: &mut C) -> Result<(), AflError> { - let testcase = corpus.next()?; - for stage in self.stages_mut() { - stage.perform(testcase.clone(), corpus)?; + /// Return the executor (mutable) + fn executor_mut(&mut self) -> &mut E; + + /// Runs the input and triggers observers and feedback + fn evaluate_input(&mut self, input: &I) -> Result { + self.reset_observers()?; + self.executor_mut().run_target(input)?; + self.post_exec_observers()?; + + let mut rate_acc = 0; + for feedback in self.feedbacks_mut() { + rate_acc += feedback.is_interesting(input)?; + } + + if rate_acc >= 25 { + let testcase = Rc::new(RefCell::new(Testcase::new(input.clone()))); + for feedback in self.feedbacks_mut() { + feedback.append_metadata(testcase.clone())?; + } + Ok(true) + } else { + for feedback in self.feedbacks_mut() { + feedback.discard_metadata()?; + } + Ok(false) } - Ok(()) } } -pub struct DefaultEngine +pub struct DefaultState where C: Corpus, + E: Executor, I: Input, { + observers: Vec>, feedbacks: Vec>>, - stages: Vec>>, + corpus: C, + executor: E, } -impl Engine for DefaultEngine +impl State for DefaultState where C: Corpus, + E: Executor, I: Input, { + fn observers(&self) -> &[Box] { + &self.observers + } + + fn observers_mut(&mut self) -> &mut Vec> { + &mut self.observers + } + fn feedbacks(&self) -> &[Box>] { &self.feedbacks } @@ -61,25 +129,98 @@ where &mut self.feedbacks } - fn stages(&self) -> &[Box>] { + fn corpus(&self) -> &C { + &self.corpus + } + + fn corpus_mut(&mut self) -> &mut C { + &mut self.corpus + } + + fn executor(&self) -> &E { + &self.executor + } + + fn executor_mut(&mut self) -> &mut E { + &mut self.executor + } +} + +impl DefaultState +where + C: Corpus, + E: Executor, + I: Input, +{ + pub fn new(corpus: C, executor: E) -> Self { + DefaultState { + observers: vec![], + feedbacks: vec![], + corpus: corpus, + executor: executor, + } + } +} + +pub trait Engine +where + S: State, + C: Corpus, + E: Executor, + I: Input, +{ + fn stages(&self) -> &[Box>]; + + fn stages_mut(&mut self) -> &mut Vec>>; + + fn add_stage(&mut self, stage: Box>) { + self.stages_mut().push(stage); + } + + fn fuzz_one(&mut self, state: &mut S) -> Result<(), AflError> { + let testcase = state.corpus_mut().next()?; + for stage in self.stages_mut() { + stage.perform(testcase.clone(), state)?; + } + Ok(()) + } +} + +pub struct DefaultEngine +where + S: State, + C: Corpus, + E: Executor, + I: Input, +{ + stages: Vec>>, +} + +impl Engine for DefaultEngine +where + S: State, + C: Corpus, + E: Executor, + I: Input, +{ + 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, + E: Executor, I: Input, { pub fn new() -> Self { - DefaultEngine { - feedbacks: vec![], - stages: vec![], - } + DefaultEngine { stages: vec![] } } } @@ -89,7 +230,7 @@ mod tests { use alloc::boxed::Box; use crate::corpus::{Corpus, InMemoryCorpus, Testcase}; - use crate::engines::{DefaultEngine, Engine}; + use crate::engines::{DefaultEngine, DefaultState, Engine}; use crate::executors::inmemory::InMemoryExecutor; use crate::executors::{Executor, ExitKind}; use crate::inputs::bytes::BytesInput; @@ -110,18 +251,21 @@ mod tests { let mut corpus = InMemoryCorpus::::new(&rand); let testcase = Testcase::new(vec![0; 4]).into(); corpus.add(testcase); - let executor = InMemoryExecutor::::new(harness).into(); + + let executor = InMemoryExecutor::::new(harness); + let mut state = DefaultState::new(corpus, executor); + let mut engine = DefaultEngine::new(); let mut mutator = DefaultScheduledMutator::new(&rand); mutator.add_mutation(mutation_bitflip); - let stage = DefaultMutationalStage::new(&rand, &executor, mutator); + let stage = DefaultMutationalStage::new(&rand, mutator); engine.add_stage(Box::new(stage)); // for i in 0..1000 { engine - .fuzz_one(&mut corpus) + .fuzz_one(&mut state) .expect(&format!("Error in iter {}", i)); } } diff --git a/src/executors/mod.rs b/src/executors/mod.rs index 5f0933d7c7..26dff889c3 100644 --- a/src/executors/mod.rs +++ b/src/executors/mod.rs @@ -3,11 +3,13 @@ pub mod inmemory; use alloc::boxed::Box; use alloc::vec::Vec; -use crate::corpus::TestcaseMetadata; +use crate::corpus::Testcase; use crate::feedbacks::Feedback; use crate::inputs::Input; use crate::observers::Observer; use crate::AflError; +use alloc::rc::Rc; +use core::cell::RefCell; pub enum ExitKind { Ok, @@ -53,19 +55,21 @@ where self.run_target(input)?; self.post_exec_observers()?; - let mut metadatas: Vec> = vec![]; let mut rate_acc = 0; for feedback in self.feedbacks_mut() { - let (rate, meta) = feedback.is_interesting(input); - rate_acc += rate; - if let Some(m) = meta { - metadatas.push(m); - } + rate_acc += feedback.is_interesting(input)?; } if rate_acc >= 25 { + let testcase = Rc::new(RefCell::new(Testcase::new(input.clone()))); + for feedback in self.feedbacks_mut() { + feedback.append_metadata(testcase.clone())?; + } Ok(true) } else { + for feedback in self.feedbacks_mut() { + feedback.discard_metadata()?; + } Ok(false) } } diff --git a/src/feedbacks/mod.rs b/src/feedbacks/mod.rs index 4a73c399d2..810af1568a 100644 --- a/src/feedbacks/mod.rs +++ b/src/feedbacks/mod.rs @@ -1,21 +1,32 @@ extern crate num; -use alloc::boxed::Box; +use alloc::rc::Rc; use alloc::vec::Vec; use core::cell::RefCell; use core::marker::PhantomData; use num::Integer; -use crate::corpus::TestcaseMetadata; +use crate::corpus::Testcase; use crate::inputs::Input; use crate::observers::MapObserver; +use crate::AflError; pub trait Feedback where I: Input, { /// is_interesting should return the "Interestingness" from 0 to 255 (percent times 2.55) - fn is_interesting(&mut self, input: &I) -> (u32, Option>); + fn is_interesting(&mut self, input: &I) -> Result; + + /// Append to the testcase the generated metadata in case of a new corpus item + fn append_metadata(&mut self, _testcase: Rc>>) -> Result<(), AflError> { + Ok(()) + } + + /// Discard the stored metadata in case that the testcase is not added to the corpus + fn discard_metadata(&mut self) -> Result<(), AflError> { + Ok(()) + } } /// A Reducer function is used to aggregate values for the novelty search @@ -88,7 +99,7 @@ where O: MapObserver, I: Input, { - fn is_interesting(&mut self, _input: &I) -> (u32, Option>) { + fn is_interesting(&mut self, _input: &I) -> Result { let mut interesting = 0; // TODO: impl. correctly, optimize @@ -103,11 +114,11 @@ where *history = reduced; interesting += 25; if interesting >= 250 { - return (255, None); + return Ok(255); } } } - (interesting, None) + Ok(interesting) } } diff --git a/src/stages/mod.rs b/src/stages/mod.rs index 31f610df4d..2df141214e 100644 --- a/src/stages/mod.rs +++ b/src/stages/mod.rs @@ -3,20 +3,24 @@ pub use mutational::DefaultMutationalStage; use crate::corpus::testcase::Testcase; use crate::corpus::Corpus; +use crate::engines::State; +use crate::executors::Executor; use crate::inputs::Input; use crate::AflError; use alloc::rc::Rc; use core::cell::RefCell; -pub trait Stage +pub trait Stage where + S: State, C: Corpus, + E: Executor, I: Input, { /// Run the stage fn perform( &mut self, testcase: Rc>>, - corpus: &mut C, + state: &mut S, ) -> Result<(), AflError>; } diff --git a/src/stages/mutational.rs b/src/stages/mutational.rs index 51fff08ff6..c024d4e36a 100644 --- a/src/stages/mutational.rs +++ b/src/stages/mutational.rs @@ -3,6 +3,7 @@ use core::cell::RefCell; use core::marker::PhantomData; use crate::corpus::testcase::Testcase; +use crate::engines::State; use crate::executors::Executor; use crate::inputs::Input; use crate::mutators::Mutator; @@ -13,12 +14,13 @@ use crate::AflError; // TODO multi mutators stage -pub trait MutationalStage: Stage + HasRand +pub trait MutationalStage: Stage + HasRand where + M: Mutator, + S: State, + E: Executor, C: Corpus, I: Input, - M: Mutator, - E: Executor, { /// The mutator registered for this stage fn mutator(&self) -> &M; @@ -26,9 +28,6 @@ where /// The mutator registered for this stage (mutable) fn mutator_mut(&mut self) -> &mut M; - /// Rc Refcell to the executor - fn executor(&self) -> &Rc>; - /// 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 { @@ -39,7 +38,7 @@ where fn perform_mutational( &mut self, testcase: Rc>>, - corpus: &mut C, + state: &mut S, ) -> Result<(), AflError> { let num = self.iterations(); let input = testcase.borrow_mut().load_input()?.clone(); @@ -47,14 +46,14 @@ where for i in 0..num { let mut input_tmp = input.clone(); self.mutator_mut() - .mutate(corpus, &mut input_tmp, i as i32)?; + .mutate(state.corpus_mut(), &mut input_tmp, i as i32)?; - let interesting = self.executor().borrow_mut().evaluate_input(&input_tmp)?; + let interesting = state.evaluate_input(&input_tmp)?; self.mutator_mut().post_exec(interesting, i as i32)?; if interesting { - corpus.add(Testcase::new(input_tmp).into()); + state.corpus_mut().add(Testcase::new(input_tmp).into()); } } Ok(()) @@ -62,28 +61,25 @@ where } /// The default mutational stage -pub struct DefaultMutationalStage +pub struct DefaultMutationalStage where + M: Mutator, C: Corpus, I: Input, R: Rand, - M: Mutator, - E: Executor, { rand: Rc>, - executor: Rc>, mutator: M, _phantom_corpus: PhantomData, _phantom_input: PhantomData, } -impl HasRand for DefaultMutationalStage +impl HasRand for DefaultMutationalStage where + M: Mutator, C: Corpus, I: Input, R: Rand, - M: Mutator, - E: Executor, { type R = R; @@ -92,13 +88,14 @@ where } } -impl MutationalStage for DefaultMutationalStage +impl MutationalStage for DefaultMutationalStage where + M: Mutator, + S: State, C: Corpus, + E: Executor, I: Input, R: Rand, - M: Mutator, - E: Executor, { /// The mutator, added to this stage fn mutator(&self) -> &M { @@ -109,42 +106,37 @@ where fn mutator_mut(&mut self) -> &mut M { &mut self.mutator } - - fn executor(&self) -> &Rc> { - &self.executor - } } -impl Stage for DefaultMutationalStage +impl Stage for DefaultMutationalStage where + M: Mutator, + S: State, C: Corpus, + E: Executor, I: Input, R: Rand, - M: Mutator, - E: Executor, { fn perform( &mut self, testcase: Rc>>, - corpus: &mut C, + state: &mut S, ) -> Result<(), AflError> { - self.perform_mutational(testcase, corpus) + self.perform_mutational(testcase, state) } } -impl DefaultMutationalStage +impl DefaultMutationalStage where + M: Mutator, C: Corpus, I: Input, R: Rand, - M: Mutator, - E: Executor, { /// Creates a new default mutational stage - pub fn new(rand: &Rc>, executor: &Rc>, mutator: M) -> Self { + pub fn new(rand: &Rc>, mutator: M) -> Self { DefaultMutationalStage { rand: Rc::clone(rand), - executor: Rc::clone(executor), mutator: mutator, _phantom_corpus: PhantomData, _phantom_input: PhantomData,