Removed the OutBHandler and MemAccessHandler abstractions and related implementations.

src/hyperlight_host/src/hypervisor/handlers.rs

@@ -16,118 +16,11 @@ limitations under the License.
 
 use std::sync::{Arc, Mutex};
 
-use tracing::{Span, instrument};
-
-#[cfg(feature = "trace_guest")]
-use super::Hypervisor;
-use crate::{Result, new_error};
-
-/// The trait representing custom logic to handle the case when
-/// a Hypervisor's virtual CPU (vCPU) informs Hyperlight the guest
-/// has initiated an outb operation.
-pub(crate) trait OutBHandlerCaller: Sync + Send {
-    /// Function that gets called when an outb operation has occurred.
-    fn call(
-        &mut self,
-        #[cfg(feature = "trace_guest")] hv: &mut dyn Hypervisor,
-        port: u16,
-        payload: u32,
-    ) -> Result<()>;
-}
-
-/// A convenient type representing a common way `OutBHandler` implementations
-/// are passed as parameters to functions
-///
-/// Note: This needs to be wrapped in a Mutex to be able to grab a mutable
-/// reference to the underlying data (i.e., handle_outb in `Sandbox` takes
-/// a &mut self).
-pub(crate) type OutBHandlerWrapper = Arc<Mutex<dyn OutBHandlerCaller>>;
-
-#[cfg(feature = "trace_guest")]
-pub(crate) type OutBHandlerFunction =
-    Box<dyn FnMut(&mut dyn Hypervisor, u16, u32) -> Result<()> + Send>;
-#[cfg(not(feature = "trace_guest"))]
-pub(crate) type OutBHandlerFunction = Box<dyn FnMut(u16, u32) -> Result<()> + Send>;
-
-/// A `OutBHandler` implementation using a `OutBHandlerFunction`
-///
-/// Note: This handler must live no longer than the `Sandbox` to which it belongs
-pub(crate) struct OutBHandler(Arc<Mutex<OutBHandlerFunction>>);
-
-impl From<OutBHandlerFunction> for OutBHandler {
-    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
-    fn from(func: OutBHandlerFunction) -> Self {
-        Self(Arc::new(Mutex::new(func)))
-    }
-}
-
-impl OutBHandlerCaller for OutBHandler {
-    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-    fn call(
-        &mut self,
-        #[cfg(feature = "trace_guest")] hv: &mut dyn Hypervisor,
-        port: u16,
-        payload: u32,
-    ) -> Result<()> {
-        let mut func = self
-            .0
-            .try_lock()
-            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
-        func(
-            #[cfg(feature = "trace_guest")]
-            hv,
-            port,
-            payload,
-        )
-    }
-}
-
-/// The trait representing custom logic to handle the case when
-/// a Hypervisor's virtual CPU (vCPU) informs Hyperlight a memory access
-/// outside the designated address space has occurred.
-pub trait MemAccessHandlerCaller: Send {
-    /// Function that gets called when unexpected memory access has occurred.
-    fn call(&mut self) -> Result<()>;
-}
-
-/// A convenient type representing a common way `MemAccessHandler` implementations
-/// are passed as parameters to functions
-///
-/// Note: This needs to be wrapped in a Mutex to be able to grab a mutable
-/// reference to the underlying data (i.e., handle_mmio_exit in `Sandbox` takes
-/// a &mut self).
-pub type MemAccessHandlerWrapper = Arc<Mutex<dyn MemAccessHandlerCaller>>;
-
-pub(crate) type MemAccessHandlerFunction = Box<dyn FnMut() -> Result<()> + Send>;
-
-/// A `MemAccessHandler` implementation using `MemAccessHandlerFunction`.
-///
-/// Note: This handler must live for as long as its Sandbox or for
-/// static in the case of its C API usage.
-pub(crate) struct MemAccessHandler(Arc<Mutex<MemAccessHandlerFunction>>);
-
-impl From<MemAccessHandlerFunction> for MemAccessHandler {
-    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
-    fn from(func: MemAccessHandlerFunction) -> Self {
-        Self(Arc::new(Mutex::new(func)))
-    }
-}
-
-impl MemAccessHandlerCaller for MemAccessHandler {
-    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-    fn call(&mut self) -> Result<()> {
-        let mut func = self
-            .0
-            .try_lock()
-            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
-        func()
-    }
-}
+use crate::Result;
 
 /// The trait representing custom logic to handle the case when
 /// a Hypervisor's virtual CPU (vCPU) informs Hyperlight a debug memory access
 /// has been requested.
-#[cfg(gdb)]
 pub trait DbgMemAccessHandlerCaller: Send {
     /// Function that gets called when a read is requested.
     fn read(&mut self, addr: usize, data: &mut [u8]) -> Result<()>;
@@ -143,5 +36,4 @@ pub trait DbgMemAccessHandlerCaller: Send {
 ///
 /// Note: This needs to be wrapped in a Mutex to be able to grab a mutable
 /// reference to the underlying data
-#[cfg(gdb)]
 pub type DbgMemAccessHandlerWrapper = Arc<Mutex<dyn DbgMemAccessHandlerCaller>>;

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -25,8 +25,8 @@ extern crate mshv_bindings3 as mshv_bindings;
 extern crate mshv_ioctls3 as mshv_ioctls;
 
 use std::fmt::{Debug, Formatter};
-use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
 
 use log::{LevelFilter, error};
 #[cfg(mshv2)]
@@ -67,7 +67,6 @@ use super::gdb::{
 };
 #[cfg(gdb)]
 use super::handlers::DbgMemAccessHandlerWrapper;
-use super::handlers::{MemAccessHandlerWrapper, OutBHandlerWrapper};
 #[cfg(feature = "init-paging")]
 use super::{
     CR0_AM, CR0_ET, CR0_MP, CR0_NE, CR0_PE, CR0_PG, CR0_WP, CR4_OSFXSR, CR4_OSXMMEXCPT, CR4_PAE,
@@ -78,9 +77,13 @@ use super::{HyperlightExit, Hypervisor, InterruptHandle, LinuxInterruptHandle, V
 use crate::HyperlightError;
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
 use crate::mem::ptr::{GuestPtr, RawPtr};
+use crate::mem::shared_mem::HostSharedMemory;
 use crate::sandbox::SandboxConfiguration;
 #[cfg(feature = "trace_guest")]
 use crate::sandbox::TraceInfo;
+use crate::sandbox::host_funcs::FunctionRegistry;
+use crate::sandbox::mem_mgr::MemMgrWrapper;
+use crate::sandbox::outb::handle_outb;
 #[cfg(crashdump)]
 use crate::sandbox::uninitialized::SandboxRuntimeConfig;
 use crate::{Result, log_then_return, new_error};
@@ -313,7 +316,8 @@ pub(crate) struct HypervLinuxDriver {
     mem_regions: Vec<MemoryRegion>,
     orig_rsp: GuestPtr,
     interrupt_handle: Arc<LinuxInterruptHandle>,
-
+    mem_mgr: Option<MemMgrWrapper<HostSharedMemory>>,
+    host_funcs: Option<Arc<Mutex<FunctionRegistry>>>,
     #[cfg(gdb)]
     debug: Option<MshvDebug>,
     #[cfg(gdb)]
@@ -447,6 +451,8 @@ impl HypervLinuxDriver {
             entrypoint: entrypoint_ptr.absolute()?,
             orig_rsp: rsp_ptr,
             interrupt_handle: interrupt_handle.clone(),
+            mem_mgr: None,
+            host_funcs: None,
             #[cfg(gdb)]
             debug,
             #[cfg(gdb)]
@@ -574,11 +580,13 @@ impl Hypervisor for HypervLinuxDriver {
         peb_addr: RawPtr,
         seed: u64,
         page_size: u32,
-        outb_hdl: OutBHandlerWrapper,
-        mem_access_hdl: MemAccessHandlerWrapper,
+        mem_mgr: MemMgrWrapper<HostSharedMemory>,
+        host_funcs: Arc<Mutex<FunctionRegistry>>,
         max_guest_log_level: Option<LevelFilter>,
         #[cfg(gdb)] dbg_mem_access_fn: DbgMemAccessHandlerWrapper,
     ) -> 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 {
@@ -601,13 +609,22 @@ impl Hypervisor for HypervLinuxDriver {
         };
         self.vcpu_fd.set_regs(&regs)?;
 
-        VirtualCPU::run(
+        // 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(
             self.as_mut_hypervisor(),
-            outb_hdl,
-            mem_access_hdl,
+            &mem_mgr,
             #[cfg(gdb)]
             dbg_mem_access_fn,
-        )?;
+        );
+
+        // Put mem_mgr back
+        self.mem_mgr = Some(mem_mgr);
+        result?;
 
         Ok(())
     }
@@ -647,8 +664,7 @@ impl Hypervisor for HypervLinuxDriver {
     fn dispatch_call_from_host(
         &mut self,
         dispatch_func_addr: RawPtr,
-        outb_handle_fn: OutBHandlerWrapper,
-        mem_access_fn: MemAccessHandlerWrapper,
+        mem_mgr: &MemMgrWrapper<HostSharedMemory>,
         #[cfg(gdb)] dbg_mem_access_fn: DbgMemAccessHandlerWrapper,
     ) -> Result<()> {
         // Reset general purpose registers, then set RIP and RSP
@@ -672,8 +688,7 @@ impl Hypervisor for HypervLinuxDriver {
         // run
         VirtualCPU::run(
             self.as_mut_hypervisor(),
-            outb_handle_fn,
-            mem_access_fn,
+            mem_mgr,
             #[cfg(gdb)]
             dbg_mem_access_fn,
         )?;
@@ -688,22 +703,47 @@ impl Hypervisor for HypervLinuxDriver {
         data: Vec<u8>,
         rip: u64,
         instruction_length: u64,
-        outb_handle_fn: OutBHandlerWrapper,
     ) -> Result<()> {
         let mut padded = [0u8; 4];
         let copy_len = data.len().min(4);
         padded[..copy_len].copy_from_slice(&data[..copy_len]);
         let val = u32::from_le_bytes(padded);
 
-        outb_handle_fn
-            .try_lock()
-            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
-            .call(
-                #[cfg(feature = "trace_guest")]
-                self,
-                port,
-                val,
-            )?;
+        #[cfg(feature = "trace_guest")]
+        {
+            // 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 should be initialized before handling IO"))?;
+            let host_funcs = self
+                .host_funcs
+                .as_ref()
+                .ok_or_else(|| new_error!("host_funcs should be initialized before handling IO"))?
+                .clone();
+
+            handle_outb(&mut mem_mgr, host_funcs, self, port, val)?;
+
+            // Put the mem_mgr back
+            self.mem_mgr = Some(mem_mgr);
+        }
+
+        #[cfg(not(feature = "trace_guest"))]
+        {
+            let mem_mgr = self
+                .mem_mgr
+                .as_mut()
+                .ok_or_else(|| new_error!("mem_mgr should be initialized before handling IO"))?;
+            let host_funcs = self
+                .host_funcs
+                .as_ref()
+                .ok_or_else(|| new_error!("host_funcs should be initialized before handling IO"))?
+                .clone();
+
+            handle_outb(mem_mgr, host_funcs, port, val)?;
+        }
 
         // update rip
         self.vcpu_fd.set_reg(&[hv_register_assoc {