diff --git a/libafl/src/corpus/cached.rs b/libafl/src/corpus/cached.rs index c9befc19e2..d9a07e1478 100644 --- a/libafl/src/corpus/cached.rs +++ b/libafl/src/corpus/cached.rs @@ -9,8 +9,8 @@ use serde::{Deserialize, Serialize}; use crate::{ Error, corpus::{ - Corpus, CorpusId, HasTestcase, Testcase, inmemory_ondisk::InMemoryOnDiskCorpus, - ondisk::OnDiskMetadataFormat, + Corpus, CorpusId, EnableDisableCorpus, HasTestcase, Testcase, + inmemory_ondisk::InMemoryOnDiskCorpus, ondisk::OnDiskMetadataFormat, }, inputs::Input, }; @@ -191,6 +191,23 @@ where } } +impl EnableDisableCorpus for CachedOnDiskCorpus +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 CachedOnDiskCorpus { /// Creates the [`CachedOnDiskCorpus`]. /// diff --git a/libafl/src/corpus/inmemory.rs b/libafl/src/corpus/inmemory.rs index d4760386bf..7a0480ba16 100644 --- a/libafl/src/corpus/inmemory.rs +++ b/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 TestcaseStorage { id } + #[cfg(not(feature = "corpus_btreemap"))] + fn insert_inner_with_id( + &mut self, + testcase: RefCell>, + 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>, + 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>, is_disabled: bool) -> CorpusId { @@ -456,6 +517,30 @@ impl Corpus for InMemoryCorpus { } } +impl EnableDisableCorpus for InMemoryCorpus { + #[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 HasTestcase for InMemoryCorpus { fn testcase(&self, id: CorpusId) -> Result>, Error> { Ok(self.get(id)?.borrow()) @@ -477,3 +562,157 @@ impl InMemoryCorpus { } } } + +#[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, Vec) { + let mut corpus = InMemoryCorpus::::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, 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(()) + } +} diff --git a/libafl/src/corpus/inmemory_ondisk.rs b/libafl/src/corpus/inmemory_ondisk.rs index 70b48bfc08..f2327776ef 100644 --- a/libafl/src/corpus/inmemory_ondisk.rs +++ b/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 EnableDisableCorpus for InMemoryOnDiskCorpus +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 HasTestcase for InMemoryOnDiskCorpus where I: Input, diff --git a/libafl/src/corpus/mod.rs b/libafl/src/corpus/mod.rs index d7670d1a10..7faf0c3ad8 100644 --- a/libafl/src/corpus/mod.rs +++ b/libafl/src/corpus/mod.rs @@ -195,6 +195,15 @@ pub trait Corpus: 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 diff --git a/libafl/src/corpus/ondisk.rs b/libafl/src/corpus/ondisk.rs index 72179a6531..e50791ff1f 100644 --- a/libafl/src/corpus/ondisk.rs +++ b/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 EnableDisableCorpus for OnDiskCorpus +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 OnDiskCorpus { /// Creates an [`OnDiskCorpus`]. ///