improve stg parsing

fuzzers/FRET/src/fuzzer.rs

@@ -43,6 +43,7 @@ use std::time::{SystemTime, UNIX_EPOCH};
 use clap::{Parser, Subcommand};
 use csv::Reader;
 use petgraph::dot::{Dot, Config};
+use crate::systemstate::stg::STGFeedbackState;
 
 // Constants ================================================================================
 
@@ -173,7 +174,7 @@ struct Cli {
     #[command(subcommand)]
     command: Commands,
 }
-#[derive(Subcommand)]
+#[derive(Subcommand,Clone)]
 enum Commands {
     /// run a single input
     Showmap {
@@ -255,6 +256,23 @@ macro_rules! do_dump_graph {
     };
 }
 
+/// Takes a state and a bool, writes out the current graph
+macro_rules! do_dump_stg {
+    ($state:expr, $cli:expr, $c:expr) => {
+        #[cfg(feature = "trace_abbs")]
+        if $cli.dump_graph {
+            let dump_path = $cli.dump_name.clone().unwrap().with_extension(if $c=="" {"stg"} else {$c});
+            println!("Dumping graph to {:?}", &dump_path);
+            if let Some(md) = $state.named_metadata_map_mut().get_mut::<STGFeedbackState>("stgfeedbackstate") {
+                let out = md.graph.map(|i,x| x.pretty_print(), |_,_| "");
+                let outs = Dot::with_config(&out, &[Config::EdgeNoLabel]).to_string();
+                let outs = outs.replace(';',"\\n");
+                fs::write(dump_path,outs).expect("Failed to write graph");
+            }
+        }
+    };
+}
+
 /// Takes a state and a bool, writes out top rated inputs
 macro_rules! do_dump_toprated {
     ($state:expr, $cli:expr, $c:expr) => {
@@ -656,7 +674,7 @@ pub fn fuzz() {
         #[cfg(not(feature = "fuzz_int"))]
         let mut stages = tuple_list!(StdMutationalStage::new(mutator));
 
-        if let Commands::Showmap { input } = cli.command {
+        if let Commands::Showmap { input } = cli.command.clone() {
             let s = input.as_os_str();
             let show_input = if s=="-" {
                     let mut buf = Vec::<u8>::new();
@@ -669,25 +687,14 @@ pub fn fuzz() {
                 };
             fuzzer.evaluate_input(&mut state, &mut executor, &mut mgr, BytesInput::new(show_input))
                 .unwrap();
-            if cli.dump_times {
+            do_dump_times!(state, &cli, "");
-                let td = cli.dump_name.clone().expect("Dump name not give").with_extension("case.time");
+            do_dump_graph!(state, &cli, "");
-                let mut file = OpenOptions::new()
+            do_dump_stg!(state, &cli, "");
-                    .read(true)
-                    .write(true)
-                    .create(true)
-                    .append(true)
-                    .open(td).expect("Could not open timedump");
-                if let Ok(ichist) = state.metadata_mut::<IcHist>() {
-                    for i in ichist.0.drain(..) {
-                        writeln!(file, "{},{}", i.0, i.1).expect("Write to dump failed");
-                    }
-                }
-            }
         } else if let Commands::Fuzz { random, time, seed } = cli.command {
             if let Some(se) = seed {
                 unsafe {
                     let mut rng = StdRng::seed_from_u64(se);
-                    for i in 0..100 {
+                    for i in 0..20 {
                         let inp = BytesInput::new(vec![rng.gen::<u8>(); MAX_INPUT_SIZE]);
                         fuzzer.evaluate_input(&mut state, &mut executor, &mut mgr, inp).unwrap();
                     }
@@ -746,6 +753,8 @@ pub fn fuzz() {
                                 do_dump_case!(state, &cli, &d);
                                 let _d = format!("{}.graph",marker);
                                 do_dump_graph!(state, &cli, &_d);
+                                let _d = format!("{}.stg",marker);
+                                do_dump_stg!(state, &cli, &_d);
                                 let d = format!("{}.toprated",marker);
                                 do_dump_toprated!(state, &cli, &d);
                             };
@@ -766,6 +775,7 @@ pub fn fuzz() {
                                 }
                                 do_dump_case!(state, &cli, "");
                                 do_dump_graph!(state, &cli, "");
+                                do_dump_stg!(state, &cli, "");
                                 do_dump_toprated!(state, &cli, "");
                             }
                         }
@@ -773,6 +783,7 @@ pub fn fuzz() {
                     do_dump_times!(state, &cli, "");
                     do_dump_case!(state, &cli, "");
                     do_dump_graph!(state, &cli, "");
+                    do_dump_stg!(state, &cli, "");
                     do_dump_toprated!(state, &cli, "");
                 },
             }

fuzzers/FRET/src/systemstate/stg.rs

@@ -56,16 +56,22 @@ use std::rc::Rc;
 use libafl_bolts::rands::Rand;
 
 //============================= Data Structures
-#[derive(Serialize, Deserialize, Clone, Debug, Default)]
+#[derive(Serialize, Deserialize, Clone, Debug, Default, Hash)]
 pub struct STGNode
 {
     base: RefinedFreeRTOSSystemState,
     abb: AtomicBasicBlock,
 }
 impl STGNode {
-    fn pretty_print(&self) -> String {
+    pub fn pretty_print(&self) -> String {
         format!("{}\n{:x}-{:x}", self.base.current_task.0.task_name, self.abb.start, self.abb.ends.iter().next().unwrap_or_else(||&0))
     }
+    fn calculate_hash(&self) -> u64 {
+        let mut s = DefaultHasher::new();
+        self.base.hash(&mut s);
+        self.abb.hash(&mut s);
+        s.finish()
+    }
 }
 impl PartialEq for STGNode {
     fn eq(&self, other: &STGNode) -> bool {
@@ -73,6 +79,49 @@ impl PartialEq for STGNode {
     }
 }
 
+/// Shared Metadata for a systemstateFeedback
+#[derive(Debug, Serialize, Deserialize, SerdeAny, Clone)]
+pub struct STGFeedbackState
+{
+    pub graph: DiGraph<STGNode, ()>,
+    index: HashMap<u64, NodeIndex>,
+    entrypoint: NodeIndex,
+    exit: NodeIndex,
+}
+
+impl Default for STGFeedbackState {
+    fn default() -> STGFeedbackState {
+        let mut graph = DiGraph::new();
+        let mut entry = STGNode::default();
+        entry.base.current_task.0.task_name="Start".to_string();
+        let mut exit = STGNode::default();
+        exit.base.current_task.0.task_name="End".to_string();
+
+        let h_entry = entry.calculate_hash();
+        let h_exit = exit.calculate_hash();
+
+        let entrypoint = graph.add_node(entry);
+        let exit = graph.add_node(exit);
+
+        let index = HashMap::from([(h_entry, entrypoint), (h_exit, exit)]);
+
+        STGFeedbackState {
+            graph,
+            index,
+            entrypoint,
+            exit
+        }
+    }
+}
+
+impl Named for STGFeedbackState
+{
+    #[inline]
+    fn name(&self) -> &str {
+        "stgfeedbackstate"
+    }
+}
+
 //============================= Graph Feedback
 
 /// A Feedback reporting novel System-State Transitions. Depends on [`QemuSystemStateObserver`]
@@ -86,24 +135,52 @@ impl StgFeedback {
     pub fn new() -> Self {
         Self {name: String::from("SysMapFeedback"), last_trace: None }
     }
-    fn trace_to_stg(trace: &Vec<RefinedFreeRTOSSystemState>, map: Vec<(String, Rc<AtomicBasicBlock>, usize, u64)>) -> DiGraph<STGNode,()>{
+    /// params:
-        let mut graph = DiGraph::new();
+    /// tarce of system states
-        let mut entry = STGNode::default();
+    /// list of all atomic basic blocks with asociated state index and ticks
-        entry.base.current_task.0.task_name="Start".to_string();
+    /// produces:
-        let mut exit = STGNode::default();
+    /// tarce of node indexes representing the path trough the graph
-        exit.base.current_task.0.task_name="End".to_string();
+    /// newly discovered node?
-        let entry = graph.add_node(entry);
+    /// side effect:
-        let exit = graph.add_node(exit);
+    /// the graph gets new nodes
-
+    fn update_stg(trace: &Vec<RefinedFreeRTOSSystemState>, map: Vec<(String, Rc<AtomicBasicBlock>, usize, u64)>, fbs: &mut STGFeedbackState) -> (Vec<NodeIndex>, bool) {
-        let mut last = entry;
+        let mut returntrace = vec![fbs.entrypoint];
-
+        let mut new_stg_edge = false;
-        for (i,state) in trace.iter().enumerate() {
+        // add all missing state+abb combinations to the graph
-            if let Some(marker) = map.iter().filter(|(_,_,index,_)| index==&i).next() {
+        for (i,state) in trace.iter().enumerate() { // Iterate states first, keep the trace order intact
-                graph.add_node(STGNode {base: state.clone(), abb: (*marker.1).clone()});
+            if let Some(marker) = map.iter().filter(|(_,_,index,_)| index==&i).next() { // 
+                let node = STGNode {base: state.clone(), abb: (*marker.1).clone()};
+                let h_node = node.calculate_hash();
+                let next_idx = if let Some(idx) = fbs.index.get(&h_node) {
+                    // alredy present
+                    *idx
+                } else {
+                    // not present
+                    let idx = fbs.graph.add_node(node);
+                    fbs.index.insert(h_node, idx);
+                    idx
+                };
+                // connect in graph if edge not present
+                if !fbs.graph.neighbors_directed(returntrace[returntrace.len()-1],Direction::Outgoing).any(|x| x == next_idx) {
+                    fbs.graph.add_edge(returntrace[returntrace.len()-1], next_idx, ());
+                    new_stg_edge = true;
+                }
+                returntrace.push(next_idx);
+                /*
+                Ideas:
+                Mark edges triggered by interrupts
+                Specify path with edges instead of nodes?
+                Form a coverage map over edges?
+                Sum up execution time per ABB
+                */
             }
         }
-
+        if !fbs.graph.neighbors_directed(returntrace[returntrace.len()-1],Direction::Outgoing).any(|x| x == fbs.exit) {
-        graph
+            fbs.graph.add_edge(returntrace[returntrace.len()-1], fbs.exit, ());
+            new_stg_edge = true;
+        }
+        returntrace.push(fbs.exit);
+        (returntrace, new_stg_edge)
     }
 }
 
@@ -127,14 +204,25 @@ where
     {
         let observer = observers.match_name::<QemuSystemStateObserver>("systemstate")
             .expect("QemuSystemStateObserver not found");
-        let abbs = trace_to_state_abb(&observer.last_run);
+        let feedbackstate = match state
-        println!("{:?}",abbs);
+            .named_metadata_map_mut()
-        let res =StgFeedback::trace_to_stg(&observer.last_run, abbs);
+            .get_mut::<STGFeedbackState>("stgfeedbackstate") {
+                Some(s) => s,
+                None => {
+                    let n=STGFeedbackState::default();
+                    state.named_metadata_map_mut().insert(n, "stgfeedbackstate");
+                    state.named_metadata_map_mut().get_mut::<STGFeedbackState>("stgfeedbackstate").unwrap()
+                }
+            };
 
-        let out = res.map(|i,x| x.pretty_print(), |_,_| "");
+        let abbs = trace_to_state_abb(&observer.last_run);
-        let outs = Dot::with_config(&out, &[Config::EdgeNoLabel]).to_string();
+        // println!("{:?}",abbs);
-        let outs = outs.replace(';',"\\n");
+        let (trace, new_edge) = StgFeedback::update_stg(&observer.last_run, abbs, feedbackstate);
-        fs::write("./mystg.dot",outs).expect("Failed to write graph");
+
+        // let out = feedbackstate.graph.map(|i,x| x.pretty_print(), |_,_| "");
+        // let outs = Dot::with_config(&out, &[Config::EdgeNoLabel]).to_string();
+        // let outs = outs.replace(';',"\\n");
+        // fs::write("./mystg.dot",outs).expect("Failed to write graph");
 
         Ok(false)
     }