refactor: move memory manager and host functions out of HyperlightVm

Remove MemoryManager and FunctionRegistry fields from HyperlightVm, passing
them as parameters to run() and dispatch_call_from_host() instead. This
eliminates duplicate ownership, simplifies the ownership model and gets
rid of a complete category of possible errors.

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

Author: ludfjig

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -76,8 +76,6 @@ pub(crate) struct HyperlightVm {
     entrypoint: u64,
     orig_rsp: GuestPtr,
     interrupt_handle: Arc<dyn InterruptHandleImpl>,
-    mem_mgr: Option<SandboxMemoryManager<HostSharedMemory>>,
-    host_funcs: Option<Arc<Mutex<FunctionRegistry>>>,
 
     sandbox_regions: Vec<MemoryRegion>, // Initially mapped regions when sandbox is created
     mmap_regions: Vec<(u32, MemoryRegion)>, // Later mapped regions (slot number, region)
@@ -193,8 +191,6 @@ impl HyperlightVm {
             orig_rsp: rsp_gp,
             interrupt_handle,
             page_size: 0, // Will be set in `initialise`
-            mem_mgr: None,
-            host_funcs: None,
 
             next_slot: mem_regions.len() as u32,
             sandbox_regions: mem_regions,
@@ -229,13 +225,11 @@ impl HyperlightVm {
         peb_addr: RawPtr,
         seed: u64,
         page_size: u32,
-        mem_mgr: SandboxMemoryManager<HostSharedMemory>,
+        mem_mgr: &mut SandboxMemoryManager<HostSharedMemory>,
         host_funcs: Arc<Mutex<FunctionRegistry>>,
         max_guest_log_level: Option<LevelFilter>,
         #[cfg(gdb)] dbg_mem_access_fn: Arc<Mutex<SandboxMemoryManager<HostSharedMemory>>>,
     ) -> Result<()> {
-        self.mem_mgr = Some(mem_mgr);
-        self.host_funcs = Some(host_funcs);
         self.page_size = page_size as usize;
 
         let max_guest_log_level: u64 = match max_guest_log_level {
@@ -259,6 +253,8 @@ impl HyperlightVm {
         self.vm.set_regs(&regs)?;
 
         self.run(
+            mem_mgr,
+            host_funcs,
             #[cfg(gdb)]
             dbg_mem_access_fn,
         )?;
@@ -269,6 +265,8 @@ impl HyperlightVm {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     pub(crate) fn dispatch_call_from_host(
         &mut self,
+        mem_mgr: &mut SandboxMemoryManager<HostSharedMemory>,
+        host_funcs: Arc<Mutex<FunctionRegistry>>,
         dispatch_func_addr: RawPtr,
         #[cfg(gdb)] dbg_mem_access_fn: Arc<Mutex<SandboxMemoryManager<HostSharedMemory>>>,
     ) -> Result<()> {
@@ -285,6 +283,8 @@ impl HyperlightVm {
         self.vm.set_fpu(&CommonFpu::default())?;
 
         self.run(
+            mem_mgr,
+            host_funcs,
             #[cfg(gdb)]
             dbg_mem_access_fn,
         )?;
@@ -325,61 +325,10 @@ impl HyperlightVm {
         &self.mmap_regions
     }
 
-    #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
-    fn handle_io(&mut self, port: u16, data: Vec<u8>) -> Result<()> {
-        if data.is_empty() {
-            log_then_return!("no data was given in IO interrupt");
-        }
-
-        #[allow(clippy::get_first)]
-        let val = u32::from_le_bytes([
-            data.get(0).copied().unwrap_or(0),
-            data.get(1).copied().unwrap_or(0),
-            data.get(2).copied().unwrap_or(0),
-            data.get(3).copied().unwrap_or(0),
-        ]);
-
-        #[cfg(feature = "mem_profile")]
-        {
-            // We need to handle the borrow checker issue where we need both:
-            // - &mut MemMgrWrapper (from self.mem_mgr.as_mut())
-            // - &mut dyn Hypervisor (from self)
-            // We'll use a temporary approach to extract the mem_mgr temporarily
-            let mem_mgr_option = self.mem_mgr.take();
-            let mut mem_mgr =
-                mem_mgr_option.ok_or_else(|| new_error!("mem_mgr not initialized"))?;
-            let host_funcs = self
-                .host_funcs
-                .as_ref()
-                .ok_or_else(|| new_error!("host_funcs not initialized"))?
-                .clone();
-
-            handle_outb(&mut mem_mgr, host_funcs, port, val, self)?;
-
-            // Put the mem_mgr back
-            self.mem_mgr = Some(mem_mgr);
-        }
-
-        #[cfg(not(feature = "mem_profile"))]
-        {
-            let mem_mgr = self
-                .mem_mgr
-                .as_mut()
-                .ok_or_else(|| new_error!("mem_mgr not initialized"))?;
-            let host_funcs = self
-                .host_funcs
-                .as_ref()
-                .ok_or_else(|| new_error!("host_funcs not initialized"))?
-                .clone();
-
-            handle_outb(mem_mgr, host_funcs, port, val)?;
-        }
-
-        Ok(())
-    }
-
     fn run(
         &mut self,
+        mem_mgr: &mut SandboxMemoryManager<HostSharedMemory>,
+        host_funcs: Arc<Mutex<FunctionRegistry>>,
         #[cfg(gdb)] dbg_mem_access_fn: Arc<Mutex<SandboxMemoryManager<HostSharedMemory>>>,
     ) -> Result<()> {
         // ===== KILL() TIMING POINT 1: Between guest function calls =====
@@ -425,7 +374,7 @@ impl HyperlightVm {
 
                 // Handle the guest trace data if any
                 #[cfg(feature = "trace_guest")]
-                if let Err(e) = self.handle_trace(&mut tc) {
+                if let Err(e) = tc.handle_trace(&self.vm.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);
@@ -469,7 +418,28 @@ impl HyperlightVm {
                 Ok(VmExit::Halt()) => {
                     break Ok(());
                 }
-                Ok(VmExit::IoOut(port, data)) => self.handle_io(port, data)?,
+                Ok(VmExit::IoOut(port, data)) => {
+                    if data.is_empty() {
+                        log_then_return!("no data was given in IO interrupt");
+                    }
+
+                    #[allow(clippy::get_first)]
+                    let val: u32 = u32::from_le_bytes([
+                        data.get(0).copied().unwrap_or(0),
+                        data.get(1).copied().unwrap_or(0),
+                        data.get(2).copied().unwrap_or(0),
+                        data.get(3).copied().unwrap_or(0),
+                    ]);
+                    handle_outb(
+                        mem_mgr,
+                        host_funcs.clone(),
+                        port,
+                        val,
+                        &self.vm.regs()?,
+                        #[cfg(feature = "mem_profile")]
+                        &mut self.trace_info,
+                    )?;
+                }
                 Ok(VmExit::MmioRead(addr)) => {
                     let all_regions = self
                         .sandbox_regions
@@ -491,17 +461,9 @@ impl HyperlightVm {
                             ));
                         }
                         None => {
-                            match &self.mem_mgr {
-                                Some(mem_mgr) => {
-                                    if !mem_mgr.check_stack_guard()? {
-                                        break Err(HyperlightError::StackOverflow());
-                                    }
-                                }
-                                None => {
-                                    break Err(new_error!("Memory manager not initialized"));
-                                }
-                            }
-
+                            // if !self.mem_mgr.check_stack_guard()? {
+                            //     break Err(HyperlightError::StackOverflow());
+                            // }
                             break Err(new_error!("MMIO READ access address {:#x}", addr));
                         }
                     }
@@ -527,15 +489,8 @@ impl HyperlightVm {
                             ));
                         }
                         None => {
-                            match &self.mem_mgr {
-                                Some(mem_mgr) => {
-                                    if !mem_mgr.check_stack_guard()? {
-                                        break Err(HyperlightError::StackOverflow());
-                                    }
-                                }
-                                None => {
-                                    break Err(new_error!("Memory manager not initialized"));
-                                }
+                            if mem_mgr.check_stack_guard()? {
+                                break Err(HyperlightError::StackOverflow());
                             }
 
                             break Err(new_error!("MMIO WRITE access address {:#x}", addr));
@@ -803,27 +758,6 @@ impl HyperlightVm {
             filename,
         ))
     }
-
-    #[cfg(feature = "mem_profile")]
-    pub(crate) fn vm_regs(&self) -> Result<CommonRegisters> {
-        self.vm.regs()
-    }
-
-    #[cfg(feature = "trace_guest")]
-    fn handle_trace(&mut self, tc: &mut crate::sandbox::trace::TraceContext) -> Result<()> {
-        let regs = self.vm.regs()?;
-        tc.handle_trace(
-            &regs,
-            self.mem_mgr.as_mut().ok_or_else(|| {
-                new_error!("Memory manager is not initialized before handling trace")
-            })?,
-        )
-    }
-
-    #[cfg(feature = "mem_profile")]
-    pub(crate) fn trace_info_mut(&mut self) -> &mut MemTraceInfo {
-        &mut self.trace_info
-    }
 }
 
 impl Drop for HyperlightVm {

src/hyperlight_host/src/hypervisor/mod.rs

@@ -492,9 +492,9 @@ pub(crate) mod tests {
         let rt_cfg: SandboxRuntimeConfig = Default::default();
         let sandbox =
             UninitializedSandbox::new(GuestBinary::FilePath(filename.clone()), Some(config))?;
-        let (mem_mgr, mut gshm) = sandbox.mgr.build();
+        let (mut hshm, _gshm) = sandbox.mgr.build();
         let mut vm = set_up_hypervisor_partition(
-            &mut gshm,
+            &mut hshm,
             &config,
             #[cfg(any(crashdump, gdb))]
             &rt_cfg,
@@ -509,14 +509,14 @@ pub(crate) mod tests {
         let guest_max_log_level = Some(log::LevelFilter::Error);
 
         #[cfg(gdb)]
-        let dbg_mem_access_fn = Arc::new(Mutex::new(mem_mgr.clone()));
+        let dbg_mem_access_fn = Arc::new(Mutex::new(hshm.clone()));
 
         // Test the initialise method
         vm.initialise(
             peb_addr,
             seed,
             page_size,
-            mem_mgr,
+            &mut hshm,
             host_funcs,
             guest_max_log_level,
             #[cfg(gdb)]

src/hyperlight_host/src/mem/layout.rs

@@ -30,8 +30,9 @@ use super::memory_region::{
 };
 #[cfg(feature = "init-paging")]
 use super::mgr::AMOUNT_OF_MEMORY_PER_PT;
-use super::shared_mem::{ExclusiveSharedMemory, GuestSharedMemory, SharedMemory};
+use super::shared_mem::{ExclusiveSharedMemory, SharedMemory};
 use crate::error::HyperlightError::{GuestOffsetIsInvalid, MemoryRequestTooBig};
+use crate::mem::shared_mem::HostSharedMemory;
 use crate::sandbox::SandboxConfiguration;
 use crate::{Result, new_error};
 
@@ -561,7 +562,7 @@ impl SandboxMemoryLayout {
 
     /// Returns the memory regions associated with this memory layout,
     /// suitable for passing to a hypervisor for mapping into memory
-    pub fn get_memory_regions(&self, shared_mem: &GuestSharedMemory) -> Result<Vec<MemoryRegion>> {
+    pub fn get_memory_regions(&self, shared_mem: &HostSharedMemory) -> Result<Vec<MemoryRegion>> {
         let mut builder = MemoryRegionVecBuilder::new(Self::BASE_ADDRESS, shared_mem.base_addr());
 
         cfg_if::cfg_if! {

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -96,8 +96,7 @@ pub struct MultiUseSandbox {
     id: u64,
     /// Whether this sandbox is poisoned
     poisoned: bool,
-    // We need to keep a reference to the host functions, even if the compiler marks it as unused. The compiler cannot detect our dynamic usages of the host function in `HyperlightFunction::call`.
-    pub(super) _host_funcs: Arc<Mutex<FunctionRegistry>>,
+    pub(super) host_funcs: Arc<Mutex<FunctionRegistry>>,
     pub(crate) mem_mgr: SandboxMemoryManager<HostSharedMemory>,
     vm: HyperlightVm,
     dispatch_ptr: RawPtr,
@@ -125,7 +124,7 @@ impl MultiUseSandbox {
         Self {
             id: SANDBOX_ID_COUNTER.fetch_add(1, Ordering::Relaxed),
             poisoned: false,
-            _host_funcs: host_funcs,
+            host_funcs,
             mem_mgr: mgr,
             vm,
             dispatch_ptr,
@@ -611,6 +610,8 @@ impl MultiUseSandbox {
             self.mem_mgr.write_guest_function_call(buffer)?;
 
             self.vm.dispatch_call_from_host(
+                &mut self.mem_mgr,
+                self.host_funcs.clone(),
                 self.dispatch_ptr.clone(),
                 #[cfg(gdb)]
                 self.dbg_mem_access_fn.clone(),

src/hyperlight_host/src/sandbox/mod.rs

@@ -174,7 +174,7 @@ mod tests {
                         .pop()
                         .unwrap_or_else(|| panic!("Failed to pop Sandbox thread {}", i));
                     let host_funcs = sandbox
-                        ._host_funcs
+                        .host_funcs
                         .try_lock()
                         .map_err(|_| new_error!("Error locking"));
 

src/hyperlight_host/src/sandbox/outb.rs

@@ -25,12 +25,14 @@ use tracing::{Span, instrument};
 use tracing_log::format_trace;
 
 use super::host_funcs::FunctionRegistry;
-#[cfg(feature = "mem_profile")]
-use crate::hypervisor::hyperlight_vm::HyperlightVm;
+use crate::hypervisor::regs::CommonRegisters;
 use crate::mem::mgr::SandboxMemoryManager;
 use crate::mem::shared_mem::HostSharedMemory;
 use crate::{HyperlightError, Result, new_error};
 
+#[cfg(feature = "mem_profile")]
+use crate::sandbox::trace::MemTraceInfo;
+
 #[instrument(err(Debug), skip_all, parent = Span::current(), level="Trace")]
 pub(super) fn outb_log(mgr: &mut SandboxMemoryManager<HostSharedMemory>) -> Result<()> {
     // This code will create either a logging record or a tracing record for the GuestLogData depending on if the host has set up a tracing subscriber.
@@ -150,7 +152,8 @@ pub(crate) fn handle_outb(
     host_funcs: Arc<Mutex<FunctionRegistry>>,
     port: u16,
     data: u32,
-    #[cfg(feature = "mem_profile")] vm: &mut HyperlightVm,
+    _regs: &CommonRegisters,
+    #[cfg(feature = "mem_profile")] trace_info: &mut MemTraceInfo,
 ) -> Result<()> {
     match port.try_into()? {
         OutBAction::Log => outb_log(mem_mgr),
@@ -185,17 +188,9 @@ pub(crate) fn handle_outb(
         #[cfg(feature = "trace_guest")]
         OutBAction::TraceBatch => Ok(()),
         #[cfg(feature = "mem_profile")]
-        OutBAction::TraceMemoryAlloc => {
-            let regs = vm.vm_regs()?;
-            let trace_info = vm.trace_info_mut();
-            trace_info.handle_trace_mem_alloc(&regs, mem_mgr)
-        }
+        OutBAction::TraceMemoryAlloc => trace_info.handle_trace_mem_alloc(_regs, mem_mgr),
         #[cfg(feature = "mem_profile")]
-        OutBAction::TraceMemoryFree => {
-            let regs = vm.vm_regs()?;
-            let trace_info = vm.trace_info_mut();
-            trace_info.handle_trace_mem_free(&regs, mem_mgr)
-        }
+        OutBAction::TraceMemoryFree => trace_info.handle_trace_mem_free(_regs, mem_mgr),
     }
 }
 #[cfg(test)]