fix mem_mgr not initialized when tracing enabled

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

Author: dblnz

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -609,24 +609,11 @@ impl Hypervisor for HypervLinuxDriver {
         };
         self.vcpu_fd.set_regs(&regs)?;
 
-        // Extract mem_mgr to avoid borrowing conflicts
-        let mem_mgr = self
-            .mem_mgr
-            .take()
-            .ok_or_else(|| new_error!("mem_mgr should be initialized"))?;
-
-        let result = VirtualCPU::run(
+        VirtualCPU::run(
             self.as_mut_hypervisor(),
-            &mem_mgr,
             #[cfg(gdb)]
             dbg_mem_access_fn,
-        );
-
-        // Put mem_mgr back
-        self.mem_mgr = Some(mem_mgr);
-        result?;
-
-        Ok(())
+        )
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
@@ -664,7 +651,6 @@ impl Hypervisor for HypervLinuxDriver {
     fn dispatch_call_from_host(
         &mut self,
         dispatch_func_addr: RawPtr,
-        mem_mgr: &MemMgrWrapper<HostSharedMemory>,
         #[cfg(gdb)] dbg_mem_access_fn: DbgMemAccessHandlerWrapper,
     ) -> Result<()> {
         // Reset general purpose registers, then set RIP and RSP
@@ -688,7 +674,6 @@ impl Hypervisor for HypervLinuxDriver {
         // run
         VirtualCPU::run(
             self.as_mut_hypervisor(),
-            mem_mgr,
             #[cfg(gdb)]
             dbg_mem_access_fn,
         )?;
@@ -1162,6 +1147,14 @@ impl Hypervisor for HypervLinuxDriver {
         Ok(())
     }
 
+    fn check_stack_guard(&self) -> Result<bool> {
+        if let Some(mgr) = self.mem_mgr.as_ref() {
+            mgr.check_stack_guard()
+        } else {
+            Err(new_error!("Memory manager is not initialized"))
+        }
+    }
+
     #[cfg(feature = "trace_guest")]
     fn read_trace_reg(&self, reg: TraceRegister) -> Result<u64> {
         let mut assoc = [hv_register_assoc {

src/hyperlight_host/src/hypervisor/hyperv_windows.rs

@@ -623,24 +623,11 @@ impl Hypervisor for HypervWindowsDriver {
         };
         self.processor.set_general_purpose_registers(&regs)?;
 
-        // Extract mem_mgr to avoid borrowing conflicts
-        let mem_mgr = self
-            .mem_mgr
-            .take()
-            .ok_or_else(|| new_error!("mem_mgr should be initialized"))?;
-
-        let result = VirtualCPU::run(
+        VirtualCPU::run(
             self.as_mut_hypervisor(),
-            &mem_mgr,
             #[cfg(gdb)]
             dbg_mem_access_hdl,
-        );
-
-        // Put mem_mgr back
-        self.mem_mgr = Some(mem_mgr);
-        result?;
-
-        Ok(())
+        )
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
@@ -662,7 +649,6 @@ impl Hypervisor for HypervWindowsDriver {
     fn dispatch_call_from_host(
         &mut self,
         dispatch_func_addr: RawPtr,
-        mem_mgr: &MemMgrWrapper<HostSharedMemory>,
         #[cfg(gdb)] dbg_mem_access_hdl: DbgMemAccessHandlerWrapper,
     ) -> Result<()> {
         // Reset general purpose registers, then set RIP and RSP
@@ -684,7 +670,6 @@ impl Hypervisor for HypervWindowsDriver {
 
         VirtualCPU::run(
             self.as_mut_hypervisor(),
-            mem_mgr,
             #[cfg(gdb)]
             dbg_mem_access_hdl,
         )?;
@@ -1094,6 +1079,14 @@ impl Hypervisor for HypervWindowsDriver {
         Ok(())
     }
 
+    fn check_stack_guard(&self) -> Result<bool> {
+        if let Some(mgr) = self.mem_mgr.as_ref() {
+            mgr.check_stack_guard()
+        } else {
+            Err(new_error!("Memory manager is not initialized"))
+        }
+    }
+
     #[cfg(feature = "trace_guest")]
     fn read_trace_reg(&self, reg: TraceRegister) -> Result<u64> {
         let regs = self.processor.get_regs()?;

src/hyperlight_host/src/hypervisor/kvm.rs

@@ -492,24 +492,11 @@ impl Hypervisor for KVMDriver {
         };
         self.vcpu_fd.set_regs(&regs)?;
 
-        // Extract mem_mgr to avoid borrowing conflicts
-        let mem_mgr = self
-            .mem_mgr
-            .take()
-            .ok_or_else(|| new_error!("mem_mgr should be initialized"))?;
-
-        let result = VirtualCPU::run(
+        VirtualCPU::run(
             self.as_mut_hypervisor(),
-            &mem_mgr,
             #[cfg(gdb)]
             dbg_mem_access_fn,
-        );
-
-        // Put mem_mgr back
-        self.mem_mgr = Some(mem_mgr);
-        result?;
-
-        Ok(())
+        )
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
@@ -555,7 +542,6 @@ impl Hypervisor for KVMDriver {
     fn dispatch_call_from_host(
         &mut self,
         dispatch_func_addr: RawPtr,
-        mem_mgr: &MemMgrWrapper<HostSharedMemory>,
         #[cfg(gdb)] dbg_mem_access_fn: DbgMemAccessHandlerWrapper,
     ) -> Result<()> {
         // Reset general purpose registers, then set RIP and RSP
@@ -578,7 +564,6 @@ impl Hypervisor for KVMDriver {
         // run
         VirtualCPU::run(
             self.as_mut_hypervisor(),
-            mem_mgr,
             #[cfg(gdb)]
             dbg_mem_access_fn,
         )?;
@@ -1012,6 +997,14 @@ impl Hypervisor for KVMDriver {
         Ok(())
     }
 
+    fn check_stack_guard(&self) -> Result<bool> {
+        if let Some(mgr) = self.mem_mgr.as_ref() {
+            mgr.check_stack_guard()
+        } else {
+            Err(new_error!("Memory manager is not initialized"))
+        }
+    }
+
     #[cfg(feature = "trace_guest")]
     fn read_trace_reg(&self, reg: TraceRegister) -> Result<u64> {
         let regs = self.vcpu_fd.get_regs()?;

src/hyperlight_host/src/hypervisor/mod.rs

@@ -170,7 +170,6 @@ pub(crate) trait Hypervisor: Debug + Send {
     fn dispatch_call_from_host(
         &mut self,
         dispatch_func_addr: RawPtr,
-        mem_mgr: &MemMgrWrapper<HostSharedMemory>,
         #[cfg(gdb)] dbg_mem_access_fn: DbgMemAccessHandlerWrapper,
     ) -> Result<()>;
 
@@ -269,6 +268,9 @@ pub(crate) trait Hypervisor: Debug + Send {
         unimplemented!()
     }
 
+    /// Check stack guard to see if the stack is still valid
+    fn check_stack_guard(&self) -> Result<bool>;
+
     /// Read a register for trace/unwind purposes
     #[cfg(feature = "trace_guest")]
     fn read_trace_reg(&self, reg: TraceRegister) -> Result<u64>;
@@ -289,7 +291,6 @@ impl VirtualCPU {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     pub(crate) fn run(
         hv: &mut dyn Hypervisor,
-        mem_mgr: &MemMgrWrapper<HostSharedMemory>,
         #[cfg(gdb)] dbg_mem_access_fn: Arc<Mutex<dyn DbgMemAccessHandlerCaller>>,
     ) -> Result<()> {
         loop {
@@ -311,7 +312,7 @@ impl VirtualCPU {
                     #[cfg(crashdump)]
                     crashdump::generate_crashdump(hv)?;
 
-                    handle_mem_access(mem_mgr)?;
+                    handle_mem_access(hv)?;
 
                     log_then_return!("MMIO access address {:#x}", addr);
                 }

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -230,7 +230,6 @@ impl MultiUseSandbox {
 
             self.vm.dispatch_call_from_host(
                 self.dispatch_ptr.clone(),
-                &self.mem_mgr,
                 #[cfg(gdb)]
                 self.dbg_mem_access_fn.clone(),
             )?;

src/hyperlight_host/src/sandbox/mem_access.rs

@@ -19,16 +19,19 @@ use std::sync::{Arc, Mutex};
 
 use tracing::{Span, instrument};
 
+#[cfg(gdb)]
 use super::mem_mgr::MemMgrWrapper;
 use crate::error::HyperlightError::StackOverflow;
+use crate::hypervisor::Hypervisor;
 #[cfg(gdb)]
 use crate::hypervisor::handlers::{DbgMemAccessHandlerCaller, DbgMemAccessHandlerWrapper};
+#[cfg(gdb)]
 use crate::mem::shared_mem::HostSharedMemory;
 use crate::{Result, log_then_return};
 
 #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-pub(crate) fn handle_mem_access(wrapper: &MemMgrWrapper<HostSharedMemory>) -> Result<()> {
-    if !wrapper.check_stack_guard()? {
+pub(crate) fn handle_mem_access(hv: &dyn Hypervisor) -> Result<()> {
+    if !hv.check_stack_guard()? {
         log_then_return!(StackOverflow());
     }