Only clear io buffer after unsuccesfull guest call.

This commit also clears poisoned Mutex after a host function panics and
introduces a new HyperlightError variant for use when a host function
panics.

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_host/src/error.rs

@@ -122,6 +122,10 @@ pub enum HyperlightError {
     #[error("HostFunction {0} was not found")]
     HostFunctionNotFound(String),
 
+    /// Host function panicked
+    #[error("Host function '{0}' panicked: {1}")]
+    HostFunctionPanic(String, String),
+
     /// Reading Writing or Seeking data failed.
     #[error("Reading Writing or Seeking data failed {0:?}")]
     IOError(#[from] std::io::Error),

src/hyperlight_host/src/func/host_functions.rs

@@ -179,9 +179,21 @@ macro_rules! impl_host_function {
                 let func = Mutex::new(func);
                 HostFunction {
                     func: Arc::new(move |args: ($($P,)*)| {
-                        func.try_lock()
-                            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
-                            (args)
+                        match func.try_lock() {
+                            Ok(mut guard) => guard(args),
+                            Err(poison_err) => {
+                                match poison_err {
+                                    // The previous call to this host function panicked, poisoning the lock.
+                                    // We can clear the poison safely.
+                                    std::sync::TryLockError::Poisoned(guard) => {
+                                        guard.into_inner()(args)
+                                    }
+                                    std::sync::TryLockError::WouldBlock => {
+                                        Err(new_error!("Error locking at {}:{}: mutex would block", file!(), line!()))
+                                    }
+                                }
+                            }
+                        }
                     })
                 }
             }

src/hyperlight_host/src/sandbox/host_funcs.rs

@@ -190,7 +190,18 @@ fn maybe_with_seccomp<T: Send>(
                             return Err(crate::HyperlightError::DisallowedSyscall);
                         }
 
-                        crate::log_then_return!("Host function {} panicked", name);
+                        let panic_msg = if let Some(s) = err.downcast_ref::<String>() {
+                            s.clone()
+                        } else if let Some(s) = err.downcast_ref::<&str>() {
+                            s.to_string()
+                        } else {
+                            "Unknown panic".to_string()
+                        };
+
+                        return Err(crate::HyperlightError::HostFunctionPanic(
+                            name.to_string(),
+                            panic_msg,
+                        ));
                     }
                 }
             })?

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -421,9 +421,14 @@ impl MultiUseSandbox {
                 .get_guest_function_call_result()
         })();
 
-        // TODO: Do we want to allow re-entrant guest function calls?
-        self.get_mgr_wrapper_mut().as_mut().clear_io_buffers();
-
+        // In the happy path we do not need to clear io-buffers from the host because:
+        // - the serialized guest function call is be zeroed out by the guest during deserialization, see call to `try_pop_shared_input_data_into::<FunctionCall>()`
+        // - the serialized guest function result is zeroed out by us (the host) during deserialization, see `get_guest_function_call_result`
+        // - any serialized host function call are zeroed out by us (the host) during deserialization, see `get_host_function_call`
+        // - any serialized host function result is zeroed out by the guest during deserialization, see `get_host_return_value`
+        if res.is_err() {
+            self.get_mgr_wrapper_mut().as_mut().clear_io_buffers();
+        }
         res
     }
 
@@ -502,6 +507,35 @@ mod tests {
     use crate::sandbox::SandboxConfiguration;
     use crate::{GuestBinary, HyperlightError, MultiUseSandbox, Result, UninitializedSandbox};
 
+    // make sure memory is cleared even if host function panics
+    #[test]
+    fn test_host_func_panic_clear_buffer() {
+        let path = simple_guest_as_string().unwrap();
+        let mut cfg = SandboxConfiguration::default();
+        cfg.set_input_data_size(20 * 1024);
+        cfg.set_output_data_size(20 * 1024);
+        let mut u_sbox = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(cfg)).unwrap();
+        u_sbox
+            .register("host_panic", |_bytes: Vec<u8>| {
+                panic!("Host function panic msg");
+                #[expect(unreachable_code, reason = "Needed for type inference")]
+                Ok(())
+            })
+            .unwrap();
+        let mut sbox = u_sbox.evolve().unwrap();
+
+        for _ in 0..30 {
+            // if io-buffers are not cleared on host func panic, this would eventually fail due to filling up input/output buffer
+            let res = sbox
+                .call::<Vec<u8>>("CallHostPanic", (vec![1; 1024],))
+                .unwrap_err();
+            assert!(matches!(
+                res,
+                crate::HyperlightError::HostFunctionPanic(func_name, payload) if func_name == "host_panic" && payload == "Host function panic msg"
+            ));
+        }
+    }
+
     /// Tests that call_guest_function_by_name restores the state correctly
     #[test]
     fn test_call_guest_function_by_name() {

src/tests/rust_guests/callbackguest/Cargo.lock

@@ -898,9 +898,27 @@ fn exec_mapped_buffer(function_call: &FunctionCall) -> Result<Vec<u8>> {
     }
 }
 
+fn call_host_panic(_: &FunctionCall) -> Result<Vec<u8>> {
+    call_host_function::<()>(
+        "host_panic",
+        Some(vec![ParameterValue::VecBytes(vec![1; 1024])]),
+        ReturnType::VecBytes,
+    )
+    .unwrap();
+    Ok(get_flatbuffer_result::<&[u8]>(vec![1; 1024].as_ref()))
+}
+
 #[no_mangle]
 #[hyperlight_guest_tracing::trace_function]
 pub extern "C" fn hyperlight_main() {
+    let call_host_panic_def = GuestFunctionDefinition::new(
+        "CallHostPanic".to_string(),
+        Vec::from(&[ParameterType::VecBytes]),
+        ReturnType::VecBytes,
+        call_host_panic as usize,
+    );
+    register_function(call_host_panic_def);
+
     let read_from_user_memory_def = GuestFunctionDefinition::new(
         "ReadFromUserMemory".to_string(),
         Vec::from(&[ParameterType::ULong, ParameterType::VecBytes]),