diff --git a/crates/libafl/src/corpus/cached.rs b/crates/libafl/src/corpus/cached.rs index 3995181491..2885a9dad4 100644 --- a/crates/libafl/src/corpus/cached.rs +++ b/crates/libafl/src/corpus/cached.rs @@ -29,21 +29,20 @@ impl CachedOnDiskCorpus where I: Input, { - fn cache_testcase<'a>( - &'a self, - testcase: &'a RefCell>, - id: CorpusId, - ) -> Result<(), Error> { - if testcase.borrow().input().is_none() { - self.load_input_into(&mut testcase.borrow_mut())?; + fn cache_testcase_input<'a>(&'a self, testcase: &'a mut Testcase) -> Result<(), Error> { + let id = testcase + .corpus_id() + .ok_or(Error::unknown("The testcase is not associated with an id"))?; + if testcase.input().is_none() { + self.inner.load_input_into(testcase)?; let mut borrowed_num = 0; while self.cached_indexes.borrow().len() >= self.cache_max_len { - let removed = self.cached_indexes.borrow_mut().pop_front().unwrap(); + let to_be_evicted = self.cached_indexes.borrow_mut().pop_front().unwrap(); - if let Ok(mut borrowed) = self.inner.get_from_all(removed)?.try_borrow_mut() { + if let Ok(mut borrowed) = self.inner.get_from_all(to_be_evicted)?.try_borrow_mut() { *borrowed.input_mut() = None; } else { - self.cached_indexes.borrow_mut().push_back(removed); + self.cached_indexes.borrow_mut().push_back(to_be_evicted); borrowed_num += 1; if self.cache_max_len == borrowed_num { break; @@ -106,16 +105,12 @@ where /// Get by id; considers only enabled testcases #[inline] fn get(&self, id: CorpusId) -> Result<&RefCell>, Error> { - let testcase = { self.inner.get(id)? }; - self.cache_testcase(testcase, id)?; - Ok(testcase) + self.inner.get(id) } /// Get by id; considers both enabled and disabled testcases #[inline] fn get_from_all(&self, id: CorpusId) -> Result<&RefCell>, Error> { - let testcase = { self.inner.get_from_all(id)? }; - self.cache_testcase(testcase, id)?; - Ok(testcase) + self.inner.get_from_all(id) } /// Current testcase scheduled @@ -169,7 +164,8 @@ where #[inline] fn load_input_into(&self, testcase: &mut Testcase) -> Result<(), Error> { - self.inner.load_input_into(testcase) + self.cache_testcase_input(testcase)?; + Ok(()) } #[inline] diff --git a/crates/libafl/src/corpus/inmemory.rs b/crates/libafl/src/corpus/inmemory.rs index d98fee9d02..8360c8f52b 100644 --- a/crates/libafl/src/corpus/inmemory.rs +++ b/crates/libafl/src/corpus/inmemory.rs @@ -257,6 +257,7 @@ impl TestcaseStorage { #[cfg(not(feature = "corpus_btreemap"))] fn insert_inner(&mut self, testcase: RefCell>, is_disabled: bool) -> CorpusId { let id = CorpusId::from(self.progressive_id); + testcase.borrow_mut().set_corpus_id(Some(id)); self.progressive_id += 1; let corpus = if is_disabled { &mut self.disabled diff --git a/crates/libafl/src/corpus/testcase.rs b/crates/libafl/src/corpus/testcase.rs index dfcaeab4eb..0b7dc8ddbb 100644 --- a/crates/libafl/src/corpus/testcase.rs +++ b/crates/libafl/src/corpus/testcase.rs @@ -52,6 +52,8 @@ pub struct Testcase { scheduled_count: usize, /// Number of executions done at discovery time executions: u64, + /// The [`CorpusId`], if added to a corpus + corpus_id: Option, /// Parent [`CorpusId`], if known parent_id: Option, /// If the testcase is "disabled" @@ -207,6 +209,18 @@ impl Testcase { self.disabled = disabled; } + /// Get the corpus id of the testcase + #[inline] + pub fn corpus_id(&self) -> Option { + self.corpus_id + } + + /// Set the corpus id of the testcase + #[inline] + pub fn set_corpus_id(&mut self, id: Option) { + self.corpus_id = id; + } + /// Get the hit feedbacks #[inline] #[cfg(feature = "track_hit_feedbacks")] @@ -250,6 +264,7 @@ impl Testcase { cached_len: None, executions: 0, scheduled_count: 0, + corpus_id: None, parent_id: None, disabled: false, objectives_found: 0, @@ -275,6 +290,7 @@ impl Testcase { cached_len: None, executions: 0, scheduled_count: 0, + corpus_id: None, parent_id: Some(parent_id), disabled: false, objectives_found: 0, @@ -302,6 +318,7 @@ impl Testcase { cached_len: None, executions: 0, scheduled_count: 0, + corpus_id: None, parent_id: None, disabled: false, objectives_found: 0, @@ -350,6 +367,7 @@ impl Default for Testcase { exec_time: None, cached_len: None, scheduled_count: 0, + corpus_id: None, parent_id: None, #[cfg(feature = "std")] file_path: None,