From 66babddb02453f7034accb0a495d67c20b2d55b5 Mon Sep 17 00:00:00 2001
From: Alwin Berger <alwin.berger@tu-dortmund.de>
Date: Mon, 21 Mar 2022 00:00:32 +0100
Subject: [PATCH] add graph trace metadata, scheduler

---
 fuzzers/wcet_qemu_sys/src/bin/fuzzer.rs     |  4 +-
 fuzzers/wcet_qemu_sys/src/sysstate/graph.rs | 70 +++++++++++++++++++--
 2 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/fuzzers/wcet_qemu_sys/src/bin/fuzzer.rs b/fuzzers/wcet_qemu_sys/src/bin/fuzzer.rs
index 6faef9e6fe..a4b05b9be5 100644
--- a/fuzzers/wcet_qemu_sys/src/bin/fuzzer.rs
+++ b/fuzzers/wcet_qemu_sys/src/bin/fuzzer.rs
@@ -1,5 +1,6 @@
 //! A singlethreaded QEMU fuzzer that can auto-restart.
 
+use wcet_qemu_sys::sysstate::graph::GraphMaximizerCorpusScheduler;
 use wcet_qemu_sys::sysstate::graph::SysMapFeedback;
 use wcet_qemu_sys::sysstate::graph::SysGraphFeedbackState;
 use libafl::stats::SimpleStats;
@@ -398,7 +399,8 @@ fn fuzz(
     // A minimization+queue policy to get testcasess from the corpus
     // let scheduler = IndexesLenTimeMinimizerCorpusScheduler::new(PowerQueueCorpusScheduler::new());
     // let scheduler = TimeStateMaximizerCorpusScheduler::new(QueueCorpusScheduler::new());
-    let scheduler = QueueCorpusScheduler::new();
+    let scheduler = GraphMaximizerCorpusScheduler::new(QueueCorpusScheduler::new());
+    // let scheduler = QueueCorpusScheduler::new();
     
 
     // A fuzzer with feedbacks and a corpus scheduler
diff --git a/fuzzers/wcet_qemu_sys/src/sysstate/graph.rs b/fuzzers/wcet_qemu_sys/src/sysstate/graph.rs
index fff4377058..418e4a9954 100644
--- a/fuzzers/wcet_qemu_sys/src/sysstate/graph.rs
+++ b/fuzzers/wcet_qemu_sys/src/sysstate/graph.rs
@@ -1,5 +1,9 @@
 
 /// Feedbacks organizing SystemStates as a graph
+use crate::worst::MaxExecsLenFavFactor;
+use libafl::corpus::MinimizerCorpusScheduler;
+use libafl::bolts::HasRefCnt;
+use libafl::bolts::AsSlice;
 use libafl::bolts::ownedref::OwnedSlice;
 use libafl::inputs::BytesInput;
 use std::path::PathBuf;
@@ -93,6 +97,39 @@ impl PartialEq for SysGraphNode {
     }
 }
 
+// Wrapper around Vec<RefinedFreeRTOSSystemState> to attach as Metadata
+#[derive(Debug, Default, Serialize, Deserialize, Clone)]
+pub struct SysGraphMetadata {
+    indices: Vec<usize>,    // Hashed enumeration of States
+    tcref: isize,
+}
+impl SysGraphMetadata {
+    pub fn new(inner: Vec<NodeIndex>) -> Self{
+        Self {indices: inner.into_iter().map(|x| x.index()).collect(), tcref: 0}
+    }
+}
+impl AsSlice<usize> for SysGraphMetadata {
+    /// Convert the slice of system-states to a slice of hashes over enumerated states
+    fn as_slice(&self) -> &[usize] {
+        self.indices.as_slice()
+    }
+}
+
+impl HasRefCnt for SysGraphMetadata {
+    fn refcnt(&self) -> isize {
+        self.tcref
+    }
+
+    fn refcnt_mut(&mut self) -> &mut isize {
+        &mut self.tcref
+    }
+}
+
+libafl::impl_serdeany!(SysGraphMetadata);
+
+pub type GraphMaximizerCorpusScheduler<CS, I, S> =
+    MinimizerCorpusScheduler<CS, MaxExecsLenFavFactor<I>, I, SysGraphMetadata, S>;
+
 //============================= Graph Feedback
 
 /// Improved System State Graph
@@ -129,9 +166,10 @@ impl SysGraphFeedbackState
         }
     }
     /// Try adding a system state path from a [Vec<RefinedFreeRTOSSystemState>], return true if the path was interesting
-    fn update(&mut self, list: &Vec<RefinedFreeRTOSSystemState>, input: &Vec<u8>) -> bool {
+    fn update(&mut self, list: &Vec<RefinedFreeRTOSSystemState>, input: &Vec<u8>) -> (bool, Vec<NodeIndex>) {
         let mut current_index = self.entrypoint;
         let mut novel = false;
+        let mut trace : Vec<NodeIndex> = vec![current_index];
         for n in list {
             let mut matching : Option<NodeIndex> = None;
             for i in self.graph.neighbors_directed(current_index, Direction::Outgoing) {
@@ -153,8 +191,9 @@ impl SysGraphFeedbackState
                     novel |= self.graph[i].unite_interesting(&n, input);
                 }
             }
+            trace.push(current_index);
         }
-        return novel;
+        return (novel, trace);
     }
 }
 impl Named for SysGraphFeedbackState
@@ -177,11 +216,12 @@ impl FeedbackState for SysGraphFeedbackState
 #[derive(Serialize, Deserialize, Clone, Debug, Default)]
 pub struct SysMapFeedback
 {
-    name: String
+    name: String,
+    last_trace: Option<Vec<NodeIndex>>,
 }
 impl SysMapFeedback {
     pub fn new() -> Self {
-        Self {name: String::from("SysMapFeedback") }
+        Self {name: String::from("SysMapFeedback"), last_trace: None }
     }
 }
 
@@ -208,7 +248,27 @@ where
             .feedback_states_mut()
             .match_name_mut::<SysGraphFeedbackState>("SysMap")
             .unwrap();
-        Ok(feedbackstate.update(&observer.last_run, &observer.last_input))
+        let ret = feedbackstate.update(&observer.last_run, &observer.last_input);
+        self.last_trace = Some(ret.1);
+        Ok(ret.0)
+    }
+
+    /// Append to the testcase the generated metadata in case of a new corpus item
+    #[inline]
+    fn append_metadata(&mut self, _state: &mut S, testcase: &mut Testcase<I>) -> Result<(), Error> {
+        let a = self.last_trace.take();
+        match a {
+            Some(s) => testcase.metadata_mut().insert(SysGraphMetadata::new(s)),
+            None => (),
+        }
+        Ok(())
+    }
+
+    /// Discard the stored metadata in case that the testcase is not added to the corpus
+    #[inline]
+    fn discard_metadata(&mut self, _state: &mut S, _input: &I) -> Result<(), Error> {
+        self.last_trace = None;
+        Ok(())
     }
 }
 impl Named for SysMapFeedback