make all guest function calls persistent

Signed-off-by: Jorge Prendes <[email protected]>

Author: jprendes

src/hyperlight_host/src/func/call_ctx.rs

@@ -183,27 +183,27 @@ mod tests {
             let sbox: MultiUseSandbox = new_uninit().unwrap().evolve(Noop::default()).unwrap();
             Self { sandbox: sbox }
         }
-        pub fn call_add_to_static_multiple_times(mut self, i: i32) -> Result<TestSandbox> {
-            let mut ctx = self.sandbox.new_call_context();
+        pub fn call_add_to_static_multiple_times(&mut self, i: i32) -> Result<()> {
+            let snapshot = self.sandbox.snapshot()?;
             let mut sum: i32 = 0;
             for n in 0..i {
-                let result = ctx.call::<i32>("AddToStatic", n);
+                let result = self.sandbox.call_guest_function_by_name::<i32>("AddToStatic", n);
                 sum += n;
                 println!("{:?}", result);
                 let result = result.unwrap();
                 assert_eq!(result, sum);
             }
-            let result = ctx.finish();
-            assert!(result.is_ok());
-            self.sandbox = result.unwrap();
-            Ok(self)
+            self.sandbox.restore(&snapshot)?;
+            Ok(())
         }
 
-        pub fn call_add_to_static(mut self, i: i32) -> Result<()> {
+        pub fn call_add_to_static(&mut self, i: i32) -> Result<()> {
+            let snapshot = self.sandbox.snapshot()?;
             for n in 0..i {
                 let result = self
                     .sandbox
                     .call_guest_function_by_name::<i32>("AddToStatic", n);
+                self.sandbox.restore(&snapshot)?;
                 println!("{:?}", result);
                 let result = result.unwrap();
                 assert_eq!(result, n);
@@ -214,14 +214,14 @@ mod tests {
 
     #[test]
     fn ensure_multiusesandbox_multi_calls_dont_reset_state() {
-        let sandbox = TestSandbox::new();
+        let mut sandbox = TestSandbox::new();
         let result = sandbox.call_add_to_static_multiple_times(5);
         assert!(result.is_ok());
     }
 
     #[test]
     fn ensure_multiusesandbox_single_calls_do_reset_state() {
-        let sandbox = TestSandbox::new();
+        let mut sandbox = TestSandbox::new();
         let result = sandbox.call_add_to_static(5);
         assert!(result.is_ok());
     }

src/hyperlight_host/src/mem/mgr.rs

@@ -561,4 +561,25 @@ impl SandboxMemoryManager<HostSharedMemory> {
             self.layout.sandbox_memory_config.get_output_data_size(),
         )
     }
+
+    pub(crate) fn clear_io_buffers(&mut self) {
+        // Clear the output data buffer
+        loop {
+            let Ok(_) = self.shared_mem.try_pop_buffer_into::<Vec<u8>>(
+                self.layout.output_data_buffer_offset,
+                self.layout.sandbox_memory_config.get_output_data_size(),
+            ) else {
+                break;
+            };
+        }
+        // Clear the input data buffer
+        loop {
+            let Ok(_) = self.shared_mem.try_pop_buffer_into::<Vec<u8>>(
+                self.layout.input_data_buffer_offset,
+                self.layout.sandbox_memory_config.get_input_data_size(),
+            ) else {
+                break;
+            };
+        }
+    }
 }

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -180,28 +180,10 @@ impl MultiUseSandbox {
         Ok(())
     }
 
-    /// Call a guest function by name, with the given return type and arguments.
-    #[instrument(err(Debug), skip(self, args), parent = Span::current())]
-    pub fn call_guest_function_by_name<Output: SupportedReturnType>(
-        &mut self,
-        func_name: &str,
-        args: impl ParameterTuple,
-    ) -> Result<Output> {
-        maybe_time_and_emit_guest_call(func_name, || {
-            let ret = self.call_guest_function_by_name_no_reset(
-                func_name,
-                Output::TYPE,
-                args.into_value(),
-            );
-            self.restore_state()?;
-            Output::from_value(ret?)
-        })
-    }
-
     /// Call a guest function by name, with the given return type and arguments.
     /// The changes made to the sandbox are persisted
     #[instrument(err(Debug), skip(self, args), parent = Span::current())]
-    pub fn persist_call_guest_function_by_name<Output: SupportedReturnType>(
+    pub fn call_guest_function_by_name<Output: SupportedReturnType + std::fmt::Debug>(
         &mut self,
         func_name: &str,
         args: impl ParameterTuple,
@@ -213,7 +195,7 @@ impl MultiUseSandbox {
                 args.into_value(),
             );
             let ret = Output::from_value(ret?);
-            self.mem_mgr.unwrap_mgr_mut().push_state()?;
+            self.mem_mgr.unwrap_mgr_mut().push_state().unwrap();
             ret
         })
     }
@@ -298,7 +280,7 @@ impl MultiUseSandbox {
     ) -> Result<ReturnValue> {
         maybe_time_and_emit_guest_call(func_name, || {
             let ret = self.call_guest_function_by_name_no_reset(func_name, ret_type, args);
-            self.restore_state()?;
+            self.mem_mgr.unwrap_mgr_mut().push_state()?;
             ret
         })
     }
@@ -318,35 +300,43 @@ impl MultiUseSandbox {
         return_type: ReturnType,
         args: Vec<ParameterValue>,
     ) -> Result<ReturnValue> {
-        let fc = FunctionCall::new(
-            function_name.to_string(),
-            Some(args),
-            FunctionCallType::Guest,
-            return_type,
-        );
-
-        let buffer: Vec<u8> = fc
-            .try_into()
-            .map_err(|_| HyperlightError::Error("Failed to serialize FunctionCall".to_string()))?;
-
-        self.get_mgr_wrapper_mut()
-            .as_mut()
-            .write_guest_function_call(&buffer)?;
-
-        self.vm.dispatch_call_from_host(
-            self.dispatch_ptr.clone(),
-            self.out_hdl.clone(),
-            self.mem_hdl.clone(),
-            #[cfg(gdb)]
-            self.dbg_mem_access_fn.clone(),
-        )?;
+        let res = (|| {
+            let fc = FunctionCall::new(
+                function_name.to_string(),
+                Some(args),
+                FunctionCallType::Guest,
+                return_type,
+            );
+
+            let buffer: Vec<u8> = fc.try_into().map_err(|_| {
+                HyperlightError::Error("Failed to serialize FunctionCall".to_string())
+            })?;
 
-        self.check_stack_guard()?;
-        check_for_guest_error(self.get_mgr_wrapper_mut())?;
+            self.get_mgr_wrapper_mut()
+                .as_mut()
+                .write_guest_function_call(&buffer)?;
 
-        self.get_mgr_wrapper_mut()
-            .as_mut()
-            .get_guest_function_call_result()
+            self.vm.dispatch_call_from_host(
+                self.dispatch_ptr.clone(),
+                self.out_hdl.clone(),
+                self.mem_hdl.clone(),
+                #[cfg(gdb)]
+                self.dbg_mem_access_fn.clone(),
+            )?;
+
+            self.check_stack_guard()?;
+            check_for_guest_error(self.get_mgr_wrapper_mut())?;
+
+            self
+                .get_mgr_wrapper_mut()
+                .as_mut()
+                .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();
+
+        res
     }
 
     /// Get a handle to the interrupt handler for this sandbox,
@@ -447,13 +437,14 @@ mod tests {
 
         let snapshot = sbox.snapshot().unwrap();
 
-        let _ = sbox.persist_call_guest_function_by_name::<i32>("AddToStatic", 5i32).unwrap();
+        let _ = sbox.call_guest_function_by_name::<i32>("AddToStatic", 5i32)
+            .unwrap();
 
-        let res: i32 = sbox.persist_call_guest_function_by_name("GetStatic", ()).unwrap();
+        let res: i32 = sbox.call_guest_function_by_name("GetStatic", ()).unwrap();
         assert_eq!(res, 5);
 
         sbox.restore(&snapshot).unwrap();
-        let res: i32 = sbox.persist_call_guest_function_by_name("GetStatic", ()).unwrap();
+        let res: i32 = sbox.call_guest_function_by_name("GetStatic", ()).unwrap();
         assert_eq!(res, 0);
     }
 

src/hyperlight_host/tests/integration_test.rs

@@ -688,18 +688,6 @@ fn execute_on_heap() {
     }
 }
 
-#[test]
-fn memory_resets_after_failed_guestcall() {
-    let mut sbox1 = new_uninit_rust().unwrap().evolve(Noop::default()).unwrap();
-    sbox1
-        .call_guest_function_by_name::<String>("AddToStaticAndFail", ())
-        .unwrap_err();
-    let res = sbox1
-        .call_guest_function_by_name::<i32>("GetStatic", ())
-        .unwrap();
-    assert_eq!(res, 0, "Expected 0, got {:?}", res);
-}
-
 // checks that a recursive function with stack allocation eventually fails with stackoverflow
 #[test]
 fn recursive_stack_allocate_overflow() {