Add option to output execution traces during fuzzing (#81)

Author: novafacing

src/arch/aarch64.rs

@@ -1056,6 +1056,16 @@ impl TracerDisassembler for Disassembler {
         Ok(())
     }
 
+    fn disassemble_to_string(&mut self, bytes: &[u8]) -> Result<String> {
+        let mut r = U8Reader::new(bytes);
+
+        if let Ok(insn) = self.decoder.decode(&mut r) {
+            Ok(insn.to_string())
+        } else {
+            bail!("Could not disassemble {:?}", bytes);
+        }
+    }
+
     fn last_was_control_flow(&self) -> bool {
         if let Some(last) = self.last.as_ref() {
             // NOTE: This is imprecise on ARM because PC is not restricted

src/arch/arm.rs

@@ -1053,6 +1053,16 @@ impl TracerDisassembler for Disassembler {
         Ok(())
     }
 
+    fn disassemble_to_string(&mut self, bytes: &[u8]) -> Result<String> {
+        let mut r = U8Reader::new(bytes);
+
+        if let Ok(insn) = self.decoder.decode(&mut r) {
+            Ok(insn.to_string())
+        } else {
+            bail!("Could not disassemble {:?}", bytes);
+        }
+    }
+
     fn last_was_control_flow(&self) -> bool {
         if let Some(last) = self.last.as_ref() {
             // NOTE: This is imprecise on ARM because PC is not restricted

src/arch/risc_v.rs

@@ -459,6 +459,16 @@ impl TracerDisassembler for Disassembler {
         Ok(())
     }
 
+    fn disassemble_to_string(&mut self, bytes: &[u8]) -> Result<String> {
+        let mut r = U8Reader::new(bytes);
+
+        if let Ok(insn) = self.decoder.decode(&mut r) {
+            Ok(insn.to_string())
+        } else {
+            bail!("Could not disassemble {:?}", bytes);
+        }
+    }
+
     fn last_was_control_flow(&self) -> bool {
         if let Some(last) = self.last.as_ref() {
             if matches!(last.opcode(), |Opcode::BEQ| Opcode::BNE

src/arch/x86.rs

@@ -985,6 +985,14 @@ impl TracerDisassembler for Disassembler {
         Ok(())
     }
 
+    fn disassemble_to_string(&mut self, bytes: &[u8]) -> Result<String> {
+        if let Ok(insn) = self.decoder.decode_slice(bytes) {
+            Ok(insn.to_string())
+        } else {
+            bail!("Could not disassemble {:?}", bytes);
+        }
+    }
+
     fn cmp(&self) -> Vec<CmpExpr> {
         let mut cmp_exprs = Vec::new();
         if self.last_was_cmp() {

src/arch/x86_64.rs

@@ -964,6 +964,14 @@ impl TracerDisassembler for Disassembler {
         Ok(())
     }
 
+    fn disassemble_to_string(&mut self, bytes: &[u8]) -> Result<String> {
+        if let Ok(insn) = self.decoder.decode_slice(bytes) {
+            Ok(insn.to_string())
+        } else {
+            bail!("Could not disassemble {:?}", bytes);
+        }
+    }
+
     fn cmp(&self) -> Vec<CmpExpr> {
         let mut cmp_exprs = Vec::new();
         if self.last_was_cmp() {

src/haps/mod.rs

@@ -58,6 +58,7 @@ impl Tsffs {
             self.post_timeout_event()?;
         }
 
+        self.execution_trace.0.clear();
         self.save_repro_bookmark_if_needed()?;
 
         debug!(self.as_conf_object(), "Resuming simulation");
@@ -152,6 +153,10 @@ impl Tsffs {
             self.post_timeout_event()?;
         }
 
+        if self.save_all_execution_traces {
+            self.save_execution_trace()?;
+        }
+
         debug!(self.as_conf_object(), "Resuming simulation");
 
         run_alone(|| {
@@ -204,6 +209,7 @@ impl Tsffs {
             self.post_timeout_event()?;
         }
 
+        self.execution_trace.0.clear();
         self.save_repro_bookmark_if_needed()?;
 
         debug!(self.as_conf_object(), "Resuming simulation");
@@ -233,6 +239,7 @@ impl Tsffs {
             self.post_timeout_event()?;
         }
 
+        self.execution_trace.0.clear();
         self.save_repro_bookmark_if_needed()?;
 
         debug!(self.as_conf_object(), "Resuming simulation");
@@ -323,6 +330,10 @@ impl Tsffs {
             self.post_timeout_event()?;
         }
 
+        if self.save_all_execution_traces {
+            self.save_execution_trace()?;
+        }
+
         debug!(self.as_conf_object(), "Resuming simulation");
 
         run_alone(|| {
@@ -416,6 +427,10 @@ impl Tsffs {
             self.post_timeout_event()?;
         }
 
+        if self.save_all_execution_traces || self.save_solution_execution_traces {
+            self.save_execution_trace()?;
+        }
+
         debug!(self.as_conf_object(), "Resuming simulation");
 
         run_alone(|| {

src/lib.rs

@@ -42,6 +42,7 @@ use libafl_targets::AFLppCmpLogMap;
 use magic::MagicNumber;
 use num_traits::FromPrimitive as _;
 use serde::{Deserialize, Serialize};
+use serde_json::to_writer;
 use simics::{
     break_simulation, class, debug, error, free_attribute, get_class, get_interface,
     get_processor_number, info, lookup_file, object_clock, run_command, run_python, simics_init,
@@ -66,15 +67,19 @@ use std::{
     alloc::{alloc_zeroed, Layout},
     cell::OnceCell,
     collections::{hash_map::Entry, BTreeSet, HashMap, HashSet},
-    fs::File,
+    fs::{create_dir_all, File},
+    hash::{DefaultHasher, Hash, Hasher},
     path::PathBuf,
     ptr::null_mut,
     str::FromStr,
     sync::mpsc::{Receiver, Sender},
     thread::JoinHandle,
     time::SystemTime,
 };
-use tracer::tsffs::{on_instruction_after, on_instruction_before};
+use tracer::{
+    tsffs::{on_instruction_after, on_instruction_before},
+    ExecutionTrace,
+};
 use typed_builder::TypedBuilder;
 use versions::{Requirement, Versioning};
 
@@ -375,6 +380,25 @@ pub(crate) struct Tsffs {
     #[class(attribute(optional, default = true))]
     /// Whether to quit on iteration limit
     pub quit_on_iteration_limit: bool,
+    #[class(attribute(optional, default = false))]
+    /// Whether to save execution traces of test cases which result in a solution
+    pub save_solution_execution_traces: bool,
+    #[class(attribute(optional, default = false))]
+    /// Whether to save execution traces of test cases which result in an interesting input
+    pub save_interesting_execution_traces: bool,
+    #[class(attribute(optional, default = false))]
+    /// Whether to save all execution traces. This will consume a very large amount of resources
+    /// and should only be used for debugging and testing purposes.
+    pub save_all_execution_traces: bool,
+    #[class(attribute(optional, default = lookup_file("%simics%")?.join("execution-traces")))]
+    #[attr_value(fallible)]
+    /// The directory to save execution traces to, if any are set to be saved. This
+    /// directory may be a SIMICS relative path prefixed with "%simics%". If not
+    /// provided, "%simics%/execution-traces" will be used by default.
+    pub execution_trace_directory: PathBuf,
+    #[class(attribute(optional, default = false))]
+    /// Whether execution traces should include just PC (vs instruction text and bytes)
+    pub execution_trace_pc_only: bool,
 
     #[attr_value(skip)]
     /// Handle for the core simulation stopped hap
@@ -440,8 +464,11 @@ pub(crate) struct Tsffs {
     edges_seen: HashSet<u64>,
     #[attr_value(skip)]
     /// A map of the new edges to their AFL indices seen since the last time the fuzzer
-    /// provided an update
+    /// provided an update. This is not cleared every execution.
     edges_seen_since_last: HashMap<u64, u64>,
+    #[attr_value(skip)]
+    /// The set of PCs comprising the current execution trace. This is cleared every execution.
+    execution_trace: ExecutionTrace,
 
     #[attr_value(skip)]
     /// The name of the fuzz snapshot, if saved
@@ -865,6 +892,29 @@ impl Tsffs {
         }
         Ok(())
     }
+
+    /// Save the current execution trace to a file
+    pub fn save_execution_trace(&mut self) -> Result<()> {
+        let mut hasher = DefaultHasher::new();
+        self.execution_trace.hash(&mut hasher);
+        let hash = hasher.finish();
+
+        if !self.execution_trace_directory.is_dir() {
+            create_dir_all(&self.execution_trace_directory)?;
+        }
+
+        let trace_path = self
+            .execution_trace_directory
+            .join(format!("{:x}.json", hash));
+
+        if !trace_path.exists() {
+            let trace_file = File::create(trace_path)?;
+
+            to_writer(trace_file, &self.execution_trace)?;
+        }
+
+        Ok(())
+    }
 }
 
 #[simics_init(name = "tsffs", class = "tsffs")]

src/log/mod.rs

@@ -74,6 +74,10 @@ impl Tsffs {
 
                         self.edges_seen_since_last.clear();
                     }
+
+                    if self.save_interesting_execution_traces {
+                        self.save_execution_trace()?;
+                    }
                 }
             }
 

src/tracer/mod.rs

@@ -6,18 +6,45 @@ use ffi::ffi;
 use libafl::prelude::CmpValues;
 use libafl_bolts::{AsMutSlice, AsSlice};
 use libafl_targets::{AFLppCmpLogOperands, AFL_CMP_TYPE_INS, CMPLOG_MAP_H};
+use serde::{Deserialize, Serialize};
 use simics::{
     api::{
         get_processor_number, sys::instruction_handle_t, AsConfObject, AttrValue, AttrValueType,
         ConfObject,
     },
     trace,
 };
-use std::{collections::HashMap, ffi::c_void, fmt::Display, num::Wrapping, str::FromStr};
+use std::{
+    collections::HashMap, ffi::c_void, fmt::Display, hash::Hash, num::Wrapping,
+    slice::from_raw_parts, str::FromStr,
+};
 use typed_builder::TypedBuilder;
 
 use crate::{arch::ArchitectureOperations, Tsffs};
 
+#[derive(Deserialize, Serialize, Debug, Default)]
+pub(crate) struct ExecutionTrace(pub HashMap<i32, Vec<ExecutionTraceEntry>>);
+
+impl Hash for ExecutionTrace {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        for (k, v) in self.0.iter() {
+            for entry in v.iter() {
+                k.hash(state);
+                entry.hash(state);
+            }
+        }
+    }
+}
+
+#[derive(TypedBuilder, Deserialize, Serialize, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub(crate) struct ExecutionTraceEntry {
+    pc: u64,
+    #[builder(default, setter(into, strip_option))]
+    insn: Option<String>,
+    #[builder(default, setter(into, strip_option))]
+    insn_bytes: Option<Vec<u8>>,
+}
+
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub(crate) enum CmpExprShift {
     Lsl,
@@ -346,6 +373,46 @@ impl Tsffs {
             }
         }
 
+        if self.coverage_enabled
+            && (self.save_all_execution_traces
+                || self.save_interesting_execution_traces
+                || self.save_solution_execution_traces)
+        {
+            if let Some(arch) = self.processors.get_mut(&processor_number) {
+                self.execution_trace
+                    .0
+                    .entry(processor_number)
+                    .or_default()
+                    .push(if self.execution_trace_pc_only {
+                        ExecutionTraceEntry::builder()
+                            .pc(arch.processor_info_v2().get_program_counter()?)
+                            .build()
+                    } else {
+                        let instruction_bytes =
+                            arch.cpu_instruction_query().get_instruction_bytes(handle)?;
+                        let instruction_bytes = unsafe {
+                            from_raw_parts(instruction_bytes.data, instruction_bytes.size)
+                        };
+
+                        if let Ok(disassembly_string) =
+                            arch.disassembler().disassemble_to_string(instruction_bytes)
+                        {
+                            ExecutionTraceEntry::builder()
+                                .pc(arch.processor_info_v2().get_program_counter()?)
+                                .insn(disassembly_string)
+                                .insn_bytes(instruction_bytes.to_vec())
+                                .build()
+                        } else {
+                            ExecutionTraceEntry::builder()
+                                .pc(arch.processor_info_v2().get_program_counter()?)
+                                .insn("(unknown)".to_string())
+                                .insn_bytes(instruction_bytes.to_vec())
+                                .build()
+                        }
+                    });
+            }
+        }
+
         Ok(())
     }
 }

src/traits/mod.rs

@@ -8,6 +8,7 @@ use anyhow::Result;
 /// and compare tracing
 pub trait TracerDisassembler {
     fn disassemble(&mut self, bytes: &[u8]) -> Result<()>;
+    fn disassemble_to_string(&mut self, bytes: &[u8]) -> Result<String>;
     fn last_was_control_flow(&self) -> bool;
     fn last_was_call(&self) -> bool;
     fn last_was_ret(&self) -> bool;