From d9c2455db6a94392c4bf05b6c8bcbedcc443e2d8 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 21 May 2025 08:20:19 +0000
Subject: [PATCH 1/8] Initial plan for issue


From 132d0f45f61181afc9ce81d19173cf38df2c7c5d Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 21 May 2025 15:05:13 +0000
Subject: [PATCH 2/8] Remove OutBHandlerWrapper type

Co-authored-by: simongdavies <1397489+simongdavies@users.noreply.github.com>
---
 src/hyperlight_host/src/hypervisor/handlers.rs   |  8 --------
 .../src/hypervisor/hyperv_linux.rs               |  9 +++++----
 .../src/hypervisor/hyperv_windows.rs             |  9 +++++----
 .../src/hypervisor/hypervisor_handler.rs         |  4 ++--
 src/hyperlight_host/src/hypervisor/kvm.rs        |  9 ++++-----
 src/hyperlight_host/src/hypervisor/mod.rs        | 16 +++++++---------
 src/hyperlight_host/src/sandbox/outb.rs          |  6 +++---
 7 files changed, 26 insertions(+), 35 deletions(-)

diff --git a/src/hyperlight_host/src/hypervisor/handlers.rs b/src/hyperlight_host/src/hypervisor/handlers.rs
index 9dbb948af..ac8106869 100644
--- a/src/hyperlight_host/src/hypervisor/handlers.rs
+++ b/src/hyperlight_host/src/hypervisor/handlers.rs
@@ -28,14 +28,6 @@ pub trait OutBHandlerCaller: Sync + Send {
     fn call(&mut self, 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 type OutBHandlerWrapper = Arc<Mutex<dyn OutBHandlerCaller>>;
-
 pub(crate) type OutBHandlerFunction = Box<dyn FnMut(u16, u32) -> Result<()> + Send>;
 
 /// A `OutBHandler` implementation using a `OutBHandlerFunction`
diff --git a/src/hyperlight_host/src/hypervisor/hyperv_linux.rs b/src/hyperlight_host/src/hypervisor/hyperv_linux.rs
index 85dc514b5..506753f49 100644
--- a/src/hyperlight_host/src/hypervisor/hyperv_linux.rs
+++ b/src/hyperlight_host/src/hypervisor/hyperv_linux.rs
@@ -25,6 +25,7 @@ extern crate mshv_bindings3 as mshv_bindings;
 extern crate mshv_ioctls3 as mshv_ioctls;
 
 use std::fmt::{Debug, Formatter};
+use std::sync::{Arc, Mutex};
 
 use log::{error, LevelFilter};
 #[cfg(mshv2)]
@@ -54,7 +55,7 @@ use super::fpu::{FP_CONTROL_WORD_DEFAULT, FP_TAG_WORD_DEFAULT, MXCSR_DEFAULT};
 use super::gdb::{DebugCommChannel, DebugMsg, DebugResponse, GuestDebug, MshvDebug};
 #[cfg(gdb)]
 use super::handlers::DbgMemAccessHandlerWrapper;
-use super::handlers::{MemAccessHandlerWrapper, OutBHandlerWrapper};
+use super::handlers::{MemAccessHandlerWrapper, OutBHandler, OutBHandlerCaller};
 use super::{
     Hypervisor, VirtualCPU, CR0_AM, CR0_ET, CR0_MP, CR0_NE, CR0_PE, CR0_PG, CR0_WP, CR4_OSFXSR,
     CR4_OSXMMEXCPT, CR4_PAE, EFER_LMA, EFER_LME, EFER_NX, EFER_SCE,
@@ -459,7 +460,7 @@ impl Hypervisor for HypervLinuxDriver {
         peb_addr: RawPtr,
         seed: u64,
         page_size: u32,
-        outb_hdl: OutBHandlerWrapper,
+        outb_hdl: Arc<Mutex<OutBHandler>>,
         mem_access_hdl: MemAccessHandlerWrapper,
         hv_handler: Option<HypervisorHandler>,
         max_guest_log_level: Option<LevelFilter>,
@@ -501,7 +502,7 @@ impl Hypervisor for HypervLinuxDriver {
     fn dispatch_call_from_host(
         &mut self,
         dispatch_func_addr: RawPtr,
-        outb_handle_fn: OutBHandlerWrapper,
+        outb_handle_fn: Arc<Mutex<OutBHandler>>,
         mem_access_fn: MemAccessHandlerWrapper,
         hv_handler: Option<HypervisorHandler>,
         #[cfg(gdb)] dbg_mem_access_fn: DbgMemAccessHandlerWrapper,
@@ -544,7 +545,7 @@ impl Hypervisor for HypervLinuxDriver {
         data: Vec<u8>,
         rip: u64,
         instruction_length: u64,
-        outb_handle_fn: OutBHandlerWrapper,
+        outb_handle_fn: Arc<Mutex<OutBHandler>>,
     ) -> Result<()> {
         let mut padded = [0u8; 4];
         let copy_len = data.len().min(4);
diff --git a/src/hyperlight_host/src/hypervisor/hyperv_windows.rs b/src/hyperlight_host/src/hypervisor/hyperv_windows.rs
index abc2b77cf..799ef115c 100644
--- a/src/hyperlight_host/src/hypervisor/hyperv_windows.rs
+++ b/src/hyperlight_host/src/hypervisor/hyperv_windows.rs
@@ -18,6 +18,7 @@ use core::ffi::c_void;
 use std::fmt;
 use std::fmt::{Debug, Formatter};
 use std::string::String;
+use std::sync::{Arc, Mutex};
 
 use hyperlight_common::mem::PAGE_SIZE_USIZE;
 use log::LevelFilter;
@@ -31,7 +32,7 @@ use windows::Win32::System::Hypervisor::{
 use super::fpu::{FP_TAG_WORD_DEFAULT, MXCSR_DEFAULT};
 #[cfg(gdb)]
 use super::handlers::DbgMemAccessHandlerWrapper;
-use super::handlers::{MemAccessHandlerWrapper, OutBHandlerWrapper};
+use super::handlers::{MemAccessHandlerWrapper, OutBHandler, OutBHandlerCaller};
 use super::surrogate_process::SurrogateProcess;
 use super::surrogate_process_manager::*;
 use super::windows_hypervisor_platform::{VMPartition, VMProcessor};
@@ -305,7 +306,7 @@ impl Hypervisor for HypervWindowsDriver {
         peb_address: RawPtr,
         seed: u64,
         page_size: u32,
-        outb_hdl: OutBHandlerWrapper,
+        outb_hdl: Arc<Mutex<OutBHandler>>,
         mem_access_hdl: MemAccessHandlerWrapper,
         hv_handler: Option<HypervisorHandler>,
         max_guest_log_level: Option<LevelFilter>,
@@ -347,7 +348,7 @@ impl Hypervisor for HypervWindowsDriver {
     fn dispatch_call_from_host(
         &mut self,
         dispatch_func_addr: RawPtr,
-        outb_hdl: OutBHandlerWrapper,
+        outb_hdl: Arc<Mutex<OutBHandler>>,
         mem_access_hdl: MemAccessHandlerWrapper,
         hv_handler: Option<HypervisorHandler>,
         #[cfg(gdb)] dbg_mem_access_hdl: DbgMemAccessHandlerWrapper,
@@ -388,7 +389,7 @@ impl Hypervisor for HypervWindowsDriver {
         data: Vec<u8>,
         rip: u64,
         instruction_length: u64,
-        outb_handle_fn: OutBHandlerWrapper,
+        outb_handle_fn: Arc<Mutex<OutBHandler>>,
     ) -> Result<()> {
         let mut padded = [0u8; 4];
         let copy_len = data.len().min(4);
diff --git a/src/hyperlight_host/src/hypervisor/hypervisor_handler.rs b/src/hyperlight_host/src/hypervisor/hypervisor_handler.rs
index 8e351708c..7f318e46d 100644
--- a/src/hyperlight_host/src/hypervisor/hypervisor_handler.rs
+++ b/src/hyperlight_host/src/hypervisor/hypervisor_handler.rs
@@ -39,7 +39,7 @@ use windows::Win32::System::Hypervisor::{WHvCancelRunVirtualProcessor, WHV_PARTI
 use super::gdb::create_gdb_thread;
 #[cfg(gdb)]
 use crate::hypervisor::handlers::DbgMemAccessHandlerWrapper;
-use crate::hypervisor::handlers::{MemAccessHandlerWrapper, OutBHandlerWrapper};
+use crate::hypervisor::handlers::{MemAccessHandlerWrapper, OutBHandler};
 #[cfg(target_os = "windows")]
 use crate::hypervisor::wrappers::HandleWrapper;
 use crate::hypervisor::Hypervisor;
@@ -184,7 +184,7 @@ pub(crate) struct HvHandlerConfig {
     pub(crate) dispatch_function_addr: Arc<Mutex<Option<RawPtr>>>,
     pub(crate) max_init_time: Duration,
     pub(crate) max_exec_time: Duration,
-    pub(crate) outb_handler: OutBHandlerWrapper,
+    pub(crate) outb_handler: Arc<Mutex<OutBHandler>>,
     pub(crate) mem_access_handler: MemAccessHandlerWrapper,
     pub(crate) max_wait_for_cancellation: Duration,
     pub(crate) max_guest_log_level: Option<LevelFilter>,
diff --git a/src/hyperlight_host/src/hypervisor/kvm.rs b/src/hyperlight_host/src/hypervisor/kvm.rs
index 3dd1cb1fc..469de6177 100644
--- a/src/hyperlight_host/src/hypervisor/kvm.rs
+++ b/src/hyperlight_host/src/hypervisor/kvm.rs
@@ -16,7 +16,6 @@ limitations under the License.
 
 use std::convert::TryFrom;
 use std::fmt::Debug;
-#[cfg(gdb)]
 use std::sync::{Arc, Mutex};
 
 use kvm_bindings::{kvm_fpu, kvm_regs, kvm_userspace_memory_region, KVM_MEM_READONLY};
@@ -30,7 +29,7 @@ use super::fpu::{FP_CONTROL_WORD_DEFAULT, FP_TAG_WORD_DEFAULT, MXCSR_DEFAULT};
 use super::gdb::{DebugCommChannel, DebugMsg, DebugResponse, GuestDebug, KvmDebug, VcpuStopReason};
 #[cfg(gdb)]
 use super::handlers::DbgMemAccessHandlerWrapper;
-use super::handlers::{MemAccessHandlerWrapper, OutBHandlerWrapper};
+use super::handlers::{MemAccessHandlerWrapper, OutBHandler, OutBHandlerCaller};
 use super::{
     HyperlightExit, Hypervisor, VirtualCPU, CR0_AM, CR0_ET, CR0_MP, CR0_NE, CR0_PE, CR0_PG, CR0_WP,
     CR4_OSFXSR, CR4_OSXMMEXCPT, CR4_PAE, EFER_LMA, EFER_LME, EFER_NX, EFER_SCE,
@@ -404,7 +403,7 @@ impl Hypervisor for KVMDriver {
         peb_addr: RawPtr,
         seed: u64,
         page_size: u32,
-        outb_hdl: OutBHandlerWrapper,
+        outb_hdl: Arc<Mutex<OutBHandler>>,
         mem_access_hdl: MemAccessHandlerWrapper,
         hv_handler: Option<HypervisorHandler>,
         max_guest_log_level: Option<LevelFilter>,
@@ -445,7 +444,7 @@ impl Hypervisor for KVMDriver {
     fn dispatch_call_from_host(
         &mut self,
         dispatch_func_addr: RawPtr,
-        outb_handle_fn: OutBHandlerWrapper,
+        outb_handle_fn: Arc<Mutex<OutBHandler>>,
         mem_access_fn: MemAccessHandlerWrapper,
         hv_handler: Option<HypervisorHandler>,
         #[cfg(gdb)] dbg_mem_access_fn: DbgMemAccessHandlerWrapper,
@@ -487,7 +486,7 @@ impl Hypervisor for KVMDriver {
         data: Vec<u8>,
         _rip: u64,
         _instruction_length: u64,
-        outb_handle_fn: OutBHandlerWrapper,
+        outb_handle_fn: Arc<Mutex<OutBHandler>>,
     ) -> Result<()> {
         // KVM does not need RIP or instruction length, as it automatically sets the RIP
 
diff --git a/src/hyperlight_host/src/hypervisor/mod.rs b/src/hyperlight_host/src/hypervisor/mod.rs
index 62cebe829..e37969ef9 100644
--- a/src/hyperlight_host/src/hypervisor/mod.rs
+++ b/src/hyperlight_host/src/hypervisor/mod.rs
@@ -67,9 +67,7 @@ use gdb::VcpuStopReason;
 
 #[cfg(gdb)]
 use self::handlers::{DbgMemAccessHandlerCaller, DbgMemAccessHandlerWrapper};
-use self::handlers::{
-    MemAccessHandlerCaller, MemAccessHandlerWrapper, OutBHandlerCaller, OutBHandlerWrapper,
-};
+use self::handlers::{MemAccessHandlerCaller, MemAccessHandlerWrapper, OutBHandler};
 use crate::hypervisor::hypervisor_handler::HypervisorHandler;
 use crate::mem::ptr::RawPtr;
 
@@ -124,7 +122,7 @@ pub(crate) trait Hypervisor: Debug + Sync + Send {
         peb_addr: RawPtr,
         seed: u64,
         page_size: u32,
-        outb_handle_fn: OutBHandlerWrapper,
+        outb_handle_fn: Arc<Mutex<OutBHandler>>,
         mem_access_fn: MemAccessHandlerWrapper,
         hv_handler: Option<HypervisorHandler>,
         guest_max_log_level: Option<LevelFilter>,
@@ -141,7 +139,7 @@ pub(crate) trait Hypervisor: Debug + Sync + Send {
     fn dispatch_call_from_host(
         &mut self,
         dispatch_func_addr: RawPtr,
-        outb_handle_fn: OutBHandlerWrapper,
+        outb_handle_fn: Arc<Mutex<OutBHandler>>,
         mem_access_fn: MemAccessHandlerWrapper,
         hv_handler: Option<HypervisorHandler>,
         #[cfg(gdb)] dbg_mem_access_fn: DbgMemAccessHandlerWrapper,
@@ -154,7 +152,7 @@ pub(crate) trait Hypervisor: Debug + Sync + Send {
         data: Vec<u8>,
         rip: u64,
         instruction_length: u64,
-        outb_handle_fn: OutBHandlerWrapper,
+        outb_handle_fn: Arc<Mutex<OutBHandler>>,
     ) -> Result<()>;
 
     /// Run the vCPU
@@ -254,7 +252,7 @@ impl VirtualCPU {
     pub fn run(
         hv: &mut dyn Hypervisor,
         hv_handler: Option<HypervisorHandler>,
-        outb_handle_fn: Arc<Mutex<dyn OutBHandlerCaller>>,
+        outb_handle_fn: Arc<Mutex<OutBHandler>>,
         mem_access_fn: Arc<Mutex<dyn MemAccessHandlerCaller>>,
         #[cfg(gdb)] dbg_mem_access_fn: Arc<Mutex<dyn DbgMemAccessHandlerCaller>>,
     ) -> Result<()> {
@@ -341,7 +339,7 @@ pub(crate) mod tests {
 
     #[cfg(gdb)]
     use super::handlers::DbgMemAccessHandlerWrapper;
-    use super::handlers::{MemAccessHandlerWrapper, OutBHandlerWrapper};
+    use super::handlers::{MemAccessHandlerWrapper, OutBHandler};
     use crate::hypervisor::hypervisor_handler::{
         HvHandlerConfig, HypervisorHandler, HypervisorHandlerAction,
     };
@@ -351,7 +349,7 @@ pub(crate) mod tests {
     use crate::{new_error, Result};
 
     pub(crate) fn test_initialise(
-        outb_hdl: OutBHandlerWrapper,
+        outb_hdl: Arc<Mutex<OutBHandler>>,
         mem_access_hdl: MemAccessHandlerWrapper,
         #[cfg(gdb)] dbg_mem_access_fn: DbgMemAccessHandlerWrapper,
     ) -> Result<()> {
diff --git a/src/hyperlight_host/src/sandbox/outb.rs b/src/hyperlight_host/src/sandbox/outb.rs
index 4dc91207f..6f04a75e5 100644
--- a/src/hyperlight_host/src/sandbox/outb.rs
+++ b/src/hyperlight_host/src/sandbox/outb.rs
@@ -26,7 +26,7 @@ use tracing_log::format_trace;
 
 use super::host_funcs::FunctionRegistry;
 use super::mem_mgr::MemMgrWrapper;
-use crate::hypervisor::handlers::{OutBHandler, OutBHandlerFunction, OutBHandlerWrapper};
+use crate::hypervisor::handlers::{OutBHandler, OutBHandlerFunction};
 use crate::mem::mgr::SandboxMemoryManager;
 use crate::mem::shared_mem::HostSharedMemory;
 use crate::{new_error, HyperlightError, Result};
@@ -184,14 +184,14 @@ fn handle_outb_impl(
 }
 
 /// Given a `MemMgrWrapper` and ` HostFuncsWrapper` -- both passed by _value_
-///  -- return an `OutBHandlerWrapper` wrapping the core OUTB handler logic.
+///  -- return an `Arc<Mutex<OutBHandler>>` wrapping the core OUTB handler logic.
 ///
 /// 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>,
     host_funcs_wrapper: Arc<Mutex<FunctionRegistry>>,
-) -> OutBHandlerWrapper {
+) -> Arc<Mutex<OutBHandler>> {
     let outb_func: OutBHandlerFunction = Box::new(move |port, payload| {
         handle_outb_impl(
             &mut mem_mgr_wrapper,

From 4c61fc467c7f4f6e6b540d422df709a8a6f7f153 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 25 Jun 2025 18:30:05 +0000
Subject: [PATCH 3/8] Remove dyn OutBHandlerCaller trait entirely

Co-authored-by: simongdavies <1397489+simongdavies@users.noreply.github.com>
---
 src/hyperlight_host/src/hypervisor/handlers.rs      | 13 +++----------
 src/hyperlight_host/src/hypervisor/hyperv_linux.rs  |  2 +-
 .../src/hypervisor/hyperv_windows.rs                |  2 +-
 src/hyperlight_host/src/hypervisor/kvm.rs           |  2 +-
 4 files changed, 6 insertions(+), 13 deletions(-)

diff --git a/src/hyperlight_host/src/hypervisor/handlers.rs b/src/hyperlight_host/src/hypervisor/handlers.rs
index ac8106869..58872139f 100644
--- a/src/hyperlight_host/src/hypervisor/handlers.rs
+++ b/src/hyperlight_host/src/hypervisor/handlers.rs
@@ -20,14 +20,6 @@ use tracing::{instrument, Span};
 
 use crate::{new_error, Result};
 
-/// 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 trait OutBHandlerCaller: Sync + Send {
-    /// Function that gets called when an outb operation has occurred.
-    fn call(&mut self, port: u16, payload: u32) -> Result<()>;
-}
-
 pub(crate) type OutBHandlerFunction = Box<dyn FnMut(u16, u32) -> Result<()> + Send>;
 
 /// A `OutBHandler` implementation using a `OutBHandlerFunction`
@@ -42,9 +34,10 @@ impl From<OutBHandlerFunction> for OutBHandler {
     }
 }
 
-impl OutBHandlerCaller for OutBHandler {
+impl OutBHandler {
+    /// Function that gets called when an outb operation has occurred.
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-    fn call(&mut self, port: u16, payload: u32) -> Result<()> {
+    pub fn call(&mut self, port: u16, payload: u32) -> Result<()> {
         let mut func = self
             .0
             .try_lock()
diff --git a/src/hyperlight_host/src/hypervisor/hyperv_linux.rs b/src/hyperlight_host/src/hypervisor/hyperv_linux.rs
index 506753f49..dc71aff64 100644
--- a/src/hyperlight_host/src/hypervisor/hyperv_linux.rs
+++ b/src/hyperlight_host/src/hypervisor/hyperv_linux.rs
@@ -55,7 +55,7 @@ use super::fpu::{FP_CONTROL_WORD_DEFAULT, FP_TAG_WORD_DEFAULT, MXCSR_DEFAULT};
 use super::gdb::{DebugCommChannel, DebugMsg, DebugResponse, GuestDebug, MshvDebug};
 #[cfg(gdb)]
 use super::handlers::DbgMemAccessHandlerWrapper;
-use super::handlers::{MemAccessHandlerWrapper, OutBHandler, OutBHandlerCaller};
+use super::handlers::{MemAccessHandlerWrapper, OutBHandler};
 use super::{
     Hypervisor, VirtualCPU, CR0_AM, CR0_ET, CR0_MP, CR0_NE, CR0_PE, CR0_PG, CR0_WP, CR4_OSFXSR,
     CR4_OSXMMEXCPT, CR4_PAE, EFER_LMA, EFER_LME, EFER_NX, EFER_SCE,
diff --git a/src/hyperlight_host/src/hypervisor/hyperv_windows.rs b/src/hyperlight_host/src/hypervisor/hyperv_windows.rs
index 799ef115c..186606a39 100644
--- a/src/hyperlight_host/src/hypervisor/hyperv_windows.rs
+++ b/src/hyperlight_host/src/hypervisor/hyperv_windows.rs
@@ -32,7 +32,7 @@ use windows::Win32::System::Hypervisor::{
 use super::fpu::{FP_TAG_WORD_DEFAULT, MXCSR_DEFAULT};
 #[cfg(gdb)]
 use super::handlers::DbgMemAccessHandlerWrapper;
-use super::handlers::{MemAccessHandlerWrapper, OutBHandler, OutBHandlerCaller};
+use super::handlers::{MemAccessHandlerWrapper, OutBHandler};
 use super::surrogate_process::SurrogateProcess;
 use super::surrogate_process_manager::*;
 use super::windows_hypervisor_platform::{VMPartition, VMProcessor};
diff --git a/src/hyperlight_host/src/hypervisor/kvm.rs b/src/hyperlight_host/src/hypervisor/kvm.rs
index 469de6177..d32e9d3f6 100644
--- a/src/hyperlight_host/src/hypervisor/kvm.rs
+++ b/src/hyperlight_host/src/hypervisor/kvm.rs
@@ -29,7 +29,7 @@ use super::fpu::{FP_CONTROL_WORD_DEFAULT, FP_TAG_WORD_DEFAULT, MXCSR_DEFAULT};
 use super::gdb::{DebugCommChannel, DebugMsg, DebugResponse, GuestDebug, KvmDebug, VcpuStopReason};
 #[cfg(gdb)]
 use super::handlers::DbgMemAccessHandlerWrapper;
-use super::handlers::{MemAccessHandlerWrapper, OutBHandler, OutBHandlerCaller};
+use super::handlers::{MemAccessHandlerWrapper, OutBHandler};
 use super::{
     HyperlightExit, Hypervisor, VirtualCPU, CR0_AM, CR0_ET, CR0_MP, CR0_NE, CR0_PE, CR0_PG, CR0_WP,
     CR4_OSFXSR, CR4_OSXMMEXCPT, CR4_PAE, EFER_LMA, EFER_LME, EFER_NX, EFER_SCE,

From 8d15b69ae8233668634a2f9e65bc3a98273e5788 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 26 Jun 2025 20:20:50 +0000
Subject: [PATCH 4/8] Remove dynamic dispatch from OutBHandlerFunction

Co-authored-by: simongdavies <1397489+simongdavies@users.noreply.github.com>
---
 .../src/hypervisor/handlers.rs                | 28 +----------
 .../src/hypervisor/hyperv_windows.rs          | 12 +++--
 src/hyperlight_host/src/hypervisor/kvm.rs     | 12 +++--
 src/hyperlight_host/src/sandbox/outb.rs       | 48 ++++++++++++++-----
 4 files changed, 56 insertions(+), 44 deletions(-)

diff --git a/src/hyperlight_host/src/hypervisor/handlers.rs b/src/hyperlight_host/src/hypervisor/handlers.rs
index 58872139f..95f225c40 100644
--- a/src/hyperlight_host/src/hypervisor/handlers.rs
+++ b/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.
diff --git a/src/hyperlight_host/src/hypervisor/hyperv_windows.rs b/src/hyperlight_host/src/hypervisor/hyperv_windows.rs
index 186606a39..34bef6779 100644
--- a/src/hyperlight_host/src/hypervisor/hyperv_windows.rs
+++ b/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(()) });
diff --git a/src/hyperlight_host/src/hypervisor/kvm.rs b/src/hyperlight_host/src/hypervisor/kvm.rs
index d32e9d3f6..96388a7db 100644
--- a/src/hyperlight_host/src/hypervisor/kvm.rs
+++ b/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(()) });
diff --git a/src/hyperlight_host/src/sandbox/outb.rs b/src/hyperlight_host/src/sandbox/outb.rs
index 6f04a75e5..5620e6f06 100644
--- a/src/hyperlight_host/src/sandbox/outb.rs
+++ b/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))
 }
 

From 7e85790995f06bb8cfef659d89a499c05fa0cffb Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 3 Jul 2025 16:43:20 +0000
Subject: [PATCH 5/8] Fix formatting in hypervisor test files

Co-authored-by: simongdavies <1397489+simongdavies@users.noreply.github.com>
Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>

Co-authored-by: jsturtevant <648372+jsturtevant@users.noreply.github.com>
---
 src/hyperlight_host/src/hypervisor/hyperv_windows.rs |  8 +++-----
 src/hyperlight_host/src/hypervisor/kvm.rs            | 10 +++++-----
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/src/hyperlight_host/src/hypervisor/hyperv_windows.rs b/src/hyperlight_host/src/hypervisor/hyperv_windows.rs
index 34bef6779..8a28a3f17 100644
--- a/src/hyperlight_host/src/hypervisor/hyperv_windows.rs
+++ b/src/hyperlight_host/src/hypervisor/hyperv_windows.rs
@@ -501,13 +501,13 @@ impl Hypervisor for HypervWindowsDriver {
 pub mod tests {
     use std::sync::{Arc, Mutex};
 
+    use hyperlight_testing::dummy_guest_as_string;
     use serial_test::serial;
 
     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]
@@ -517,10 +517,8 @@ pub mod tests {
         let (hshm, gshm) = sandbox.mgr.build();
         drop(gshm);
         let host_funcs = sandbox.host_funcs.clone();
-        
-        let outb_handler = {
-            crate::sandbox::outb::outb_handler_wrapper(hshm.clone(), host_funcs)
-        };
+
+        let outb_handler = { 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(()) });
             Arc::new(Mutex::new(MemAccessHandler::from(func)))
diff --git a/src/hyperlight_host/src/hypervisor/kvm.rs b/src/hyperlight_host/src/hypervisor/kvm.rs
index 96388a7db..504ecb21e 100644
--- a/src/hyperlight_host/src/hypervisor/kvm.rs
+++ b/src/hyperlight_host/src/hypervisor/kvm.rs
@@ -637,13 +637,14 @@ impl Hypervisor for KVMDriver {
 mod tests {
     use std::sync::{Arc, Mutex};
 
+    use hyperlight_testing::dummy_guest_as_string;
+
     #[cfg(gdb)]
     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 {}
@@ -674,10 +675,9 @@ mod tests {
         let (hshm, gshm) = sandbox.mgr.build();
         drop(gshm);
         let host_funcs = sandbox.host_funcs.clone();
-        
-        let outb_handler: Arc<Mutex<OutBHandler>> = {
-            crate::sandbox::outb::outb_handler_wrapper(hshm.clone(), host_funcs)
-        };
+
+        let outb_handler: Arc<Mutex<OutBHandler>> =
+            { 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(()) });
             Arc::new(Mutex::new(MemAccessHandler::from(func)))

From b50e0b050f8a944433270f4d446453dfeeff2bd4 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 21 Jul 2025 19:37:53 +0000
Subject: [PATCH 6/8] Addressing PR comment

Co-authored-by: simongdavies <1397489+simongdavies@users.noreply.github.com>
---
 src/tests/rust_guests/witguest/interface.wasm | Bin 0 -> 5725 bytes
 1 file changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 src/tests/rust_guests/witguest/interface.wasm

diff --git a/src/tests/rust_guests/witguest/interface.wasm b/src/tests/rust_guests/witguest/interface.wasm
new file mode 100644
index 0000000000000000000000000000000000000000..3dcf263cc86e53d15475f50636f2111c9c42c7fd
GIT binary patch
literal 5725
zcmeHL(RSKK5S>8+jBR2kZsNFhS~p3PG>tIAn2@Ax;`~NVAA}G>2}r7TWxzJhC-g)5
zE&Zp??r;D#@X&`m^c+2n*n4O0?(AxY6*yv81%&dSRXi`?dmod+h8u;_3#AaHffxF+
zKQducdO{|i8zo%>07IkH8#un@u<5XIi6@56V4yHzKmo&JjnwOGY&H!xr)-*R&e*ir
zykyg6^NJ0t)mg3X@?e*Z$EK%^&3ZQL*{o-?p3QnT>)EVlv!2a*HtX4}XS1HodN%9X
ztY`OV$pa_xz1fmY{KKeRI7#4y(sBUi&m9ztA1+ETVN2z7#x<aXL!;z{>5!7DO@|F7
zbwd%=m;oF34srY$KAUhWiPEqulOVP_Q8d6!4c)+v3EK-uv~f!V*Kln@bqTOrxUG>(
zxXdD_4XkQp4C5@)YT_-8OktXlN!z%ifh)MmxTGYrtC0yzvIxoC*2oCxge^@iC7C-K
zxqyod(j%Wejiiuf5%T#)BNAjLId0!c@UDgg2orXfazNP&eY}V9RO9=wpK%=oLgIal
z3WD(iI51&n;V;v8;NiFUNoe;&ILw5PVi`ms9^prn$dm{0AVWz!kq#t&r~M~Lr($mr
zK8&mx{yq9rQ*(~t*n~R^PA?c+dYV6=KViy4cxb}CFDX_^*?WZkh-r`E@fsSX@<;S<
znDzvom~eMRw4jYo(Z6QaGkCT-SU~&#6Z)4-dk)Xn&}eEuqd#Wa3wW`HMpLuUzhGJo
zYO87Ln4F-WGOZ5vH8hH5Q)c<4YCyw;!xaln+m;9}y*8dw?h1Mrn$R@iXcf2B)Z_-9
zp)Xj!1uYW}R*<Qps^j+({e44+ZA1HuuuhUV6D#WJT^9HXV}br_p2iz^W5T0s<T;<C
z&)j~|@e%M}r%x4h{8fVn7}<?hMcmT$9Dmb{A{19Jbos-#no)+b3AYy$npHQ#RN1OP
zCG&czl7{DWoK<~$JqqN>+_MN<I|!BeZ)!T_zfrfSEh;?RjYO0tu4f<yfbR;$APl4}
zv}s7S(t{7w$tVq}E2HH1WgjeHU-_8A8}tkuOZDa*e;9hQA9XX#niI!f*tK*O;r-NE
zTq;-5D8!;(AgVB%i6_$}wA65k4&Hom{YXglLhD)60{0NAnK7dP<)&UEZ7--AZz^XM
z^l%<UJy;nx3UwbjiCKX5v;5bQ3V6Pew|4$#khk`{wX45Od27dU-rDol{#9GMvAVV6
z`~OuBH`exW%)4sdRbxxgbtCVpc~|A@LEcsKuKHEG>go$KeZ2i|w??`hRx`Is`XaeW
zcTF{BlmK?OW5@mE_@32`T!B|pIH|-*)J<J45x7l{t|x@owK@~5(&gQ9qhTC{^x?Q%
LuQg83YW05rma01#

literal 0
HcmV?d00001


From e38f7073e5af4af457208e1e05c9a44ad392ce77 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 21 Jul 2025 21:02:58 +0000
Subject: [PATCH 7/8] Remove double Arc<Mutex<>> wrapping to optimize
 OutBHandler performance

- Optimized OutBHandler to avoid double Arc<Mutex<>> wrapping when possible
- Added FunctionRegistryStorage enum to handle both owned and shared FunctionRegistry
- Uses Arc::try_unwrap to extract FunctionRegistry when it's no longer shared
- Falls back to original shared behavior when FunctionRegistry has multiple references
- Added handle_outb_impl_direct for direct access path without extra locking
- This eliminates potential deadlock issues and improves performance in common case

Co-authored-by: syntactically <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_host/src/sandbox/outb.rs | 134 +++++++++++++++++++++---
 1 file changed, 119 insertions(+), 15 deletions(-)

diff --git a/src/hyperlight_host/src/sandbox/outb.rs b/src/hyperlight_host/src/sandbox/outb.rs
index 5620e6f06..cb511ccfe 100644
--- a/src/hyperlight_host/src/sandbox/outb.rs
+++ b/src/hyperlight_host/src/sandbox/outb.rs
@@ -33,33 +33,59 @@ 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
+/// Note: To avoid double Arc<Mutex<>> wrapping, we store FunctionRegistry directly
+/// when possible, falling back to Arc<Mutex<>> only when necessary for sharing.
 pub(crate) struct OutBHandler {
     mem_mgr_wrapper: MemMgrWrapper<HostSharedMemory>,
-    host_funcs_wrapper: Arc<Mutex<FunctionRegistry>>,
+    host_funcs: FunctionRegistryStorage,
+}
+
+/// Storage for FunctionRegistry that avoids double Arc<Mutex<>> wrapping when possible
+enum FunctionRegistryStorage {
+    /// Direct ownership when FunctionRegistry doesn't need to be shared
+    Owned(FunctionRegistry),
+    /// Shared ownership when FunctionRegistry needs to be shared with other components
+    Shared(Arc<Mutex<FunctionRegistry>>),
 }
 
 impl OutBHandler {
-    /// Create a new OutBHandler with the required dependencies
+    /// Create a new OutBHandler with owned FunctionRegistry to avoid double Arc<Mutex<>> wrapping
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
-    pub fn new(
+    pub fn new_with_owned_registry(
+        mem_mgr_wrapper: MemMgrWrapper<HostSharedMemory>,
+        host_funcs: FunctionRegistry,
+    ) -> Self {
+        Self {
+            mem_mgr_wrapper,
+            host_funcs: FunctionRegistryStorage::Owned(host_funcs),
+        }
+    }
+
+    /// Create a new OutBHandler with shared FunctionRegistry (fallback for compatibility)
+    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
+    pub fn new_with_shared_registry(
         mem_mgr_wrapper: MemMgrWrapper<HostSharedMemory>,
         host_funcs_wrapper: Arc<Mutex<FunctionRegistry>>,
     ) -> Self {
         Self {
             mem_mgr_wrapper,
-            host_funcs_wrapper,
+            host_funcs: FunctionRegistryStorage::Shared(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,
-        )
+        match &mut self.host_funcs {
+            FunctionRegistryStorage::Owned(ref registry) => {
+                // Direct access - no extra locking needed
+                handle_outb_impl_direct(&mut self.mem_mgr_wrapper, registry, port, payload)
+            }
+            FunctionRegistryStorage::Shared(arc_mutex) => {
+                // Shared access - requires locking (original behavior)
+                handle_outb_impl(&mut self.mem_mgr_wrapper, arc_mutex.clone(), port, payload)
+            }
+        }
     }
 }
 
@@ -190,10 +216,62 @@ pub(crate) fn handle_outb_impl(
             let call = mem_mgr.as_mut().get_host_function_call()?; // pop output buffer
             let name = call.function_name.clone();
             let args: Vec<ParameterValue> = call.parameters.unwrap_or(vec![]);
+
+            // Lock the FunctionRegistry for the minimal time needed to call the function
+            // Note: This creates a brief double-lock situation (OutBHandler + FunctionRegistry)
+            // but it's unavoidable given the current architecture where FunctionRegistry
+            // needs to be shared between OutBHandler and other components
             let res = host_funcs
-                .try_lock()
-                .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
+                .lock()
+                .map_err(|e| {
+                    new_error!(
+                        "Error locking FunctionRegistry at {}:{}: {}",
+                        file!(),
+                        line!(),
+                        e
+                    )
+                })?
                 .call_host_function(&name, args)?;
+
+            mem_mgr
+                .as_mut()
+                .write_response_from_host_method_call(&res)?; // push input buffers
+
+            Ok(())
+        }
+        OutBAction::Abort => outb_abort(mem_mgr, data),
+        OutBAction::DebugPrint => {
+            let ch: char = match char::from_u32(data) {
+                Some(c) => c,
+                None => {
+                    return Err(new_error!("Invalid character for logging: {}", data));
+                }
+            };
+
+            eprint!("{}", ch);
+            Ok(())
+        }
+    }
+}
+
+/// Handles OutB operations from the guest (direct access version - no double locking).
+#[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+pub(crate) fn handle_outb_impl_direct(
+    mem_mgr: &mut MemMgrWrapper<HostSharedMemory>,
+    host_funcs: &FunctionRegistry,
+    port: u16,
+    data: u32,
+) -> Result<()> {
+    match port.try_into()? {
+        OutBAction::Log => outb_log(mem_mgr.as_mut()),
+        OutBAction::CallFunction => {
+            let call = mem_mgr.as_mut().get_host_function_call()?; // pop output buffer
+            let name = call.function_name.clone();
+            let args: Vec<ParameterValue> = call.parameters.unwrap_or(vec![]);
+
+            // Direct call - no locking needed since we have direct access to FunctionRegistry
+            let res = host_funcs.call_host_function(&name, args)?;
+
             mem_mgr
                 .as_mut()
                 .write_response_from_host_method_call(&res)?; // push input buffers
@@ -218,14 +296,40 @@ pub(crate) fn handle_outb_impl(
 /// Given a `MemMgrWrapper` and ` HostFuncsWrapper` -- both passed by _value_
 ///  -- return an `Arc<Mutex<OutBHandler>>` wrapping the core OUTB handler logic.
 ///
-/// TODO: pass at least the `host_funcs_wrapper` param by reference.
+/// This function attempts to optimize away double Arc<Mutex<>> wrapping by checking
+/// if the FunctionRegistry can be moved out of its Arc<Mutex<>> (i.e., when it's
+/// no longer shared with other components).
 #[instrument(skip_all, parent = Span::current(), level= "Trace")]
 pub(crate) fn outb_handler_wrapper(
     mem_mgr_wrapper: MemMgrWrapper<HostSharedMemory>,
     host_funcs_wrapper: Arc<Mutex<FunctionRegistry>>,
 ) -> Arc<Mutex<OutBHandler>> {
-    let outb_hdl = OutBHandler::new(mem_mgr_wrapper, host_funcs_wrapper);
-    Arc::new(Mutex::new(outb_hdl))
+    // Try to extract the FunctionRegistry from Arc<Mutex<>> to avoid double wrapping
+    match Arc::try_unwrap(host_funcs_wrapper) {
+        Ok(mutex) => {
+            // Successfully extracted Arc, now extract from Mutex
+            match mutex.into_inner() {
+                Ok(function_registry) => {
+                    // Success! We can use direct ownership
+                    let outb_hdl =
+                        OutBHandler::new_with_owned_registry(mem_mgr_wrapper, function_registry);
+                    Arc::new(Mutex::new(outb_hdl))
+                }
+                Err(poisoned) => {
+                    // Mutex was poisoned, fall back to shared access
+                    let function_registry = poisoned.into_inner();
+                    let outb_hdl =
+                        OutBHandler::new_with_owned_registry(mem_mgr_wrapper, function_registry);
+                    Arc::new(Mutex::new(outb_hdl))
+                }
+            }
+        }
+        Err(arc_mutex) => {
+            // Arc has multiple references, fall back to shared access
+            let outb_hdl = OutBHandler::new_with_shared_registry(mem_mgr_wrapper, arc_mutex);
+            Arc::new(Mutex::new(outb_hdl))
+        }
+    }
 }
 
 #[cfg(test)]

From b91970c86d783ccee650770705c97205676a508d Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 21 Jul 2025 23:03:17 +0000
Subject: [PATCH 8/8] Remove opportunistic Arc<Mutex<>> extraction to ensure
 consistent FunctionRegistry sharing

- Simplified OutBHandler to consistently use Arc<Mutex<FunctionRegistry>>
- Removed FunctionRegistryStorage enum and dual-path optimization
- Eliminated handle_outb_impl_direct function
- Streamlined outb_handler_wrapper to always use shared FunctionRegistry
- Maintains necessary sharing between sandbox and OutBHandler components

Co-authored-by: syntactically <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_host/src/sandbox/outb.rs | 127 +++---------------------
 1 file changed, 15 insertions(+), 112 deletions(-)

diff --git a/src/hyperlight_host/src/sandbox/outb.rs b/src/hyperlight_host/src/sandbox/outb.rs
index cb511ccfe..f6b5ae284 100644
--- a/src/hyperlight_host/src/sandbox/outb.rs
+++ b/src/hyperlight_host/src/sandbox/outb.rs
@@ -33,59 +33,33 @@ 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
-/// Note: To avoid double Arc<Mutex<>> wrapping, we store FunctionRegistry directly
-/// when possible, falling back to Arc<Mutex<>> only when necessary for sharing.
 pub(crate) struct OutBHandler {
     mem_mgr_wrapper: MemMgrWrapper<HostSharedMemory>,
-    host_funcs: FunctionRegistryStorage,
-}
-
-/// Storage for FunctionRegistry that avoids double Arc<Mutex<>> wrapping when possible
-enum FunctionRegistryStorage {
-    /// Direct ownership when FunctionRegistry doesn't need to be shared
-    Owned(FunctionRegistry),
-    /// Shared ownership when FunctionRegistry needs to be shared with other components
-    Shared(Arc<Mutex<FunctionRegistry>>),
+    host_funcs: Arc<Mutex<FunctionRegistry>>,
 }
 
 impl OutBHandler {
-    /// Create a new OutBHandler with owned FunctionRegistry to avoid double Arc<Mutex<>> wrapping
-    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
-    pub fn new_with_owned_registry(
-        mem_mgr_wrapper: MemMgrWrapper<HostSharedMemory>,
-        host_funcs: FunctionRegistry,
-    ) -> Self {
-        Self {
-            mem_mgr_wrapper,
-            host_funcs: FunctionRegistryStorage::Owned(host_funcs),
-        }
-    }
-
-    /// Create a new OutBHandler with shared FunctionRegistry (fallback for compatibility)
+    /// Create a new OutBHandler
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
-    pub fn new_with_shared_registry(
+    pub fn new(
         mem_mgr_wrapper: MemMgrWrapper<HostSharedMemory>,
-        host_funcs_wrapper: Arc<Mutex<FunctionRegistry>>,
+        host_funcs: Arc<Mutex<FunctionRegistry>>,
     ) -> Self {
         Self {
             mem_mgr_wrapper,
-            host_funcs: FunctionRegistryStorage::Shared(host_funcs_wrapper),
+            host_funcs,
         }
     }
 
     /// 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<()> {
-        match &mut self.host_funcs {
-            FunctionRegistryStorage::Owned(ref registry) => {
-                // Direct access - no extra locking needed
-                handle_outb_impl_direct(&mut self.mem_mgr_wrapper, registry, port, payload)
-            }
-            FunctionRegistryStorage::Shared(arc_mutex) => {
-                // Shared access - requires locking (original behavior)
-                handle_outb_impl(&mut self.mem_mgr_wrapper, arc_mutex.clone(), port, payload)
-            }
-        }
+        handle_outb_impl(
+            &mut self.mem_mgr_wrapper,
+            self.host_funcs.clone(),
+            port,
+            payload,
+        )
     }
 }
 
@@ -217,10 +191,6 @@ pub(crate) fn handle_outb_impl(
             let name = call.function_name.clone();
             let args: Vec<ParameterValue> = call.parameters.unwrap_or(vec![]);
 
-            // Lock the FunctionRegistry for the minimal time needed to call the function
-            // Note: This creates a brief double-lock situation (OutBHandler + FunctionRegistry)
-            // but it's unavoidable given the current architecture where FunctionRegistry
-            // needs to be shared between OutBHandler and other components
             let res = host_funcs
                 .lock()
                 .map_err(|e| {
@@ -254,82 +224,15 @@ pub(crate) fn handle_outb_impl(
     }
 }
 
-/// Handles OutB operations from the guest (direct access version - no double locking).
-#[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-pub(crate) fn handle_outb_impl_direct(
-    mem_mgr: &mut MemMgrWrapper<HostSharedMemory>,
-    host_funcs: &FunctionRegistry,
-    port: u16,
-    data: u32,
-) -> Result<()> {
-    match port.try_into()? {
-        OutBAction::Log => outb_log(mem_mgr.as_mut()),
-        OutBAction::CallFunction => {
-            let call = mem_mgr.as_mut().get_host_function_call()?; // pop output buffer
-            let name = call.function_name.clone();
-            let args: Vec<ParameterValue> = call.parameters.unwrap_or(vec![]);
-
-            // Direct call - no locking needed since we have direct access to FunctionRegistry
-            let res = host_funcs.call_host_function(&name, args)?;
-
-            mem_mgr
-                .as_mut()
-                .write_response_from_host_method_call(&res)?; // push input buffers
-
-            Ok(())
-        }
-        OutBAction::Abort => outb_abort(mem_mgr, data),
-        OutBAction::DebugPrint => {
-            let ch: char = match char::from_u32(data) {
-                Some(c) => c,
-                None => {
-                    return Err(new_error!("Invalid character for logging: {}", data));
-                }
-            };
-
-            eprint!("{}", ch);
-            Ok(())
-        }
-    }
-}
-
-/// Given a `MemMgrWrapper` and ` HostFuncsWrapper` -- both passed by _value_
-///  -- return an `Arc<Mutex<OutBHandler>>` wrapping the core OUTB handler logic.
-///
-/// This function attempts to optimize away double Arc<Mutex<>> wrapping by checking
-/// if the FunctionRegistry can be moved out of its Arc<Mutex<>> (i.e., when it's
-/// no longer shared with other components).
+/// Given a `MemMgrWrapper` and `Arc<Mutex<FunctionRegistry>>` -- both passed by _value_
+/// -- return an `Arc<Mutex<OutBHandler>>` wrapping the core OUTB handler logic.
 #[instrument(skip_all, parent = Span::current(), level= "Trace")]
 pub(crate) fn outb_handler_wrapper(
     mem_mgr_wrapper: MemMgrWrapper<HostSharedMemory>,
     host_funcs_wrapper: Arc<Mutex<FunctionRegistry>>,
 ) -> Arc<Mutex<OutBHandler>> {
-    // Try to extract the FunctionRegistry from Arc<Mutex<>> to avoid double wrapping
-    match Arc::try_unwrap(host_funcs_wrapper) {
-        Ok(mutex) => {
-            // Successfully extracted Arc, now extract from Mutex
-            match mutex.into_inner() {
-                Ok(function_registry) => {
-                    // Success! We can use direct ownership
-                    let outb_hdl =
-                        OutBHandler::new_with_owned_registry(mem_mgr_wrapper, function_registry);
-                    Arc::new(Mutex::new(outb_hdl))
-                }
-                Err(poisoned) => {
-                    // Mutex was poisoned, fall back to shared access
-                    let function_registry = poisoned.into_inner();
-                    let outb_hdl =
-                        OutBHandler::new_with_owned_registry(mem_mgr_wrapper, function_registry);
-                    Arc::new(Mutex::new(outb_hdl))
-                }
-            }
-        }
-        Err(arc_mutex) => {
-            // Arc has multiple references, fall back to shared access
-            let outb_hdl = OutBHandler::new_with_shared_registry(mem_mgr_wrapper, arc_mutex);
-            Arc::new(Mutex::new(outb_hdl))
-        }
-    }
+    let outb_hdl = OutBHandler::new(mem_mgr_wrapper, host_funcs_wrapper);
+    Arc::new(Mutex::new(outb_hdl))
 }
 
 #[cfg(test)]