From 4fa50ebfaf68f0c4110feca3399f0f2a415c6574 Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Mon, 16 Nov 2020 11:18:44 +0100 Subject: [PATCH] lifetime fixes and generator --- src/engines/mod.rs | 7 +--- src/generators/mod.rs | 87 +++++++++++++++++++++++++++++++++++++++ src/lib.rs | 1 + src/mutators/scheduled.rs | 16 +++---- src/stages/mutational.rs | 2 +- 5 files changed, 99 insertions(+), 14 deletions(-) create mode 100644 src/generators/mod.rs diff --git a/src/engines/mod.rs b/src/engines/mod.rs index 5f714c94dc..6b71cc7965 100644 --- a/src/engines/mod.rs +++ b/src/engines/mod.rs @@ -86,8 +86,6 @@ where mod tests { use alloc::boxed::Box; - use alloc::rc::Rc; - use core::cell::RefCell; use crate::corpus::{Corpus, InMemoryCorpus, Testcase}; use crate::engines::{DefaultEngine, Engine}; @@ -106,13 +104,12 @@ mod tests { #[test] fn test_engine() { - let rand: Rc<_> = DefaultRand::new(0).into(); + let rand = DefaultRand::new(0).into(); let mut corpus = InMemoryCorpus::::new(&rand); let testcase = Testcase::new(vec![0; 4]).into(); corpus.add(testcase); - let executor: Rc>> = - InMemoryExecutor::new(harness).into(); + let executor = InMemoryExecutor::::new(harness).into(); let mut engine = DefaultEngine::new(); let mut mutator = DefaultScheduledMutator::new(&rand); mutator.add_mutation(mutation_bitflip); diff --git a/src/generators/mod.rs b/src/generators/mod.rs new file mode 100644 index 0000000000..60813d59b4 --- /dev/null +++ b/src/generators/mod.rs @@ -0,0 +1,87 @@ +use crate::inputs::bytes::BytesInput; +use crate::inputs::Input; +use crate::utils::{HasRand, Rand}; +use crate::AflError; +use alloc::rc::Rc; +use core::cell::RefCell; +use core::cmp::min; + +pub trait Generator: HasRand +where + I: Input, +{ + /// Generate a new input + fn generate(&mut self) -> Result; + + /// Generate a new dummy input + fn generate_dummy(&self) -> I; +} + +const DUMMY_BYTES_SIZE: usize = 64; + +pub struct RandBytesGenerator { + rand: Rc>, + max_size: usize, +} + +impl HasRand 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(); + Ok(BytesInput::new(random_bytes)) + } + + fn generate_dummy(&self) -> BytesInput { + let size = min(self.max_size, DUMMY_BYTES_SIZE); + BytesInput::new(vec![0; size]) + } +} + +pub struct RandPrintablesGenerator { + rand: Rc>, + max_size: usize, +} + +impl HasRand for RandPrintablesGenerator +where + R: Rand, +{ + type R = R; + + fn rand(&self) -> &Rc> { + &self.rand + } +} + +impl Generator for RandPrintablesGenerator +where + R: Rand, +{ + fn generate(&mut self) -> Result { + let size = self.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]) + .collect(); + Ok(BytesInput::new(random_bytes)) + } + + fn generate_dummy(&self) -> BytesInput { + let size = min(self.max_size, DUMMY_BYTES_SIZE); + BytesInput::new(vec!['0' as u8; size]) + } +} diff --git a/src/lib.rs b/src/lib.rs index 3e3d7c868f..c143804eda 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,6 +11,7 @@ pub mod inputs; pub mod monitors; pub mod mutators; pub mod observers; +pub mod generators; pub mod stages; pub mod utils; diff --git a/src/mutators/scheduled.rs b/src/mutators/scheduled.rs index 78beb269ea..d68430d700 100644 --- a/src/mutators/scheduled.rs +++ b/src/mutators/scheduled.rs @@ -66,7 +66,7 @@ where } } -pub struct DefaultScheduledMutator<'a, C, I, R> +pub struct DefaultScheduledMutator where C: Corpus, I: Input, @@ -76,7 +76,7 @@ where mutations: Vec>, } -impl<'a, C, I, R> HasRand for DefaultScheduledMutator<'_, C, I, R> +impl HasRand for DefaultScheduledMutator where C: Corpus, I: Input, @@ -89,7 +89,7 @@ where } } -impl<'a, C, I, R> Mutator for DefaultScheduledMutator<'_, C, I, R> +impl Mutator for DefaultScheduledMutator where C: Corpus, I: Input, @@ -100,7 +100,7 @@ where } } -impl<'a, C, I, R> ComposedByMutations for DefaultScheduledMutator<'_, C, I, R> +impl ComposedByMutations for DefaultScheduledMutator where C: Corpus, I: Input, @@ -122,7 +122,7 @@ where } } -impl<'a, C, I, R> ScheduledMutator for DefaultScheduledMutator<'_, C, I, R> +impl ScheduledMutator for DefaultScheduledMutator where C: Corpus, I: Input, @@ -131,7 +131,7 @@ where // Just use the default methods } -impl<'a, C, I, R> DefaultScheduledMutator<'a, C, I, R> +impl DefaultScheduledMutator where C: Corpus, I: Input, @@ -303,7 +303,7 @@ where } } -impl<'a, C, I, R> HavocBytesMutator> +impl HavocBytesMutator> where C: Corpus, I: Input + HasBytesVec, @@ -311,7 +311,7 @@ where { /// Create a new HavocBytesMutator instance wrapping DefaultScheduledMutator pub fn new_default(rand: &Rc>) -> Self { - let mut scheduled = DefaultScheduledMutator::<'a, C, I, R>::new(rand); + let mut scheduled = DefaultScheduledMutator::::new(rand); scheduled.add_mutation(mutation_bitflip); HavocBytesMutator { scheduled: scheduled, diff --git a/src/stages/mutational.rs b/src/stages/mutational.rs index 749548b0c7..80ad11abf9 100644 --- a/src/stages/mutational.rs +++ b/src/stages/mutational.rs @@ -11,7 +11,7 @@ use crate::stages::Stage; use crate::utils::{HasRand, Rand}; use crate::AflError; -// TODO create HasMutatorsVec trait +// TODO multi mutators stage pub trait MutationalStage: Stage + HasRand where