bring back the previous call_guest_function_by_name behavior

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

Author: jprendes

README.md

@@ -58,7 +58,7 @@ fn main() -> hyperlight_host::Result<()> {
     let message = "Hello, World! I am executing inside of a VM :)\n".to_string();
     // in order to call a function it first must be defined in the guest and exposed so that
     // the host can call it
-    multi_use_sandbox.call_guest_function_by_name::<i32>(
+    multi_use_sandbox.run::<i32>(
         "PrintOutput",
         message,
     )?;

src/hyperlight_host/src/func/utils.rs

@@ -19,7 +19,8 @@ limitations under the License.
 /// for may parameter tuples.
 ///
 /// Usage:
-/// ```rust
+/// ```rust,ignore
+/// #// ignore as this is not a public macro, and would fail to compile
 /// use hyperlight_host::func::for_each_tuple;
 ///
 /// macro_rules! my_macro {

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -67,6 +67,7 @@ pub struct MultiUseSandbox {
     dispatch_ptr: RawPtr,
     #[cfg(gdb)]
     dbg_mem_access_fn: DbgMemAccessHandlerWrapper,
+    snapshot: Option<Snapshot>,
 }
 
 impl MultiUseSandbox {
@@ -91,6 +92,7 @@ impl MultiUseSandbox {
             dispatch_ptr,
             #[cfg(gdb)]
             dbg_mem_access_fn,
+            snapshot: None,
         }
     }
 
@@ -104,7 +106,7 @@ impl MultiUseSandbox {
             .unwrap_mgr_mut()
             .snapshot(self.id, mapped_regions_vec)?;
         Ok(Snapshot {
-            inner: memory_snapshot,
+            inner: Arc::new(memory_snapshot),
         })
     }
 
@@ -133,17 +135,43 @@ impl MultiUseSandbox {
             unsafe { self.vm.map_region(region)? };
         }
 
+        // The restored snapshot is now our most current snapshot
+        self.snapshot = Some(snapshot.clone());
+
         Ok(())
     }
 
     /// Call a guest function by name, with the given return type and arguments.
-    /// The changes made to the sandbox are persisted
+    /// The changes made to the sandbox are **not** persisted
+    #[deprecated(
+        since = "0.8.0",
+        note = "Deprecated in favour or call_guest and snapshot/restore."
+    )]
     #[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> {
+        let snapshot = match &self.snapshot {
+            Some(snapshot) => snapshot.clone(),
+            None => self.snapshot()?,
+        };
+        let res = self.run(func_name, args);
+        self.restore(&snapshot)?;
+        res
+    }
+
+    /// 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 run<Output: SupportedReturnType>(
+        &mut self,
+        func_name: &str,
+        args: impl ParameterTuple,
+    ) -> Result<Output> {
+        // Reset snapshot since we are mutating the sandbox state
+        self.snapshot = None;
         maybe_time_and_emit_guest_call(func_name, || {
             let ret = self.call_guest_function_by_name_no_reset(
                 func_name,
@@ -326,9 +354,33 @@ mod tests {
     use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags, MemoryRegionType};
     #[cfg(target_os = "linux")]
     use crate::mem::shared_mem::{ExclusiveSharedMemory, GuestSharedMemory, SharedMemory as _};
-    use crate::sandbox::{Callable, SandboxConfiguration};
+    use crate::sandbox::{Callable as _, SandboxConfiguration};
     use crate::{GuestBinary, HyperlightError, MultiUseSandbox, Result, UninitializedSandbox};
 
+    /// Tests that call_guest_function_by_name restores the state correctly
+    #[test]
+    fn test_call_guest_function_by_name() {
+        let mut sbox: MultiUseSandbox = {
+            let path = simple_guest_as_string().unwrap();
+            let u_sbox = UninitializedSandbox::new(GuestBinary::FilePath(path), None).unwrap();
+            u_sbox.evolve()
+        }
+        .unwrap();
+
+        let snapshot = sbox.snapshot().unwrap();
+
+        let _ = sbox.run::<i32>("AddToStatic", 5i32).unwrap();
+        let res: i32 = sbox.run("GetStatic", ()).unwrap();
+        assert_eq!(res, 5);
+        
+        sbox.restore(&snapshot).unwrap();
+        #[allow(deprecated)]
+        let _ = sbox.call_guest_function_by_name::<i32>("AddToStatic", 5i32).unwrap();
+        #[allow(deprecated)]
+        let res: i32 = sbox.call_guest_function_by_name("GetStatic", ()).unwrap();
+        assert_eq!(res, 0);
+    }
+
     // Tests to ensure that many (1000) function calls can be made in a call context with a small stack (1K) and heap(14K).
     // This test effectively ensures that the stack is being properly reset after each call and we are not leaking memory in the Guest.
     #[test]
@@ -378,15 +430,13 @@ mod tests {
 
         let snapshot = sbox.snapshot().unwrap();
 
-        let _ = sbox
-            .call_guest_function_by_name::<i32>("AddToStatic", 5i32)
-            .unwrap();
+        let _ = sbox.run::<i32>("AddToStatic", 5i32).unwrap();
 
-        let res: i32 = sbox.call_guest_function_by_name("GetStatic", ()).unwrap();
+        let res: i32 = sbox.run("GetStatic", ()).unwrap();
         assert_eq!(res, 5);
 
         sbox.restore(&snapshot).unwrap();
-        let res: i32 = sbox.call_guest_function_by_name("GetStatic", ()).unwrap();
+        let res: i32 = sbox.run("GetStatic", ()).unwrap();
         assert_eq!(res, 0);
     }
 

src/hyperlight_host/src/sandbox/snapshot.rs

@@ -14,11 +14,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+use std::sync::Arc;
+
 use crate::mem::shared_mem_snapshot::SharedMemorySnapshot;
 
 /// A snapshot capturing the state of the memory in a `MultiUseSandbox`.
 #[derive(Clone)]
 pub struct Snapshot {
-    /// TODO: Use Arc<SharedMemorySnapshot>
-    pub(crate) inner: SharedMemorySnapshot,
+    pub(crate) inner: Arc<SharedMemorySnapshot>,
 }