From 20e5500d93720976353499f2cc141da7517766c1 Mon Sep 17 00:00:00 2001
From: Andrea Fioraldi <andreafioraldi@gmail.com>
Date: Fri, 12 Nov 2021 14:50:50 +0100
Subject: [PATCH] ClosureStage

---
 libafl/src/observers/cmp.rs |  2 +-
 libafl/src/stages/mod.rs    | 49 +++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/libafl/src/observers/cmp.rs b/libafl/src/observers/cmp.rs
index d306da7509..79b2025588 100644
--- a/libafl/src/observers/cmp.rs
+++ b/libafl/src/observers/cmp.rs
@@ -5,7 +5,7 @@ use alloc::{
     vec::Vec,
 };
 
-use serde::{de::DeserializeOwned, Deserialize, Serialize, Serializer};
+use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 use crate::{
     bolts::{ownedref::OwnedRefMut, tuples::Named, AsSlice},
diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs
index e2d740f68a..33362df5c8 100644
--- a/libafl/src/stages/mod.rs
+++ b/libafl/src/stages/mod.rs
@@ -27,6 +27,8 @@ pub use concolic::ConcolicTracingStage;
 pub use concolic::SimpleConcolicMutationalStage;
 
 use crate::Error;
+use core::{convert::From, marker::PhantomData};
+
 /// A stage is one step in the fuzzing process.
 /// Multiple stages will be scheduled one by one for each input.
 pub trait Stage<E, EM, S, Z> {
@@ -89,3 +91,50 @@ where
             .perform_all(fuzzer, executor, state, manager, corpus_idx)
     }
 }
+
+pub struct ClosureStage<CB, E, EM, S, Z>
+where
+    CB: FnMut(&mut Z, &mut E, &mut S, &mut EM, usize) -> Result<(), Error>,
+{
+    closure: CB,
+    phantom: PhantomData<(E, EM, S, Z)>,
+}
+
+impl<CB, E, EM, S, Z> Stage<E, EM, S, Z> for ClosureStage<CB, E, EM, S, Z>
+where
+    CB: FnMut(&mut Z, &mut E, &mut S, &mut EM, usize) -> Result<(), Error>,
+{
+    fn perform(
+        &mut self,
+        fuzzer: &mut Z,
+        executor: &mut E,
+        state: &mut S,
+        manager: &mut EM,
+        corpus_idx: usize,
+    ) -> Result<(), Error> {
+        (self.closure)(fuzzer, executor, state, manager, corpus_idx)
+    }
+}
+
+impl<CB, E, EM, S, Z> ClosureStage<CB, E, EM, S, Z>
+where
+    CB: FnMut(&mut Z, &mut E, &mut S, &mut EM, usize) -> Result<(), Error>,
+{
+    #[must_use]
+    pub fn new(closure: CB) -> Self {
+        Self {
+            closure,
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl<CB, E, EM, S, Z> From<CB> for ClosureStage<CB, E, EM, S, Z>
+where
+    CB: FnMut(&mut Z, &mut E, &mut S, &mut EM, usize) -> Result<(), Error>,
+{
+    #[must_use]
+    fn from(closure: CB) -> Self {
+        Self::new(closure)
+    }
+}