diff --git a/libafl/src/events/llmp.rs b/libafl/src/events/llmp.rs index 2c2af29130..051167b732 100644 --- a/libafl/src/events/llmp.rs +++ b/libafl/src/events/llmp.rs @@ -22,6 +22,7 @@ use crate::{ executors::{Executor, HasObservers}, inputs::Input, observers::ObserversTuple, + corpus::CorpusScheduler, state::IfInteresting, stats::Stats, Error, @@ -252,14 +253,16 @@ where } // Handle arriving events in the client - fn handle_in_client( + fn handle_in_client( &mut self, state: &mut S, _sender_id: u32, event: Event, _executor: &mut E, + scheduler: &CS, ) -> Result<(), Error> where + CS: CorpusScheduler, E: Executor + HasObservers, OT: ObserversTuple, { @@ -315,10 +318,11 @@ where } } - fn process(&mut self, state: &mut S, executor: &mut E) -> Result + fn process(&mut self, state: &mut S, executor: &mut E, scheduler: &CS) -> Result where + CS: CorpusScheduler, E: Executor + HasObservers, - OT: ObserversTuple, + OT: ObserversTuple { // TODO: Get around local event copy by moving handle_in_client let mut events = vec![]; @@ -342,7 +346,7 @@ where }; let count = events.len(); events.drain(..).try_for_each(|(sender_id, event)| { - self.handle_in_client(state, sender_id, event, executor) + self.handle_in_client(state, sender_id, event, executor, scheduler) })?; Ok(count) } @@ -426,12 +430,13 @@ where .send_buf(_LLMP_TAG_RESTART, &state_corpus_serialized) } - fn process(&mut self, state: &mut S, executor: &mut E) -> Result + fn process(&mut self, state: &mut S, executor: &mut E, scheduler: &CS) -> Result where + CS: CorpusScheduler, E: Executor + HasObservers, - OT: ObserversTuple, + OT: ObserversTuple { - self.llmp_mgr.process(state, executor) + self.llmp_mgr.process(state, executor, scheduler) } fn fire(&mut self, state: &mut S, event: Event) -> Result<(), Error> { diff --git a/libafl/src/events/mod.rs b/libafl/src/events/mod.rs index 0ff73ead86..934a3163c7 100644 --- a/libafl/src/events/mod.rs +++ b/libafl/src/events/mod.rs @@ -12,6 +12,7 @@ use serde::{Deserialize, Serialize}; use crate::{ executors::{Executor, HasObservers}, inputs::Input, + corpus::{ CorpusScheduler}, observers::ObserversTuple, Error, }; @@ -159,8 +160,9 @@ where /// Lookup for incoming events and process them. /// Return the number of processes events or an error - fn process(&mut self, state: &mut S, executor: &mut E) -> Result + fn process(&mut self, state: &mut S, executor: &mut E, scheduler: &CS) -> Result where + CS: CorpusScheduler, E: Executor + HasObservers, OT: ObserversTuple; @@ -203,8 +205,9 @@ impl EventManager for NopEventManager where I: Input, { - fn process(&mut self, _state: &mut S, _executor: &mut E) -> Result + fn process(&mut self, _state: &mut S, _executor: &mut E, _scheduler: &CS) -> Result where + CS: CorpusScheduler, E: Executor + HasObservers, OT: ObserversTuple, { diff --git a/libafl/src/events/simple.rs b/libafl/src/events/simple.rs index 5fb217ec37..b3f3fce830 100644 --- a/libafl/src/events/simple.rs +++ b/libafl/src/events/simple.rs @@ -7,6 +7,7 @@ use crate::{ executors::{Executor, HasObservers}, inputs::Input, observers::ObserversTuple, + corpus::CorpusScheduler, stats::Stats, Error, }; @@ -30,10 +31,11 @@ where I: Input, ST: Stats, //CE: CustomEvent, { - fn process(&mut self, state: &mut S, _executor: &mut E) -> Result + fn process(&mut self, state: &mut S, _executor: &mut E, _scheduler: &CS) -> Result where + CS: CorpusScheduler, E: Executor + HasObservers, - OT: ObserversTuple, + OT: ObserversTuple { let count = self.events.len(); while self.events.len() > 0 { diff --git a/libafl/src/fuzzer.rs b/libafl/src/fuzzer.rs index da4b20e3f3..127794458a 100644 --- a/libafl/src/fuzzer.rs +++ b/libafl/src/fuzzer.rs @@ -12,12 +12,13 @@ use crate::{ use core::marker::PhantomData; /// Holds a set of stages -pub trait HasStages +pub trait HasStages where - ST: StagesTuple, + ST: StagesTuple, E: Executor, EM: EventManager, I: Input, + Self: Sized, { fn stages(&self) -> &ST; @@ -47,7 +48,7 @@ pub trait Fuzzer { pub struct StdFuzzer where CS: CorpusScheduler, - ST: StagesTuple, + ST: StagesTuple, E: Executor, EM: EventManager, I: Input, @@ -57,10 +58,10 @@ where phantom: PhantomData<(E, EM, I, OT, S)>, } -impl HasStages for StdFuzzer +impl HasStages for StdFuzzer where CS: CorpusScheduler, - ST: StagesTuple, + ST: StagesTuple, E: Executor, EM: EventManager, I: Input, @@ -77,7 +78,7 @@ where impl HasCorpusScheduler for StdFuzzer where CS: CorpusScheduler, - ST: StagesTuple, + ST: StagesTuple, E: Executor, EM: EventManager, I: Input, @@ -95,7 +96,7 @@ impl Fuzzer for StdFuzzer, S: HasExecutions, - ST: StagesTuple, + ST: StagesTuple, EM: EventManager, E: Executor + HasObservers, OT: ObserversTuple, @@ -104,9 +105,9 @@ where fn fuzz_one(&self, state: &mut S, executor: &mut E, manager: &mut EM) -> Result { let idx = self.scheduler().next(state)?; - self.stages().perform_all(state, executor, manager, idx)?; + self.stages().perform_all(state, executor, manager, self.scheduler(), idx)?; - manager.process(state, executor)?; + manager.process(state, executor, self.scheduler())?; Ok(idx) } @@ -133,7 +134,7 @@ where impl StdFuzzer where CS: CorpusScheduler, - ST: StagesTuple, + ST: StagesTuple, E: Executor, EM: EventManager, I: Input, diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs index abb2aadce4..03ada9567d 100644 --- a/libafl/src/stages/mod.rs +++ b/libafl/src/stages/mod.rs @@ -7,7 +7,7 @@ use crate::{ /// A stage is one step in the fuzzing process. /// Multiple stages will be scheduled one by one for each input. -pub trait Stage +pub trait Stage where EM: EventManager, E: Executor, @@ -19,11 +19,12 @@ where state: &mut S, executor: &mut E, manager: &mut EM, + scheduler: &CS, corpus_idx: usize, ) -> Result<(), Error>; } -pub trait StagesTuple +pub trait StagesTuple where EM: EventManager, E: Executor, @@ -34,25 +35,26 @@ where state: &mut S, executor: &mut E, manager: &mut EM, + scheduler: &CS, corpus_idx: usize, ) -> Result<(), Error>; } -impl StagesTuple for () +impl StagesTuple for () where EM: EventManager, E: Executor, I: Input, { - fn perform_all(&self, _: &mut S, _: &mut E, _: &mut EM, _: usize) -> Result<(), Error> { + fn perform_all(&self, _: &mut S, _: &mut E, _: &mut EM, _: &CS, _: usize) -> Result<(), Error> { Ok(()) } } -impl StagesTuple for (Head, Tail) +impl StagesTuple for (Head, Tail) where - Head: Stage, - Tail: StagesTuple + TupleList, + Head: Stage, + Tail: StagesTuple + TupleList, EM: EventManager, E: Executor, I: Input, @@ -62,9 +64,10 @@ where state: &mut S, executor: &mut E, manager: &mut EM, + scheduler: &CS, corpus_idx: usize, ) -> Result<(), Error> { - self.0.perform(state, executor, manager, corpus_idx)?; - self.1.perform_all(state, executor, manager, corpus_idx) + self.0.perform(state, executor, manager, scheduler, corpus_idx)?; + self.1.perform_all(state, executor, manager, scheduler, corpus_idx) } } diff --git a/libafl/src/stages/mutational.rs b/libafl/src/stages/mutational.rs index bf891f510f..7a2731a058 100644 --- a/libafl/src/stages/mutational.rs +++ b/libafl/src/stages/mutational.rs @@ -1,7 +1,7 @@ use core::marker::PhantomData; use crate::{ - corpus::Corpus, + corpus::{Corpus, CorpusScheduler}, events::EventManager, executors::{Executor, HasObservers}, inputs::Input, @@ -18,7 +18,7 @@ use crate::{ /// A Mutational stage is the stage in a fuzzing run that mutates inputs. /// Mutational stages will usually have a range of mutations that are /// being applied to the input one by one, between executions. -pub trait MutationalStage: Stage +pub trait MutationalStage: Stage where M: Mutator, I: Input, @@ -27,6 +27,7 @@ where EM: EventManager, E: Executor + HasObservers, OT: ObserversTuple, + CS: CorpusScheduler { /// The mutator registered for this stage fn mutator(&self) -> &M; @@ -43,6 +44,7 @@ where state: &mut S, executor: &mut E, manager: &mut EM, + scheduler: &CS, corpus_idx: usize, ) -> Result<(), Error> { let num = self.iterations(state); @@ -67,7 +69,7 @@ pub static DEFAULT_MUTATIONAL_MAX_ITERATIONS: u64 = 128; /// The default mutational stage #[derive(Clone, Debug)] -pub struct StdMutationalStage +pub struct StdMutationalStage where M: Mutator, I: Input, @@ -76,14 +78,15 @@ where EM: EventManager, E: Executor + HasObservers, OT: ObserversTuple, + CS: CorpusScheduler, R: Rand, { mutator: M, - phantom: PhantomData<(C, E, EM, I, OT, R, S)>, + phantom: PhantomData<(C, CS, E, EM, I, OT, R, S)>, } -impl MutationalStage - for StdMutationalStage +impl MutationalStage + for StdMutationalStage where M: Mutator, I: Input, @@ -92,6 +95,7 @@ where EM: EventManager, E: Executor + HasObservers, OT: ObserversTuple, + CS: CorpusScheduler, R: Rand, { /// The mutator, added to this stage @@ -112,7 +116,7 @@ where } } -impl Stage for StdMutationalStage +impl Stage for StdMutationalStage where M: Mutator, I: Input, @@ -121,6 +125,7 @@ where EM: EventManager, E: Executor + HasObservers, OT: ObserversTuple, + CS: CorpusScheduler, R: Rand, { #[inline] @@ -129,13 +134,14 @@ where state: &mut S, executor: &mut E, manager: &mut EM, + scheduler: &CS, corpus_idx: usize, ) -> Result<(), Error> { - self.perform_mutational(state, executor, manager, corpus_idx) + self.perform_mutational(state, executor, manager, scheduler, corpus_idx) } } -impl StdMutationalStage +impl StdMutationalStage where M: Mutator, I: Input, @@ -144,6 +150,7 @@ where EM: EventManager, E: Executor + HasObservers, OT: ObserversTuple, + CS: CorpusScheduler, R: Rand, { /// Creates a new default mutational stage diff --git a/libafl/src/state/mod.rs b/libafl/src/state/mod.rs index 7467fe3cd7..e1d29ecb37 100644 --- a/libafl/src/state/mod.rs +++ b/libafl/src/state/mod.rs @@ -600,7 +600,7 @@ where phantom: PhantomData, }, )?; - manager.process(self, executor)?; + manager.process(self, executor, scheduler)?; Ok(()) } } @@ -680,7 +680,7 @@ where phantom: PhantomData, }, )?; - manager.process(self, executor)?; + manager.process(self, executor, scheduler)?; Ok(()) }