From 3b0883721e795fe4b257424b943f9a8d9964fa3c Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Mon, 22 Feb 2021 14:13:00 +0100 Subject: [PATCH] rand and queue corpus schedulers --- libafl/src/corpus/inmemory.rs | 98 ++++++++--------- libafl/src/corpus/mod.rs | 118 +++++++------------- libafl/src/corpus/queue.rs | 143 ++++++------------------- libafl/src/executors/inprocess.rs | 10 +- libafl/src/mutators/mutations.rs | 20 +++- libafl/src/mutators/token_mutations.rs | 2 +- libafl/src/stages/mod.rs | 11 +- libafl/src/stages/mutational.rs | 3 +- libafl/src/state/mod.rs | 14 +-- 9 files changed, 155 insertions(+), 264 deletions(-) diff --git a/libafl/src/corpus/inmemory.rs b/libafl/src/corpus/inmemory.rs index 83b63edfa1..22a28ff67a 100644 --- a/libafl/src/corpus/inmemory.rs +++ b/libafl/src/corpus/inmemory.rs @@ -1,74 +1,74 @@ //! In-memory corpus, keeps all test cases in memory at all times -use alloc::{borrow::ToOwned, vec::Vec}; -use core::{cell::RefCell, marker::PhantomData}; +use alloc::vec::Vec; +use core::cell::RefCell; use serde::{Deserialize, Serialize}; -use crate::{ - corpus::Corpus, corpus::HasTestcaseVec, corpus::Testcase, inputs::Input, utils::Rand, Error, -}; +use crate::{corpus::Corpus, corpus::Testcase, inputs::Input, Error}; -/// A corpus handling all important fuzzing in memory. -#[derive(Serialize, Deserialize, Clone, Debug)] +/// A corpus handling all in memory. +#[derive(Default, Serialize, Deserialize, Clone, Debug)] #[serde(bound = "I: serde::de::DeserializeOwned")] -pub struct InMemoryCorpus +pub struct InMemoryCorpus where I: Input, - R: Rand, { entries: Vec>>, - pos: usize, - phantom: PhantomData, + current: Option, } -impl HasTestcaseVec for InMemoryCorpus +impl Corpus for InMemoryCorpus where I: Input, - R: Rand, { - fn entries(&self) -> &[RefCell>] { - &self.entries - } - fn entries_mut(&mut self) -> &mut Vec>> { - &mut self.entries - } -} - -impl Corpus for InMemoryCorpus -where - I: Input, - R: Rand, -{ - /// Gets the next entry + /// Returns the number of elements #[inline] - fn next(&mut self, rand: &mut R) -> Result<(&RefCell>, usize), Error> { - if self.count() == 0 { - Err(Error::Empty("No entries in corpus".to_owned())) + fn count(&self) -> usize { + self.entries.len() + } + + /// Add an entry to the corpus and return its index + #[inline] + fn add(&mut self, testcase: Testcase) -> Result { + self.entries.push(RefCell::new(testcase)); + Ok(self.entries.len() - 1) + } + + /// Replaces the testcase at the given idx + #[inline] + fn replace(&mut self, idx: usize, testcase: Testcase) -> Result<(), Error> { + if idx >= self.entries.len() { + return Err(Error::KeyNotFound(format!("Index {} out of bounds", idx))); + } + self.entries[idx] = RefCell::new(testcase); + Ok(()) + } + + /// Removes an entry from the corpus, returning it if it was present. + #[inline] + fn remove(&mut self, idx: usize) -> Result>, Error> { + if idx >= self.entries.len() { + Ok(None) } else { - let len = { self.entries().len() }; - let id = rand.below(len as u64) as usize; - self.pos = id; - Ok((self.get(id), id)) + Ok(Some(self.entries.remove(idx).into_inner())) } } - /// Returns the testacase we currently use + /// Get by id #[inline] - fn current_testcase(&self) -> (&RefCell>, usize) { - (self.get(self.pos), self.pos) + fn get(&self, idx: usize) -> Result<&RefCell>, Error> { + Ok(&self.entries[idx]) } -} -impl InMemoryCorpus -where - I: Input, - R: Rand, -{ - pub fn new() -> Self { - Self { - entries: vec![], - pos: 0, - phantom: PhantomData, - } + /// Current testcase scheduled + #[inline] + fn current(&self) -> &Option { + &self.current + } + + /// Current testcase scheduled (mut) + #[inline] + fn current_mut(&mut self) -> &mut Option { + &mut self.current } } diff --git a/libafl/src/corpus/mod.rs b/libafl/src/corpus/mod.rs index 1bea657106..6ed7f18bb1 100644 --- a/libafl/src/corpus/mod.rs +++ b/libafl/src/corpus/mod.rs @@ -3,11 +3,21 @@ pub mod testcase; pub use testcase::Testcase; -use alloc::vec::Vec; -use core::cell::RefCell; -use serde::{Deserialize, Serialize}; +pub mod inmemory; +pub use inmemory::InMemoryCorpus; -use crate::{inputs::Input, Error}; +pub mod queue; +pub use queue::QueueCorpusScheduler; + +use core::cell::RefCell; +use core::marker::PhantomData; + +use crate::{ + inputs::Input, + state::{HasCorpus, HasRand}, + utils::Rand, + Error, +}; /// Corpus with all current testcases pub trait Corpus: serde::Serialize + serde::de::DeserializeOwned @@ -46,7 +56,12 @@ where } /// Replaces the testcase at the given idx - fn on_replace(&self, _state: &mut S, _idx: usize, _testcase: &Testcase) -> Result<(), Error> { + fn on_replace( + &self, + _state: &mut S, + _idx: usize, + _testcase: &Testcase, + ) -> Result<(), Error> { Ok(()) } @@ -65,89 +80,34 @@ where fn next(&self, state: &mut S) -> Result; } -/* -pub struct RandCorpusScheduler {} +pub struct RandCorpusScheduler +where + S: HasCorpus + HasRand, + C: Corpus, + I: Input, + R: Rand, +{ + phantom: PhantomData<(C, I, R, S)>, +} -impl CorpusScheduler for RandCorpusScheduler { +impl CorpusScheduler for RandCorpusScheduler +where + S: HasCorpus + HasRand, + C: Corpus, + I: Input, + R: Rand, +{ /// Gets the next entry at random - fn next(state: &mut S) -> Result - where - S: HasCorpus + HasRand, - C: Corpus, - I: Input, - R: Rand, - { + fn next(&self, state: &mut S) -> Result { if state.corpus().count() == 0 { Err(Error::Empty("No entries in corpus".to_owned())) } else { let len = state.corpus().count(); let id = state.rand_mut().below(len as u64) as usize; + *state.corpus_mut().current_mut() = Some(id); Ok(id) } } } -*/ -#[derive(Default, Serialize, Deserialize, Clone, Debug)] -#[serde(bound = "I: serde::de::DeserializeOwned")] -pub struct InMemoryCorpus -where - I: Input, -{ - entries: Vec>>, - current: Option, -} - -impl Corpus for InMemoryCorpus -where - I: Input, -{ - /// Returns the number of elements - #[inline] - fn count(&self) -> usize { - self.entries.len() - } - - /// Add an entry to the corpus and return its index - #[inline] - fn add(&mut self, testcase: Testcase) -> Result { - self.entries.push(RefCell::new(testcase)); - Ok(self.entries.len() - 1) - } - - /// Replaces the testcase at the given idx - #[inline] - fn replace(&mut self, idx: usize, testcase: Testcase) -> Result<(), Error> { - if idx >= self.entries.len() { - return Err(Error::KeyNotFound(format!("Index {} out of bounds", idx))); - } - self.entries[idx] = RefCell::new(testcase); - Ok(()) - } - - /// Removes an entry from the corpus, returning it if it was present. - #[inline] - fn remove(&mut self, idx: usize) -> Result>, Error> { - if idx >= self.entries.len() { - Ok(None) - } else { - Ok(Some(self.entries.remove(idx).into_inner())) - } - } - - /// Get by id - #[inline] - fn get(&self, idx: usize) -> Result<&RefCell>, Error> { - Ok(&self.entries[idx]) - } - - /// Current testcase scheduled - fn current(&self) -> &Option { - &self.current - } - - /// Current testcase scheduled (mut) - fn current_mut(&mut self) -> &mut Option { - &mut self.current - } -} +pub type StdCorpusScheduler = RandCorpusScheduler; diff --git a/libafl/src/corpus/queue.rs b/libafl/src/corpus/queue.rs index 1e90a6e0ce..12fa682e14 100644 --- a/libafl/src/corpus/queue.rs +++ b/libafl/src/corpus/queue.rs @@ -1,132 +1,52 @@ //! The queue corpus implements an afl-like queue mechanism -use alloc::{borrow::ToOwned, vec::Vec}; -use core::{cell::RefCell, marker::PhantomData}; -use serde::{Deserialize, Serialize}; +use alloc::borrow::ToOwned; +use core::marker::PhantomData; use crate::{ - corpus::Corpus, corpus::HasTestcaseVec, corpus::Testcase, inputs::Input, utils::Rand, Error, + corpus::{Corpus, CorpusScheduler}, + inputs::Input, + state::HasCorpus, + Error, }; -/// A Queue-like corpus, wrapping an existing Corpus instance -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(bound = "I: serde::de::DeserializeOwned")] -pub struct QueueCorpus +pub struct QueueCorpusScheduler where - C: Corpus, + S: HasCorpus, + C: Corpus, I: Input, - R: Rand, { - corpus: C, - pos: usize, - cycles: u64, - phantom: PhantomData<(I, R)>, + phantom: PhantomData<(C, I, S)>, } -impl HasTestcaseVec for QueueCorpus +impl CorpusScheduler for QueueCorpusScheduler where - C: Corpus, + S: HasCorpus, + C: Corpus, I: Input, - R: Rand, { - #[inline] - fn entries(&self) -> &[RefCell>] { - self.corpus.entries() - } - #[inline] - fn entries_mut(&mut self) -> &mut Vec>> { - self.corpus.entries_mut() - } -} - -impl Corpus for QueueCorpus -where - C: Corpus, - I: Input, - R: Rand, -{ - /// Returns the number of elements - #[inline] - fn count(&self) -> usize { - self.corpus.count() - } - - #[inline] - fn add(&mut self, entry: Testcase) -> usize { - self.corpus.add(entry) - } - - /// Removes an entry from the corpus, returning it if it was present. - #[inline] - fn remove(&mut self, entry: &Testcase) -> Option> { - self.corpus.remove(entry) - } - - /// Gets a random entry - #[inline] - fn random_entry(&self, rand: &mut R) -> Result<(&RefCell>, usize), Error> { - self.corpus.random_entry(rand) - } - - /// Returns the testacase we currently use - #[inline] - fn current_testcase(&self) -> (&RefCell>, usize) { - (self.get(self.pos - 1), self.pos - 1) - } - - /// Gets the next entry - #[inline] - fn next(&mut self, _rand: &mut R) -> Result<(&RefCell>, usize), Error> { - self.pos += 1; - if self.corpus.count() == 0 { - return Err(Error::Empty("Corpus".to_owned())); - } - if self.pos > self.corpus.count() { - // TODO: Always loop or return informational error? - self.pos = 1; - self.cycles += 1; - } - Ok((&self.corpus.entries()[self.pos - 1], self.pos - 1)) - } -} - -impl QueueCorpus -where - C: Corpus, - I: Input, - R: Rand, -{ - pub fn new(corpus: C) -> Self { - Self { - corpus: corpus, - phantom: PhantomData, - cycles: 0, - pos: 0, + /// Gets the next entry at random + fn next(&self, state: &mut S) -> Result { + if state.corpus().count() == 0 { + Err(Error::Empty("No entries in corpus".to_owned())) + } else { + let id = match state.corpus().current() { + Some(cur) => { + if *cur + 1 > state.corpus().count() { + 0 + } else { + *cur + 1 + } + } + None => 0, + }; + *state.corpus_mut().current_mut() = Some(id); + Ok(id) } } - - #[inline] - pub fn cycles(&self) -> u64 { - self.cycles - } - - #[inline] - pub fn pos(&self) -> usize { - self.pos - } - - // TODO maybe impl HasCorpus - #[inline] - pub fn corpus(&self) -> &C { - &self.corpus - } - - #[inline] - pub fn corpus_mut(&mut self) -> &mut C { - &mut self.corpus - } } +/* #[cfg(test)] #[cfg(feature = "std")] mod tests { @@ -170,3 +90,4 @@ mod tests { assert_eq!(filename, "fancyfile"); } } +*/ diff --git a/libafl/src/executors/inprocess.rs b/libafl/src/executors/inprocess.rs index 825e52ca73..65222ab004 100644 --- a/libafl/src/executors/inprocess.rs +++ b/libafl/src/executors/inprocess.rs @@ -249,7 +249,10 @@ pub mod unix_signals { .is_interesting_all(&input, observers, ExitKind::Crash) .expect("In crash handler objectives failure.".into()); if obj_fitness > 0 { - state.solutions_mut().add(Testcase::new(input.clone())).expect("In crash handler solutions failure.".into()); + state + .solutions_mut() + .add(Testcase::new(input.clone())) + .expect("In crash handler solutions failure.".into()); mgr.fire( state, Event::Objective { @@ -302,7 +305,10 @@ pub mod unix_signals { .is_interesting_all(&input, observers, ExitKind::Crash) .expect("In timeout handler objectives failure.".into()); if obj_fitness > 0 { - state.solutions_mut().add(Testcase::new(input.clone())).expect("In timeout handler solutions failure.".into()); + state + .solutions_mut() + .add(Testcase::new(input.clone())) + .expect("In timeout handler solutions failure.".into()); mgr.fire( state, Event::Objective { diff --git a/libafl/src/mutators/mutations.rs b/libafl/src/mutators/mutations.rs index ef45a83085..730efad2a1 100644 --- a/libafl/src/mutators/mutations.rs +++ b/libafl/src/mutators/mutations.rs @@ -711,7 +711,13 @@ where } } - let other_size = state.corpus().get(idx)?.borrow_mut().load_input()?.bytes().len(); + let other_size = state + .corpus() + .get(idx)? + .borrow_mut() + .load_input()? + .bytes() + .len(); if other_size < 2 { return Ok(MutationResult::Skipped); } @@ -761,8 +767,14 @@ where return Ok(MutationResult::Skipped); } } - - let other_size = state.corpus().get(idx)?.borrow_mut().load_input()?.bytes().len(); + + let other_size = state + .corpus() + .get(idx)? + .borrow_mut() + .load_input()? + .bytes() + .len(); if other_size < 2 { return Ok(MutationResult::Skipped); } @@ -838,7 +850,7 @@ where let split_at = state .rand_mut() .between(first_diff as u64, last_diff as u64) as usize; - + let mut other_testcase = state.corpus().get(idx)?.borrow_mut(); let other = other_testcase.load_input()?; input diff --git a/libafl/src/mutators/token_mutations.rs b/libafl/src/mutators/token_mutations.rs index 505dc9d54b..252b17fcfb 100644 --- a/libafl/src/mutators/token_mutations.rs +++ b/libafl/src/mutators/token_mutations.rs @@ -53,7 +53,7 @@ where meta.unwrap().tokens.len() }; let token_idx = state.rand_mut().below(tokens_len as u64) as usize; - + let size = input.bytes().len(); let off = state.rand_mut().below((size + 1) as u64) as usize; diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs index 27f88c1bfb..01e729ac0f 100644 --- a/libafl/src/stages/mod.rs +++ b/libafl/src/stages/mod.rs @@ -1,7 +1,7 @@ pub mod mutational; pub use mutational::StdMutationalStage; -use crate::{bolts::tuples::TupleList, Error}; +use crate::{bolts::tuples::TupleList, Error}; /// A stage is one step in the fuzzing process. /// Multiple stages will be scheduled one by one for each input. @@ -29,14 +29,7 @@ pub trait StagesTuple { } impl StagesTuple for () { - fn perform_all( - &self, - _: &F, - _: &mut S, - _: &mut E, - _: &mut EM, - _: usize, - ) -> Result<(), Error> { + fn perform_all(&self, _: &F, _: &mut S, _: &mut E, _: &mut EM, _: usize) -> Result<(), Error> { Ok(()) } } diff --git a/libafl/src/stages/mutational.rs b/libafl/src/stages/mutational.rs index a77aaeb6e5..d69c2e3600 100644 --- a/libafl/src/stages/mutational.rs +++ b/libafl/src/stages/mutational.rs @@ -59,8 +59,7 @@ where let fitness = state.evaluate_input(input_mut, executor, manager)?; - self.mutator() - .post_exec(fuzzer, state, fitness, i as i32)?; + self.mutator().post_exec(fuzzer, state, fitness, i as i32)?; } Ok(()) } diff --git a/libafl/src/state/mod.rs b/libafl/src/state/mod.rs index 38eb7b0c86..0bdc3dd1d1 100644 --- a/libafl/src/state/mod.rs +++ b/libafl/src/state/mod.rs @@ -583,14 +583,14 @@ where executor.post_exec_observers()?; let observers = executor.observers(); - let fitness = self - .feedbacks_mut() - .is_interesting_all(&input, observers, exit_kind.clone())?; + let fitness = + self.feedbacks_mut() + .is_interesting_all(&input, observers, exit_kind.clone())?; - let is_solution = - self.objectives_mut() - .is_interesting_all(&input, observers, exit_kind)? - > 0; + let is_solution = self + .objectives_mut() + .is_interesting_all(&input, observers, exit_kind)? + > 0; Ok((fitness, is_solution)) }