Allow dyn in StagesTuple, add Current Testcase API, Untraitify Progress (#1915)

* Move into_vec to extra trait

* fix no_std

* First step towards stages tuples as vecs

* Allow dyn in StagesTuple, add Current Testcase API, un-traitify Progress, cleanups

* Move from generics to impl keyword, more replacements with better API

* rename fn

* Fix additional stages, more cleanup, rename progress to retries

* Fix more fixes

* Fixes

* Rename ProgressHelper -> RestartHelper

* Fix sugar, python, add perform_restartable

* fixes

* remove prelude bs

* rename to restart_progress_should_run

* more cleanup, remove tests I don't understand (sorry)

* fix docs

* more fix

* fix miri

* unsafe safety annotations

* more comments

* last docs

* Mixed_attributes only allowed for bindgen

bindings/pylibafl/Cargo.toml

@@ -6,11 +6,13 @@ edition = "2021"
 [dependencies]
 pyo3 = { version = "0.18.3", features = ["extension-module"] }
 pyo3-log = "0.8.1"
-libafl_qemu = { path = "../../libafl_qemu", version = "0.11.2", features = ["python"] }
 libafl_sugar = { path = "../../libafl_sugar", version = "0.11.2", features = ["python"] }
 libafl = { path = "../../libafl", version = "0.11.2", features = ["python"] }
 libafl_bolts = { path = "../../libafl_bolts", version = "0.11.2", features = ["python"] }
 
+[target.'cfg(target_os = "linux")'.dependencies]
+libafl_qemu = { path = "../../libafl_qemu", version = "0.11.2", features = ["python"] }
+
 [build-dependencies]
 pyo3-build-config = { version = "0.17" }
 

bindings/pylibafl/src/lib.rs

@@ -1,8 +1,3 @@
-use libafl;
-use libafl_bolts;
-#[cfg(target_os = "linux")]
-use libafl_qemu;
-use libafl_sugar;
 use pyo3::{prelude::*, types::PyDict};
 
 const LIBAFL_CODE: &str = r#"

fuzzers/baby_fuzzer_minimizing/src/main.rs

@@ -57,7 +57,7 @@ pub fn main() -> Result<(), Error> {
         // RNG
         StdRand::with_seed(current_nanos()),
         // Corpus that will be evolved, we keep it in memory for performance
-        InMemoryOnDiskCorpus::new(&corpus_dir).unwrap(),
+        InMemoryOnDiskCorpus::new(corpus_dir).unwrap(),
         // Corpus in which we store solutions (crashes in this example),
         // on disk so the user can get them after stopping the fuzzer
         OnDiskCorpus::new(&solution_dir).unwrap(),

fuzzers/baby_fuzzer_tokens/src/main.rs

@@ -1,6 +1,6 @@
 #[cfg(windows)]
 use std::ptr::write_volatile;
-use std::{fs, io::Read, path::PathBuf, ptr::write};
+use std::{fs, io::Read, path::PathBuf};
 
 use libafl::{
     corpus::{InMemoryCorpus, OnDiskCorpus},

fuzzers/baby_fuzzer_wasm/src/lib.rs

@@ -11,7 +11,7 @@ use libafl::{
     mutators::{havoc_mutations, StdScheduledMutator},
     observers::StdMapObserver,
     schedulers::QueueScheduler,
-    stages::StdMutationalStage,
+    stages::{ExecutionCountRestartHelperMetadata, StdMutationalStage},
     state::{HasSolutions, StdState},
     Fuzzer, StdFuzzer,
 };
@@ -23,7 +23,7 @@ use web_sys::{Performance, Window};
 
 use crate::utils::set_panic_hook;
 
-// defined for internal use by libafl
+// Defined for internal use by LibAFL
 #[no_mangle]
 #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
 pub extern "C" fn external_current_millis() -> u64 {
@@ -39,8 +39,14 @@ pub extern "C" fn external_current_millis() -> u64 {
 pub fn fuzz() {
     set_panic_hook();
 
+    // We need to register the types as LibAFL doesn't support `SerdeAny`
+    // auto registration in non-standard environments.
+    //
+    // # Safety
+    // No concurrency in WASM so these accesses are not racing.
     unsafe {
         RegistryBuilder::register::<MapFeedbackMetadata<u8>>();
+        RegistryBuilder::register::<ExecutionCountRestartHelperMetadata>();
     }
 
     let mut signals = [0u8; 64];

fuzzers/fuzzbench_forkserver_cmplog/src/main.rs

@@ -9,7 +9,7 @@ use std::{
 
 use clap::{Arg, ArgAction, Command};
 use libafl::{
-    corpus::{Corpus, HasCurrentCorpusIdx, InMemoryOnDiskCorpus, OnDiskCorpus},
+    corpus::{Corpus, InMemoryOnDiskCorpus, OnDiskCorpus},
     events::SimpleEventManager,
     executors::forkserver::ForkserverExecutor,
     feedback_or,
@@ -29,7 +29,7 @@ use libafl::{
         calibrate::CalibrationStage, mutational::MultiMutationalStage,
         power::StdPowerMutationalStage, ColorizationStage, IfStage,
     },
-    state::{HasCorpus, HasMetadata, StdState},
+    state::{HasCorpus, HasCurrentTestcase, HasMetadata, StdState},
     Error,
 };
 use libafl_bolts::{
@@ -373,14 +373,8 @@ fn fuzz(
                   state: &mut StdState<_, InMemoryOnDiskCorpus<_>, _, _>,
                   _event_manager: &mut _|
          -> Result<bool, Error> {
-            let Some(corpus_id) = state.current_corpus_idx()? else {
+            let testcase = state.current_testcase()?;
-                return Err(Error::illegal_state(
+            let res = testcase.scheduled_count() == 1; // let's try on the 2nd trial
-                    "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
 
             Ok(res)
         };

fuzzers/qemu_coverage/src/fuzzer.rs

@@ -8,12 +8,11 @@ use std::{env, fs::DirEntry, io, path::PathBuf, process};
 use clap::{builder::Str, Parser};
 use libafl::{
     corpus::{Corpus, NopCorpus},
-    events::{launcher::Launcher, EventConfig, EventRestarter},
+    events::{launcher::Launcher, EventConfig, EventRestarter, LlmpRestartingEventManager},
     executors::ExitKind,
     fuzzer::StdFuzzer,
     inputs::{BytesInput, HasTargetBytes},
     monitors::MultiMonitor,
-    prelude::LlmpRestartingEventManager,
     schedulers::QueueScheduler,
     state::{HasCorpus, StdState},
     Error,

libafl/src/corpus/ondisk.rs

@@ -42,7 +42,7 @@ pub struct OnDiskMetadata<'a> {
     /// The exec time for this [`Testcase`]
     pub exec_time: &'a Option<Duration>,
     /// The amount of executions for this [`Testcase`]
-    pub executions: &'a usize,
+    pub executions: &'a u64,
 }
 
 /// A corpus able to store [`Testcase`]s to disk, and load them from disk, when they are being used.

libafl/src/corpus/testcase.rs

@@ -59,7 +59,7 @@ where
     /// Cached len of the input, if any
     cached_len: Option<usize>,
     /// Number of executions done at discovery time
-    executions: usize,
+    executions: u64,
     /// Number of fuzzing iterations of this particular input updated in `perform_mutational`
     scheduled_count: usize,
     /// Parent [`CorpusId`], if known
@@ -174,13 +174,13 @@ where
 
     /// Get the executions
     #[inline]
-    pub fn executions(&self) -> &usize {
+    pub fn executions(&self) -> &u64 {
         &self.executions
     }
 
     /// Get the executions (mutable)
     #[inline]
-    pub fn executions_mut(&mut self) -> &mut usize {
+    pub fn executions_mut(&mut self) -> &mut u64 {
         &mut self.executions
     }
 
@@ -258,7 +258,7 @@ where
 
     /// Create a new Testcase instance given an [`Input`] and the number of executions
     #[inline]
-    pub fn with_executions(mut input: I, executions: usize) -> Self {
+    pub fn with_executions(mut input: I, executions: u64) -> Self {
         input.wrapped_as_testcase();
         Self {
             input: Some(input),
@@ -541,7 +541,7 @@ pub mod pybind {
         }
 
         #[getter]
-        fn executions(&self) -> usize {
+        fn executions(&self) -> u64 {
             *self.inner.as_ref().executions()
         }
 

libafl/src/events/llmp.rs

@@ -251,7 +251,7 @@ where
                 if id == client_id {
                     // do not update executions for forwarded messages, otherwise we loose the total order
                     // as a forwarded msg with a lower executions may arrive after a stats msg with an higher executions
-                    client.update_executions(*executions as u64, *time);
+                    client.update_executions(*executions, *time);
                 }
                 monitor.display(event.name(), id);
                 Ok(BrokerEventResult::Forward)
@@ -264,7 +264,7 @@ where
                 // TODO: The monitor buffer should be added on client add.
                 monitor.client_stats_insert(client_id);
                 let client = monitor.client_stats_mut_for(client_id);
-                client.update_executions(*executions as u64, *time);
+                client.update_executions(*executions, *time);
                 monitor.display(event.name(), client_id);
                 Ok(BrokerEventResult::Handled)
             }
@@ -294,7 +294,7 @@ where
                 let client = monitor.client_stats_mut_for(client_id);
 
                 // Update the normal monitor for this client
-                client.update_executions(*executions as u64, *time);
+                client.update_executions(*executions, *time);
 
                 // Update the performance monitor for this client
                 client.update_introspection_monitor((**introspection_monitor).clone());

libafl/src/events/mod.rs

@@ -296,7 +296,7 @@ where
         /// The time of generation of the event
         time: Duration,
         /// The executions of this client
-        executions: usize,
+        executions: u64,
         /// The original sender if, if forwarded
         forward_id: Option<ClientId>,
     },
@@ -305,7 +305,7 @@ where
         /// The time of generation of the [`Event`]
         time: Duration,
         /// The executions of this client
-        executions: usize,
+        executions: u64,
         /// [`PhantomData`]
         phantom: PhantomData<I>,
     },
@@ -324,7 +324,7 @@ where
         /// The time of generation of the event
         time: Duration,
         /// The executions of this client
-        executions: usize,
+        executions: u64,
         /// Current performance statistics
         introspection_monitor: Box<ClientPerfMonitor>,
 

libafl/src/events/simple.rs

@@ -216,7 +216,7 @@ where
                     .update_corpus_size(*corpus_size as u64);
                 monitor
                     .client_stats_mut_for(ClientId(0))
-                    .update_executions(*executions as u64, *time);
+                    .update_executions(*executions, *time);
                 monitor.display(event.name(), ClientId(0));
                 Ok(BrokerEventResult::Handled)
             }
@@ -229,7 +229,7 @@ where
                 monitor.client_stats_insert(ClientId(0));
                 let client = monitor.client_stats_mut_for(ClientId(0));
 
-                client.update_executions(*executions as u64, *time);
+                client.update_executions(*executions, *time);
 
                 monitor.display(event.name(), ClientId(0));
                 Ok(BrokerEventResult::Handled)
@@ -257,7 +257,7 @@ where
                 // TODO: The monitor buffer should be added on client add.
                 monitor.client_stats_insert(ClientId(0));
                 let client = monitor.client_stats_mut_for(ClientId(0));
-                client.update_executions(*executions as u64, *time);
+                client.update_executions(*executions, *time);
                 client.update_introspection_monitor((**introspection_monitor).clone());
                 monitor.display(event.name(), ClientId(0));
                 Ok(BrokerEventResult::Handled)

libafl/src/events/tcp.rs

@@ -333,7 +333,7 @@ where
                 monitor.client_stats_insert(id);
                 let client = monitor.client_stats_mut_for(id);
                 client.update_corpus_size(*corpus_size as u64);
-                client.update_executions(*executions as u64, *time);
+                client.update_executions(*executions, *time);
                 monitor.display(event.name(), id);
                 Ok(BrokerEventResult::Forward)
             }
@@ -345,7 +345,7 @@ where
                 // TODO: The monitor buffer should be added on client add.
                 monitor.client_stats_insert(client_id);
                 let client = monitor.client_stats_mut_for(client_id);
-                client.update_executions(*executions as u64, *time);
+                client.update_executions(*executions, *time);
                 monitor.display(event.name(), client_id);
                 Ok(BrokerEventResult::Handled)
             }
@@ -375,7 +375,7 @@ where
                 let client = monitor.client_stats_mut_for(client_id);
 
                 // Update the normal monitor for this client
-                client.update_executions(*executions as u64, *time);
+                client.update_executions(*executions, *time);
 
                 // Update the performance monitor for this client
                 client.update_introspection_monitor((**introspection_monitor).clone());

libafl/src/executors/hooks/inprocess.rs

@@ -198,7 +198,7 @@ impl ExecutorHook for InProcessHooks {
             (*data).timeout_handler = self.timeout_handler;
         }
 
-        #[cfg(feature = "std")]
+        #[cfg(all(feature = "std", not(all(miri, target_vendor = "apple"))))]
         self.timer_mut().set_timer();
     }
 
@@ -212,7 +212,7 @@ impl ExecutorHook for InProcessHooks {
         _input: &I,
     ) {
         // timeout stuff
-        #[cfg(feature = "std")]
+        #[cfg(all(feature = "std", not(all(miri, target_vendor = "apple"))))]
         self.timer_mut().unset_timer();
     }
 }

libafl/src/executors/inprocess_fork/inner.rs

@@ -70,7 +70,7 @@ where
             .field("observers", &self.observers)
             .field("shmem_provider", &self.shmem_provider)
             .field("itimerspec", &self.itimerspec)
-            .finish()
+            .finish_non_exhaustive()
     }
 
     #[cfg(not(target_os = "linux"))]
@@ -81,7 +81,7 @@ where
             .field("observers", &self.observers)
             .field("shmem_provider", &self.shmem_provider)
             .field("itimerval", &self.itimerval)
-            .finish();
+            .finish_non_exhaustive();
     }
 }
 
@@ -278,7 +278,7 @@ where
         })
     }
 
-    /// Creates a new [`GenericInProcessForkExecutor`], non linux
+    /// Creates a new [`GenericInProcessForkExecutorInner`], non linux
     #[cfg(not(target_os = "linux"))]
     #[allow(clippy::too_many_arguments)]
     pub fn with_hooks(

libafl/src/feedbacks/map.rs

@@ -401,7 +401,7 @@ where
     fn init_state(&mut self, state: &mut S) -> Result<(), Error> {
         // Initialize `MapFeedbackMetadata` with an empty vector and add it to the state.
         // The `MapFeedbackMetadata` would be resized on-demand in `is_interesting`
-        state.add_named_metadata(MapFeedbackMetadata::<T>::default(), &self.name);
+        state.add_named_metadata(&self.name, MapFeedbackMetadata::<T>::default());
         Ok(())
     }
 

libafl/src/feedbacks/new_hash_feedback.rs

@@ -92,8 +92,8 @@ where
 {
     fn init_state(&mut self, state: &mut S) -> Result<(), Error> {
         state.add_named_metadata(
-            NewHashFeedbackMetadata::with_capacity(self.capacity),
             &self.name,
+            NewHashFeedbackMetadata::with_capacity(self.capacity),
         );
         Ok(())
     }

libafl/src/lib.rs

@@ -135,8 +135,12 @@ pub unsafe extern "C" fn external_current_millis() -> u64 {
 #[cfg(test)]
 mod tests {
 
+    #[cfg(miri)]
+    use libafl_bolts::serdeany::RegistryBuilder;
     use libafl_bolts::{rands::StdRand, tuples::tuple_list};
 
+    #[cfg(miri)]
+    use crate::stages::ExecutionCountRestartHelperMetadata;
     use crate::{
         corpus::{Corpus, InMemoryCorpus, Testcase},
         events::NopEventManager,
@@ -155,6 +159,13 @@ mod tests {
     #[test]
     #[allow(clippy::similar_names)]
     fn test_fuzzer() {
+        // # Safety
+        // No concurrency per testcase
+        #[cfg(miri)]
+        unsafe {
+            RegistryBuilder::register::<ExecutionCountRestartHelperMetadata>();
+        }
+
         let rand = StdRand::with_seed(0);
 
         let mut corpus = InMemoryCorpus::<BytesInput>::new();

libafl/src/mutators/tuneable.rs

@@ -381,6 +381,8 @@ mod test {
 
     #[test]
     fn test_tuning() {
+        // # Safety
+        // No concurrency per testcase
         #[cfg(any(not(feature = "serdeany_autoreg"), miri))]
         unsafe {
             TuneableScheduledMutatorMetadata::register();
@@ -404,6 +406,8 @@ mod test {
 
     #[test]
     fn test_mutation_distribution() {
+        // # Safety
+        // No concurrency per testcase
         #[cfg(any(not(feature = "serdeany_autoreg"), miri))]
         unsafe {
             TuneableScheduledMutatorMetadata::register();

libafl/src/observers/cmp.rs

@@ -232,12 +232,7 @@ where
         S: HasMetadata,
     {
         #[allow(clippy::option_if_let_else)] // we can't mutate state in a closure
-        let meta = if let Some(meta) = state.metadata_map_mut().get_mut::<M>() {
+        let meta = state.metadata_or_insert_with(|| M::new_metadata());
-            meta
-        } else {
-            state.add_metadata(M::new_metadata());
-            state.metadata_map_mut().get_mut::<M>().unwrap()
-        };
 
         let usable_count = self.usable_count();
         let cmp_observer_data = self.cmp_observer_data();

libafl/src/schedulers/probabilistic_sampling.rs

@@ -228,6 +228,8 @@ mod tests {
 
     #[test]
     fn test_prob_sampling() {
+        // # Safety
+        // No concurrency per testcase
         #[cfg(any(not(feature = "serdeany_autoreg"), miri))]
         unsafe {
             super::ProbabilityMetadata::register();

libafl/src/schedulers/weighted.rs

@@ -114,13 +114,9 @@ where
     /// Create a new [`WeightedScheduler`]
     #[must_use]
     pub fn with_schedule(state: &mut S, map_observer: &O, strat: Option<PowerSchedule>) -> Self {
-        if !state.has_metadata::<SchedulerMetadata>() {
+        let _ = state.metadata_or_insert_with(|| SchedulerMetadata::new(strat));
-            state.add_metadata(SchedulerMetadata::new(strat));
+        let _ = state.metadata_or_insert_with(WeightedScheduleMetadata::new);
-        }
 
-        if !state.has_metadata::<WeightedScheduleMetadata>() {
-            state.add_metadata(WeightedScheduleMetadata::new());
-        }
         Self {
             strat,
             map_observer_name: map_observer.name().to_string(),

libafl/src/stages/calibrate.rs

@@ -12,7 +12,7 @@ use num_traits::Bounded;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    corpus::{Corpus, HasCurrentCorpusIdx, SchedulerTestcaseMetadata},
+    corpus::{Corpus, SchedulerTestcaseMetadata},
     events::{Event, EventFirer, LogSeverity},
     executors::{Executor, ExitKind, HasObservers},
     feedbacks::{map::MapFeedbackMetadata, HasObserverName},
@@ -20,8 +20,11 @@ use crate::{
     monitors::{AggregatorOps, UserStats, UserStatsValue},
     observers::{MapObserver, ObserversTuple, UsesObserver},
     schedulers::powersched::SchedulerMetadata,
-    stages::Stage,
+    stages::{ExecutionCountRestartHelper, Stage},
-    state::{HasCorpus, HasExecutions, HasMetadata, HasNamedMetadata, State, UsesState},
+    state::{
+        HasCorpus, HasCurrentTestcase, HasExecutions, HasMetadata, HasNamedMetadata, State,
+        UsesState,
+    },
     Error,
 };
 
@@ -68,7 +71,9 @@ pub struct CalibrationStage<O, OT, S> {
     map_observer_name: String,
     map_name: String,
     stage_max: usize,
+    /// If we should track stability
     track_stability: bool,
+    restart_helper: ExecutionCountRestartHelper,
     phantom: PhantomData<(O, OT, S)>,
 }
 
@@ -92,8 +97,6 @@ where
     E::State: HasCorpus + HasMetadata + HasNamedMetadata + HasExecutions,
     Z: Evaluator<E, EM, State = E::State>,
 {
-    type Progress = (); // TODO stage may be resumed, but how?
-
     #[inline]
     #[allow(
         clippy::let_and_return,
@@ -107,25 +110,19 @@ where
         state: &mut E::State,
         mgr: &mut EM,
     ) -> Result<(), Error> {
-        let Some(corpus_idx) = state.current_corpus_idx()? else {
-            return Err(Error::illegal_state(
-                "state is not currently processing a corpus index",
-            ));
-        };
-
         // Run this stage only once for each corpus entry and only if we haven't already inspected it
         {
-            let corpus = state.corpus().get(corpus_idx)?.borrow();
+            let testcase = state.current_testcase()?;
             // println!("calibration; corpus.scheduled_count() : {}", corpus.scheduled_count());
 
-            if corpus.scheduled_count() > 0 {
+            if testcase.scheduled_count() > 0 {
                 return Ok(());
             }
         }
 
         let mut iter = self.stage_max;
 
-        let input = state.corpus().cloned_input_for_id(corpus_idx)?;
+        let input = state.current_input_cloned()?;
 
         // Run once to get the initial calibration map
         executor.observers_mut().pre_exec_all(state, &input)?;
@@ -162,8 +159,11 @@ where
         let mut i = 1;
         let mut has_errors = false;
 
+        // If we restarted after a timeout or crash, do less iterations.
+        iter -= usize::try_from(self.restart_helper.execs_since_progress_start(state)?)?;
+
         while i < iter {
-            let input = state.corpus().cloned_input_for_id(corpus_idx)?;
+            let input = state.current_input_cloned()?;
 
             executor.observers_mut().pre_exec_all(state, &input)?;
             start = current_time();
@@ -268,7 +268,7 @@ where
             psmeta.set_bitmap_size_log(psmeta.bitmap_size_log() + libm::log2(bitmap_size as f64));
             psmeta.set_bitmap_entries(psmeta.bitmap_entries() + 1);
 
-            let mut testcase = state.corpus().get(corpus_idx)?.borrow_mut();
+            let mut testcase = state.current_testcase_mut()?;
 
             testcase.set_exec_time(total_time / (iter as u32));
             // log::trace!("time: {:#?}", testcase.exec_time());
@@ -301,7 +301,7 @@ where
             data.set_handicap(handicap);
         }
 
-        *state.executions_mut() += i;
+        *state.executions_mut() += u64::try_from(i).unwrap();
 
         // Send the stability event to the broker
         if unstable_found {
@@ -327,6 +327,16 @@ where
 
         Ok(())
     }
+
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        // TODO: Make sure this is the correct way / there may be a better way?
+        self.restart_helper.restart_progress_should_run(state)
+    }
+
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        // TODO: Make sure this is the correct way / there may be a better way?
+        self.restart_helper.clear_restart_progress(state)
+    }
 }
 
 impl<O, OT, S> CalibrationStage<O, OT, S>
@@ -347,6 +357,7 @@ where
             map_name: map_feedback.name().to_string(),
             stage_max: CAL_STAGE_START,
             track_stability: true,
+            restart_helper: ExecutionCountRestartHelper::default(),
             phantom: PhantomData,
         }
     }
@@ -363,6 +374,7 @@ where
             map_name: map_feedback.name().to_string(),
             stage_max: CAL_STAGE_START,
             track_stability: false,
+            restart_helper: ExecutionCountRestartHelper::default(),
             phantom: PhantomData,
         }
     }

libafl/src/stages/colorization.rs

@@ -6,18 +6,17 @@ use alloc::{
 };
 use core::{cmp::Ordering, fmt::Debug, marker::PhantomData, ops::Range};
 
-use libafl_bolts::{rands::Rand, tuples::MatchName};
+use libafl_bolts::{rands::Rand, tuples::MatchName, Named};
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    corpus::{Corpus, HasCurrentCorpusIdx},
     events::EventFirer,
     executors::{Executor, HasObservers},
     inputs::HasBytesVec,
     mutators::mutations::buffer_copy,
     observers::{MapObserver, ObserversTuple},
-    stages::Stage,
+    stages::{RetryRestartHelper, Stage},
-    state::{HasCorpus, HasMetadata, HasRand, UsesState},
+    state::{HasCorpus, HasCurrentTestcase, HasMetadata, HasNamedMetadata, HasRand, UsesState},
     Error,
 };
 
@@ -68,17 +67,24 @@ where
     type State = E::State;
 }
 
+impl<EM, O, E, Z> Named for ColorizationStage<EM, O, E, Z>
+where
+    E: UsesState,
+{
+    fn name(&self) -> &str {
+        &self.map_observer_name
+    }
+}
+
 impl<E, EM, O, Z> Stage<E, EM, Z> for ColorizationStage<EM, O, E, Z>
 where
     EM: UsesState<State = E::State> + EventFirer,
     E: HasObservers + Executor<EM, Z>,
-    E::State: HasCorpus + HasMetadata + HasRand,
+    E::State: HasCorpus + HasMetadata + HasRand + HasNamedMetadata,
     E::Input: HasBytesVec,
     O: MapObserver,
     Z: UsesState<State = E::State>,
 {
-    type Progress = (); // TODO this stage needs resume
-
     #[inline]
     #[allow(clippy::let_and_return)]
     fn perform(
@@ -93,6 +99,16 @@ where
 
         Ok(())
     }
+
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        // TODO this stage needs a proper resume
+        RetryRestartHelper::restart_progress_should_run(state, self, 3)
+    }
+
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        // TODO this stage needs a proper resume
+        RetryRestartHelper::clear_restart_progress(state, self)
+    }
 }
 
 /// Store the taint and the input
@@ -152,13 +168,7 @@ where
         manager: &mut EM,
         name: &str,
     ) -> Result<E::Input, Error> {
-        let Some(corpus_idx) = state.current_corpus_idx()? else {
+        let mut input = state.current_input_cloned()?;
-            return Err(Error::illegal_state(
-                "state is not currently processing a corpus index",
-            ));
-        };
-
-        let mut input = state.corpus().cloned_input_for_id(corpus_idx)?;
         // The backup of the input
         let backup = input.clone();
         // This is the buffer we'll randomly mutate during type_replace

libafl/src/stages/concolic.rs

@@ -4,23 +4,23 @@
 
 use alloc::string::String;
 #[cfg(feature = "concolic_mutation")]
-use alloc::{borrow::ToOwned, string::ToString, vec::Vec};
+use alloc::{string::ToString, vec::Vec};
+#[cfg(feature = "concolic_mutation")]
 use core::marker::PhantomData;
 
-use libafl_bolts::tuples::MatchName;
+use libafl_bolts::{tuples::MatchName, Named};
 
-use super::{RetryProgress, RetryingStage, Stage, TracingStage};
 #[cfg(all(feature = "concolic_mutation", feature = "introspection"))]
 use crate::monitors::PerfFeature;
 #[cfg(all(feature = "introspection", feature = "concolic_mutation"))]
 use crate::state::HasClientPerfMonitor;
-#[cfg(feature = "concolic_mutation")]
-use crate::state::State;
 use crate::{
-    corpus::{Corpus, HasCurrentCorpusIdx},
     executors::{Executor, HasObservers},
     observers::concolic::ConcolicObserver,
-    state::{HasCorpus, HasExecutions, HasMetadata, HasNamedMetadata, UsesState},
+    stages::{RetryRestartHelper, Stage, TracingStage},
+    state::{
+        HasCorpus, HasCurrentTestcase, HasExecutions, HasMetadata, HasNamedMetadata, UsesState,
+    },
     Error,
 };
 #[cfg(feature = "concolic_mutation")]
@@ -28,7 +28,10 @@ use crate::{
     inputs::HasBytesVec,
     mark_feature_time,
     observers::concolic::{ConcolicMetadata, SymExpr, SymExprRef},
-    start_timer, Evaluator,
+    stages::ExecutionCountRestartHelper,
+    start_timer,
+    state::State,
+    Evaluator,
 };
 
 /// Wraps a [`TracingStage`] to add concolic observing.
@@ -45,6 +48,12 @@ where
     type State = TE::State;
 }
 
+impl<EM, TE, Z> Named for ConcolicTracingStage<EM, TE, Z> {
+    fn name(&self) -> &str {
+        "ConcolicTracingStage"
+    }
+}
+
 impl<E, EM, TE, Z> Stage<E, EM, Z> for ConcolicTracingStage<EM, TE, Z>
 where
     E: UsesState<State = TE::State>,
@@ -53,8 +62,6 @@ where
     TE::State: HasExecutions + HasCorpus + HasNamedMetadata,
     Z: UsesState<State = TE::State>,
 {
-    type Progress = RetryProgress;
-
     #[inline]
     fn perform(
         &mut self,
@@ -63,16 +70,7 @@ where
         state: &mut TE::State,
         manager: &mut EM,
     ) -> Result<(), Error> {
-        let Some(corpus_idx) = state.current_corpus_idx()? else {
+        self.inner.trace(fuzzer, state, manager)?;
-            return Err(Error::illegal_state(
-                "state is not currently processing a corpus index",
-            ));
-        };
-        if Self::Progress::should_skip(state, &self.inner, corpus_idx)? {
-            return Ok(());
-        }
-
-        self.inner.trace(fuzzer, state, manager, corpus_idx)?;
         if let Some(observer) = self
             .inner
             .executor()
@@ -81,27 +79,25 @@ where
         {
             let metadata = observer.create_metadata_from_current_map();
             state
-                .corpus_mut()
+                .current_testcase_mut()?
-                .get(corpus_idx)
-                .unwrap()
-                .borrow_mut()
                 .metadata_map_mut()
                 .insert(metadata);
         }
         Ok(())
     }
-}
 
-impl<EM, TE, Z> RetryingStage for ConcolicTracingStage<EM, TE, Z> {
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
-    fn max_retries(&self) -> usize {
+        RetryRestartHelper::restart_progress_should_run(state, self, 3)
-        self.inner.max_retries()
+    }
+
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        RetryRestartHelper::clear_restart_progress(state, self)
     }
 }
 
 impl<EM, TE, Z> ConcolicTracingStage<EM, TE, Z> {
     /// Creates a new default tracing stage using the given [`Executor`], observing traces from a
-    /// [`ConcolicObserver`] with the given name. The [`RetryingStage::max_retries`] is
+    /// [`ConcolicObserver`] with the given name.
-    /// used from the provided inner stage.
     pub fn new(inner: TracingStage<EM, TE, Z>, observer_name: String) -> Self {
         Self {
             inner,
@@ -353,9 +349,12 @@ fn generate_mutations(iter: impl Iterator<Item = (SymExprRef, SymExpr)>) -> Vec<
 }
 
 /// A mutational stage that uses Z3 to solve concolic constraints attached to the [`crate::corpus::Testcase`] by the [`ConcolicTracingStage`].
+#[cfg(feature = "concolic_mutation")]
 #[derive(Clone, Debug)]
 pub struct SimpleConcolicMutationalStage<Z> {
-    _phantom: PhantomData<Z>,
+    /// The helper keeps track of progress for timeouting/restarting targets
+    restart_helper: ExecutionCountRestartHelper,
+    phantom: PhantomData<Z>,
 }
 
 #[cfg(feature = "concolic_mutation")]
@@ -373,10 +372,8 @@ where
     EM: UsesState<State = Z::State>,
     Z: Evaluator<E, EM>,
     Z::Input: HasBytesVec,
-    Z::State: State + HasExecutions + HasCorpus,
+    Z::State: State + HasExecutions + HasCorpus + HasMetadata,
 {
-    type Progress = (); // TODO we need a resume for this type
-
     #[inline]
     fn perform(
         &mut self,
@@ -385,30 +382,25 @@ where
         state: &mut Z::State,
         manager: &mut EM,
     ) -> Result<(), Error> {
-        let Some(corpus_idx) = state.current_corpus_idx()? else {
+        {
-            return Err(Error::illegal_state(
-                "state is not currently processing a corpus index",
-            ));
-        };
-
             start_timer!(state);
-        let testcase = state.corpus().get(corpus_idx)?.clone();
             mark_feature_time!(state, PerfFeature::GetInputFromCorpus);
+        }
+        let testcase = state.current_testcase()?.clone();
 
-        let mutations =
+        let mutations = testcase.metadata::<ConcolicMetadata>().ok().map(|meta| {
-            if let Some(meta) = testcase.borrow().metadata_map().get::<ConcolicMetadata>() {
             start_timer!(state);
-                let mutations = generate_mutations(meta.iter_messages());
+            let mutations = { generate_mutations(meta.iter_messages()) };
             mark_feature_time!(state, PerfFeature::Mutate);
-                Some(mutations)
+            mutations
-            } else {
+        });
-                None
+
-            };
+        let post_restart_skip_cnt =
+            usize::try_from(self.restart_helper.execs_since_progress_start(state)?)?;
 
         if let Some(mutations) = mutations {
-            let input = { testcase.borrow().input().as_ref().unwrap().clone() };
+            for mutation in mutations.into_iter().skip(post_restart_skip_cnt) {
-            for mutation in mutations {
+                let mut input_copy = state.current_input_cloned()?;
-                let mut input_copy = input.to_owned();
                 for (index, new_byte) in mutation {
                     input_copy.bytes_mut()[index] = new_byte;
                 }
@@ -418,12 +410,24 @@ where
         }
         Ok(())
     }
+
+    #[inline]
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        self.restart_helper.restart_progress_should_run(state)
+    }
+
+    #[inline]
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        self.restart_helper.clear_restart_progress(state)
+    }
 }
 
+#[cfg(feature = "concolic_mutation")]
 impl<Z> Default for SimpleConcolicMutationalStage<Z> {
     fn default() -> Self {
         Self {
-            _phantom: PhantomData,
+            restart_helper: ExecutionCountRestartHelper::default(),
+            phantom: PhantomData,
         }
     }
 }

libafl/src/stages/dump.rs

@@ -52,8 +52,6 @@ where
     Z: UsesState,
     Z::State: HasCorpus + HasSolutions + HasRand + HasMetadata,
 {
-    type Progress = (); // if this fails, we have bigger problems
-
     #[inline]
     fn perform(
         &mut self,
@@ -115,6 +113,18 @@ where
 
         Ok(())
     }
+
+    #[inline]
+    fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
+        // Not executing the target, so restart safety is not needed
+        Ok(true)
+    }
+
+    #[inline]
+    fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> {
+        // Not executing the target, so restart safety is not needed
+        Ok(())
+    }
 }
 
 impl<CB, EM, Z> DumpToDiskStage<CB, EM, Z>

libafl/src/stages/generalization.rs

@@ -6,7 +6,7 @@ use alloc::{
 };
 use core::{fmt::Debug, marker::PhantomData};
 
-use libafl_bolts::AsSlice;
+use libafl_bolts::{AsSlice, Named};
 
 use crate::{
     corpus::{Corpus, HasCurrentCorpusIdx},
@@ -15,9 +15,9 @@ use crate::{
     inputs::{BytesInput, GeneralizedInputMetadata, GeneralizedItem, HasBytesVec, UsesInput},
     mark_feature_time,
     observers::{MapObserver, ObserversTuple},
-    stages::Stage,
+    stages::{RetryRestartHelper, Stage},
     start_timer,
-    state::{HasCorpus, HasExecutions, HasMetadata, UsesState},
+    state::{HasCorpus, HasExecutions, HasMetadata, HasNamedMetadata, UsesState},
     Error,
 };
 #[cfg(feature = "introspection")]
@@ -47,6 +47,12 @@ pub struct GeneralizationStage<EM, O, OT, Z> {
     phantom: PhantomData<(EM, O, OT, Z)>,
 }
 
+impl<EM, O, OT, Z> Named for GeneralizationStage<EM, O, OT, Z> {
+    fn name(&self) -> &str {
+        "GeneralizationStage"
+    }
+}
+
 impl<EM, O, OT, Z> UsesState for GeneralizationStage<EM, O, OT, Z>
 where
     EM: UsesState,
@@ -60,12 +66,11 @@ where
     O: MapObserver,
     E: Executor<EM, Z> + HasObservers,
     E::Observers: ObserversTuple<E::State>,
-    E::State: UsesInput<Input = BytesInput> + HasExecutions + HasMetadata + HasCorpus,
+    E::State:
+        UsesInput<Input = BytesInput> + HasExecutions + HasMetadata + HasCorpus + HasNamedMetadata,
     EM: UsesState<State = E::State>,
     Z: UsesState<State = E::State>,
 {
-    type Progress = (); // TODO this stage needs a resume
-
     #[inline]
     #[allow(clippy::too_many_lines)]
     fn perform(
@@ -312,6 +317,18 @@ where
 
         Ok(())
     }
+
+    #[inline]
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        // TODO: We need to be able to resume better if something crashes or times out
+        RetryRestartHelper::restart_progress_should_run(state, self, 3)
+    }
+
+    #[inline]
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        // TODO: We need to be able to resume better if something crashes or times out
+        RetryRestartHelper::clear_restart_progress(state, self)
+    }
 }
 
 impl<EM, O, OT, Z> GeneralizationStage<EM, O, OT, Z>

libafl/src/stages/logics.rs

@@ -3,36 +3,31 @@
 use core::marker::PhantomData;
 
 use crate::{
-    stages::{HasCurrentStage, HasNestedStageStatus, Stage, StageProgress, StagesTuple},
+    stages::{HasCurrentStage, HasNestedStageStatus, Stage, StagesTuple},
     state::UsesState,
     Error,
 };
 
 /// Progress for nested stages. This merely enters/exits the inner stage's scope.
 #[derive(Debug)]
-pub struct NestedStageProgress;
+pub struct NestedStageRestartHelper;
 
-impl<S, ST> StageProgress<S, ST> for NestedStageProgress
+impl NestedStageRestartHelper {
-where
+    fn restart_progress_should_run<S, ST>(state: &mut S, _stage: &ST) -> Result<bool, Error>
+    where
         S: HasNestedStageStatus,
-{
+    {
-    fn initialize_progress(state: &mut S, _stage: &ST) -> Result<(), Error> {
         state.enter_inner_stage()?;
-        Ok(())
+        Ok(true)
     }
 
-    fn clear_progress(state: &mut S, _stage: &ST) -> Result<(), Error> {
+    fn clear_restart_progress<S, ST>(state: &mut S, _stage: &ST) -> Result<(), Error>
+    where
+        S: HasNestedStageStatus,
+    {
         state.exit_inner_stage()?;
         Ok(())
     }
-
-    fn progress<'a>(_state: &'a S, _stage: &ST) -> Result<&'a Self, Error> {
-        unimplemented!("NestedStageProgress should not be queried")
-    }
-
-    fn progress_mut<'a>(_state: &'a mut S, _stage: &ST) -> Result<&'a mut Self, Error> {
-        unimplemented!("NestedStageProgress should not be queried")
-    }
 }
 
 #[derive(Debug)]
@@ -70,8 +65,6 @@ where
     Z: UsesState<State = E::State>,
     E::State: HasNestedStageStatus,
 {
-    type Progress = NestedStageProgress;
-
     fn perform(
         &mut self,
         fuzzer: &mut Z,
@@ -86,6 +79,14 @@ where
 
         Ok(())
     }
+
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        NestedStageRestartHelper::restart_progress_should_run(state, self)
+    }
+
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        NestedStageRestartHelper::clear_restart_progress(state, self)
+    }
 }
 
 impl<CB, E, EM, ST, Z> WhileStage<CB, E, EM, ST, Z>
@@ -142,8 +143,6 @@ where
     Z: UsesState<State = E::State>,
     E::State: HasNestedStageStatus,
 {
-    type Progress = NestedStageProgress;
-
     fn perform(
         &mut self,
         fuzzer: &mut Z,
@@ -157,6 +156,14 @@ where
         }
         Ok(())
     }
+
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        NestedStageRestartHelper::restart_progress_should_run(state, self)
+    }
+
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        NestedStageRestartHelper::clear_restart_progress(state, self)
+    }
 }
 
 impl<CB, E, EM, ST, Z> IfStage<CB, E, EM, ST, Z>
@@ -217,8 +224,6 @@ where
     Z: UsesState<State = E::State>,
     E::State: HasNestedStageStatus,
 {
-    type Progress = NestedStageProgress;
-
     fn perform(
         &mut self,
         fuzzer: &mut Z,
@@ -252,6 +257,14 @@ where
 
         Ok(())
     }
+
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        NestedStageRestartHelper::restart_progress_should_run(state, self)
+    }
+
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        NestedStageRestartHelper::clear_restart_progress(state, self)
+    }
 }
 
 impl<CB, E, EM, ST1, ST2, Z> IfElseStage<CB, E, EM, ST1, ST2, Z>
@@ -305,8 +318,6 @@ where
     Z: UsesState<State = E::State>,
     E::State: HasNestedStageStatus,
 {
-    type Progress = NestedStageProgress;
-
     fn perform(
         &mut self,
         fuzzer: &mut Z,
@@ -320,6 +331,14 @@ where
             Ok(())
         }
     }
+
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        NestedStageRestartHelper::restart_progress_should_run(state, self)
+    }
+
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        NestedStageRestartHelper::clear_restart_progress(state, self)
+    }
 }
 
 impl<E, EM, ST, Z> OptionalStage<E, EM, ST, Z>
@@ -356,178 +375,3 @@ where
         }
     }
 }
-
-#[cfg(test)]
-mod test {
-    use core::{cell::RefCell, marker::PhantomData};
-
-    use libafl_bolts::{tuples::tuple_list, Error};
-
-    use crate::{
-        inputs::NopInput,
-        stages::{
-            test::{test_resume, test_resume_stages},
-            ClosureStage, IfElseStage, IfStage, Stage, WhileStage,
-        },
-        state::{test::test_std_state, State, UsesState},
-    };
-
-    #[test]
-    fn check_resumability_while() {
-        let once = RefCell::new(true);
-        let (completed, stages) = test_resume_stages();
-        let whilestage = WhileStage::new(|_, _, _, _| Ok(once.replace(false)), stages);
-        let resetstage = ClosureStage::new(|_, _, _, _| {
-            once.replace(true);
-            Ok(())
-        });
-
-        let mut state = test_std_state::<NopInput>();
-
-        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(|_, _, _, _| Ok(once.replace(false)), stages);
-        let resetstage = ClosureStage::new(|_, _, _, _| {
-            once.replace(true);
-            Ok(())
-        });
-
-        let mut state = test_std_state::<NopInput>();
-
-        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(
-            |_, _, _, _| Ok(true),
-            tuple_list!(IfStage::new(
-                |_, _, _, _| Ok(true),
-                tuple_list!(IfStage::new(
-                    |_, _, _, _| Ok(true),
-                    tuple_list!(IfStage::new(
-                        |_, _, _, _| Ok(true),
-                        tuple_list!(IfStage::new(|_, _, _, _| Ok(true), stages),),
-                    ),),
-                ))
-            )),
-        );
-
-        let mut state = test_std_state::<NopInput>();
-
-        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<S> UsesState for PanicStage<S>
-    where
-        S: State,
-    {
-        type State = S;
-    }
-
-    impl<E, EM, Z> Stage<E, EM, Z> for PanicStage<E::State>
-    where
-        E: UsesState,
-        EM: UsesState<State = E::State>,
-        Z: UsesState<State = E::State>,
-    {
-        type Progress = ();
-
-        fn perform(
-            &mut self,
-            _fuzzer: &mut Z,
-            _executor: &mut E,
-            _state: &mut Self::State,
-            _manager: &mut EM,
-        ) -> Result<(), Error> {
-            panic!("Test failed; panic stage should never be executed.");
-        }
-    }
-
-    #[test]
-    fn check_resumability_if_else_if() {
-        let once = RefCell::new(true);
-        let (completed, stages) = test_resume_stages();
-        let ifstage = IfElseStage::new(
-            |_, _, _, _| Ok(once.replace(false)),
-            stages,
-            tuple_list!(PanicStage::new()),
-        );
-        let resetstage = ClosureStage::new(|_, _, _, _| {
-            once.replace(true);
-            Ok(())
-        });
-
-        let mut state = test_std_state::<NopInput>();
-
-        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(
-            |_, _, _, _| Ok(once.replace(true)),
-            tuple_list!(PanicStage::new()),
-            stages,
-        );
-        let resetstage = ClosureStage::new(|_, _, _, _| {
-            once.replace(false);
-            Ok(())
-        });
-
-        let mut state = test_std_state::<NopInput>();
-
-        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(
-            |_, _, _, _| Ok(false),
-            tuple_list!(PanicStage::new()),
-            tuple_list!(IfElseStage::new(
-                |_, _, _, _| Ok(false),
-                tuple_list!(PanicStage::new()),
-                tuple_list!(IfElseStage::new(
-                    |_, _, _, _| Ok(false),
-                    tuple_list!(PanicStage::new()),
-                    tuple_list!(IfElseStage::new(
-                        |_, _, _, _| Ok(false),
-                        tuple_list!(PanicStage::new()),
-                        tuple_list!(IfElseStage::new(
-                            |_, _, _, _| Ok(false),
-                            tuple_list!(PanicStage::new()),
-                            stages,
-                        )),
-                    )),
-                )),
-            )),
-        );
-
-        let mut state = test_std_state::<NopInput>();
-
-        test_resume(&completed, &mut state, tuple_list!(ifstage));
-    }
-}

libafl/src/stages/mod.rs

@@ -4,19 +4,24 @@ A well-known [`Stage`], for example, is the mutational stage, running multiple [
 Other stages may enrich [`crate::corpus::Testcase`]s with metadata.
 */
 
-use core::marker::PhantomData;
+use alloc::{boxed::Box, vec::Vec};
+use core::{any, marker::PhantomData};
 
 pub use calibrate::CalibrationStage;
 pub use colorization::*;
 #[cfg(feature = "std")]
 pub use concolic::ConcolicTracingStage;
-#[cfg(feature = "std")]
+#[cfg(all(feature = "std", feature = "concolic_mutation"))]
 pub use concolic::SimpleConcolicMutationalStage;
 #[cfg(feature = "std")]
 pub use dump::*;
 pub use generalization::GeneralizationStage;
 use hashbrown::HashSet;
-use libafl_bolts::{impl_serdeany, tuples::HasConstLen};
+use libafl_bolts::{
+    impl_serdeany,
+    tuples::{HasConstLen, IntoVec},
+    Named,
+};
 pub use logics::*;
 pub use mutational::{MutationalStage, StdMutationalStage};
 pub use power::{PowerMutationalStage, StdPowerMutationalStage};
@@ -31,8 +36,8 @@ pub use tmin::{
 };
 pub use tracing::{ShadowTracingStage, TracingStage};
 pub use tuneable::*;
+use tuple_list::NonEmptyTuple;
 
-use self::push::PushStage;
 use crate::{
     corpus::{CorpusId, HasCurrentCorpusIdx},
     events::{EventFirer, EventRestarter, HasEventManagerId, ProgressReporter},
@@ -40,8 +45,9 @@ use crate::{
     inputs::UsesInput,
     observers::ObserversTuple,
     schedulers::Scheduler,
+    stages::push::PushStage,
     state::{
-        HasCorpus, HasExecutions, HasLastReportTime, HasMetadata, HasNamedMetadata, HasRand,
+        HasCorpus, HasExecutions, HasLastReportTime, HasMetadata, HasNamedMetadata, HasRand, State,
         UsesState,
     },
     Error, EvaluatorObservers, ExecutesInput, ExecutionProcessor, HasScheduler,
@@ -77,13 +83,21 @@ where
     EM: UsesState<State = Self::State>,
     Z: UsesState<State = Self::State>,
 {
-    // TODO: default this to () when associated_type_defaults is stable
+    /// This method will be called before every call to [`Stage::perform`].
-    // TODO: see RFC 2532: https://github.com/rust-lang/rust/issues/29661
+    /// Initialize the restart tracking for this stage, _if it is not yet initialized_.
-    // type Status: ResumableStageStatus = ();
+    /// On restart, this will be called again.
-    /// The resumption data for this stage. Set to () if resuming is not necessary/possible.
+    /// As long as [`Stage::clear_restart_progress`], all subsequent calls happen on restart.
-    type Progress: StageProgress<Self::State, Self>;
+    /// Returns `true`, if the stage's [`Stage::perform`] method should run, else `false`.
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error>;
 
-    /// Run the stage
+    /// Clear the current status tracking of the associated stage
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error>;
+
+    /// Run the stage.
+    ///
+    /// Before a call to perform, [`Stage::restart_progress_should_run`] will be (must be!) called.
+    /// After returning (so non-target crash or timeout in a restarting case), [`Stage::clear_restart_progress`] gets called.
+    /// A call to [`Stage::perform_restartable`] will do these things implicitly.
     fn perform(
         &mut self,
         fuzzer: &mut Z,
@@ -91,6 +105,20 @@ where
         state: &mut Self::State,
         manager: &mut EM,
     ) -> Result<(), Error>;
+
+    /// Run the stage, calling [`Stage::restart_progress_should_run`] and [`Stage::clear_restart_progress`] appropriately
+    fn perform_restartable(
+        &mut self,
+        fuzzer: &mut Z,
+        executor: &mut E,
+        state: &mut Self::State,
+        manager: &mut EM,
+    ) -> Result<(), Error> {
+        if self.restart_progress_should_run(state)? {
+            self.perform(fuzzer, executor, state, manager)?;
+        }
+        self.clear_restart_progress(state)
+    }
 }
 
 /// A tuple holding all `Stages` used for fuzzing.
@@ -158,9 +186,9 @@ where
             Some(idx) if idx == Self::LEN => {
                 // perform the stage, but don't set it
                 let stage = &mut self.0;
-                Head::Progress::initialize_progress(state, stage)?;
+
-                stage.perform(fuzzer, executor, state, manager)?;
+                stage.perform_restartable(fuzzer, executor, state, manager)?;
-                Head::Progress::clear_progress(state, stage)?;
+
                 state.clear_stage()?;
             }
             Some(idx) if idx > Self::LEN => {
@@ -169,10 +197,10 @@ where
             // this is None, but the match can't deduce that
             _ => {
                 state.set_stage(Self::LEN)?;
+
                 let stage = &mut self.0;
-                Head::Progress::initialize_progress(state, stage)?;
+                stage.perform_restartable(fuzzer, executor, state, manager)?;
-                stage.perform(fuzzer, executor, state, manager)?;
+
-                Head::Progress::clear_progress(state, stage)?;
                 state.clear_stage()?;
             }
         }
@@ -182,6 +210,71 @@ where
     }
 }
 
+impl<Head, Tail, E, EM, Z>
+    IntoVec<Box<dyn Stage<E, EM, Z, State = Head::State, Input = Head::Input>>> for (Head, Tail)
+where
+    Head: Stage<E, EM, Z> + 'static,
+    Tail: StagesTuple<E, EM, Head::State, Z>
+        + HasConstLen
+        + IntoVec<Box<dyn Stage<E, EM, Z, State = Head::State, Input = Head::Input>>>,
+    E: UsesState<State = Head::State>,
+    EM: UsesState<State = Head::State>,
+    Z: UsesState<State = Head::State>,
+    Head::State: HasCurrentStage,
+{
+    fn into_vec(self) -> Vec<Box<dyn Stage<E, EM, Z, State = Head::State, Input = Head::Input>>> {
+        let (head, tail) = self.uncons();
+        let mut ret = tail.0.into_vec();
+        ret.insert(0, Box::new(head));
+        ret
+    }
+}
+
+impl<Tail, E, EM, Z> IntoVec<Box<dyn Stage<E, EM, Z, State = Tail::State, Input = Tail::Input>>>
+    for (Tail,)
+where
+    Tail: UsesState + IntoVec<Box<dyn Stage<E, EM, Z, State = Tail::State, Input = Tail::Input>>>,
+    Z: UsesState<State = Tail::State>,
+    EM: UsesState<State = Tail::State>,
+    E: UsesState<State = Tail::State>,
+{
+    fn into_vec(self) -> Vec<Box<dyn Stage<E, EM, Z, State = Tail::State, Input = Tail::Input>>> {
+        self.0.into_vec()
+    }
+}
+
+impl<E, EM, Z> IntoVec<Box<dyn Stage<E, EM, Z, State = Z::State, Input = Z::Input>>>
+    for Vec<Box<dyn Stage<E, EM, Z, State = Z::State, Input = Z::Input>>>
+where
+    Z: UsesState,
+    EM: UsesState<State = Z::State>,
+    E: UsesState<State = Z::State>,
+{
+    fn into_vec(self) -> Vec<Box<dyn Stage<E, EM, Z, State = Z::State, Input = Z::Input>>> {
+        self
+    }
+}
+
+impl<E, EM, S, Z> StagesTuple<E, EM, S, Z>
+    for Vec<Box<dyn Stage<E, EM, Z, State = S, Input = S::Input>>>
+where
+    E: UsesState<State = S>,
+    EM: UsesState<State = S>,
+    Z: UsesState<State = S>,
+    S: UsesInput + HasCurrentStage + State,
+{
+    fn perform_all(
+        &mut self,
+        fuzzer: &mut Z,
+        executor: &mut E,
+        state: &mut S,
+        manager: &mut EM,
+    ) -> Result<(), Error> {
+        self.iter_mut()
+            .try_for_each(|x| x.perform_restartable(fuzzer, executor, state, manager))
+    }
+}
+
 /// A [`Stage`] that will call a closure
 #[derive(Debug)]
 pub struct ClosureStage<CB, E, EM, Z>
@@ -201,15 +294,24 @@ where
     type State = E::State;
 }
 
+impl<CB, E, EM, Z> Named for ClosureStage<CB, E, EM, Z>
+where
+    CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM) -> Result<(), Error>,
+    E: UsesState,
+{
+    fn name(&self) -> &str {
+        any::type_name::<Self>()
+    }
+}
+
 impl<CB, E, EM, Z> Stage<E, EM, Z> for ClosureStage<CB, E, EM, Z>
 where
     CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM) -> Result<(), Error>,
     E: UsesState,
     EM: UsesState<State = E::State>,
     Z: UsesState<State = E::State>,
+    E::State: HasNamedMetadata,
 {
-    type Progress = ();
-
     fn perform(
         &mut self,
         fuzzer: &mut Z,
@@ -219,6 +321,17 @@ where
     ) -> Result<(), Error> {
         (self.closure)(fuzzer, executor, state, manager)
     }
+
+    #[inline]
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        // Make sure we don't get stuck crashing on a single closure
+        RetryRestartHelper::restart_progress_should_run(state, self, 3)
+    }
+
+    #[inline]
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        RetryRestartHelper::clear_restart_progress(state, self)
+    }
 }
 
 /// A stage that takes a closure
@@ -292,8 +405,6 @@ where
         + EvaluatorObservers<OT>
         + HasScheduler<Scheduler = CS>,
 {
-    type Progress = (); // TODO implement resume
-
     fn perform(
         &mut self,
         fuzzer: &mut Z,
@@ -336,6 +447,193 @@ where
         self.push_stage
             .deinit(fuzzer, state, event_mgr, executor.observers_mut())
     }
+
+    #[inline]
+    fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
+        // TODO: Proper restart handling - call post_exec at the right time, etc...
+        Ok(true)
+    }
+
+    #[inline]
+    fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> {
+        Ok(())
+    }
+}
+
+/// Progress which permits a fixed amount of resumes per round of fuzzing. If this amount is ever
+/// exceeded, the input will no longer be executed by this stage.
+#[derive(Clone, Deserialize, Serialize, Debug)]
+pub struct RetryRestartHelper {
+    tries_remaining: Option<usize>,
+    skipped: HashSet<CorpusId>,
+}
+
+impl_serdeany!(RetryRestartHelper);
+
+impl RetryRestartHelper {
+    /// Initializes (or counts down in) the progress helper, giving it the amount of max retries
+    ///
+    /// Returns `true` if the stage should run
+    pub fn restart_progress_should_run<S, ST>(
+        state: &mut S,
+        stage: &ST,
+        max_retries: usize,
+    ) -> Result<bool, Error>
+    where
+        S: HasNamedMetadata + HasCurrentCorpusIdx,
+        ST: Named,
+    {
+        let corpus_idx = state.current_corpus_idx()?.ok_or_else(|| {
+            Error::illegal_state(
+                "No current_corpus_idx set in State, but called RetryRestartHelper::should_skip",
+            )
+        })?;
+
+        let initial_tries_remaining = max_retries + 1;
+        let metadata = state.named_metadata_or_insert_with(stage.name(), || Self {
+            tries_remaining: Some(initial_tries_remaining),
+            skipped: HashSet::new(),
+        });
+        let tries_remaining = metadata
+            .tries_remaining
+            .unwrap_or(initial_tries_remaining)
+            .checked_sub(1)
+            .ok_or_else(|| {
+                Error::illegal_state(
+                    "Attempted further retries after we had already gotten to none remaining.",
+                )
+            })?;
+
+        metadata.tries_remaining = Some(tries_remaining);
+
+        Ok(if tries_remaining == 0 {
+            metadata.skipped.insert(corpus_idx);
+            false
+        } else if metadata.skipped.contains(&corpus_idx) {
+            // skip this testcase, we already retried it often enough...
+            false
+        } else {
+            true
+        })
+    }
+
+    /// Clears the progress
+    pub fn clear_restart_progress<S, ST>(state: &mut S, stage: &ST) -> Result<(), Error>
+    where
+        S: HasNamedMetadata,
+        ST: Named,
+    {
+        state
+            .named_metadata_mut::<Self>(stage.name())?
+            .tries_remaining = None;
+        Ok(())
+    }
+}
+
+/// Trait for types which track the current stage
+pub trait HasCurrentStage {
+    /// Set the current stage; we have started processing this stage
+    fn set_stage(&mut self, idx: usize) -> Result<(), Error>;
+
+    /// Clear the current stage; we are done processing this stage
+    fn clear_stage(&mut self) -> Result<(), Error>;
+
+    /// Fetch the current stage -- typically used after a state recovery or transfer
+    fn current_stage(&self) -> Result<Option<usize>, 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: HasCurrentStage {
+    /// 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.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ExecutionCountRestartHelperMetadata {
+    /// How many executions we had when we started this stage initially (this round)
+    started_at_execs: u64,
+}
+
+/// A tool shed of functions to be used for stages that try to run for `n` iterations.
+///
+/// # Note
+/// This helper assumes resumable mutational stages are not nested.
+/// If you want to nest them, you will have to switch all uses of `metadata` in this helper to `named_metadata` instead.
+#[derive(Debug, Default, Clone)]
+pub struct ExecutionCountRestartHelper {
+    /// At what exec count this Stage was started (cache)
+    /// Only used as cache for the value stored in [`MutationalStageMetadata`].
+    started_at_execs: Option<u64>,
+}
+
+impl ExecutionCountRestartHelper {
+    /// Create a new [`ExecutionCountRestartHelperMetadata`]
+    #[must_use]
+    pub fn new() -> Self {
+        Self {
+            started_at_execs: None,
+        }
+    }
+
+    /// The execs done since start of this [`Stage`]/helper
+    pub fn execs_since_progress_start<S>(&mut self, state: &mut S) -> Result<u64, Error>
+    where
+        S: HasMetadata + HasExecutions,
+    {
+        let started_at_execs = if let Some(started_at_execs) = self.started_at_execs {
+            started_at_execs
+        } else {
+            state
+                .metadata::<ExecutionCountRestartHelperMetadata>()
+                .map(|x| {
+                    self.started_at_execs = Some(x.started_at_execs);
+                    x.started_at_execs
+                })
+                .map_err(|err| {
+                    Error::illegal_state(format!(
+                        "The ExecutionCountRestartHelperMetadata should have been set at this point - {err}"
+                    ))
+                })?
+        };
+        Ok(state.executions() - started_at_execs)
+    }
+
+    /// Initialize progress for the stage this wrapper wraps.
+    pub fn restart_progress_should_run<S>(&mut self, state: &mut S) -> Result<bool, Error>
+    where
+        S: HasMetadata + HasExecutions,
+    {
+        let executions = *state.executions();
+        let metadata = state.metadata_or_insert_with(|| ExecutionCountRestartHelperMetadata {
+            started_at_execs: executions,
+        });
+        self.started_at_execs = Some(metadata.started_at_execs);
+        Ok(true)
+    }
+
+    /// Clear progress for the stage this wrapper wraps.
+    pub fn clear_restart_progress<S>(&mut self, state: &mut S) -> Result<(), Error>
+    where
+        S: HasMetadata,
+    {
+        self.started_at_execs = None;
+        let _metadata = state.remove_metadata::<ExecutionCountRestartHelperMetadata>();
+        debug_assert!(_metadata.is_some(), "Called clear_restart_progress, but restart_progress_should_run was not called before (or did mutational stages get nested?)");
+        Ok(())
+    }
 }
 
 /// `Stage` Python bindings
@@ -344,6 +642,7 @@ where
 pub mod pybind {
     use alloc::vec::Vec;
 
+    use libafl_bolts::Named;
     use pyo3::prelude::*;
 
     use crate::{
@@ -352,7 +651,8 @@ pub mod pybind {
         executors::pybind::PythonExecutor,
         fuzzer::pybind::{PythonStdFuzzer, PythonStdFuzzerWrapper},
         stages::{
-            mutational::pybind::PythonStdMutationalStage, HasCurrentStage, Stage, StagesTuple,
+            mutational::pybind::PythonStdMutationalStage, HasCurrentStage, RetryRestartHelper,
+            Stage, StagesTuple,
         },
         state::{
             pybind::{PythonStdState, PythonStdStateWrapper},
@@ -377,9 +677,13 @@ pub mod pybind {
         type State = PythonStdState;
     }
 
-    impl Stage<PythonExecutor, PythonEventManager, PythonStdFuzzer> for PyObjectStage {
+    impl Named for PyObjectStage {
-        type Progress = (); // we don't support resumption in python, and maybe can't?
+        fn name(&self) -> &str {
+            "PyObjectStage"
+        }
+    }
 
+    impl Stage<PythonExecutor, PythonEventManager, PythonStdFuzzer> for PyObjectStage {
         #[inline]
         fn perform(
             &mut self,
@@ -410,6 +714,15 @@ pub mod pybind {
             })?;
             Ok(())
         }
+
+        fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+            // we don't support resumption in python, and maybe can't?
+            RetryRestartHelper::restart_progress_should_run(state, self, 2)
+        }
+
+        fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+            RetryRestartHelper::clear_restart_progress(state, self)
+        }
     }
 
     #[derive(Clone, Debug)]
@@ -472,10 +785,13 @@ pub mod pybind {
         type State = PythonStdState;
     }
 
-    impl Stage<PythonExecutor, PythonEventManager, PythonStdFuzzer> for PythonStage {
+    impl Named for PythonStage {
-        // TODO if we implement resumption for StdMutational, we need to apply it here
+        fn name(&self) -> &str {
-        type Progress = ();
+            "PythonStage"
+        }
+    }
 
+    impl Stage<PythonExecutor, PythonEventManager, PythonStdFuzzer> for PythonStage {
         #[inline]
         #[allow(clippy::let_and_return)]
         fn perform(
@@ -486,9 +802,24 @@ pub mod pybind {
             manager: &mut PythonEventManager,
         ) -> Result<(), Error> {
             unwrap_me_mut!(self.wrapper, s, {
-                s.perform(fuzzer, executor, state, manager)
+                s.perform_restartable(fuzzer, executor, state, manager)
             })
         }
+
+        #[inline]
+        fn restart_progress_should_run(
+            &mut self,
+            state: &mut PythonStdState,
+        ) -> Result<bool, Error> {
+            // TODO we need to apply MutationalStage-like resumption here.
+            // For now, make sure we don't get stuck crashing on a single test
+            RetryRestartHelper::restart_progress_should_run(state, self, 3)
+        }
+
+        #[inline]
+        fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+            RetryRestartHelper::clear_restart_progress(state, self)
+        }
     }
 
     #[derive(Clone, Debug)]
@@ -532,7 +863,7 @@ pub mod pybind {
                 } else {
                     state.set_stage(i)?;
                 }
-                s.perform(fuzzer, executor, state, manager)?;
+                s.perform_restartable(fuzzer, executor, state, manager)?;
                 state.clear_stage()?;
             }
             Ok(())
@@ -547,173 +878,17 @@ pub mod pybind {
     }
 }
 
-/// Trait for status tracking of stages which stash data to resume
-pub trait StageProgress<S, ST>
-where
-    ST: ?Sized,
-{
-    /// Initialize the current status tracking for this stage, if it is not yet initialised
-    fn initialize_progress(state: &mut S, stage: &ST) -> Result<(), Error>;
-
-    /// Clear the current status tracking of the associated stage
-    fn clear_progress(state: &mut S, stage: &ST) -> Result<(), Error>;
-
-    /// Get the current status tracking of this stage
-    fn progress<'a>(state: &'a S, stage: &ST) -> Result<&'a Self, Error>;
-
-    /// Get the current status tracking of this stage, mutably
-    fn progress_mut<'a>(state: &'a mut S, stage: &ST) -> Result<&'a mut Self, Error>;
-}
-
-impl<S, ST> StageProgress<S, ST> for () {
-    fn initialize_progress(_state: &mut S, _stage: &ST) -> Result<(), Error> {
-        Ok(())
-    }
-
-    fn clear_progress(_state: &mut S, _stage: &ST) -> Result<(), Error> {
-        Ok(())
-    }
-
-    fn progress<'a>(_state: &'a S, _stage: &ST) -> Result<&'a Self, Error> {
-        unimplemented!("The empty tuple resumable stage status should never be queried")
-    }
-
-    fn progress_mut<'a>(_state: &'a mut S, _stage: &ST) -> Result<&'a mut Self, Error> {
-        unimplemented!("The empty tuple resumable stage status should never be queried")
-    }
-}
-
-/// Progress which permits a fixed amount of resumes per round of fuzzing. If this amount is ever
-/// exceeded, the input will no longer be executed by this stage.
-#[derive(Clone, Deserialize, Serialize, Debug)]
-pub struct RetryProgress {
-    tries_remaining: Option<usize>,
-    skipped: HashSet<CorpusId>,
-}
-
-impl_serdeany!(RetryProgress);
-
-/// Stage which specifies a certain amount of retries over which a scheduled input is attempted. To
-/// be used in combination with [`RetryProgress`].
-pub trait RetryingStage {
-    /// The number of times each testcase may be retries.
-    fn max_retries(&self) -> usize;
-}
-
-impl<S, ST> StageProgress<S, ST> for RetryProgress
-where
-    S: HasNamedMetadata,
-    ST: RetryingStage,
-{
-    fn initialize_progress(state: &mut S, stage: &ST) -> Result<(), Error> {
-        if let Ok(metadata) = state.named_metadata_mut::<Self>(core::any::type_name_of_val(stage)) {
-            if let Some(ref mut remaining) = metadata.tries_remaining {
-                *remaining = remaining.checked_sub(1).ok_or_else(|| {
-                    Error::illegal_state(
-                        "Attempted further retries after we had already gotten to none remaining.",
-                    )
-                })?;
-            } else {
-                metadata.tries_remaining = Some(stage.max_retries() + 1);
-            }
-        } else {
-            state.add_named_metadata(
-                Self {
-                    tries_remaining: Some(stage.max_retries() + 1),
-                    skipped: HashSet::new(),
-                },
-                core::any::type_name_of_val(stage),
-            );
-        }
-        Ok(())
-    }
-
-    fn clear_progress(state: &mut S, stage: &ST) -> Result<(), Error> {
-        let metadata = state.named_metadata_mut::<Self>(core::any::type_name_of_val(stage))?;
-        metadata.tries_remaining = None;
-        Ok(())
-    }
-
-    fn progress<'a>(state: &'a S, stage: &ST) -> Result<&'a Self, Error> {
-        state.named_metadata::<Self>(core::any::type_name_of_val(stage))
-    }
-
-    fn progress_mut<'a>(state: &'a mut S, stage: &ST) -> Result<&'a mut Self, Error> {
-        state.named_metadata_mut::<Self>(core::any::type_name_of_val(stage))
-    }
-}
-
-impl RetryProgress {
-    /// Whether we should skip the provided corpus entry.
-    pub fn should_skip<S, ST>(
-        state: &mut S,
-        stage: &ST,
-        corpus_idx: CorpusId,
-    ) -> Result<bool, Error>
-    where
-        S: HasNamedMetadata,
-        ST: RetryingStage,
-    {
-        let progress = Self::progress_mut(state, stage)?;
-        if progress.skipped.contains(&corpus_idx) {
-            return Ok(true);
-        }
-        let remaining = progress.tries_remaining.as_mut().ok_or_else(||
-            Error::illegal_state(
-                "Attempted to check if we should skip a testcase without having initialised the number of tries remaining.",
-            ))?;
-        if *remaining == 0 {
-            progress.skipped.insert(corpus_idx);
-            return Ok(true);
-        }
-        Ok(false)
-    }
-}
-
-/// Trait for types which track the current stage
-pub trait HasCurrentStage {
-    /// Set the current stage; we have started processing this stage
-    fn set_stage(&mut self, idx: usize) -> Result<(), Error>;
-
-    /// Clear the current stage; we are done processing this stage
-    fn clear_stage(&mut self) -> Result<(), Error>;
-
-    /// Fetch the current stage -- typically used after a state recovery or transfer
-    fn current_stage(&self) -> Result<Option<usize>, 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: HasCurrentStage {
-    /// 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>;
-}
-
 #[cfg(test)]
 pub mod test {
-    use alloc::rc::Rc;
+    use core::marker::PhantomData;
-    use core::{cell::RefCell, marker::PhantomData};
 
-    use libafl_bolts::{impl_serdeany, Error};
+    use libafl_bolts::{impl_serdeany, Error, Named};
     use serde::{Deserialize, Serialize};
-    use tuple_list::{tuple_list, tuple_list_type};
 
     use crate::{
-        corpus::{Corpus, Testcase},
+        corpus::{Corpus, HasCurrentCorpusIdx, Testcase},
-        events::NopEventManager,
-        executors::test::NopExecutor,
-        fuzzer::test::NopFuzzer,
         inputs::NopInput,
-        stages::{RetryProgress, RetryingStage, Stage, StageProgress, StagesTuple},
+        stages::{RetryRestartHelper, Stage},
         state::{test::test_std_state, HasCorpus, HasMetadata, State, UsesState},
     };
 
@@ -722,12 +897,6 @@ pub mod test {
         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,
@@ -735,34 +904,35 @@ pub mod test {
 
     impl_serdeany!(TestProgress);
 
-    impl<S, ST> StageProgress<S, ST> for TestProgress
+    impl TestProgress {
+        #[allow(clippy::unnecessary_wraps)]
+        fn restart_progress_should_run<S, ST>(state: &mut S, _stage: &ST) -> Result<bool, Error>
         where
             S: HasMetadata,
         {
-        fn initialize_progress(state: &mut S, _stage: &ST) -> Result<(), Error> {
             // check if we're resuming
-            if !state.has_metadata::<Self>() {
+            let metadata = state.metadata_or_insert_with(|| Self { count: 0 });
-                state.add_metadata(Self { count: 0 });
+
-            }
+            metadata.count += 1;
-            Ok(())
+            assert!(
+                metadata.count == 1,
+                "Test failed; we resumed a succeeded stage!"
+            );
+
+            Ok(true)
         }
 
-        fn clear_progress(state: &mut S, _stage: &ST) -> Result<(), Error> {
+        fn clear_restart_progress<S, ST>(state: &mut S, _stage: &ST) -> Result<(), Error>
-            if state.metadata_map_mut().remove::<Self>().is_none() {
+        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(())
         }
-
-        fn progress<'a>(state: &'a S, _stage: &ST) -> Result<&'a Self, Error> {
-            state.metadata()
-        }
-
-        fn progress_mut<'a>(state: &'a mut S, _stage: &ST) -> Result<&'a mut Self, Error> {
-            state.metadata_mut()
-        }
     }
 
     impl<S> UsesState for ResumeSucceededStage<S>
@@ -779,139 +949,39 @@ pub mod test {
         Z: UsesState,
         Z::State: HasMetadata,
     {
-        type Progress = TestProgress;
-
         fn perform(
             &mut self,
             _fuzzer: &mut Z,
             _executor: &mut E,
-            state: &mut Self::State,
+            _state: &mut Self::State,
             _manager: &mut EM,
         ) -> Result<(), Error> {
-            // metadata is attached by the status
-            let meta = Self::Progress::progress_mut(state, self)?;
-            meta.count += 1;
-            assert!(
-                meta.count == 1,
-                "Test failed; we resumed a succeeded stage!"
-            );
-
             Ok(())
         }
+
+        fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+            TestProgress::restart_progress_should_run(state, self)
         }
 
-    impl<S> UsesState for ResumeFailedStage<S>
+        fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
-    where
+            TestProgress::clear_restart_progress(state, self)
-        S: State,
-    {
-        type State = S;
-    }
-
-    impl<E, EM, Z> Stage<E, EM, Z> for ResumeFailedStage<Z::State>
-    where
-        E: UsesState<State = Z::State>,
-        EM: UsesState<State = Z::State>,
-        Z: UsesState,
-        Z::State: HasMetadata,
-    {
-        type Progress = TestProgress;
-
-        fn perform(
-            &mut self,
-            _fuzzer: &mut Z,
-            _executor: &mut E,
-            state: &mut Self::State,
-            _manager: &mut EM,
-        ) -> Result<(), Error> {
-            // metadata is attached by the status
-            let meta = Self::Progress::progress_mut(state, self)?;
-            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(())
-        }
-    }
-
-    #[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>, S, NopFuzzer<S>>,
-        S: State,
-    {
-        #[cfg(any(not(feature = "serdeany_autoreg"), miri))]
-        unsafe {
-            TestProgress::register();
-        }
-
-        let mut fuzzer = NopFuzzer::new();
-        let mut executor = NopExecutor::new();
-        let mut manager = NopEventManager::new();
-
-        for _ in 0..2 {
-            completed.replace(false);
-            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 test_tries_progress() -> Result<(), Error> {
+        // # Safety
+        // No concurrency per testcase
         #[cfg(any(not(feature = "serdeany_autoreg"), miri))]
         unsafe {
-            RetryProgress::register();
+            RetryRestartHelper::register();
         }
 
         struct StageWithOneTry;
 
-        impl RetryingStage for StageWithOneTry {
+        impl Named for StageWithOneTry {
-            fn max_retries(&self) -> usize {
+            fn name(&self) -> &str {
-                1
+                "TestStage"
             }
         }
 
@@ -920,38 +990,48 @@ pub mod test {
 
         let corpus_idx = state.corpus_mut().add(Testcase::new(NopInput {}))?;
 
+        state.set_corpus_idx(corpus_idx)?;
+
         for _ in 0..10 {
             // used normally, no retries means we never skip
-            RetryProgress::initialize_progress(&mut state, &stage)?;
+            assert!(RetryRestartHelper::restart_progress_should_run(
-            assert!(!RetryProgress::should_skip(&mut state, &stage, corpus_idx)?);
+                &mut state, &stage, 1
-            RetryProgress::clear_progress(&mut state, &stage)?;
+            )?);
+            RetryRestartHelper::clear_restart_progress(&mut state, &stage)?;
         }
 
         for _ in 0..10 {
             // used normally, only one retry means we never skip
-            RetryProgress::initialize_progress(&mut state, &stage)?;
+            assert!(RetryRestartHelper::restart_progress_should_run(
-            assert!(!RetryProgress::should_skip(&mut state, &stage, corpus_idx)?);
+                &mut state, &stage, 2
-            RetryProgress::initialize_progress(&mut state, &stage)?;
+            )?);
-            assert!(!RetryProgress::should_skip(&mut state, &stage, corpus_idx)?);
+            assert!(RetryRestartHelper::restart_progress_should_run(
-            RetryProgress::clear_progress(&mut state, &stage)?;
+                &mut state, &stage, 2
+            )?);
+            RetryRestartHelper::clear_restart_progress(&mut state, &stage)?;
         }
 
-        RetryProgress::initialize_progress(&mut state, &stage)?;
+        assert!(RetryRestartHelper::restart_progress_should_run(
-        assert!(!RetryProgress::should_skip(&mut state, &stage, corpus_idx)?);
+            &mut state, &stage, 2
+        )?);
         // task failed, let's resume
-        RetryProgress::initialize_progress(&mut state, &stage)?;
         // we still have one more try!
-        assert!(!RetryProgress::should_skip(&mut state, &stage, corpus_idx)?);
+        assert!(RetryRestartHelper::restart_progress_should_run(
-        // task failed, let's resume
+            &mut state, &stage, 2
-        RetryProgress::initialize_progress(&mut state, &stage)?;
+        )?);
-        // out of retries, so now we skip
+
-        assert!(RetryProgress::should_skip(&mut state, &stage, corpus_idx)?);
+        // task failed, let's resume
-        RetryProgress::clear_progress(&mut state, &stage)?;
+        // out of retries, so now we skip
+        assert!(!RetryRestartHelper::restart_progress_should_run(
+            &mut state, &stage, 2
+        )?);
+        RetryRestartHelper::clear_restart_progress(&mut state, &stage)?;
 
-        RetryProgress::initialize_progress(&mut state, &stage)?;
         // we previously exhausted this testcase's retries, so we skip
-        assert!(RetryProgress::should_skip(&mut state, &stage, corpus_idx)?);
+        assert!(!RetryRestartHelper::restart_progress_should_run(
-        RetryProgress::clear_progress(&mut state, &stage)?;
+            &mut state, &stage, 2
+        )?);
+        RetryRestartHelper::clear_restart_progress(&mut state, &stage)?;
 
         Ok(())
     }

libafl/src/stages/mutational.rs

@@ -1,9 +1,9 @@
 //| The [`MutationalStage`] is the default stage used during fuzzing.
 //! For the current input, it will perform a range of random mutations, and then run them in the executor.
 
-use core::marker::PhantomData;
+use core::{any::type_name, marker::PhantomData};
 
-use libafl_bolts::rands::Rand;
+use libafl_bolts::{rands::Rand, Named};
 
 use crate::{
     corpus::{Corpus, CorpusId, HasCurrentCorpusIdx, Testcase},
@@ -11,9 +11,9 @@ use crate::{
     inputs::Input,
     mark_feature_time,
     mutators::{MultiMutator, MutationResult, Mutator},
-    stages::Stage,
+    stages::{ExecutionCountRestartHelper, RetryRestartHelper, Stage},
     start_timer,
-    state::{HasCorpus, HasRand, UsesState},
+    state::{HasCorpus, HasExecutions, HasMetadata, HasNamedMetadata, HasRand, UsesState},
     Error,
 };
 #[cfg(feature = "introspection")]
@@ -104,7 +104,10 @@ where
     fn mutator_mut(&mut self) -> &mut M;
 
     /// Gets the number of iterations this mutator should run for.
-    fn iterations(&self, state: &mut Z::State, corpus_idx: CorpusId) -> Result<u64, Error>;
+    fn iterations(&self, state: &mut Z::State) -> Result<u64, Error>;
+
+    /// Gets the number of executions this mutator already did since it got first called in this fuzz round.
+    fn execs_since_progress_start(&mut self, state: &mut Z::State) -> Result<u64, Error>;
 
     /// Runs this (mutational) stage for the given testcase
     #[allow(clippy::cast_possible_wrap)] // more than i32 stages on 32 bit system - highly unlikely...
@@ -121,7 +124,7 @@ where
             ));
         };
 
-        let num = self.iterations(state, corpus_idx)?;
+        let num = self.iterations(state)? - self.execs_since_progress_start(state)?;
 
         start_timer!(state);
         let mut testcase = state.corpus().get(corpus_idx)?.borrow_mut();
@@ -163,8 +166,12 @@ pub static DEFAULT_MUTATIONAL_MAX_ITERATIONS: u64 = 128;
 /// The default mutational stage
 #[derive(Clone, Debug)]
 pub struct StdMutationalStage<E, EM, I, M, Z> {
+    /// The mutator(s) to use
     mutator: M,
+    /// The maximum amount of iterations we should do each round
     max_iterations: u64,
+    /// The progress helper for this mutational stage
+    restart_helper: ExecutionCountRestartHelper,
     #[allow(clippy::type_complexity)]
     phantom: PhantomData<(E, EM, I, Z)>,
 }
@@ -175,7 +182,7 @@ where
     EM: UsesState<State = Z::State>,
     M: Mutator<I, Z::State>,
     Z: Evaluator<E, EM>,
-    Z::State: HasCorpus + HasRand,
+    Z::State: HasCorpus + HasRand + HasExecutions + HasMetadata,
     I: MutatedTransform<Self::Input, Self::State> + Clone,
 {
     /// The mutator, added to this stage
@@ -191,9 +198,13 @@ where
     }
 
     /// Gets the number of iterations as a random number
-    fn iterations(&self, state: &mut Z::State, _corpus_idx: CorpusId) -> Result<u64, Error> {
+    fn iterations(&self, state: &mut Z::State) -> Result<u64, Error> {
         Ok(1 + state.rand_mut().below(self.max_iterations))
     }
+
+    fn execs_since_progress_start(&mut self, state: &mut <Z>::State) -> Result<u64, Error> {
+        self.restart_helper.execs_since_progress_start(state)
+    }
 }
 
 impl<E, EM, I, M, Z> UsesState for StdMutationalStage<E, EM, I, M, Z>
@@ -213,11 +224,9 @@ where
     EM: UsesState<State = Z::State>,
     M: Mutator<I, Z::State>,
     Z: Evaluator<E, EM>,
-    Z::State: HasCorpus + HasRand,
+    Z::State: HasCorpus + HasRand + HasMetadata + HasExecutions,
     I: MutatedTransform<Self::Input, Self::State> + Clone,
 {
-    type Progress = (); // TODO should this stage be resumed?
-
     #[inline]
     #[allow(clippy::let_and_return)]
     fn perform(
@@ -234,6 +243,14 @@ where
 
         ret
     }
+
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        self.restart_helper.restart_progress_should_run(state)
+    }
+
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        self.restart_helper.clear_restart_progress(state)
+    }
 }
 
 impl<E, EM, M, Z> StdMutationalStage<E, EM, Z::Input, M, Z>
@@ -273,12 +290,13 @@ where
         Self {
             mutator,
             max_iterations,
+            restart_helper: ExecutionCountRestartHelper::default(),
             phantom: PhantomData,
         }
     }
 }
 
-/// The default mutational stage
+/// A mutational stage that operates on multiple inputs, as returned by [`MultiMutator::multi_mutate`].
 #[derive(Clone, Debug)]
 pub struct MultiMutationalStage<E, EM, I, M, Z> {
     mutator: M,
@@ -297,16 +315,39 @@ where
     type State = Z::State;
 }
 
-impl<E, EM, I, M, Z> Stage<E, EM, Z> for MultiMutationalStage<E, EM, I, M, Z>
+impl<E, EM, I, M, Z> Named for MultiMutationalStage<E, EM, I, M, Z>
 where
     E: UsesState<State = Z::State>,
     EM: UsesState<State = Z::State>,
     M: MultiMutator<I, Z::State>,
     Z: Evaluator<E, EM>,
     Z::State: HasCorpus + HasRand,
+{
+    fn name(&self) -> &str {
+        type_name::<Self>()
+    }
+}
+
+impl<E, EM, I, M, Z> Stage<E, EM, Z> for MultiMutationalStage<E, EM, I, M, Z>
+where
+    E: UsesState<State = Z::State>,
+    EM: UsesState<State = Z::State>,
+    M: MultiMutator<I, Z::State>,
+    Z: Evaluator<E, EM>,
+    Z::State: HasCorpus + HasRand + HasNamedMetadata,
     I: MutatedTransform<Self::Input, Self::State> + Clone,
 {
-    type Progress = (); // TODO implement resume
+    #[inline]
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        // TODO: add proper crash/timeout handling
+        // For now, Make sure we don't get stuck crashing on a single testcase
+        RetryRestartHelper::restart_progress_should_run(state, self, 3)
+    }
+
+    #[inline]
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        RetryRestartHelper::clear_restart_progress(state, self)
+    }
 
     #[inline]
     #[allow(clippy::let_and_return)]
@@ -353,7 +394,7 @@ where
     Z: Evaluator<E, EM>,
     Z::State: HasCorpus + HasRand,
 {
-    /// Creates a new default mutational stage
+    /// Creates a new [`MultiMutationalStage`]
     pub fn new(mutator: M) -> Self {
         Self::transforming(mutator)
     }

libafl/src/stages/power.rs

@@ -3,20 +3,22 @@
 use core::{fmt::Debug, marker::PhantomData};
 
 use crate::{
-    corpus::{Corpus, CorpusId},
     executors::{Executor, HasObservers},
     fuzzer::Evaluator,
     mutators::Mutator,
     schedulers::{testcase_score::CorpusPowerTestcaseScore, TestcaseScore},
-    stages::{mutational::MutatedTransform, MutationalStage, Stage},
+    stages::{mutational::MutatedTransform, ExecutionCountRestartHelper, MutationalStage, Stage},
-    state::{HasCorpus, HasMetadata, HasRand, UsesState},
+    state::{HasCorpus, HasCurrentTestcase, HasExecutions, HasMetadata, HasRand, UsesState},
     Error,
 };
 
 /// The mutational stage using power schedules
 #[derive(Clone, Debug)]
 pub struct PowerMutationalStage<E, F, EM, I, M, Z> {
+    /// The mutators we use
     mutator: M,
+    /// Helper for restarts
+    restart_helper: ExecutionCountRestartHelper,
     #[allow(clippy::type_complexity)]
     phantom: PhantomData<(E, F, EM, I, Z)>,
 }
@@ -34,7 +36,7 @@ where
     EM: UsesState<State = E::State>,
     F: TestcaseScore<E::State>,
     M: Mutator<I, E::State>,
-    E::State: HasCorpus + HasMetadata + HasRand,
+    E::State: HasCorpus + HasMetadata + HasRand + HasExecutions,
     Z: Evaluator<E, EM, State = E::State>,
     I: MutatedTransform<E::Input, E::State> + Clone,
 {
@@ -52,13 +54,17 @@ where
 
     /// Gets the number of iterations as a random number
     #[allow(clippy::cast_sign_loss)]
-    fn iterations(&self, state: &mut E::State, corpus_idx: CorpusId) -> Result<u64, Error> {
+    fn iterations(&self, state: &mut E::State) -> Result<u64, Error> {
         // Update handicap
-        let mut testcase = state.corpus().get(corpus_idx)?.borrow_mut();
+        let mut testcase = state.current_testcase_mut()?;
-        let score = F::compute(state, &mut *testcase)? as u64;
+        let score = F::compute(state, &mut testcase)? as u64;
 
         Ok(score)
     }
+
+    fn execs_since_progress_start(&mut self, state: &mut <Z>::State) -> Result<u64, Error> {
+        self.restart_helper.execs_since_progress_start(state)
+    }
 }
 
 impl<E, F, EM, I, M, Z> Stage<E, EM, Z> for PowerMutationalStage<E, F, EM, I, M, Z>
@@ -67,12 +73,10 @@ where
     EM: UsesState<State = E::State>,
     F: TestcaseScore<E::State>,
     M: Mutator<I, E::State>,
-    E::State: HasCorpus + HasMetadata + HasRand,
+    E::State: HasCorpus + HasMetadata + HasRand + HasExecutions,
     Z: Evaluator<E, EM, State = E::State>,
     I: MutatedTransform<E::Input, E::State> + Clone,
 {
-    type Progress = (); // TODO should we resume this stage?
-
     #[inline]
     #[allow(clippy::let_and_return)]
     fn perform(
@@ -85,6 +89,14 @@ where
         let ret = self.perform_mutational(fuzzer, executor, state, manager);
         ret
     }
+
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        self.restart_helper.restart_progress_should_run(state)
+    }
+
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        self.restart_helper.clear_restart_progress(state)
+    }
 }
 
 impl<E, F, EM, M, Z> PowerMutationalStage<E, F, EM, E::Input, M, Z>
@@ -116,6 +128,7 @@ where
         Self {
             mutator,
             phantom: PhantomData,
+            restart_helper: ExecutionCountRestartHelper::default(),
         }
     }
 }

libafl/src/stages/stats.rs

@@ -62,8 +62,6 @@ where
     Z: UsesState<State = E::State>,
     E::State: HasImported + HasCorpus + HasMetadata,
 {
-    type Progress = (); // this stage does not require resume
-
     fn perform(
         &mut self,
         _fuzzer: &mut Z,
@@ -133,6 +131,18 @@ where
 
         Ok(())
     }
+
+    #[inline]
+    fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
+        // Not running the target so we wont't crash/timeout and, hence, don't need to restore anything
+        Ok(true)
+    }
+
+    #[inline]
+    fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> {
+        // Not running the target so we wont't crash/timeout and, hence, don't need to restore anything
+        Ok(())
+    }
 }
 
 impl<E, EM, Z> AflStatsStage<E, EM, Z>

libafl/src/stages/string.rs

@@ -104,8 +104,6 @@ where
     EM: UsesState<State = S>,
     Z: UsesState<State = S>,
 {
-    type Progress = (); // this stage does not need to be resumed
-
     fn perform(
         &mut self,
         _fuzzer: &mut Z,
@@ -132,4 +130,16 @@ where
 
         Ok(())
     }
+
+    #[inline]
+    fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
+        // Stage does not run the target. No reset helper needed.
+        Ok(true)
+    }
+
+    #[inline]
+    fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> {
+        // Stage does not run the target. No reset helper needed.
+        Ok(())
+    }
 }

libafl/src/stages/sync.rs

@@ -7,7 +7,7 @@ use std::{
     time::SystemTime,
 };
 
-use libafl_bolts::{current_time, shmem::ShMemProvider};
+use libafl_bolts::{current_time, shmem::ShMemProvider, Named};
 use serde::{Deserialize, Serialize};
 
 #[cfg(feature = "introspection")]
@@ -18,8 +18,8 @@ use crate::{
     executors::{Executor, ExitKind, HasObservers},
     fuzzer::{Evaluator, EvaluatorObservers, ExecutionProcessor},
     inputs::{Input, InputConverter, UsesInput},
-    stages::Stage,
+    stages::{RetryRestartHelper, Stage},
-    state::{HasCorpus, HasExecutions, HasMetadata, HasRand, State, UsesState},
+    state::{HasCorpus, HasExecutions, HasMetadata, HasNamedMetadata, HasRand, State, UsesState},
     Error,
 };
 
@@ -59,16 +59,23 @@ where
     type State = E::State;
 }
 
+impl<CB, E, EM, Z> Named for SyncFromDiskStage<CB, E, EM, Z>
+where
+    E: UsesState,
+{
+    fn name(&self) -> &str {
+        self.sync_dir.to_str().unwrap()
+    }
+}
+
 impl<CB, E, EM, Z> Stage<E, EM, Z> for SyncFromDiskStage<CB, E, EM, Z>
 where
     CB: FnMut(&mut Z, &mut Z::State, &Path) -> Result<<Z::State as UsesInput>::Input, Error>,
     E: UsesState<State = Z::State>,
     EM: UsesState<State = Z::State>,
     Z: Evaluator<E, EM>,
-    Z::State: HasCorpus + HasRand + HasMetadata,
+    Z::State: HasCorpus + HasRand + HasMetadata + HasNamedMetadata,
 {
-    type Progress = (); // TODO load from directory should be resumed
-
     #[inline]
     fn perform(
         &mut self,
@@ -103,6 +110,18 @@ where
 
         Ok(())
     }
+
+    #[inline]
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        // TODO: Needs proper crash handling for when an imported testcase crashes
+        // For now, Make sure we don't get stuck crashing on this testcase
+        RetryRestartHelper::restart_progress_should_run(state, self, 3)
+    }
+
+    #[inline]
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        RetryRestartHelper::clear_restart_progress(state, self)
+    }
 }
 
 impl<CB, E, EM, Z> SyncFromDiskStage<CB, E, EM, Z>
@@ -254,8 +273,6 @@ where
     ICB: InputConverter<From = DI, To = S::Input>,
     DI: Input,
 {
-    type Progress = (); // TODO this should be resumed in the case that a testcase causes a crash
-
     #[inline]
     fn perform(
         &mut self,
@@ -312,6 +329,18 @@ where
         state.introspection_monitor_mut().finish_stage();
         Ok(())
     }
+
+    #[inline]
+    fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
+        // No restart handling needed - does not execute the target.
+        Ok(true)
+    }
+
+    #[inline]
+    fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> {
+        // Not needed - does not execute the target.
+        Ok(())
+    }
 }
 
 impl<IC, ICB, DI, S, SP> SyncFromBrokerStage<IC, ICB, DI, S, SP>

libafl/src/stages/tmin.rs

@@ -7,7 +7,7 @@ use ahash::RandomState;
 use libafl_bolts::{HasLen, Named};
 
 use crate::{
-    corpus::{Corpus, CorpusId, HasCurrentCorpusIdx, Testcase},
+    corpus::{Corpus, HasCurrentCorpusIdx, Testcase},
     events::EventFirer,
     executors::{Executor, ExitKind, HasObservers},
     feedbacks::{Feedback, FeedbackFactory, HasObserverName},
@@ -16,9 +16,12 @@ use crate::{
     mutators::{MutationResult, Mutator},
     observers::{MapObserver, ObserversTuple},
     schedulers::{RemovableScheduler, Scheduler},
-    stages::Stage,
+    stages::{ExecutionCountRestartHelper, Stage},
     start_timer,
-    state::{HasCorpus, HasExecutions, HasMaxSize, HasSolutions, State, UsesState},
+    state::{
+        HasCorpus, HasCurrentTestcase, HasExecutions, HasMaxSize, HasMetadata, HasSolutions, State,
+        UsesState,
+    },
     Error, ExecutesInput, ExecutionProcessor, HasFeedback, HasScheduler,
 };
 #[cfg(feature = "introspection")]
@@ -51,7 +54,7 @@ where
     fn mutator_mut(&mut self) -> &mut M;
 
     /// Gets the number of iterations this mutator should run for.
-    fn iterations(&self, state: &mut CS::State, corpus_idx: CorpusId) -> Result<usize, Error>;
+    fn iterations(&self, state: &mut CS::State) -> Result<usize, Error>;
 
     /// Runs this (mutational) stage for new objectives
     #[allow(clippy::cast_possible_wrap)] // more than i32 stages on 32 bit system - highly unlikely...
@@ -70,10 +73,11 @@ where
 
         let orig_max_size = state.max_size();
         // basically copy-pasted from mutational.rs
-        let num = self.iterations(state, base_corpus_idx)?;
+        let num = self.iterations(state)?
+            - usize::try_from(self.execs_since_progress_start(state)?).unwrap();
 
         start_timer!(state);
-        let mut base = state.corpus().cloned_input_for_id(base_corpus_idx)?;
+        let mut base = state.current_input_cloned()?;
         let base_hash = RandomState::with_seeds(0, 0, 0, 0).hash_one(&base);
         mark_feature_time!(state, PerfFeature::GetInputFromCorpus);
 
@@ -173,14 +177,22 @@ where
 
         Ok(())
     }
+
+    /// Gets the number of executions this mutator already did since it got first called in this fuzz round.
+    fn execs_since_progress_start(&mut self, state: &mut Z::State) -> Result<u64, Error>;
 }
 
 /// The default corpus entry minimising mutational stage
 #[derive(Clone, Debug)]
 pub struct StdTMinMutationalStage<CS, E, EM, F1, F2, FF, M, OT, Z> {
+    /// The mutator(s) this stage uses
     mutator: M,
+    /// The factory
     factory: FF,
+    /// The runs (=iterations) we are supposed to do
     runs: usize,
+    /// The progress helper for this stage, keeping track of resumes after timeouts/crashes
+    restart_helper: ExecutionCountRestartHelper,
     #[allow(clippy::type_complexity)]
     phantom: PhantomData<(CS, E, EM, F1, F2, OT, Z)>,
 }
@@ -200,7 +212,7 @@ impl<CS, E, EM, F1, F2, FF, M, OT, Z> Stage<E, EM, Z>
     for StdTMinMutationalStage<CS, E, EM, F1, F2, FF, M, OT, Z>
 where
     CS: Scheduler + RemovableScheduler,
-    CS::State: HasCorpus + HasSolutions + HasExecutions + HasMaxSize + HasCorpus,
+    CS::State: HasCorpus + HasSolutions + HasExecutions + HasMaxSize + HasCorpus + HasMetadata,
     <CS::State as UsesInput>::Input: HasLen + Hash,
     E: Executor<EM, Z> + HasObservers<Observers = OT, State = CS::State>,
     EM: EventFirer<State = CS::State>,
@@ -214,8 +226,6 @@ where
         + HasFeedback<Feedback = F1>
         + HasScheduler<Scheduler = CS>,
 {
-    type Progress = (); // TODO this stage desperately needs a resume
-
     fn perform(
         &mut self,
         fuzzer: &mut Z,
@@ -230,6 +240,14 @@ where
 
         Ok(())
     }
+
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        self.restart_helper.restart_progress_should_run(state)
+    }
+
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        self.restart_helper.clear_restart_progress(state)
+    }
 }
 
 impl<CS, E, EM, F1, F2, FF, M, OT, Z> FeedbackFactory<F2, Z::State, OT>
@@ -256,7 +274,7 @@ where
     <CS::State as UsesInput>::Input: HasLen + Hash,
     M: Mutator<CS::Input, CS::State>,
     OT: ObserversTuple<CS::State>,
-    CS::State: HasCorpus + HasSolutions + HasExecutions + HasMaxSize,
+    CS::State: HasCorpus + HasSolutions + HasExecutions + HasMaxSize + HasMetadata,
     Z: ExecutionProcessor<OT, State = CS::State>
         + ExecutesInput<E, EM>
         + HasFeedback<Feedback = F1>
@@ -275,9 +293,13 @@ where
     }
 
     /// Gets the number of iterations from a fixed number of runs
-    fn iterations(&self, _state: &mut CS::State, _corpus_idx: CorpusId) -> Result<usize, Error> {
+    fn iterations(&self, _state: &mut CS::State) -> Result<usize, Error> {
         Ok(self.runs)
     }
+
+    fn execs_since_progress_start(&mut self, state: &mut <Z>::State) -> Result<u64, Error> {
+        self.restart_helper.execs_since_progress_start(state)
+    }
 }
 
 impl<CS, E, EM, F1, F2, FF, M, OT, Z> StdTMinMutationalStage<CS, E, EM, F1, F2, FF, M, OT, Z>
@@ -287,12 +309,13 @@ where
     Z: ExecutionProcessor<OT, State = CS::State>,
     CS::State: HasCorpus,
 {
-    /// Creates a new minimising mutational stage that will minimize provided corpus entries
+    /// Creates a new minimizing mutational stage that will minimize provided corpus entries
     pub fn new(mutator: M, factory: FF, runs: usize) -> Self {
         Self {
             mutator,
             factory,
             runs,
+            restart_helper: ExecutionCountRestartHelper::default(),
             phantom: PhantomData,
         }
     }

libafl/src/stages/tracing.rs

@@ -2,14 +2,15 @@
 
 use core::{fmt::Debug, marker::PhantomData};
 
+use libafl_bolts::Named;
+
 use crate::{
-    corpus::{Corpus, CorpusId, HasCurrentCorpusIdx},
     executors::{Executor, HasObservers, ShadowExecutor},
     mark_feature_time,
     observers::ObserversTuple,
-    stages::{RetryProgress, RetryingStage, Stage},
+    stages::{RetryRestartHelper, Stage},
     start_timer,
-    state::{HasCorpus, HasExecutions, HasNamedMetadata, State, UsesState},
+    state::{HasCorpus, HasCurrentTestcase, HasExecutions, HasNamedMetadata, State, UsesState},
     Error,
 };
 #[cfg(feature = "introspection")]
@@ -38,18 +39,17 @@ where
     EM: UsesState<State = TE::State>,
     Z: UsesState<State = TE::State>,
 {
-    /// Perform tracing on the given [`CorpusId`]. Useful for if wrapping [`TracingStage`] with your
+    /// Perform tracing on the given `CorpusId`. Useful for if wrapping [`TracingStage`] with your
-    /// own stage and you need to manage [`super::StageProgress`] differently; see
+    /// own stage and you need to manage [`super::NestedStageRestartHelper`] differently; see
     /// [`super::ConcolicTracingStage`]'s implementation as an example of usage.
     pub fn trace(
         &mut self,
         fuzzer: &mut Z,
         state: &mut TE::State,
         manager: &mut EM,
-        corpus_idx: CorpusId,
     ) -> Result<(), Error> {
         start_timer!(state);
-        let input = state.corpus().cloned_input_for_id(corpus_idx)?;
+        let input = state.current_input_cloned()?;
 
         mark_feature_time!(state, PerfFeature::GetInputFromCorpus);
 
@@ -85,8 +85,6 @@ where
     EM: UsesState<State = TE::State>,
     Z: UsesState<State = TE::State>,
 {
-    type Progress = RetryProgress;
-
     #[inline]
     fn perform(
         &mut self,
@@ -95,24 +93,21 @@ where
         state: &mut TE::State,
         manager: &mut EM,
     ) -> Result<(), Error> {
-        let Some(corpus_idx) = state.current_corpus_idx()? else {
+        self.trace(fuzzer, state, manager)
-            return Err(Error::illegal_state(
-                "state is not currently processing a corpus index",
-            ));
-        };
-        if Self::Progress::should_skip(state, self, corpus_idx)? {
-            return Ok(());
     }
 
-        self.trace(fuzzer, state, manager, corpus_idx)?;
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        RetryRestartHelper::restart_progress_should_run(state, self, self.max_retries)
+    }
 
-        Ok(())
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        RetryRestartHelper::clear_restart_progress(state, self)
     }
 }
 
-impl<EM, TE, Z> RetryingStage for TracingStage<EM, TE, Z> {
+impl<EM, TE, Z> Named for TracingStage<EM, TE, Z> {
-    fn max_retries(&self) -> usize {
+    fn name(&self) -> &str {
-        self.max_retries
+        "TracingStage"
     }
 }
 
@@ -160,6 +155,15 @@ where
     type State = E::State;
 }
 
+impl<E, EM, SOT, Z> Named for ShadowTracingStage<E, EM, SOT, Z>
+where
+    E: UsesState,
+{
+    fn name(&self) -> &str {
+        "ShadowTracingStage"
+    }
+}
+
 impl<E, EM, SOT, Z> Stage<ShadowExecutor<E, SOT>, EM, Z> for ShadowTracingStage<E, EM, SOT, Z>
 where
     E: Executor<EM, Z> + HasObservers,
@@ -168,8 +172,6 @@ where
     Z: UsesState<State = E::State>,
     E::State: State + HasExecutions + HasCorpus + HasNamedMetadata + Debug,
 {
-    type Progress = RetryProgress;
-
     #[inline]
     fn perform(
         &mut self,
@@ -178,17 +180,8 @@ where
         state: &mut E::State,
         manager: &mut EM,
     ) -> Result<(), Error> {
-        let Some(corpus_idx) = state.current_corpus_idx()? else {
-            return Err(Error::illegal_state(
-                "state is not currently processing a corpus index",
-            ));
-        };
-        if Self::Progress::should_skip(state, self, corpus_idx)? {
-            return Ok(());
-        }
-
         start_timer!(state);
-        let input = state.corpus().cloned_input_for_id(corpus_idx)?;
+        let input = state.current_input_cloned()?;
 
         mark_feature_time!(state, PerfFeature::GetInputFromCorpus);
 
@@ -216,11 +209,13 @@ where
 
         Ok(())
     }
-}
 
-impl<E, EM, SOT, Z> RetryingStage for ShadowTracingStage<E, EM, SOT, Z> {
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
-    fn max_retries(&self) -> usize {
+        RetryRestartHelper::restart_progress_should_run(state, self, self.max_retries)
-        self.max_retries
+    }
+
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        RetryRestartHelper::clear_restart_progress(state, self)
     }
 }
 

libafl/src/stages/tuneable.rs

@@ -7,15 +7,15 @@ use libafl_bolts::{current_time, impl_serdeany, rands::Rand};
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    corpus::{Corpus, CorpusId, HasCurrentCorpusIdx},
+    corpus::{Corpus, HasCurrentCorpusIdx},
     mark_feature_time,
     mutators::{MutationResult, Mutator},
     stages::{
         mutational::{MutatedTransform, MutatedTransformPost, DEFAULT_MUTATIONAL_MAX_ITERATIONS},
-        MutationalStage, Stage,
+        ExecutionCountRestartHelper, MutationalStage, Stage,
     },
     start_timer,
-    state::{HasCorpus, HasMetadata, HasNamedMetadata, HasRand, UsesState},
+    state::{HasCorpus, HasExecutions, HasMetadata, HasNamedMetadata, HasRand, UsesState},
     Error, Evaluator,
 };
 #[cfg(feature = "introspection")]
@@ -150,8 +150,12 @@ where
 /// A [`crate::stages::MutationalStage`] where the mutator iteration can be tuned at runtime
 #[derive(Clone, Debug)]
 pub struct TuneableMutationalStage<E, EM, I, M, Z> {
+    /// The mutator we use
     mutator: M,
+    /// The name of this stage
     name: String,
+    /// The progress helper we use to keep track of progress across restarts
+    restart_helper: ExecutionCountRestartHelper,
     phantom: PhantomData<(E, EM, I, Z)>,
 }
 
@@ -161,7 +165,7 @@ where
     EM: UsesState<State = Z::State>,
     M: Mutator<I, Z::State>,
     Z: Evaluator<E, EM>,
-    Z::State: HasCorpus + HasRand + HasNamedMetadata + HasMetadata,
+    Z::State: HasCorpus + HasRand + HasNamedMetadata + HasMetadata + HasExecutions,
     I: MutatedTransform<Z::Input, Z::State> + Clone,
 {
     /// Runs this (mutational) stage for the given `testcase`
@@ -222,7 +226,7 @@ where
             }
             (None, None) => {
                 // fall back to random
-                let iters = self.iterations(state, corpus_idx)?;
+                let iters = self.iterations(state)? - self.execs_since_progress_start(state)?;
                 for i in 1..=iters {
                     self.perform_mutation(fuzzer, executor, state, manager, &input, i)?;
                 }
@@ -245,12 +249,16 @@ where
 
     /// Gets the number of iterations as a random number
     #[allow(clippy::cast_possible_truncation)]
-    fn iterations(&self, state: &mut Z::State, _corpus_idx: CorpusId) -> Result<u64, Error> {
+    fn iterations(&self, state: &mut Z::State) -> Result<u64, Error> {
         Ok(
             // fall back to random
             1 + state.rand_mut().below(DEFAULT_MUTATIONAL_MAX_ITERATIONS),
         )
     }
+
+    fn execs_since_progress_start(&mut self, state: &mut <Z>::State) -> Result<u64, Error> {
+        self.restart_helper.execs_since_progress_start(state)
+    }
 }
 
 impl<E, EM, I, M, Z> UsesState for TuneableMutationalStage<E, EM, I, M, Z>
@@ -259,7 +267,7 @@ where
     EM: UsesState<State = Z::State>,
     M: Mutator<I, Z::State>,
     Z: Evaluator<E, EM>,
-    Z::State: HasCorpus + HasRand,
+    Z::State: HasCorpus + HasRand + HasExecutions,
     I: MutatedTransform<Z::Input, Z::State> + Clone,
 {
     type State = Z::State;
@@ -271,11 +279,9 @@ where
     EM: UsesState<State = Z::State>,
     M: Mutator<I, Z::State>,
     Z: Evaluator<E, EM>,
-    Z::State: HasCorpus + HasRand + HasNamedMetadata + HasMetadata,
+    Z::State: HasCorpus + HasRand + HasNamedMetadata + HasMetadata + HasExecutions,
     I: MutatedTransform<Z::Input, Z::State> + Clone,
 {
-    type Progress = (); // TODO should this stage be resumed?
-
     #[inline]
     #[allow(clippy::let_and_return)]
     fn perform(
@@ -292,6 +298,14 @@ where
 
         ret
     }
+
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        self.restart_helper.restart_progress_should_run(state)
+    }
+
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        self.restart_helper.clear_restart_progress(state)
+    }
 }
 
 impl<E, EM, I, M, Z> TuneableMutationalStage<E, EM, I, M, Z>
@@ -300,7 +314,7 @@ where
     EM: UsesState<State = Z::State>,
     M: Mutator<I, Z::State>,
     Z: Evaluator<E, EM>,
-    Z::State: HasCorpus + HasRand + HasNamedMetadata + HasMetadata,
+    Z::State: HasCorpus + HasRand + HasNamedMetadata + HasMetadata + HasExecutions,
     I: MutatedTransform<Z::Input, Z::State> + Clone,
 {
     /// Creates a new default tuneable mutational stage
@@ -473,12 +487,11 @@ where
     /// Creates a new tranforming mutational stage
     #[must_use]
     pub fn transforming(state: &mut Z::State, mutator: M, name: &str) -> Self {
-        if !state.has_named_metadata::<TuneableMutationalStageMetadata>(name) {
+        let _ = state.named_metadata_or_insert_with(name, TuneableMutationalStageMetadata::default);
-            state.add_named_metadata(TuneableMutationalStageMetadata::default(), name);
-        }
         Self {
             mutator,
             name: name.to_string(),
+            restart_helper: ExecutionCountRestartHelper::default(),
             phantom: PhantomData,
         }
     }

libafl/src/state/mod.rs

@@ -1,7 +1,8 @@
 //! The fuzzer, and state are the core pieces of every good fuzzer
 
-use alloc::vec::Vec;
+use alloc::{boxed::Box, vec::Vec};
 use core::{
+    borrow::BorrowMut,
     cell::{Ref, RefMut},
     fmt::Debug,
     marker::PhantomData,
@@ -170,7 +171,27 @@ pub trait HasMetadata {
         self.metadata_map_mut().insert(meta);
     }
 
+    /// Gets metadata, or inserts it using the given construction function `default`
+    fn metadata_or_insert_with<M>(&mut self, default: impl FnOnce() -> M) -> &mut M
+    where
+        M: SerdeAny,
+    {
+        self.metadata_map_mut().or_insert_with::<M>(default)
+    }
+
+    /// Remove a metadata from the metadata map
+    #[inline]
+    fn remove_metadata<M>(&mut self) -> Option<Box<M>>
+    where
+        M: SerdeAny,
+    {
+        self.metadata_map_mut().remove::<M>()
+    }
+
     /// Check for a metadata
+    ///
+    /// # Note
+    /// For performance reasons, you likely want to use [`Self::metadata_or_insert_with`] instead
     #[inline]
     fn has_metadata<M>(&self) -> bool
     where
@@ -211,14 +232,39 @@ pub trait HasNamedMetadata {
 
     /// Add a metadata to the metadata map
     #[inline]
-    fn add_named_metadata<M>(&mut self, meta: M, name: &str)
+    fn add_named_metadata<M>(&mut self, name: &str, meta: M)
     where
         M: SerdeAny,
     {
-        self.named_metadata_map_mut().insert(meta, name);
+        self.named_metadata_map_mut().insert(name, meta);
+    }
+
+    /// Add a metadata to the metadata map
+    #[inline]
+    fn remove_named_metadata<M>(&mut self, name: &str) -> Option<Box<M>>
+    where
+        M: SerdeAny,
+    {
+        self.named_metadata_map_mut().remove::<M>(name)
+    }
+
+    /// Gets metadata, or inserts it using the given construction function `default`
+    fn named_metadata_or_insert_with<M>(
+        &mut self,
+        name: &str,
+        default: impl FnOnce() -> M,
+    ) -> &mut M
+    where
+        M: SerdeAny,
+    {
+        self.named_metadata_map_mut()
+            .or_insert_with::<M>(name, default)
     }
 
     /// Check for a metadata
+    ///
+    /// # Note
+    /// You likely want to use [`Self::named_metadata_or_insert_with`] for performance reasons.
     #[inline]
     fn has_named_metadata<M>(&self, name: &str) -> bool
     where
@@ -255,10 +301,10 @@ pub trait HasNamedMetadata {
 /// Trait for the execution counter
 pub trait HasExecutions {
     /// The executions counter
-    fn executions(&self) -> &usize;
+    fn executions(&self) -> &u64;
 
     /// The executions counter (mutable)
-    fn executions_mut(&mut self) -> &mut usize;
+    fn executions_mut(&mut self) -> &mut u64;
 }
 
 /// Trait for some stats of AFL
@@ -301,7 +347,7 @@ pub struct StdState<I, C, R, SC> {
     /// RNG instance
     rand: R,
     /// How many times the executor ran the harness/target
-    executions: usize,
+    executions: u64,
     /// At what time the fuzzing started
     start_time: Duration,
     /// the number of new paths that imported from other fuzzers
@@ -406,7 +452,10 @@ where
     R: Rand,
 {
     /// To get the testcase
-    fn testcase(&self, id: CorpusId) -> Result<Ref<Testcase<<Self as UsesInput>::Input>>, Error> {
+    fn testcase(
+        &self,
+        id: CorpusId,
+    ) -> Result<Ref<'_, Testcase<<Self as UsesInput>::Input>>, Error> {
         Ok(self.corpus().get(id)?.borrow())
     }
 
@@ -414,7 +463,7 @@ where
     fn testcase_mut(
         &self,
         id: CorpusId,
-    ) -> Result<RefMut<Testcase<<Self as UsesInput>::Input>>, Error> {
+    ) -> Result<RefMut<'_, Testcase<<Self as UsesInput>::Input>>, Error> {
         Ok(self.corpus().get(id)?.borrow_mut())
     }
 }
@@ -470,13 +519,13 @@ impl<I, C, R, SC> HasNamedMetadata for StdState<I, C, R, SC> {
 impl<I, C, R, SC> HasExecutions for StdState<I, C, R, SC> {
     /// The executions counter
     #[inline]
-    fn executions(&self) -> &usize {
+    fn executions(&self) -> &u64 {
         &self.executions
     }
 
     /// The executions counter (mutable)
     #[inline]
-    fn executions_mut(&mut self) -> &mut usize {
+    fn executions_mut(&mut self) -> &mut u64 {
         &mut self.executions
     }
 }
@@ -549,6 +598,64 @@ impl<I, C, R, SC> HasCurrentCorpusIdx for StdState<I, C, R, SC> {
     }
 }
 
+/// Has information about the current [`Testcase`] we are fuzzing
+pub trait HasCurrentTestcase<I>
+where
+    I: Input,
+{
+    /// Gets the current [`Testcase`] we are fuzzing
+    ///
+    /// Will return [`Error::key_not_found`] if no `corpus_idx` is currently set.
+    fn current_testcase(&self) -> Result<Ref<'_, Testcase<I>>, Error>;
+    //fn current_testcase(&self) -> Result<&Testcase<I>, Error>;
+
+    /// Gets the current [`Testcase`] we are fuzzing (mut)
+    ///
+    /// Will return [`Error::key_not_found`] if no `corpus_idx` is currently set.
+    fn current_testcase_mut(&self) -> Result<RefMut<'_, Testcase<I>>, Error>;
+    //fn current_testcase_mut(&self) -> Result<&mut Testcase<I>, Error>;
+
+    /// Gets a cloned representation of the current [`Testcase`].
+    ///
+    /// Will return [`Error::key_not_found`] if no `corpus_idx` is currently set.
+    ///
+    /// # Note
+    /// This allocates memory and copies the contents!
+    /// For performance reasons, if you just need to access the testcase, use [`Self::current_testcase`] instead.
+    fn current_input_cloned(&self) -> Result<I, Error>;
+}
+
+impl<I, T> HasCurrentTestcase<I> for T
+where
+    I: Input,
+    T: HasCorpus + HasCurrentCorpusIdx + UsesInput<Input = I>,
+{
+    fn current_testcase(&self) -> Result<Ref<'_, Testcase<I>>, Error> {
+        let Some(corpus_id) = self.current_corpus_idx()? else {
+            return Err(Error::key_not_found(
+                "We are not currently processing a testcase",
+            ));
+        };
+
+        Ok(self.corpus().get(corpus_id)?.borrow())
+    }
+
+    fn current_testcase_mut(&self) -> Result<RefMut<'_, Testcase<I>>, Error> {
+        let Some(corpus_id) = self.current_corpus_idx()? else {
+            return Err(Error::illegal_state(
+                "We are not currently processing a testcase",
+            ));
+        };
+
+        Ok(self.corpus().get(corpus_id)?.borrow_mut())
+    }
+
+    fn current_input_cloned(&self) -> Result<I, Error> {
+        let mut testcase = self.current_testcase_mut()?;
+        Ok(testcase.borrow_mut().load_input(self.corpus())?.clone())
+    }
+}
+
 impl<I, C, R, SC> HasCurrentStage for StdState<I, C, R, SC> {
     fn set_stage(&mut self, idx: usize) -> Result<(), Error> {
         // ensure we are in the right frame
@@ -1115,7 +1222,7 @@ impl<I, C, R, SC> HasScalabilityMonitor for StdState<I, C, R, SC> {
 #[derive(Debug, Serialize, Deserialize, Default)]
 pub struct NopState<I> {
     metadata: SerdeAnyMap,
-    execution: usize,
+    execution: u64,
     rand: StdRand,
     phantom: PhantomData<I>,
 }
@@ -1141,11 +1248,11 @@ where
 }
 
 impl<I> HasExecutions for NopState<I> {
-    fn executions(&self) -> &usize {
+    fn executions(&self) -> &u64 {
         &self.execution
     }
 
-    fn executions_mut(&mut self) -> &mut usize {
+    fn executions_mut(&mut self) -> &mut u64 {
         &mut self.execution
     }
 }
@@ -1239,11 +1346,7 @@ pub mod test {
     use libafl_bolts::rands::StdRand;
 
     use super::StdState;
-    use crate::{
+    use crate::{corpus::InMemoryCorpus, inputs::Input};
-        corpus::InMemoryCorpus,
-        inputs::{Input, NopInput},
-        stages::test::{test_resume, test_resume_stages},
-    };
 
     #[must_use]
     pub fn test_std_state<I: Input>() -> StdState<I, InMemoryCorpus<I>, StdRand, InMemoryCorpus<I>>
@@ -1257,14 +1360,6 @@ pub mod test {
         )
         .expect("couldn't instantiate the test state")
     }
-
-    #[test]
-    fn resume_simple() {
-        let mut state = test_std_state::<NopInput>();
-        let (completed, stages) = test_resume_stages();
-
-        test_resume(&completed, &mut state, stages);
-    }
 }
 
 #[cfg(feature = "python")]
@@ -1360,7 +1455,7 @@ pub mod pybind {
             self.inner.as_ref().solutions().clone()
         }
 
-        fn executions(&self) -> usize {
+        fn executions(&self) -> u64 {
             *self.inner.as_ref().executions()
         }
 

libafl_bolts/src/lib.rs

@@ -442,7 +442,7 @@ impl Display for Error {
                 display_error_backtrace(f, b)
             }
             Self::KeyNotFound(s, b) => {
-                write!(f, "Key `{0}` not in Corpus", &s)?;
+                write!(f, "Key: `{0}` - not found", &s)?;
                 display_error_backtrace(f, b)
             }
             Self::Empty(s, b) => {

libafl_bolts/src/serdeany.rs

@@ -67,7 +67,7 @@ where
 pub mod serdeany_registry {
 
     use alloc::boxed::Box;
-    use core::{any::TypeId, fmt};
+    use core::{any::TypeId, fmt, hash::BuildHasherDefault};
 
     use hashbrown::{
         hash_map::{Keys, Values, ValuesMut},
@@ -75,6 +75,7 @@ pub mod serdeany_registry {
     };
     use serde::{Deserialize, Serialize};
 
+    use super::SerdeAny;
     use crate::{
         anymap::{pack_type_id, unpack_type_id},
         hash_std,
@@ -180,7 +181,7 @@ pub mod serdeany_registry {
     #[allow(clippy::unsafe_derive_deserialize)]
     #[derive(Debug, Serialize, Deserialize)]
     pub struct SerdeAnyMap {
-        map: HashMap<u128, Box<dyn crate::serdeany::SerdeAny>>,
+        map: HashMap<u128, Box<dyn SerdeAny>>,
     }
 
     // Cloning by serializing and deserializing. It ain't fast, but it's honest work.
@@ -257,7 +258,24 @@ pub mod serdeany_registry {
 
         /// Insert a boxed element into the map.
         #[inline]
-        pub fn insert_boxed<T>(&mut self, t: Box<T>)
+        pub fn insert_boxed<T>(&mut self, value: Box<T>)
+        where
+            T: crate::serdeany::SerdeAny,
+        {
+            self.entry::<T>().insert(value);
+        }
+
+        /// Get an entry to an element in this map.
+        #[inline]
+        #[allow(unused_qualifications)]
+        pub fn entry<T>(
+            &mut self,
+        ) -> hashbrown::hash_map::Entry<
+            '_,
+            u128,
+            Box<dyn SerdeAny + 'static>,
+            BuildHasherDefault<ahash::AHasher>,
+        >
         where
             T: crate::serdeany::SerdeAny,
         {
@@ -271,11 +289,28 @@ pub mod serdeany_registry {
                                 .get(&id)
                                 .is_some()
                         },
-                        "Type {} was inserted without registration! Call {}::register or use serde_autoreg.",
+                        "Type {} was inserted without registration! Call RegistryBuilder::register::<{}>() or use serde_autoreg.",
                         core::any::type_name::<T>(),
                         core::any::type_name::<T>()
                     );
-            self.map.insert(id, t);
+            self.map.entry(id)
+        }
+
+        /// Gets a value by type, or inserts it using the given construction function `default`
+        pub fn or_insert_with<T>(&mut self, default: impl FnOnce() -> T) -> &mut T
+        where
+            T: SerdeAny,
+        {
+            self.or_insert_with_boxed::<T>(|| Box::new(default()))
+        }
+
+        /// Gets a value by type, or inserts it using the given construction function `default` (returning a boxed value)
+        pub fn or_insert_with_boxed<T>(&mut self, default: impl FnOnce() -> Box<T>) -> &mut T
+        where
+            T: SerdeAny + 'static,
+        {
+            let ret = self.entry::<T>().or_insert_with(|| default());
+            ret.as_mut().as_any_mut().downcast_mut::<T>().unwrap()
         }
 
         /// Returns the count of elements in this map.
@@ -350,6 +385,21 @@ pub mod serdeany_registry {
             }
         }
 
+        /// Remove an element by type and name
+        #[must_use]
+        #[inline]
+        pub fn remove<T>(&mut self, name: &str) -> Option<Box<T>>
+        where
+            T: crate::serdeany::SerdeAny,
+        {
+            match self.map.get_mut(&unpack_type_id(TypeId::of::<T>())) {
+                None => None,
+                Some(h) => h
+                    .remove(&hash_std(name.as_bytes()))
+                    .map(|x| x.as_any_boxed().downcast::<T>().unwrap()),
+            }
+        }
+
         /// Get an element of a given type contained in this map by [`TypeId`].
         #[must_use]
         #[allow(unused_qualifications)]
@@ -532,7 +582,25 @@ pub mod serdeany_registry {
         /// Insert an element into this map.
         #[inline]
         #[allow(unused_qualifications)]
-        pub fn insert<T>(&mut self, val: T, name: &str)
+        pub fn insert<T>(&mut self, name: &str, val: T)
+        where
+            T: crate::serdeany::SerdeAny,
+        {
+            self.entry::<T>(name).insert(Box::new(val));
+        }
+
+        /// Get an entry to an element into this map.
+        #[inline]
+        #[allow(unused_qualifications)]
+        pub fn entry<T>(
+            &mut self,
+            name: &str,
+        ) -> hashbrown::hash_map::Entry<
+            '_,
+            u64,
+            Box<dyn SerdeAny + 'static>,
+            BuildHasherDefault<ahash::AHasher>,
+        >
         where
             T: crate::serdeany::SerdeAny,
         {
@@ -546,17 +614,36 @@ pub mod serdeany_registry {
                                 .get(&id)
                                 .is_some()
                         },
-                        "Type {} was inserted without registration! Call {}::register or use serde_autoreg.",
+                        "Type {} was inserted without registration! Call RegistryBuilder::register::<{}>() or use serde_autoreg.",
                         core::any::type_name::<T>(),
                         core::any::type_name::<T>()
                     );
-            if !self.map.contains_key(&id) {
-                self.map.insert(id, HashMap::default());
-            }
             self.map
-                .get_mut(&id)
+                .entry(id)
-                .unwrap()
+                .or_default()
-                .insert(hash_std(name.as_bytes()), Box::new(val));
+                .entry(hash_std(name.as_bytes()))
+        }
+
+        /// Gets a value by name, or inserts it using the given construction function `default`
+        pub fn or_insert_with<T>(&mut self, name: &str, default: impl FnOnce() -> T) -> &mut T
+        where
+            T: SerdeAny,
+        {
+            let ret = self.entry::<T>(name).or_insert_with(|| Box::new(default()));
+            ret.as_mut().as_any_mut().downcast_mut::<T>().unwrap()
+        }
+
+        /// Gets a value by name, or inserts it using the given construction function `default` (returning a boxed value)
+        pub fn or_insert_with_boxed<T>(
+            &mut self,
+            name: &str,
+            default: impl FnOnce() -> Box<T>,
+        ) -> &mut T
+        where
+            T: SerdeAny + 'static,
+        {
+            let ret = self.entry::<T>(name).or_insert_with(|| default());
+            ret.as_mut().as_any_mut().downcast_mut::<T>().unwrap()
         }
 
         /// Returns the `len` of this map.

libafl_concolic/symcc_runtime/src/lib.rs

@@ -30,8 +30,7 @@
 #![allow(
     clippy::module_name_repetitions,
     clippy::missing_panics_doc,
-    clippy::pub_underscore_fields,
+    clippy::pub_underscore_fields
-    clippy::mixed_attributes_style
 )]
 
 pub mod filter;
@@ -40,6 +39,7 @@ pub mod tracing;
 // The following exports are used by the `export_runtime` macro. They are therefore exported, but hidden from docs, as they are not supposed to be used directly by the user.
 #[doc(hidden)]
 #[cfg(target_os = "linux")]
+#[allow(clippy::mixed_attributes_style)]
 pub mod cpp_runtime {
     #![allow(non_upper_case_globals)]
     #![allow(non_camel_case_types)]

libafl_qemu/src/edges.rs

@@ -498,13 +498,7 @@ where
         }
     }
     let state = state.expect("The gen_unique_edge_ids hook works only for in-process fuzzing");
-    if state.metadata_map().get::<QemuEdgesMapMetadata>().is_none() {
+    let meta = state.metadata_or_insert_with(QemuEdgesMapMetadata::new);
-        state.add_metadata(QemuEdgesMapMetadata::new());
-    }
-    let meta = state
-        .metadata_map_mut()
-        .get_mut::<QemuEdgesMapMetadata>()
-        .unwrap();
 
     match meta.map.entry((src, dest)) {
         Entry::Occupied(e) => {

libafl_sugar/Cargo.toml

@@ -36,12 +36,14 @@ pyo3-build-config = { version = "0.18", optional = true }
 libafl = { path = "../libafl", version = "0.11.2" }
 libafl_bolts = { path = "../libafl_bolts", version = "0.11.2" }
 libafl_targets = { path = "../libafl_targets", version = "0.11.2" }
-libafl_qemu = { path = "../libafl_qemu", version = "0.11.2" }
 
 typed-builder = "0.16" # Implement the builder pattern at compiletime
 pyo3 = { version = "0.18", optional = true }
 log = "0.4.20"
 
+[target.'cfg(target_os = "linux")'.dependencies]
+libafl_qemu = { path = "../libafl_qemu", version = "0.11.2" }
+
 [lib]
 name = "libafl_sugar"
 crate-type = ["cdylib", "rlib"]

libafl_targets/src/cmps/stages/aflpptracing.rs

@@ -4,15 +4,16 @@ use core::marker::PhantomData;
 #[cfg(feature = "introspection")]
 use libafl::state::HasClientPerfMonitor;
 use libafl::{
-    corpus::{Corpus, HasCurrentCorpusIdx},
     executors::{Executor, HasObservers},
     inputs::{BytesInput, UsesInput},
     observers::ObserversTuple,
-    stages::{colorization::TaintMetadata, Stage},
+    stages::{colorization::TaintMetadata, RetryRestartHelper, Stage},
-    state::{HasCorpus, HasExecutions, HasMetadata, UsesState},
+    state::{
+        HasCorpus, HasCurrentTestcase, HasExecutions, HasMetadata, HasNamedMetadata, UsesState,
+    },
     Error,
 };
-use libafl_bolts::tuples::MatchName;
+use libafl_bolts::{tuples::MatchName, Named};
 
 use crate::cmps::observers::AFLppCmpLogObserver;
 
@@ -32,16 +33,21 @@ where
     type State = TE::State;
 }
 
+impl<EM, TE, Z> Named for AFLppCmplogTracingStage<EM, TE, Z> {
+    fn name(&self) -> &str {
+        "AFLppCmplogTracingStage"
+    }
+}
+
 impl<E, EM, TE, Z> Stage<E, EM, Z> for AFLppCmplogTracingStage<EM, TE, Z>
 where
     E: UsesState<State = TE::State>,
     TE: Executor<EM, Z> + HasObservers,
-    TE::State: HasExecutions + HasCorpus + HasMetadata + UsesInput<Input = BytesInput>,
+    TE::State:
+        HasExecutions + HasCorpus + HasMetadata + UsesInput<Input = BytesInput> + HasNamedMetadata,
     EM: UsesState<State = TE::State>,
     Z: UsesState<State = TE::State>,
 {
-    type Progress = (); // TODO this needs resumption
-
     #[inline]
     fn perform(
         &mut self,
@@ -51,11 +57,7 @@ where
         manager: &mut EM,
     ) -> Result<(), Error> {
         // First run with the un-mutated input
-        let corpus_idx = state.current_corpus_idx()?.ok_or_else(|| {
+        let unmutated_input = state.current_input_cloned()?;
-            Error::illegal_state("state is not currently processing a corpus index")
-        })?;
-
-        let unmutated_input = state.corpus().cloned_input_for_id(corpus_idx)?;
 
         if let Some(name) = &self.cmplog_observer_name {
             if let Some(ob) = self
@@ -121,6 +123,16 @@ where
 
         Ok(())
     }
+
+    fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
+        // TODO: this may need better resumption? (Or is it always used with a forkserver?)
+        RetryRestartHelper::restart_progress_should_run(state, self, 3)
+    }
+
+    fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
+        // TODO: this may need better resumption? (Or is it always used with a forkserver?)
+        RetryRestartHelper::clear_restart_progress(state, self)
+    }
 }
 
 impl<EM, TE, Z> AFLppCmplogTracingStage<EM, TE, Z> {

libafl_targets/src/lib.rs

@@ -14,8 +14,7 @@
     clippy::missing_panics_doc,
     clippy::missing_docs_in_private_items,
     clippy::module_name_repetitions,
-    clippy::pub_underscore_fields,
+    clippy::pub_underscore_fields
-    clippy::mixed_attributes_style
 )]
 #![cfg_attr(not(test), warn(
     missing_debug_implementations,
@@ -90,6 +89,7 @@ pub use sancov_cmp::*;
 
 /// Module containing bindings to the various sanitizer interface headers
 #[cfg(feature = "sanitizer_interfaces")]
+#[allow(clippy::mixed_attributes_style)]
 pub mod sanitizer_ifaces {
     #![allow(non_snake_case)]
     #![allow(non_camel_case_types)]

libafl_targets/src/sancov_pcguard.rs

@@ -7,12 +7,17 @@ use core::simd::num::SimdUint;
 #[cfg(any(feature = "sancov_ngram4", feature = "sancov_ctx"))]
 use libafl::executors::{hooks::ExecutorHook, HasObservers};
 
+#[cfg(any(
+    feature = "pointer_maps",
+    feature = "sancov_pcguard_edges",
+    feature = "sancov_pcguard_hitcounts"
+))]
+use crate::coverage::EDGES_MAP;
+use crate::coverage::MAX_EDGES_NUM;
 #[cfg(feature = "pointer_maps")]
 use crate::coverage::{EDGES_MAP_PTR, EDGES_MAP_PTR_NUM};
-use crate::{
+#[cfg(feature = "sancov_ngram4")]
-    coverage::{EDGES_MAP, MAX_EDGES_NUM},
+use crate::EDGES_MAP_SIZE;
-    EDGES_MAP_SIZE,
-};
 
 #[cfg(all(feature = "sancov_pcguard_edges", feature = "sancov_pcguard_hitcounts"))]
 #[cfg(not(any(doc, feature = "clippy")))]
@@ -160,6 +165,7 @@ extern "C" {
 #[no_mangle]
 #[allow(unused_assignments)]
 pub unsafe extern "C" fn __sanitizer_cov_trace_pc_guard(guard: *mut u32) {
+    #[allow(unused_mut)]
     let mut pos = *guard as usize;
 
     #[cfg(any(feature = "sancov_ngram4", feature = "sancov_ngram8"))]