Skippable stage, generator wrapper for Grimoire (#748)

* Skippable stage, generator wrapper for Grimoire

* more fancy wrapper

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<G, S> {
+    bytes_generator: G,
+    phantom: PhantomData<S>,
+}
+
+impl<G, S> GeneralizedInputBytesGenerator<G, S>
+where
+    S: HasRand,
+    G: Generator<BytesInput, S>,
+{
+    /// Creates a new [`GeneralizedInputBytesGenerator`] by wrapping a bytes generator.
+    pub fn new(bytes_generator: G) -> Self {
+        Self {
+            bytes_generator,
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl<G, S> From<G> for GeneralizedInputBytesGenerator<G, S>
+where
+    S: HasRand,
+    G: Generator<BytesInput, S>,
+{
+    fn from(bytes_generator: G) -> Self {
+        Self::new(bytes_generator)
+    }
+}
+
+impl<G, S> Generator<GeneralizedInput, S> for GeneralizedInputBytesGenerator<G, S>
+where
+    S: HasRand,
+    G: Generator<BytesInput, S>,
+{
+    fn generate(&mut self, state: &mut S) -> Result<GeneralizedInput, Error> {
+        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<S>

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<u8>,
+    pub(crate) bytes: Vec<u8>,
 }
 
 impl Input for BytesInput {

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<BytesInput> 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]

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<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, S, ST, Z>
+where
+    CD: FnMut(&mut S) -> SkippableStageDecision,
+    ST: Stage<E, EM, S, Z>,
+{
+    wrapped_stage: ST,
+    condition: CD,
+    phantom: PhantomData<(E, EM, S, Z)>,
+}
+
+impl<CD, E, EM, S, ST, Z> SkippableStage<CD, E, EM, S, ST, Z>
+where
+    CD: FnMut(&mut S) -> SkippableStageDecision,
+    ST: Stage<E, EM, S, Z>,
+{
+    /// Create a new [`SkippableStage`]
+    pub fn new(wrapped_stage: ST, condition: CD) -> Self {
+        Self {
+            wrapped_stage,
+            condition,
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl<CD, E, EM, S, ST, Z> Stage<E, EM, S, Z> for SkippableStage<CD, E, EM, S, ST, Z>
+where
+    CD: FnMut(&mut S) -> SkippableStageDecision,
+    ST: Stage<E, EM, S, Z>,
+{
+    /// 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)]