add stg edge feedback

fuzzers/FRET/src/fuzzer.rs

@@ -2,42 +2,19 @@
 //!
 use core::time::Duration;
 use std::{env, path::PathBuf, process::{self, abort}, io::{Read, Write}, fs::{self, OpenOptions}, cmp::{min, max}, mem::transmute_copy, collections::btree_map::Range, ptr::addr_of_mut, ffi::OsStr};
-
 use hashbrown::HashMap;
 use libafl_bolts::{
-    core_affinity::Cores,
+    core_affinity::Cores, current_nanos, rands::StdRand, shmem::{ShMemProvider, StdShMemProvider}, tuples::tuple_list, AsMutSlice, AsSlice
-    current_nanos,
-    rands::StdRand,
-    shmem::{ShMemProvider, StdShMemProvider},
-    tuples::tuple_list,
-    AsSlice,
-    AsMutSlice
 };
 use libafl::{
-    corpus::{Corpus, InMemoryCorpus, OnDiskCorpus},
+    corpus::{Corpus, InMemoryCorpus, OnDiskCorpus}, events::{launcher::Launcher, EventConfig}, executors::{ExitKind, TimeoutExecutor}, feedback_or, feedback_or_fast, feedbacks::{CrashFeedback, MaxMapFeedback, TimeoutFeedback}, fuzzer::{Fuzzer, StdFuzzer}, inputs::{BytesInput, HasTargetBytes}, monitors::MultiMonitor, observers::VariableMapObserver, prelude::{havoc_mutations, minimizer::TopRatedsMetadata, CorpusId, Generator, HasBytesVec, HitcountsMapObserver, RandBytesGenerator, SimpleEventManager, SimpleMonitor, SimpleRestartingEventManager, StdScheduledMutator}, schedulers::{IndexesLenTimeMinimizerScheduler, QueueScheduler}, stages::StdMutationalStage, state::{HasCorpus, HasMetadata, HasNamedMetadata, StdState}, Error, Evaluator
-    events::EventConfig,
-    events::launcher::Launcher,
-    executors::{ExitKind, TimeoutExecutor},
-    feedback_or,
-    feedback_or_fast,
-    feedbacks::{CrashFeedback, MaxMapFeedback, TimeoutFeedback},
-    fuzzer::{Fuzzer, StdFuzzer},
-    inputs::{BytesInput, HasTargetBytes},
-    monitors::MultiMonitor,
-    observers::{VariableMapObserver},
-    schedulers::{IndexesLenTimeMinimizerScheduler, QueueScheduler},
-    state::{HasCorpus, StdState, HasMetadata, HasNamedMetadata},
-    Error,
-    prelude::{SimpleMonitor, SimpleEventManager, RandBytesGenerator, Generator, SimpleRestartingEventManager, HasBytesVec, minimizer::TopRatedsMetadata, havoc_mutations, StdScheduledMutator, HitcountsMapObserver, CorpusId}, Evaluator, stages::StdMutationalStage,
 };
 use libafl_qemu::{
-    edges::{self, edges_map_mut_slice, MAX_EDGES_NUM}, edges::QemuEdgeCoverageHelper, elf::EasyElf, emu::Emulator, GuestPhysAddr, QemuExecutor,
+    edges::{self, edges_map_mut_slice, QemuEdgeCoverageHelper, MAX_EDGES_NUM}, elf::EasyElf, emu::{libafl_qemu_remove_native_breakpoint, libafl_qemu_set_native_breakpoint, Emulator}, GuestAddr, GuestPhysAddr, QemuExecutor, QemuHooks, QemuInstrumentationFilter, Regs
-    QemuHooks, Regs, QemuInstrumentationFilter, GuestAddr,
-    emu::libafl_qemu_set_native_breakpoint, emu::libafl_qemu_remove_native_breakpoint,
 };
 use rand::{SeedableRng, StdRng, Rng};
 use crate::{
-    clock::{ClockTimeFeedback, IcHist, QemuClockIncreaseFeedback, QemuClockObserver, FUZZ_START_TIMESTAMP}, mutational::{InterruptShiftStage, MINIMUM_INTER_ARRIVAL_TIME}, qemustate::QemuStateRestoreHelper, systemstate::{self, feedbacks::{DumpSystraceFeedback, NovelSystemStateFeedback}, graph::{GraphMaximizerCorpusScheduler, SysGraphFeedbackState, SysMapFeedback}, helpers::QemuSystemStateHelper, observers::QemuSystemStateObserver, schedulers::{GenerationScheduler, LongestTraceScheduler}, stg::StgFeedback}, worst::{AlwaysTrueFeedback, ExecTimeIncFeedback, TimeMaximizerCorpusScheduler, TimeStateMaximizerCorpusScheduler} 
+    clock::{ClockTimeFeedback, IcHist, QemuClockIncreaseFeedback, QemuClockObserver, FUZZ_START_TIMESTAMP}, mutational::{InterruptShiftStage, MINIMUM_INTER_ARRIVAL_TIME}, qemustate::QemuStateRestoreHelper, systemstate::{self, feedbacks::{DumpSystraceFeedback, NovelSystemStateFeedback}, graph::{GraphMaximizerCorpusScheduler, SysGraphFeedbackState, SysMapFeedback}, helpers::QemuSystemStateHelper, observers::QemuSystemStateObserver, schedulers::{GenerationScheduler, LongestTraceScheduler}, stg::{stg_map_mut_slice, StgFeedback, MAX_STG_NUM}}, worst::{AlwaysTrueFeedback, ExecTimeIncFeedback, TimeMaximizerCorpusScheduler, TimeStateMaximizerCorpusScheduler} 
 };
 use std::time::{SystemTime, UNIX_EPOCH};
 use clap::{Parser, Subcommand};
@@ -544,6 +521,13 @@ pub fn fuzz() {
         #[cfg(feature = "observer_hitcounts")]
         let edges_observer = HitcountsMapObserver::new(edges_observer);
 
+        // #[cfg(feature = "feed_stg")]
+        let stg_observer = unsafe { VariableMapObserver::from_mut_slice(
+            "stg",
+            stg_map_mut_slice(),
+            addr_of_mut!(MAX_STG_NUM)
+        )};
+
         // Create an observation channel to keep track of the execution time
         let clock_time_observer = QemuClockObserver::new("clocktime");
 
@@ -573,7 +557,7 @@ pub fn fuzz() {
             // Feedback to reward any input which increses the execution time
             ExecTimeIncFeedback::new()
         );
-        #[cfg(all(feature = "systemstate",not(any(feature = "feed_systemgraph",feature = "feed_systemtrace"))))]
+        #[cfg(all(feature = "systemstate"))]
         let mut feedback = feedback_or!(
             feedback,
             DumpSystraceFeedback::with_dump(if cli.dump_traces {cli.dump_name.clone().map(|x| x.with_extension("trace.ron"))} else {None})
@@ -583,6 +567,11 @@ pub fn fuzz() {
             feedback,
             StgFeedback::default()
         );
+        #[cfg(feature = "feed_stg")]
+        let mut feedback = feedback_or!(
+            feedback,
+            MaxMapFeedback::tracking(&stg_observer, true, true)
+        );
         #[cfg(feature = "feed_systemtrace")]
         let mut feedback = feedback_or!(
             feedback,
@@ -647,7 +636,7 @@ pub fn fuzz() {
         #[cfg(not(feature = "systemstate"))]
         let observer_list = tuple_list!(edges_observer, clock_time_observer);
         #[cfg(feature = "systemstate")]
-        let observer_list = tuple_list!(edges_observer, clock_time_observer, systemstate_observer);
+        let observer_list = tuple_list!(edges_observer, clock_time_observer, systemstate_observer, stg_observer);
 
         // Create a QEMU in-process executor
         let executor = QemuExecutor::new(

fuzzers/FRET/src/systemstate/feedbacks.rs

@@ -251,10 +251,10 @@ where
                     std::fs::write(s,ron::to_string(&observer.last_run).expect("Error serializing hashmap")).expect("Can not dump to file");
                     self.dumpfile = None
                 },
-            None => if !self.dump_metadata {println!("{:?}\n{:?}",observer.last_run,names);}
+            None => if self.dump_metadata {println!("{:?}\n{:?}",observer.last_run,names);}
         };
         if self.dump_metadata {self.last_trace=Some(observer.last_run.clone());}
-        Ok(!self.dump_metadata)
+        Ok(false)
     }
     /// Append to the testcase the generated metadata in case of a new corpus item
     #[inline]

fuzzers/FRET/src/systemstate/mod.rs

@@ -148,6 +148,7 @@ impl Hash for RefinedFreeRTOSSystemState {
     fn hash<H: Hasher>(&self, state: &mut H) {
         self.current_task.hash(state);
         self.ready_list_after.hash(state);
+        self.delay_list_after.hash(state);
         // self.last_pc.hash(state);
     }
 }
@@ -155,6 +156,18 @@ impl RefinedFreeRTOSSystemState {
     fn get_time(&self) -> u64 {
         self.end_tick-self.start_tick
     }
+
+    pub fn print_lists(&self) -> String {
+        let mut ret = String::from("+");
+        for j in self.ready_list_after.iter() {
+            ret.push_str(format!(" {}#{}", j.0.task_name, j.1).as_str());
+        }
+        ret.push_str("\n-");
+        for j in self.delay_list_after.iter() {
+            ret.push_str(format!(" {}#{}", j.0.task_name, j.1).as_str());
+        }
+        ret
+    }
 }
 
 // Wrapper around Vec<RefinedFreeRTOSSystemState> to attach as Metadata
@@ -206,6 +219,17 @@ pub struct AtomicBasicBlock {
     start: GuestAddr,
     ends: HashSet<GuestAddr>,
 }
+
+impl Hash for AtomicBasicBlock {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        // Use a combination of the start address and the set of ending addresses to compute the hash value
+        self.start.hash(state);
+        let mut keys : Vec<_> = self.ends.iter().collect();
+        keys.sort();
+        keys.hash(state);
+    }
+}
+
 impl fmt::Display for AtomicBasicBlock {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let mut ends_str = String::new();

fuzzers/FRET/src/systemstate/stg.rs

@@ -2,6 +2,8 @@
 use libafl::SerdeAny;
 /// Feedbacks organizing SystemStates as a graph
 use libafl::inputs::HasBytesVec;
+use libafl_bolts::ownedref::OwnedMutSlice;
+use petgraph::graph::EdgeIndex;
 use std::fs;
 use libafl_bolts::rands::RandomSeed;
 use libafl_bolts::rands::StdRand;
@@ -37,6 +39,7 @@ use libafl::state::MaybeHasClientPerfMonitor;
 use libafl::feedbacks::Feedback;
 use libafl_bolts::Named;
 use libafl::Error;
+use libafl_qemu::edges::EDGES_MAP_SIZE;
 use hashbrown::HashMap;
 use libafl::{executors::ExitKind, inputs::Input, observers::ObserversTuple, state::HasMetadata};
 use serde::{Deserialize, Serialize};
@@ -64,7 +67,7 @@ pub struct STGNode
 }
 impl STGNode {
     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))
+        format!("{}\n{:x}-{:x}\n{}", self.base.current_task.0.task_name, self.abb.start, self.abb.ends.iter().next().unwrap_or_else(||&0), self.base.print_lists())
     }
     fn calculate_hash(&self) -> u64 {
         let mut s = DefaultHasher::new();
@@ -124,6 +127,12 @@ impl Named for STGFeedbackState
 
 //============================= Graph Feedback
 
+pub static mut STG_MAP: [u8; EDGES_MAP_SIZE] = [0; EDGES_MAP_SIZE];
+pub static mut MAX_STG_NUM: usize = 0;
+pub unsafe fn stg_map_mut_slice<'a>() -> OwnedMutSlice<'a, u8> {
+    OwnedMutSlice::from_raw_parts_mut(STG_MAP.as_mut_ptr(), STG_MAP.len())
+}
+
 /// A Feedback reporting novel System-State Transitions. Depends on [`QemuSystemStateObserver`]
 #[derive(Serialize, Deserialize, Clone, Debug, Default)]
 pub struct StgFeedback
@@ -131,6 +140,21 @@ pub struct StgFeedback
     name: String,
     last_trace: Option<Vec<NodeIndex>>,
 }
+const INTEREST_EDGE : bool = true;
+const INTEREST_NODE : bool = true;
+fn set_observer_map(trace : &Vec<EdgeIndex>) {
+    unsafe {
+        for i in 0..MAX_STG_NUM {
+            STG_MAP[i] = 0;
+        }
+        for i in trace {
+            if MAX_STG_NUM < i.index() {
+                MAX_STG_NUM = i.index();
+            }
+            STG_MAP[i.index()]+=1;
+        }
+    }
+}
 impl StgFeedback {
     pub fn new() -> Self {
         Self {name: String::from("SysMapFeedback"), last_trace: None }
@@ -143,9 +167,10 @@ impl StgFeedback {
     /// newly discovered node?
     /// side effect:
     /// the graph gets new nodes
-    fn update_stg(trace: &Vec<RefinedFreeRTOSSystemState>, map: Vec<(String, Rc<AtomicBasicBlock>, usize, u64)>, fbs: &mut STGFeedbackState) -> (Vec<NodeIndex>, bool) {
+    fn update_stg(trace: &Vec<RefinedFreeRTOSSystemState>, map: Vec<(String, Rc<AtomicBasicBlock>, usize, u64)>, fbs: &mut STGFeedbackState) -> (Vec<NodeIndex>, Vec<EdgeIndex>, bool) {
-        let mut returntrace = vec![fbs.entrypoint];
+        let mut return_node_trace = vec![fbs.entrypoint];
-        let mut new_stg_edge = false;
+        let mut return_edge_trace = vec![];
+        let mut interesting = false;
         // add all missing state+abb combinations to the graph
         for (i,state) in trace.iter().enumerate() { // Iterate states first, keep the trace order intact
             if let Some(marker) = map.iter().filter(|(_,_,index,_)| index==&i).next() { // 
@@ -158,14 +183,19 @@ impl StgFeedback {
                     // not present
                     let idx = fbs.graph.add_node(node);
                     fbs.index.insert(h_node, idx);
+                    interesting |= INTEREST_NODE;
                     idx
                 };
                 // connect in graph if edge not present
-                if !fbs.graph.neighbors_directed(returntrace[returntrace.len()-1],Direction::Outgoing).any(|x| x == next_idx) {
+                let e = fbs.graph.edges_directed(return_node_trace[return_node_trace.len()-1],Direction::Outgoing).find(|x| petgraph::visit::EdgeRef::target(x) == next_idx);
-                    fbs.graph.add_edge(returntrace[returntrace.len()-1], next_idx, ());
+                if let Some(e_) = e {
-                    new_stg_edge = true;
+                    return_edge_trace.push(petgraph::visit::EdgeRef::id(&e_));
+                } else {
+                    let e_ = fbs.graph.add_edge(return_node_trace[return_node_trace.len()-1], next_idx, ());
+                    return_edge_trace.push(e_);
+                    interesting |= INTEREST_EDGE;
                 }
-                returntrace.push(next_idx);
+                return_node_trace.push(next_idx);
                 /*
                 Ideas:
                 Mark edges triggered by interrupts
@@ -175,12 +205,16 @@ impl StgFeedback {
                 */
             }
         }
-        if !fbs.graph.neighbors_directed(returntrace[returntrace.len()-1],Direction::Outgoing).any(|x| x == fbs.exit) {
+        // every path terminates at the end
-            fbs.graph.add_edge(returntrace[returntrace.len()-1], fbs.exit, ());
+        if !fbs.graph.neighbors_directed(return_node_trace[return_node_trace.len()-1],Direction::Outgoing).any(|x| x == fbs.exit) {
-            new_stg_edge = true;
+            let e_ = fbs.graph.add_edge(return_node_trace[return_node_trace.len()-1], fbs.exit, ());
+            return_edge_trace.push(e_);
+            interesting |= INTEREST_EDGE;
         }
-        returntrace.push(fbs.exit);
+        return_node_trace.push(fbs.exit);
-        (returntrace, new_stg_edge)
+        #[cfg(feature = "feed_stg")]
+        set_observer_map(&return_edge_trace);
+        (return_node_trace, return_edge_trace, interesting)
     }
 }
 
@@ -217,7 +251,7 @@ where
 
         let abbs = trace_to_state_abb(&observer.last_run);
         // println!("{:?}",abbs);
-        let (trace, new_edge) = StgFeedback::update_stg(&observer.last_run, abbs, feedbackstate);
+        let (trace, _, new_edge) = StgFeedback::update_stg(&observer.last_run, abbs, feedbackstate);
 
         // let out = feedbackstate.graph.map(|i,x| x.pretty_print(), |_,_| "");
         // let outs = Dot::with_config(&out, &[Config::EdgeNoLabel]).to_string();