From 382673b1737b41da2cf85027b06823cc03eea6dd Mon Sep 17 00:00:00 2001
From: Langston Barrett <langston.barrett@gmail.com>
Date: Fri, 3 May 2024 13:18:14 -0400
Subject: [PATCH] `GenStage`: A stage for generational (e.g., black-box)
 fuzzers (#2137)

* GenStage: A stage for generational (e.g., black-box) fuzzers

* mv stages/{gen,generation}.rs

* Fix doc link

* `GenStage`: Alphabetize type parameters

* Fighting rustdoc
---
 libafl/src/stages/generation.rs | 72 +++++++++++++++++++++++++++++++++
 libafl/src/stages/mod.rs        |  2 +
 2 files changed, 74 insertions(+)
 create mode 100644 libafl/src/stages/generation.rs

diff --git a/libafl/src/stages/generation.rs b/libafl/src/stages/generation.rs
new file mode 100644
index 0000000000..8df777a28a
--- /dev/null
+++ b/libafl/src/stages/generation.rs
@@ -0,0 +1,72 @@
+//! A [`Stage`] that generates a single input via a
+//! [`crate::generators::Generator`] and evaluates it using the fuzzer, possibly
+//! adding it to the corpus.
+
+use core::marker::PhantomData;
+
+use crate::{
+    generators::Generator,
+    inputs::UsesInput,
+    stages::Stage,
+    state::{HasCorpus, HasRand, UsesState},
+    Error, Evaluator,
+};
+
+/// A [`Stage`] that generates a single input via a [`Generator`] and evaluates
+/// it using the fuzzer, possibly adding it to the corpus.
+///
+/// This stage can be used to construct black-box (e.g., grammar-based) fuzzers.
+#[derive(Debug)]
+pub struct GenStage<G, Z>(G, PhantomData<Z>)
+where
+    Z: UsesState,
+    G: Generator<<<Z as UsesState>::State as UsesInput>::Input, Z::State>;
+
+impl<G, Z> GenStage<G, Z>
+where
+    Z: UsesState,
+    G: Generator<<<Z as UsesState>::State as UsesInput>::Input, Z::State>,
+{
+    /// Create a new [`GenStage`].
+    pub fn new(g: G) -> Self {
+        Self(g, PhantomData)
+    }
+}
+
+impl<G, Z> UsesState for GenStage<G, Z>
+where
+    Z: UsesState,
+    G: Generator<<<Z as UsesState>::State as UsesInput>::Input, Z::State>,
+{
+    type State = Z::State;
+}
+
+impl<E, EM, Z, G> Stage<E, EM, Z> for GenStage<G, Z>
+where
+    E: UsesState<State = Z::State>,
+    EM: UsesState<State = Z::State>,
+    Z: Evaluator<E, EM>,
+    Z::State: HasCorpus + HasRand,
+    G: Generator<<<Z as UsesState>::State as UsesInput>::Input, Z::State>,
+{
+    #[inline]
+    fn perform(
+        &mut self,
+        fuzzer: &mut Z,
+        executor: &mut E,
+        state: &mut Z::State,
+        manager: &mut EM,
+    ) -> Result<(), Error> {
+        let input = self.0.generate(state)?;
+        fuzzer.evaluate_input(state, executor, manager, input)?;
+        Ok(())
+    }
+
+    fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
+        Ok(true)
+    }
+
+    fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> {
+        Ok(())
+    }
+}
diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs
index 29fc9af677..d6e711a3c0 100644
--- a/libafl/src/stages/mod.rs
+++ b/libafl/src/stages/mod.rs
@@ -63,6 +63,8 @@ pub mod concolic;
 #[cfg(feature = "std")]
 pub mod dump;
 pub mod generalization;
+/// The [`generation::GenStage`] generates a single input and evaluates it.
+pub mod generation;
 pub mod logics;
 pub mod power;
 pub mod stats;