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/accounting.rs

@@ -16,7 +16,7 @@ use crate::{
     corpus::{Corpus, CorpusId},
     observers::CanTrack,
     schedulers::{
-        Scheduler,
+        RemovableScheduler, Scheduler,
         minimizer::{DEFAULT_SKIP_NON_FAVORED_PROB, IsFavoredMetadata, MinimizerScheduler},
     },
     state::{HasCorpus, HasRand},
@@ -112,7 +112,7 @@ pub struct CoverageAccountingScheduler<'a, CS, I, O> {
 
 impl<CS, I, O, S> Scheduler<I, S> for CoverageAccountingScheduler<'_, CS, I, O>
 where
-    CS: Scheduler<I, S>,
+    CS: Scheduler<I, S> + RemovableScheduler<I, S>,
     S: HasCorpus<I> + HasMetadata + HasRand,
     I: HasLen,
     O: CanTrack,

crates/libafl/src/schedulers/minimizer.rs

@@ -84,7 +84,7 @@ where
     CS: RemovableScheduler<I, S> + Scheduler<I, S>,
     F: TestcasePenalty<I, S>,
     M: for<'a> AsIter<'a, Item = usize> + SerdeAny + HasRefCnt,
-    S: HasCorpus<I> + HasMetadata + HasRand,
+    S: HasCorpus<I> + HasMetadata,
 {
     /// Replaces the [`Testcase`] at the given [`CorpusId`]
     fn on_replace(
@@ -189,7 +189,7 @@ where
 
 impl<CS, F, I, M, O, S> Scheduler<I, S> for MinimizerScheduler<CS, F, I, M, O>
 where
-    CS: Scheduler<I, S>,
+    CS: Scheduler<I, S> + RemovableScheduler<I, S>,
     F: TestcasePenalty<I, S>,
     M: for<'a> AsIter<'a, Item = usize> + SerdeAny + HasRefCnt,
     S: HasCorpus<I> + HasMetadata + HasRand,
@@ -319,9 +319,12 @@ where
     }
 
     /// Cull the [`Corpus`] using the [`MinimizerScheduler`]
-    pub fn cull<S>(&self, state: &S) -> Result<(), Error>
+    pub fn cull<S>(&mut self, state: &mut S) -> Result<(), Error>
     where
-        S: HasCorpus<I> + HasMetadata,
+        S: HasCorpus<I> + HasMetadata + HasRand,
+        CS: RemovableScheduler<I, S>,
+        CS: Scheduler<I, S>,
+        F: TestcasePenalty<I, S>,
     {
         let Some(top_rated) = state.metadata_map().get::<TopRatedsMetadata>() else {
             return Ok(());
@@ -331,18 +334,34 @@ where
 
         for (key, id) in &top_rated.map {
             if !acc.contains(key) {
-                let mut entry = state.corpus().get(*id)?.borrow_mut();
-                let meta = entry.metadata_map().get::<M>().ok_or_else(|| {
-                    Error::key_not_found(format!(
-                        "{} needed for MinimizerScheduler not found in testcase #{id}",
-                        type_name::<M>()
-                    ))
-                })?;
-                for elem in meta.as_iter() {
-                    acc.insert(*elem);
-                }
+                let entry = state.corpus().get(*id);
+                match entry {
+                    Ok(entry) => {
+                        let mut entry = entry.borrow_mut();
+                        let meta = entry.metadata_map().get::<M>().ok_or_else(|| {
+                            Error::key_not_found(format!(
+                                "{} needed for MinimizerScheduler not found in testcase #{id}",
+                                type_name::<M>()
+                            ))
+                        })?;
+                        for elem in meta.as_iter() {
+                            acc.insert(*elem);
+                        }
 
-                entry.add_metadata(IsFavoredMetadata {});
+                        entry.add_metadata(IsFavoredMetadata {});
+                    }
+                    Err(Error::KeyNotFound(_, _)) => {
+                        // Fallback for disabled and removed testcases.
+                        log::info!(
+                            "Removing testcase #{id} from top rateds as it was no longer in the corpus!"
+                        );
+                        self.on_remove(state, *id, &None)?;
+                        return self.cull(state);
+                    }
+                    Err(err) => {
+                        return Err(err);
+                    }
+                }
             }
         }
 

crates/libafl/src/schedulers/mod.rs

@@ -78,10 +78,13 @@ where
     let current_id = *state.corpus().current();
 
     let mut depth = match current_id {
-        Some(parent_idx) => state
-            .testcase(parent_idx)?
-            .metadata::<SchedulerTestcaseMetadata>()?
-            .depth(),
+        Some(parent_idx) => state.corpus().get_from_all(parent_idx).map_or_else(
+            |_| {
+                log::warn!("Parent testcase with id {parent_idx} not found! Removed?");
+                Ok::<u64, Error>(0)
+            },
+            |v| Ok(v.borrow().metadata::<SchedulerTestcaseMetadata>()?.depth()),
+        )?,
         None => 0,
     };
 

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(())
+    }
+}