Calibration Stage Fix for crashing testcases

crates/libafl/src/corpus/inmemory.rs

@@ -286,8 +286,9 @@ impl<I> TestcaseStorage<I> {
         id
     }
 
+    /// Insert a testcase with a specific `CorpusId`. Handle with care!
     #[cfg(not(feature = "corpus_btreemap"))]
-    fn insert_inner_with_id(
+    pub fn insert_inner_with_id(
         &mut self,
         testcase: RefCell<Testcase<I>>,
         is_disabled: bool,
@@ -325,8 +326,9 @@ impl<I> TestcaseStorage<I> {
         Ok(())
     }
 
+    /// Insert a testcase with a specific `CorpusId`. Handle with care!
     #[cfg(feature = "corpus_btreemap")]
-    fn insert_inner_with_id(
+    pub fn insert_inner_with_id(
         &mut self,
         testcase: RefCell<Testcase<I>>,
         is_disabled: bool,

crates/libafl/src/corpus/mod.rs

@@ -212,10 +212,12 @@ pub trait HasCurrentCorpusId {
     /// Set the current corpus index; we have started processing this corpus entry
     fn set_corpus_id(&mut self, id: CorpusId) -> Result<(), Error>;
 
-    /// Clear the current corpus index; we are done with this entry
+    /// Clear the current corpus index; we are done with this entry.
+    /// This can also be used by as stage to signal that subsequent stages should be skipped.
     fn clear_corpus_id(&mut self) -> Result<(), Error>;
 
-    /// Fetch the current corpus index -- typically used after a state recovery or transfer
+    /// Fetch the current corpus index -- typically used after a state recovery or transfer.
+    /// If it is `None`, the corpus scheduler should be called to schedule the next testcase.
     fn current_corpus_id(&self) -> Result<Option<CorpusId>, Error>;
 }
 

crates/libafl/src/corpus/testcase.rs

@@ -199,7 +199,7 @@ impl<I> Testcase<I> {
 
     /// Get `disabled`
     #[inline]
-    pub fn disabled(&mut self) -> bool {
+    pub fn disabled(&self) -> bool {
         self.disabled
     }
 

crates/libafl/src/events/simple.rs

@@ -1,5 +1,4 @@
 //! A very simple event manager, that just supports log outputs, but no multiprocessing
-
 use alloc::vec::Vec;
 #[cfg(feature = "std")]
 use core::sync::atomic::{Ordering, compiler_fence};
@@ -189,7 +188,6 @@ impl<I, MT, S> SimpleEventManager<I, MT, S>
 where
     I: Debug,
     MT: Monitor,
-    S: Stoppable,
 {
     /// Creates a new [`SimpleEventManager`].
     pub fn new(monitor: MT) -> Self {

crates/libafl/src/executors/forkserver.rs

@@ -119,26 +119,30 @@ const FAILED_TO_START_FORKSERVER_MSG: &str = "Failed to start forkserver";
 fn report_error_and_exit(status: i32) -> Result<(), Error> {
     /* Report on the error received via the forkserver controller and exit */
     match status {
-    FS_ERROR_MAP_SIZE =>
-        Err(Error::unknown(
-            format!(
-            "{AFL_MAP_SIZE_ENV_VAR} is not set and fuzzing target reports that the required size is very large. Solution: Run the fuzzing target stand-alone with the environment variable AFL_DEBUG=1 set and set the value for __afl_final_loc in the {AFL_MAP_SIZE_ENV_VAR} environment variable for afl-fuzz."))),
-    FS_ERROR_MAP_ADDR =>
-        Err(Error::unknown(
-            "the fuzzing target reports that hardcoded map address might be the reason the mmap of the shared memory failed. Solution: recompile the target with either afl-clang-lto and do not set AFL_LLVM_MAP_ADDR or recompile with afl-clang-fast.".to_string())),
-    FS_ERROR_SHM_OPEN =>
-        Err(Error::unknown("the fuzzing target reports that the shm_open() call failed.".to_string())),
-    FS_ERROR_SHMAT =>
-        Err(Error::unknown("the fuzzing target reports that the shmat() call failed.".to_string())),
-    FS_ERROR_MMAP =>
-        Err(Error::unknown("the fuzzing target reports that the mmap() call to the shared memory failed.".to_string())),
-    FS_ERROR_OLD_CMPLOG =>
-        Err(Error::unknown(
-            "the -c cmplog target was instrumented with an too old AFL++ version, you need to recompile it.".to_string())),
-    FS_ERROR_OLD_CMPLOG_QEMU =>
-        Err(Error::unknown("The AFL++ QEMU/FRIDA loaders are from an older version, for -c you need to recompile it.".to_string())),
-    _ =>
-        Err(Error::unknown(format!("unknown error code {status} from fuzzing target!"))),
+        FS_ERROR_MAP_SIZE => Err(Error::unknown(format!(
+            "{AFL_MAP_SIZE_ENV_VAR} is not set and fuzzing target reports that the required size is very large. Solution: Run the fuzzing target stand-alone with the environment variable AFL_DEBUG=1 set and set the value for __afl_final_loc in the {AFL_MAP_SIZE_ENV_VAR} environment variable for afl-fuzz."
+        ))),
+        FS_ERROR_MAP_ADDR => Err(Error::unknown(
+            "the fuzzing target reports that hardcoded map address might be the reason the mmap of the shared memory failed. Solution: recompile the target with either afl-clang-lto and do not set AFL_LLVM_MAP_ADDR or recompile with afl-clang-fast.",
+        )),
+        FS_ERROR_SHM_OPEN => Err(Error::unknown(
+            "the fuzzing target reports that the shm_open() call failed.",
+        )),
+        FS_ERROR_SHMAT => Err(Error::unknown(
+            "the fuzzing target reports that the shmat() call failed.",
+        )),
+        FS_ERROR_MMAP => Err(Error::unknown(
+            "the fuzzing target reports that the mmap() call to the shared memory failed.",
+        )),
+        FS_ERROR_OLD_CMPLOG => Err(Error::unknown(
+            "the -c cmplog target was instrumented with an too old AFL++ version, you need to recompile it.",
+        )),
+        FS_ERROR_OLD_CMPLOG_QEMU => Err(Error::unknown(
+            "The AFL++ QEMU/FRIDA loaders are from an older version, for -c you need to recompile it.",
+        )),
+        _ => Err(Error::unknown(format!(
+            "unknown error code {status} from fuzzing target!"
+        ))),
     }
 }
 
@@ -422,7 +426,9 @@ impl Forkserver {
         };
 
         if env::var(SHM_ENV_VAR).is_err() {
-            return Err(Error::unknown("__AFL_SHM_ID not set. It is necessary to set this env, otherwise the forkserver cannot communicate with the fuzzer".to_string()));
+            return Err(Error::unknown(
+                "__AFL_SHM_ID not set. It is necessary to set this env, otherwise the forkserver cannot communicate with the fuzzer",
+            ));
         }
 
         let afl_debug = if let Ok(afl_debug) = env::var("AFL_DEBUG") {

crates/libafl/src/schedulers/testcase_score.rs

@@ -1,5 +1,5 @@
 //! The `TestcaseScore` is an evaluator providing scores of corpus items.
-use alloc::string::{String, ToString};
+use alloc::string::String;
 
 use libafl_bolts::{HasLen, HasRefCnt};
 use num_traits::Zero;
@@ -102,7 +102,7 @@ where
         let mut perf_score = 100.0;
         let q_exec_us = entry
             .exec_time()
-            .ok_or_else(|| Error::key_not_found("exec_time not set".to_string()))?
+            .ok_or_else(|| Error::key_not_found("exec_time not set when computing corpus power. This happens if CalibrationStage fails to set it or is not added to stages."))?
             .as_nanos() as f64;
 
         let avg_exec_us = psmeta.exec_time().as_nanos() as f64 / psmeta.cycles() as f64;
@@ -290,7 +290,7 @@ where
 
         let q_exec_us = entry
             .exec_time()
-            .ok_or_else(|| Error::key_not_found("exec_time not set".to_string()))?
+            .ok_or_else(|| Error::key_not_found("exec_time not set when computing corpus weight. This happens if CalibrationStage fails to set it or is not added to stages."))?
             .as_nanos() as f64;
         let favored = entry.has_metadata::<IsFavoredMetadata>();
 

crates/libafl/src/stages/calibrate.rs

@@ -14,7 +14,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::{
     Error, HasMetadata, HasNamedMetadata, HasScheduler,
-    corpus::{Corpus, HasCurrentCorpusId, SchedulerTestcaseMetadata},
+    corpus::{Corpus, EnableDisableCorpus, HasCurrentCorpusId, SchedulerTestcaseMetadata},
     events::{Event, EventFirer, EventWithStats, LogSeverity},
     executors::{Executor, ExitKind, HasObservers},
     feedbacks::{HasObserverHandle, map::MapFeedbackMetadata},
@@ -48,6 +48,11 @@ pub struct UnstableEntriesMetadata {
 }
 impl_serdeany!(UnstableEntriesMetadata);
 
+/// Metadata to mark a testcase as disabled in the calibration stage due to crash or timeout.
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct DisabledInCalibrationStageMetadata;
+impl_serdeany!(DisabledInCalibrationStageMetadata);
+
 impl UnstableEntriesMetadata {
     #[must_use]
     /// Create a new [`struct@UnstableEntriesMetadata`]
@@ -376,19 +381,33 @@ where
 
 impl<C, I, O, OT, S> Restartable<S> for CalibrationStage<C, I, O, OT, S>
 where
-    S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
+    S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId + HasCorpus<I> + HasCurrentTestcase<I>,
+    S::Corpus: EnableDisableCorpus,
 {
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         // Calibration stage disallow restarts
         // If a testcase that causes crash/timeout in the queue, we need to remove it from the queue immediately.
-        RetryCountRestartHelper::no_retry(state, &self.name)
-
-        // todo
-        // remove this guy from corpus queue
+        let retry = RetryCountRestartHelper::no_retry(state, &self.name)?;
+        if !retry {
+            let id = state
+                .current_corpus_id()?
+                .ok_or_else(|| Error::illegal_state("No current corpus id"))?;
+            log::info!("Disabling crashing/timeouting testcase {id} during calibration");
+            let insert_result = state
+                .current_testcase_mut()?
+                .try_add_metadata(DisabledInCalibrationStageMetadata);
+            if let Err(err) = insert_result {
+                log::warn!("Calibration stage called on already disabled testcase {id}: {err:?}.");
+            } else {
+                state.corpus_mut().disable(id)?;
+                self.clear_progress(state)?;
+                return Err(Error::skip_remaining_stages());
+            }
+        }
+        Ok(retry)
     }
 
     fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
-        // TODO: Make sure this is the correct way / there may be a better way?
         RetryCountRestartHelper::clear_progress(state, &self.name)
     }
 }
@@ -436,3 +455,61 @@ impl<C, I, O, OT, S> Named for CalibrationStage<C, I, O, OT, S> {
         &self.name
     }
 }
+
+#[cfg(test)]
+mod tests {
+    #[cfg(not(feature = "serdeany_autoreg"))]
+    use libafl_bolts::serdeany::RegistryBuilder;
+    use libafl_bolts::{Error, rands::StdRand};
+
+    #[cfg(not(feature = "serdeany_autoreg"))]
+    use super::DisabledInCalibrationStageMetadata;
+    use crate::{
+        corpus::{Corpus, HasCurrentCorpusId, InMemoryCorpus, Testcase},
+        feedbacks::{MaxMapFeedback, StateInitializer},
+        inputs::NopInput,
+        observers::StdMapObserver,
+        stages::{CalibrationStage, Restartable},
+        state::{HasCorpus, StdState},
+    };
+
+    #[test]
+    fn test_calibration_restart() -> Result<(), Error> {
+        #[cfg(not(feature = "serdeany_autoreg"))]
+        RegistryBuilder::register::<DisabledInCalibrationStageMetadata>();
+
+        // Setup
+        let mut state = StdState::new(
+            StdRand::with_seed(0),
+            InMemoryCorpus::new(),
+            InMemoryCorpus::new(),
+            &mut (),
+            &mut (),
+        )?;
+
+        let input = NopInput {};
+        let testcase = Testcase::new(input);
+        let id = state.corpus_mut().add(testcase)?;
+        state.set_corpus_id(id)?;
+
+        let observer = StdMapObserver::owned("map", vec![0u8; 16]);
+        let mut feedback = MaxMapFeedback::new(&observer);
+        feedback.init_state(&mut state)?;
+        let mut stage: CalibrationStage<_, _, _, (), _> = CalibrationStage::new(&feedback);
+
+        // 1. First try (should restart)
+        assert!(stage.should_restart(&mut state)?);
+
+        // 2. Second call - should return Error::SkipRemainingStages
+        match stage.should_restart(&mut state) {
+            Err(Error::SkipRemainingStages) => (),
+            res => panic!("Expected SkipRemainingStages, got {:?}", res),
+        }
+
+        // Verify testcase is disabled
+        assert!(state.corpus().get(id).is_err()); // Should be error because it's disabled
+        assert!(state.corpus().get_from_all(id).is_ok()); // Should be ok
+
+        Ok(())
+    }
+}

crates/libafl/src/stages/mod.rs

@@ -177,7 +177,14 @@ where
 
                 let stage = &mut self.0;
 
-                stage.perform_restartable(fuzzer, executor, state, manager)?;
+                match stage.perform_restartable(fuzzer, executor, state, manager) {
+                    Ok(()) => {}
+                    Err(Error::SkipRemainingStages) => {
+                        state.clear_stage_id()?;
+                        return Ok(());
+                    }
+                    Err(e) => return Err(e),
+                }
 
                 state.clear_stage_id()?;
             }
@@ -190,7 +197,14 @@ where
 
                 let stage = &mut self.0;
 
-                stage.perform_restartable(fuzzer, executor, state, manager)?;
+                match stage.perform_restartable(fuzzer, executor, state, manager) {
+                    Ok(()) => {}
+                    Err(Error::SkipRemainingStages) => {
+                        state.clear_stage_id()?;
+                        return Ok(());
+                    }
+                    Err(e) => return Err(e),
+                }
 
                 state.clear_stage_id()?;
             }
@@ -207,7 +221,10 @@ where
         }
 
         // Execute the remaining stages
-        self.1.perform_all(fuzzer, executor, state, manager)
+        match self.1.perform_all(fuzzer, executor, state, manager) {
+            Ok(()) | Err(Error::SkipRemainingStages) => Ok(()),
+            Err(e) => Err(e),
+        }
     }
 }
 
@@ -290,14 +307,28 @@ where
         state: &mut S,
         manager: &mut EM,
     ) -> Result<(), Error> {
-        self.iter_mut().try_for_each(|stage| {
-            if state.stop_requested() {
-                state.discard_stop_request();
-                manager.on_shutdown()?;
-                return Err(Error::shutting_down());
-            }
-            stage.perform_restartable(fuzzer, executor, state, manager)
-        })
+        self.iter_mut()
+            .try_for_each(|stage| {
+                if state.stop_requested() {
+                    state.discard_stop_request();
+                    manager.on_shutdown()?;
+                    return Err(Error::shutting_down());
+                }
+                match stage.perform_restartable(fuzzer, executor, state, manager) {
+                    Ok(()) => Ok(()),
+                    Err(Error::SkipRemainingStages) => {
+                        // Skip the remaining stages
+                        // We return an error to stop the iterator, but we want to return Ok(()) from perform_all
+
+                        Err(Error::SkipRemainingStages)
+                    }
+                    Err(e) => Err(e),
+                }
+            })
+            .or_else(|e| match e {
+                Error::SkipRemainingStages => Ok(()),
+                _ => Err(e),
+            })
     }
 }
 

crates/libafl/src/state/mod.rs

@@ -566,8 +566,7 @@ pub trait HasCurrentStageId {
 /// Trait for types which track nested stages. Stages which themselves contain stage tuples should
 /// ensure that they constrain the state with this trait accordingly.
 pub trait HasNestedStage: HasCurrentStageId {
-    /// Enter a stage scope, potentially resuming to an inner stage status. Returns Ok(true) if
-    /// resumed.
+    /// Enter a stage scope, potentially resuming to an inner stage status.
     fn enter_inner_stage(&mut self) -> Result<(), Error>;
 
     /// Exit a stage scope

crates/libafl_core/src/lib.rs

@@ -159,6 +159,8 @@ pub enum Error {
     IllegalArgument(String, ErrorBacktrace),
     /// The performed action is not supported on the current platform
     Unsupported(String, ErrorBacktrace),
+    /// Raise this from a stage to skip the remaining stages for a given input, not really an error.
+    SkipRemainingStages,
     /// Shutting down, not really an error.
     ShuttingDown,
     /// OS error, wrapping a [`io::Error`]
@@ -336,6 +338,12 @@ impl Error {
     {
         Error::Runtime(arg.into(), ErrorBacktrace::capture())
     }
+
+    /// Skip the remaining stages for this input
+    #[must_use]
+    pub fn skip_remaining_stages() -> Self {
+        Error::SkipRemainingStages
+    }
 }
 
 impl core::error::Error for Error {
@@ -422,6 +430,7 @@ impl Display for Error {
                 write!(f, "Encountered an invalid input: {0}", &s)?;
                 display_error_backtrace(f, b)
             }
+            Self::SkipRemainingStages => write!(f, "Skip remaining stages"),
         }
     }
 }