Fix grimoire when used with on_replace/on_remove (#1075)

libafl/src/inputs/generalized.rs

@@ -5,7 +5,7 @@ use alloc::vec::Vec;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    corpus::{Corpus, CorpusId, Testcase},
+    corpus::{CorpusId, Testcase},
     impl_serdeany,
     inputs::BytesInput,
     stages::mutational::{MutatedTransform, MutatedTransformPost},
@@ -23,7 +23,7 @@ pub enum GeneralizedItem {
 }
 
 /// Metadata regarding the generalised content of an input
-#[derive(Serialize, Deserialize, Clone, Debug, Default, PartialEq, Eq, Hash)]
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash)]
 pub struct GeneralizedInputMetadata {
     generalized: Vec<GeneralizedItem>,
 }
@@ -112,14 +112,16 @@ where
         _state: &S,
         corpus_idx: CorpusId,
     ) -> Result<Self, Error> {
-        base.metadata()
+        let meta = base
+            .metadata()
             .get::<GeneralizedInputMetadata>()
             .ok_or_else(|| {
                 Error::key_not_found(format!(
                     "Couldn't find the GeneralizedInputMetadata for corpus entry {corpus_idx}",
                 ))
             })
-            .cloned()
+            .cloned()?;
+        Ok(meta)
     }
 
     fn try_transform_into(self, _state: &S) -> Result<(BytesInput, Self::Post), Error> {
@@ -127,20 +129,4 @@ where
     }
 }
 
-impl<S> MutatedTransformPost<S> for GeneralizedInputMetadata
+impl<S> MutatedTransformPost<S> for GeneralizedInputMetadata where S: HasCorpus {}
-where
-    S: HasCorpus,
-{
-    fn post_exec(
-        self,
-        state: &mut S,
-        _stage_idx: i32,
-        corpus_idx: Option<CorpusId>,
-    ) -> Result<(), Error> {
-        if let Some(corpus_idx) = corpus_idx {
-            let mut testcase = state.corpus().get(corpus_idx)?.borrow_mut();
-            testcase.metadata_mut().insert(self);
-        }
-        Ok(())
-    }
-}

libafl/src/mutators/grimoire.rs

@@ -9,7 +9,7 @@ use crate::{
     corpus::Corpus,
     inputs::{GeneralizedInputMetadata, GeneralizedItem},
     mutators::{token_mutations::Tokens, MutationResult, Mutator},
-    stages::generalization::GeneralizedIndexesMetadata,
+    random_corpus_id,
     state::{HasCorpus, HasMetadata, HasRand},
     Error,
 };
@@ -22,60 +22,43 @@ fn extend_with_random_generalized<S>(
     state: &mut S,
     items: &mut Vec<GeneralizedItem>,
     gap_indices: &mut Vec<usize>,
-) -> Result<(), Error>
+) -> Result<MutationResult, Error>
 where
     S: HasMetadata + HasRand + HasCorpus,
 {
-    let rand_idx = state.rand_mut().next() as usize;
+    let idx = random_corpus_id!(state.corpus(), state.rand_mut());
-
-    let idx = {
-        let meta = state.metadata_mut().get_mut::<GeneralizedIndexesMetadata>().ok_or_else(|| {
-            Error::key_not_found("GeneralizedIndexesMetadata needed by extend_with_random_generalized() not found, make sure that you have GeneralizationStage in")
-        })?;
-
-        *meta
-            .indexes
-            .iter()
-            .nth(rand_idx % meta.indexes.len())
-            .unwrap()
-    };
 
     if state.rand_mut().below(100) > CHOOSE_SUBINPUT_PROB {
         if state.rand_mut().below(100) < 50 {
             let rand1 = state.rand_mut().next() as usize;
             let rand2 = state.rand_mut().next() as usize;
 
-            let mut other_testcase = state.corpus().get(idx)?.borrow_mut();
+            let other_testcase = state.corpus().get(idx)?.borrow();
-            if let Some(other) = other_testcase
+            if let Some(other) = other_testcase.metadata().get::<GeneralizedInputMetadata>() {
-                .metadata_mut()
+                let gen = other.generalized();
-                .get::<GeneralizedInputMetadata>()
-            {
-                if other.generalized_len() > 0 {
-                    let gen = other.generalized();
 
-                    for (i, _) in gen
+                for (i, _) in gen
-                        .iter()
+                    .iter()
-                        .enumerate()
+                    .enumerate()
-                        .filter(|&(_, x)| *x == GeneralizedItem::Gap)
+                    .filter(|&(_, x)| *x == GeneralizedItem::Gap)
-                    {
+                {
-                        gap_indices.push(i);
+                    gap_indices.push(i);
-                    }
-                    let min_idx = gap_indices[rand1 % gap_indices.len()];
-                    let max_idx = gap_indices[rand2 % gap_indices.len()];
-                    let (mut min_idx, max_idx) = (min(min_idx, max_idx), max(min_idx, max_idx));
-
-                    gap_indices.clear();
-
-                    if items.last() == Some(&GeneralizedItem::Gap) {
-                        min_idx += 1;
-                    }
-                    items.extend_from_slice(&gen[min_idx..=max_idx]);
-
-                    debug_assert!(items.first() == Some(&GeneralizedItem::Gap));
-                    debug_assert!(items.last() == Some(&GeneralizedItem::Gap));
-
-                    return Ok(());
                 }
+                let min_idx = gap_indices[rand1 % gap_indices.len()];
+                let max_idx = gap_indices[rand2 % gap_indices.len()];
+                let (mut min_idx, max_idx) = (min(min_idx, max_idx), max(min_idx, max_idx));
+
+                gap_indices.clear();
+
+                if items.last() == Some(&GeneralizedItem::Gap) {
+                    min_idx += 1;
+                }
+                items.extend_from_slice(&gen[min_idx..=max_idx]);
+
+                debug_assert!(items.first() == Some(&GeneralizedItem::Gap));
+                debug_assert!(items.last() == Some(&GeneralizedItem::Gap));
+
+                return Ok(MutationResult::Mutated);
             }
         }
 
@@ -93,28 +76,29 @@ where
                 debug_assert!(items.first() == Some(&GeneralizedItem::Gap));
                 debug_assert!(items.last() == Some(&GeneralizedItem::Gap));
 
-                return Ok(());
+                return Ok(MutationResult::Mutated);
             }
         }
     }
 
-    let mut other_testcase = state.corpus().get(idx)?.borrow_mut();
+    let other_testcase = state.corpus().get(idx)?.borrow();
-    let other = other_testcase
+    if let Some(other) = other_testcase.metadata().get::<GeneralizedInputMetadata>() {
-        .metadata_mut()
+        let gen = other.generalized();
-        .get::<GeneralizedInputMetadata>()
-        .unwrap();
-    let gen = other.generalized();
 
-    if items.last() == Some(&GeneralizedItem::Gap) && gen.first() == Some(&GeneralizedItem::Gap) {
+        if items.last() == Some(&GeneralizedItem::Gap) && gen.first() == Some(&GeneralizedItem::Gap)
-        items.extend_from_slice(&gen[1..]);
+        {
+            items.extend_from_slice(&gen[1..]);
+        } else {
+            items.extend_from_slice(gen);
+        }
+
+        debug_assert!(items.first() == Some(&GeneralizedItem::Gap));
+        debug_assert!(items.last() == Some(&GeneralizedItem::Gap));
+
+        Ok(MutationResult::Mutated)
     } else {
-        items.extend_from_slice(gen);
+        Ok(MutationResult::Skipped)
     }
-
-    debug_assert!(items.first() == Some(&GeneralizedItem::Gap));
-    debug_assert!(items.last() == Some(&GeneralizedItem::Gap));
-
-    Ok(())
 }
 
 /// Extend the generalized input with another random one from the corpus
@@ -137,9 +121,7 @@ where
             state,
             generalised_meta.generalized_mut(),
             &mut self.gap_indices,
-        )?;
+        )
-
-        Ok(MutationResult::Mutated)
     }
 }
 
@@ -202,7 +184,11 @@ where
             self.scratch.extend_from_slice(&gen[selected + 1..]);
             gen.truncate(selected);
 
-            extend_with_random_generalized(state, gen, &mut self.gap_indices)?;
+            if extend_with_random_generalized(state, gen, &mut self.gap_indices)?
+                == MutationResult::Skipped
+            {
+                gen.push(GeneralizedItem::Gap);
+            }
 
             gen.extend_from_slice(&self.scratch);
             self.scratch.clear();
@@ -358,10 +344,10 @@ where
     ) -> Result<MutationResult, Error> {
         let gen = generalised_meta.generalized_mut();
 
-        for (i, _) in gen
+        for i in gen
             .iter()
             .enumerate()
-            .filter(|&(_, x)| *x == GeneralizedItem::Gap)
+            .filter_map(|(i, x)| (*x == GeneralizedItem::Gap).then_some(i))
         {
             self.gap_indices.push(i);
         }
@@ -369,6 +355,7 @@ where
             self.gap_indices[state.rand_mut().below(self.gap_indices.len() as u64) as usize];
         let max_idx =
             self.gap_indices[state.rand_mut().below(self.gap_indices.len() as u64) as usize];
+
         let (min_idx, max_idx) = (min(min_idx, max_idx), max(min_idx, max_idx));
 
         self.gap_indices.clear();

libafl/src/stages/calibrate.rs

@@ -106,8 +106,16 @@ where
         mgr: &mut EM,
         corpus_idx: CorpusId,
     ) -> Result<(), Error> {
-        // Run this stage only once for each corpus entry
+        // Run this stage only once for each corpus entry and only if we haven't already inspected it
-        if state.corpus().get(corpus_idx)?.borrow_mut().fuzz_level() > 0 {
+        if state.corpus().get(corpus_idx)?.borrow().fuzz_level() > 0
+            || state
+                .corpus()
+                .get(corpus_idx)?
+                .borrow()
+                .metadata()
+                .get::<SchedulerTestcaseMetaData>()
+                .map_or(false, |meta| meta.bitmap_size() != 0)
+        {
             return Ok(());
         }
 

libafl/src/stages/generalization.rs

@@ -6,9 +6,6 @@ use alloc::{
 };
 use core::{fmt::Debug, marker::PhantomData};
 
-use hashbrown::HashSet;
-use serde::{Deserialize, Serialize};
-
 #[cfg(feature = "introspection")]
 use crate::monitors::PerfFeature;
 use crate::{
@@ -27,23 +24,6 @@ use crate::{
 
 const MAX_GENERALIZED_LEN: usize = 8192;
 
-/// A state metadata holding the set of indexes related to the generalized corpus entries
-#[derive(Debug, Default, Serialize, Deserialize)]
-pub struct GeneralizedIndexesMetadata {
-    /// The set of indexes
-    pub indexes: HashSet<CorpusId>,
-}
-
-crate::impl_serdeany!(GeneralizedIndexesMetadata);
-
-impl GeneralizedIndexesMetadata {
-    /// Create the metadata
-    #[must_use]
-    pub fn new() -> Self {
-        Self::default()
-    }
-}
-
 fn increment_by_offset(_list: &[Option<u8>], idx: usize, off: u8) -> usize {
     idx + 1 + off as usize
 }
@@ -97,31 +77,12 @@ where
         manager: &mut EM,
         corpus_idx: CorpusId,
     ) -> Result<(), Error> {
-        if state
-            .metadata()
-            .get::<GeneralizedIndexesMetadata>()
-            .is_none()
-        {
-            state.add_metadata(GeneralizedIndexesMetadata::new());
-        }
-
         let (mut payload, original, novelties) = {
             start_timer!(state);
             state.corpus().get(corpus_idx)?.borrow_mut().load_input()?;
             mark_feature_time!(state, PerfFeature::GetInputFromCorpus);
             let mut entry = state.corpus().get(corpus_idx)?.borrow_mut();
 
-            if entry.metadata().contains::<GeneralizedInputMetadata>() {
-                drop(entry);
-                state
-                    .metadata_mut()
-                    .get_mut::<GeneralizedIndexesMetadata>()
-                    .unwrap()
-                    .indexes
-                    .insert(corpus_idx);
-                return Ok(());
-            }
-
             let input = entry.input_mut().as_mut().unwrap();
 
             let payload: Vec<_> = input.bytes().iter().map(|&x| Some(x)).collect();
@@ -327,19 +288,12 @@ where
             {
                 let meta = GeneralizedInputMetadata::generalized_from_options(&payload);
 
-                debug_assert!(meta.generalized().first() == Some(&GeneralizedItem::Gap));
+                assert!(meta.generalized().first() == Some(&GeneralizedItem::Gap));
-                debug_assert!(meta.generalized().last() == Some(&GeneralizedItem::Gap));
+                assert!(meta.generalized().last() == Some(&GeneralizedItem::Gap));
 
                 let mut entry = state.corpus().get(corpus_idx)?.borrow_mut();
                 entry.metadata_mut().insert(meta);
             }
-
-            state
-                .metadata_mut()
-                .get_mut::<GeneralizedIndexesMetadata>()
-                .unwrap()
-                .indexes
-                .insert(corpus_idx);
         }
 
         Ok(())