replace evolving a MultiUseSandbox with persist_call_guest_function_by_name

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

Author: jprendes

src/hyperlight_host/src/func/call_ctx.rs

@@ -56,20 +56,6 @@ impl MultiUseGuestCallContext {
         self.sbox.restore_state()?;
         Ok(self.sbox)
     }
-    /// Close out the context and get back the internally-stored
-    /// `MultiUseSandbox`.
-    ///
-    /// Note that this method is pub(crate) and does not reset the state of the
-    /// sandbox.
-    ///
-    /// It is intended to be used when evolving a MultiUseSandbox to a new state
-    /// and is not intended to be called publicly. It allows the state of the guest to be altered
-    /// during the evolution of one sandbox state to another, enabling the new state created
-    /// to be captured and stored in the Sandboxes state stack.
-    ///
-    pub(crate) fn finish_no_reset(self) -> MultiUseSandbox {
-        self.sbox
-    }
 }
 
 impl Callable for MultiUseGuestCallContext {

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -35,8 +35,7 @@ use crate::hypervisor::{Hypervisor, InterruptHandle};
 use crate::mem::ptr::RawPtr;
 use crate::mem::shared_mem::HostSharedMemory;
 use crate::metrics::maybe_time_and_emit_guest_call;
-use crate::sandbox_state::sandbox::{EvolvableSandbox, Sandbox};
-use crate::sandbox_state::transition::MultiUseContextCallback;
+use crate::sandbox_state::sandbox::Sandbox;
 use crate::{HyperlightError, Result};
 
 /// A sandbox that supports being used Multiple times.
@@ -189,6 +188,26 @@ impl MultiUseSandbox {
         })
     }
 
+    /// 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>(
+        &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(),
+            );
+            let ret = Output::from_value(ret?);
+            self.mem_mgr.unwrap_mgr_mut().push_state()?;
+            ret
+        })
+    }
+
     /// This function is kept here for fuzz testing the parameter and return types
     #[cfg(feature = "fuzzing")]
     #[instrument(err(Debug), skip(self, args), parent = Span::current())]
@@ -279,48 +298,16 @@ impl std::fmt::Debug for MultiUseSandbox {
     }
 }
 
-impl<'a, F>
-    EvolvableSandbox<
-        MultiUseSandbox,
-        MultiUseSandbox,
-        MultiUseContextCallback<'a, MultiUseSandbox, F>,
-    > for MultiUseSandbox
-where
-    F: FnOnce(&mut MultiUseGuestCallContext) -> Result<()> + 'a,
-{
-    /// The purpose of this function is to allow multiple states to be associated with a single MultiUseSandbox.
-    ///
-    /// An implementation such as HyperlightJs or HyperlightWasm can use this to call guest functions to load JS or WASM code and then evolve the sandbox causing state to be captured.
-    /// The new MultiUseSandbox can then be used to call guest functions to execute the loaded code.
-    ///
-    /// The evolve function creates a new MultiUseCallContext which is then passed to a callback function  allowing the
-    /// callback function to call guest functions as part of the evolve process, once the callback function  is complete
-    /// the context is finished using a crate internal method that does not restore the prior state of the Sandbox.
-    /// It then creates a mew  memory snapshot on the snapshot stack and returns the MultiUseSandbox
-    #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
-    fn evolve(
-        self,
-        transition_func: MultiUseContextCallback<'a, MultiUseSandbox, F>,
-    ) -> Result<MultiUseSandbox> {
-        let mut ctx = self.new_call_context();
-        transition_func.call(&mut ctx)?;
-        let mut sbox = ctx.finish_no_reset();
-        sbox.mem_mgr.unwrap_mgr_mut().push_state()?;
-        Ok(sbox)
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use std::sync::{Arc, Barrier};
     use std::thread;
 
     use hyperlight_testing::simple_guest_as_string;
 
-    use crate::func::call_ctx::MultiUseGuestCallContext;
     use crate::sandbox::{Callable, SandboxConfiguration};
     use crate::sandbox_state::sandbox::EvolvableSandbox;
-    use crate::sandbox_state::transition::{MultiUseContextCallback, Noop};
+    use crate::sandbox_state::transition::Noop;
     use crate::{GuestBinary, HyperlightError, MultiUseSandbox, Result, UninitializedSandbox};
 
     // Tests to ensure that many (1000) function calls can be made in a call context with a small stack (1K) and heap(14K).
@@ -375,18 +362,13 @@ mod tests {
 
         let snapshot = sbox.snapshot().unwrap();
 
-        let func = Box::new(|call_ctx: &mut MultiUseGuestCallContext| {
-            call_ctx.call::<i32>("AddToStatic", 5i32)?;
-            Ok(())
-        });
-        let transition_func = MultiUseContextCallback::from(func);
-        let mut sbox = sbox.evolve(transition_func).unwrap();
+        sbox.persist_call_guest_function_by_name::<i32>("AddToStatic", 5i32).unwrap();
 
-        let res: i32 = sbox.call_guest_function_by_name("GetStatic", ()).unwrap();
+        let res: i32 = sbox.persist_call_guest_function_by_name("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.persist_call_guest_function_by_name("GetStatic", ()).unwrap();
         assert_eq!(res, 0);
     }
 

src/hyperlight_host/src/sandbox_state/transition.rs

@@ -16,11 +16,7 @@ limitations under the License.
 
 use std::marker::PhantomData;
 
-use tracing::{Span, instrument};
-
 use super::sandbox::Sandbox;
-use crate::Result;
-use crate::func::call_ctx::MultiUseGuestCallContext;
 
 /// Metadata about an evolution. Any `Sandbox` implementation
 /// that also implements `EvolvableSandbox` can decide the following
@@ -82,60 +78,3 @@ impl<Cur: Sandbox, Next: Sandbox> Default for Noop<Cur, Next> {
 }
 
 impl<Cur: Sandbox, Next: Sandbox> TransitionMetadata<Cur, Next> for Noop<Cur, Next> {}
-
-/// A `TransitionMetadata` that calls a callback. The callback function takes
-/// a mutable reference to a `MultiUseGuestCallContext` and returns a `Result<()>`
-/// to signify success or failure of the function.
-///
-/// The function use the context to call guest functions.
-///
-/// Construct one of these by passing your callback to
-/// `MultiUseContextCallback::from`, as in the following code (assuming `MySandbox`
-/// is a `Sandbox` implementation):
-///
-/// ```ignore
-/// let my_cb_fn: dyn FnOnce(&mut MultiUseGuestCallContext) -> Result<()> = |_sbox| {
-///     println!("hello world!");
-/// };
-/// let mutating_cb = MultiUseContextCallback::from(my_cb_fn);
-/// ```
-pub struct MultiUseContextCallback<'func, Cur: Sandbox, F>
-where
-    F: FnOnce(&mut MultiUseGuestCallContext) -> Result<()> + 'func,
-{
-    cur_ph: PhantomData<Cur>,
-    fn_life_ph: PhantomData<&'func ()>,
-    cb: F,
-}
-
-impl<Cur: Sandbox, Next: Sandbox, F> TransitionMetadata<Cur, Next>
-    for MultiUseContextCallback<'_, Cur, F>
-where
-    F: FnOnce(&mut MultiUseGuestCallContext) -> Result<()>,
-{
-}
-
-impl<Cur: Sandbox, F> MultiUseContextCallback<'_, Cur, F>
-where
-    F: FnOnce(&mut MultiUseGuestCallContext) -> Result<()>,
-{
-    /// Invokes the callback on the provided guest context
-    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
-    pub fn call(self, cur: &mut MultiUseGuestCallContext) -> Result<()> {
-        (self.cb)(cur)
-    }
-}
-
-impl<'a, Cur: Sandbox, F> From<F> for MultiUseContextCallback<'a, Cur, F>
-where
-    F: FnOnce(&mut MultiUseGuestCallContext) -> Result<()> + 'a,
-{
-    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
-    fn from(val: F) -> Self {
-        MultiUseContextCallback {
-            cur_ph: PhantomData,
-            fn_life_ph: PhantomData,
-            cb: val,
-        }
-    }
-}