diff --git a/src/corpus/mod.rs b/src/corpus/mod.rs index 2b13225a7e..aa43f1d628 100644 --- a/src/corpus/mod.rs +++ b/src/corpus/mod.rs @@ -1,55 +1,37 @@ pub mod testcase; -pub use testcase::{Testcase, SimpleTestcase}; +pub use testcase::Testcase; -use crate::utils::Rand; +use crate::utils::{Rand, HasRand}; +use crate::inputs::Input; use crate::AflError; use std::path::PathBuf; +pub trait HasEntriesVec { + /// Get the entries vector field + fn entries(&self) -> &Vec>>; + + /// Get the entries vector field (mutable) + fn entries_mut(&mut self) -> &mut Vec>>; +} + /// Corpus with all current testcases -pub trait Corpus { - /// Returns the number of elements - fn count(&self) -> usize; - - fn add(&mut self, entry: Box); - - /// Removes an entry from the corpus, returning it if it was present. - fn remove(&mut self, entry: &dyn Testcase) -> Option>; - - /// Gets a random entry - fn random_entry(&mut self) -> Result<&Box, AflError>; - - /// Gets the next entry - fn get(&mut self) -> Result<&Box, AflError>; -} - -pub struct BaseCorpus<'a, RandT: Rand> { - rand: &'a mut RandT, - entries: Vec>, - dir_path: PathBuf, -} - -impl Corpus for BaseCorpus<'_, RandT> { +pub trait Corpus : HasEntriesVec + HasRand { /// Returns the number of elements fn count(&self) -> usize { - self.entries.len() + self.entries().len() } - fn add(&mut self, mut entry: Box) { - if entry.get_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.set_filename(filename); - } - self.entries.push(entry); + /// Add an entry to the corpus + fn add(&mut self, mut entry: Box>) { + self.entries_mut().push(entry); } /// Removes an entry from the corpus, returning it if it was present. - fn remove(&mut self, entry: &dyn Testcase) -> Option> { + fn remove(&mut self, entry: &Testcase) -> Option>> { let mut i: usize = 0; let mut found = false; - for x in &self.entries { + for x in self.entries() { i = i + 1; if x.as_ref() as *const _ == entry as *const _ { found = true; @@ -59,24 +41,99 @@ impl Corpus for BaseCorpus<'_, RandT> { if !found { return None; } - Some(self.entries.remove(i)) + Some(self.entries_mut().remove(i)) } /// Gets a random entry - fn random_entry(&mut self) -> Result<&Box, AflError> { - let id = self.rand.below(self.entries.len() as u64) as usize; - Ok(self.entries.get_mut(id).unwrap()) + fn random_entry(&mut self) -> Result<&Box>, AflError> { + let id = self.rand_mut().below(self.entries().len() as u64) as usize; + Ok(self.entries_mut().get_mut(id).unwrap()) } - /// Gets the next entry - fn get(&mut self) -> Result<&Box, AflError> { + /// Gets the next entry (random by default) + fn get(&mut self) -> Result<&Box>, AflError> { self.random_entry() } } -impl BaseCorpus<'_, RandT> { - pub fn new<'a>(rand: &'a mut RandT, dir_path: PathBuf) -> BaseCorpus<'a, RandT> { - BaseCorpus { +pub struct InMemoryCorpus<'a, InputT: Input, RandT: Rand> { + rand: &'a mut RandT, + entries: Vec>>, +} + +impl HasEntriesVec for InMemoryCorpus<'_, InputT, RandT> { + fn entries(&self) -> &Vec>> { + &self.entries + } + fn entries_mut(&mut self) -> &mut Vec>>{ + &mut self.entries + } +} + +impl HasRand for InMemoryCorpus<'_, InputT, RandT> { + fn rand(&self) -> &Box { + &self.rand + } + fn rand_mut(&mut self) -> &mut Box { + &mut self.rand + } +} + +impl Corpus for InMemoryCorpus<'_, InputT, RandT> { + // Just use the default implementation +} + +impl InMemoryCorpus<'_, InputT, RandT> { + pub fn new<'a>(rand: &'a mut RandT) -> Self { + InMemoryCorpus { + rand: rand, + entries: vec![], + } + } +} + +pub struct OnDiskCorpus<'a, InputT: Input, RandT: Rand> { + rand: &'a mut RandT, + entries: Vec>>, + dir_path: PathBuf, +} + +impl HasEntriesVec for OnDiskCorpus<'_, InputT, RandT> { + fn entries(&self) -> &Vec>> { + &self.entries + } + fn entries_mut(&mut self) -> &mut Vec>>{ + &mut self.entries + } +} + +impl HasRand for OnDiskCorpus<'_, InputT, RandT> { + fn rand(&self) -> &Box { + &self.rand + } + fn rand_mut(&mut self) -> &mut Box { + &mut self.rand + } +} + +impl Corpus for OnDiskCorpus<'_, InputT, RandT> { + /// Add an entry and save it to disk + fn add(&mut self, mut entry: Box>) { + if entry.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() = filename; + } + self.entries.push(entry); + } + + // TODO save and remove files, cache, etc..., ATM use just InMemoryCorpus +} + +impl OnDiskCorpus<'_, InputT, RandT> { + pub fn new<'a>(rand: &'a mut RandT, dir_path: PathBuf) -> Self { + OnDiskCorpus { dir_path: dir_path, entries: vec![], rand: rand, @@ -84,35 +141,35 @@ impl BaseCorpus<'_, RandT> { } } -/// A queue-like corpus -pub struct QueueCorpus<'a, RandT: Rand> { - base: BaseCorpus<'a, RandT>, +/// A Queue-like corpus, wrapping an existing Corpus instance +pub struct QueueCorpus<'a, InputT: Input, RandT: Rand, CorpusT: Corpus> { + corpus: CorpusT, pos: usize, cycles: u64, } -impl Corpus for QueueCorpus<'_, RandT> { +impl<'a, InputT: Input, RandT: Rand, CorpusT: Corpus> Corpus for QueueCorpus<'_, InputT, RandT, CorpusT> { /// Returns the number of elements fn count(&self) -> usize { - self.base.count() + self.corpus.count() } - fn add(&mut self, entry: Box) { - self.base.add(entry); + fn add(&mut self, entry: Box>) { + self.corpus.add(entry); } /// Removes an entry from the corpus, returning it if it was present. - fn remove(&mut self, entry: &dyn Testcase) -> Option> { - self.base.remove(entry) + fn remove(&mut self, entry: &Testcase) -> Option>> { + self.corpus.remove(entry) } /// Gets a random entry - fn random_entry(&mut self) -> Result<&Box, AflError> { - self.base.random_entry() + fn random_entry(&mut self) -> Result<&Box>, AflError> { + self.corpus.random_entry() } /// Gets the next entry - fn get(&mut self) -> Result<&Box, AflError> { + fn get(&mut self) -> Result<&Box>, AflError> { if self.count() == 0 { return Err(AflError::Empty("Testcases".to_string())); } @@ -121,24 +178,24 @@ impl Corpus for QueueCorpus<'_, RandT> { self.cycles = self.cycles + 1; self.pos = 0; } - Ok(self.base.entries.get_mut(self.pos).unwrap()) + Ok(self.corpus.entries.get_mut(self.pos).unwrap()) } } -impl QueueCorpus<'_, RandT> { - pub fn new<'a>(rand: &'a mut RandT, dir_path: PathBuf) -> QueueCorpus<'a, RandT> { +impl<'a, InputT: Input, RandT: Rand, CorpusT: Corpus> QueueCorpus<'_, InputT, RandT, CorpusT> { + pub fn new(corpus: CorpusT) -> Self { QueueCorpus { - base: BaseCorpus::new(rand, dir_path), + corpus: corpus, cycles: 0, pos: 0, } } - pub fn get_cycles(&self) -> u64 { + pub fn cycles(&self) -> u64 { self.cycles } - pub fn get_pos(&self) -> usize { + pub fn pos(&self) -> usize { self.pos } } @@ -146,8 +203,7 @@ impl QueueCorpus<'_, RandT> { #[cfg(test)] mod tests { use crate::corpus::Corpus; - use crate::corpus::QueueCorpus; - use crate::corpus::SimpleTestcase; + use crate::corpus::{QueueCorpus, OnDiskCorpus}; use crate::corpus::Testcase; use crate::inputs::bytes::BytesInput; use crate::utils::Xoshiro256StarRand; @@ -158,13 +214,13 @@ mod tests { fn test_queuecorpus() { let mut rand = Xoshiro256StarRand::new(); - let mut q = QueueCorpus::new(&mut rand, PathBuf::from("fancy/path")); + 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(SimpleTestcase::new(i)); + let mut t = Box::new(Testcase::new(i)); t.set_filename(PathBuf::from("fancyfile")); q.add(t); let filename = q.get().unwrap().get_filename().unwrap().to_owned(); assert_eq!(filename, q.get().unwrap().get_filename().unwrap().to_owned()); - assert_eq!(filename, PathBuf::from("fancyfile")); + assert_eq!(filename, PathBuf::from("fancy/path/fancyfile")); } } diff --git a/src/corpus/testcase.rs b/src/corpus/testcase.rs index 86384f1b5d..65328db237 100644 --- a/src/corpus/testcase.rs +++ b/src/corpus/testcase.rs @@ -6,63 +6,82 @@ use std::path::PathBuf; pub trait TestcaseMetadata {} -pub trait Testcase { +/* +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<&Box, AflError>; + /// Get the input, if any + fn input(&self) -> &Option>; - fn get_input(&self) -> Option<& Box>; + /// Get the input, if any (mutable) + fn input_mut(&mut self) -> &mut Option>; - fn is_on_disk(&self) -> bool; + /// Get the filename, if any + fn filename(&self) -> &Option; - fn get_filename(&self) -> Option<& PathBuf>; - - fn set_filename(&mut self, filename: PathBuf); - - fn get_metadatas(&mut self) -> &mut HashMap>; + /// Get the filename, if any (mutable) + fn filename_mut(&mut self, filename: PathBuf) -> &mut &Option; + /// Get all the metadatas into an HashMap + fn metadatas(&mut self) -> &mut HashMap>; } +*/ #[derive(Default)] -pub struct SimpleTestcase { - input: Option>, - // is_on_disk: bool, // not needed, look at the Option +pub struct Testcase { + input: Option>, filename: Option, metadatas: HashMap>, } -impl Testcase for SimpleTestcase { - fn load_input(&mut self) -> Result<&Box, AflError> { +impl Testcase { + /// 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> { // TODO: Implement cache to disk self.input.as_ref().ok_or(AflError::NotImplemented("load_input".to_string())) } - fn get_input(&self) -> Option<& Box> { - self.input.as_ref() + /// Get the input, if any + pub fn input(&self) -> &Option> { + &self.input } - fn is_on_disk(&self) -> bool { - !self.input.is_some() && self.filename.is_some() + /// Get the input, if any (mutable) + pub fn input_mut(&mut self) -> &mut Option> { + &mut self.input } - fn get_filename(&self) -> Option<& PathBuf> { - self.filename.as_ref() + /// Get the filename, if any + pub fn filename(&self) -> &Option { + &self.filename } - fn set_filename(&mut self, filename: PathBuf) { - self.filename = Some(filename) + /// Get the filename, if any (mutable) + pub fn filename_mut(&mut self, filename: PathBuf) -> &mut Option { + &mut self.filename } - fn get_metadatas(&mut self) -> &mut HashMap> { + /// Get all the metadatas into an HashMap + pub fn metadatas(&mut self) -> &mut HashMap> { &mut self.metadatas } -} -impl SimpleTestcase { - pub fn new(input: Box) -> Self { - SimpleTestcase { + /// Create a new DefaultTestcase instace given an input + pub fn new(input: Box) -> Self { + Testcase { input: Some(input), filename: None, metadatas: HashMap::default(), } } + + /// Create a new DefaultTestcase instace given an input and a filename + pub fn new_with_filename(input: Box, filename: &PathBuf) -> Self { + Testcase { + input: Some(input), + filename: filename, + metadatas: HashMap::default(), + } + } } diff --git a/src/mutators/mod.rs b/src/mutators/mod.rs index 8887039322..f045fa1299 100644 --- a/src/mutators/mod.rs +++ b/src/mutators/mod.rs @@ -1,21 +1,20 @@ use crate::inputs::Input; -use crate::utils::Rand; +use crate::utils::{Rand, HasRand}; use crate::corpus::Corpus; use crate::AflError; pub mod scheduled; -pub use scheduled::ScheduledMutator; - -pub trait Mutator { - - fn rand(&mut self) -> &mut Box; +pub use scheduled::{ComposedByMutations, ScheduledMutator, HavocBytesMutator}; +pub trait Mutator : HasRand { + /// Mutate a given input fn mutate(&mut self, input: &mut InputT, 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> { Ok(()) } - fn corpus(&mut self) -> &mut Option>; - + /// Get the associated corpus, if any + fn corpus(&mut self) -> &mut Option>>; } diff --git a/src/mutators/scheduled.rs b/src/mutators/scheduled.rs index a474512122..cd26df644c 100644 --- a/src/mutators/scheduled.rs +++ b/src/mutators/scheduled.rs @@ -1,5 +1,5 @@ use crate::mutators::Mutator; -use crate::utils::Rand; +use crate::utils::{Rand, HasRand}; use crate::corpus::Corpus; use crate::inputs::{Input, HasBytesVec}; use crate::AflError; @@ -20,10 +20,10 @@ pub trait ComposedByMutations { fn add_mutation(&mut self, mutation: MutationFunction); } -pub trait ScheduledMutator: Mutator + ComposedByMutations { +pub trait ScheduledMutator: Mutator + ComposedByMutations { /// Computer the number of iterations used to apply stacked mutations fn iterations(&mut self, _input: &InputT) -> u64 { - 1 << (1 + self.rand().below(7)) + 1 << (1 + self.rand_mut().below(7)) } /// Get the next mutation to apply @@ -34,7 +34,7 @@ pub trait ScheduledMutator: Mutator + ComposedByMutation } let idx; { - idx = self.rand().below(count) as usize; + idx = self.rand_mut().below(count) as usize; } self.mutation_by_idx(idx) } @@ -50,13 +50,22 @@ pub trait ScheduledMutator: Mutator + ComposedByMutation } } -pub struct DefaultScheduledMutator { - rand: Box, - corpus: Option>, +pub struct DefaultScheduledMutator { + rand: Box, + corpus: Option>>, mutations: Vec> } -impl ComposedByMutations for DefaultScheduledMutator { +impl HasRand for DefaultScheduledMutator { + fn rand(&self) -> &Box { + &self.rand + } + fn rand_mut(&mut self) -> &mut Box { + &mut self.rand + } +} + +impl ComposedByMutations for DefaultScheduledMutator { fn mutation_by_idx(&self, index: usize) -> Result, AflError> { if index >= self.mutations.len() { return Err(AflError::Unknown("oob".to_string())); @@ -73,15 +82,11 @@ impl ComposedByMutations for DefaultScheduledMutator ScheduledMutator for DefaultScheduledMutator { +impl ScheduledMutator for DefaultScheduledMutator { // Just use the default methods } -impl Mutator for DefaultScheduledMutator { - fn rand(&mut self) -> &mut Box { - &mut self.rand - } - +impl Mutator for DefaultScheduledMutator { fn mutate(&mut self, input: &mut InputT, _stage_idx: i32) -> Result<(), AflError> { self.scheduled_mutate(input, _stage_idx) } @@ -91,9 +96,9 @@ impl Mutator for DefaultScheduledMutator { } } -impl DefaultScheduledMutator { +impl DefaultScheduledMutator { /// Create a new DefaultScheduledMutator instance without mutations and corpus - pub fn new(rand: Box) -> Self { + pub fn new(rand: Box) -> Self { DefaultScheduledMutator { rand: rand, corpus: None, @@ -102,7 +107,7 @@ impl DefaultScheduledMutator { } /// Create a new DefaultScheduledMutator instance specifying mutations and corpus too - pub fn new_all(rand: Box, corpus: Option>, mutations: Vec>) -> Self { + pub fn new_all(rand: Box, corpus: Option>, mutations: Vec>) -> Self { DefaultScheduledMutator { rand: rand, corpus: corpus, @@ -112,23 +117,19 @@ impl DefaultScheduledMutator { } /// Bitflip mutation for inputs with a bytes vector -pub fn mutation_bitflip, InputT: Input + HasBytesVec>(mutator: &mut MutatorT, input: &mut InputT) -> Result<(), AflError> { - let bit = mutator.rand().below(input.bytes().len() as u64) as usize; +pub fn mutation_bitflip, InputT: Input + HasBytesVec, RandT: Rand>(mutator: &mut MutatorT, input: &mut InputT) -> Result<(), AflError> { + let bit = mutator.rand_mut().below(input.bytes().len() as u64) as usize; input.bytes_mut()[bit >> 3] ^= (128 >> (bit & 7)) as u8; Ok(()) } /// Schedule some selected byte level mutations given a ScheduledMutator type -pub struct HavocBytesMutator, InputT: Input + HasBytesVec> { +pub struct HavocBytesMutator> { scheduled: ScheduledMutatorT, _phantom: PhantomData } -impl, InputT: Input + HasBytesVec> Mutator for HavocBytesMutator { - fn rand(&mut self) -> &mut Box { - self.scheduled.rand() - } - +impl> Mutator for HavocBytesMutator { fn mutate(&mut self, input: &mut InputT, stage_idx: i32) -> Result<(), AflError> { self.scheduled.mutate(input, stage_idx) } @@ -138,7 +139,16 @@ impl, InputT: Input + HasBytesVec> M } } -impl, InputT: Input + HasBytesVec> HavocBytesMutator { +impl> HasRand for HavocBytesMutator { + fn rand(&self) -> &Box { + self.scheduled.rand() + } + fn rand_mut(&mut self) -> &mut Box { + self.scheduled.rand_mut() + } +} + +impl> HavocBytesMutator { /// Create a new HavocBytesMutator instance given a ScheduledMutator to wrap pub fn new(mut scheduled: ScheduledMutatorT) -> Self { scheduled.add_mutation(mutation_bitflip); @@ -149,9 +159,9 @@ impl, InputT: Input + HasBytesVec> H } } -impl HavocBytesMutator, InputT> { +impl HavocBytesMutator> { /// Create a new HavocBytesMutator instance wrapping DefaultScheduledMutator - pub fn new_default(rand: Box) -> Self { + pub fn new_default(rand: Box) -> Self { let mut scheduled = DefaultScheduledMutator::new(rand); scheduled.add_mutation(mutation_bitflip); HavocBytesMutator { diff --git a/src/stages/mod.rs b/src/stages/mod.rs index 08e8c61c9e..c4f3e4962d 100644 --- a/src/stages/mod.rs +++ b/src/stages/mod.rs @@ -3,5 +3,5 @@ use crate::inputs::Input; use crate::AflError; /// Stages pub trait Stage { - fn perform(&mut self, input: &dyn Input, entry: &mut dyn Testcase) -> Result<(), AflError>; + fn perform(&mut self, input: &dyn Input, entry: &mut Testcase) -> Result<(), AflError>; } diff --git a/src/utils.rs b/src/utils.rs index e53e4fd51e..ed4e927d21 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -39,6 +39,15 @@ pub trait Rand: Debug { } } +/// Has a Rand box field +pub trait HasRand { + /// Get the hold Rand instance + fn rand(&self) -> &Box; + + /// Get the hold Rand instance (mutable) + fn rand_mut(&mut self) -> &mut Box; +} + const HASH_CONST: u64 = 0xa5b35705; /// XXH3 Based, hopefully speedy, rnd implementation