Add IntelPT tracing module to libafl_qemu systemmode with KVM (#2774)

* intelpt module

.github/workflows/ubuntu-prepare/action.yml

@@ -5,7 +5,7 @@ runs:
   steps:
     - name: Install and cache deps
       shell: bash
-      run: sudo apt-get update && sudo apt-get install -y curl lsb-release wget software-properties-common gnupg ninja-build shellcheck pax-utils nasm libsqlite3-dev libc6-dev libgtk-3-dev gcc g++ gcc-arm-none-eabi gcc-arm-linux-gnueabi g++-arm-linux-gnueabi libslirp-dev libz3-dev build-essential cmake
+      run: sudo apt-get update && sudo apt-get install -y curl lsb-release wget software-properties-common gnupg ninja-build shellcheck pax-utils nasm libsqlite3-dev libc6-dev libgtk-3-dev gcc g++ gcc-arm-none-eabi gcc-arm-linux-gnueabi g++-arm-linux-gnueabi libslirp-dev libz3-dev build-essential
     - uses: dtolnay/rust-toolchain@stable
     - name: install just
       uses: extractions/setup-just@v2

libafl_qemu/Cargo.toml

@@ -88,6 +88,9 @@ slirp = [
   "libafl_qemu_sys/slirp",
 ] # build qemu with host libslirp (for user networking)
 
+intel_pt = ["systemmode", "x86_64", "dep:libafl_intelpt"]
+intel_pt_export_raw = ["intel_pt", "libafl_intelpt/export_raw"]
+
 # Requires the binary's build.rs to call `build_libafl_qemu`
 shared = ["libafl_qemu_sys/shared"]
 
@@ -101,6 +104,7 @@ libafl_bolts = { workspace = true, features = ["std", "derive"] }
 libafl_targets = { workspace = true, default-features = true }
 libafl_qemu_sys = { workspace = true }
 libafl_derive = { workspace = true, default-features = true }
+libafl_intelpt = { workspace = true, default-features = true, optional = true }
 
 serde = { workspace = true, default-features = false, features = [
   "alloc",

libafl_qemu/libafl_qemu_build/src/build.rs

@@ -1,8 +1,8 @@
-use core::str::FromStr;
 use std::{
     env, fs,
     path::{Path, PathBuf},
     process::Command,
+    str::FromStr,
 };
 
 use which::which;

libafl_qemu/src/modules/mod.rs

@@ -17,8 +17,7 @@ pub use usermode::*;
 
 #[cfg(feature = "systemmode")]
 pub mod systemmode;
-#[cfg(feature = "systemmode")]
+#[cfg(all(feature = "systemmode", feature = "intel_pt"))]
-#[expect(unused_imports)]
 pub use systemmode::*;
 
 pub mod edges;

libafl_qemu/src/modules/systemmode/intel_pt.rs

@@ -0,0 +1,155 @@
+use std::{
+    fmt::Debug,
+    ops::{Range, RangeInclusive},
+};
+
+use libafl::{HasMetadata, observers::ObserversTuple};
+pub use libafl_intelpt::SectionInfo;
+use libafl_intelpt::{Image, IntelPT, IntelPTBuilder};
+use libafl_qemu_sys::{CPUArchStatePtr, GuestAddr};
+use num_traits::SaturatingAdd;
+use typed_builder::TypedBuilder;
+
+use crate::{
+    EmulatorModules, NewThreadHook, Qemu, QemuParams,
+    modules::{AddressFilter, EmulatorModule, EmulatorModuleTuple, ExitKind},
+};
+
+#[derive(Debug, TypedBuilder)]
+pub struct IntelPTModule<T = u8> {
+    #[builder(setter(skip), default)]
+    pt: Option<IntelPT>,
+    #[builder(default = IntelPTModule::default_pt_builder())]
+    intel_pt_builder: IntelPTBuilder,
+    #[builder(setter(transform = |sections: &[SectionInfo]| {
+        let mut i = Image::new(None).unwrap();
+        i.add_files_cached(sections, None).unwrap();
+        i
+    }))]
+    image: Image,
+    map_ptr: *mut T,
+    map_len: usize,
+}
+
+impl IntelPTModule {
+    pub fn default_pt_builder() -> IntelPTBuilder {
+        IntelPT::builder().exclude_kernel(false)
+    }
+}
+
+impl<I, S, T> EmulatorModule<I, S> for IntelPTModule<T>
+where
+    I: Unpin,
+    S: Unpin + HasMetadata,
+    T: SaturatingAdd + From<u8> + Debug + 'static,
+{
+    fn pre_qemu_init<ET>(
+        &mut self,
+        emulator_modules: &mut EmulatorModules<ET, I, S>,
+        _qemu_params: &mut QemuParams,
+    ) where
+        ET: EmulatorModuleTuple<I, S>,
+    {
+        emulator_modules
+            .thread_creation(NewThreadHook::Function(intel_pt_new_thread::<ET, I, S, T>))
+            .unwrap();
+        // fixme: consider implementing a clean emulator_modules.thread_teradown
+    }
+
+    fn pre_exec<ET>(
+        &mut self,
+        _qemu: Qemu,
+        _emulator_modules: &mut EmulatorModules<ET, I, S>,
+        _state: &mut S,
+        _input: &I,
+    ) where
+        ET: EmulatorModuleTuple<I, S>,
+    {
+        let pt = self.pt.as_mut().expect("Intel PT module not initialized.");
+        pt.enable_tracing().unwrap();
+    }
+
+    fn post_exec<OT, ET>(
+        &mut self,
+        _qemu: Qemu,
+        _emulator_modules: &mut EmulatorModules<ET, I, S>,
+        _state: &mut S,
+        _input: &I,
+        _observers: &mut OT,
+        _exit_kind: &mut ExitKind,
+    ) where
+        OT: ObserversTuple<I, S>,
+        ET: EmulatorModuleTuple<I, S>,
+    {
+        let pt = self.pt.as_mut().expect("Intel PT module not initialized.");
+        pt.disable_tracing().unwrap();
+
+        let _ = pt
+            .decode_traces_into_map(&mut self.image, self.map_ptr, self.map_len)
+            .inspect_err(|e| log::warn!("Intel PT trace decode failed: {e}"));
+
+        #[cfg(feature = "intel_pt_export_raw")]
+        {
+            let _ = pt
+                .dump_last_trace_to_file()
+                .inspect_err(|e| log::warn!("Intel PT trace save to file failed: {e}"));
+        }
+    }
+}
+
+impl<T> AddressFilter for IntelPTModule<T>
+where
+    T: Debug + 'static,
+{
+    fn register(&mut self, address_range: &Range<GuestAddr>) {
+        let pt = self.pt.as_mut().unwrap();
+        let mut filters = pt.ip_filters();
+        let range_inclusive =
+            RangeInclusive::new(address_range.start as usize, address_range.end as usize - 1);
+        filters.push(range_inclusive);
+        pt.set_ip_filters(&filters).unwrap()
+    }
+
+    fn allowed(&self, address: &GuestAddr) -> bool {
+        let pt = self.pt.as_ref().unwrap();
+        for f in pt.ip_filters() {
+            if f.contains(&(*address as usize)) {
+                return true;
+            }
+        }
+        false
+    }
+}
+
+pub fn intel_pt_new_thread<ET, I, S, T>(
+    emulator_modules: &mut EmulatorModules<ET, I, S>,
+    _state: Option<&mut S>,
+    _env: CPUArchStatePtr,
+    tid: u32,
+) -> bool
+where
+    I: Unpin,
+    S: HasMetadata + Unpin,
+    ET: EmulatorModuleTuple<I, S>,
+    T: Debug + 'static,
+{
+    let intel_pt_module = emulator_modules
+        .modules_mut()
+        .match_first_type_mut::<IntelPTModule<T>>()
+        .unwrap();
+
+    if intel_pt_module.pt.is_some() {
+        panic!("Intel PT module already initialized, only single core VMs are supported ATM.");
+    }
+
+    let pt = intel_pt_module
+        .intel_pt_builder
+        .clone()
+        .pid(Some(tid as i32))
+        .build()
+        .unwrap();
+
+    intel_pt_module.pt = Some(pt);
+
+    true
+}

libafl_qemu/src/modules/systemmode/mod.rs

@@ -1 +1,2 @@
-
+#[cfg(feature = "intel_pt")]
+pub mod intel_pt;

scripts/parallellize_cargo_check.py

@@ -23,7 +23,7 @@ command = (
     "--exclude-features=prelude,python,sancov_pcguard_edges,arm,aarch64,i386,be,systemmode,whole_archive "
     "--no-dev-deps --exclude libafl_libfuzzer --exclude libafl_qemu --exclude libafl_qemu_sys --print-command-list;"
     "DOCS_RS=1 cargo hack check -p libafl_qemu -p libafl_qemu_sys --each-feature --clean-per-run "
-    "--exclude-features=prelude,python,sancov_pcguard_edges,arm,aarch64,i386,be,systemmode,whole_archive,slirp "
+    "--exclude-features=prelude,python,sancov_pcguard_edges,arm,aarch64,i386,be,systemmode,whole_archive,slirp,intel_pt,intel_pt_export_raw "
     "--no-dev-deps --features usermode --print-command-list"
 )