From 7aadf31246050b1a1b4848f05b5393b00ad69756 Mon Sep 17 00:00:00 2001
From: "Dongjia \"toka\" Zhang" <tokazerkje@outlook.com>
Date: Tue, 13 Sep 2022 09:39:17 +0200
Subject: [PATCH] Add track_stability option to CalibrationStage  (#781)

* add

* Update gramatron.rs

* Update emu.rs

* try

* clp
---
 libafl/src/mutators/gramatron.rs |  1 +
 libafl/src/stages/calibrate.rs   | 78 +++++++++++++++++++++-----------
 libafl_qemu/src/emu.rs           |  6 +--
 3 files changed, 53 insertions(+), 32 deletions(-)

diff --git a/libafl/src/mutators/gramatron.rs b/libafl/src/mutators/gramatron.rs
index 85c9bd25a0..f20edae97a 100644
--- a/libafl/src/mutators/gramatron.rs
+++ b/libafl/src/mutators/gramatron.rs
@@ -79,6 +79,7 @@ crate::impl_serdeany!(GramatronIdxMapMetadata);
 impl GramatronIdxMapMetadata {
     /// Creates a new [`struct@GramatronIdxMapMetadata`].
     #[must_use]
+    #[allow(clippy::or_fun_call)]
     pub fn new(input: &GramatronInput) -> Self {
         let mut map = HashMap::default();
         for i in 0..input.terminals().len() {
diff --git a/libafl/src/stages/calibrate.rs b/libafl/src/stages/calibrate.rs
index b51bf49039..5741c61a2a 100644
--- a/libafl/src/stages/calibrate.rs
+++ b/libafl/src/stages/calibrate.rs
@@ -73,6 +73,7 @@ where
     map_observer_name: String,
     map_name: String,
     stage_max: usize,
+    track_stability: bool,
     phantom: PhantomData<(I, O, OT, S)>,
 }
 
@@ -142,12 +143,13 @@ where
             .ok_or_else(|| Error::key_not_found("MapObserver not found".to_string()))?
             .to_vec();
 
+        let mut unstable_entries: Vec<usize> = vec![];
+        let map_len: usize = map_first.len();
         // Run CAL_STAGE_START - 1 times, increase by 2 for every time a new
         // run is found to be unstable, with CAL_STAGE_MAX total runs.
         let mut i = 1;
         let mut has_errors = false;
-        let mut unstable_entries: Vec<usize> = vec![];
-        let map_len: usize = map_first.len();
+
         while i < iter {
             let input = state
                 .corpus()
@@ -182,35 +184,37 @@ where
                 .observers_mut()
                 .post_exec_all(state, &input, &exit_kind)?;
 
-            let map = &executor
-                .observers()
-                .match_name::<O>(&self.map_observer_name)
-                .ok_or_else(|| Error::key_not_found("MapObserver not found".to_string()))?
-                .to_vec();
+            if self.track_stability {
+                let map = &executor
+                    .observers()
+                    .match_name::<O>(&self.map_observer_name)
+                    .ok_or_else(|| Error::key_not_found("MapObserver not found".to_string()))?
+                    .to_vec();
 
-            let history_map = &mut state
-                .named_metadata_mut()
-                .get_mut::<MapFeedbackMetadata<O::Entry>>(&self.map_name)
-                .unwrap()
-                .history_map;
+                let history_map = &mut state
+                    .named_metadata_mut()
+                    .get_mut::<MapFeedbackMetadata<O::Entry>>(&self.map_name)
+                    .unwrap()
+                    .history_map;
 
-            if history_map.len() < map_len {
-                history_map.resize(map_len, O::Entry::default());
-            }
+                if history_map.len() < map_len {
+                    history_map.resize(map_len, O::Entry::default());
+                }
 
-            for (idx, (first, (cur, history))) in map_first
-                .iter()
-                .zip(map.iter().zip(history_map.iter_mut()))
-                .enumerate()
-            {
-                if *first != *cur && *history != O::Entry::max_value() {
-                    *history = O::Entry::max_value();
-                    unstable_entries.push(idx);
-                };
-            }
+                for (idx, (first, (cur, history))) in map_first
+                    .iter()
+                    .zip(map.iter().zip(history_map.iter_mut()))
+                    .enumerate()
+                {
+                    if *first != *cur && *history != O::Entry::max_value() {
+                        *history = O::Entry::max_value();
+                        unstable_entries.push(idx);
+                    };
+                }
 
-            if !unstable_entries.is_empty() && iter < CAL_STAGE_MAX {
-                iter += 2;
+                if !unstable_entries.is_empty() && iter < CAL_STAGE_MAX {
+                    iter += 2;
+                }
             }
             i += 1;
         }
@@ -302,6 +306,26 @@ where
             map_observer_name: map_feedback.observer_name().to_string(),
             map_name: map_feedback.name().to_string(),
             stage_max: CAL_STAGE_START,
+            track_stability: true,
+            phantom: PhantomData,
+        }
+    }
+
+    /// Create a new [`CalibrationStage`], but without checking stability.
+    #[must_use]
+    pub fn ignore_stability<N, R>(map_feedback: &MapFeedback<I, N, O, R, S, O::Entry>) -> Self
+    where
+        O::Entry:
+            PartialEq + Default + Copy + 'static + Serialize + serde::de::DeserializeOwned + Debug,
+        R: Reducer<O::Entry>,
+        for<'it> O: AsIter<'it, Item = O::Entry>,
+        N: IsNovel<O::Entry>,
+    {
+        Self {
+            map_observer_name: map_feedback.observer_name().to_string(),
+            map_name: map_feedback.name().to_string(),
+            stage_max: CAL_STAGE_START,
+            track_stability: false,
             phantom: PhantomData,
         }
     }
diff --git a/libafl_qemu/src/emu.rs b/libafl_qemu/src/emu.rs
index f2018eb084..170941c86e 100644
--- a/libafl_qemu/src/emu.rs
+++ b/libafl_qemu/src/emu.rs
@@ -433,11 +433,7 @@ extern "C" fn gdb_cmd(buf: *const u8, len: usize, data: *const ()) -> i32 {
         let closure = &mut *(data as *mut Box<dyn for<'r> FnMut(&Emulator, &'r str) -> bool>);
         let cmd = std::str::from_utf8_unchecked(std::slice::from_raw_parts(buf, len));
         let emu = Emulator::new_empty();
-        if closure(&emu, cmd) {
-            1
-        } else {
-            0
-        }
+        i32::from(closure(&emu, cmd))
     }
 }