diff --git a/libafl/src/generators/mod.rs b/libafl/src/generators/mod.rs index 6a7dbe1c13..d8d92ca8b1 100644 --- a/libafl/src/generators/mod.rs +++ b/libafl/src/generators/mod.rs @@ -5,7 +5,7 @@ use core::{cmp::min, marker::PhantomData}; use crate::{ bolts::rands::Rand, - inputs::{bytes::BytesInput, Input}, + inputs::{bytes::BytesInput, GeneralizedInput, Input}, state::HasRand, Error, }; @@ -33,6 +33,51 @@ where fn generate_dummy(&self, state: &mut S) -> I; } +/// A Generator that produces [`GeneralizedInput`]s from a wrapped [`BytesInput`] generator +#[derive(Clone, Debug)] +pub struct GeneralizedInputBytesGenerator { + bytes_generator: G, + phantom: PhantomData, +} + +impl GeneralizedInputBytesGenerator +where + S: HasRand, + G: Generator, +{ + /// Creates a new [`GeneralizedInputBytesGenerator`] by wrapping a bytes generator. + pub fn new(bytes_generator: G) -> Self { + Self { + bytes_generator, + phantom: PhantomData, + } + } +} + +impl From for GeneralizedInputBytesGenerator +where + S: HasRand, + G: Generator, +{ + fn from(bytes_generator: G) -> Self { + Self::new(bytes_generator) + } +} + +impl Generator for GeneralizedInputBytesGenerator +where + S: HasRand, + G: Generator, +{ + fn generate(&mut self, state: &mut S) -> Result { + Ok(self.bytes_generator.generate(state)?.into()) + } + + fn generate_dummy(&self, state: &mut S) -> GeneralizedInput { + self.bytes_generator.generate_dummy(state).into() + } +} + #[derive(Clone, Debug)] /// Generates random bytes pub struct RandBytesGenerator diff --git a/libafl/src/inputs/bytes.rs b/libafl/src/inputs/bytes.rs index 2ab274d9fc..9e3b8353ed 100644 --- a/libafl/src/inputs/bytes.rs +++ b/libafl/src/inputs/bytes.rs @@ -20,7 +20,7 @@ use crate::{ #[derive(Serialize, Deserialize, Clone, Debug, Default, PartialEq, Eq, Hash)] pub struct BytesInput { /// The raw input bytes - bytes: Vec, + pub(crate) bytes: Vec, } impl Input for BytesInput { diff --git a/libafl/src/inputs/generalized.rs b/libafl/src/inputs/generalized.rs index 50a5030c53..a32040b9c2 100644 --- a/libafl/src/inputs/generalized.rs +++ b/libafl/src/inputs/generalized.rs @@ -12,7 +12,7 @@ use serde::{Deserialize, Serialize}; use crate::Error; use crate::{ bolts::{ownedref::OwnedSlice, HasLen}, - inputs::{HasBytesVec, HasTargetBytes, Input}, + inputs::{BytesInput, HasBytesVec, HasTargetBytes, Input}, }; /// An item of the generalized input @@ -122,6 +122,18 @@ impl From<&[u8]> for GeneralizedInput { } } +impl From for GeneralizedInput { + fn from(bytes_input: BytesInput) -> Self { + Self::new(bytes_input.bytes) + } +} + +impl From<&BytesInput> for GeneralizedInput { + fn from(bytes_input: &BytesInput) -> Self { + bytes_input.bytes().into() + } +} + impl GeneralizedInput { /// Creates a new bytes input using the given bytes #[must_use] diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs index 844aec861b..0a32c006e3 100644 --- a/libafl/src/stages/mod.rs +++ b/libafl/src/stages/mod.rs @@ -260,6 +260,77 @@ 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 +where + CD: FnMut(&mut S) -> SkippableStageDecision, + ST: Stage, +{ + wrapped_stage: ST, + condition: CD, + phantom: PhantomData<(E, EM, S, Z)>, +} + +impl SkippableStage +where + CD: FnMut(&mut S) -> SkippableStageDecision, + ST: Stage, +{ + /// Create a new [`SkippableStage`] + pub fn new(wrapped_stage: ST, condition: CD) -> Self { + Self { + wrapped_stage, + condition, + phantom: PhantomData, + } + } +} + +impl Stage for SkippableStage +where + CD: FnMut(&mut S) -> SkippableStageDecision, + ST: Stage, +{ + /// Run the stage + #[inline] + fn perform( + &mut self, + fuzzer: &mut Z, + executor: &mut E, + state: &mut S, + manager: &mut EM, + corpus_idx: usize, + ) -> 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)]