[trace] Walk page table entries to get the tracing data

- With the new Copy-on-Write changes, the guest virtual addresses and
  guest physical addresses are no longer identity mapped, so we need a
  way to do this translation when a new guest trace batch arrives from
  the guest

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

Author: dblnz

src/hyperlight_common/src/flatbuffer_wrappers/guest_trace_data.rs

@@ -40,6 +40,15 @@ use crate::flatbuffers::hyperlight::generated::{
     OpenSpanTypeArgs as FbOpenSpanTypeArgs,
 };
 
+/// TODO: Change these constant to be configurable at runtime by the guest
+/// Maybe use a weak symbol that the guest can override at link time?
+///
+/// Pre-calculated capacity for the encoder buffer
+/// This is to avoid reallocations in the guest
+/// If the next event would exceed this size, the encoder will flush the current buffer to the host
+/// before encoding the new event.
+pub const MAX_TRACE_DATA_SIZE: usize = 4096;
+
 /// Key-Value pair structure used in tracing spans/events
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct EventKeyValue {

src/hyperlight_common/src/vmem.rs

@@ -222,8 +222,8 @@ pub struct Mapping {
 ///   are being remapped, TLB invalidation may need to be performed
 ///   afterwards.
 pub use arch::map;
-/// This function is not presently used for anything, but is useful
-/// for debugging
+/// This function is presently used for reading the tracing data, also
+/// it is useful for debugging
 ///
 /// # Safety
 /// This function traverses page table data structures, and should not

src/hyperlight_guest_tracing/src/state.rs

@@ -20,7 +20,7 @@ use alloc::vec::Vec;
 use core::sync::atomic::{AtomicU64, Ordering};
 
 use hyperlight_common::flatbuffer_wrappers::guest_trace_data::{
-    EventsBatchEncoder, EventsEncoder, GuestEvent,
+    EventsBatchEncoder, EventsEncoder, GuestEvent, MAX_TRACE_DATA_SIZE,
 };
 use hyperlight_common::outb::OutBAction;
 use tracing_core::Event;
@@ -43,12 +43,6 @@ pub(crate) struct GuestState {
     stack: Vec<u64>,
 }
 
-/// TODO: Change these constant to be configurable at runtime by the guest
-/// Maybe use a weak symbol that the guest can override at link time?
-///
-/// Pre-calculated capacity for the encoder buffer
-/// This is to avoid reallocations in the guest
-const ENCODER_CAPACITY: usize = 4096;
 /// Start with a stack capacity for active spans
 const ACTIVE_SPANS_CAPACITY: usize = 64;
 
@@ -69,7 +63,7 @@ fn send_to_host(data: &[u8]) {
 
 impl GuestState {
     pub(crate) fn new(guest_start_tsc: u64) -> Self {
-        let mut encoder = EventsBatchEncoder::new(ENCODER_CAPACITY, send_to_host);
+        let mut encoder = EventsBatchEncoder::new(MAX_TRACE_DATA_SIZE, send_to_host);
         encoder.encode(&GuestEvent::GuestStart {
             tsc: guest_start_tsc,
         });

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -191,6 +191,8 @@ pub enum RunVmError {
     DebugHandler(#[from] HandleDebugError),
     #[error("Execution was cancelled by the host")]
     ExecutionCancelledByHost,
+    #[error("Failed to access page: {0}")]
+    PageTableAccess(AccessPageTableError),
     #[cfg(feature = "trace_guest")]
     #[error("Failed to get registers: {0}")]
     GetRegs(RegisterError),
@@ -754,10 +756,18 @@ impl HyperlightVm {
                     tc.end_host_trace();
                     // Handle the guest trace data if any
                     let regs = self.vm.regs().map_err(RunVmError::GetRegs)?;
-                    if let Err(e) = tc.handle_trace(&regs, mem_mgr) {
-                        // If no trace data is available, we just log a message and continue
-                        // Is this the right thing to do?
-                        log::debug!("Error handling guest trace: {:?}", e);
+
+                    // Only parse the trace if it has reported
+                    if tc.has_trace_data(&regs) {
+                        let root_pt = self.get_root_pt().map_err(RunVmError::PageTableAccess)?;
+
+                        // If something goes wrong with parsing the trace data, we log the error and
+                        // continue execution instead of returning an error since this is not critical
+                        // to correct execution of the guest
+                        tc.handle_trace(&regs, mem_mgr, root_pt)
+                            .unwrap_or_else(|e| {
+                                tracing::error!("Cannot handle trace data: {}", e);
+                            });
                     }
                 }
                 result

src/hyperlight_host/src/mem/mgr.rs

@@ -411,6 +411,117 @@ impl SandboxMemoryManager<HostSharedMemory> {
 
         Ok(())
     }
+
+    /// Read guest memory at a Guest Virtual Address (GVA) by walking the
+    /// page tables to translate GVA → GPA, then reading from the correct
+    /// backing memory (shared_mem or scratch_mem).
+    ///
+    /// This is necessary because with Copy-on-Write (CoW) the guest's
+    /// virtual pages are backed by physical pages in the scratch
+    /// region rather than being identity-mapped.
+    ///
+    /// # Arguments
+    /// * `gva` - The Guest Virtual Address to read from
+    /// * `len` - The number of bytes to read
+    /// * `root_pt` - The root page table physical address (CR3)
+    #[cfg(feature = "trace_guest")]
+    pub(crate) fn read_guest_memory_by_gva(
+        &mut self,
+        gva: u64,
+        len: usize,
+        root_pt: u64,
+    ) -> Result<Vec<u8>> {
+        use hyperlight_common::vmem::PAGE_SIZE;
+
+        use crate::sandbox::snapshot::{SharedMemoryPageTableBuffer, access_gpa};
+
+        let scratch_size = self.scratch_mem.mem_size();
+
+        self.shared_mem.with_exclusivity(|snap| {
+            self.scratch_mem.with_exclusivity(|scratch| {
+                let pt_buf = SharedMemoryPageTableBuffer::new(snap, scratch, scratch_size, root_pt);
+
+                // Walk page tables to get all mappings that cover the GVA range
+                let mappings: Vec<_> = unsafe {
+                    hyperlight_common::vmem::virt_to_phys(&pt_buf, gva, len as u64)
+                }
+                .collect();
+
+                if mappings.is_empty() {
+                    return Err(new_error!(
+                        "No page table mappings found for GVA {:#x} (len {})",
+                        gva,
+                        len,
+                    ));
+                }
+
+                // Resulting vector of bytes to return
+                let mut result = Vec::with_capacity(len);
+                let mut current_gva = gva;
+
+                for mapping in &mappings {
+                    // The page table walker should only return valid mappings
+                    // that cover our current read position.
+                    if mapping.virt_base > current_gva {
+                        return Err(new_error!(
+                            "Page table walker returned mapping with virt_base {:#x} > current read position {:#x}",
+                            mapping.virt_base,
+                            current_gva,
+                        ));
+                    }
+
+                    // Calculate the offset within this page where to start copying
+                    let page_offset = (current_gva - mapping.virt_base) as usize;
+
+                    let bytes_remaining = len - result.len();
+                    let available_in_page = PAGE_SIZE - page_offset;
+                    let bytes_to_copy = bytes_remaining.min(available_in_page);
+
+                    // Translate the GPA to host memory
+                    let gpa = mapping.phys_base + page_offset as u64;
+                    let (mem, offset) = access_gpa(snap, scratch, scratch_size, gpa)
+                        .ok_or_else(|| {
+                            new_error!(
+                                "Failed to resolve GPA {:#x} to host memory (GVA {:#x})",
+                                gpa,
+                                gva
+                            )
+                        })?;
+
+                    let slice = mem
+                        .as_slice()
+                        .get(offset..offset + bytes_to_copy)
+                        .ok_or_else(|| {
+                            new_error!(
+                                "GPA {:#x} resolved to out-of-bounds host offset {} (need {} bytes)",
+                                gpa,
+                                offset,
+                                bytes_to_copy
+                            )
+                        })?;
+
+                    result.extend_from_slice(slice);
+                    current_gva += bytes_to_copy as u64;
+                }
+
+                if result.len() != len {
+                    tracing::error!(
+                        "Page table walker returned mappings that don't cover the full requested length: got {}, expected {}",
+                        result.len(),
+                        len,
+                    );
+                    return Err(new_error!(
+                        "Could not read full GVA range: got {} of {} bytes {:?}",
+                        result.len(),
+                        len,
+                        mappings
+                    ));
+                }
+
+                Ok(result)
+            })
+        })??
+    }
 }
 
 #[cfg(test)]

src/hyperlight_host/src/sandbox/snapshot.rs

@@ -174,7 +174,7 @@ fn hash(memory: &[u8], regions: &[MemoryRegion]) -> Result<[u8; 32]> {
     Ok(hasher.finalize().into())
 }
 
-fn access_gpa<'a>(
+pub(crate) fn access_gpa<'a>(
     snap: &'a ExclusiveSharedMemory,
     scratch: &'a ExclusiveSharedMemory,
     scratch_size: usize,
@@ -197,7 +197,7 @@ pub(crate) struct SharedMemoryPageTableBuffer<'a> {
     root: u64,
 }
 impl<'a> SharedMemoryPageTableBuffer<'a> {
-    fn new(
+    pub(crate) fn new(
         snap: &'a ExclusiveSharedMemory,
         scratch: &'a ExclusiveSharedMemory,
         scratch_size: usize,

src/hyperlight_host/src/sandbox/trace/context.rs

@@ -18,7 +18,7 @@ use std::collections::HashMap;
 use std::time::{Duration, Instant, SystemTime};
 
 use hyperlight_common::flatbuffer_wrappers::guest_trace_data::{
-    EventKeyValue, EventsBatchDecoder, EventsDecoder, GuestEvent,
+    EventKeyValue, EventsBatchDecoder, EventsDecoder, GuestEvent, MAX_TRACE_DATA_SIZE,
 };
 use hyperlight_common::outb::OutBAction;
 use opentelemetry::global::BoxedSpan;
@@ -28,73 +28,56 @@ use tracing::span::{EnteredSpan, Span};
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 
 use crate::hypervisor::regs::CommonRegisters;
-use crate::mem::layout::SandboxMemoryLayout;
 use crate::mem::mgr::SandboxMemoryManager;
-use crate::mem::shared_mem::{HostSharedMemory, SharedMemory};
-use crate::{HyperlightError, Result, new_error};
+use crate::mem::shared_mem::HostSharedMemory;
+use crate::{Result, new_error};
 
 /// Type that helps get the data from the guest provided the registers and memory access
 struct EventsBatch {
     events: Vec<GuestEvent>,
 }
 
-impl
-    TryFrom<(
-        &CommonRegisters,
-        &mut SandboxMemoryManager<HostSharedMemory>,
-    )> for EventsBatch
-{
-    type Error = HyperlightError;
-    fn try_from(
-        (regs, mem_mgr): (
-            &CommonRegisters,
-            &mut SandboxMemoryManager<HostSharedMemory>,
-        ),
+impl EventsBatch {
+    /// Extract a batch of guest trace events from guest memory.
+    ///
+    /// The guest passes the trace data pointer as a Guest Virtual Address (GVA)
+    /// in register r9. With Copy-on-Write enabled, this GVA may not be
+    /// identity-mapped to its physical address, so we walk the guest page
+    /// tables to translate GVA → GPA before reading the data.
+    ///
+    /// # Arguments
+    /// * `regs` - The guest registers (r8 = magic, r9 = GVA pointer, r10 = length)
+    /// * `mem_mgr` - The sandbox memory manager with access to shared and scratch memory
+    /// * `root_pt` - The root page table physical address (CR3) for GVA translation
+    fn from_regs(
+        regs: &CommonRegisters,
+        mem_mgr: &mut SandboxMemoryManager<HostSharedMemory>,
+        root_pt: u64,
     ) -> Result<Self> {
         let magic_no = regs.r8;
-        let trace_data_ptr = regs.r9 as usize;
+        let trace_data_gva = regs.r9;
         let trace_data_len = regs.r10 as usize;
 
+        // Validate the magic number to ensure the guest is providing trace data
         if magic_no != OutBAction::TraceBatch as u64 {
             return Err(new_error!("A TraceBatch is not present"));
         }
 
-        // Extract the GuestTraceData from guest memory
-        // This involves:
-        // 1. Using a mutable reference to the memory manager to get exclusive access to the shared memory.
-        //   This is necessary to ensure that no other part of the code is accessing the memory
-        //   while we are reading from it.
-        // 2. Getting immutable access to the slice of memory that contains the GuestTraceData
-        // 3. Parsing the slice into a GuestTraceData structure
-        //
-        // Error handling is done at each step to ensure that any issues are properly reported.
-        // This includes logging errors for easier debugging.
-        //
-        // The reason for using `with_exclusivity` is to ensure that we have exclusive access
-        // and avoid allocating new memory, which needs to be correctly aligned for the
-        // flatbuffer parsing.
-        let events = mem_mgr.shared_mem.with_exclusivity(|mem| {
-            let buf_slice = mem
-                .as_slice()
-                // Adjust the pointer to be relative to the base address of the sandbox memory
-                .get(
-                    trace_data_ptr - SandboxMemoryLayout::BASE_ADDRESS
-                        ..trace_data_ptr - SandboxMemoryLayout::BASE_ADDRESS + trace_data_len,
-                )
-                // Convert the slice to a Result to handle the case where the slice is out of
-                // bounds and return a proper error message and log the error.
-                .ok_or_else(|| {
-                    tracing::error!("Failed to get guest trace batch slice from guest memory");
-                    new_error!("Failed to get guest trace batch slice from guest memory")
-                })?;
+        // Validate the length to prevent reading excessive memory
+        if trace_data_len == 0 || trace_data_len > MAX_TRACE_DATA_SIZE {
+            return Err(new_error!("Invalid TraceBatch length: {}", trace_data_len));
+        }
 
-            let events = EventsBatchDecoder {}.decode(buf_slice).map_err(|e| {
-                tracing::error!("Failed to deserialize guest trace events: {:?}", e);
-                new_error!("Failed to deserialize guest trace events: {:?}", e)
-            })?;
+        // Read the trace data from guest memory by walking the page tables
+        // to translate the GVA to physical addresses. This is necessary
+        // because with CoW, guest virtual pages are backed by physical
+        // pages in the scratch region rather than being identity-mapped.
+        let buf = mem_mgr.read_guest_memory_by_gva(trace_data_gva, trace_data_len, root_pt)?;
 
-            Ok::<Vec<GuestEvent>, HyperlightError>(events)
-        })??;
+        let events = EventsBatchDecoder {}.decode(&buf).map_err(|e| {
+            tracing::error!("Failed to deserialize guest trace events: {:?}", e);
+            new_error!("Failed to deserialize guest trace events: {:?}", e)
+        })?;
 
         Ok(EventsBatch { events })
     }
@@ -215,13 +198,19 @@ impl TraceContext {
             + Duration::from_micros(rel_start_us as u64))
     }
 
+    /// Check if the registers indicate that there is trace data to be handled.
+    pub fn has_trace_data(&self, regs: &CommonRegisters) -> bool {
+        regs.r8 == OutBAction::TraceBatch as u64
+    }
+
     pub fn handle_trace(
         &mut self,
         regs: &CommonRegisters,
         mem_mgr: &mut SandboxMemoryManager<HostSharedMemory>,
+        root_pt: u64,
     ) -> Result<()> {
         // Get the guest sent info
-        let trace_batch = EventsBatch::try_from((regs, mem_mgr))?;
+        let trace_batch = EventsBatch::from_regs(regs, mem_mgr, root_pt)?;
 
         self.handle_trace_impl(trace_batch.events)
     }