[hyperlight_host/trace] Handle the trace records sent by a guest

- When an `Out` instruction is intercepted, the Hypervisor checks for
  the frame Id, to verify what type of exit it is.
  Based on this, when a trace record type is received, we copy the array
  of trace records from the guest's memory, calculate the timestamp
  based on the cycles returned by the guest's RDTSC and write the record
  to the trace file.

Signed-off-by: Doru Blânzeanu <dblnz@pm.me>

Author: dblnz

src/hyperlight_host/Cargo.toml

@@ -39,6 +39,7 @@ tracing = { version = "0.1.41", features = ["log"] }
 tracing-log = "0.2.0"
 tracing-core = "0.1.34"
 hyperlight-common = { workspace = true, default-features = true, features = [ "std" ] }
+hyperlight-guest-tracing = { workspace = true, default-features = true, optional = true }
 vmm-sys-util = "0.14.0"
 crossbeam-channel = "0.5.15"
 thiserror = "2.0.12"
@@ -131,7 +132,7 @@ executable_heap = []
 print_debug = []
 # Dumps the VM state to a file on unexpected errors or crashes. The path of the file will be printed on stdout and logged.
 crashdump = ["dep:chrono"]
-trace_guest = []
+trace_guest = ["hyperlight-common/trace_guest", "dep:hyperlight-guest-tracing"]
 # 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", "hyperlight-common/unwind_guest" ]

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -48,7 +48,7 @@ use mshv_bindings::{
     hv_partition_property_code_HV_PARTITION_PROPERTY_SYNTHETIC_PROC_FEATURES,
     hv_partition_synthetic_processor_features,
 };
-#[cfg(feature = "unwind_guest")]
+#[cfg(feature = "trace_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,
@@ -58,7 +58,7 @@ use tracing::{Span, instrument};
 #[cfg(crashdump)]
 use {super::crashdump, std::path::Path};
 
-#[cfg(feature = "unwind_guest")]
+#[cfg(feature = "trace_guest")]
 use super::TraceRegister;
 use super::fpu::{FP_CONTROL_WORD_DEFAULT, FP_TAG_WORD_DEFAULT, MXCSR_DEFAULT};
 #[cfg(gdb)]
@@ -532,7 +532,7 @@ impl Debug for HypervLinuxDriver {
     }
 }
 
-#[cfg(feature = "unwind_guest")]
+#[cfg(feature = "trace_guest")]
 impl From<TraceRegister> for hv_register_name {
     fn from(r: TraceRegister) -> Self {
         match r {
@@ -739,6 +739,12 @@ impl Hypervisor for HypervLinuxDriver {
                 libc::EINTR,
             )))
         } else {
+            #[cfg(feature = "trace_guest")]
+            if self.trace_info.guest_start_epoch.is_none() {
+                // Set the guest start epoch to the current time, before running the vcpu
+                crate::debug!("HyperV - Guest Start Epoch set");
+                self.trace_info.guest_start_epoch = Some(std::time::Instant::now());
+            }
             // Note: if a `InterruptHandle::kill()` called while this thread is **here**
             // Then the vcpu will run, but we will keep sending signals to this thread
             // to interrupt it until `running` is set to false. The `vcpu_fd::run()` call will
@@ -1092,7 +1098,7 @@ impl Hypervisor for HypervLinuxDriver {
         Ok(())
     }
 
-    #[cfg(feature = "unwind_guest")]
+    #[cfg(feature = "trace_guest")]
     fn read_trace_reg(&self, reg: TraceRegister) -> Result<u64> {
         let mut assoc = [hv_register_assoc {
             name: reg.into(),
@@ -1107,6 +1113,10 @@ impl Hypervisor for HypervLinuxDriver {
     fn trace_info_as_ref(&self) -> &TraceInfo {
         &self.trace_info
     }
+    #[cfg(feature = "trace_guest")]
+    fn trace_info_as_mut(&mut self) -> &mut TraceInfo {
+        &mut self.trace_info
+    }
 }
 
 impl Drop for HypervLinuxDriver {

src/hyperlight_host/src/hypervisor/hyperv_windows.rs

@@ -41,7 +41,7 @@ use {
     std::sync::Mutex,
 };
 
-#[cfg(feature = "unwind_guest")]
+#[cfg(feature = "trace_guest")]
 use super::TraceRegister;
 use super::fpu::{FP_TAG_WORD_DEFAULT, MXCSR_DEFAULT};
 use super::handlers::{MemAccessHandlerWrapper, OutBHandlerWrapper};
@@ -723,6 +723,12 @@ impl Hypervisor for HypervWindowsDriver {
                 Reserved: Default::default(),
             }
         } else {
+            #[cfg(feature = "trace_guest")]
+            if self.trace_info.guest_start_epoch.is_none() {
+                // Set the guest start epoch to the current time, before running the vcpu
+                crate::debug!("MSHV - Guest Start Epoch set");
+                self.trace_info.guest_start_epoch = Some(std::time::Instant::now());
+            }
             self.processor.run()?
         };
         self.interrupt_handle
@@ -1037,7 +1043,7 @@ impl Hypervisor for HypervWindowsDriver {
         Ok(())
     }
 
-    #[cfg(feature = "unwind_guest")]
+    #[cfg(feature = "trace_guest")]
     fn read_trace_reg(&self, reg: TraceRegister) -> Result<u64> {
         let regs = self.processor.get_regs()?;
         match reg {
@@ -1053,6 +1059,10 @@ impl Hypervisor for HypervWindowsDriver {
     fn trace_info_as_ref(&self) -> &TraceInfo {
         &self.trace_info
     }
+    #[cfg(feature = "trace_guest")]
+    fn trace_info_as_mut(&mut self) -> &mut TraceInfo {
+        &mut self.trace_info
+    }
 }
 
 impl Drop for HypervWindowsDriver {

src/hyperlight_host/src/hypervisor/kvm.rs

@@ -29,7 +29,7 @@ use tracing::{Span, instrument};
 #[cfg(crashdump)]
 use {super::crashdump, std::path::Path};
 
-#[cfg(feature = "unwind_guest")]
+#[cfg(feature = "trace_guest")]
 use super::TraceRegister;
 use super::fpu::{FP_CONTROL_WORD_DEFAULT, FP_TAG_WORD_DEFAULT, MXCSR_DEFAULT};
 #[cfg(gdb)]
@@ -627,6 +627,13 @@ impl Hypervisor for KVMDriver {
         {
             Err(kvm_ioctls::Error::new(libc::EINTR))
         } else {
+            #[cfg(feature = "trace_guest")]
+            if self.trace_info.guest_start_epoch.is_none() {
+                // Set the guest start epoch to the current time, before running the vcpu
+                crate::debug!("KVM - Guest Start Epoch set");
+                self.trace_info.guest_start_epoch = Some(std::time::Instant::now());
+            }
+
             // Note: if a `InterruptHandle::kill()` called while this thread is **here**
             // Then the vcpu will run, but we will keep sending signals to this thread
             // to interrupt it until `running` is set to false. The `vcpu_fd::run()` call will
@@ -949,7 +956,7 @@ impl Hypervisor for KVMDriver {
         Ok(())
     }
 
-    #[cfg(feature = "unwind_guest")]
+    #[cfg(feature = "trace_guest")]
     fn read_trace_reg(&self, reg: TraceRegister) -> Result<u64> {
         let regs = self.vcpu_fd.get_regs()?;
         Ok(match reg {
@@ -965,6 +972,10 @@ impl Hypervisor for KVMDriver {
     fn trace_info_as_ref(&self) -> &TraceInfo {
         &self.trace_info
     }
+    #[cfg(feature = "trace_guest")]
+    fn trace_info_as_mut(&mut self) -> &mut TraceInfo {
+        &mut self.trace_info
+    }
 }
 
 impl Drop for KVMDriver {

src/hyperlight_host/src/hypervisor/mod.rs

@@ -270,12 +270,15 @@ pub(crate) trait Hypervisor: Debug + Sync + Send {
     }
 
     /// Read a register for trace/unwind purposes
-    #[cfg(feature = "unwind_guest")]
+    #[cfg(feature = "trace_guest")]
     fn read_trace_reg(&self, reg: TraceRegister) -> Result<u64>;
 
     /// Get a reference of the trace info for the guest
     #[cfg(feature = "trace_guest")]
     fn trace_info_as_ref(&self) -> &TraceInfo;
+    /// Get a mutable reference of the trace info for the guest
+    #[cfg(feature = "trace_guest")]
+    fn trace_info_as_mut(&mut self) -> &mut TraceInfo;
 }
 
 /// A virtual CPU that can be run until an exit occurs

src/hyperlight_host/src/sandbox/mod.rs

@@ -102,6 +102,17 @@ pub(crate) struct TraceInfo {
     /// early as the creation of the sandbox being traced.
     #[allow(dead_code)]
     pub epoch: std::time::Instant,
+    /// The frequency of the timestamp counter.
+    #[allow(dead_code)]
+    pub tsc_freq: u64,
+    /// The epoch at which the guest started, if it has started.
+    /// This is used to calculate the time spent in the guest relative to the
+    /// time of the host.
+    #[allow(dead_code)]
+    pub guest_start_epoch: Option<std::time::Instant>,
+    /// The start guest time, in TSC cycles, for the current guest.
+    #[allow(dead_code)]
+    pub guest_start_tsc: Option<u64>,
     /// The file to which the trace is being written
     #[allow(dead_code)]
     pub file: Arc<Mutex<std::fs::File>>,
@@ -136,8 +147,13 @@ impl TraceInfo {
             let cache = framehop::x86_64::CacheX86_64::new();
             (unwinder, Arc::new(Mutex::new(cache)))
         };
+        let tsc_freq = Self::calculate_tsc_freq()?;
+
         let ret = Self {
             epoch: std::time::Instant::now(),
+            tsc_freq,
+            guest_start_epoch: None,
+            guest_start_tsc: None,
             file: Arc::new(Mutex::new(std::fs::File::create_new(path)?)),
             #[cfg(feature = "unwind_guest")]
             unwind_module,
@@ -153,6 +169,24 @@ impl TraceInfo {
         })?;
         Ok(ret)
     }
+
+    /// Calculate the TSC frequency based on the RDTSC instruction.
+    fn calculate_tsc_freq() -> crate::Result<u64> {
+        if !hyperlight_guest_tracing::invariant_tsc::has_invariant_tsc() {
+            return Err(crate::new_error!(
+                "Invariant TSC is not supported on this platform"
+            ));
+        }
+        let start = hyperlight_guest_tracing::invariant_tsc::read_tsc();
+        let start_time = std::time::Instant::now();
+        // Sleep for 1 second to get a good sample
+        std::thread::sleep(std::time::Duration::from_secs(1));
+        let end = hyperlight_guest_tracing::invariant_tsc::read_tsc();
+        let end_time = std::time::Instant::now();
+        let elapsed = end_time.duration_since(start_time).as_secs_f64();
+
+        Ok(((end - start) as f64 / elapsed) as u64)
+    }
 }
 
 pub(crate) trait WrapperGetter {

src/hyperlight_host/src/sandbox/outb.rs

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-#[cfg(feature = "unwind_guest")]
+#[cfg(feature = "trace_guest")]
 use std::io::Write;
 use std::sync::{Arc, Mutex};
 
@@ -26,14 +26,16 @@ use hyperlight_common::flatbuffer_wrappers::function_types::ParameterValue;
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData;
 use hyperlight_common::outb::{Exception, OutBAction};
+#[cfg(feature = "trace_guest")]
+use hyperlight_guest_tracing::TraceRecord;
 use log::{Level, Record};
 use tracing::{Span, instrument};
 use tracing_log::format_trace;
 
 use super::host_funcs::FunctionRegistry;
 use super::mem_mgr::MemMgrWrapper;
 use crate::hypervisor::handlers::{OutBHandler, OutBHandlerFunction, OutBHandlerWrapper};
-#[cfg(feature = "unwind_guest")]
+#[cfg(feature = "trace_guest")]
 use crate::mem::layout::SandboxMemoryLayout;
 use crate::mem::mgr::SandboxMemoryManager;
 use crate::mem::shared_mem::HostSharedMemory;
@@ -211,6 +213,53 @@ pub(super) fn record_trace_frame<F: FnOnce(&mut std::fs::File)>(
     Ok(())
 }
 
+#[cfg(feature = "trace_guest")]
+pub(super) fn record_guest_trace_frame<F: FnOnce(&mut std::fs::File)>(
+    trace_info: &TraceInfo,
+    frame_id: u64,
+    cycles: u64,
+    write_frame: F,
+) -> Result<()> {
+    let Ok(mut out) = trace_info.file.lock() else {
+        return Ok(());
+    };
+    // frame structure:
+    // 16 bytes timestamp
+
+    // The number of cycles spent in the guest relative to the first received trace record
+    let cycles_spent = cycles
+        - trace_info
+            .guest_start_tsc
+            .as_ref()
+            .map_or_else(|| 0, |c| *c);
+    // Convert cycles to microseconds based on the TSC frequency
+    let micros = cycles_spent as f64 / trace_info.tsc_freq as f64 * 1_000_000f64;
+
+    // Convert to a Duration
+    let guest_duration = std::time::Duration::from_micros(micros as u64);
+
+    // Calculate the time when the guest started execution relative to the host epoch
+    // Note: This is relative to the time saved when the `TraceInfo` was created (before the
+    // Hypervisor is created).
+    let guest_start_time = trace_info
+        .guest_start_epoch
+        .as_ref()
+        .unwrap_or(&trace_info.epoch)
+        .saturating_duration_since(trace_info.epoch);
+
+    // Calculate the timestamp when the actual frame was recorded relative to the host epoch
+    let timestamp = guest_start_time
+        .checked_add(guest_duration)
+        .unwrap_or(guest_duration);
+
+    let _ = out.write_all(&timestamp.as_micros().to_ne_bytes());
+    // 8 bytes frame type id
+    let _ = out.write_all(&frame_id.to_ne_bytes());
+    // frame data
+    write_frame(&mut out);
+    Ok(())
+}
+
 /// Handles OutB operations from the guest.
 #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
 fn handle_outb_impl(
@@ -287,6 +336,54 @@ fn handle_outb_impl(
                 write_stack(f, &stack);
             })
         }
+        #[cfg(feature = "trace_guest")]
+        OutBAction::TraceRecord => {
+            let Ok(len) = _hv.read_trace_reg(crate::hypervisor::TraceRegister::RAX) else {
+                return Ok(());
+            };
+            let Ok(ptr) = _hv.read_trace_reg(crate::hypervisor::TraceRegister::RCX) else {
+                return Ok(());
+            };
+            let mut buffer = vec![0u8; len as usize * std::mem::size_of::<TraceRecord>()];
+            let buffer = &mut buffer[..];
+
+            // Read the trace records from the guest memory
+            // TODO: maybe this can be done without copying?
+            mem_mgr
+                .as_ref()
+                .shared_mem
+                .copy_to_slice(buffer, ptr as usize - SandboxMemoryLayout::BASE_ADDRESS)
+                // .read::<u64>((addr - SandboxMemoryLayout::BASE_ADDRESS as u64) as usize)
+                .map_err(|e| {
+                    new_error!(
+                        "Failed to copy trace records from guest memory to host: {:?}",
+                        e
+                    )
+                })?;
+
+            let traces = unsafe {
+                std::slice::from_raw_parts(buffer.as_ptr() as *const TraceRecord, len as usize)
+            };
+
+            {
+                let trace_info = _hv.trace_info_as_mut();
+                // Store the start guest cycles if not already set
+                // This is the `entrypoint` of the guest execution
+                // This should be set only once, at the start of the guest execution
+                if trace_info.guest_start_tsc.is_none() && !traces.is_empty() {
+                    trace_info.guest_start_tsc = Some(traces[0].cycles);
+                }
+            }
+
+            for record in traces {
+                record_guest_trace_frame(_hv.trace_info_as_ref(), 4u64, record.cycles, |f| {
+                    let _ = f.write_all(&record.msg_len.to_ne_bytes());
+                    let _ = f.write_all(&record.msg[..record.msg_len]);
+                })?
+            }
+
+            Ok(())
+        }
     }
 }