Remove dynamic dispatch from OutBHandlerFunction

Co-authored-by: simongdavies <1397489+simongdavies@users.noreply.github.com>

Author: Copilot

src/hyperlight_host/src/hypervisor/handlers.rs

@@ -18,34 +18,10 @@ use std::sync::{Arc, Mutex};
 
 use tracing::{instrument, Span};
 
+/// Re-export OutBHandler from the outb module where it naturally belongs
+pub(crate) use crate::sandbox::outb::OutBHandler;
 use crate::{new_error, Result};
 
-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 OutBHandler {
-    /// Function that gets called when an outb operation has occurred.
-    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-    pub fn call(&mut self, port: u16, payload: u32) -> Result<()> {
-        let mut func = self
-            .0
-            .try_lock()
-            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
-        func(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.

src/hyperlight_host/src/hypervisor/hyperv_windows.rs

@@ -505,15 +505,21 @@ pub mod tests {
 
     use crate::hypervisor::handlers::{MemAccessHandler, OutBHandler};
     use crate::hypervisor::tests::test_initialise;
+    use crate::sandbox::uninitialized::{GuestBinary, UninitializedSandbox};
     use crate::Result;
+    use hyperlight_testing::dummy_guest_as_string;
 
     #[test]
     #[serial]
     fn test_init() {
+        let filename = dummy_guest_as_string().expect("Guest Binary Missing");
+        let sandbox = UninitializedSandbox::new(GuestBinary::FilePath(filename), None).unwrap();
+        let (hshm, gshm) = sandbox.mgr.build();
+        drop(gshm);
+        let host_funcs = sandbox.host_funcs.clone();
+        
         let outb_handler = {
-            let func: Box<dyn FnMut(u16, u32) -> Result<()> + Send> =
-                Box::new(|_, _| -> Result<()> { Ok(()) });
-            Arc::new(Mutex::new(OutBHandler::from(func)))
+            crate::sandbox::outb::outb_handler_wrapper(hshm.clone(), host_funcs)
         };
         let mem_access_handler = {
             let func: Box<dyn FnMut() -> Result<()> + Send> = Box::new(|| -> Result<()> { Ok(()) });

src/hyperlight_host/src/hypervisor/kvm.rs

@@ -641,7 +641,9 @@ mod tests {
     use crate::hypervisor::handlers::DbgMemAccessHandlerCaller;
     use crate::hypervisor::handlers::{MemAccessHandler, OutBHandler};
     use crate::hypervisor::tests::test_initialise;
+    use crate::sandbox::uninitialized::{GuestBinary, UninitializedSandbox};
     use crate::Result;
+    use hyperlight_testing::dummy_guest_as_string;
 
     #[cfg(gdb)]
     struct DbgMemAccessHandler {}
@@ -667,10 +669,14 @@ mod tests {
             return;
         }
 
+        let filename = dummy_guest_as_string().expect("Guest Binary Missing");
+        let sandbox = UninitializedSandbox::new(GuestBinary::FilePath(filename), None).unwrap();
+        let (hshm, gshm) = sandbox.mgr.build();
+        drop(gshm);
+        let host_funcs = sandbox.host_funcs.clone();
+        
         let outb_handler: Arc<Mutex<OutBHandler>> = {
-            let func: Box<dyn FnMut(u16, u32) -> Result<()> + Send> =
-                Box::new(|_, _| -> Result<()> { Ok(()) });
-            Arc::new(Mutex::new(OutBHandler::from(func)))
+            crate::sandbox::outb::outb_handler_wrapper(hshm.clone(), host_funcs)
         };
         let mem_access_handler = {
             let func: Box<dyn FnMut() -> Result<()> + Send> = Box::new(|| -> Result<()> { Ok(()) });

src/hyperlight_host/src/sandbox/outb.rs

@@ -26,11 +26,43 @@ use tracing_log::format_trace;
 
 use super::host_funcs::FunctionRegistry;
 use super::mem_mgr::MemMgrWrapper;
-use crate::hypervisor::handlers::{OutBHandler, OutBHandlerFunction};
 use crate::mem::mgr::SandboxMemoryManager;
 use crate::mem::shared_mem::HostSharedMemory;
 use crate::{new_error, HyperlightError, Result};
 
+/// A `OutBHandler` implementation that contains the required data directly
+///
+/// Note: This handler must live no longer than the `Sandbox` to which it belongs
+pub(crate) struct OutBHandler {
+    mem_mgr_wrapper: MemMgrWrapper<HostSharedMemory>,
+    host_funcs_wrapper: Arc<Mutex<FunctionRegistry>>,
+}
+
+impl OutBHandler {
+    /// Create a new OutBHandler with the required dependencies
+    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
+    pub fn new(
+        mem_mgr_wrapper: MemMgrWrapper<HostSharedMemory>,
+        host_funcs_wrapper: Arc<Mutex<FunctionRegistry>>,
+    ) -> Self {
+        Self {
+            mem_mgr_wrapper,
+            host_funcs_wrapper,
+        }
+    }
+
+    /// Function that gets called when an outb operation has occurred.
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+    pub fn call(&mut self, port: u16, payload: u32) -> Result<()> {
+        handle_outb_impl(
+            &mut self.mem_mgr_wrapper,
+            self.host_funcs_wrapper.clone(),
+            port,
+            payload,
+        )
+    }
+}
+
 #[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.
@@ -146,7 +178,7 @@ fn outb_abort(mem_mgr: &mut MemMgrWrapper<HostSharedMemory>, data: u32) -> Resul
 
 /// Handles OutB operations from the guest.
 #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-fn handle_outb_impl(
+pub(crate) fn handle_outb_impl(
     mem_mgr: &mut MemMgrWrapper<HostSharedMemory>,
     host_funcs: Arc<Mutex<FunctionRegistry>>,
     port: u16,
@@ -189,18 +221,10 @@ fn handle_outb_impl(
 /// TODO: pass at least the `host_funcs_wrapper` param by reference.
 #[instrument(skip_all, parent = Span::current(), level= "Trace")]
 pub(crate) fn outb_handler_wrapper(
-    mut mem_mgr_wrapper: MemMgrWrapper<HostSharedMemory>,
+    mem_mgr_wrapper: MemMgrWrapper<HostSharedMemory>,
     host_funcs_wrapper: Arc<Mutex<FunctionRegistry>>,
 ) -> Arc<Mutex<OutBHandler>> {
-    let outb_func: OutBHandlerFunction = Box::new(move |port, payload| {
-        handle_outb_impl(
-            &mut mem_mgr_wrapper,
-            host_funcs_wrapper.clone(),
-            port,
-            payload,
-        )
-    });
-    let outb_hdl = OutBHandler::from(outb_func);
+    let outb_hdl = OutBHandler::new(mem_mgr_wrapper, host_funcs_wrapper);
     Arc::new(Mutex::new(outb_hdl))
 }