[WIP] Enable guest tracing

Cargo.toml

@@ -10,11 +10,14 @@ members = [
     "src/hyperlight_guest",
     "src/hyperlight_host",
     "src/hyperlight_guest_capi",
+    "src/hyperlight_guest_tracing",
+    "src/hyperlight_guest_tracing_macro",
     "src/hyperlight_testing",
     "fuzz",
     "src/hyperlight_guest_bin",
     "src/hyperlight_component_util",
     "src/hyperlight_component_macro",
+    "src/trace_dump",
 ]
 # Guests have custom linker flags, so we need to exclude them from the workspace
 exclude = [
@@ -39,6 +42,8 @@ hyperlight-host = { path = "src/hyperlight_host", version = "0.7.0", default-fea
 hyperlight-guest = { path = "src/hyperlight_guest", version = "0.7.0", default-features = false }
 hyperlight-guest-bin = { path = "src/hyperlight_guest_bin", version = "0.7.0", default-features = false }
 hyperlight-testing = { path = "src/hyperlight_testing", default-features = false }
+hyperlight-guest-tracing = { path = "src/hyperlight_guest_tracing", default-features = false }
+hyperlight-guest-tracing-macro = { path = "src/hyperlight_guest_tracing_macro", default-features = false }
 hyperlight-component-util = { path = "src/hyperlight_component_util", version = "0.7.0", default-features = false }
 hyperlight-component-macro = { path = "src/hyperlight_component_macro", version = "0.7.0", default-features = false }
 

Justfile

@@ -38,11 +38,11 @@ witguest-wit:
     cargo install --locked wasm-tools
     cd src/tests/rust_guests/witguest && wasm-tools component wit guest.wit -w -o interface.wasm
 
-build-rust-guests target=default-target: (witguest-wit)
-    cd src/tests/rust_guests/callbackguest && cargo build --profile={{ if target == "debug" { "dev" } else { target } }}
-    cd src/tests/rust_guests/simpleguest && cargo build --profile={{ if target == "debug" { "dev" } else { target } }} 
-    cd src/tests/rust_guests/dummyguest && cargo build --profile={{ if target == "debug" { "dev" } else { target } }} 
-    cd src/tests/rust_guests/witguest && cargo build --profile={{ if target == "debug" { "dev" } else { target } }}
+build-rust-guests target=default-target features="": (witguest-wit)
+    cd src/tests/rust_guests/callbackguest && cargo build {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" { "dev" } else { target } }}
+    cd src/tests/rust_guests/simpleguest && cargo build {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" { "dev" } else { target } }} 
+    cd src/tests/rust_guests/dummyguest && cargo build {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" { "dev" } else { target } }} 
+    cd src/tests/rust_guests/witguest && cargo build {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" { "dev" } else { target } }}
 
 @move-rust-guests target=default-target:
     cp {{ callbackguest_source }}/{{ target }}/callbackguest* {{ rust_guests_bin_dir }}/{{ target }}/

src/hyperlight_common/Cargo.toml

@@ -25,11 +25,14 @@ spin = "0.10.0"
 [features]
 default = ["tracing"]
 fuzzing = ["dep:arbitrary"]
+trace_guest = []
+unwind_guest = []
+mem_profile = []
 std = []
 
 [dev-dependencies]
 hyperlight-testing = { workspace = true }
 
 [lib]
 bench = false # see https://bheisler.github.io/criterion.rs/book/faq.html#cargo-bench-gives-unrecognized-option-errors-for-valid-command-line-options
-doctest = false # reduce noise in test output
+doctest = false # reduce noise in test output

src/hyperlight_common/src/outb.rs

@@ -90,11 +90,22 @@ impl TryFrom<u8> for Exception {
 /// - CallFunction: makes a call to a host function,
 /// - Abort: aborts the execution of the guest,
 /// - DebugPrint: prints a message to the host
+/// - TraceRecordStack: records the stack trace of the guest
+/// - TraceMemoryAlloc: records memory allocation events
+/// - TraceMemoryFree: records memory deallocation events
 pub enum OutBAction {
     Log = 99,
     CallFunction = 101,
     Abort = 102,
     DebugPrint = 103,
+    #[cfg(feature = "unwind_guest")]
+    TraceRecordStack = 104,
+    #[cfg(feature = "mem_profile")]
+    TraceMemoryAlloc = 105,
+    #[cfg(feature = "mem_profile")]
+    TraceMemoryFree = 106,
+    #[cfg(feature = "trace_guest")]
+    TraceRecord = 107,
 }
 
 impl TryFrom<u16> for OutBAction {
@@ -105,6 +116,14 @@ impl TryFrom<u16> for OutBAction {
             101 => Ok(OutBAction::CallFunction),
             102 => Ok(OutBAction::Abort),
             103 => Ok(OutBAction::DebugPrint),
+            #[cfg(feature = "unwind_guest")]
+            104 => Ok(OutBAction::TraceRecordStack),
+            #[cfg(feature = "mem_profile")]
+            105 => Ok(OutBAction::TraceMemoryAlloc),
+            #[cfg(feature = "mem_profile")]
+            106 => Ok(OutBAction::TraceMemoryFree),
+            #[cfg(feature = "trace_guest")]
+            107 => Ok(OutBAction::TraceRecord),
             _ => Err(anyhow::anyhow!("Invalid OutBAction value: {}", val)),
         }
     }

src/hyperlight_guest/Cargo.toml

@@ -15,3 +15,9 @@ Provides only the essential building blocks for interacting with the host enviro
 anyhow = { version = "1.0.98", default-features = false }
 serde_json = { version = "1.0", default-features = false, features = ["alloc"] }
 hyperlight-common = { workspace = true }
+hyperlight-guest-tracing = { workspace = true, optional = true }
+hyperlight-guest-tracing-macro = { workspace = true}
+
+[features]
+default = []
+trace_guest = ["dep:hyperlight-guest-tracing"]

src/hyperlight_guest/src/exit.rs

@@ -21,17 +21,20 @@ use hyperlight_common::outb::OutBAction;
 
 /// Halt the execution of the guest and returns control to the host.
 #[inline(never)]
+#[hyperlight_guest_tracing_macro::trace_function]
 pub fn halt() {
     unsafe { asm!("hlt", options(nostack)) }
 }
 
 /// Exits the VM with an Abort OUT action and code 0.
 #[unsafe(no_mangle)]
+#[hyperlight_guest_tracing_macro::trace_function]
 pub extern "C" fn abort() -> ! {
     abort_with_code(&[0, 0xFF])
 }
 
 /// Exits the VM with an Abort OUT action and a specific code.
+#[hyperlight_guest_tracing_macro::trace_function]
 pub fn abort_with_code(code: &[u8]) -> ! {
     outb(OutBAction::Abort as u16, code);
     outb(OutBAction::Abort as u16, &[0xFF]); // send abort terminator (if not included in code)
@@ -42,6 +45,7 @@ pub fn abort_with_code(code: &[u8]) -> ! {
 ///
 /// # Safety
 /// This function is unsafe because it dereferences a raw pointer.
+#[hyperlight_guest_tracing_macro::trace_function]
 pub unsafe fn abort_with_code_and_message(code: &[u8], message_ptr: *const c_char) -> ! {
     unsafe {
         // Step 1: Send abort code (typically 1 byte, but `code` allows flexibility)
@@ -63,6 +67,7 @@ pub unsafe fn abort_with_code_and_message(code: &[u8], message_ptr: *const c_cha
 
 /// OUT bytes to the host through multiple exits.
 pub(crate) fn outb(port: u16, data: &[u8]) {
+    hyperlight_guest_tracing_macro::flush!();
     unsafe {
         let mut i = 0;
         while i < data.len() {

src/hyperlight_guest/src/guest_handle/host_comm.rs

@@ -35,6 +35,7 @@ use crate::exit::out32;
 
 impl GuestHandle {
     /// Get user memory region as bytes.
+    #[hyperlight_guest_tracing_macro::trace_function]
     pub fn read_n_bytes_from_user_memory(&self, num: u64) -> Result<Vec<u8>> {
         let peb_ptr = self.peb().unwrap();
         let user_memory_region_ptr = unsafe { (*peb_ptr).init_data.ptr as *mut u8 };
@@ -83,6 +84,7 @@ impl GuestHandle {
     ///
     /// Note: The function return value must be obtained by calling
     /// `get_host_return_value`.
+    #[hyperlight_guest_tracing_macro::trace_function]
     pub fn call_host_function_without_returning_result(
         &self,
         function_name: &str,
@@ -114,6 +116,7 @@ impl GuestHandle {
     /// sends it to the host, and then retrieves the return value.
     ///
     /// The return value is deserialized into the specified type `T`.
+    #[hyperlight_guest_tracing_macro::trace_function]
     pub fn call_host_function<T: TryFrom<ReturnValue>>(
         &self,
         function_name: &str,
@@ -155,6 +158,7 @@ impl GuestHandle {
     }
 
     /// Log a message with the specified log level, source, caller, source file, and line number.
+    #[hyperlight_guest_tracing_macro::trace_function]
     pub fn log_message(
         &self,
         log_level: LogLevel,

src/hyperlight_guest/src/guest_handle/io.rs

@@ -27,6 +27,7 @@ use crate::error::{HyperlightGuestError, Result};
 
 impl GuestHandle {
     /// Pops the top element from the shared input data buffer and returns it as a T
+    #[hyperlight_guest_tracing_macro::trace_function]
     pub fn try_pop_shared_input_data_into<T>(&self) -> Result<T>
     where
         T: for<'a> TryFrom<&'a [u8]>,
@@ -88,6 +89,7 @@ impl GuestHandle {
     }
 
     /// Pushes the given data onto the shared output data buffer.
+    #[hyperlight_guest_tracing_macro::trace_function]
     pub fn push_shared_output_data(&self, data: Vec<u8>) -> Result<()> {
         let peb_ptr = self.peb().unwrap();
         let output_stack_size = unsafe { (*peb_ptr).output_stack.size as usize };

src/hyperlight_guest_bin/Cargo.toml

@@ -17,10 +17,14 @@ and third-party code used by our C-API needed to build a native hyperlight-guest
 default = ["libc", "printf"]
 libc = [] # compile musl libc
 printf = [] # compile printf
+trace_guest = ["hyperlight-common/trace_guest", "dep:hyperlight-guest-tracing", "hyperlight-guest/trace_guest"]
+mem_profile = ["hyperlight-common/unwind_guest","hyperlight-common/mem_profile"]
 
 [dependencies]
 hyperlight-guest = { workspace = true, default-features = false }
 hyperlight-common = { workspace = true, default-features = false }
+hyperlight-guest-tracing = { workspace = true, optional = true }
+hyperlight-guest-tracing-macro = { workspace = true }
 buddy_system_allocator = "0.11.0"
 log = { version = "0.4", default-features = false }
 spin = "0.10.0"

src/hyperlight_guest_bin/src/guest_function/call.rs

@@ -27,6 +27,7 @@ use crate::{GUEST_HANDLE, REGISTERED_GUEST_FUNCTIONS};
 
 type GuestFunc = fn(&FunctionCall) -> Result<Vec<u8>>;
 
+#[hyperlight_guest_tracing_macro::trace_function]
 pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result<Vec<u8>> {
     // Validate this is a Guest Function Call
     if function_call.function_call_type() != FunctionCallType::Guest {
@@ -99,6 +100,7 @@ fn internal_dispatch_function() -> Result<()> {
 // This is implemented as a separate function to make sure that epilogue in the internal_dispatch_function is called before the halt()
 // which if it were included in the internal_dispatch_function cause the epilogue to not be called because the halt() would not return
 // when running in the hypervisor.
+#[hyperlight_guest_tracing_macro::trace_function]
 pub(crate) extern "C" fn dispatch_function() {
     let _ = internal_dispatch_function();
     halt();

src/hyperlight_guest_bin/src/lib.rs

@@ -28,8 +28,11 @@ use guest_function::register::GuestFunctionRegister;
 use guest_logger::init_logger;
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::mem::HyperlightPEB;
+#[cfg(feature = "mem_profile")]
+use hyperlight_common::outb::OutBAction;
 use hyperlight_guest::exit::{abort_with_code_and_message, halt};
 use hyperlight_guest::guest_handle::handle::GuestHandle;
+use hyperlight_guest_tracing_macro::{trace, trace_function};
 use log::LevelFilter;
 use spin::Once;
 
@@ -53,9 +56,69 @@ pub mod guest_logger;
 pub mod host_comm;
 pub mod memory;
 
+// 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 = unsafe { 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 = unsafe { 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 = unsafe { 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
+    }
+}
+
 // === Globals ===
+#[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());
 
 pub(crate) static mut GUEST_HANDLE: GuestHandle = GuestHandle::new();
 pub(crate) static mut REGISTERED_GUEST_FUNCTIONS: GuestFunctionRegister =
@@ -92,6 +155,7 @@ unsafe extern "C" {
 static INIT: Once = Once::new();
 
 #[unsafe(no_mangle)]
+#[trace_function]
 pub extern "C" fn entrypoint(peb_address: u64, seed: u64, ops: u64, max_log_level: u64) {
     if peb_address == 0 {
         panic!("PEB address is null");
@@ -128,7 +192,11 @@ pub extern "C" fn entrypoint(peb_address: u64, seed: u64, ops: u64, max_log_leve
 
             let heap_start = (*peb_ptr).guest_heap.ptr as usize;
             let heap_size = (*peb_ptr).guest_heap.size 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);
@@ -137,6 +205,7 @@ pub extern "C" fn entrypoint(peb_address: u64, seed: u64, ops: u64, max_log_leve
 
             (*peb_ptr).guest_function_dispatch_ptr = dispatch_function as usize as u64;
 
+            trace!("hyperlight_main");
             hyperlight_main();
         }
     });

src/hyperlight_guest_tracing/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "hyperlight-guest-tracing"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+homepage.workspace = true
+repository.workspace = true
+readme.workspace = true
+description = """Provides the tracing functionality for the hyperlight guest."""
+
+[dependencies]
+hyperlight-common = { workspace = true, default-features = false, features = ["trace_guest"] }
+spin = "0.10.0"
+
+[lints]
+workspace = true