From 2e711360c15742ccf7e8deee549f4c5e1e59a3b3 Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Wed, 4 Nov 2020 15:19:59 +0100 Subject: [PATCH] pass Testcase inside Rc+RefCell --- src/corpus/mod.rs | 59 ++++++++++++++++--------------- src/corpus/testcase.rs | 25 +++++++------ src/engines/aflengine.rs | 3 +- src/engines/mod.rs | 13 ++++++- src/executors/inmemory.rs | 1 + src/executors/mod.rs | 1 - src/inputs/mod.rs | 3 +- src/mutators/mod.rs | 3 +- src/mutators/scheduled.rs | 2 +- src/stages/mod.rs | 18 ++++++++-- src/stages/mutational.rs | 74 ++++++++++++++++++++++++++++++++++----- 11 files changed, 145 insertions(+), 57 deletions(-) diff --git a/src/corpus/mod.rs b/src/corpus/mod.rs index 90ebe8f297..0715a68096 100644 --- a/src/corpus/mod.rs +++ b/src/corpus/mod.rs @@ -1,5 +1,5 @@ pub mod testcase; -pub use testcase::Testcase; +pub use testcase::{Testcase, TestcaseMetadata}; use crate::utils::{Rand, HasRand}; use crate::inputs::Input; @@ -7,13 +7,15 @@ use crate::AflError; use std::path::PathBuf; use std::marker::PhantomData; +use std::cell::RefCell; +use std::rc::Rc; pub trait HasEntriesVec where I: Input { /// Get the entries vector field - fn entries(&self) -> &Vec>>; + fn entries(&self) -> &Vec>>>; /// Get the entries vector field (mutable) - fn entries_mut(&mut self) -> &mut Vec>>; + fn entries_mut(&mut self) -> &mut Vec>>>; } /// Corpus with all current testcases @@ -25,17 +27,17 @@ pub trait Corpus : HasEntriesVec + HasRand where I: Input { /// Add an entry to the corpus #[allow(unused_mut)] - fn add(&mut self, mut entry: Box>) { + fn add(&mut self, mut entry: Rc>>) { self.entries_mut().push(entry); } /// Removes an entry from the corpus, returning it if it was present. - fn remove(&mut self, entry: &Testcase) -> Option>> { + fn remove(&mut self, entry: &Testcase) -> Option>>> { let mut i: usize = 0; let mut found = false; for x in self.entries() { i = i + 1; - if x.as_ref() as *const _ == entry as *const _ { + if &*x.borrow() as *const _ == entry as *const _ { // TODO check if correct found = true; break; } @@ -47,28 +49,28 @@ pub trait Corpus : HasEntriesVec + HasRand where I: Input { } /// Gets a random entry - fn random_entry(&mut self) -> Result<&Box>, AflError> { + fn random_entry(&mut self) -> Result<&Rc>>, AflError> { let len = { self.entries().len() }; let id = self.rand_mut().below(len as u64) as usize; Ok(self.entries_mut().get_mut(id).unwrap()) } /// Gets the next entry (random by default) - fn get(&mut self) -> Result<&Box>, AflError> { + fn get(&mut self) -> Result<&Rc>>, AflError> { self.random_entry() } } pub struct InMemoryCorpus<'a, I, R> where I: Input, R: Rand { rand: &'a mut R, - entries: Vec>> + entries: Vec>>> } impl HasEntriesVec for InMemoryCorpus<'_, I, R> where I: Input, R: Rand { - fn entries(&self) -> &Vec>> { + fn entries(&self) -> &Vec>>> { &self.entries } - fn entries_mut(&mut self) -> &mut Vec>>{ + fn entries_mut(&mut self) -> &mut Vec>>>{ &mut self.entries } } @@ -99,15 +101,15 @@ impl<'a, I, R> InMemoryCorpus<'a, I, R> where I: Input, R: Rand { pub struct OnDiskCorpus<'a, I, R> where I: Input, R: Rand { rand: &'a mut R, - entries: Vec>>, + entries: Vec>>>, dir_path: PathBuf, } impl HasEntriesVec for OnDiskCorpus<'_, I, R> where I: Input, R: Rand { - fn entries(&self) -> &Vec>> { + fn entries(&self) -> &Vec>>> { &self.entries } - fn entries_mut(&mut self) -> &mut Vec>>{ + fn entries_mut(&mut self) -> &mut Vec>>>{ &mut self.entries } } @@ -125,12 +127,12 @@ impl HasRand for OnDiskCorpus<'_, I, R> where I: Input, R: Rand { impl Corpus for OnDiskCorpus<'_, I, R> where I: Input, R: Rand { /// Add an entry and save it to disk - fn add(&mut self, mut entry: Box>) { - if *entry.filename() == None { + fn add(&mut self, entry: Rc>>) { + if *entry.borrow().filename() == None { // TODO walk entry metadatas to ask for pices of filename (e.g. :havoc in AFL) let filename = &(String::from("id:") + &self.entries.len().to_string()); let filename = self.dir_path.join(filename); - *entry.filename_mut() = Some(filename); + *entry.borrow_mut().filename_mut() = Some(filename); } self.entries.push(entry); } @@ -157,10 +159,10 @@ pub struct QueueCorpus where I: Input, C: Corpus { } impl<'a, I, C> HasEntriesVec for QueueCorpus where I: Input, C: Corpus { - fn entries(&self) -> &Vec>> { + fn entries(&self) -> &Vec>>> { self.corpus.entries() } - fn entries_mut(&mut self) -> &mut Vec>>{ + fn entries_mut(&mut self) -> &mut Vec>>>{ self.corpus.entries_mut() } } @@ -182,22 +184,22 @@ impl<'a, I, C> Corpus for QueueCorpus where I: Input, C: Corpus { self.corpus.count() } - fn add(&mut self, entry: Box>) { + fn add(&mut self, entry: Rc>>) { self.corpus.add(entry); } /// Removes an entry from the corpus, returning it if it was present. - fn remove(&mut self, entry: &Testcase) -> Option>> { + fn remove(&mut self, entry: &Testcase) -> Option>>> { self.corpus.remove(entry) } /// Gets a random entry - fn random_entry(&mut self) -> Result<&Box>, AflError> { + fn random_entry(&mut self) -> Result<&Rc>>, AflError> { self.corpus.random_entry() } /// Gets the next entry - fn get(&mut self) -> Result<&Box>, AflError> { + fn get(&mut self) -> Result<&Rc>>, AflError> { if self.corpus.count() == 0 { return Err(AflError::Empty("Testcases".to_string())); } @@ -238,18 +240,19 @@ mod tests { use crate::utils::Xoshiro256StarRand; use std::path::PathBuf; + use std::cell::RefCell; + use std::rc::Rc; #[test] fn test_queuecorpus() { let mut rand = Xoshiro256StarRand::new(); let mut q = QueueCorpus::new(OnDiskCorpus::new(&mut rand, PathBuf::from("fancy/path"))); - let i = Box::new(BytesInput::new(vec![0; 4])); - let mut t = Box::new(Testcase::new(i)); - *t.filename_mut() = Some(PathBuf::from("fancyfile")); + let i = BytesInput::new(vec![0; 4]); + let t = Rc::new(RefCell::new(Testcase::new_with_filename(i, PathBuf::from("fancyfile")))); q.add(t); - let filename = q.get().unwrap().filename().as_ref().unwrap().to_owned(); - assert_eq!(filename, q.get().unwrap().filename().as_ref().unwrap().to_owned()); + let filename = q.get().unwrap().borrow().filename().as_ref().unwrap().to_owned(); + assert_eq!(filename, q.get().unwrap().borrow().filename().as_ref().unwrap().to_owned()); assert_eq!(filename, PathBuf::from("fancy/path/fancyfile")); } } diff --git a/src/corpus/testcase.rs b/src/corpus/testcase.rs index 1a37139aec..44d4d30718 100644 --- a/src/corpus/testcase.rs +++ b/src/corpus/testcase.rs @@ -9,13 +9,13 @@ pub trait TestcaseMetadata {} /* pub trait TestcaseTrait { /// Make sure to return a valid input instance loading it from disk if not in memory - fn load_input(&mut self) -> Result<&Box, AflError>; + fn load_input(&mut self) -> Result<&I, AflError>; /// Get the input, if any - fn input(&self) -> &Option>; + fn input(&self) -> &Option; /// Get the input, if any (mutable) - fn input_mut(&mut self) -> &mut Option>; + fn input_mut(&mut self) -> &mut Option; /// Get the filename, if any fn filename(&self) -> &Option; @@ -30,28 +30,31 @@ pub trait TestcaseTrait { #[derive(Default)] pub struct Testcase where I: Input { - input: Option>, + input: Option, // TODO remove box filename: Option, metadatas: HashMap>, } impl Testcase where I: Input { /// Make sure to return a valid input instance loading it from disk if not in memory - pub fn load_input(&mut self) -> Result<&Box, AflError> { + pub fn load_input(&mut self) -> Result<&I, AflError> { // TODO: Implement cache to disk - self.input.as_ref().ok_or(AflError::NotImplemented("load_input".to_string())) + match self.input.as_ref() { + Some(i) => Ok(i), + None => Err(AflError::NotImplemented("load_input".to_string())) + } } /// Get the input, if any - pub fn input(&self) -> &Option> { + pub fn input(&self) -> &Option { &self.input } /// Get the input, if any (mutable) - pub fn input_mut(&mut self) -> &mut Option> { + pub fn input_mut(&mut self) -> &mut Option { &mut self.input } /// Set the input - pub fn set_input(&mut self, input: Option>) { + pub fn set_input(&mut self, input: Option) { self.input = input; } @@ -74,7 +77,7 @@ impl Testcase where I: Input { } /// Create a new DefaultTestcase instace given an input - pub fn new(input: Box) -> Self { + pub fn new(input: I) -> Self { Testcase { input: Some(input), filename: None, @@ -83,7 +86,7 @@ impl Testcase where I: Input { } /// Create a new DefaultTestcase instace given an input and a filename - pub fn new_with_filename(input: Box, filename: PathBuf) -> Self { + pub fn new_with_filename(input: I, filename: PathBuf) -> Self { Testcase { input: Some(input), filename: Some(filename), diff --git a/src/engines/aflengine.rs b/src/engines/aflengine.rs index 28635e72f6..5722d7cd6e 100644 --- a/src/engines/aflengine.rs +++ b/src/engines/aflengine.rs @@ -7,7 +7,7 @@ use crate::feedbacks::Feedback; use crate::monitors::Monitor; use crate::stages::Stage; use crate::utils::Rand; - +/* pub struct AflEngine<'a, I: Input> { pub rand: &'a mut dyn Rand, pub feedbacks: Vec>>, @@ -27,3 +27,4 @@ pub struct AflEngine<'a, I: Input> { } impl Engine<'_> for AflEngine<'_, I> {} +*/ \ No newline at end of file diff --git a/src/engines/mod.rs b/src/engines/mod.rs index 3df483b1eb..ba01edd530 100644 --- a/src/engines/mod.rs +++ b/src/engines/mod.rs @@ -1,3 +1,14 @@ pub mod aflengine; -pub trait Engine<'a> {} +use crate::AflError; +use crate::inputs::Input; +use crate::corpus::testcase::Testcase; + +use std::cell::RefCell; +use std::rc::Rc; + +pub trait Engine<'a, I> where I: Input { + + fn execute(&mut self, input: &mut I, entry: Rc>>) -> Result; + +} diff --git a/src/executors/inmemory.rs b/src/executors/inmemory.rs index a2c092bedc..b6e8c8c0c3 100644 --- a/src/executors/inmemory.rs +++ b/src/executors/inmemory.rs @@ -177,6 +177,7 @@ mod tests { use crate::observers::Observer; use crate::AflError; + #[derive(Clone)] struct NopInput {} impl Input for NopInput { fn serialize(&self) -> Result<&[u8], AflError> { diff --git a/src/executors/mod.rs b/src/executors/mod.rs index a99b92f967..7b2c74d921 100644 --- a/src/executors/mod.rs +++ b/src/executors/mod.rs @@ -1,5 +1,4 @@ pub mod inmemory; -pub use inmemory::{InMemoryExecutor}; use crate::inputs::Input; use crate::observers::Observer; diff --git a/src/inputs/mod.rs b/src/inputs/mod.rs index a305e4beef..79ec14e0d7 100644 --- a/src/inputs/mod.rs +++ b/src/inputs/mod.rs @@ -5,10 +5,11 @@ use std::fs::File; use std::io::Read; use std::io::Write; use std::path::PathBuf; +use std::clone::Clone; use crate::AflError; -pub trait Input { +pub trait Input : Clone { fn to_file(&self, path: &PathBuf) -> Result<(), AflError> { let mut file = File::create(path)?; file.write_all(self.serialize()?)?; diff --git a/src/mutators/mod.rs b/src/mutators/mod.rs index 9bb40b727a..9385afb6e3 100644 --- a/src/mutators/mod.rs +++ b/src/mutators/mod.rs @@ -4,7 +4,6 @@ use crate::corpus::Corpus; use crate::AflError; pub mod scheduled; -pub use scheduled::{ComposedByMutations, ScheduledMutator, HavocBytesMutator}; pub trait HasOptionCorpus where I: Input { type C : Corpus; @@ -19,7 +18,7 @@ pub trait HasOptionCorpus where I: Input { fn set_corpus(&mut self, corpus: Option>); } -pub trait Mutator : HasRand + HasOptionCorpus where I: Input { +pub trait Mutator : HasRand where I: Input { /// Mutate a given input fn mutate(&mut self, input: &mut I, stage_idx: i32) -> Result<(), AflError>; diff --git a/src/mutators/scheduled.rs b/src/mutators/scheduled.rs index b2ff7f41e9..212f1c78d6 100644 --- a/src/mutators/scheduled.rs +++ b/src/mutators/scheduled.rs @@ -20,7 +20,7 @@ pub trait ComposedByMutations where I: Input { fn add_mutation(&mut self, mutation: MutationFunction); } -pub trait ScheduledMutator: Mutator + ComposedByMutations where I: Input { +pub trait ScheduledMutator: Mutator + HasOptionCorpus + ComposedByMutations where I: Input { /// Compute the number of iterations used to apply stacked mutations fn iterations(&mut self, _input: &I) -> u64 { 1 << (1 + self.rand_mut().below(7)) diff --git a/src/stages/mod.rs b/src/stages/mod.rs index f24ddac29c..06c25616c3 100644 --- a/src/stages/mod.rs +++ b/src/stages/mod.rs @@ -1,7 +1,19 @@ +pub mod mutational; + use crate::corpus::Testcase; use crate::inputs::Input; +use crate::engines::Engine; use crate::AflError; -/// Stages -pub trait Stage { - fn perform(&mut self, input: &dyn Input, entry: &mut Testcase) -> Result<(), AflError>; + +pub trait HasEngine<'a, I> where I: Input { + type E : Engine<'a, I>; + + fn engine(&self) -> &Self::E; + + fn engine_mut(&mut self) -> &mut Self::E; +} + +pub trait Stage<'a, I> : HasEngine<'a, I> where I: Input { + /// Run the stage + fn perform(&mut self, entry: &mut Testcase) -> Result<(), AflError>; } diff --git a/src/stages/mutational.rs b/src/stages/mutational.rs index 29cade474d..5de64cdf8e 100644 --- a/src/stages/mutational.rs +++ b/src/stages/mutational.rs @@ -1,16 +1,74 @@ -use std::Vec; +use crate::AflError; use crate::mutators::Mutator; use crate::inputs::Input; +use crate::utils::{Rand, HasRand}; +use crate::stages::{Stage, HasEngine}; +use crate::corpus::testcase::Testcase; +use crate::engines::Engine; -pub struct MutationalStage { - mutators: Vec>; -} +use std::cell::RefCell; +use std::rc::Rc; -impl Stage for MutationalStage { +pub trait MutationalStage<'a, I> : Stage<'a, I> + HasRand where I: Input { + fn mutators(&self) -> &Vec>>; - fn Perform(&mut self, input: &Input, entry: &mut Entry) -> Result<(), AflError> { - // TODO: Implement me - Err(AflError::NotImplemented("Stage does not perform yet")); + fn mutators_mut(&mut self) -> &mut Vec>>; + + fn add_mutator(&mut self, mutator: Box>) { + self.mutators_mut().push(mutator); } + fn iterations(&mut self) -> usize { + 1 + self.rand_mut().below(128) as usize + } + + fn perform_mutational(&mut self, entry: Rc>>) -> Result<(), AflError> { + let num = self.iterations(); + let mut input = entry.borrow_mut().load_input()?.clone(); + + for i in 0..num { + for m in self.mutators_mut() { + m.mutate(&mut input, i as i32)?; + } + + let interesting = self.engine_mut().execute(&mut input, entry.clone())?; + + for m in self.mutators_mut() { + m.post_exec(interesting, i as i32)?; + } + + input = entry.borrow_mut().load_input()?.clone(); + } + + Ok(()) + } +} + +pub struct DefaultMutationalStage<'a, I, R, E> where I: Input, R: Rand, E: Engine<'a, I> { + rand: &'a mut R, + engine: &'a mut E, + mutators: Vec>> +} + +impl<'a, I, R, E> HasRand for DefaultMutationalStage<'a, I, R, E> where I: Input, R: Rand, E: Engine<'a, I> { + type R = R; + + fn rand(&self) -> &Self::R { + &self.rand + } + fn rand_mut(&mut self) -> &mut Self::R { + &mut self.rand + } +} + +impl<'a, I, R, E> HasEngine<'a, I> for DefaultMutationalStage<'a, I, R, E> where I: Input, R: Rand, E: Engine<'a, I> { + type E = E; + + fn engine(&self) -> &Self::E { + self.engine + } + + fn engine_mut(&mut self) -> &mut Self::E { + self.engine + } } \ No newline at end of file