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 mod owned;
|
||||||
pub use owned::StagesOwnedList;
|
pub use owned::StagesOwnedList;
|
||||||
|
|
||||||
|
pub mod logics;
|
||||||
|
pub use logics::*;
|
||||||
|
|
||||||
pub mod tuneable;
|
pub mod tuneable;
|
||||||
pub use 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
|
/// `Stage` Python bindings
|
||||||
#[cfg(feature = "python")]
|
#[cfg(feature = "python")]
|
||||||
#[allow(missing_docs)]
|
#[allow(missing_docs)]
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user