Merge pull request #475 from plotchy/corpus_cmp

feat: Replace testcase in corpus when comparisons are closer

src/feedback.rs

@@ -17,8 +17,8 @@ use libafl::{
     state::{HasCorpus, State},
     Error,
 };
-use libafl_bolts::Named;
-use serde::{de::DeserializeOwned, Serialize};
+use libafl_bolts::{impl_serdeany, Named};
+use serde::{de::DeserializeOwned, Deserialize, Serialize};
 use tracing::debug;
 
 /// Implements the feedback mechanism needed by ItyFuzz.
@@ -497,6 +497,7 @@ where
     S0: State
         + HasInfantStateState<Loc, Addr, VS, CI>
         + HasExecutionResult<Loc, Addr, VS, Out, CI>
+        + HasMetadata
         + UsesInput<Input = I0>,
     SC: Scheduler<State = InfantStateState<Loc, Addr, VS, CI>> + HasVote<InfantStateState<Loc, Addr, VS, CI>>,
     VS: Default + VMStateT + 'static,
@@ -528,6 +529,15 @@ where
         let mut cmp_interesting = false;
         let cov_interesting = false;
 
+        // Clear the result of any cached cmp analysis
+        if let Some(metadata) = state.metadata_map_mut().get_mut::<CmpMetadata>() {
+            metadata.set_cmp_interesting(false);
+        } else {
+            let mut metadata = CmpMetadata::new();
+            metadata.set_cmp_interesting(false);
+            state.metadata_map_mut().insert(metadata);
+        }
+
         // check if the current distance is smaller than the min_map
         for i in 0..MAP_SIZE {
             if self.current_map[i] < self.min_map[i] {
@@ -536,6 +546,15 @@ where
             }
         }
 
+        // cache the result of this testcase's cmp analysis
+        if let Some(metadata) = state.metadata_map_mut().get_mut::<CmpMetadata>() {
+            metadata.set_cmp_interesting(cmp_interesting);
+        } else {
+            let mut metadata = CmpMetadata::new();
+            metadata.set_cmp_interesting(cmp_interesting);
+            state.metadata_map_mut().insert(metadata);
+        }
+
         // if the current distance is smaller than the min_map, vote for the state
         if cmp_interesting {
             debug!("Voted for {} because of CMP", input.get_state_idx());
@@ -573,3 +592,24 @@ where
         Ok(false)
     }
 }
+
+/// Metadata for Coverage Comparisons
+///
+/// This is metadata attached to the global fuzz state
+#[derive(Clone, Debug, Serialize, Deserialize, Default)]
+pub struct CmpMetadata {
+    /// Used to cache the result of the last run's comparison mapping
+    pub cmp_interesting: bool,
+}
+
+impl CmpMetadata {
+    pub fn new() -> Self {
+        Default::default()
+    }
+
+    pub fn set_cmp_interesting(&mut self, cmp_interesting: bool) {
+        self.cmp_interesting = cmp_interesting;
+    }
+}
+
+impl_serdeany!(CmpMetadata);

src/fuzzer.rs

@@ -43,6 +43,7 @@ use tracing::info;
 
 use crate::{
     evm::{host::JMP_MAP, solution, utils::prettify_concise_inputs},
+    feedback::CmpMetadata,
     generic_vm::{vm_executor::MAP_SIZE, vm_state::VMStateT},
     input::{ConciseSerde, SolutionTx, VMInputT},
     minimizer::SequentialMinimizer,
@@ -167,11 +168,12 @@ where
     }
 
     /// Determine if a testcase should be replaced based on the minimizer map
-    /// If the new testcase has a higher fav factor, replace the old one
+    /// If the new testcase has a higher fav factor or any better comparison
+    /// results, replace the old one
     /// Returns None if the testcase should not be replaced
-    /// Returns Some((hash, new_fav_factor, testcase_idx)) if the testcase
-    /// should be replaced
-    pub fn should_replace(&self, input: &I, coverage: &[u8; MAP_SIZE]) -> Option<(u64, f64, usize)> {
+    /// Returns Some((hash, new_fav_factor, testcase_idx)) if
+    /// the testcase should be replaced
+    pub fn should_replace(&self, input: &I, coverage: &[u8; MAP_SIZE], state: &S) -> Option<(u64, f64, usize)> {
         let mut hasher = DefaultHasher::new();
         coverage.hash(&mut hasher);
         let hash = hasher.finish();
@@ -182,6 +184,13 @@ where
             if new_fav_factor > *fav_factor {
                 return Some((hash, new_fav_factor, *testcase_idx));
             }
+
+            // check if the testcase had a better comparison result
+            if let Some(metadata) = state.metadata_map().get::<CmpMetadata>() {
+                if metadata.cmp_interesting {
+                    return Some((hash, new_fav_factor, *testcase_idx));
+                }
+            }
         }
         None
     }
@@ -489,7 +498,7 @@ where
             // not interesting input, just check whether we should replace it due to better fav factor
             ExecuteInputResult::None => {
                 self.objective.discard_metadata(state, &input)?;
-                match self.should_replace(&input, unsafe { &JMP_MAP }) {
+                match self.should_replace(&input, unsafe { &JMP_MAP }, state) {
                     Some((hash, new_fav_factor, old_testcase_idx)) => {
                         let testcase = Testcase::new(input.clone());
                         let prev = state.corpus_mut().replace(old_testcase_idx.into(), testcase)?;