Add fast trace collection & use it for instrumenting guest memory operations

Cargo.toml

@@ -12,6 +12,7 @@ members = [
     "src/hyperlight_guest_capi",
     "src/hyperlight_testing",
     "src/hyperlight_host/fuzz",
+    "src/trace_dump",
 ]
 # Because hyperlight-guest has custom linker flags,
 # we exclude it from the default-members list

src/hyperlight_guest/Cargo.toml

@@ -17,6 +17,7 @@ default = ["libc", "printf", "alloca"]
 libc = [] # compile musl libc
 printf = [] # compile printf
 alloca = [] # compile alloca wrapper
+mem_profile = []
 
 [dependencies]
 anyhow = { version = "1.0.45", default-features = false }

src/hyperlight_guest/src/entrypoint.rs

@@ -131,7 +131,11 @@ pub extern "win64" fn entrypoint(peb_address: u64, seed: u64, ops: u64, max_log_
 
             let heap_start = (*peb_ptr).guestheapData.guestHeapBuffer as usize;
             let heap_size = (*peb_ptr).guestheapData.guestHeapSize as usize;
-            HEAP_ALLOCATOR
+            #[cfg(not(feature = "mem_profile"))]
+            let heap_allocator = &HEAP_ALLOCATOR;
+            #[cfg(feature = "mem_profile")]
+            let heap_allocator = &HEAP_ALLOCATOR.0;
+            heap_allocator
                 .try_lock()
                 .expect("Failed to access HEAP_ALLOCATOR")
                 .init(heap_start, heap_size);

src/hyperlight_guest/src/host_function_call.rs

@@ -37,6 +37,13 @@ pub enum OutBAction {
     Log = 99,
     CallFunction = 101,
     Abort = 102,
+    #[cfg(feature = "mem_profile")]
+    #[allow(dead_code)]
+    TraceRecordStack = 103,
+    #[cfg(feature = "mem_profile")]
+    TraceMemoryAlloc = 104,
+    #[cfg(feature = "mem_profile")]
+    TraceMemoryFree = 105,
 }
 
 pub fn get_host_value_return_as_void() -> Result<()> {

src/hyperlight_guest/src/lib.rs

@@ -83,8 +83,66 @@ fn panic(info: &core::panic::PanicInfo) -> ! {
 }
 
 // Globals
+#[cfg(feature = "mem_profile")]
+struct ProfiledLockedHeap<const ORDER: usize>(LockedHeap<ORDER>);
+#[cfg(feature = "mem_profile")]
+unsafe impl<const ORDER: usize> alloc::alloc::GlobalAlloc for ProfiledLockedHeap<ORDER> {
+    unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 {
+        let addr = self.0.alloc(layout);
+        unsafe {
+            core::arch::asm!("out dx, al",
+                in("dx") OutBAction::TraceMemoryAlloc as u16,
+                in("rax") layout.size() as u64,
+                in("rcx") addr as u64);
+        }
+        addr
+    }
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: core::alloc::Layout) {
+        unsafe {
+            core::arch::asm!("out dx, al",
+                in("dx") OutBAction::TraceMemoryFree as u16,
+                in("rax") layout.size() as u64,
+                in("rcx") ptr as u64);
+        }
+        self.0.dealloc(ptr, layout)
+    }
+    unsafe fn alloc_zeroed(&self, layout: core::alloc::Layout) -> *mut u8 {
+        let addr = self.0.alloc_zeroed(layout);
+        unsafe {
+            core::arch::asm!("out dx, al",
+                in("dx") OutBAction::TraceMemoryAlloc as u16,
+                in("rax") layout.size() as u64,
+                in("rcx") addr as u64);
+        }
+        addr
+    }
+    unsafe fn realloc(
+        &self,
+        ptr: *mut u8,
+        layout: core::alloc::Layout,
+        new_size: usize,
+    ) -> *mut u8 {
+        let new_ptr = self.0.realloc(ptr, layout, new_size);
+        unsafe {
+            core::arch::asm!("out dx, al",
+                in("dx") OutBAction::TraceMemoryFree as u16,
+                in("rax") layout.size() as u64,
+                in("rcx") ptr);
+            core::arch::asm!("out dx, al",
+                in("dx") OutBAction::TraceMemoryAlloc as u16,
+                in("rax") new_size as u64,
+                in("rcx") new_ptr);
+        }
+        new_ptr
+    }
+}
+#[cfg(not(feature = "mem_profile"))]
 #[global_allocator]
 pub(crate) static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::empty();
+#[cfg(feature = "mem_profile")]
+#[global_allocator]
+pub(crate) static HEAP_ALLOCATOR: ProfiledLockedHeap<32> =
+    ProfiledLockedHeap(LockedHeap::<32>::empty());
 
 ///cbindgen:ignore
 #[no_mangle]

src/hyperlight_host/Cargo.toml

@@ -22,6 +22,9 @@ workspace = true
 
 [dependencies]
 goblin = { version = "0.9" }
+framehop = { version = "0.13.1", optional = true }
+fallible-iterator = { version = "0.3.0", optional = true }
+blake3 = { version = "1.5.5", optional = true }
 rand = { version = "0.8.5" }
 cfg-if = { version = "1.0.0" }
 libc = { version = "0.2.167" }
@@ -48,6 +51,8 @@ strum = { version = "0.26", features = ["derive"] }
 tempfile = { version = "3.10", optional = true }
 serde_yaml = "0.9"
 anyhow = "1.0"
+uuid = { version = "1.4.1", features = ["v4"] }
+
 
 [target.'cfg(windows)'.dependencies]
 windows = { version = "0.58", features = [
@@ -76,7 +81,6 @@ kvm-bindings = { version = "0.10.0", features = ["fam-wrappers"], optional = tru
 kvm-ioctls = { version = "0.19.0", optional = true }
 
 [dev-dependencies]
-uuid = { version = "1.4.1", features = ["v4"] }
 signal-hook-registry = "1.4.1"
 envy = { version = "0.4.2" }
 serde = "1.0"
@@ -120,6 +124,14 @@ executable_heap = []
 # This feature enables printing of debug information to stdout in debug builds
 print_debug = []
 crashdump = ["dep:tempfile"] # Dumps the VM state to a file on unexpected errors or crashes. The path of the file will be printed on stdout and logged. This feature can only be used in debug builds.
+# This feature enables the generation of trace files.  Traces do not
+# include any particularly expensive instrumentation unless other
+# features are enabled.
+trace_guest = []
+# This feature enables unwinding the guest stack from the host, in
+# order to produce stack traces for debugging or profiling.
+unwind_guest = [ "trace_guest", "dep:framehop", "dep:fallible-iterator" ]
+mem_profile = [ "unwind_guest", "dep:blake3" ]
 kvm = ["dep:kvm-bindings", "dep:kvm-ioctls"]
 mshv = ["dep:mshv-bindings", "dep:mshv-ioctls"]
 inprocess = []

src/hyperlight_host/src/hypervisor/handlers.rs

@@ -18,14 +18,24 @@ use std::sync::{Arc, Mutex};
 
 use tracing::{instrument, Span};
 
+#[cfg(feature = "trace_guest")]
+use super::Hypervisor;
+#[cfg(feature = "trace_guest")]
+use crate::sandbox::TraceInfo;
 use crate::{new_error, Result};
 
 /// The trait representing custom logic to handle the case when
 /// a Hypervisor's virtual CPU (vCPU) informs Hyperlight the guest
 /// has initiated an outb operation.
-pub trait OutBHandlerCaller: Sync + Send {
+pub(crate) trait OutBHandlerCaller: Sync + Send {
     /// Function that gets called when an outb operation has occurred.
-    fn call(&mut self, port: u16, payload: u64) -> Result<()>;
+    fn call(
+        &mut self,
+        #[cfg(feature = "trace_guest")] hv: &mut dyn Hypervisor,
+        #[cfg(feature = "trace_guest")] trace_info: TraceInfo,
+        port: u16,
+        payload: u64,
+    ) -> Result<()>;
 }
 
 /// A convenient type representing a common way `OutBHandler` implementations
@@ -34,8 +44,12 @@ pub trait OutBHandlerCaller: Sync + Send {
 /// Note: This needs to be wrapped in a Mutex to be able to grab a mutable
 /// reference to the underlying data (i.e., handle_outb in `Sandbox` takes
 /// a &mut self).
-pub type OutBHandlerWrapper = Arc<Mutex<dyn OutBHandlerCaller>>;
+pub(crate) type OutBHandlerWrapper = Arc<Mutex<dyn OutBHandlerCaller>>;
 
+#[cfg(feature = "trace_guest")]
+pub(crate) type OutBHandlerFunction =
+    Box<dyn FnMut(&mut dyn Hypervisor, TraceInfo, u16, u64) -> Result<()> + Send>;
+#[cfg(not(feature = "trace_guest"))]
 pub(crate) type OutBHandlerFunction = Box<dyn FnMut(u16, u64) -> Result<()> + Send>;
 
 /// A `OutBHandler` implementation using a `OutBHandlerFunction`
@@ -52,12 +66,25 @@ impl From<OutBHandlerFunction> for OutBHandler {
 
 impl OutBHandlerCaller for OutBHandler {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-    fn call(&mut self, port: u16, payload: u64) -> Result<()> {
+    fn call(
+        &mut self,
+        #[cfg(feature = "trace_guest")] hv: &mut dyn Hypervisor,
+        #[cfg(feature = "trace_guest")] trace_info: TraceInfo,
+        port: u16,
+        payload: u64,
+    ) -> Result<()> {
         let mut func = self
             .0
             .try_lock()
             .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
-        func(port, payload)
+        func(
+            #[cfg(feature = "trace_guest")]
+            hv,
+            #[cfg(feature = "trace_guest")]
+            trace_info,
+            port,
+            payload,
+        )
     }
 }
 

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -24,11 +24,18 @@ use mshv_bindings::{
     hv_register_name_HV_X64_REGISTER_RIP, hv_register_value, mshv_user_mem_region,
     FloatingPointUnit, SegmentRegister, SpecialRegisters, StandardRegisters,
 };
+#[cfg(feature = "unwind_guest")]
+use mshv_bindings::{
+    hv_register_name, hv_register_name_HV_X64_REGISTER_RAX, hv_register_name_HV_X64_REGISTER_RBP,
+    hv_register_name_HV_X64_REGISTER_RCX, hv_register_name_HV_X64_REGISTER_RSP,
+};
 use mshv_ioctls::{Mshv, VcpuFd, VmFd};
 use tracing::{instrument, Span};
 
 use super::fpu::{FP_CONTROL_WORD_DEFAULT, FP_TAG_WORD_DEFAULT, MXCSR_DEFAULT};
 use super::handlers::{MemAccessHandlerWrapper, OutBHandlerWrapper};
+#[cfg(feature = "unwind_guest")]
+use super::TraceRegister;
 use super::{
     Hypervisor, VirtualCPU, CR0_AM, CR0_ET, CR0_MP, CR0_NE, CR0_PE, CR0_PG, CR0_WP, CR4_OSFXSR,
     CR4_OSXMMEXCPT, CR4_PAE, EFER_LMA, EFER_LME, EFER_NX, EFER_SCE,
@@ -37,6 +44,8 @@ use crate::hypervisor::hypervisor_handler::HypervisorHandler;
 use crate::hypervisor::HyperlightExit;
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
 use crate::mem::ptr::{GuestPtr, RawPtr};
+#[cfg(feature = "trace_guest")]
+use crate::sandbox::TraceInfo;
 use crate::{log_then_return, new_error, Result};
 
 /// Determine whether the HyperV for Linux hypervisor API is present
@@ -163,6 +172,19 @@ impl Debug for HypervLinuxDriver {
     }
 }
 
+#[cfg(feature = "unwind_guest")]
+impl From<TraceRegister> for hv_register_name {
+    fn from(r: TraceRegister) -> Self {
+        match r {
+            TraceRegister::RAX => hv_register_name_HV_X64_REGISTER_RAX,
+            TraceRegister::RCX => hv_register_name_HV_X64_REGISTER_RCX,
+            TraceRegister::RIP => hv_register_name_HV_X64_REGISTER_RIP,
+            TraceRegister::RSP => hv_register_name_HV_X64_REGISTER_RSP,
+            TraceRegister::RBP => hv_register_name_HV_X64_REGISTER_RBP,
+        }
+    }
+}
+
 impl Hypervisor for HypervLinuxDriver {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     fn initialise(
@@ -173,6 +195,7 @@ impl Hypervisor for HypervLinuxDriver {
         outb_hdl: OutBHandlerWrapper,
         mem_access_hdl: MemAccessHandlerWrapper,
         hv_handler: Option<HypervisorHandler>,
+        #[cfg(feature = "trace_guest")] trace_info: TraceInfo,
     ) -> Result<()> {
         let regs = StandardRegisters {
             rip: self.entrypoint,
@@ -194,6 +217,8 @@ impl Hypervisor for HypervLinuxDriver {
             hv_handler,
             outb_hdl,
             mem_access_hdl,
+            #[cfg(feature = "trace_guest")]
+            trace_info,
         )?;
 
         // reset RSP to what it was before initialise
@@ -212,6 +237,7 @@ impl Hypervisor for HypervLinuxDriver {
         outb_handle_fn: OutBHandlerWrapper,
         mem_access_fn: MemAccessHandlerWrapper,
         hv_handler: Option<HypervisorHandler>,
+        #[cfg(feature = "trace_guest")] trace_info: TraceInfo,
     ) -> Result<()> {
         // Reset general purpose registers except RSP, then set RIP
         let rsp_before = self.vcpu_fd.get_regs()?.rsp;
@@ -238,6 +264,8 @@ impl Hypervisor for HypervLinuxDriver {
             hv_handler,
             outb_handle_fn,
             mem_access_fn,
+            #[cfg(feature = "trace_guest")]
+            trace_info,
         )?;
 
         // reset RSP to what it was before function call
@@ -257,12 +285,20 @@ impl Hypervisor for HypervLinuxDriver {
         rip: u64,
         instruction_length: u64,
         outb_handle_fn: OutBHandlerWrapper,
+        #[cfg(feature = "trace_guest")] trace_info: TraceInfo,
     ) -> Result<()> {
         let payload = data[..8].try_into()?;
         outb_handle_fn
             .try_lock()
             .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
-            .call(port, u64::from_le_bytes(payload))?;
+            .call(
+                #[cfg(feature = "trace_guest")]
+                self,
+                #[cfg(feature = "trace_guest")]
+                trace_info,
+                port,
+                u64::from_le_bytes(payload),
+            )?;
 
         // update rip
         self.vcpu_fd.set_reg(&[hv_register_assoc {
@@ -350,6 +386,17 @@ impl Hypervisor for HypervLinuxDriver {
         Ok(result)
     }
 
+    #[cfg(feature = "unwind_guest")]
+    fn read_trace_reg(&self, reg: TraceRegister) -> Result<u64> {
+        let mut assoc = [hv_register_assoc {
+            name: reg.into(),
+            ..Default::default()
+        }];
+        self.vcpu_fd.get_reg(&mut assoc)?;
+        // safety: all registers that we currently support are 64-bit
+        unsafe { Ok(assoc[0].value.reg64) }
+    }
+
     #[instrument(skip_all, parent = Span::current(), level = "Trace")]
     fn as_mut_hypervisor(&mut self) -> &mut dyn Hypervisor {
         self as &mut dyn Hypervisor