Refactor stages (#3002)

* refactor

* miss

* lol

* revivet tests

* aa

* ?

* a

* fuck

* fuck

* a

* m

* fuck

.github/workflows/build_and_test.yml

@@ -237,7 +237,6 @@ jobs:
     needs:
       - fuzzers-preflight
     strategy:
-      fail-fast: false
       matrix:
         os: [ ubuntu-24.04 ]
         fuzzer:

fuzzers/baby/tutorial/src/lib.rs

@@ -186,7 +186,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
 
     // It's important, that we store the state before restarting!
     // Else, the parent will not respawn a new child and quit.
-    restarting_mgr.on_restart(&mut state)?;
+    restarting_mgr.mgr_on_restart(&mut state)?;
 
     Ok(())
 }

fuzzers/binary_only/qemu_launcher/src/instance.rs

@@ -469,7 +469,7 @@ impl<M: Monitor> Instance<'_, M> {
 
                 // It's important, that we store the state before restarting!
                 // Else, the parent will not respawn a new child and quit.
-                self.mgr.on_restart(state)?;
+                self.mgr.mgr_on_restart(state)?;
             }
         } else {
             fuzzer.fuzz_loop(stages, executor, state, &mut self.mgr)?;

fuzzers/forkserver/libafl-fuzz/Justfile

@@ -53,7 +53,7 @@ build_libafl_fuzz_fuzzbench:
 
 test_instr: build_afl build_libafl_fuzz
     #!/bin/bash
-    AFL_PATH={{AFL_DIR}} {{AFL_CC_PATH}} ./test/test-instr.c -o ./test/out-instr
+    AFL_PATH={{AFL_DIR}} {{AFL_CC_PATH}} -O0 ./test/test-instr.c -o ./test/out-instr
 
     export LIBAFL_DEBUG_OUTPUT=1 
     export AFL_CORES=0

fuzzers/forkserver/libafl-fuzz/src/main.rs

@@ -256,14 +256,15 @@ struct Opt {
     #[arg(short = 't', default_value_t = 1000)]
     hang_timeout: u64,
 
+    #[arg(short = 'd')]
+    debug_child: bool,
+
     // Environment Variables
     #[clap(skip)]
     bench_just_one: bool,
     #[clap(skip)]
     bench_until_crash: bool,
 
-    #[clap(skip)]
-    debug_child: bool,
     #[clap(skip)]
     is_persistent: bool,
     #[clap(skip)]

fuzzers/full_system/nyx_launcher/src/instance.rs

@@ -250,7 +250,7 @@ impl<M: Monitor> Instance<'_, M> {
 
             // It's important, that we store the state before restarting!
             // Else, the parent will not respawn a new child and quit.
-            self.mgr.on_restart(state)?;
+            self.mgr.mgr_on_restart(state)?;
         } else {
             fuzzer.fuzz_loop(stages, executor, state, &mut self.mgr)?;
         }

fuzzers/inprocess/libfuzzer_libpng/src/lib.rs

@@ -222,7 +222,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
 
     // It's important, that we store the state before restarting!
     // Else, the parent will not respawn a new child and quit.
-    restarting_mgr.on_restart(&mut state)?;
+    restarting_mgr.mgr_on_restart(&mut state)?;
 
     Ok(())
 }

fuzzers/inprocess/libfuzzer_libpng_cmin/src/lib.rs

@@ -236,7 +236,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
 
     // It's important, that we store the state before restarting!
     // Else, the parent will not respawn a new child and quit.
-    restarting_mgr.on_restart(&mut state)?;
+    restarting_mgr.mgr_on_restart(&mut state)?;
 
     Ok(())
 }

fuzzers/inprocess/libfuzzer_libpng_norestart/src/lib.rs

@@ -274,7 +274,7 @@ pub extern "C" fn libafl_main() {
             &mut restarting_mgr,
             opt.loop_iters,
         )?;
-        restarting_mgr.on_restart(&mut state)?;
+        restarting_mgr.mgr_on_restart(&mut state)?;
 
         Ok(())
     };

fuzzers/inprocess/libfuzzer_libpng_tcp_manager/src/lib.rs

@@ -219,7 +219,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
 
     // It's important, that we store the state before restarting!
     // Else, the parent will not respawn a new child and quit.
-    restarting_mgr.on_restart(&mut state)?;
+    restarting_mgr.mgr_on_restart(&mut state)?;
 
     Ok(())
 }

fuzzers/inprocess/libfuzzer_windows_asan/src/lib.rs

@@ -184,7 +184,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
 
     // It's important, that we store the state before restarting!
     // Else, the parent will not respawn a new child and quit.
-    restarting_mgr.on_restart(&mut state)?;
+    restarting_mgr.mgr_on_restart(&mut state)?;
 
     Ok(())
 }

libafl/src/events/centralized.rs

@@ -272,9 +272,9 @@ where
     SP: ShMemProvider<ShMem = SHM>,
 {
     #[inline]
-    fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
         self.client.await_safe_to_unmap_blocking();
-        self.inner.on_restart(state)?;
+        self.inner.mgr_on_restart(state)?;
         Ok(())
     }
 }

libafl/src/events/llmp/restarting.rs

@@ -63,10 +63,9 @@ use crate::{
     inputs::Input,
     monitors::Monitor,
     observers::TimeObserver,
-    stages::HasCurrentStageId,
     state::{
-        HasCurrentTestcase, HasExecutions, HasImported, HasLastReportTime, HasSolutions,
+        HasCurrentStageId, HasCurrentTestcase, HasExecutions, HasImported, HasLastReportTime,
-        MaybeHasClientPerfMonitor, Stoppable,
+        HasSolutions, MaybeHasClientPerfMonitor, Stoppable,
     },
     Error,
 };
@@ -257,7 +256,7 @@ where
     SP: ShMemProvider<ShMem = SHM>,
 {
     /// 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> {
+    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
         state.on_restart()?;
 
         if let Some(sr) = &mut self.staterestorer {

libafl/src/events/mod.rs

@@ -108,7 +108,7 @@ pub struct EventManagerId(
 use crate::events::multi_machine::NodeId;
 #[cfg(feature = "introspection")]
 use crate::monitors::stats::ClientPerfStats;
-use crate::{observers::TimeObserver, stages::HasCurrentStageId};
+use crate::{observers::TimeObserver, state::HasCurrentStageId};
 
 /// The log event severity
 #[derive(Serialize, Deserialize, Debug, Clone, Copy)]
@@ -559,10 +559,10 @@ pub trait EventRestarter<S> {
     /// You *must* ensure that [`HasCurrentStageId::on_restart`] will be invoked in this method, by you
     /// or an internal [`EventRestarter`], before the state is saved for recovery.
     /// [`std_on_restart`] is the standard implementation that you can call.
-    fn on_restart(&mut self, state: &mut S) -> Result<(), Error>;
+    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error>;
 }
 
-/// Default implementation of [`EventRestarter::on_restart`] for implementors with the given
+/// Default implementation of [`EventRestarter::mgr_on_restart`] for implementors with the given
 /// constraints
 pub fn std_on_restart<EM, S>(restarter: &mut EM, state: &mut S) -> Result<(), Error>
 where
@@ -644,7 +644,7 @@ impl<S> EventRestarter<S> for NopEventManager
 where
     S: HasCurrentStageId,
 {
-    fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
         std_on_restart(self, state)
     }
 }
@@ -769,8 +769,8 @@ where
     EM: EventRestarter<S>,
 {
     #[inline]
-    fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
-        self.inner.on_restart(state)
+        self.inner.mgr_on_restart(state)
     }
 }
 

libafl/src/events/simple.rs

@@ -32,8 +32,9 @@ use crate::{
         SendExiting,
     },
     monitors::{stats::ClientStatsManager, Monitor},
-    stages::HasCurrentStageId,
+    state::{
-    state::{HasExecutions, HasLastReportTime, MaybeHasClientPerfMonitor, Stoppable},
+        HasCurrentStageId, HasExecutions, HasLastReportTime, MaybeHasClientPerfMonitor, Stoppable,
+    },
     Error, HasMetadata,
 };
 #[cfg(feature = "std")]
@@ -111,7 +112,7 @@ impl<I, MT, S> EventRestarter<S> for SimpleEventManager<I, MT, S>
 where
     S: HasCurrentStageId,
 {
-    fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
         std_on_restart(self, state)
     }
 }
@@ -325,7 +326,7 @@ where
     MT: Monitor,
 {
     /// 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> {
+    fn mgr_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

libafl/src/events/tcp.rs

@@ -48,10 +48,9 @@ use crate::{
     },
     inputs::Input,
     monitors::{stats::ClientStatsManager, Monitor},
-    stages::HasCurrentStageId,
     state::{
-        HasCurrentTestcase, HasExecutions, HasImported, HasLastReportTime, HasSolutions,
+        HasCurrentStageId, HasCurrentTestcase, HasExecutions, HasImported, HasLastReportTime,
-        MaybeHasClientPerfMonitor, Stoppable,
+        HasSolutions, MaybeHasClientPerfMonitor, Stoppable,
     },
     Error, HasMetadata,
 };
@@ -627,7 +626,7 @@ impl<EMH, I, S> EventRestarter<S> for TcpEventManager<EMH, I, S>
 where
     S: HasCurrentStageId,
 {
-    fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
         std_on_restart(self, state)
     }
 }
@@ -861,7 +860,7 @@ where
     SP: ShMemProvider<ShMem = SHM>,
 {
     /// 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> {
+    fn mgr_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

libafl/src/executors/inprocess/mod.rs

@@ -430,7 +430,7 @@ pub fn run_observers_and_save_state<E, EM, I, OF, S, Z>(
     }
 
     // Serialize the state and wait safely for the broker to read pending messages
-    event_mgr.on_restart(state).unwrap();
+    event_mgr.mgr_on_restart(state).unwrap();
 
     log::info!("Bye!");
 }

libafl/src/executors/mod.rs

@@ -157,7 +157,8 @@ pub fn common_signals() -> Vec<Signal> {
 }
 
 #[cfg(test)]
-mod test {
+/// Tester for executor
+pub mod test {
     use core::marker::PhantomData;
 
     use libafl_bolts::{AsSlice, Error};

libafl/src/fuzzer/mod.rs

@@ -24,11 +24,11 @@ use crate::{
     mark_feature_time,
     observers::ObserversTuple,
     schedulers::Scheduler,
-    stages::{HasCurrentStageId, StagesTuple},
+    stages::StagesTuple,
     start_timer,
     state::{
-        HasCorpus, HasCurrentTestcase, HasExecutions, HasImported, HasLastFoundTime,
+        HasCorpus, HasCurrentStageId, HasCurrentTestcase, HasExecutions, HasImported,
-        HasLastReportTime, HasSolutions, MaybeHasClientPerfMonitor, Stoppable,
+        HasLastFoundTime, HasLastReportTime, HasSolutions, MaybeHasClientPerfMonitor, Stoppable,
     },
     Error, HasMetadata,
 };
@@ -207,7 +207,7 @@ pub trait Fuzzer<E, EM, I, S, ST> {
     /// because each stage could run the harness for multiple times)
     ///
     /// If you use this fn in a restarting scenario to only run for `n` iterations,
-    /// before exiting, make sure you call `event_mgr.on_restart(&mut state)?;`.
+    /// before exiting, make sure you call `event_mgr.mgr_on_restart(&mut state)?;`.
     /// This way, the state will be available in the next, respawned, iteration.
     fn fuzz_one(
         &mut self,
@@ -233,7 +233,7 @@ pub trait Fuzzer<E, EM, I, S, ST> {
     /// because each stage could run the harness for multiple times)
     ///
     /// If you use this fn in a restarting scenario to only run for `n` iterations,
-    /// before exiting, make sure you call `event_mgr.on_restart(&mut state)?;`.
+    /// before exiting, make sure you call `event_mgr.mgr_on_restart(&mut state)?;`.
     /// This way, the state will be available in the next, respawned, iteration.
     fn fuzz_loop_for(
         &mut self,

libafl/src/stages/logics.rs

@@ -3,7 +3,8 @@
 use core::marker::PhantomData;
 
 use crate::{
-    stages::{HasNestedStageStatus, Stage, StageId, StagesTuple},
+    stages::{Stage, StageId, StagesTuple},
+    state::HasNestedStage,
     Error,
 };
 
@@ -14,7 +15,7 @@ pub struct NestedStageRetryCountRestartHelper;
 impl NestedStageRetryCountRestartHelper {
     fn should_restart<S, ST>(state: &mut S, _stage: &ST) -> Result<bool, Error>
     where
-        S: HasNestedStageStatus,
+        S: HasNestedStage,
     {
         state.enter_inner_stage()?;
         Ok(true)
@@ -22,7 +23,7 @@ impl NestedStageRetryCountRestartHelper {
 
     fn clear_progress<S, ST>(state: &mut S, _stage: &ST) -> Result<(), Error>
     where
-        S: HasNestedStageStatus,
+        S: HasNestedStage,
     {
         state.exit_inner_stage()?;
         Ok(())
@@ -41,7 +42,7 @@ impl<CB, E, EM, ST, S, Z> Stage<E, EM, S, Z> for WhileStage<CB, E, EM, ST, S, Z>
 where
     CB: FnMut(&mut Z, &mut E, &mut S, &mut EM) -> Result<bool, Error>,
     ST: StagesTuple<E, EM, S, Z>,
-    S: HasNestedStageStatus,
+    S: HasNestedStage,
 {
     fn perform(
         &mut self,
@@ -95,7 +96,7 @@ impl<CB, E, EM, ST, S, Z> Stage<E, EM, S, Z> for IfStage<CB, E, EM, ST, S, Z>
 where
     CB: FnMut(&mut Z, &mut E, &mut S, &mut EM) -> Result<bool, Error>,
     ST: StagesTuple<E, EM, S, Z>,
-    S: HasNestedStageStatus,
+    S: HasNestedStage,
 {
     fn perform(
         &mut self,
@@ -150,7 +151,7 @@ where
     CB: FnMut(&mut Z, &mut E, &mut S, &mut EM) -> Result<bool, Error>,
     ST1: StagesTuple<E, EM, S, Z>,
     ST2: StagesTuple<E, EM, S, Z>,
-    S: HasNestedStageStatus,
+    S: HasNestedStage,
 {
     fn perform(
         &mut self,
@@ -161,10 +162,12 @@ where
     ) -> Result<(), Error> {
         let current = state.current_stage_id()?;
 
+        // this is None if you didn't recover from restart
+        // because should_restart() which is called right before this will create a new stage stack
         let fresh = current.is_none();
-        let closure_return = fresh && (self.closure)(fuzzer, executor, state, manager)?;
+        let closure_res = fresh && (self.closure)(fuzzer, executor, state, manager)?;
 
-        if current == Some(StageId(0)) || closure_return {
+        if current == Some(StageId(0)) || closure_res {
             if fresh {
                 state.set_current_stage_id(StageId(0))?;
             }
@@ -220,7 +223,7 @@ pub struct OptionalStage<E, EM, ST, S, Z> {
 impl<E, EM, ST, S, Z> Stage<E, EM, S, Z> for OptionalStage<E, EM, ST, S, Z>
 where
     ST: StagesTuple<E, EM, S, Z>,
-    S: HasNestedStageStatus,
+    S: HasNestedStage,
 {
     fn perform(
         &mut self,
@@ -273,3 +276,344 @@ impl<E, EM, ST, S, Z> OptionalStage<E, EM, ST, S, Z> {
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+    use alloc::rc::Rc;
+    use core::{cell::RefCell, marker::PhantomData};
+
+    use libafl_bolts::{
+        impl_serdeany,
+        tuples::{tuple_list, tuple_list_type},
+        Error,
+    };
+    use serde::{Deserialize, Serialize};
+
+    #[cfg(any(not(feature = "serdeany_autoreg"), miri))]
+    use crate::stages::RetryCountRestartHelper;
+    use crate::{
+        events::NopEventManager,
+        executors::test::NopExecutor,
+        stages::{
+            ClosureStage, CorpusId, HasCurrentCorpusId, IfElseStage, IfStage, Stage, StagesTuple,
+            WhileStage,
+        },
+        state::{HasCurrentStageId, StdState},
+        HasMetadata, NopFuzzer,
+    };
+
+    #[derive(Debug)]
+    pub struct ResumeSucceededStage<S> {
+        phantom: PhantomData<S>,
+    }
+    #[derive(Debug)]
+    pub struct ResumeFailedStage<S> {
+        completed: Rc<RefCell<bool>>,
+        phantom: PhantomData<S>,
+    }
+    #[derive(Serialize, Deserialize, Debug)]
+    pub struct TestProgress {
+        count: usize,
+    }
+
+    impl_serdeany!(TestProgress);
+
+    impl TestProgress {
+        #[expect(clippy::unnecessary_wraps)]
+        fn should_restart<S, ST>(state: &mut S, _stage: &ST) -> Result<bool, Error>
+        where
+            S: HasMetadata,
+        {
+            // check if we're resuming
+            let _metadata = state.metadata_or_insert_with(|| Self { count: 0 });
+            Ok(true)
+        }
+
+        fn clear_progress<S, ST>(state: &mut S, _stage: &ST) -> Result<(), Error>
+        where
+            S: HasMetadata,
+        {
+            if state.remove_metadata::<Self>().is_none() {
+                return Err(Error::illegal_state(
+                    "attempted to clear status metadata when none was present",
+                ));
+            }
+            Ok(())
+        }
+    }
+
+    impl<E, EM, S, Z> Stage<E, EM, S, Z> for ResumeSucceededStage<S>
+    where
+        S: HasMetadata,
+    {
+        fn perform(
+            &mut self,
+            _fuzzer: &mut Z,
+            _executor: &mut E,
+            state: &mut S,
+            _manager: &mut EM,
+        ) -> Result<(), Error> {
+            // metadata is attached by the status
+            let meta = state.metadata_mut::<TestProgress>().unwrap();
+            meta.count += 1;
+            assert!(
+                meta.count == 1,
+                "Test failed; we resumed a succeeded stage!"
+            );
+            Ok(())
+        }
+
+        fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
+            TestProgress::should_restart(state, self)
+        }
+
+        fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
+            TestProgress::clear_progress(state, self)
+        }
+    }
+
+    impl<E, EM, S, Z> Stage<E, EM, S, Z> for ResumeFailedStage<S>
+    where
+        S: HasMetadata,
+    {
+        fn perform(
+            &mut self,
+            _fuzzer: &mut Z,
+            _executor: &mut E,
+            state: &mut S,
+            _manager: &mut EM,
+        ) -> Result<(), Error> {
+            // metadata is attached by the status
+            let meta = state.metadata_mut::<TestProgress>().unwrap();
+            meta.count += 1;
+            if meta.count == 1 {
+                return Err(Error::shutting_down());
+            } else if meta.count > 2 {
+                panic!("Resume was somehow corrupted?")
+            } else {
+                self.completed.replace(true);
+            }
+            Ok(())
+        }
+
+        fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
+            TestProgress::should_restart(state, self)
+        }
+
+        fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
+            TestProgress::clear_progress(state, self)
+        }
+    }
+
+    #[must_use]
+    #[allow(clippy::type_complexity)]
+    pub fn test_resume_stages<S>() -> (
+        Rc<RefCell<bool>>,
+        tuple_list_type!(ResumeSucceededStage<S>, ResumeFailedStage<S>),
+    ) {
+        let completed = Rc::new(RefCell::new(false));
+        (
+            completed.clone(),
+            tuple_list!(
+                ResumeSucceededStage {
+                    phantom: PhantomData
+                },
+                ResumeFailedStage {
+                    completed,
+                    phantom: PhantomData
+                },
+            ),
+        )
+    }
+
+    pub fn test_resume<ST, S>(completed: &Rc<RefCell<bool>>, state: &mut S, mut stages: ST)
+    where
+        ST: StagesTuple<NopExecutor<S>, NopEventManager, S, NopFuzzer>,
+        S: HasCurrentStageId + HasCurrentCorpusId,
+    {
+        #[cfg(any(not(feature = "serdeany_autoreg"), miri))]
+        unsafe {
+            TestProgress::register();
+            RetryCountRestartHelper::register();
+        }
+
+        let mut fuzzer = NopFuzzer::new();
+        let mut executor = NopExecutor::new();
+        let mut manager = NopEventManager::new();
+        for _ in 0..2 {
+            completed.replace(false);
+            // fake one, just any number so retryhelper won't fail.
+            // in reality you always have corpus id set by stdfuzzer
+            state.set_corpus_id(CorpusId::from(0_usize)).unwrap();
+            let Err(e) = stages.perform_all(&mut fuzzer, &mut executor, state, &mut manager) else {
+                panic!("Test failed; stages should fail the first time.")
+            };
+            assert!(
+                matches!(e, Error::ShuttingDown),
+                "Unexpected error encountered."
+            );
+            assert!(!*completed.borrow(), "Unexpectedly complete?");
+            state
+                .on_restart()
+                .expect("Couldn't notify state of restart.");
+            assert!(
+                stages
+                    .perform_all(&mut fuzzer, &mut executor, state, &mut manager)
+                    .is_ok(),
+                "Test failed; stages should pass the second time."
+            );
+            assert!(
+                *completed.borrow(),
+                "Test failed; we did not set completed."
+            );
+        }
+    }
+
+    #[test]
+    fn check_resumability_while() {
+        let once = RefCell::new(true);
+
+        let (completed, stages) = test_resume_stages();
+        let whilestage = WhileStage::new(
+            |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(once.replace(false)),
+            stages,
+        );
+        let resetstage = ClosureStage::new(|_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| {
+            once.replace(true);
+            Ok(())
+        });
+        let mut state = StdState::nop().unwrap();
+        test_resume(&completed, &mut state, tuple_list!(whilestage, resetstage));
+    }
+
+    #[test]
+    fn check_resumability_if() {
+        let once = RefCell::new(true);
+        let (completed, stages) = test_resume_stages();
+        let ifstage = IfStage::new(
+            |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(once.replace(false)),
+            stages,
+        );
+        let resetstage = ClosureStage::new(|_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| {
+            once.replace(true);
+            Ok(())
+        });
+        let mut state = StdState::nop().unwrap();
+        test_resume(&completed, &mut state, tuple_list!(ifstage, resetstage));
+    }
+
+    #[test]
+    fn check_resumability_if_deep() {
+        let (completed, stages) = test_resume_stages();
+        let ifstage = IfStage::new(
+            |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(true),
+            tuple_list!(IfStage::new(
+                |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(true),
+                tuple_list!(IfStage::new(
+                    |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(true),
+                    tuple_list!(IfStage::new(
+                        |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(true),
+                        tuple_list!(IfStage::new(
+                            |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(true),
+                            stages
+                        ),),
+                    ),),
+                ))
+            )),
+        );
+        let mut state = StdState::nop().unwrap();
+        test_resume(&completed, &mut state, tuple_list!(ifstage));
+    }
+
+    #[derive(Debug)]
+    pub struct PanicStage<S> {
+        phantom: PhantomData<S>,
+    }
+    impl<S> PanicStage<S> {
+        pub fn new() -> Self {
+            Self {
+                phantom: PhantomData,
+            }
+        }
+    }
+    impl<E, EM, S, Z> Stage<E, EM, S, Z> for PanicStage<S> {
+        fn perform(
+            &mut self,
+            _fuzzer: &mut Z,
+            _executor: &mut E,
+            _state: &mut S,
+            _manager: &mut EM,
+        ) -> Result<(), Error> {
+            panic!("Test failed; panic stage should never be executed.");
+        }
+
+        fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
+            Ok(true)
+        }
+
+        fn clear_progress(&mut self, _state: &mut S) -> Result<(), Error> {
+            Ok(())
+        }
+    }
+
+    #[test]
+    fn check_resumability_if_else_if() {
+        let once = RefCell::new(true);
+        let (completed, stages) = test_resume_stages();
+        let ifstage = IfElseStage::new(
+            |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(once.replace(false)),
+            stages,
+            tuple_list!(PanicStage::new()),
+        );
+        let resetstage = ClosureStage::new(|_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| {
+            once.replace(true);
+            Ok(())
+        });
+        let mut state = StdState::nop().unwrap();
+        test_resume(&completed, &mut state, tuple_list!(ifstage, resetstage));
+    }
+
+    #[test]
+    fn check_resumability_if_else_else() {
+        let once = RefCell::new(false);
+        let (completed, stages) = test_resume_stages();
+        let ifstage = IfElseStage::new(
+            |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(once.replace(true)),
+            tuple_list!(PanicStage::new()),
+            stages,
+        );
+        let resetstage = ClosureStage::new(|_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| {
+            once.replace(false);
+            Ok(())
+        });
+        let mut state = StdState::nop().unwrap();
+        test_resume(&completed, &mut state, tuple_list!(ifstage, resetstage));
+    }
+    #[test]
+    fn check_resumability_if_else_else_deep() {
+        let (completed, stages) = test_resume_stages();
+        let ifstage = IfElseStage::new(
+            |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(false),
+            tuple_list!(PanicStage::new()),
+            tuple_list!(IfElseStage::new(
+                |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(false),
+                tuple_list!(PanicStage::new()),
+                tuple_list!(IfElseStage::new(
+                    |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(false),
+                    tuple_list!(PanicStage::new()),
+                    tuple_list!(IfElseStage::new(
+                        |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(false),
+                        tuple_list!(PanicStage::new()),
+                        tuple_list!(IfElseStage::new(
+                            |_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| Ok(false),
+                            tuple_list!(PanicStage::new()),
+                            stages,
+                        )),
+                    )),
+                )),
+            )),
+        );
+        let mut state = StdState::nop().unwrap();
+        test_resume(&completed, &mut state, tuple_list!(ifstage));
+    }
+}

libafl/src/stages/mod.rs

@@ -49,7 +49,7 @@ pub use verify_timeouts::{TimeoutsToVerify, VerifyTimeoutsStage};
 use crate::{
     corpus::{CorpusId, HasCurrentCorpusId},
     events::SendExiting,
-    state::{HasExecutions, Stoppable},
+    state::{HasCurrentStageId, HasExecutions, Stoppable},
     Error, HasNamedMetadata,
 };
 
@@ -98,8 +98,6 @@ pub trait Stage<E, EM, S, Z> {
     ///
     /// Before a call to perform, [`Stage::should_restart`] will be (must be!) called.
     /// After returning (so non-target crash or timeout in a restarting case), [`Stage::clear_progress`] gets called.
-    /// A call to [`Stage::perform_restartable`] will do these things implicitly.
-    /// DON'T call this function directly except from `preform_restartable` !!
     fn perform(
         &mut self,
         fuzzer: &mut Z,
@@ -107,20 +105,6 @@ pub trait Stage<E, EM, S, Z> {
         state: &mut S,
         manager: &mut EM,
     ) -> Result<(), Error>;
-
-    /// Run the stage, calling [`Stage::should_restart`] and [`Stage::clear_progress`] appropriately
-    fn perform_restartable(
-        &mut self,
-        fuzzer: &mut Z,
-        executor: &mut E,
-        state: &mut S,
-        manager: &mut EM,
-    ) -> Result<(), Error> {
-        if self.should_restart(state)? {
-            self.perform(fuzzer, executor, state, manager)?;
-        }
-        self.clear_progress(state)
-    }
 }
 
 /// A tuple holding all `Stages` used for fuzzing.
@@ -182,7 +166,10 @@ where
 
                 let stage = &mut self.0;
 
-                stage.perform_restartable(fuzzer, executor, state, manager)?;
+                if stage.should_restart(state)? {
+                    stage.perform(fuzzer, executor, state, manager)?;
+                }
+                stage.clear_progress(state)?;
 
                 state.clear_stage_id()?;
             }
@@ -194,7 +181,11 @@ where
                 state.set_current_stage_id(StageId(Self::LEN))?;
 
                 let stage = &mut self.0;
-                stage.perform_restartable(fuzzer, executor, state, manager)?;
+
+                if stage.should_restart(state)? {
+                    stage.perform(fuzzer, executor, state, manager)?;
+                }
+                stage.clear_progress(state)?;
 
                 state.clear_stage_id()?;
             }
@@ -261,13 +252,16 @@ where
         state: &mut S,
         manager: &mut EM,
     ) -> Result<(), Error> {
-        self.iter_mut().try_for_each(|x| {
+        self.iter_mut().try_for_each(|stage| {
             if state.stop_requested() {
                 state.discard_stop_request();
                 manager.on_shutdown()?;
                 return Err(Error::shutting_down());
             }
-            x.perform_restartable(fuzzer, executor, state, manager)
+            if stage.should_restart(state)? {
+                stage.perform(fuzzer, executor, state, manager)?;
+            }
+            stage.clear_progress(state)
         })
     }
 }
@@ -418,34 +412,6 @@ impl fmt::Display for StageId {
     }
 }
 
-/// Trait for types which track the current stage
-pub trait HasCurrentStageId {
-    /// Set the current stage; we have started processing this stage
-    fn set_current_stage_id(&mut self, id: StageId) -> Result<(), Error>;
-
-    /// Clear the current stage; we are done processing this stage
-    fn clear_stage_id(&mut self) -> Result<(), Error>;
-
-    /// Fetch the current stage -- typically used after a state recovery or transfer
-    fn current_stage_id(&self) -> Result<Option<StageId>, Error>;
-
-    /// Notify of a reset from which we may recover
-    fn on_restart(&mut self) -> Result<(), Error> {
-        Ok(())
-    }
-}
-
-/// Trait for types which track nested stages. Stages which themselves contain stage tuples should
-/// ensure that they constrain the state with this trait accordingly.
-pub trait HasNestedStageStatus: HasCurrentStageId {
-    /// Enter a stage scope, potentially resuming to an inner stage status. Returns Ok(true) if
-    /// resumed.
-    fn enter_inner_stage(&mut self) -> Result<(), Error>;
-
-    /// Exit a stage scope
-    fn exit_inner_stage(&mut self) -> Result<(), Error>;
-}
-
 impl_serdeany!(ExecutionCountRestartHelperMetadata);
 
 /// `SerdeAny` metadata used to keep track of executions since start for a given stage.
@@ -528,87 +494,16 @@ impl ExecutionCountRestartHelper {
 #[cfg(test)]
 mod test {
     use alloc::borrow::Cow;
-    use core::marker::PhantomData;
 
-    use libafl_bolts::{impl_serdeany, Error, Named};
+    use libafl_bolts::{Error, Named};
-    use serde::{Deserialize, Serialize};
 
     use crate::{
         corpus::{Corpus, HasCurrentCorpusId, Testcase},
         inputs::NopInput,
-        stages::{RetryCountRestartHelper, Stage},
+        stages::RetryCountRestartHelper,
         state::{HasCorpus, StdState},
-        HasMetadata,
     };
 
-    /// A stage that succeeds to resume
-    #[derive(Debug)]
-    pub struct ResumeSucceededStage<S> {
-        phantom: PhantomData<S>,
-    }
-
-    /// A progress state for testing
-    #[derive(Serialize, Deserialize, Debug)]
-    pub struct TestProgress {
-        count: usize,
-    }
-
-    impl_serdeany!(TestProgress);
-
-    impl TestProgress {
-        #[expect(clippy::unnecessary_wraps)]
-        fn should_restart<S, ST>(state: &mut S, _stage: &ST) -> Result<bool, Error>
-        where
-            S: HasMetadata,
-        {
-            // check if we're resuming
-            let metadata = state.metadata_or_insert_with(|| Self { count: 0 });
-
-            metadata.count += 1;
-            assert!(
-                metadata.count == 1,
-                "Test failed; we resumed a succeeded stage!"
-            );
-
-            Ok(true)
-        }
-
-        fn clear_progress<S, ST>(state: &mut S, _stage: &ST) -> Result<(), Error>
-        where
-            S: HasMetadata,
-        {
-            if state.remove_metadata::<Self>().is_none() {
-                return Err(Error::illegal_state(
-                    "attempted to clear status metadata when none was present",
-                ));
-            }
-            Ok(())
-        }
-    }
-
-    impl<E, EM, S, Z> Stage<E, EM, S, Z> for ResumeSucceededStage<S>
-    where
-        S: HasMetadata,
-    {
-        fn perform(
-            &mut self,
-            _fuzzer: &mut Z,
-            _executor: &mut E,
-            _state: &mut S,
-            _manager: &mut EM,
-        ) -> Result<(), Error> {
-            Ok(())
-        }
-
-        fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
-            TestProgress::should_restart(state, self)
-        }
-
-        fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
-            TestProgress::clear_progress(state, self)
-        }
-    }
-
     /// Test to test retries in stages
     #[test]
     fn test_tries_progress() -> Result<(), Error> {

libafl/src/stages/time_tracker.rs

@@ -40,7 +40,12 @@ where
         self.inner.perform(fuzzer, executor, state, manager)?;
         let after_run = current_time();
         self.count += after_run - before_run;
-        *state.metadata_mut::<T>()? = T::from(self.count);
+
+        if let Ok(meta) = state.metadata_mut::<T>() {
+            *meta = T::from(self.count);
+        } else {
+            state.add_metadata::<T>(T::from(self.count));
+        }
         Ok(())
     }
 
@@ -51,15 +56,4 @@ where
     fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
         self.inner.clear_progress(state)
     }
-
-    fn perform_restartable(
-        &mut self,
-        fuzzer: &mut Z,
-        executor: &mut E,
-        state: &mut S,
-        manager: &mut M,
-    ) -> Result<(), Error> {
-        self.inner
-            .perform_restartable(fuzzer, executor, state, manager)
-    }
 }

libafl/src/state/mod.rs

@@ -35,7 +35,7 @@ use crate::{
     fuzzer::{Evaluator, ExecuteInputResult},
     generators::Generator,
     inputs::{Input, NopInput},
-    stages::{HasCurrentStageId, HasNestedStageStatus, StageId},
+    stages::StageId,
     Error, HasMetadata, HasNamedMetadata,
 };
 
@@ -547,7 +547,35 @@ impl<C, I, R, SC> HasCurrentStageId for StdState<C, I, R, SC> {
     }
 }
 
-impl<C, I, R, SC> HasNestedStageStatus for StdState<C, I, R, SC> {
+/// Trait for types which track the current stage
+pub trait HasCurrentStageId {
+    /// Set the current stage; we have started processing this stage
+    fn set_current_stage_id(&mut self, id: StageId) -> Result<(), Error>;
+
+    /// Clear the current stage; we are done processing this stage
+    fn clear_stage_id(&mut self) -> Result<(), Error>;
+
+    /// Fetch the current stage -- typically used after a state recovery or transfer
+    fn current_stage_id(&self) -> Result<Option<StageId>, Error>;
+
+    /// Notify of a reset from which we may recover
+    fn on_restart(&mut self) -> Result<(), Error> {
+        Ok(())
+    }
+}
+
+/// Trait for types which track nested stages. Stages which themselves contain stage tuples should
+/// ensure that they constrain the state with this trait accordingly.
+pub trait HasNestedStage: HasCurrentStageId {
+    /// Enter a stage scope, potentially resuming to an inner stage status. Returns Ok(true) if
+    /// resumed.
+    fn enter_inner_stage(&mut self) -> Result<(), Error>;
+
+    /// Exit a stage scope
+    fn exit_inner_stage(&mut self) -> Result<(), Error>;
+}
+
+impl<C, I, R, SC> HasNestedStage for StdState<C, I, R, SC> {
     fn enter_inner_stage(&mut self) -> Result<(), Error> {
         self.stage_stack.enter_inner_stage()
     }

libafl/src/state/stack.rs

@@ -3,7 +3,10 @@ use alloc::vec::Vec;
 use libafl_bolts::Error;
 use serde::{Deserialize, Serialize};
 
-use crate::stages::{HasCurrentStageId, HasNestedStageStatus, StageId};
+use crate::{
+    stages::StageId,
+    state::{HasCurrentStageId, HasNestedStage},
+};
 
 /// A stack to keep track of which stage is executing
 #[derive(Serialize, Deserialize, Clone, Debug, Default)]
@@ -41,7 +44,7 @@ impl HasCurrentStageId for StageStack {
     }
 }
 
-impl HasNestedStageStatus for StageStack {
+impl HasNestedStage for StageStack {
     fn enter_inner_stage(&mut self) -> Result<(), Error> {
         self.stage_depth += 1;
         Ok(())

libafl_libfuzzer/log

@@ -0,0 +1,4 @@
+warning: profiles for the non root package will be ignored, specify profiles at the workspace root:
+package:   /home/toka/LibAFL/bindings/pylibafl/Cargo.toml
+workspace: /home/toka/LibAFL/Cargo.toml
+   Compiling libafl_libfuzzer v0.15.0 (/home/toka/LibAFL/libafl_libfuzzer)

libafl_libfuzzer/runtime/src/fuzz.rs

@@ -11,8 +11,8 @@ use libafl::{
     },
     executors::ExitKind,
     monitors::{tui::TuiMonitor, Monitor, MultiMonitor},
-    stages::{HasCurrentStageId, StagesTuple},
+    stages::StagesTuple,
-    state::{HasExecutions, HasLastReportTime, HasSolutions, Stoppable},
+    state::{HasCurrentStageId, HasExecutions, HasLastReportTime, HasSolutions, Stoppable},
     Error, Fuzzer, HasMetadata,
 };
 use libafl_bolts::{

libafl_libfuzzer/runtime/src/report.rs

@@ -5,8 +5,8 @@ use libafl::{
     executors::HasObservers,
     feedbacks::MapFeedbackMetadata,
     monitors::SimpleMonitor,
-    stages::{HasCurrentStageId, StagesTuple},
+    stages::StagesTuple,
-    state::{HasExecutions, HasLastReportTime, Stoppable},
+    state::{HasCurrentStageId, HasExecutions, HasLastReportTime, Stoppable},
     Error, Fuzzer, HasMetadata, HasNamedMetadata,
 };
 

libafl_sugar/src/forkserver.rs

@@ -257,7 +257,7 @@ impl ForkserverBytesCoverageSugar<'_> {
                         &mut mgr,
                         iters,
                     )?;
-                    mgr.on_restart(&mut state)?;
+                    mgr.mgr_on_restart(&mut state)?;
                     std::process::exit(0);
                 } else {
                     fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -278,7 +278,7 @@ impl ForkserverBytesCoverageSugar<'_> {
                         &mut mgr,
                         iters,
                     )?;
-                    mgr.on_restart(&mut state)?;
+                    mgr.mgr_on_restart(&mut state)?;
                     std::process::exit(0);
                 } else {
                     fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;

libafl_sugar/src/inmemory.rs

@@ -281,7 +281,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.on_restart(&mut state)?;
+                        mgr.mgr_on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -296,7 +296,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.on_restart(&mut state)?;
+                        mgr.mgr_on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -318,7 +318,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.on_restart(&mut state)?;
+                        mgr.mgr_on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -333,7 +333,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.on_restart(&mut state)?;
+                        mgr.mgr_on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;

libafl_sugar/src/qemu.rs

@@ -319,7 +319,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.on_restart(&mut state)?;
+                        mgr.mgr_on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -340,7 +340,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.on_restart(&mut state)?;
+                        mgr.mgr_on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -431,7 +431,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.on_restart(&mut state)?;
+                        mgr.mgr_on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -452,7 +452,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.on_restart(&mut state)?;
+                        mgr.mgr_on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;