Calibration Stage Fix for crashing testcases

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,31 @@ 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)?;
+            }
+        }
+        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 +453,51 @@ impl<C, I, O, OT, S> Named for CalibrationStage<C, I, O, OT, S> {
         &self.name
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use libafl_bolts::{Error, rands::StdRand};
+
+    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> {
+        // 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 false and disable testcase
+        assert!(!stage.should_restart(&mut state)?);
+
+        // 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_tinyinst/src/executor.rs

@@ -93,7 +93,7 @@ where
             RunResult::OTHER_ERROR => Err(Error::unknown(
                 "Tinyinst RunResult is other error".to_string(),
             )),
-            _ => Err(Error::unknown("Tinyinst RunResult is unknown".to_string())),
+            _ => Err(Error::unknown("Tinyinst RunResult is unknown")),
         }
     }
 }