Logic stages (#1148)
* IfStage * fmt clp * constructor * fmt * else * fmt * while logic * delete skippable * fmt --------- Co-authored-by: Andrea Fioraldi <andreafioraldi@gmail.com>
This commit is contained in:
Dongjia "toka" Zhang
GitHub
parent
306cdcd800
commit
8245c7eda9
151
libafl/src/stages/logics.rs
Normal file
151
libafl/src/stages/logics.rs
Normal file
@ -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<CB, E, EM, ST, Z>
|
||||
where
|
||||
CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
|
||||
E: UsesState,
|
||||
EM: UsesState<State = E::State>,
|
||||
ST: StagesTuple<E, EM, E::State, Z>,
|
||||
Z: UsesState<State = E::State>,
|
||||
{
|
||||
closure: CB,
|
||||
stages: ST,
|
||||
phantom: PhantomData<(E, EM, Z)>,
|
||||
}
|
||||
|
||||
impl<CB, E, EM, ST, Z> UsesState for WhileStage<CB, E, EM, ST, Z>
|
||||
where
|
||||
CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
|
||||
E: UsesState,
|
||||
EM: UsesState<State = E::State>,
|
||||
ST: StagesTuple<E, EM, E::State, Z>,
|
||||
Z: UsesState<State = E::State>,
|
||||
{
|
||||
type State = E::State;
|
||||
}
|
||||
|
||||
impl<CB, E, EM, ST, Z> Stage<E, EM, Z> for WhileStage<CB, E, EM, ST, Z>
|
||||
where
|
||||
CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
|
||||
E: UsesState,
|
||||
EM: UsesState<State = E::State>,
|
||||
ST: StagesTuple<E, EM, E::State, Z>,
|
||||
Z: UsesState<State = E::State>,
|
||||
{
|
||||
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<CB, E, EM, ST, Z> WhileStage<CB, E, EM, ST, Z>
|
||||
where
|
||||
CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
|
||||
E: UsesState,
|
||||
EM: UsesState<State = E::State>,
|
||||
ST: StagesTuple<E, EM, E::State, Z>,
|
||||
Z: UsesState<State = E::State>,
|
||||
{
|
||||
/// 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<CB, E, EM, ST, Z>
|
||||
where
|
||||
CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
|
||||
E: UsesState,
|
||||
EM: UsesState<State = E::State>,
|
||||
ST: StagesTuple<E, EM, E::State, Z>,
|
||||
Z: UsesState<State = E::State>,
|
||||
{
|
||||
closure: CB,
|
||||
if_stages: ST,
|
||||
else_stages: ST,
|
||||
phantom: PhantomData<(E, EM, Z)>,
|
||||
}
|
||||
|
||||
impl<CB, E, EM, ST, Z> UsesState for IfElseStage<CB, E, EM, ST, Z>
|
||||
where
|
||||
CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
|
||||
E: UsesState,
|
||||
EM: UsesState<State = E::State>,
|
||||
ST: StagesTuple<E, EM, E::State, Z>,
|
||||
Z: UsesState<State = E::State>,
|
||||
{
|
||||
type State = E::State;
|
||||
}
|
||||
|
||||
impl<CB, E, EM, ST, Z> Stage<E, EM, Z> for IfElseStage<CB, E, EM, ST, Z>
|
||||
where
|
||||
CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
|
||||
E: UsesState,
|
||||
EM: UsesState<State = E::State>,
|
||||
ST: StagesTuple<E, EM, E::State, Z>,
|
||||
Z: UsesState<State = E::State>,
|
||||
{
|
||||
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<CB, E, EM, ST, Z> IfElseStage<CB, E, EM, ST, Z>
|
||||
where
|
||||
CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
|
||||
E: UsesState,
|
||||
EM: UsesState<State = E::State>,
|
||||
ST: StagesTuple<E, EM, E::State, Z>,
|
||||
Z: UsesState<State = E::State>,
|
||||
{
|
||||
/// Constructor
|
||||
pub fn new(closure: CB, if_stages: ST, else_stages: ST) -> Self {
|
||||
Self {
|
||||
closure,
|
||||
if_stages,
|
||||
else_stages,
|
||||
phantom: PhantomData,
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -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<bool> 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<CD, E, EM, ST, Z> {
|
||||
wrapped_stage: ST,
|
||||
condition: CD,
|
||||
phantom: PhantomData<(E, EM, Z)>,
|
||||
}
|
||||
|
||||
impl<CD, E, EM, ST, Z> SkippableStage<CD, E, EM, ST, Z>
|
||||
where
|
||||
CD: FnMut(&mut ST::State) -> SkippableStageDecision,
|
||||
ST: Stage<E, EM, Z>,
|
||||
E: UsesState<State = ST::State>,
|
||||
EM: UsesState<State = ST::State>,
|
||||
Z: UsesState<State = ST::State>,
|
||||
{
|
||||
/// Create a new [`SkippableStage`]
|
||||
pub fn new(wrapped_stage: ST, condition: CD) -> Self {
|
||||
Self {
|
||||
wrapped_stage,
|
||||
condition,
|
||||
phantom: PhantomData,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<CD, E, EM, ST, Z> UsesState for SkippableStage<CD, E, EM, ST, Z>
|
||||
where
|
||||
CD: FnMut(&mut ST::State) -> SkippableStageDecision,
|
||||
ST: Stage<E, EM, Z>,
|
||||
E: UsesState<State = ST::State>,
|
||||
EM: UsesState<State = ST::State>,
|
||||
Z: UsesState<State = ST::State>,
|
||||
{
|
||||
type State = ST::State;
|
||||
}
|
||||
|
||||
impl<CD, E, EM, ST, Z> Stage<E, EM, Z> for SkippableStage<CD, E, EM, ST, Z>
|
||||
where
|
||||
CD: FnMut(&mut ST::State) -> SkippableStageDecision,
|
||||
ST: Stage<E, EM, Z>,
|
||||
E: UsesState<State = ST::State>,
|
||||
EM: UsesState<State = ST::State>,
|
||||
Z: UsesState<State = ST::State>,
|
||||
{
|
||||
/// 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)]
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user