diff --git a/src/corpus/mod.rs b/src/corpus/mod.rs index 81fc1d3cae..3a4517bced 100644 --- a/src/corpus/mod.rs +++ b/src/corpus/mod.rs @@ -58,7 +58,7 @@ where /// Gets a random entry fn random_entry(&mut self) -> Result>>, AflError> { let len = { self.entries().len() }; - let id = self.rand_mut().below(len as u64) as usize; + let id = self.rand_below(len as u64) as usize; Ok(self.entries()[id].clone()) } @@ -68,16 +68,16 @@ where } } -pub struct InMemoryCorpus<'a, I, R> +pub struct InMemoryCorpus where I: Input, R: Rand, { - rand: &'a mut R, + rand: Rc>, entries: Vec>>>, } -impl HasEntriesVec for InMemoryCorpus<'_, I, R> +impl HasEntriesVec for InMemoryCorpus where I: Input, R: Rand, @@ -90,22 +90,19 @@ where } } -impl HasRand for InMemoryCorpus<'_, I, R> +impl HasRand for InMemoryCorpus where I: Input, R: Rand, { type R = R; - fn rand(&self) -> &Self::R { + fn rand(&self) -> &Rc> { &self.rand } - fn rand_mut(&mut self) -> &mut Self::R { - &mut self.rand - } } -impl Corpus for InMemoryCorpus<'_, I, R> +impl Corpus for InMemoryCorpus< I, R> where I: Input, R: Rand, @@ -113,30 +110,30 @@ where // Just use the default implementation } -impl<'a, I, R> InMemoryCorpus<'a, I, R> +impl InMemoryCorpus where I: Input, R: Rand, { - pub fn new(rand: &'a mut R) -> Self { + pub fn new(rand: &Rc>) -> Self { InMemoryCorpus { - rand: rand, + rand: Rc::clone(rand), entries: vec![], } } } -pub struct OnDiskCorpus<'a, I, R> +pub struct OnDiskCorpus where I: Input, R: Rand, { - rand: &'a mut R, + rand: Rc>, entries: Vec>>>, dir_path: PathBuf, } -impl HasEntriesVec for OnDiskCorpus<'_, I, R> +impl HasEntriesVec for OnDiskCorpus where I: Input, R: Rand, @@ -149,22 +146,19 @@ where } } -impl HasRand for OnDiskCorpus<'_, I, R> +impl HasRand for OnDiskCorpus where I: Input, R: Rand, { type R = R; - fn rand(&self) -> &Self::R { + fn rand(&self) -> &Rc> { &self.rand } - fn rand_mut(&mut self) -> &mut Self::R { - &mut self.rand - } } -impl Corpus for OnDiskCorpus<'_, I, R> +impl Corpus for OnDiskCorpus where I: Input, R: Rand, @@ -183,16 +177,16 @@ where // TODO save and remove files, cache, etc..., ATM use just InMemoryCorpus } -impl<'a, I, R> OnDiskCorpus<'a, I, R> +impl OnDiskCorpus where I: Input, R: Rand, { - pub fn new(rand: &'a mut R, dir_path: PathBuf) -> Self { + pub fn new(rand: &Rc>, dir_path: PathBuf) -> Self { OnDiskCorpus { + rand: Rc::clone(rand), dir_path: dir_path, entries: vec![], - rand: rand, } } } @@ -229,12 +223,9 @@ where { type R = C::R; - fn rand(&self) -> &Self::R { + fn rand(&self) -> &Rc> { self.corpus.rand() } - fn rand_mut(&mut self) -> &mut Self::R { - self.corpus.rand_mut() - } } impl<'a, I, C> Corpus for QueueCorpus @@ -313,8 +304,8 @@ mod tests { #[test] fn test_queuecorpus() { - let mut rand = Xoshiro256StarRand::new(); - let mut q = QueueCorpus::new(OnDiskCorpus::new(&mut rand, PathBuf::from("fancy/path"))); + let rand = Xoshiro256StarRand::new_rc(); + let mut q = QueueCorpus::new(OnDiskCorpus::new(&rand, PathBuf::from("fancy/path"))); let i = BytesInput::new(vec![0; 4]); let t = Rc::new(RefCell::new(Testcase::new_with_filename( i, diff --git a/src/engines/mod.rs b/src/engines/mod.rs index bcbdfec074..c9ff590356 100644 --- a/src/engines/mod.rs +++ b/src/engines/mod.rs @@ -175,7 +175,6 @@ where #[cfg(test)] mod tests { - use crate::corpus::testcase::Testcase; use crate::corpus::InMemoryCorpus; use crate::engines::{DefaultEngine, Engine}; use crate::executors::inmemory::InMemoryExecutor; @@ -185,24 +184,22 @@ mod tests { use crate::utils::Xoshiro256StarRand; use std::cell::RefCell; - use std::path::PathBuf; use std::rc::Rc; - fn harness(_executor: &dyn Executor, buf: &[u8]) -> ExitKind { + fn harness(_executor: &dyn Executor, _buf: &[u8]) -> ExitKind { ExitKind::Ok } #[test] fn test_engine() { - let mut rand = Xoshiro256StarRand::new(); - let mut corpus = InMemoryCorpus::::new(&mut rand); + let rand = Rc::new(RefCell::new(Xoshiro256StarRand::new())); + let mut corpus = InMemoryCorpus::::new(&rand); let mut executor = InMemoryExecutor::new(harness); let mut engine = DefaultEngine::new(&mut corpus, &mut executor); - let mut rand1 = Xoshiro256StarRand::new(); - let mut stage = Box::new(DefaultMutationalStage::new(&mut rand1, &mut engine)); + let mut stage = Box::new(DefaultMutationalStage::new(&rand, &mut engine)); engine.add_stage(stage); engine.fuzz_one().unwrap(); - let mut stage1 = Box::new(DefaultMutationalStage::new(&mut rand1, &mut engine)); + let mut stage1 = Box::new(DefaultMutationalStage::new(&rand, &mut engine)); engine.fuzz_one().unwrap(); } } diff --git a/src/mutators/scheduled.rs b/src/mutators/scheduled.rs index d0f97552fe..e4114e2ee0 100644 --- a/src/mutators/scheduled.rs +++ b/src/mutators/scheduled.rs @@ -5,6 +5,8 @@ use crate::utils::{HasRand, Rand}; use crate::AflError; use std::marker::PhantomData; +use std::rc::Rc; +use std::cell::RefCell; /// The generic function type that identifies mutations type MutationFunction = fn(&mut M, &mut I) -> Result<(), AflError>; @@ -29,7 +31,7 @@ where { /// Compute the number of iterations used to apply stacked mutations fn iterations(&mut self, _input: &I) -> u64 { - 1 << (1 + self.rand_mut().below(7)) + 1 << (1 + self.rand_below(7)) } /// Get the next mutation to apply @@ -40,7 +42,7 @@ where } let idx; { - idx = self.rand_mut().below(count) as usize; + idx = self.rand().borrow_mut().below(count) as usize; } self.mutation_by_idx(idx) } @@ -62,7 +64,7 @@ where R: Rand, C: Corpus, { - rand: &'a mut R, + rand: Rc>, corpus: Option>, mutations: Vec>, } @@ -75,12 +77,9 @@ where { type R = R; - fn rand(&self) -> &Self::R { + fn rand(&self) -> &Rc> { &self.rand } - fn rand_mut(&mut self) -> &mut Self::R { - &mut self.rand - } } impl HasOptionCorpus for DefaultScheduledMutator<'_, I, R, C> @@ -153,9 +152,9 @@ where C: Corpus, { /// Create a new DefaultScheduledMutator instance without mutations and corpus - pub fn new(rand: &'a mut R) -> Self { + pub fn new(rand: &Rc>) -> Self { DefaultScheduledMutator { - rand: rand, + rand: Rc::clone(rand), corpus: None, mutations: vec![], } @@ -163,12 +162,12 @@ where /// Create a new DefaultScheduledMutator instance specifying mutations and corpus too pub fn new_all( - rand: &'a mut R, + rand: &Rc>, corpus: Option>, mutations: Vec>, ) -> Self { DefaultScheduledMutator { - rand: rand, + rand: Rc::clone(rand), corpus: corpus, mutations: mutations, } @@ -181,7 +180,7 @@ where M: Mutator, I: Input + HasBytesVec, { - let bit = mutator.rand_mut().below(input.bytes().len() as u64) as usize; + let bit = mutator.rand().borrow_mut().below(input.bytes().len() as u64) as usize; input.bytes_mut()[bit >> 3] ^= (128 >> (bit & 7)) as u8; Ok(()) } @@ -203,11 +202,8 @@ where { type R = S::R; - fn rand(&self) -> &Self::R { - self.scheduled.rand() - } - fn rand_mut(&mut self) -> &mut Self::R { - self.scheduled.rand_mut() + fn rand(&self) -> &Rc> { + &self.scheduled.rand() } } @@ -263,7 +259,7 @@ where C: Corpus, { /// Create a new HavocBytesMutator instance wrapping DefaultScheduledMutator - pub fn new_default(rand: &'a mut R) -> Self { + pub fn new_default(rand: &Rc>) -> Self { let mut scheduled = DefaultScheduledMutator::<'a, I, R, C>::new(rand); scheduled.add_mutation(mutation_bitflip); HavocBytesMutator { diff --git a/src/stages/mutational.rs b/src/stages/mutational.rs index cdad856d23..83d716b7ab 100644 --- a/src/stages/mutational.rs +++ b/src/stages/mutational.rs @@ -24,7 +24,7 @@ where } fn iterations(&mut self) -> usize { - 1 + self.rand_mut().below(128) as usize + 1 + self.rand().borrow_mut().below(128) as usize } fn perform_mutational(&mut self, entry: Rc>>) -> Result<(), AflError> { @@ -54,7 +54,7 @@ where R: Rand, E: Evaluator, { - rand: &'a mut R, + rand: Rc>, eval: &'a mut E, mutators: Vec>>, } @@ -67,12 +67,10 @@ where { type R = R; - fn rand(&self) -> &Self::R { + fn rand(&self) -> &Rc> { &self.rand } - fn rand_mut(&mut self) -> &mut Self::R { - &mut self.rand - } + } impl<'a, I, R, E> HasEvaluator for DefaultMutationalStage<'a, I, R, E> @@ -124,9 +122,9 @@ where R: Rand, E: Evaluator, { - pub fn new(rand: &'a mut R, eval: &'a mut E) -> Self { + pub fn new(rand: &Rc>, eval: &'a mut E) -> Self { DefaultMutationalStage { - rand: rand, + rand: Rc::clone(rand), eval: eval, mutators: vec![], } diff --git a/src/utils.rs b/src/utils.rs index 2c9da6763a..258ff64662 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,6 +1,8 @@ //! Utility functions for AFL use std::debug_assert; +use std::rc::Rc; +use std::cell::RefCell; use std::fmt::Debug; use xxhash_rust::xxh3::xxh3_64_with_seed; @@ -44,10 +46,21 @@ pub trait HasRand { type R: Rand; /// Get the hold Rand instance - fn rand(&self) -> &Self::R; + fn rand(&self) -> &Rc>; - /// Get the hold Rand instance (mutable) - fn rand_mut(&mut self) -> &mut Self::R; + // Gets the next 64 bit value + fn rand_next(&self) -> u64 { + self.rand().borrow_mut().next() + } + // Gets a value below the given 64 bit val (inclusive) + fn rand_below(&self, upper_bound_excl: u64) -> u64 { + self.rand().borrow_mut().below(upper_bound_excl) + } + + // Gets a value between the given lower bound (inclusive) and upper bound (inclusive) + fn rand_between(&self, lower_bound_incl: u64, upper_bound_incl: u64) -> u64 { + self.rand().borrow_mut().between(lower_bound_incl, upper_bound_incl) + } } const HASH_CONST: u64 = 0xa5b35705; @@ -96,6 +109,10 @@ impl Xoshiro256StarRand { ret.set_seed(0); // TODO: Proper random seed? ret } + + pub fn new_rc() -> Rc> { + Rc::new(RefCell::new(Xoshiro256StarRand::new())) + } } /// Get the next higher power of two @@ -111,7 +128,7 @@ pub fn next_pow2(val: u64) -> u64 { #[cfg(test)] mod tests { - use crate::utils::{next_pow2, Rand, Xoshiro256StarRand}; + use crate::utils::{next_pow2, Rand, HasRand, Xoshiro256StarRand}; #[test] fn test_rand() { @@ -123,6 +140,35 @@ mod tests { assert!(rand.between(11, 20) > 10); } + use std::rc::Rc; + use std::cell::RefCell; + struct HasRandTest + where + R: Rand, + { + rand: Rc> + } + + impl HasRand for HasRandTest + where + R: Rand + { + type R = R; + + fn rand(&self) -> &Rc> { + &self.rand + } + } + + fn test_has_rand() { + let rand = Xoshiro256StarRand::new_rc(); + let has_rand = HasRandTest{rand: Rc::clone(&rand)}; + + assert!(has_rand.rand_below(100) < 100); + assert_eq!(has_rand.rand_below(1), 0); + + } + #[test] fn test_next_pow2() { assert_eq!(next_pow2(0), 0);