diff --git a/libafl/src/stages/logics.rs b/libafl/src/stages/logics.rs new file mode 100644 index 0000000000..49347910f7 --- /dev/null +++ b/libafl/src/stages/logics.rs @@ -0,0 +1,151 @@ +//! Stage wrappers that add logics to stage list + +use core::marker::PhantomData; + +use crate::{ + corpus::CorpusId, + stages::{Stage, StagesTuple}, + state::UsesState, + Error, +}; + +#[derive(Debug)] +/// Perform the stage while closure evaluates to true +pub struct WhileStage +where + CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result, + E: UsesState, + EM: UsesState, + ST: StagesTuple, + Z: UsesState, +{ + closure: CB, + stages: ST, + phantom: PhantomData<(E, EM, Z)>, +} + +impl UsesState for WhileStage +where + CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result, + E: UsesState, + EM: UsesState, + ST: StagesTuple, + Z: UsesState, +{ + type State = E::State; +} + +impl Stage for WhileStage +where + CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result, + E: UsesState, + EM: UsesState, + ST: StagesTuple, + Z: UsesState, +{ + fn perform( + &mut self, + fuzzer: &mut Z, + executor: &mut E, + state: &mut E::State, + manager: &mut EM, + corpus_idx: CorpusId, + ) -> Result<(), Error> { + while (self.closure)(fuzzer, executor, state, manager, corpus_idx)? { + self.stages + .perform_all(fuzzer, executor, state, manager, corpus_idx)?; + } + Ok(()) + } +} + +impl WhileStage +where + CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result, + E: UsesState, + EM: UsesState, + ST: StagesTuple, + Z: UsesState, +{ + /// Constructor + pub fn new(closure: CB, stages: ST) -> Self { + Self { + closure, + stages, + phantom: PhantomData, + } + } +} + +/// Perform the stage if closure evaluates to true +#[derive(Debug)] +pub struct IfElseStage +where + CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result, + E: UsesState, + EM: UsesState, + ST: StagesTuple, + Z: UsesState, +{ + closure: CB, + if_stages: ST, + else_stages: ST, + phantom: PhantomData<(E, EM, Z)>, +} + +impl UsesState for IfElseStage +where + CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result, + E: UsesState, + EM: UsesState, + ST: StagesTuple, + Z: UsesState, +{ + type State = E::State; +} + +impl Stage for IfElseStage +where + CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result, + E: UsesState, + EM: UsesState, + ST: StagesTuple, + Z: UsesState, +{ + fn perform( + &mut self, + fuzzer: &mut Z, + executor: &mut E, + state: &mut E::State, + manager: &mut EM, + corpus_idx: CorpusId, + ) -> Result<(), Error> { + if (self.closure)(fuzzer, executor, state, manager, corpus_idx)? { + self.if_stages + .perform_all(fuzzer, executor, state, manager, corpus_idx)?; + } else { + self.else_stages + .perform_all(fuzzer, executor, state, manager, corpus_idx)?; + } + Ok(()) + } +} + +impl IfElseStage +where + CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result, + E: UsesState, + EM: UsesState, + ST: StagesTuple, + Z: UsesState, +{ + /// Constructor + pub fn new(closure: CB, if_stages: ST, else_stages: ST) -> Self { + Self { + closure, + if_stages, + else_stages, + phantom: PhantomData, + } + } +} diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs index 32f0295213..bec7eb5c07 100644 --- a/libafl/src/stages/mod.rs +++ b/libafl/src/stages/mod.rs @@ -30,6 +30,9 @@ pub use generalization::GeneralizationStage; pub mod owned; pub use owned::StagesOwnedList; +pub mod logics; +pub use logics::*; + pub mod tuneable; pub use tuneable::*; @@ -297,90 +300,6 @@ where } } -/// The decision if the [`SkippableStage`] should be skipped -#[derive(Debug, Clone, Copy, Eq, PartialEq)] -pub enum SkippableStageDecision { - /// Return to indicate that this [`Stage`] should be executed - Perform, - /// Return to indicate that this [`Stage`] should be skipped - Skip, -} - -impl From for SkippableStageDecision { - fn from(b: bool) -> SkippableStageDecision { - if b { - SkippableStageDecision::Perform - } else { - SkippableStageDecision::Skip - } - } -} - -/// The [`SkippableStage`] wraps any [`Stage`] so that it can be skipped, according to a condition. -#[derive(Debug, Clone)] -pub struct SkippableStage { - wrapped_stage: ST, - condition: CD, - phantom: PhantomData<(E, EM, Z)>, -} - -impl SkippableStage -where - CD: FnMut(&mut ST::State) -> SkippableStageDecision, - ST: Stage, - E: UsesState, - EM: UsesState, - Z: UsesState, -{ - /// Create a new [`SkippableStage`] - pub fn new(wrapped_stage: ST, condition: CD) -> Self { - Self { - wrapped_stage, - condition, - phantom: PhantomData, - } - } -} - -impl UsesState for SkippableStage -where - CD: FnMut(&mut ST::State) -> SkippableStageDecision, - ST: Stage, - E: UsesState, - EM: UsesState, - Z: UsesState, -{ - type State = ST::State; -} - -impl Stage for SkippableStage -where - CD: FnMut(&mut ST::State) -> SkippableStageDecision, - ST: Stage, - E: UsesState, - EM: UsesState, - Z: UsesState, -{ - /// Run the stage - #[inline] - fn perform( - &mut self, - fuzzer: &mut Z, - executor: &mut E, - state: &mut ST::State, - manager: &mut EM, - corpus_idx: CorpusId, - ) -> Result<(), Error> { - let condition = &mut self.condition; - if condition(state) == SkippableStageDecision::Perform { - self.wrapped_stage - .perform(fuzzer, executor, state, manager, corpus_idx) - } else { - Ok(()) - } - } -} - /// `Stage` Python bindings #[cfg(feature = "python")] #[allow(missing_docs)]