Add trait to abstract over sandboxes in guest function calling

This is useful to allow host-side code to work on both core Hyperlight
sandboxes and wrapped sandboxes like Hyperlight Wasm.

Signed-off-by: Lucy Menon <[email protected]>

Author: syntactically

src/hyperlight_host/benches/benchmarks.rs

@@ -19,7 +19,7 @@ use std::sync::{Arc, Mutex};
 use criterion::{criterion_group, criterion_main, Criterion};
 use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType};
 use hyperlight_host::func::HostFunction2;
-use hyperlight_host::sandbox::{MultiUseSandbox, UninitializedSandbox};
+use hyperlight_host::sandbox::{Callable, MultiUseSandbox, UninitializedSandbox};
 use hyperlight_host::sandbox_state::sandbox::EvolvableSandbox;
 use hyperlight_host::sandbox_state::transition::Noop;
 use hyperlight_host::GuestBinary;

src/hyperlight_host/examples/func_ctx/main.rs

@@ -16,7 +16,7 @@ limitations under the License.
 
 use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType};
 use hyperlight_host::func::call_ctx::MultiUseGuestCallContext;
-use hyperlight_host::sandbox::{MultiUseSandbox, UninitializedSandbox};
+use hyperlight_host::sandbox::{Callable, MultiUseSandbox, UninitializedSandbox};
 use hyperlight_host::sandbox_state::sandbox::EvolvableSandbox;
 use hyperlight_host::sandbox_state::transition::Noop;
 use hyperlight_host::{new_error, GuestBinary, Result};

src/hyperlight_host/examples/logging/main.rs

@@ -19,6 +19,7 @@ use std::sync::{Arc, Mutex};
 
 use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType};
 use hyperlight_host::sandbox::uninitialized::UninitializedSandbox;
+use hyperlight_host::sandbox::Callable;
 use hyperlight_host::sandbox_state::sandbox::EvolvableSandbox;
 use hyperlight_host::sandbox_state::transition::Noop;
 use hyperlight_host::{GuestBinary, MultiUseSandbox, Result};

src/hyperlight_host/examples/metrics/main.rs

@@ -20,6 +20,7 @@ use std::thread::{spawn, JoinHandle};
 
 use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType};
 use hyperlight_host::sandbox::uninitialized::UninitializedSandbox;
+use hyperlight_host::sandbox::Callable;
 use hyperlight_host::sandbox_state::sandbox::EvolvableSandbox;
 use hyperlight_host::sandbox_state::transition::Noop;
 use hyperlight_host::{set_metrics_registry, GuestBinary, MultiUseSandbox, Result};

src/hyperlight_host/examples/otlp_tracing/main.rs

@@ -29,6 +29,7 @@ use std::sync::{Arc, Mutex};
 use std::thread::{self, spawn, JoinHandle};
 
 use hyperlight_host::sandbox::uninitialized::UninitializedSandbox;
+use hyperlight_host::sandbox::Callable;
 use hyperlight_host::sandbox_state::sandbox::EvolvableSandbox;
 use hyperlight_host::sandbox_state::transition::Noop;
 use hyperlight_host::{GuestBinary, MultiUseSandbox, Result as HyperlightResult};

src/hyperlight_host/examples/tracing/main.rs

@@ -21,6 +21,7 @@ use std::sync::{Arc, Mutex};
 use std::thread::{spawn, JoinHandle};
 
 use hyperlight_host::sandbox::uninitialized::UninitializedSandbox;
+use hyperlight_host::sandbox::Callable;
 use hyperlight_host::sandbox_state::sandbox::EvolvableSandbox;
 use hyperlight_host::sandbox_state::transition::Noop;
 use hyperlight_host::{GuestBinary, MultiUseSandbox, Result};

src/hyperlight_host/src/func/call_ctx.rs

@@ -20,6 +20,7 @@ use hyperlight_common::flatbuffer_wrappers::function_types::{
 use tracing::{instrument, Span};
 
 use super::guest_dispatch::call_function_on_guest;
+use crate::sandbox::Callable;
 use crate::{MultiUseSandbox, Result};
 /// A context for calling guest functions.
 ///
@@ -50,6 +51,31 @@ impl MultiUseGuestCallContext {
         Self { sbox }
     }
 
+    /// Close out the context and get back the internally-stored
+    /// `MultiUseSandbox`. Future contexts opened by the returned sandbox
+    /// will have guest state restored.
+    #[instrument(err(Debug), skip(self), parent = Span::current())]
+    pub fn finish(mut self) -> Result<MultiUseSandbox> {
+        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 {
     /// Call the guest function called `func_name` with the given arguments
     /// `args`, and expect the return value have the same type as
     /// `func_ret_type`.
@@ -61,7 +87,7 @@ impl MultiUseGuestCallContext {
     /// If you want  to reset state, call `finish()` on this `MultiUseGuestCallContext`
     /// and get a new one from the resulting `MultiUseSandbox`
     #[instrument(err(Debug),skip(self, args),parent = Span::current())]
-    pub fn call(
+    fn call(
         &mut self,
         func_name: &str,
         func_ret_type: ReturnType,
@@ -74,29 +100,6 @@ impl MultiUseGuestCallContext {
 
         call_function_on_guest(&mut self.sbox, func_name, func_ret_type, args)
     }
-
-    /// Close out the context and get back the internally-stored
-    /// `MultiUseSandbox`. Future contexts opened by the returned sandbox
-    /// will have guest state restored.
-    #[instrument(err(Debug), skip(self), parent = Span::current())]
-    pub fn finish(mut self) -> Result<MultiUseSandbox> {
-        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
-    }
 }
 
 #[cfg(test)]
@@ -109,6 +112,7 @@ mod tests {
     };
     use hyperlight_testing::simple_guest_as_string;
 
+    use crate::sandbox::Callable;
     use crate::sandbox_state::sandbox::EvolvableSandbox;
     use crate::sandbox_state::transition::Noop;
     use crate::{GuestBinary, HyperlightError, MultiUseSandbox, Result, UninitializedSandbox};

src/hyperlight_host/src/func/guest_dispatch.rs

@@ -113,8 +113,8 @@ mod tests {
     use super::*;
     use crate::func::call_ctx::MultiUseGuestCallContext;
     use crate::func::host_functions::HostFunction0;
-    use crate::sandbox::is_hypervisor_present;
     use crate::sandbox::uninitialized::GuestBinary;
+    use crate::sandbox::{is_hypervisor_present, Callable};
     use crate::sandbox_state::sandbox::EvolvableSandbox;
     use crate::sandbox_state::transition::Noop;
     use crate::{new_error, HyperlightError, MultiUseSandbox, Result, UninitializedSandbox};

src/hyperlight_host/src/sandbox/callable.rs

@@ -0,0 +1,16 @@
+use hyperlight_common::flatbuffer_wrappers::function_types::{
+    ParameterValue, ReturnType, ReturnValue,
+};
+
+use crate::Result;
+
+/// Trait used by the macros to paper over the differences between hyperlight and hyperlight-wasm
+pub trait Callable {
+    /// Call a guest function dynamically
+    fn call(
+        &mut self,
+        func_name: &str,
+        func_ret_type: ReturnType,
+        args: Option<Vec<ParameterValue>>,
+    ) -> Result<ReturnValue>;
+}

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -263,7 +263,7 @@ mod tests {
     use hyperlight_testing::simple_guest_as_string;
 
     use crate::func::call_ctx::MultiUseGuestCallContext;
-    use crate::sandbox::SandboxConfiguration;
+    use crate::sandbox::{Callable, SandboxConfiguration};
     use crate::sandbox_state::sandbox::{DevolvableSandbox, EvolvableSandbox};
     use crate::sandbox_state::transition::{MultiUseContextCallback, Noop};
     use crate::{GuestBinary, MultiUseSandbox, UninitializedSandbox};