Only clear io buffer after unsuccesfull guest call (#811)

* Only clear io buffer after unsuccesfull guest call.

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

Undo stuff that breaks unwinding

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

* Add error path to test

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

---------

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

Author: ludfjig

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -425,9 +425,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 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
     }
 
@@ -497,6 +502,7 @@ mod tests {
     use std::sync::{Arc, Barrier};
     use std::thread;
 
+    use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
     use hyperlight_testing::simple_guest_as_string;
 
     #[cfg(target_os = "linux")]
@@ -506,6 +512,29 @@ mod tests {
     use crate::sandbox::SandboxConfiguration;
     use crate::{GuestBinary, HyperlightError, MultiUseSandbox, Result, UninitializedSandbox};
 
+    /// Make sure input/output buffers are properly reset after guest call (with host call)
+    #[test]
+    fn io_buffer_reset() {
+        let mut cfg = SandboxConfiguration::default();
+        cfg.set_input_data_size(4096);
+        cfg.set_output_data_size(4096);
+        let path = simple_guest_as_string().unwrap();
+        let mut sandbox =
+            UninitializedSandbox::new(GuestBinary::FilePath(path), Some(cfg)).unwrap();
+        sandbox.register("HostAdd", |a: i32, b: i32| a + b).unwrap();
+        let mut sandbox = sandbox.evolve().unwrap();
+
+        // will exhaust io if leaky. Tests both success and error paths
+        for _ in 0..1000 {
+            let result = sandbox.call::<i32>("Add", (5i32, 10i32)).unwrap();
+            assert_eq!(result, 15);
+            let result = sandbox.call::<i32>("AddToStaticAndFail", ()).unwrap_err();
+            assert!(
+                matches!(result, HyperlightError::GuestError (code, msg ) if code == ErrorCode::GuestError && msg == "Crash on purpose")
+            );
+        }
+    }
+
     /// Tests that call_guest_function_by_name restores the state correctly
     #[test]
     fn test_call_guest_function_by_name() {