Fix caching logic for CachedOnDiskCorpus (#3408)

Author: henryhchchc

crates/libafl/src/corpus/cached.rs

@@ -29,21 +29,20 @@ impl<I> CachedOnDiskCorpus<I>
 where
     I: Input,
 {
-    fn cache_testcase<'a>(
-        &'a self,
-        testcase: &'a RefCell<Testcase<I>>,
-        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<I>) -> 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<Testcase<I>>, 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<Testcase<I>>, 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<I>) -> Result<(), Error> {
-        self.inner.load_input_into(testcase)
+        self.cache_testcase_input(testcase)?;
+        Ok(())
     }
 
     #[inline]

crates/libafl/src/corpus/inmemory.rs

@@ -257,6 +257,7 @@ impl<I> TestcaseStorage<I> {
     #[cfg(not(feature = "corpus_btreemap"))]
     fn insert_inner(&mut self, testcase: RefCell<Testcase<I>>, 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

crates/libafl/src/corpus/testcase.rs

@@ -52,6 +52,8 @@ pub struct Testcase<I> {
     scheduled_count: usize,
     /// Number of executions done at discovery time
     executions: u64,
+    /// The [`CorpusId`], if added to a corpus
+    corpus_id: Option<CorpusId>,
     /// Parent [`CorpusId`], if known
     parent_id: Option<CorpusId>,
     /// If the testcase is "disabled"
@@ -207,6 +209,18 @@ impl<I> Testcase<I> {
         self.disabled = disabled;
     }
 
+    /// Get the corpus id of the testcase
+    #[inline]
+    pub fn corpus_id(&self) -> Option<CorpusId> {
+        self.corpus_id
+    }
+
+    /// Set the corpus id of the testcase
+    #[inline]
+    pub fn set_corpus_id(&mut self, id: Option<CorpusId>) {
+        self.corpus_id = id;
+    }
+
     /// Get the hit feedbacks
     #[inline]
     #[cfg(feature = "track_hit_feedbacks")]
@@ -250,6 +264,7 @@ impl<I> Testcase<I> {
             cached_len: None,
             executions: 0,
             scheduled_count: 0,
+            corpus_id: None,
             parent_id: None,
             disabled: false,
             objectives_found: 0,
@@ -275,6 +290,7 @@ impl<I> Testcase<I> {
             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<I> Testcase<I> {
             cached_len: None,
             executions: 0,
             scheduled_count: 0,
+            corpus_id: None,
             parent_id: None,
             disabled: false,
             objectives_found: 0,
@@ -350,6 +367,7 @@ impl<I> Default for Testcase<I> {
             exec_time: None,
             cached_len: None,
             scheduled_count: 0,
+            corpus_id: None,
             parent_id: None,
             #[cfg(feature = "std")]
             file_path: None,