Enable tracing for hyperlight-wasm guests

Cargo.toml

@@ -13,4 +13,4 @@ repository = "https://github.com/hyperlight-dev/hyperlight-wasm"
 readme = "README.md"
 
 [workspace.dependencies]
-hyperlight-host = { version = "0.7.0", git = "https://github.com/hyperlight-dev/hyperlight", rev = "ea6fa8f", default-features = false, features = ["executable_heap", "init-paging"] }
+hyperlight-host = { version = "0.7.0", git = "https://github.com/hyperlight-dev/hyperlight", rev = "652affb5811a524db8a9a5a232820fb2e7cda7f2", default-features = false, features = ["executable_heap", "init-paging"] }

src/hyperlight_wasm/Cargo.toml

@@ -58,7 +58,7 @@ windows = { version = "0.61", features = ["Win32_System_Threading"] }
 page_size = "0.6.0"
 
 [dev-dependencies]
-hyperlight-component-macro = { version = "0.7.0", git = "https://github.com/hyperlight-dev/hyperlight", rev = "b61265e4aa9e2ecf8d648b994022caeea0205352" }
+hyperlight-component-macro = { version = "0.7.0", git = "https://github.com/hyperlight-dev/hyperlight", rev = "652affb5811a524db8a9a5a232820fb2e7cda7f2" }
 examples_common = { path = "../examples_common" }
 criterion = { version = "0.7.0", features = ["html_reports"] }
 crossbeam-queue = "0.3"
@@ -76,10 +76,14 @@ goblin = "0.10.0"
 tar = "0.4.44"
 
 [features]
-default = ["function_call_metrics", "kvm", "mshv2"]
+default = ["function_call_metrics", "kvm", "mshv3"]
+gdb = ["hyperlight-host/gdb"]
 function_call_metrics = ["hyperlight-host/function_call_metrics"]
 seccomp = ["hyperlight-host/seccomp"]
 print_debug = ["hyperlight-host/print_debug"]
+trace_guest = ["hyperlight-host/trace_guest"]
+mem_profile = ["hyperlight-host/mem_profile"]
+unwind_guest = ["hyperlight-host/unwind_guest"]
 crashdump = ["hyperlight-host/crashdump"]
 kvm = ["hyperlight-host/kvm"]
 mshv2 = ["hyperlight-host/mshv2"]

src/hyperlight_wasm/build.rs

@@ -114,7 +114,7 @@ fn build_wasm_runtime() -> PathBuf {
     let mut env_vars = env::vars().collect::<Vec<_>>();
     env_vars.retain(|(key, _)| !key.starts_with("CARGO_"));
 
-    let cmd = cmd
+    let mut cmd = cmd
         .arg("build")
         .arg("--profile")
         .arg(cargo_profile)
@@ -125,6 +125,17 @@ fn build_wasm_runtime() -> PathBuf {
         .env_clear()
         .envs(env_vars);
 
+    // If trace_guest feature is enabled, pass it as an argument to the build command
+    if env::var("CARGO_FEATURE_TRACE_GUEST").is_ok() {
+        cmd = cmd.arg("--features").arg("trace_guest");
+    }
+    if env::var("CARGO_FEATURE_MEM_PROFILE").is_ok() {
+        cmd = cmd.arg("--features").arg("mem_profile");
+    }
+    if env::var("CARGO_FEATURE_UNWIND_GUEST").is_ok() {
+        cmd = cmd.arg("--features").arg("unwind_guest");
+    }
+
     let status = cmd
         .status()
         .unwrap_or_else(|e| panic!("could not run cargo build: {}", e));

src/hyperlight_wasm/src/lib.rs

@@ -39,6 +39,8 @@ pub type Result<T> = hyperlight_host::Result<T>;
 pub use hyperlight_host::is_hypervisor_present;
 /// Create a generic HyperlightError
 pub use hyperlight_host::new_error;
+// A snapshot of the memory of a sandbox at a given point in time.
+pub use hyperlight_host::sandbox::snapshot::Snapshot;
 
 /// Get the build information for this version of hyperlight-wasm
 pub fn get_build_info() -> BuildInfo {

src/hyperlight_wasm/src/sandbox/loaded_wasm_sandbox.rs

@@ -14,14 +14,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+use std::fmt::Debug;
 use std::sync::Arc;
 
 use hyperlight_host::func::{ParameterTuple, SupportedReturnType};
 // re-export the InterruptHandle trait as it's part of the public API
 pub use hyperlight_host::hypervisor::InterruptHandle;
 use hyperlight_host::sandbox::Callable;
-use hyperlight_host::sandbox_state::sandbox::{DevolvableSandbox, Sandbox};
-use hyperlight_host::sandbox_state::transition::Noop;
+use hyperlight_host::sandbox::snapshot::Snapshot;
 use hyperlight_host::{MultiUseSandbox, Result, log_then_return, new_error};
 
 use super::metrics::METRIC_TOTAL_LOADED_WASM_SANDBOXES;
@@ -37,16 +37,15 @@ use crate::sandbox::metrics::{METRIC_ACTIVE_LOADED_WASM_SANDBOXES, METRIC_SANDBO
 /// memory context. If you want to "reset" the memory context, create
 /// a new `LoadedWasmSandbox` -- either from another `WasmSandbox` or by
 /// calling `my_loaded_wasm_sandbox.devolve()?.evolve()?`
-#[derive(Debug)]
 pub struct LoadedWasmSandbox {
     // inner is an Option<MultiUseSandbox> as we need to take ownership of it
     // We implement drop on the LoadedWasmSandbox to decrement the count of Sandboxes when it is dropped
     // because of this we cannot implement drop without making inner an Option (alternatively we could make MultiUseSandbox Copy but that would introduce other issues)
     inner: Option<MultiUseSandbox>,
+    // The state the sandbox was in before loading a wasm module. Used for transitioning back to a `WasmSandbox` (unloading the wasm module).
+    wasm_sandbox_snapshot: Snapshot,
 }
 
-impl Sandbox for LoadedWasmSandbox {}
-
 impl LoadedWasmSandbox {
     /// Call the function in the guest with the name `fn_name`, passing
     /// parameters `params`.
@@ -64,17 +63,49 @@ impl LoadedWasmSandbox {
             None => log_then_return!("No inner MultiUseSandbox to call_guest_function"),
         }
     }
+
+    /// Take a snapshot of the current state of the sandbox.
+    pub fn snapshot(&mut self) -> Result<Snapshot> {
+        match &mut self.inner {
+            Some(inner) => inner.snapshot(),
+            None => log_then_return!("No inner MultiUseSandbox to snapshot"),
+        }
+    }
+
+    /// Restore the state of the sandbox to the state captured in the given snapshot.
+    pub fn restore(&mut self, snapshot: &Snapshot) -> Result<()> {
+        match &mut self.inner {
+            Some(inner) => inner.restore(snapshot),
+            None => log_then_return!("No inner MultiUseSandbox to restore"),
+        }
+    }
+
     /// unload the wasm module and return a `WasmSandbox` that can be used to load another module
-    pub fn unload_module(self) -> Result<WasmSandbox> {
-        self.devolve(Noop::default()).inspect(|_| {
-            metrics::counter!(METRIC_SANDBOX_UNLOADS).increment(1);
-        })
+    pub fn unload_module(mut self) -> Result<WasmSandbox> {
+        self.inner
+            .as_mut()
+            .ok_or_else(|| new_error!("No inner MultiUseSandbox to unload"))?
+            .restore(&self.wasm_sandbox_snapshot)?;
+
+        self.inner
+            .take()
+            .ok_or_else(|| new_error!("No inner MultiUseSandbox to unload"))
+            .and_then(WasmSandbox::new)
+            .inspect(|_| {
+                metrics::counter!(METRIC_SANDBOX_UNLOADS).increment(1);
+            })
     }
 
-    pub(super) fn new(inner: MultiUseSandbox) -> Result<LoadedWasmSandbox> {
+    pub(super) fn new(
+        inner: MultiUseSandbox,
+        wasm_sandbox_snapshot: Snapshot,
+    ) -> Result<LoadedWasmSandbox> {
         metrics::gauge!(METRIC_ACTIVE_LOADED_WASM_SANDBOXES).increment(1);
         metrics::counter!(METRIC_TOTAL_LOADED_WASM_SANDBOXES).increment(1);
-        Ok(LoadedWasmSandbox { inner: Some(inner) })
+        Ok(LoadedWasmSandbox {
+            inner: Some(inner),
+            wasm_sandbox_snapshot,
+        })
     }
 
     /// Get a handle to the interrupt handler for this sandbox,
@@ -100,26 +131,20 @@ impl Callable for LoadedWasmSandbox {
     }
 }
 
-/// Capability to transform a `LoadedWasmSandbox` back down to a
-/// `WasmSandbox`
-impl DevolvableSandbox<LoadedWasmSandbox, WasmSandbox, Noop<LoadedWasmSandbox, WasmSandbox>>
-    for LoadedWasmSandbox
-{
-    fn devolve(mut self, _: Noop<LoadedWasmSandbox, WasmSandbox>) -> Result<WasmSandbox> {
-        let new_inner: MultiUseSandbox = match self.inner.take() {
-            Some(inner) => inner.devolve(Noop::default())?,
-            None => log_then_return!("No inner MultiUseSandbox to devolve"),
-        };
-        Ok(WasmSandbox::new(new_inner))
-    }
-}
-
 impl Drop for LoadedWasmSandbox {
     fn drop(&mut self) {
         metrics::gauge!(METRIC_ACTIVE_LOADED_WASM_SANDBOXES).decrement(1);
     }
 }
 
+impl Debug for LoadedWasmSandbox {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LoadedWasmSandbox")
+            .field("inner", &self.inner)
+            .finish()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;

src/hyperlight_wasm/src/sandbox/proto_wasm_sandbox.rs

@@ -14,15 +14,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use hyperlight_host::func::call_ctx::MultiUseGuestCallContext;
 use hyperlight_host::func::{HostFunction, ParameterTuple, Registerable, SupportedReturnType};
-use hyperlight_host::sandbox::Callable;
 #[cfg(all(feature = "seccomp", target_os = "linux"))]
 use hyperlight_host::sandbox::ExtraAllowedSyscall;
 use hyperlight_host::sandbox::config::SandboxConfiguration;
-use hyperlight_host::sandbox_state::sandbox::{EvolvableSandbox, Sandbox};
-use hyperlight_host::sandbox_state::transition::{MultiUseContextCallback, Noop};
-use hyperlight_host::{GuestBinary, MultiUseSandbox, Result, UninitializedSandbox, new_error};
+use hyperlight_host::{GuestBinary, Result, UninitializedSandbox, new_error};
 
 use super::metrics::{METRIC_ACTIVE_PROTO_WASM_SANDBOXES, METRIC_TOTAL_PROTO_WASM_SANDBOXES};
 use super::sandbox_builder::SandboxBuilder;
@@ -39,8 +35,6 @@ pub struct ProtoWasmSandbox {
     pub(super) inner: Option<UninitializedSandbox>,
 }
 
-impl Sandbox for ProtoWasmSandbox {}
-
 impl Registerable for ProtoWasmSandbox {
     fn register_host_function<Args: ParameterTuple, Output: SupportedReturnType>(
         &mut self,
@@ -95,27 +89,20 @@ impl ProtoWasmSandbox {
     /// The returned `WasmSandbox` can be then be cached and used to load a different Wasm module.
     ///
     pub fn load_runtime(mut self) -> Result<WasmSandbox> {
-        let multi_use_sandbox: MultiUseSandbox = match self.inner.take() {
-            Some(s) => s.evolve(Noop::default())?,
+        let mut sandbox = match self.inner.take() {
+            Some(s) => s.evolve()?,
             None => return Err(new_error!("No inner sandbox found.")),
         };
 
-        let func = Box::new(move |call_ctx: &mut MultiUseGuestCallContext| {
-            let res: i32 = call_ctx.call("InitWasmRuntime", ())?;
-            if res != 0 {
-                return Err(new_error!(
-                    "InitWasmRuntime Failed  with error code {:?}",
-                    res
-                ));
-            }
-            Ok(())
-        });
-
-        let transition_func = MultiUseContextCallback::from(func);
-
-        let new_sbox: MultiUseSandbox = multi_use_sandbox.evolve(transition_func)?;
+        let res: i32 = sandbox.call_guest_function_by_name("InitWasmRuntime", ())?;
+        if res != 0 {
+            return Err(new_error!(
+                "InitWasmRuntime Failed  with error code {:?}",
+                res
+            ));
+        }
 
-        Ok(WasmSandbox::new(new_sbox))
+        WasmSandbox::new(sandbox)
     }
 
     /// Register the given host function `host_func` with `self` under