Avoid passing the guest dispatch pointer in memory

In preparation for the snapshot becoming immutable, this gets rid of
on of the last places that assume identity mapping & require early
memory writes in the sandbox.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>

Author: syntactically

src/hyperlight_common/src/mem.rs

@@ -31,7 +31,6 @@ pub struct GuestMemoryRegion {
 #[derive(Debug, Clone, Copy)]
 #[repr(C)]
 pub struct HyperlightPEB {
-    pub guest_function_dispatch_ptr: u64,
     pub input_stack: GuestMemoryRegion,
     pub output_stack: GuestMemoryRegion,
     pub init_data: GuestMemoryRegion,

src/hyperlight_guest_bin/src/lib.rs

@@ -178,8 +178,13 @@ unsafe extern "C" {
 /// Architecture-nonspecific initialisation: set up the heap,
 /// coordinate some addresses and configuration with the host, and run
 /// user initialisation
-pub(crate) extern "C" fn generic_init(peb_address: u64, seed: u64, ops: u64, max_log_level: u64) {
-    let peb_ptr = unsafe {
+pub(crate) extern "C" fn generic_init(
+    peb_address: u64,
+    seed: u64,
+    ops: u64,
+    max_log_level: u64,
+) -> u64 {
+    unsafe {
         GUEST_HANDLE = GuestHandle::init(peb_address as *mut HyperlightPEB);
         #[allow(static_mut_refs)]
         let peb_ptr = GUEST_HANDLE.peb().unwrap();
@@ -202,8 +207,6 @@ pub(crate) extern "C" fn generic_init(peb_address: u64, seed: u64, ops: u64, max
     let guest_start_tsc = hyperlight_guest_tracing::invariant_tsc::read_tsc();
 
     unsafe {
-        (*peb_ptr).guest_function_dispatch_ptr = dispatch_function as usize as u64;
-
         let srand_seed = (((peb_address << 8) ^ (seed >> 4)) >> 32) as u32;
         // Set the seed for the random number generator for C code using rand;
         srand(srand_seed);
@@ -231,6 +234,8 @@ pub(crate) extern "C" fn generic_init(peb_address: u64, seed: u64, ops: u64, max
     unsafe {
         hyperlight_main();
     }
+
+    dispatch_function as usize as u64
 }
 
 #[cfg(feature = "macros")]

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -68,6 +68,7 @@ use crate::metrics::{METRIC_ERRONEOUS_VCPU_KICKS, METRIC_GUEST_CANCELLATION};
 use crate::sandbox::SandboxConfiguration;
 use crate::sandbox::host_funcs::FunctionRegistry;
 use crate::sandbox::outb::{HandleOutbError, handle_outb};
+use crate::sandbox::snapshot::NextAction;
 #[cfg(feature = "mem_profile")]
 use crate::sandbox::trace::MemTraceInfo;
 #[cfg(crashdump)]
@@ -85,7 +86,7 @@ pub(crate) struct HyperlightVm {
     #[cfg(not(gdb))]
     vm: Box<dyn VirtualMachine>,
     page_size: usize,
-    entrypoint: Option<u64>, // only present if this vm has not yet been initialised
+    entrypoint: NextAction, // only present if this vm has not yet been initialised
     rsp_gva: u64,
     interrupt_handle: Arc<dyn InterruptHandleImpl>,
 
@@ -120,6 +121,8 @@ pub enum DispatchGuestCallError {
     Run(#[from] RunVmError),
     #[error("Failed to setup registers: {0}")]
     SetupRegs(RegisterError),
+    #[error("VM was uninitialized")]
+    Uninitialized,
 }
 
 impl DispatchGuestCallError {
@@ -129,7 +132,7 @@ impl DispatchGuestCallError {
             // These errors poison the sandbox because they can leave it in an inconsistent state
             // by returning before the guest can unwind properly
             DispatchGuestCallError::Run(_) => true,
-            DispatchGuestCallError::SetupRegs(_) => false,
+            DispatchGuestCallError::SetupRegs(_) | DispatchGuestCallError::Uninitialized => false,
         }
     }
 
@@ -338,7 +341,7 @@ impl HyperlightVm {
         snapshot_mem: GuestSharedMemory,
         scratch_mem: GuestSharedMemory,
         _pml4_addr: u64,
-        entrypoint: Option<u64>,
+        entrypoint: NextAction,
         rsp_gva: u64,
         #[cfg_attr(target_os = "windows", allow(unused_variables))] config: &SandboxConfiguration,
         #[cfg(gdb)] gdb_conn: Option<DebugCommChannel<DebugResponse, DebugMsg>>,
@@ -438,9 +441,9 @@ impl HyperlightVm {
             ret.send_dbg_msg(DebugResponse::InterruptHandle(ret.interrupt_handle.clone()))?;
             // Add breakpoint to the entry point address, if we are going to initialise
             ret.vm.set_debug(true).map_err(VmError::Debug)?;
-            if let Some(entrypoint) = entrypoint {
+            if let NextAction::Initialise(initialise) = entrypoint {
                 ret.vm
-                    .add_hw_breakpoint(entrypoint)
+                    .add_hw_breakpoint(initialise)
                     .map_err(CreateHyperlightVmError::AddHwBreakpoint)?;
             }
         }
@@ -462,7 +465,7 @@ impl HyperlightVm {
         guest_max_log_level: Option<LevelFilter>,
         #[cfg(gdb)] dbg_mem_access_fn: Arc<Mutex<SandboxMemoryManager<HostSharedMemory>>>,
     ) -> std::result::Result<(), InitializeError> {
-        let Some(entrypoint) = self.entrypoint else {
+        let NextAction::Initialise(initialise) = self.entrypoint else {
             return Ok(());
         };
 
@@ -474,7 +477,7 @@ impl HyperlightVm {
         };
 
         let regs = CommonRegisters {
-            rip: entrypoint,
+            rip: initialise,
             // We usually keep the top of the stack 16-byte
             // aligned. However, the ABI requirement is that the stack
             // be aligned _before a call instruction_, which means
@@ -508,6 +511,7 @@ impl HyperlightVm {
             return Err(InitializeError::InvalidStackPointer(regs.rsp));
         }
         self.rsp_gva = regs.rsp;
+        self.entrypoint = NextAction::Call(regs.rax);
 
         Ok(())
     }
@@ -631,6 +635,16 @@ impl HyperlightVm {
         self.rsp_gva = gva;
     }
 
+    /// Get the current entrypoint action
+    pub(crate) fn get_entrypoint(&self) -> NextAction {
+        self.entrypoint
+    }
+
+    /// Set the current entrypoint action
+    pub(crate) fn set_entrypoint(&mut self, entrypoint: NextAction) {
+        self.entrypoint = entrypoint
+    }
+
     /// Dispatch a call from the host to the guest using the given pointer
     /// to the dispatch function _in the guest's address space_.
     ///
@@ -641,14 +655,16 @@ impl HyperlightVm {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     pub(crate) fn dispatch_call_from_host(
         &mut self,
-        dispatch_func_addr: RawPtr,
         mem_mgr: &mut SandboxMemoryManager<HostSharedMemory>,
         host_funcs: &Arc<Mutex<FunctionRegistry>>,
         #[cfg(gdb)] dbg_mem_access_fn: Arc<Mutex<SandboxMemoryManager<HostSharedMemory>>>,
     ) -> std::result::Result<(), DispatchGuestCallError> {
+        let NextAction::Call(dispatch_func_addr) = self.entrypoint else {
+            return Err(DispatchGuestCallError::Uninitialized);
+        };
         // set RIP and RSP, reset others
         let regs = CommonRegisters {
-            rip: dispatch_func_addr.into(),
+            rip: dispatch_func_addr,
             // We usually keep the top of the stack 16-byte
             // aligned. However, the ABI requirement is that the stack
             // be aligned _before a call instruction_, which means
@@ -755,13 +771,13 @@ impl HyperlightVm {
             match exit_reason {
                 #[cfg(gdb)]
                 Ok(VmExit::Debug { dr6, exception }) => {
+                    let initialise = match self.entrypoint {
+                        NextAction::Initialise(initialise) => initialise,
+                        _ => 0,
+                    };
                     // Handle debug event (breakpoints)
-                    let stop_reason = arch::vcpu_stop_reason(
-                        self.vm.as_mut(),
-                        dr6,
-                        self.entrypoint.unwrap_or(0),
-                        exception,
-                    )?;
+                    let stop_reason =
+                        arch::vcpu_stop_reason(self.vm.as_mut(), dr6, initialise, exception)?;
                     if let Err(e) = self.handle_debug(dbg_mem_access_fn.clone(), stop_reason) {
                         break Err(e.into());
                     }
@@ -1133,14 +1149,19 @@ impl HyperlightVm {
                     .and_then(|name| name.to_os_string().into_string().ok())
             });
 
+            let initialise = match self.entrypoint {
+                NextAction::Initialise(initialise) => initialise,
+                _ => 0,
+            };
+
             // Include dynamically mapped regions
             // TODO: include the snapshot and scratch regions
             let regions: Vec<MemoryRegion> = self.get_mapped_regions().cloned().collect();
             Ok(Some(crashdump::CrashDumpContext::new(
                 regions,
                 regs,
                 xsave.to_vec(),
-                self.entrypoint.unwrap_or(0),
+                initialise,
                 self.rt_cfg.binary_path.clone(),
                 filename,
             )))

src/hyperlight_host/src/mem/layout.rs

@@ -90,7 +90,6 @@ pub(crate) struct SandboxMemoryLayout {
     /// The following fields are offsets to the actual PEB struct fields.
     /// They are used when writing the PEB struct itself
     peb_offset: usize,
-    peb_guest_dispatch_function_ptr_offset: usize, // set by guest in guest entrypoint
     peb_input_data_offset: usize,
     peb_output_data_offset: usize,
     peb_init_data_offset: usize,
@@ -129,10 +128,6 @@ impl Debug for SandboxMemoryLayout {
             .field("PEB Address", &format_args!("{:#x}", self.peb_address))
             .field("PEB Offset", &format_args!("{:#x}", self.peb_offset))
             .field("Code Size", &format_args!("{:#x}", self.code_size))
-            .field(
-                "Guest Dispatch Function Pointer Offset",
-                &format_args!("{:#x}", self.peb_guest_dispatch_function_ptr_offset),
-            )
             .field(
                 "Input Data Offset",
                 &format_args!("{:#x}", self.peb_input_data_offset),
@@ -216,8 +211,6 @@ impl SandboxMemoryLayout {
         let guest_code_offset = 0;
         // The following offsets are to the fields of the PEB struct itself!
         let peb_offset = code_size.next_multiple_of(PAGE_SIZE_USIZE);
-        let peb_guest_dispatch_function_ptr_offset =
-            peb_offset + offset_of!(HyperlightPEB, guest_function_dispatch_ptr);
         let peb_input_data_offset = peb_offset + offset_of!(HyperlightPEB, input_stack);
         let peb_output_data_offset = peb_offset + offset_of!(HyperlightPEB, output_stack);
         let peb_init_data_offset = peb_offset + offset_of!(HyperlightPEB, init_data);
@@ -238,7 +231,6 @@ impl SandboxMemoryLayout {
         Ok(Self {
             peb_offset,
             heap_size,
-            peb_guest_dispatch_function_ptr_offset,
             peb_input_data_offset,
             peb_output_data_offset,
             peb_init_data_offset,
@@ -344,13 +336,6 @@ impl SandboxMemoryLayout {
         .next_multiple_of(hyperlight_common::vmem::PAGE_SIZE) as u64
     }
 
-    /// Get the offset in guest memory to where the guest dispatch function
-    /// pointer is written
-    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
-    pub(super) fn get_dispatch_function_pointer_offset(&self) -> usize {
-        self.peb_guest_dispatch_function_ptr_offset
-    }
-
     /// Get the offset in guest memory to the heap size
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     fn get_heap_size_offset(&self) -> usize {

src/hyperlight_host/src/mem/mgr.rs

@@ -24,11 +24,10 @@ use tracing::{Span, instrument};
 
 use super::layout::SandboxMemoryLayout;
 use super::memory_region::MemoryRegion;
-use super::ptr::{GuestPtr, RawPtr};
-use super::ptr_offset::Offset;
+use super::ptr::RawPtr;
 use super::shared_mem::{ExclusiveSharedMemory, GuestSharedMemory, HostSharedMemory, SharedMemory};
 use crate::hypervisor::regs::CommonSpecialRegisters;
-use crate::sandbox::snapshot::Snapshot;
+use crate::sandbox::snapshot::{NextAction, Snapshot};
 use crate::{Result, new_error};
 
 /// A struct that is responsible for laying out and managing the memory
@@ -44,7 +43,7 @@ pub(crate) struct SandboxMemoryManager<S> {
     /// Pointer to where to load memory from
     pub(crate) load_addr: RawPtr,
     /// Offset for the execution entrypoint from `load_addr`
-    pub(crate) entrypoint_offset: Option<Offset>,
+    pub(crate) entrypoint: NextAction,
     /// How many memory regions were mapped after sandbox creation
     pub(crate) mapped_rgns: u64,
     /// Buffer for accumulating guest abort messages
@@ -152,14 +151,14 @@ where
         shared_mem: S,
         scratch_mem: S,
         load_addr: RawPtr,
-        entrypoint_offset: Option<Offset>,
+        entrypoint: NextAction,
     ) -> Self {
         Self {
             layout,
             shared_mem,
             scratch_mem,
             load_addr,
-            entrypoint_offset,
+            entrypoint,
             mapped_rgns: 0,
             abort_buffer: Vec::new(),
         }
@@ -176,12 +175,6 @@ where
         &mut self.shared_mem
     }
 
-    /// Get `SharedMemory` in `self` as a mutable reference
-    #[cfg(any(gdb, test))]
-    pub(crate) fn get_scratch_mem_mut(&mut self) -> &mut S {
-        &mut self.scratch_mem
-    }
-
     /// Create a snapshot with the given mapped regions
     pub(crate) fn snapshot(
         &mut self,
@@ -190,6 +183,7 @@ where
         root_pt_gpa: u64,
         rsp_gva: u64,
         sregs: CommonSpecialRegisters,
+        entrypoint: NextAction,
     ) -> Result<Snapshot> {
         Snapshot::new(
             &mut self.shared_mem,
@@ -201,6 +195,7 @@ where
             root_pt_gpa,
             rsp_gva,
             sregs,
+            entrypoint,
         )
     }
 }
@@ -212,14 +207,13 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
         shared_mem.copy_from_slice(s.memory(), 0)?;
         let scratch_mem = ExclusiveSharedMemory::new(s.layout().get_scratch_size())?;
         let load_addr: RawPtr = RawPtr::try_from(layout.get_guest_code_address())?;
-        let entrypoint_gva = s.preinitialise();
-        let entrypoint_offset = entrypoint_gva.map(|x| (x - u64::from(&load_addr)).into());
+        let entrypoint = s.entrypoint();
         Ok(Self::new(
             layout,
             shared_mem,
             scratch_mem,
             load_addr,
-            entrypoint_offset,
+            entrypoint,
         ))
     }
 
@@ -257,7 +251,7 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
             scratch_mem: hscratch,
             layout: self.layout,
             load_addr: self.load_addr.clone(),
-            entrypoint_offset: self.entrypoint_offset,
+            entrypoint: self.entrypoint,
             mapped_rgns: self.mapped_rgns,
             abort_buffer: self.abort_buffer,
         };
@@ -266,7 +260,7 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
             scratch_mem: gscratch,
             layout: self.layout,
             load_addr: self.load_addr.clone(),
-            entrypoint_offset: self.entrypoint_offset,
+            entrypoint: self.entrypoint,
             mapped_rgns: self.mapped_rgns,
             abort_buffer: Vec::new(), // Guest doesn't need abort buffer
         };
@@ -278,20 +272,6 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
 }
 
 impl SandboxMemoryManager<HostSharedMemory> {
-    /// Get the address of the dispatch function in memory
-    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-    pub(crate) fn get_pointer_to_dispatch_function(&self) -> Result<u64> {
-        let guest_dispatch_function_ptr = self
-            .shared_mem
-            .read::<u64>(self.layout.get_dispatch_function_pointer_offset())?;
-
-        // This pointer is written by the guest library but is accessible to
-        // the guest engine so we should bounds check it before we return it.
-
-        let guest_ptr = GuestPtr::try_from(RawPtr::from(guest_dispatch_function_ptr))?;
-        guest_ptr.absolute()
-    }
-
     /// Reads a host function call from memory
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     pub(crate) fn get_host_function_call(&mut self) -> Result<FunctionCall> {

src/hyperlight_host/src/mem/shared_mem.rs

@@ -334,7 +334,10 @@ impl ExclusiveSharedMemory {
     #[cfg(target_os = "linux")]
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     pub fn new(min_size_bytes: usize) -> Result<Self> {
-        use libc::{MAP_ANONYMOUS, MAP_PRIVATE, MAP_FAILED, PROT_READ, PROT_WRITE, c_int, mmap, off_t, size_t};
+        use libc::{
+            MAP_ANONYMOUS, MAP_FAILED, MAP_PRIVATE, PROT_READ, PROT_WRITE, c_int, mmap, off_t,
+            size_t,
+        };
         #[cfg(not(miri))]
         use libc::{MAP_NORESERVE, PROT_NONE, mprotect};