diff --git a/afl/src/corpus/mod.rs b/afl/src/corpus/mod.rs index d92ab7e347..e260379b4a 100644 --- a/afl/src/corpus/mod.rs +++ b/afl/src/corpus/mod.rs @@ -4,6 +4,7 @@ pub use testcase::{Testcase, TestcaseMetadata}; use alloc::borrow::ToOwned; use alloc::vec::Vec; use core::marker::PhantomData; +use core::cell::RefCell; use core::ptr; #[cfg(feature = "std")] @@ -18,10 +19,10 @@ where I: Input, { /// Get the entries vector field - fn entries(&self) -> &[Testcase]; + fn entries(&self) -> &[RefCell>]; /// 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 @@ -37,7 +38,7 @@ where /// Add an entry to the corpus and return its index fn add(&mut self, testcase: Testcase) -> usize { - self.entries_mut().push(testcase); + self.entries_mut().push(RefCell::new(testcase)); self.entries().len() - 1 } @@ -49,25 +50,25 @@ where idx ))); } - self.entries_mut()[idx] = testcase; + self.entries_mut()[idx] = RefCell::new(testcase); Ok(()) } /// Get by id - fn get(&self, idx: usize) -> &Testcase { + fn get(&self, idx: usize) -> &RefCell> { &self.entries()[idx] } /// Removes an entry from the corpus, returning it if it was present. fn remove(&mut self, entry: &Testcase) -> Option> { - match self.entries().iter().position(|x| ptr::eq(x, entry)) { - Some(i) => Some(self.entries_mut().remove(i)), + match self.entries().iter().position(|x| ptr::eq(x.as_ptr(), entry)) { + Some(i) => Some(self.entries_mut().remove(i).into_inner()), None => None, } } /// Gets a random entry - fn random_entry(&self, rand: &mut R) -> Result<(&Testcase, usize), AflError> { + fn random_entry(&self, rand: &mut R) -> Result<(&RefCell>, usize), AflError> { if self.count() == 0 { Err(AflError::Empty("No entries in corpus".to_owned())) } else { @@ -78,7 +79,7 @@ where } /// Returns the testcase for the given idx, with loaded input - fn load_testcase(&mut self, idx: usize) -> Result<(), AflError> { + /*fn load_testcase(&mut self, idx: usize) -> Result<(), AflError> { let testcase = self.get(idx); // Ensure testcase is loaded match testcase.input() { @@ -97,14 +98,14 @@ where _ => (), } Ok(()) - } + }*/ // TODO: IntoIter /// Gets the next entry - fn next(&mut self, rand: &mut R) -> Result<(&Testcase, usize), AflError>; + fn next(&mut self, rand: &mut R) -> Result<(&RefCell>, usize), AflError>; /// Returns the testacase we currently use - fn current_testcase(&self) -> (&Testcase, usize); + fn current_testcase(&self) -> (&RefCell>, usize); } pub struct InMemoryCorpus @@ -112,7 +113,7 @@ where I: Input, R: Rand, { - entries: Vec>, + entries: Vec>>, pos: usize, phantom: PhantomData, } @@ -122,10 +123,10 @@ where I: Input, R: Rand, { - fn entries(&self) -> &[Testcase] { + fn entries(&self) -> &[RefCell>] { &self.entries } - fn entries_mut(&mut self) -> &mut Vec> { + fn entries_mut(&mut self) -> &mut Vec>> { &mut self.entries } } @@ -136,7 +137,7 @@ where R: Rand, { /// Gets the next entry - fn next(&mut self, rand: &mut R) -> Result<(&Testcase, usize), AflError> { + fn next(&mut self, rand: &mut R) -> Result<(&RefCell>, usize), AflError> { if self.count() == 0 { Err(AflError::Empty("No entries in corpus".to_owned())) } else { @@ -148,7 +149,7 @@ where } /// Returns the testacase we currently use - fn current_testcase(&self) -> (&Testcase, usize) { + fn current_testcase(&self) -> (&RefCell>, usize) { (self.get(self.pos), self.pos) } } @@ -173,7 +174,7 @@ where I: Input, R: Rand, { - entries: Vec>, + entries: Vec>>, dir_path: PathBuf, pos: usize, phantom: PhantomData, @@ -185,10 +186,10 @@ where I: Input, R: Rand, { - fn entries(&self) -> &[Testcase] { + fn entries(&self) -> &[RefCell>] { &self.entries } - fn entries_mut(&mut self) -> &mut Vec> { + fn entries_mut(&mut self) -> &mut Vec>> { &mut self.entries } } @@ -210,16 +211,16 @@ where } _ => {} } - self.entries.push(entry); + self.entries.push(RefCell::new(entry)); self.entries.len() - 1 } - fn current_testcase(&self) -> (&Testcase, usize) { + fn current_testcase(&self) -> (&RefCell>, usize) { (self.get(self.pos), self.pos) } /// Gets the next entry - fn next(&mut self, rand: &mut R) -> Result<(&Testcase, usize), AflError> { + fn next(&mut self, rand: &mut R) -> Result<(&RefCell>, usize), AflError> { if self.count() == 0 { Err(AflError::Empty("No entries in corpus".to_owned())) } else { @@ -268,10 +269,10 @@ where I: Input, R: Rand, { - fn entries(&self) -> &[Testcase] { + fn entries(&self) -> &[RefCell>] { self.corpus.entries() } - fn entries_mut(&mut self) -> &mut Vec> { + fn entries_mut(&mut self) -> &mut Vec>> { self.corpus.entries_mut() } } @@ -297,17 +298,17 @@ where } /// Gets a random entry - fn random_entry(&self, rand: &mut R) -> Result<(&Testcase, usize), AflError> { + fn random_entry(&self, rand: &mut R) -> Result<(&RefCell>, usize), AflError> { self.corpus.random_entry(rand) } /// Returns the testacase we currently use - fn current_testcase(&self) -> (&Testcase, usize) { + fn current_testcase(&self) -> (&RefCell>, usize) { (self.get(self.pos - 1), self.pos - 1) } /// Gets the next entry - fn next(&mut self, _rand: &mut R) -> Result<(&Testcase, usize), AflError> { + fn next(&mut self, _rand: &mut R) -> Result<(&RefCell>, usize), AflError> { self.pos += 1; if self.corpus.count() == 0 { return Err(AflError::Empty("Corpus".to_owned())); @@ -417,6 +418,7 @@ mod tests { .next(&mut rand) .unwrap() .0 + .borrow() .filename() .as_ref() .unwrap() @@ -426,6 +428,7 @@ mod tests { q.next(&mut rand) .unwrap() .0 + .borrow() .filename() .as_ref() .unwrap() diff --git a/afl/src/corpus/testcase.rs b/afl/src/corpus/testcase.rs index daca24ef74..a30d45e262 100644 --- a/afl/src/corpus/testcase.rs +++ b/afl/src/corpus/testcase.rs @@ -75,9 +75,11 @@ where I: Input, { /// Returns this testcase with a loaded input - pub fn load_from_disk(filename: &str) -> Result { - let input = I::from_file(filename)?; - Ok(Testcase::new(input)) + pub fn load_input(&mut self) -> Result<&I, AflError> { + if self.input.is_none() { + self.input = Some(I::from_file(self.filename.as_ref().unwrap())?); + } + Ok(self.input.as_ref().unwrap()) } /// Get the input, if any diff --git a/afl/src/engines/mod.rs b/afl/src/engines/mod.rs index 4545fc5329..1e4542d36c 100644 --- a/afl/src/engines/mod.rs +++ b/afl/src/engines/mod.rs @@ -293,19 +293,10 @@ where corpus: &mut C, engine: &mut Engine, ) -> Result { - let (testcase, idx) = corpus.next(rand)?; - match testcase.input() { - None => { - // Load from disk. - corpus.load_testcase(idx)?; - } - _ => (), - }; - - let input = corpus.get(idx).input().as_ref().unwrap(); + let (_, idx) = corpus.next(rand)?; for stage in self.stages_mut() { - stage.perform(rand, state, corpus, engine, &input)?; + stage.perform(rand, state, corpus, engine, idx)?; } engine.events_manager_mut().process(state, corpus)?; diff --git a/afl/src/executors/inmemory.rs b/afl/src/executors/inmemory.rs index 9f23506ec9..e2fd5ceb63 100644 --- a/afl/src/executors/inmemory.rs +++ b/afl/src/executors/inmemory.rs @@ -4,7 +4,6 @@ use core::ptr; use crate::executors::{Executor, ExitKind}; use crate::inputs::Input; use crate::observers::observer_serde::NamedSerdeAnyMap; -use crate::observers::Observer; use crate::AflError; type HarnessFunction = fn(&dyn Executor, &[u8]) -> ExitKind; diff --git a/afl/src/mutators/mutations.rs b/afl/src/mutators/mutations.rs index 032efc1f9f..64c3aa5bce 100644 --- a/afl/src/mutators/mutations.rs +++ b/afl/src/mutators/mutations.rs @@ -686,13 +686,13 @@ where R: Rand, { // We don't want to use the testcase we're already using for splicing - let (other_testcase, _) = corpus.random_entry(rand)?.clone(); - // TODO: Load let other = Testcase::load_from_disk(other_test)?; + let (other_testcase, idx) = corpus.random_entry(rand)?; + if idx == corpus.current_testcase().1 { + return Ok(MutationResult::Skipped); + } // println!("Input: {:?}, other input: {:?}", input.bytes(), other.bytes()); - let other = match other_testcase.input() { - Some(i) => i, - None => return Ok(MutationResult::Skipped), //TODO - }; + let mut other_ref = other_testcase.borrow_mut(); + let other = other_ref.load_input()?; let mut counter = 0; let (first_diff, last_diff) = loop { diff --git a/afl/src/mutators/scheduled.rs b/afl/src/mutators/scheduled.rs index eb342a8ad6..62ee7df245 100644 --- a/afl/src/mutators/scheduled.rs +++ b/afl/src/mutators/scheduled.rs @@ -265,7 +265,7 @@ mod tests { let (testcase, _) = corpus .next(&mut rand) .expect("Corpus did not contain entries"); - let mut input = testcase.input().as_ref().unwrap().clone(); + let mut input = testcase.borrow_mut().load_input().unwrap().clone(); rand.set_seed(5); let mut mutator = StdScheduledMutator::new(); diff --git a/afl/src/stages/mod.rs b/afl/src/stages/mod.rs index 49fd2d81fa..97881719ef 100644 --- a/afl/src/stages/mod.rs +++ b/afl/src/stages/mod.rs @@ -24,6 +24,6 @@ where state: &mut State, corpus: &mut C, engine: &mut Engine, - input: &I, + corpus_idx: usize, ) -> Result<(), AflError>; } diff --git a/afl/src/stages/mutational.rs b/afl/src/stages/mutational.rs index 164dff94ce..9c455d6d4f 100644 --- a/afl/src/stages/mutational.rs +++ b/afl/src/stages/mutational.rs @@ -40,11 +40,11 @@ where state: &mut State, corpus: &mut C, engine: &mut Engine, - input: &I, + corpus_idx: usize, ) -> Result<(), AflError> { let num = self.iterations(rand); for i in 0..num { - let mut input_mut = input.clone(); + let mut input_mut = corpus.get(corpus_idx).borrow_mut().load_input()?.clone(); self.mutator_mut() .mutate(rand, corpus, &mut input_mut, i as i32)?; @@ -120,9 +120,9 @@ where state: &mut State, corpus: &mut C, engine: &mut Engine, - input: &I, + corpus_idx: usize, ) -> Result<(), AflError> { - self.perform_mutational(rand, state, corpus, engine, input) + self.perform_mutational(rand, state, corpus, engine, corpus_idx) } }