diff --git a/afl/src/corpus/mod.rs b/afl/src/corpus/mod.rs index b8a33d55c4..ea7467a95c 100644 --- a/afl/src/corpus/mod.rs +++ b/afl/src/corpus/mod.rs @@ -2,10 +2,10 @@ pub mod testcase; pub use testcase::{Testcase, TestcaseMetadata}; use alloc::borrow::ToOwned; -use alloc::rc::Rc; use alloc::vec::Vec; -use core::cell::RefCell; use core::marker::PhantomData; +use core::ptr; + #[cfg(feature = "std")] use std::path::PathBuf; @@ -18,10 +18,10 @@ where I: Input, { /// Get the entries vector field - fn entries(&self) -> &[Rc>>]; + fn entries(&self) -> &[Testcase]; /// 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 @@ -36,57 +36,74 @@ where } /// Add an entry to the corpus - fn add(&mut self, testcase: Rc>>) { + fn add(&mut self, testcase: Testcase) { self.entries_mut().push(testcase); } + /// Replaces the testcase at the given idx + fn replace(&mut self, idx: usize, testcase: Testcase) -> Result<(), AflError> { + if self.entries_mut().len() < idx { + return Err(AflError::KeyNotFound(format!( + "Index {} out of bounds", + idx + ))); + } + self.entries_mut()[idx] = testcase; + Ok(()) + } + + /// Get by id + fn get(&self, idx: usize) -> &Testcase { + &self.entries()[idx] + } + /// Removes an entry from the corpus, returning it if it was present. - fn remove(&mut self, entry: &Testcase) -> Option>>> { - let mut i: usize = 0; - let mut found = false; - for x in self.entries() { - i = i + 1; - if &*x.borrow() as *const _ == entry as *const _ { - // TODO check if correct - found = true; - break; - } + 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)), + None => None, } - if !found { - return None; - } - Some(self.entries_mut().remove(i)) } /// Gets a random entry - fn random_entry(&self, rand: &mut R) -> Result<(Rc>>, usize), AflError> { + fn random_entry(&self, rand: &mut R) -> Result<(&Testcase, usize), AflError> { if self.count() == 0 { Err(AflError::Empty("No entries in corpus".to_owned())) } else { let len = { self.entries().len() }; let id = rand.below(len as u64) as usize; - Ok((self.entries()[id].clone(), id)) + Ok((self.get(id), id)) } } + /// Returns the testcase for the given idx, with loaded input + fn load_testcase(&mut self, idx: usize) -> Result<(), AflError> { + let testcase = self.get(idx); + // Ensure testcase is loaded + match testcase.input() { + None => { + let new_testcase = match testcase.filename() { + Some(filename) => Testcase::load_from_disk(filename)?, + None => { + return Err(AflError::IllegalState( + "Neither input, nor filename specified for testcase".into(), + )) + } + }; + + self.replace(idx, new_testcase)?; + } + _ => (), + } + Ok(()) + } + // TODO: IntoIter /// Gets the next entry - fn next(&mut self, rand: &mut R) -> Result<(Rc>>, usize), AflError> { - self.random_entry(rand) - } -} + fn next(&mut self, rand: &mut R) -> Result<(&Testcase, usize), AflError>; -pub trait HasCorpus -where - C: Corpus, - I: Input, - R: Rand, -{ - /// Get the corpus field - fn corpus(&self) -> &C; - - /// Get thecorpus field (mutable) - fn corpus_mut(&mut self) -> &mut C; + /// Returns the testacase we currently use + fn current_testcase(&self) -> (&Testcase, usize); } pub struct InMemoryCorpus @@ -94,7 +111,8 @@ where I: Input, R: Rand, { - entries: Vec>>>, + entries: Vec>, + pos: usize, phantom: PhantomData, } @@ -103,10 +121,10 @@ where I: Input, R: Rand, { - fn entries(&self) -> &[Rc>>] { + fn entries(&self) -> &[Testcase] { &self.entries } - fn entries_mut(&mut self) -> &mut Vec>>> { + fn entries_mut(&mut self) -> &mut Vec> { &mut self.entries } } @@ -116,7 +134,22 @@ where I: Input, R: Rand, { - // Just use the default implementation + /// Gets the next entry + fn next(&mut self, rand: &mut R) -> Result<(&Testcase, usize), AflError> { + if self.count() == 0 { + Err(AflError::Empty("No entries in corpus".to_owned())) + } else { + let len = { self.entries().len() }; + let id = rand.below(len as u64) as usize; + self.pos = id; + Ok((self.get(id), id)) + } + } + + /// Returns the testacase we currently use + fn current_testcase(&self) -> (&Testcase, usize) { + (self.get(self.pos), self.pos) + } } impl InMemoryCorpus @@ -127,6 +160,7 @@ where pub fn new() -> Self { InMemoryCorpus { entries: vec![], + pos: 0, phantom: PhantomData, } } @@ -138,8 +172,9 @@ where I: Input, R: Rand, { - entries: Vec>>>, + entries: Vec>, dir_path: PathBuf, + pos: usize, phantom: PhantomData, } @@ -149,10 +184,10 @@ where I: Input, R: Rand, { - fn entries(&self) -> &[Rc>>] { + fn entries(&self) -> &[Testcase] { &self.entries } - fn entries_mut(&mut self) -> &mut Vec>>> { + fn entries_mut(&mut self) -> &mut Vec> { &mut self.entries } } @@ -164,16 +199,35 @@ where R: Rand, { /// Add an entry and save it to disk - fn add(&mut self, entry: Rc>>) { - if *entry.borrow().filename() == None { - // TODO walk entry metadatas to ask for pices of filename (e.g. :havoc in AFL) - let filename = self.dir_path.join(format!("id_{}", &self.entries.len())); - let filename_str = filename.to_str().expect("Invalid Path"); - *entry.borrow_mut().filename_mut() = Some(filename_str.into()); + fn add(&mut self, mut entry: Testcase) { + match entry.filename() { + None => { + // TODO walk entry metadatas to ask for pices of filename (e.g. :havoc in AFL) + let filename = self.dir_path.join(format!("id_{}", &self.entries.len())); + let filename_str = filename.to_str().expect("Invalid Path"); + entry.set_filename(filename_str.into()); + } + _ => {} } self.entries.push(entry); } + fn current_testcase(&self) -> (&Testcase, usize) { + (self.get(self.pos), self.pos) + } + + /// Gets the next entry + fn next(&mut self, rand: &mut R) -> Result<(&Testcase, usize), AflError> { + if self.count() == 0 { + Err(AflError::Empty("No entries in corpus".to_owned())) + } else { + let len = { self.entries().len() }; + let id = rand.below(len as u64) as usize; + self.pos = id; + Ok((self.get(id), id)) + } + } + // TODO save and remove files, cache, etc..., ATM use just InMemoryCorpus } @@ -187,6 +241,7 @@ where OnDiskCorpus { dir_path: dir_path, entries: vec![], + pos: 0, phantom: PhantomData, } } @@ -211,10 +266,10 @@ where I: Input, R: Rand, { - fn entries(&self) -> &[Rc>>] { + fn entries(&self) -> &[Testcase] { self.corpus.entries() } - fn entries_mut(&mut self) -> &mut Vec>>> { + fn entries_mut(&mut self) -> &mut Vec> { self.corpus.entries_mut() } } @@ -230,22 +285,27 @@ where self.corpus.count() } - fn add(&mut self, entry: Rc>>) { + fn add(&mut self, entry: Testcase) { self.corpus.add(entry); } /// Removes an entry from the corpus, returning it if it was present. - fn remove(&mut self, entry: &Testcase) -> Option>>> { + fn remove(&mut self, entry: &Testcase) -> Option> { self.corpus.remove(entry) } /// Gets a random entry - fn random_entry(&self, rand: &mut R) -> Result<(Rc>>, usize), AflError> { + fn random_entry(&self, rand: &mut R) -> Result<(&Testcase, usize), AflError> { self.corpus.random_entry(rand) } + /// Returns the testacase we currently use + fn current_testcase(&self) -> (&Testcase, usize) { + (self.get(self.pos - 1), self.pos - 1) + } + /// Gets the next entry - fn next(&mut self, _rand: &mut R) -> Result<(Rc>>, usize), AflError> { + fn next(&mut self, _rand: &mut R) -> Result<(&Testcase, usize), AflError> { self.pos += 1; if self.corpus.count() == 0 { return Err(AflError::Empty("Corpus".to_owned())); @@ -255,7 +315,7 @@ where self.pos = 1; self.cycles += 1; } - Ok((self.corpus.entries()[self.pos - 1].clone(), self.pos - 1)) + Ok((&self.corpus.entries()[self.pos - 1], self.pos - 1)) } } @@ -341,7 +401,6 @@ mod tests { use crate::inputs::bytes::BytesInput; use crate::utils::StdRand; - use alloc::rc::Rc; use std::path::PathBuf; #[test] @@ -350,14 +409,12 @@ mod tests { let mut q = QueueCorpus::new(OnDiskCorpus::::new(PathBuf::from( "fancy/path", ))); - let t: Rc<_> = - Testcase::with_filename(BytesInput::new(vec![0 as u8; 4]), "fancyfile".into()).into(); + let t = Testcase::with_filename(BytesInput::new(vec![0 as u8; 4]), "fancyfile".into()); q.add(t); let filename = q .next(&mut rand) .unwrap() .0 - .borrow() .filename() .as_ref() .unwrap() @@ -367,7 +424,6 @@ 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 b2e518d7d8..b4c6d6b696 100644 --- a/afl/src/corpus/testcase.rs +++ b/afl/src/corpus/testcase.rs @@ -51,17 +51,10 @@ impl Testcase where I: Input, { - /// Make sure to return a valid input instance loading it from disk if not in memory - pub fn load_input(&mut self) -> Result<&I, AflError> { - if self.input.is_none() { - let input = I::from_file( - self.filename - .as_ref() - .ok_or(AflError::EmptyOptional("filename not specified".into()))?, - )?; - self.input = Some(input); - } - Ok(self.input.as_ref().unwrap()) + /// 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)) } /// Get the input, if any diff --git a/afl/src/engines/mod.rs b/afl/src/engines/mod.rs index d8fec2ed31..cb90ac6793 100644 --- a/afl/src/engines/mod.rs +++ b/afl/src/engines/mod.rs @@ -8,7 +8,7 @@ use core::fmt::Debug; use core::marker::PhantomData; use hashbrown::HashMap; -use crate::corpus::{Corpus, HasCorpus, Testcase}; +use crate::corpus::{Corpus, Testcase}; use crate::events::{EventManager, LoadInitialEvent, UpdateStatsEvent}; use crate::executors::Executor; use crate::feedbacks::Feedback; @@ -26,7 +26,7 @@ pub trait StateMetadata: Debug { fn name(&self) -> &'static str; } -pub trait State: HasCorpus +pub trait State where C: Corpus, E: Executor, @@ -111,10 +111,7 @@ where fn executor_mut(&mut self) -> &mut E; /// Runs the input and triggers observers and feedback - fn evaluate_input( - &mut self, - input: I, - ) -> Result<(bool, Option>>>), AflError> { + fn evaluate_input(&mut self, input: &I) -> Result { self.reset_observers()?; self.executor_mut().run_target(&input)?; self.set_executions(self.executions() + 1); @@ -124,29 +121,37 @@ where for feedback in self.feedbacks_mut() { fitness += feedback.is_interesting(&input)?; } - - if fitness > 0 { - let testcase: Rc> = Testcase::new(input).into(); - for feedback in self.feedbacks_mut() { - feedback.append_metadata(testcase.clone())?; - } - testcase.borrow_mut().set_fitness(fitness); - self.corpus_mut().add(testcase.clone()); - Ok((true, Some(testcase))) - } else { - for feedback in self.feedbacks_mut() { - feedback.discard_metadata()?; - } - Ok((false, None)) - } + Ok(fitness) } - fn load_initial_input(&mut self, input: I) -> Result<(), AflError> { - // Inefficent clone, but who cares this is done once at init - let (_, testcase) = self.evaluate_input(input.clone())?; - if testcase.is_none() { - let testcase = Testcase::new(input).into(); - self.corpus_mut().add(testcase); + /// Resets all current feedbacks + fn discard_input(&mut self, input: &I) -> Result<(), AflError> { + // TODO: This could probably be automatic in the feedback somehow? + for feedback in self.feedbacks_mut() { + feedback.discard_metadata(input)?; + } + Ok(()) + } + + /// Creates a new testcase, appending the metadata from each feedback + fn input_to_testcase(&mut self, input: I, fitness: u32) -> Result, AflError> { + let mut testcase = Testcase::new(input); + testcase.set_fitness(fitness); + for feedback in self.feedbacks_mut() { + feedback.append_metadata(&mut testcase)?; + } + + Ok(testcase) + } + + /// Adds this input to the corpus, if it's intersting + fn add_input(&mut self, corpus: &mut C, input: I) -> Result<(), AflError> { + let fitness = self.evaluate_input(&input)?; + if fitness > 0 { + let testcase = self.input_to_testcase(input, fitness)?; + corpus.add(testcase); + } else { + self.discard_input(&input)?; } Ok(()) } @@ -155,6 +160,7 @@ where pub fn generate_initial_inputs( rand: &mut R, state: &mut S, + corpus: &mut C, generator: &mut G, events: &mut EM, num: usize, @@ -170,7 +176,7 @@ where { for _ in 0..num { let input = generator.generate(rand)?; - state.load_initial_input(input)?; + state.add_input(corpus, input)?; fire_event!(events, LoadInitialEvent)?; } events.process(state)?; @@ -190,26 +196,8 @@ where // additional_corpuses: HashMap<&'static str, Box>, observers: Vec>>, feedbacks: Vec>>, - corpus: C, executor: E, - phantom: PhantomData, -} - -impl HasCorpus for StdState -where - C: Corpus, - E: Executor, - I: Input, - R: Rand, -{ - fn corpus(&self) -> &C { - &self.corpus - } - - /// Get thecorpus field (mutable) - fn corpus_mut(&mut self) -> &mut C { - &mut self.corpus - } + phantom: PhantomData<(C, R)>, } impl State for StdState @@ -274,14 +262,13 @@ where I: Input, R: Rand, { - pub fn new(corpus: C, executor: E) -> Self { + pub fn new(executor: E) -> Self { StdState { executions: 0, start_time: current_milliseconds(), metadatas: HashMap::default(), observers: vec![], feedbacks: vec![], - corpus: corpus, executor: executor, phantom: PhantomData, } @@ -309,20 +296,38 @@ where &mut self, rand: &mut R, state: &mut S, + corpus: &mut C, events: &mut EM, ) -> Result { - let (testcase, idx) = state.corpus_mut().next(rand)?; + 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(); + for stage in self.stages_mut() { - stage.perform(rand, state, events, testcase.clone())?; + stage.perform(rand, state, corpus, events, &input)?; } + events.process(state)?; Ok(idx) } - fn fuzz_loop(&mut self, rand: &mut R, state: &mut S, events: &mut EM) -> Result<(), AflError> { + fn fuzz_loop( + &mut self, + rand: &mut R, + state: &mut S, + corpus: &mut C, + events: &mut EM, + ) -> Result<(), AflError> { let mut last = current_milliseconds(); loop { - self.fuzz_one(rand, state, events)?; + self.fuzz_one(rand, state, corpus, events)?; let cur = current_milliseconds(); if cur - last > 60 * 100 { last = cur; @@ -414,7 +419,7 @@ mod tests { corpus.add(testcase); let executor = InMemoryExecutor::::new(harness); - let mut state = StdState::new(corpus, executor); + let mut state = StdState::new(executor); let mut events_manager = LoggerEventManager::new(stderr()); @@ -428,7 +433,7 @@ mod tests { for i in 0..1000 { engine - .fuzz_one(&mut rand, &mut state, &mut events_manager) + .fuzz_one(&mut rand, &mut state, &mut corpus, &mut events_manager) .expect(&format!("Error in iter {}", i)); } } diff --git a/afl/src/events/mod.rs b/afl/src/events/mod.rs index 0faec55cf0..8df9f0a849 100644 --- a/afl/src/events/mod.rs +++ b/afl/src/events/mod.rs @@ -8,8 +8,6 @@ pub mod shmem_translated; #[cfg(feature = "std")] pub use crate::events::llmp::LLMP; -use alloc::rc::Rc; -use core::cell::RefCell; //use core::any::TypeId; #[cfg(feature = "std")] use std::io::Write; @@ -87,7 +85,7 @@ pub struct NewTestcaseEvent where I: Input, { - testcase: Rc>>, + testcase: Testcase, } impl Event for NewTestcaseEvent @@ -103,11 +101,11 @@ impl NewTestcaseEvent where I: Input, { - pub fn new(testcase: Rc>>) -> Self { + pub fn new(testcase: Testcase) -> Self { NewTestcaseEvent { testcase: testcase } } - pub fn testcase(&self) -> &Rc>> { + pub fn testcase(&self) -> &Testcase { &self.testcase } } @@ -170,10 +168,10 @@ where for event in &self.events { writeln!( &mut self.writer, - "#{}\t[{}] corp: {} exec/s: {}", + "#{}\t[{}] exec/s: {}", state.executions(), event, - state.corpus().entries().len(), + //TODO: Count corpus.entries().len(), state.executions_over_seconds() )?; } diff --git a/afl/src/feedbacks/mod.rs b/afl/src/feedbacks/mod.rs index 2f11da8508..f77f3debe3 100644 --- a/afl/src/feedbacks/mod.rs +++ b/afl/src/feedbacks/mod.rs @@ -18,12 +18,12 @@ where fn is_interesting(&mut self, input: &I) -> Result; /// Append to the testcase the generated metadata in case of a new corpus item - fn append_metadata(&mut self, _testcase: Rc>>) -> Result<(), AflError> { + fn append_metadata(&mut self, _testcase: &mut Testcase) -> Result<(), AflError> { Ok(()) } /// Discard the stored metadata in case that the testcase is not added to the corpus - fn discard_metadata(&mut self) -> Result<(), AflError> { + fn discard_metadata(&mut self, _input: &I) -> Result<(), AflError> { Ok(()) } } @@ -230,16 +230,16 @@ where Ok(interesting) } - fn append_metadata(&mut self, testcase: Rc>>) -> Result<(), AflError> { + fn append_metadata(&mut self, testcase: &mut Testcase) -> Result<(), AflError> { let meta = Box::new(MapNoveltiesMetadata::new(core::mem::take( &mut self.novelties, ))); - testcase.borrow_mut().add_metadata(meta); + testcase.add_metadata(meta); Ok(()) } /// Discard the stored metadata in case that the testcase is not added to the corpus - fn discard_metadata(&mut self) -> Result<(), AflError> { + fn discard_metadata(&mut self, _input: &I) -> Result<(), AflError> { self.novelties.clear(); Ok(()) } diff --git a/afl/src/mutators/mod.rs b/afl/src/mutators/mod.rs index 40dc42d6b5..9067e7f8b0 100644 --- a/afl/src/mutators/mod.rs +++ b/afl/src/mutators/mod.rs @@ -4,11 +4,7 @@ pub use scheduled::HavocBytesMutator; pub use scheduled::ScheduledMutator; pub use scheduled::StdScheduledMutator; -use alloc::rc::Rc; -use core::cell::RefCell; - use crate::corpus::Corpus; -use crate::corpus::Testcase; use crate::inputs::Input; use crate::utils::Rand; use crate::AflError; @@ -23,7 +19,7 @@ where fn mutate( &mut self, rand: &mut R, - corpus: &mut C, + corpus: &C, input: &mut I, stage_idx: i32, ) -> Result<(), AflError>; @@ -31,8 +27,8 @@ where /// Post-process given the outcome of the execution fn post_exec( &mut self, - _is_interesting: bool, - _new_testcase: Option>>>, + _is_interesting: u32, + _input: &I, _stage_idx: i32, ) -> Result<(), AflError> { Ok(()) diff --git a/afl/src/mutators/scheduled.rs b/afl/src/mutators/scheduled.rs index 84c91c6279..45a85f12d2 100644 --- a/afl/src/mutators/scheduled.rs +++ b/afl/src/mutators/scheduled.rs @@ -15,7 +15,7 @@ pub enum MutationResult { // TODO maybe the mutator arg is not needed /// The generic function type that identifies mutations type MutationFunction = - fn(&mut M, &mut R, &mut C, &mut I) -> Result; + fn(&mut M, &mut R, &C, &mut I) -> Result; pub trait ComposedByMutations where @@ -66,7 +66,7 @@ where fn scheduled_mutate( &mut self, rand: &mut R, - corpus: &mut C, + corpus: &C, input: &mut I, _stage_idx: i32, ) -> Result<(), AflError> { @@ -96,7 +96,7 @@ where fn mutate( &mut self, rand: &mut R, - corpus: &mut C, + corpus: &C, input: &mut I, _stage_idx: i32, ) -> Result<(), AflError> { @@ -155,7 +155,7 @@ where pub fn mutation_bitflip( _mutator: &mut M, rand: &mut R, - _corpus: &mut C, + _corpus: &C, input: &mut I, ) -> Result where @@ -179,7 +179,7 @@ where pub fn mutation_byteflip( _mutator: &mut M, rand: &mut R, - _corpus: &mut C, + _corpus: &C, input: &mut I, ) -> Result where @@ -203,7 +203,7 @@ where pub fn mutation_byteinc( _mutator: &mut M, rand: &mut R, - _corpus: &mut C, + _corpus: &C, input: &mut I, ) -> Result where @@ -227,7 +227,7 @@ where pub fn mutation_bytedec( _mutator: &mut M, rand: &mut R, - _corpus: &mut C, + _corpus: &C, input: &mut I, ) -> Result where @@ -251,7 +251,7 @@ where pub fn mutation_byteneg( _mutator: &mut M, rand: &mut R, - _corpus: &mut C, + _corpus: &C, input: &mut I, ) -> Result where @@ -292,7 +292,7 @@ fn locate_diffs(this: &[u8], other: &[u8]) -> (i64, i64) { pub fn mutation_splice( _mutator: &mut M, rand: &mut R, - corpus: &mut C, + corpus: &C, input: &mut I, ) -> Result where @@ -301,17 +301,15 @@ where I: Input + HasBytesVec, R: Rand, { - // We don't want to use the testcase we're already using for splicing - let (other_rr, _) = corpus.random_entry(rand)?.clone(); - let mut other_testcase = match other_rr.try_borrow_mut() { - Ok(x) => x, - Err(_) => { - return Ok(MutationResult::Skipped); - } - }; - - let other = other_testcase.load_input()?; + // TODO: Don't reuse the current testcase! + // (We don't want to use the testcase we're already using for splicing) + let (other_testcase, _) = corpus.random_entry(rand)?; + // TODO: let other = other_testcase.load_input()?; // 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 counter = 0; let (first_diff, last_diff) = loop { @@ -362,7 +360,7 @@ where fn mutate( &mut self, rand: &mut R, - corpus: &mut C, + corpus: &C, input: &mut I, stage_idx: i32, ) -> Result<(), AflError> { @@ -380,7 +378,7 @@ where /// Create a new HavocBytesMutator instance given a ScheduledMutator to wrap pub fn new(mut scheduled: SM) -> Self { scheduled.add_mutation(mutation_bitflip); - //scheduled.add_mutation(mutation_splice); + scheduled.add_mutation(mutation_splice); HavocBytesMutator { scheduled: scheduled, phantom: PhantomData, @@ -432,13 +430,15 @@ where #[cfg(test)] mod tests { use crate::inputs::BytesInput; - use crate::mutators::scheduled::{mutation_splice, StdScheduledMutator}; + use crate::mutators::scheduled::StdScheduledMutator; use crate::utils::{Rand, XKCDRand}; use crate::{ corpus::{Corpus, InMemoryCorpus, Testcase}, inputs::HasBytesVec, }; + use super::mutation_splice; + #[test] fn test_mut_splice() { // With the current impl, seed of 1 will result in a split at pos 2. @@ -447,11 +447,10 @@ mod tests { corpus.add(Testcase::new(vec!['a' as u8, 'b' as u8, 'c' as u8]).into()); corpus.add(Testcase::new(vec!['d' as u8, 'e' as u8, 'f' as u8]).into()); - let (testcase_rr, _) = corpus + let (testcase, _) = corpus .next(&mut rand) .expect("Corpus did not contain entries"); - let mut testcase = testcase_rr.borrow_mut(); - let mut input = testcase.load_input().expect("No input in testcase").clone(); + let mut input = testcase.input().as_ref().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 17c8110a78..fe16267980 100644 --- a/afl/src/stages/mod.rs +++ b/afl/src/stages/mod.rs @@ -1,10 +1,6 @@ pub mod mutational; pub use mutational::StdMutationalStage; -use alloc::rc::Rc; -use core::cell::RefCell; - -use crate::corpus::testcase::Testcase; use crate::corpus::Corpus; use crate::engines::State; use crate::events::EventManager; @@ -27,7 +23,8 @@ where &mut self, rand: &mut R, state: &mut S, + corpus: &C, events: &mut EM, - testcase: Rc>>, + input: &I, ) -> Result<(), AflError>; } diff --git a/afl/src/stages/mutational.rs b/afl/src/stages/mutational.rs index 154d3beeda..49433ce4cb 100644 --- a/afl/src/stages/mutational.rs +++ b/afl/src/stages/mutational.rs @@ -1,17 +1,15 @@ -use alloc::rc::Rc; -use core::cell::RefCell; use core::marker::PhantomData; -use crate::corpus::testcase::Testcase; use crate::engines::State; -use crate::events::{EventManager, NewTestcaseEvent}; +use crate::events::EventManager; use crate::executors::Executor; use crate::inputs::Input; use crate::mutators::Mutator; use crate::stages::Corpus; use crate::stages::Stage; use crate::utils::Rand; -use crate::{fire_event, AflError}; +use crate::AflError; +use crate::{events::NewTestcaseEvent, fire_event}; // TODO multi mutators stage @@ -42,22 +40,27 @@ where &mut self, rand: &mut R, state: &mut S, + corpus: &C, events: &mut EM, - testcase: Rc>>, + input: &I, ) -> Result<(), AflError> { let num = self.iterations(rand); for i in 0..num { - let mut input = testcase.borrow_mut().load_input()?.clone(); + let mut input_mut = input.clone(); self.mutator_mut() - .mutate(rand, state.corpus_mut(), &mut input, i as i32)?; + .mutate(rand, corpus, &mut input_mut, i as i32)?; - let (interesting, new_testcase) = state.evaluate_input(input)?; + let interesting = state.evaluate_input(&input_mut)?; self.mutator_mut() - .post_exec(interesting, new_testcase.clone(), i as i32)?; + .post_exec(interesting, &input_mut, i as i32)?; - if !new_testcase.is_none() { - fire_event!(events, NewTestcaseEvent, new_testcase.unwrap())?; + if interesting > 0 { + let new_testcase = state.input_to_testcase(input_mut, interesting)?; + fire_event!(events, NewTestcaseEvent, new_testcase)?; + //state.corpus_mut().add(new_testcase); // TODO: Probably no longer needed, once events work + } else { + state.discard_input(&input_mut)?; } } Ok(()) @@ -115,10 +118,11 @@ where &mut self, rand: &mut R, state: &mut S, + corpus: &C, events: &mut EM, - testcase: Rc>>, + input: &I, ) -> Result<(), AflError> { - self.perform_mutational(rand, state, events, testcase) + self.perform_mutational(rand, state, corpus, events, input) } }