Refcnt for MapIndexesMetadata (#348)

* refcnt for MapIndexesMetadata

* fix clippy

fuzzers/frida_libpng/src/fuzzer.rs

@@ -13,8 +13,8 @@ use libafl::{
         tuples::{tuple_list, Merge},
     },
     corpus::{
-        ondisk::OnDiskMetadataFormat, Corpus, IndexesLenTimeMinimizerCorpusScheduler, OnDiskCorpus,
-        QueueCorpusScheduler,
+        ondisk::OnDiskMetadataFormat, CachedOnDiskCorpus, Corpus,
+        IndexesLenTimeMinimizerCorpusScheduler, OnDiskCorpus, QueueCorpusScheduler,
     },
     events::{llmp::LlmpRestartingEventManager, EventConfig},
     executors::{
@@ -347,7 +347,7 @@ unsafe fn fuzz(
                 // RNG
                 StdRand::with_seed(current_nanos()),
                 // Corpus that will be evolved, we keep it in memory for performance
-                OnDiskCorpus::new(PathBuf::from("./corpus_discovered")).unwrap(),
+                CachedOnDiskCorpus::new(PathBuf::from("./corpus_discovered"), 64).unwrap(),
                 // Corpus in which we store solutions (crashes in this example),
                 // on disk so the user can get them after stopping the fuzzer
                 OnDiskCorpus::new_save_meta(

libafl/src/bolts/mod.rs

@@ -38,6 +38,11 @@ pub trait HasLen {
     }
 }
 
+pub trait HasRefCnt {
+    fn refcnt(&self) -> isize;
+    fn refcnt_mut(&mut self) -> &mut isize;
+}
+
 /// Current time
 #[cfg(feature = "std")]
 #[must_use]

libafl/src/corpus/minimizer.rs

@@ -2,7 +2,7 @@
 // with testcases only from a subset of the total corpus.
 
 use crate::{
-    bolts::{rands::Rand, serdeany::SerdeAny, AsSlice, HasLen},
+    bolts::{rands::Rand, serdeany::SerdeAny, AsSlice, HasLen, HasRefCnt},
     corpus::{Corpus, CorpusScheduler, Testcase},
     feedbacks::MapIndexesMetadata,
     inputs::Input,
@@ -84,7 +84,7 @@ where
     CS: CorpusScheduler<I, S>,
     F: FavFactor<I>,
     I: Input,
-    M: AsSlice<usize> + SerdeAny,
+    M: AsSlice<usize> + SerdeAny + HasRefCnt,
     S: HasCorpus<C, I> + HasMetadata,
     C: Corpus<I>,
 {
@@ -98,7 +98,7 @@ where
     CS: CorpusScheduler<I, S>,
     F: FavFactor<I>,
     I: Input,
-    M: AsSlice<usize> + SerdeAny,
+    M: AsSlice<usize> + SerdeAny + HasRefCnt,
     S: HasCorpus<C, I> + HasMetadata + HasRand<R>,
     C: Corpus<I>,
     R: Rand,
@@ -148,13 +148,14 @@ where
     CS: CorpusScheduler<I, S>,
     F: FavFactor<I>,
     I: Input,
-    M: AsSlice<usize> + SerdeAny,
+    M: AsSlice<usize> + SerdeAny + HasRefCnt,
     S: HasCorpus<C, I> + HasMetadata + HasRand<R>,
     C: Corpus<I>,
     R: Rand,
 {
     /// Update the `Corpus` score using the `MinimizerCorpusScheduler`
     #[allow(clippy::unused_self)]
+    #[allow(clippy::cast_possible_wrap)]
     pub fn update_score(&self, state: &mut S, idx: usize) -> Result<(), Error> {
         // Create a new top rated meta if not existing
         if state.metadata().get::<TopRatedsMetadata>().is_none() {
@@ -165,7 +166,7 @@ where
         {
             let mut entry = state.corpus().get(idx)?.borrow_mut();
             let factor = F::compute(&mut *entry)?;
-            let meta = entry.metadata().get::<M>().ok_or_else(|| {
+            let meta = entry.metadata_mut().get_mut::<M>().ok_or_else(|| {
                 Error::KeyNotFound(format!(
                     "Metadata needed for MinimizerCorpusScheduler not found in testcase #{}",
                     idx
@@ -179,22 +180,52 @@ where
                     .map
                     .get(elem)
                 {
-                    if factor > F::compute(&mut *state.corpus().get(*old_idx)?.borrow_mut())? {
+                    let mut old = state.corpus().get(*old_idx)?.borrow_mut();
+                    if factor > F::compute(&mut *old)? {
                         continue;
                     }
-                }
 
-                new_favoreds.push((*elem, idx));
+                    let must_remove = {
+                        let old_meta = old.metadata_mut().get_mut::<M>().ok_or_else(|| {
+                            Error::KeyNotFound(format!(
+                                "Metadata needed for MinimizerCorpusScheduler not found in testcase #{}",
+                                old_idx
+                            ))
+                        })?;
+                        *old_meta.refcnt_mut() -= 1;
+                        old_meta.refcnt() <= 0
+                    };
+
+                    if must_remove {
+                        drop(old.metadata_mut().remove::<M>());
                     }
                 }
 
-        for pair in new_favoreds {
+                new_favoreds.push(*elem);
+            }
+
+            *meta.refcnt_mut() = new_favoreds.len() as isize;
+        }
+
+        if new_favoreds.is_empty() {
+            drop(
+                state
+                    .corpus()
+                    .get(idx)?
+                    .borrow_mut()
+                    .metadata_mut()
+                    .remove::<M>(),
+            );
+            return Ok(());
+        }
+
+        for elem in new_favoreds {
             state
                 .metadata_mut()
                 .get_mut::<TopRatedsMetadata>()
                 .unwrap()
                 .map
-                .insert(pair.0, pair.1);
+                .insert(elem, idx);
         }
         Ok(())
     }

libafl/src/feedbacks/map.rs

@@ -9,7 +9,7 @@ use num::Integer;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    bolts::{tuples::Named, AsSlice},
+    bolts::{tuples::Named, AsSlice, HasRefCnt},
     corpus::Testcase,
     events::{Event, EventFirer},
     executors::ExitKind,
@@ -76,6 +76,8 @@ where
 pub struct MapIndexesMetadata {
     /// The list of indexes.
     pub list: Vec<usize>,
+    /// A refcount used to know when remove this meta
+    pub tcref: isize,
 }
 
 crate::impl_serdeany!(MapIndexesMetadata);
@@ -87,11 +89,21 @@ impl AsSlice<usize> for MapIndexesMetadata {
     }
 }
 
+impl HasRefCnt for MapIndexesMetadata {
+    fn refcnt(&self) -> isize {
+        self.tcref
+    }
+
+    fn refcnt_mut(&mut self) -> &mut isize {
+        &mut self.tcref
+    }
+}
+
 impl MapIndexesMetadata {
     /// Creates a new [`struct@MapIndexesMetadata`].
     #[must_use]
     pub fn new(list: Vec<usize>) -> Self {
-        Self { list }
+        Self { list, tcref: 0 }
     }
 }