From 26aace6073ae3dd3ab8d1c4a384c39e6f734cb51 Mon Sep 17 00:00:00 2001 From: Addison Crump Date: Thu, 16 Feb 2023 00:14:26 +0100 Subject: [PATCH] Fix grimoire when used with on_replace/on_remove (#1075) --- libafl/src/inputs/generalized.rs | 28 ++----- libafl/src/mutators/grimoire.rs | 117 +++++++++++++--------------- libafl/src/stages/calibrate.rs | 12 ++- libafl/src/stages/generalization.rs | 50 +----------- 4 files changed, 71 insertions(+), 136 deletions(-) diff --git a/libafl/src/inputs/generalized.rs b/libafl/src/inputs/generalized.rs index 8fa46021fd..03b30e40b4 100644 --- a/libafl/src/inputs/generalized.rs +++ b/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, } @@ -112,14 +112,16 @@ where _state: &S, corpus_idx: CorpusId, ) -> Result { - base.metadata() + let meta = base + .metadata() .get::() .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 MutatedTransformPost for GeneralizedInputMetadata -where - S: HasCorpus, -{ - fn post_exec( - self, - state: &mut S, - _stage_idx: i32, - corpus_idx: Option, - ) -> 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(()) - } -} +impl MutatedTransformPost for GeneralizedInputMetadata where S: HasCorpus {} diff --git a/libafl/src/mutators/grimoire.rs b/libafl/src/mutators/grimoire.rs index f7825d0231..2ea854cdec 100644 --- a/libafl/src/mutators/grimoire.rs +++ b/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( state: &mut S, items: &mut Vec, gap_indices: &mut Vec, -) -> Result<(), Error> +) -> Result where S: HasMetadata + HasRand + HasCorpus, { - let rand_idx = state.rand_mut().next() as usize; - - let idx = { - let meta = state.metadata_mut().get_mut::().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() - }; + let idx = random_corpus_id!(state.corpus(), state.rand_mut()); 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(); - if let Some(other) = other_testcase - .metadata_mut() - .get::() - { - if other.generalized_len() > 0 { - let gen = other.generalized(); + let other_testcase = state.corpus().get(idx)?.borrow(); + if let Some(other) = other_testcase.metadata().get::() { + let gen = other.generalized(); - for (i, _) in gen - .iter() - .enumerate() - .filter(|&(_, x)| *x == GeneralizedItem::Gap) - { - 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(()); + for (i, _) in gen + .iter() + .enumerate() + .filter(|&(_, x)| *x == GeneralizedItem::Gap) + { + 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(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 = other_testcase - .metadata_mut() - .get::() - .unwrap(); - let gen = other.generalized(); + let other_testcase = state.corpus().get(idx)?.borrow(); + if let Some(other) = other_testcase.metadata().get::() { + let gen = other.generalized(); - if items.last() == Some(&GeneralizedItem::Gap) && gen.first() == Some(&GeneralizedItem::Gap) { - items.extend_from_slice(&gen[1..]); + if items.last() == Some(&GeneralizedItem::Gap) && gen.first() == Some(&GeneralizedItem::Gap) + { + 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 { 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(); diff --git a/libafl/src/stages/calibrate.rs b/libafl/src/stages/calibrate.rs index 6511772ea3..96102b3b54 100644 --- a/libafl/src/stages/calibrate.rs +++ b/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 - if state.corpus().get(corpus_idx)?.borrow_mut().fuzz_level() > 0 { + // 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().fuzz_level() > 0 + || state + .corpus() + .get(corpus_idx)? + .borrow() + .metadata() + .get::() + .map_or(false, |meta| meta.bitmap_size() != 0) + { return Ok(()); } diff --git a/libafl/src/stages/generalization.rs b/libafl/src/stages/generalization.rs index d525c0137b..c51ef85004 100644 --- a/libafl/src/stages/generalization.rs +++ b/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, -} - -crate::impl_serdeany!(GeneralizedIndexesMetadata); - -impl GeneralizedIndexesMetadata { - /// Create the metadata - #[must_use] - pub fn new() -> Self { - Self::default() - } -} - fn increment_by_offset(_list: &[Option], 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::() - .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::() { - drop(entry); - state - .metadata_mut() - .get_mut::() - .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)); - debug_assert!(meta.generalized().last() == Some(&GeneralizedItem::Gap)); + assert!(meta.generalized().first() == 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::() - .unwrap() - .indexes - .insert(corpus_idx); } Ok(())