Add disable/enable to Corpus (#3193)

* Add enable and disable methods for Corpus

* Add insert_inner_with_id to fix disable/enable & test

Since we need to insert an 'existing' testcase with a certain id, let's
use a private inner function for it.

It's not the most posh way to keep consistency, but as showed in the
test it works 'good enough'.

* Implement disable/enable for libafl_libfuzzer/corpus

* fix clippy issues and fix cfg[not"corpus_btreemap"]

* Move enable/disable from Corpus to a trait

* Rename HasCorpusEnablementOperations to EnableDisableCorpus

Unless we come up with a better idea. Naming is hard.

* fmt the changes

libafl/src/corpus/cached.rs

@@ -9,8 +9,8 @@ use serde::{Deserialize, Serialize};
 use crate::{
     Error,
     corpus::{
-        Corpus, CorpusId, HasTestcase, Testcase, inmemory_ondisk::InMemoryOnDiskCorpus,
+        Corpus, CorpusId, EnableDisableCorpus, HasTestcase, Testcase,
-        ondisk::OnDiskMetadataFormat,
+        inmemory_ondisk::InMemoryOnDiskCorpus, ondisk::OnDiskMetadataFormat,
     },
     inputs::Input,
 };
@@ -191,6 +191,23 @@ where
     }
 }
 
+impl<I> EnableDisableCorpus for CachedOnDiskCorpus<I>
+where
+    I: Input,
+{
+    #[inline]
+    fn disable(&mut self, id: CorpusId) -> Result<(), Error> {
+        self.cached_indexes.borrow_mut().retain(|e| *e != id);
+        self.inner.disable(id)
+    }
+
+    #[inline]
+    fn enable(&mut self, id: CorpusId) -> Result<(), Error> {
+        self.cached_indexes.borrow_mut().retain(|e| *e != id);
+        self.inner.enable(id)
+    }
+}
+
 impl<I> CachedOnDiskCorpus<I> {
     /// Creates the [`CachedOnDiskCorpus`].
     ///

libafl/src/corpus/inmemory.rs

@@ -5,7 +5,7 @@ use core::cell::{Ref, RefCell, RefMut};
 
 use serde::{Deserialize, Serialize};
 
-use super::HasTestcase;
+use super::{EnableDisableCorpus, HasTestcase};
 use crate::{
     Error,
     corpus::{Corpus, CorpusId, Testcase},
@@ -285,6 +285,67 @@ impl<I> TestcaseStorage<I> {
         id
     }
 
+    #[cfg(not(feature = "corpus_btreemap"))]
+    fn insert_inner_with_id(
+        &mut self,
+        testcase: RefCell<Testcase<I>>,
+        is_disabled: bool,
+        id: CorpusId,
+    ) -> Result<(), Error> {
+        if self.progressive_id < id.into() {
+            return Err(Error::illegal_state(
+                "trying to insert a testcase with an id bigger than the internal Id counter",
+            ));
+        }
+        let corpus = if is_disabled {
+            &mut self.disabled
+        } else {
+            &mut self.enabled
+        };
+        let prev = if let Some(last_id) = corpus.last_id {
+            corpus.map.get_mut(&last_id).unwrap().next = Some(id);
+            Some(last_id)
+        } else {
+            None
+        };
+        if corpus.first_id.is_none() {
+            corpus.first_id = Some(id);
+        }
+        corpus.last_id = Some(id);
+        corpus.insert_key(id);
+        corpus.map.insert(
+            id,
+            TestcaseStorageItem {
+                testcase,
+                prev,
+                next: None,
+            },
+        );
+        Ok(())
+    }
+
+    #[cfg(feature = "corpus_btreemap")]
+    fn insert_inner_with_id(
+        &mut self,
+        testcase: RefCell<Testcase<I>>,
+        is_disabled: bool,
+        id: CorpusId,
+    ) -> Result<(), Error> {
+        if self.progressive_id < id.into() {
+            return Err(Error::illegal_state(
+                "trying to insert a testcase with an id bigger than the internal Id counter",
+            ));
+        }
+        let corpus = if is_disabled {
+            &mut self.disabled
+        } else {
+            &mut self.enabled
+        };
+        corpus.insert_key(id);
+        corpus.map.insert(id, testcase);
+        Ok(())
+    }
+
     /// Insert a testcase assigning a `CorpusId` to it
     #[cfg(feature = "corpus_btreemap")]
     fn insert_inner(&mut self, testcase: RefCell<Testcase<I>>, is_disabled: bool) -> CorpusId {
@@ -456,6 +517,30 @@ impl<I> Corpus<I> for InMemoryCorpus<I> {
     }
 }
 
+impl<I> EnableDisableCorpus for InMemoryCorpus<I> {
+    #[inline]
+    fn disable(&mut self, id: CorpusId) -> Result<(), Error> {
+        if let Some(testcase) = self.storage.enabled.remove(id) {
+            self.storage.insert_inner_with_id(testcase, true, id)
+        } else {
+            Err(Error::key_not_found(format!(
+                "Index {id} not found in enabled testcases"
+            )))
+        }
+    }
+
+    #[inline]
+    fn enable(&mut self, id: CorpusId) -> Result<(), Error> {
+        if let Some(testcase) = self.storage.disabled.remove(id) {
+            self.storage.insert_inner_with_id(testcase, false, id)
+        } else {
+            Err(Error::key_not_found(format!(
+                "Index {id} not found in disabled testcases"
+            )))
+        }
+    }
+}
+
 impl<I> HasTestcase<I> for InMemoryCorpus<I> {
     fn testcase(&self, id: CorpusId) -> Result<Ref<Testcase<I>>, Error> {
         Ok(self.get(id)?.borrow())
@@ -477,3 +562,157 @@ impl<I> InMemoryCorpus<I> {
         }
     }
 }
+
+#[cfg(test)]
+#[cfg(not(feature = "corpus_btreemap"))]
+mod tests {
+    use super::*;
+    use crate::{
+        Error,
+        corpus::Testcase,
+        inputs::{HasMutatorBytes, bytes::BytesInput},
+    };
+
+    /// Helper function to create a corpus with predefined test cases
+    #[cfg(not(feature = "corpus_btreemap"))]
+    fn setup_corpus() -> (InMemoryCorpus<BytesInput>, Vec<CorpusId>) {
+        let mut corpus = InMemoryCorpus::<BytesInput>::new();
+        let mut ids = Vec::new();
+
+        // Add initial test cases with distinct byte patterns ([1,2,3],[2,3,4],[3,4,5])
+        for i in 0..3u8 {
+            let input = BytesInput::new(vec![i + 1, i + 2, i + 3]);
+            let tc_id = corpus.add(Testcase::new(input)).unwrap();
+            ids.push(tc_id);
+        }
+
+        (corpus, ids)
+    }
+
+    /// Helper function to verify corpus counts
+    #[cfg(not(feature = "corpus_btreemap"))]
+    fn assert_corpus_counts(corpus: &InMemoryCorpus<BytesInput>, enabled: usize, disabled: usize) {
+        let total = enabled + disabled; // if a testcase is not in the enabled map, then it's in the disabled one.
+        assert_eq!(corpus.count(), enabled, "Wrong number of enabled testcases");
+        assert_eq!(
+            corpus.count_disabled(),
+            disabled,
+            "Wrong number of disabled testcases"
+        );
+        assert_eq!(corpus.count_all(), total, "Wrong total number of testcases");
+    }
+
+    #[test]
+    #[cfg(not(feature = "corpus_btreemap"))]
+    fn test_corpus_basic_operations() {
+        let (corpus, ids) = setup_corpus();
+        assert_corpus_counts(&corpus, 3, 0);
+
+        for id in &ids {
+            assert!(corpus.get(*id).is_ok(), "Failed to get testcase {id:?}");
+            assert!(
+                corpus.get_from_all(*id).is_ok(),
+                "Failed to get testcase from all {id:?}"
+            );
+        }
+
+        // Non-existent ID should fail
+        let invalid_id = CorpusId(999);
+        assert!(corpus.get(invalid_id).is_err());
+        assert!(corpus.get_from_all(invalid_id).is_err());
+    }
+
+    #[test]
+    #[cfg(not(feature = "corpus_btreemap"))]
+    fn test_corpus_disable_enable() -> Result<(), Error> {
+        let (mut corpus, ids) = setup_corpus();
+        let invalid_id = CorpusId(999);
+
+        corpus.disable(ids[1])?;
+        assert_corpus_counts(&corpus, 2, 1);
+
+        // Verify disabled testcase is not in enabled list but is in all list
+        assert!(
+            corpus.get(ids[1]).is_err(),
+            "Disabled testcase should not be accessible via get()"
+        );
+        assert!(
+            corpus.get_from_all(ids[1]).is_ok(),
+            "Disabled testcase should be accessible via get_from_all()"
+        );
+
+        // Other testcases are still accessible
+        assert!(corpus.get(ids[0]).is_ok());
+        assert!(corpus.get(ids[2]).is_ok());
+
+        corpus.enable(ids[1])?;
+        assert_corpus_counts(&corpus, 3, 0);
+
+        // Verify all testcases are accessible from the enabled map again
+        for id in &ids {
+            assert!(corpus.get(*id).is_ok());
+        }
+
+        // Corner cases
+        assert!(
+            corpus.disable(ids[1]).is_ok(),
+            "Should be able to disable testcase"
+        );
+        assert!(
+            corpus.disable(ids[1]).is_err(),
+            "Should not be able to disable already disabled testcase"
+        );
+        assert!(
+            corpus.enable(ids[0]).is_err(),
+            "Should not be able to enable already enabled testcase"
+        );
+        assert!(
+            corpus.disable(invalid_id).is_err(),
+            "Should not be able to disable non-existent testcase"
+        );
+        assert!(
+            corpus.enable(invalid_id).is_err(),
+            "Should not be able to enable non-existent testcase"
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    #[cfg(not(feature = "corpus_btreemap"))]
+    fn test_corpus_operations_after_disabled() -> Result<(), Error> {
+        let (mut corpus, ids) = setup_corpus();
+
+        corpus.disable(ids[0])?;
+        assert_corpus_counts(&corpus, 2, 1);
+
+        let removed = corpus.remove(ids[0])?;
+        let removed_data = removed.input().as_ref().unwrap().mutator_bytes();
+        assert_eq!(
+            removed_data,
+            &vec![1, 2, 3],
+            "Removed testcase has incorrect data"
+        );
+        assert_corpus_counts(&corpus, 2, 0);
+
+        let removed = corpus.remove(ids[1])?;
+        let removed_data = removed.input().as_ref().unwrap().mutator_bytes();
+        assert_eq!(
+            removed_data,
+            &vec![2, 3, 4],
+            "Removed testcase has incorrect data"
+        );
+        assert_corpus_counts(&corpus, 1, 0);
+
+        // Not possible to get removed testcases
+        assert!(corpus.get(ids[0]).is_err());
+        assert!(corpus.get_from_all(ids[0]).is_err());
+        assert!(corpus.get(ids[1]).is_err());
+        assert!(corpus.get_from_all(ids[1]).is_err());
+
+        // Only the third testcase should remain
+        assert!(corpus.get(ids[2]).is_ok());
+
+        Ok(())
+    }
+}

libafl/src/corpus/inmemory_ondisk.rs

@@ -20,7 +20,7 @@ use libafl_bolts::compress::GzipCompressor;
 use serde::{Deserialize, Serialize};
 
 use super::{
-    HasTestcase,
+    EnableDisableCorpus, HasTestcase,
     ondisk::{OnDiskMetadata, OnDiskMetadataFormat},
 };
 use crate::{
@@ -214,6 +214,29 @@ where
     }
 }
 
+impl<I> EnableDisableCorpus for InMemoryOnDiskCorpus<I>
+where
+    I: Input,
+{
+    #[inline]
+    fn disable(&mut self, id: CorpusId) -> Result<(), Error> {
+        self.inner.disable(id)?;
+        // Ensure testcase is saved to disk correctly with its new status
+        let testcase_cell = &mut self.get_from_all(id).unwrap().borrow_mut();
+        self.save_testcase(testcase_cell, Some(id))?;
+        Ok(())
+    }
+
+    #[inline]
+    fn enable(&mut self, id: CorpusId) -> Result<(), Error> {
+        self.inner.enable(id)?;
+        // Ensure testcase is saved to disk correctly with its new status
+        let testcase_cell = &mut self.get_from_all(id).unwrap().borrow_mut();
+        self.save_testcase(testcase_cell, Some(id))?;
+        Ok(())
+    }
+}
+
 impl<I> HasTestcase<I> for InMemoryOnDiskCorpus<I>
 where
     I: Input,

libafl/src/corpus/mod.rs

@@ -195,6 +195,15 @@ pub trait Corpus<I>: Sized {
     }
 }
 
+/// Marker trait for corpus implementations that actually support enable/disable functionality
+pub trait EnableDisableCorpus {
+    /// Disables a testcase, moving it to the disabled map
+    fn disable(&mut self, id: CorpusId) -> Result<(), Error>;
+
+    /// Enables a testcase, moving it to the enabled map
+    fn enable(&mut self, id: CorpusId) -> Result<(), Error>;
+}
+
 /// Trait for types which track the current corpus index
 pub trait HasCurrentCorpusId {
     /// Set the current corpus index; we have started processing this corpus entry

libafl/src/corpus/ondisk.rs

@@ -17,7 +17,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::{
     Error,
-    corpus::{CachedOnDiskCorpus, Corpus, CorpusId, HasTestcase, Testcase},
+    corpus::{CachedOnDiskCorpus, Corpus, CorpusId, EnableDisableCorpus, HasTestcase, Testcase},
     inputs::Input,
 };
 
@@ -188,6 +188,21 @@ where
     }
 }
 
+impl<I> EnableDisableCorpus for OnDiskCorpus<I>
+where
+    I: Input,
+{
+    #[inline]
+    fn disable(&mut self, id: CorpusId) -> Result<(), Error> {
+        self.inner.disable(id)
+    }
+
+    #[inline]
+    fn enable(&mut self, id: CorpusId) -> Result<(), Error> {
+        self.inner.enable(id)
+    }
+}
+
 impl<I> OnDiskCorpus<I> {
     /// Creates an [`OnDiskCorpus`].
     ///