added corpus access to perform

src/engines/mod.rs

@@ -1,5 +1,7 @@
+//! The engine is the core piece of every good fuzzer
 extern crate alloc;
-use crate::corpus::testcase::Testcase;
+
+use crate::corpus::Corpus;
 use crate::feedbacks::Feedback;
 use crate::inputs::Input;
 use crate::stages::Stage;
@@ -8,8 +10,9 @@ use crate::AflError;
 use alloc::rc::Rc;
 use core::cell::RefCell;
 
-pub trait Engine<I>
+pub trait Engine<C, I>
 where
+    C: Corpus<I>,
     I: Input,
 {
     fn feedbacks(&self) -> &Vec<Box<dyn Feedback<I>>>;
@@ -20,32 +23,34 @@ where
         self.feedbacks_mut().push(feedback);
     }
 
-    fn stages(&self) -> &Vec<Box<dyn Stage<I>>>;
+    fn stages(&self) -> &Vec<Box<dyn Stage<C, I>>>;
 
-    fn stages_mut(&mut self) -> &mut Vec<Box<dyn Stage<I>>>;
+    fn stages_mut(&mut self) -> &mut Vec<Box<dyn Stage<C, I>>>;
 
-    fn add_stage(&mut self, stage: Box<dyn Stage<I>>) {
+    fn add_stage(&mut self, stage: Box<dyn Stage<C, I>>) {
         self.stages_mut().push(stage);
     }
 
-    fn fuzz_one(&mut self, testcase: &Rc<RefCell<Testcase<I>>>) -> Result<(), AflError> {
+    fn fuzz_one(&mut self, corpus: &mut C) -> Result<(), AflError> {
         for stage in self.stages_mut() {
-            stage.perform(&testcase)?;
+            stage.perform(corpus)?;
         }
         Ok(())
     }
 }
 
-pub struct DefaultEngine<I>
+pub struct DefaultEngine<C, I>
 where
+    C: Corpus<I>,
     I: Input,
 {
     feedbacks: Vec<Box<dyn Feedback<I>>>,
-    stages: Vec<Box<dyn Stage<I>>>,
+    stages: Vec<Box<dyn Stage<C, I>>>,
 }
 
-impl<I> Engine<I> for DefaultEngine<I>
+impl<C, I> Engine<C, I> for DefaultEngine<C, I>
 where
+    C: Corpus<I>,
     I: Input,
 {
     fn feedbacks(&self) -> &Vec<Box<dyn Feedback<I>>> {
@@ -56,17 +61,18 @@ where
         &mut self.feedbacks
     }
 
-    fn stages(&self) -> &Vec<Box<dyn Stage<I>>> {
+    fn stages(&self) -> &Vec<Box<dyn Stage<C, I>>> {
         &self.stages
     }
 
-    fn stages_mut(&mut self) -> &mut Vec<Box<dyn Stage<I>>> {
+    fn stages_mut(&mut self) -> &mut Vec<Box<dyn Stage<C, I>>> {
         &mut self.stages
     }
 }
 
-impl<I> DefaultEngine<I>
+impl<C, I> DefaultEngine<C, I>
 where
+    C: Corpus<I>,
     I: Input,
 {
     pub fn new() -> Self {
@@ -92,6 +98,8 @@ mod tests {
         mutation_bitflip, ComposedByMutations, DefaultScheduledMutator,
     };
     use crate::stages::mutational::DefaultMutationalStage;
+    use alloc::rc::Rc;
+    use core::cell::RefCell;
 
     use crate::utils::Xoshiro256StarRand;
 
@@ -106,12 +114,16 @@ mod tests {
         let mut corpus = InMemoryCorpus::<BytesInput, _>::new(&rand);
         let testcase = Testcase::new_rr(BytesInput::new(vec![0; 4]));
         corpus.add(testcase);
-        let executor = InMemoryExecutor::new_rr(harness);
+        let executor: Rc<RefCell<InMemoryExecutor<BytesInput>>> = InMemoryExecutor::new_rr(harness);
         let mut engine = DefaultEngine::new();
         let mut mutator = DefaultScheduledMutator::new(&rand);
         mutator.add_mutation(mutation_bitflip);
         let stage = DefaultMutationalStage::new(&rand, &executor, mutator);
         engine.add_stage(Box::new(stage));
-        engine.fuzz_one(&corpus.next().unwrap()).unwrap();
+        for i in 0..1000 {
+            engine
+                .fuzz_one(&mut corpus)
+                .expect(&format!("Error in iter {}", i));
+        }
     }
 }

src/stages/mod.rs

@@ -1,18 +1,16 @@
 extern crate alloc;
 pub mod mutational;
+use crate::corpus::Corpus;
 pub use mutational::DefaultMutationalStage;
 
-use crate::corpus::Testcase;
 use crate::inputs::Input;
 use crate::AflError;
 
-use alloc::rc::Rc;
+pub trait Stage<C, I>
-use core::cell::RefCell;
-
-pub trait Stage<I>
 where
+    C: Corpus<I>,
     I: Input,
 {
     /// Run the stage
-    fn perform(&mut self, entry: &Rc<RefCell<Testcase<I>>>) -> Result<(), AflError>;
+    fn perform(&mut self, corpus: &mut C) -> Result<(), AflError>;
 }

src/stages/mutational.rs

@@ -3,6 +3,7 @@ use crate::corpus::testcase::Testcase;
 use crate::executors::Executor;
 use crate::inputs::Input;
 use crate::mutators::Mutator;
+use crate::stages::Corpus;
 use crate::stages::Stage;
 use crate::utils::{HasRand, Rand};
 use crate::AflError;
@@ -13,8 +14,9 @@ use core::marker::PhantomData;
 
 // TODO create HasMutatorsVec trait
 
-pub trait MutationalStage<I, M, E>: Stage<I> + HasRand
+pub trait MutationalStage<C, I, M, E>: Stage<C, I> + HasRand
 where
+    C: Corpus<I>,
     I: Input,
     M: Mutator<I, R = Self::R>,
     E: Executor<I>,
@@ -34,9 +36,9 @@ where
         1 + self.rand_below(128) as usize
     }
 
-    // TODO: We need a way to get testcases for splicing
     /// Runs this (mutational) stage for the given testcase
-    fn perform_mutational(&mut self, testcase: &Rc<RefCell<Testcase<I>>>) -> Result<(), AflError> {
+    fn perform_mutational(&mut self, corpus: &mut C) -> Result<(), AflError> {
+        let testcase = corpus.next()?;
         let num = self.iterations();
         let input = testcase.borrow_mut().load_input()?.clone();
 
@@ -47,6 +49,10 @@ where
             let interesting = self.executor().borrow_mut().evaluate_input(&input_tmp)?;
 
             self.mutator_mut().post_exec(interesting, i as i32)?;
+
+            if interesting {
+                corpus.add(Testcase::new_rr(input_tmp));
+            }
         }
         Ok(())
     }
@@ -80,8 +86,9 @@ where
     }
 }
 
-impl<I, R, M, E> MutationalStage<I, M, E> for DefaultMutationalStage<I, R, M, E>
+impl<C, I, R, M, E> MutationalStage<C, I, M, E> for DefaultMutationalStage<I, R, M, E>
 where
+    C: Corpus<I>,
     I: Input,
     R: Rand,
     M: Mutator<I, R = R>,
@@ -102,15 +109,16 @@ where
     }
 }
 
-impl<I, R, M, E> Stage<I> for DefaultMutationalStage<I, R, M, E>
+impl<C, I, R, M, E> Stage<C, I> for DefaultMutationalStage<I, R, M, E>
 where
+    C: Corpus<I>,
     I: Input,
     R: Rand,
     M: Mutator<I, R = R>,
     E: Executor<I>,
 {
-    fn perform(&mut self, testcase: &Rc<RefCell<Testcase<I>>>) -> Result<(), AflError> {
+    fn perform(&mut self, corpus: &mut C) -> Result<(), AflError> {
-        self.perform_mutational(testcase)
+        self.perform_mutational(corpus)
     }
 }