diff --git a/src/hyperlight_host/src/hypervisor/hyperlight_vm.rs b/src/hyperlight_host/src/hypervisor/hyperlight_vm.rs index c836d3d80..e89a40def 100644 --- a/src/hyperlight_host/src/hypervisor/hyperlight_vm.rs +++ b/src/hyperlight_host/src/hypervisor/hyperlight_vm.rs @@ -381,13 +381,10 @@ impl HyperlightVm { { Ok(VmExit::Cancelled()) } else { - #[cfg(feature = "trace_guest")] - tc.setup_guest_trace(Span::current().context()); - // ==== KILL() TIMING POINT 3: Before calling run() ==== // If kill() is called and ran to completion BEFORE this line executes: // - Will still do a VM entry, but signals will be sent until VM exits - let result = self.vm.run_vcpu(); + let result = self.vm.run_vcpu(&tc); // End current host trace by closing the current span that captures traces // happening when a guest exits and re-enters. diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/kvm.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/kvm.rs index b5ca402fe..c270b0015 100644 --- a/src/hyperlight_host/src/hypervisor/virtual_machine/kvm.rs +++ b/src/hyperlight_host/src/hypervisor/virtual_machine/kvm.rs @@ -28,6 +28,7 @@ use crate::hypervisor::gdb::DebuggableVm; use crate::hypervisor::regs::{CommonFpu, CommonRegisters, CommonSpecialRegisters}; use crate::hypervisor::virtual_machine::{VirtualMachine, VmExit}; use crate::mem::memory_region::MemoryRegion; +use crate::sandbox::trace::context::TraceContext as SandboxTraceContext; use crate::{Result, new_error}; /// Return `true` if the KVM API is available, version 12, and has UserMemory capability, or `false` otherwise @@ -104,7 +105,10 @@ impl VirtualMachine for KvmVm { Ok(()) } - fn run_vcpu(&mut self) -> Result { + fn run_vcpu(&mut self, tc: &SandboxTraceContext) -> Result { + #[cfg(feature = "trace_guest")] + tc.setup_guest_trace(Span::current().context()); + match self.vcpu_fd.run() { Ok(VcpuExit::Hlt) => Ok(VmExit::Halt()), Ok(VcpuExit::IoOut(port, data)) => Ok(VmExit::IoOut(port, data.to_vec())), diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs index 55af78ebb..985fa09ed 100644 --- a/src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs +++ b/src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs @@ -22,6 +22,7 @@ use tracing::{Span, instrument}; use crate::Result; use crate::hypervisor::regs::{CommonFpu, CommonRegisters, CommonSpecialRegisters}; use crate::mem::memory_region::MemoryRegion; +use crate::sandbox::trace::context::TraceContext as SandboxTraceContext; /// KVM (Kernel-based Virtual Machine) functionality (linux) #[cfg(kvm)] @@ -147,7 +148,7 @@ pub(crate) trait VirtualMachine: Debug + Send { /// Runs the vCPU until it exits. /// Note: this function should not emit any traces or spans as it is called after guest span is setup - fn run_vcpu(&mut self) -> Result; + fn run_vcpu(&mut self, tc: &SandboxTraceContext) -> Result; /// Get regs #[allow(dead_code)] diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/mshv.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/mshv.rs index c0e352604..1cad380fa 100644 --- a/src/hyperlight_host/src/hypervisor/virtual_machine/mshv.rs +++ b/src/hyperlight_host/src/hypervisor/virtual_machine/mshv.rs @@ -36,6 +36,7 @@ use crate::hypervisor::gdb::DebuggableVm; use crate::hypervisor::regs::{CommonFpu, CommonRegisters, CommonSpecialRegisters}; use crate::hypervisor::virtual_machine::{VirtualMachine, VmExit}; use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags}; +use crate::sandbox::trace::context::TraceContext as SandboxTraceContext; use crate::{Result, new_error}; /// Determine whether the HyperV for Linux hypervisor API is present @@ -103,7 +104,9 @@ impl VirtualMachine for MshvVm { Ok(()) } - fn run_vcpu(&mut self) -> Result { + fn run_vcpu(&mut self, tc: &SandboxTraceContext) -> Result { + #[cfg(feature = "trace_guest")] + tc.setup_guest_trace(Span::current().context()); const HALT_MESSAGE: hv_message_type = hv_message_type_HVMSG_X64_HALT; const IO_PORT_INTERCEPT_MESSAGE: hv_message_type = hv_message_type_HVMSG_X64_IO_PORT_INTERCEPT; diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/whp.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/whp.rs index 7effa3e9f..728983f48 100644 --- a/src/hyperlight_host/src/hypervisor/virtual_machine/whp.rs +++ b/src/hyperlight_host/src/hypervisor/virtual_machine/whp.rs @@ -34,6 +34,7 @@ use crate::hypervisor::surrogate_process_manager::get_surrogate_process_manager; use crate::hypervisor::virtual_machine::{VirtualMachine, VmExit}; use crate::hypervisor::wrappers::HandleWrapper; use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags}; +use crate::sandbox::trace::context::TraceContext as SandboxTraceContext; use crate::{Result, log_then_return, new_error}; #[allow(dead_code)] // Will be used for runtime hypervisor detection @@ -205,7 +206,9 @@ impl VirtualMachine for WhpVm { } #[expect(non_upper_case_globals, reason = "Windows API constant are lower case")] - fn run_vcpu(&mut self) -> Result { + fn run_vcpu(&mut self, tc: &SandboxTraceContext) -> Result { + #[cfg(feature = "trace_guest")] + tc.setup_guest_trace(Span::current().context()); let mut exit_context: WHV_RUN_VP_EXIT_CONTEXT = Default::default(); unsafe {