Resumable stages redux

fuzzers/baby_fuzzer_minimizing/src/main.rs

@@ -137,13 +137,9 @@ pub fn main() -> Result<(), Error> {
     let mut executor = InProcessExecutor::new(&mut harness, (), &mut fuzzer, &mut state, &mut mgr)?;
 
     state.load_initial_inputs_forced(&mut fuzzer, &mut executor, &mut mgr, &[solution_dir])?;
-    stages.perform_all(
-        &mut fuzzer,
-        &mut executor,
-        &mut state,
-        &mut mgr,
-        CorpusId::from(0_usize),
-    )?;
+
+    state.set_corpus_idx(CorpusId::from(0usize))?;
+    stages.perform_all(&mut fuzzer, &mut executor, &mut state, &mut mgr)?;
 
     Ok(())
 }

fuzzers/fuzzbench_forkserver_cmplog/src/main.rs

@@ -9,7 +9,7 @@ use std::{
 
 use clap::{Arg, ArgAction, Command};
 use libafl::{
-    corpus::{Corpus, CorpusId, InMemoryOnDiskCorpus, OnDiskCorpus},
+    corpus::{Corpus, HasCurrentCorpusIdx, InMemoryOnDiskCorpus, OnDiskCorpus},
     events::SimpleEventManager,
     executors::forkserver::{ForkserverExecutor, TimeoutForkserverExecutor},
     feedback_or,
@@ -374,9 +374,14 @@ fn fuzz(
         let cb = |_fuzzer: &mut _,
                   _executor: &mut _,
                   state: &mut StdState<_, InMemoryOnDiskCorpus<_>, _, _>,
-                  _event_manager: &mut _,
-                  corpus_id: CorpusId|
-         -> Result<bool, libafl::Error> {
+                  _event_manager: &mut _|
+         -> Result<bool, Error> {
+            let Some(corpus_id) = state.current_corpus_idx()? else {
+                return Err(Error::illegal_state(
+                    "state is not currently processing a corpus index",
+                ));
+            };
+
             let corpus = state.corpus().get(corpus_id)?.borrow();
             let res = corpus.scheduled_count() == 1; // let's try on the 2nd trial
 

libafl/src/corpus/mod.rs

@@ -148,6 +148,18 @@ pub trait Corpus: UsesInput + Serialize + for<'de> Deserialize<'de> {
     }
 }
 
+/// Trait for types which track the current corpus index
+pub trait HasCurrentCorpusIdx {
+    /// Set the current corpus index; we have started processing this corpus entry
+    fn set_corpus_idx(&mut self, idx: CorpusId) -> Result<(), Error>;
+
+    /// Clear the current corpus index; we are done with this entry
+    fn clear_corpus_idx(&mut self) -> Result<(), Error>;
+
+    /// Fetch the current corpus index -- typically used after a state recovery or transfer
+    fn current_corpus_idx(&self) -> Result<Option<CorpusId>, Error>;
+}
+
 /// [`Iterator`] over the ids of a [`Corpus`]
 #[derive(Debug)]
 pub struct CorpusIdIterator<'a, C>

libafl/src/events/llmp.rs

@@ -962,6 +962,8 @@ where
 
     /// Reset the single page (we reuse it over and over from pos 0), then send the current state to the next runner.
     fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+        state.on_restart()?;
+
         // First, reset the page to 0 so the next iteration can read read from the beginning of this page
         self.staterestorer.reset();
         self.staterestorer.save(&(

libafl/src/events/mod.rs

@@ -128,7 +128,7 @@ pub struct EventManagerId(
 
 #[cfg(feature = "introspection")]
 use crate::monitors::ClientPerfMonitor;
-use crate::{inputs::UsesInput, state::UsesState};
+use crate::{inputs::UsesInput, stages::HasCurrentStage, state::UsesState};
 
 /// The log event severity
 #[derive(Serialize, Deserialize, Debug, Clone, Copy)]
@@ -548,8 +548,11 @@ where
 /// Restartable trait
 pub trait EventRestarter: UsesState {
     /// For restarting event managers, implement a way to forward state to their next peers.
+    /// You *must* ensure that [`HasCurrentStage::on_restart`] will be invoked in this method, by you
+    /// or an internal [`EventRestarter`], before the state is saved for recovery.
     #[inline]
-    fn on_restart(&mut self, _state: &mut Self::State) -> Result<(), Error> {
+    fn on_restart(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        state.on_restart()?;
         self.await_restart_safe();
         Ok(())
     }
@@ -605,7 +608,7 @@ pub trait HasCustomBufHandlers: UsesState {
 }
 
 /// An eventmgr for tests, and as placeholder if you really don't need an event manager.
-#[derive(Copy, Clone, Debug, Default)]
+#[derive(Copy, Clone, Debug)]
 pub struct NopEventManager<S> {
     phantom: PhantomData<S>,
 }
@@ -620,6 +623,12 @@ impl<S> NopEventManager<S> {
     }
 }
 
+impl<S> Default for NopEventManager<S> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl<S> UsesState for NopEventManager<S>
 where
     S: State,

libafl/src/events/simple.rs

@@ -353,6 +353,8 @@ where
 {
     /// Reset the single page (we reuse it over and over from pos 0), then send the current state to the next runner.
     fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+        state.on_restart()?;
+
         // First, reset the page to 0 so the next iteration can read read from the beginning of this page
         self.staterestorer.reset();
         self.staterestorer.save(&(

libafl/src/events/tcp.rs

@@ -849,13 +849,16 @@ where
 
     /// Reset the single page (we reuse it over and over from pos 0), then send the current state to the next runner.
     fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+        state.on_restart()?;
+
         // First, reset the page to 0 so the next iteration can read read from the beginning of this page
         self.staterestorer.reset();
         self.staterestorer.save(&if self.save_state {
             Some((state, self.tcp_mgr.client_id))
         } else {
             None
         })?;
+
         self.await_restart_safe();
         Ok(())
     }

libafl/src/executors/command.rs

@@ -690,10 +690,10 @@ mod tests {
             command::{CommandExecutor, InputLocation},
             Executor,
         },
+        fuzzer::test::NopFuzzer,
         inputs::BytesInput,
         monitors::SimpleMonitor,
-        state::NopState,
-        NopFuzzer,
+        state::test::NopState,
     };
 
     #[test]

libafl/src/executors/inprocess.rs

@@ -2186,9 +2186,9 @@ mod tests {
     use crate::{
         events::NopEventManager,
         executors::{inprocess::InProcessHandlers, Executor, ExitKind, InProcessExecutor},
+        fuzzer::test::NopFuzzer,
         inputs::{NopInput, UsesInput},
-        state::NopState,
-        NopFuzzer,
+        state::test::NopState,
     };
 
     impl UsesInput for () {
@@ -2226,8 +2226,8 @@ mod tests {
         use crate::{
             events::SimpleEventManager,
             executors::{inprocess::InChildProcessHandlers, InProcessForkExecutor},
-            state::NopState,
-            NopFuzzer,
+            fuzzer::test::NopFuzzer,
+            state::test::NopState,
         };
 
         let provider = StdShMemProvider::new().unwrap();

libafl/src/executors/mod.rs

@@ -1,50 +1,42 @@
 //! Executors take input, and run it in the target.
 
-pub mod inprocess;
-pub use inprocess::InProcessExecutor;
-#[cfg(all(feature = "std", feature = "fork", unix))]
-pub use inprocess::InProcessForkExecutor;
+use core::fmt::Debug;
 
-pub mod differential;
+pub use combined::CombinedExecutor;
+#[cfg(all(feature = "std", any(unix, doc)))]
+pub use command::CommandExecutor;
 pub use differential::DiffExecutor;
-
-/// Timeout executor.
-/// Not possible on `no-std` Windows or `no-std`, but works for unix
-#[cfg(any(unix, feature = "std"))]
-pub mod timeout;
-#[cfg(any(unix, feature = "std"))]
-pub use timeout::TimeoutExecutor;
-
-#[cfg(all(feature = "std", feature = "fork", unix))]
-pub mod forkserver;
 #[cfg(all(feature = "std", feature = "fork", unix))]
 pub use forkserver::{Forkserver, ForkserverExecutor, TimeoutForkserverExecutor};
-
-pub mod combined;
-pub use combined::CombinedExecutor;
-
-pub mod shadow;
+pub use inprocess::InProcessExecutor;
+#[cfg(all(feature = "std", feature = "fork", unix))]
+pub use inprocess::InProcessForkExecutor;
+use serde::{Deserialize, Serialize};
 pub use shadow::ShadowExecutor;
-
-pub mod with_observers;
+#[cfg(any(unix, feature = "std"))]
+pub use timeout::TimeoutExecutor;
 pub use with_observers::WithObservers;
 
-#[cfg(all(feature = "std", any(unix, doc)))]
-pub mod command;
-use core::{fmt::Debug, marker::PhantomData};
-
-#[cfg(all(feature = "std", any(unix, doc)))]
-pub use command::CommandExecutor;
-use libafl_bolts::AsSlice;
-use serde::{Deserialize, Serialize};
-
 use crate::{
-    inputs::HasTargetBytes,
     observers::{ObserversTuple, UsesObservers},
-    state::{HasExecutions, State, UsesState},
+    state::UsesState,
     Error,
 };
 
+pub mod combined;
+#[cfg(all(feature = "std", any(unix, doc)))]
+pub mod command;
+pub mod differential;
+#[cfg(all(feature = "std", feature = "fork", unix))]
+pub mod forkserver;
+pub mod inprocess;
+pub mod shadow;
+/// Timeout executor.
+/// Not possible on `no-std` Windows or `no-std`, but works for unix
+#[cfg(any(unix, feature = "std"))]
+pub mod timeout;
+pub mod with_observers;
+
 /// How an execution finished.
 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
 #[cfg_attr(
@@ -149,58 +141,78 @@ where
     fn post_run_reset(&mut self) {}
 }
 
-/// A simple executor that does nothing.
-/// If intput len is 0, `run_target` will return Err
-#[derive(Debug)]
-struct NopExecutor<S> {
-    phantom: PhantomData<S>,
-}
+#[cfg(test)]
+pub mod test {
+    use core::marker::PhantomData;
 
-impl<S> UsesState for NopExecutor<S>
-where
-    S: State,
-{
-    type State = S;
-}
+    use libafl_bolts::{AsSlice, Error};
 
-impl<EM, S, Z> Executor<EM, Z> for NopExecutor<S>
-where
-    EM: UsesState<State = S>,
-    S: State + HasExecutions,
-    S::Input: HasTargetBytes,
-    Z: UsesState<State = S>,
-{
-    fn run_target(
-        &mut self,
-        _fuzzer: &mut Z,
-        state: &mut Self::State,
-        _mgr: &mut EM,
-        input: &Self::Input,
-    ) -> Result<ExitKind, Error> {
-        *state.executions_mut() += 1;
+    use crate::{
+        events::NopEventManager,
+        executors::{Executor, ExitKind},
+        fuzzer::test::NopFuzzer,
+        inputs::{BytesInput, HasTargetBytes},
+        state::{test::NopState, HasExecutions, State, UsesState},
+    };
+
+    /// A simple executor that does nothing.
+    /// If intput len is 0, `run_target` will return Err
+    #[derive(Debug)]
+    pub struct NopExecutor<S> {
+        phantom: PhantomData<S>,
+    }
+
+    impl<S> NopExecutor<S> {
+        #[must_use]
+        pub fn new() -> Self {
+            Self {
+                phantom: PhantomData,
+            }
+        }
+    }
 
-        if input.target_bytes().as_slice().is_empty() {
-            Err(Error::empty("Input Empty"))
-        } else {
-            Ok(ExitKind::Ok)
+    impl<S> Default for NopExecutor<S> {
+        fn default() -> Self {
+            Self::new()
         }
     }
-}
 
-#[cfg(test)]
-mod test {
-    use core::marker::PhantomData;
+    impl<S> UsesState for NopExecutor<S>
+    where
+        S: State,
+    {
+        type State = S;
+    }
 
-    use super::{Executor, NopExecutor};
-    use crate::{events::NopEventManager, inputs::BytesInput, state::NopState, NopFuzzer};
+    impl<EM, S, Z> Executor<EM, Z> for NopExecutor<S>
+    where
+        EM: UsesState<State = S>,
+        S: State + HasExecutions,
+        S::Input: HasTargetBytes,
+        Z: UsesState<State = S>,
+    {
+        fn run_target(
+            &mut self,
+            _fuzzer: &mut Z,
+            state: &mut Self::State,
+            _mgr: &mut EM,
+            input: &Self::Input,
+        ) -> Result<ExitKind, Error> {
+            *state.executions_mut() += 1;
+
+            if input.target_bytes().as_slice().is_empty() {
+                Err(Error::empty("Input Empty"))
+            } else {
+                Ok(ExitKind::Ok)
+            }
+        }
+    }
 
     #[test]
     fn nop_executor() {
         let empty_input = BytesInput::new(vec![]);
         let nonempty_input = BytesInput::new(vec![1u8]);
-        let mut executor = NopExecutor {
-            phantom: PhantomData,
-        };
+        let mut executor = NopExecutor::new();
         let mut fuzzer = NopFuzzer::new();
 
         let mut state = NopState::new();

libafl/src/feedbacks/differential.rs

@@ -163,7 +163,7 @@ mod tests {
         feedbacks::{differential::DiffResult, DiffFeedback, Feedback},
         inputs::{BytesInput, UsesInput},
         observers::Observer,
-        state::{NopState, State, UsesState},
+        state::{test::NopState, State, UsesState},
     };
 
     #[derive(Debug)]