From 053827f94ffcf27fef0a326d06873bbb66d5b9d7 Mon Sep 17 00:00:00 2001 From: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com> Date: Thu, 5 Nov 2026 12:39:36 -0800 Subject: [PATCH 1/2] Rebuild the guest runtime when the WIT world file changes The guest bindings macro reads the WIT_WORLD file at compile time, but neither build script tracked it. Editing the file kept stale guest bindings until another change forced a rebuild. Track the file in both build scripts. Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com> --- src/hyperlight_wasm/build.rs | 6 ++++++ src/hyperlight_wasm_runtime/build.rs | 4 +++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/hyperlight_wasm/build.rs b/src/hyperlight_wasm/build.rs index eda934a2..bd072b24 100644 --- a/src/hyperlight_wasm/build.rs +++ b/src/hyperlight_wasm/build.rs @@ -127,6 +127,12 @@ fn build_wasm_runtime() -> PathBuf { println!("cargo::rerun-if-changed={}", runtime_dir.display()); println!("cargo::rerun-if-env-changed=WIT_WORLD"); println!("cargo::rerun-if-env-changed=WIT_WORLD_NAME"); + if let Some(path) = env::var_os("WIT_WORLD") { + println!( + "cargo::rerun-if-changed={}", + runtime_dir.join(path).display() + ); + } // the PROFILE env var unfortunately only gives us 1 bit of "dev or release" let cargo_profile = if profile == "debug" { "dev" } else { "release" }; diff --git a/src/hyperlight_wasm_runtime/build.rs b/src/hyperlight_wasm_runtime/build.rs index b8277603..af5f387a 100644 --- a/src/hyperlight_wasm_runtime/build.rs +++ b/src/hyperlight_wasm_runtime/build.rs @@ -90,7 +90,9 @@ fn main() { println!("cargo::rerun-if-env-changed=WIT_WORLD"); println!("cargo::rerun-if-env-changed=WIT_WORLD_NAME"); println!("cargo::rustc-check-cfg=cfg(component)"); - if env::var_os("WIT_WORLD").is_some() { + if let Some(path) = env::var_os("WIT_WORLD") { + let path = Path::new(&env::var_os("CARGO_MANIFEST_DIR").unwrap()).join(path); + println!("cargo::rerun-if-changed={}", path.display()); println!("cargo::rustc-cfg=component"); } From f6af4493c2769d84038d831fcd51f06cb66d5f77 Mon Sep 17 00:00:00 2001 From: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com> Date: Mon, 9 Nov 2026 13:35:51 -0800 Subject: [PATCH 2/2] Add persistent wasm snapshots Add typed Snapshot and Snapshot. They save to and load from OCI layouts. SandboxBuilder::from_snapshot builds a sandbox from one. Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com> --- CHANGELOG.md | 12 + Cargo.lock | 1 + Justfile | 1 + docs/observability.md | 2 +- src/hyperlight_wasm/Cargo.toml | 3 +- src/hyperlight_wasm/build.rs | 38 +- src/hyperlight_wasm/src/lib.rs | 8 +- .../src/sandbox/loaded_wasm_sandbox.rs | 36 +- src/hyperlight_wasm/src/sandbox/mod.rs | 1 + .../src/sandbox/sandbox_builder.rs | 130 ++- src/hyperlight_wasm/src/sandbox/snapshot.rs | 918 ++++++++++++++++++ .../src/sandbox/wasm_sandbox.rs | 109 ++- .../rust_guests/rust_wasm_samples/src/lib.rs | 8 + 13 files changed, 1169 insertions(+), 98 deletions(-) create mode 100644 src/hyperlight_wasm/src/sandbox/snapshot.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 15d3fe67..489e40c3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,18 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ## [Prerelease] - Unreleased +### Added + +* Add `Snapshot::save`, `Snapshot::load` and `Snapshot::checked_load` to save snapshots to disk and load them in another process. +* Add `SandboxBuilder::from_snapshot` to create a sandbox from a snapshot. +* Add `WasmSandbox::snapshot` and `WasmSandbox::restore`. + +### Changed + +* Change `LoadedWasmSandbox::snapshot`, `LoadedWasmSandbox::restore` and `WasmSandbox::load_from_snapshot` to use `Snapshot` instead of `Arc`. +* Make `unload_module` return to the snapshot's runtime state after `LoadedWasmSandbox::restore` or `WasmSandbox::load_from_snapshot`. +* Count sandboxes created from snapshots in `sandbox_loads_total`. + ## [v0.15.0] - 2026-08 ### Changed diff --git a/Cargo.lock b/Cargo.lock index b19931cc..8025f2d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2070,6 +2070,7 @@ dependencies = [ "serde", "serde_json", "tar", + "tempfile", "tokio", "toml", "tracing", diff --git a/Justfile b/Justfile index 90c63613..7a04ccae 100644 --- a/Justfile +++ b/Justfile @@ -178,6 +178,7 @@ examples-components target=default-target features="": (build-rust-component-exa run-examples-components target=default-target features="": {{ wit-world }} cargo run {{ if features =="" {''} else {"--no-default-features -F kvm -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example component_example + {{ wit-world }} cargo test -p hyperlight-wasm {{ if features =="" {''} else {"--no-default-features -F kvm -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --lib snapshot::tests::component {{ wit-world-c }} cargo run {{ if features =="" {''} else {"--no-default-features -F kvm -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example c-component {{ wit-world }} {{ wit-world-name-greeter }} cargo run {{ if features =="" {''} else {"--no-default-features -F kvm -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example component_greeter_example diff --git a/docs/observability.md b/docs/observability.md index bbea5ae8..a0ceacfd 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -15,7 +15,7 @@ The following metrics are provided and are enabled by default: * `proto_wasm_sandboxes_total` - A counter indicating the total number of proto wasm sandboxes created during the lifetime of the process * `wasm_sandboxes_total` - A counter indicating the total number of wasm sandboxes created during the lifetime of the process * `loaded_wasm_sandboxes_total` - A counter indicating the total number of loaded wasm sandboxes created during the lifetime of the process -* `sandbox_loads_total` - A counter indicating how many times a wasm sandbox has been loaded into a loaded wasm sandbox during the lifetime of the process +* `sandbox_loads_total` - A counter indicating how many loaded wasm sandboxes have been created by loading a module or restoring a snapshot during the lifetime of the process * `sandbox_unloads_total` - A counter indicating how many times a loaded wasm sandbox has been unloaded into a wasm sandbox during the lifetime of the process diff --git a/src/hyperlight_wasm/Cargo.toml b/src/hyperlight_wasm/Cargo.toml index e5dfe646..83c42925 100644 --- a/src/hyperlight_wasm/Cargo.toml +++ b/src/hyperlight_wasm/Cargo.toml @@ -67,6 +67,7 @@ cfg-if = "1" metrics = "0.24.5" env_logger = "0.11.10" hyperlight-wasm-runtime.workspace = true +serde = { version = "1.0", features = ["derive"] } [target.'cfg(windows)'.dependencies] windows = { version = "0.62", features = ["Win32_System_Threading"] } @@ -91,10 +92,10 @@ tracing = "0.1.44" tracing-subscriber = {version = "0.3.23", features = ["std", "env-filter"]} tracing-opentelemetry = "0.33.0" uuid = { version = "1.23.3", features = ["v4"] } -serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" semver = "1.0.28" bytes = { version = "1", default-features = false } +tempfile = "3" [build-dependencies] cfg_aliases = "0.2.1" diff --git a/src/hyperlight_wasm/build.rs b/src/hyperlight_wasm/build.rs index bd072b24..980418b8 100644 --- a/src/hyperlight_wasm/build.rs +++ b/src/hyperlight_wasm/build.rs @@ -108,7 +108,9 @@ fn find_target_dir() -> PathBuf { target_dir.to_path_buf() } -fn build_wasm_runtime() -> PathBuf { +/// Build the guest runtime and return its binary path and component interface hash. +/// The hash is `Some` when `WIT_WORLD` selects component mode and `None` for core modules. +fn build_wasm_runtime() -> (PathBuf, Option) { let profile = env::var_os("PROFILE").unwrap(); // Get the current target directory. @@ -127,12 +129,19 @@ fn build_wasm_runtime() -> PathBuf { println!("cargo::rerun-if-changed={}", runtime_dir.display()); println!("cargo::rerun-if-env-changed=WIT_WORLD"); println!("cargo::rerun-if-env-changed=WIT_WORLD_NAME"); - if let Some(path) = env::var_os("WIT_WORLD") { - println!( - "cargo::rerun-if-changed={}", - runtime_dir.join(path).display() - ); - } + let component_interface = env::var_os("WIT_WORLD").map(|path| { + let path = runtime_dir.join(path); + println!("cargo::rerun-if-changed={}", path.display()); + let bytes = fs::read(path).expect("could not read WIT_WORLD"); + let mut hash = blake3::Hasher::new(); + hash.update(blake3::hash(&bytes).as_bytes()); + // One file can hold several worlds. Unset selects the last one. + if let Ok(world) = env::var("WIT_WORLD_NAME") { + hash.update(b"\x01"); + hash.update(world.as_bytes()); + } + hash.finalize().to_string() + }); // the PROFILE env var unfortunately only gives us 1 bit of "dev or release" let cargo_profile = if profile == "debug" { "dev" } else { "release" }; @@ -190,7 +199,7 @@ fn build_wasm_runtime() -> PathBuf { ); } - path + (path, component_interface) } else { panic!( "could not find hyperlight-wasm-runtime after building it (expected {:?})", @@ -200,7 +209,12 @@ fn build_wasm_runtime() -> PathBuf { } fn main() -> Result<()> { - let built_runtime = build_wasm_runtime(); + let (built_runtime, component_interface) = build_wasm_runtime(); + + println!("cargo::rustc-check-cfg=cfg(component)"); + if component_interface.is_some() { + println!("cargo::rustc-cfg=component"); + } let out_dir = env::var_os("OUT_DIR").unwrap(); let wasm_runtime_resource = Path::new(&out_dir).join("hyperlight-wasm-runtime"); @@ -208,8 +222,10 @@ fn main() -> Result<()> { fs::copy(&built_runtime, &wasm_runtime_resource)?; let dest_path = Path::new(&out_dir).join("wasm_runtime_resource.rs"); let contents = format!( - "pub (super) static WASM_RUNTIME: [u8; include_bytes!({name:?}).len()] = *include_bytes!({name:?});", - name = wasm_runtime_resource.as_os_str() + "pub (super) static WASM_RUNTIME: [u8; include_bytes!({name:?}).len()] = *include_bytes!({name:?}); + pub (super) const WASM_RUNTIME_COMPONENT_INTERFACE: Option<&str> = {component_interface:?};", + name = wasm_runtime_resource.as_os_str(), + component_interface = component_interface.as_deref(), ); fs::write(dest_path, contents).unwrap(); diff --git a/src/hyperlight_wasm/src/lib.rs b/src/hyperlight_wasm/src/lib.rs index c10f493b..a170c162 100644 --- a/src/hyperlight_wasm/src/lib.rs +++ b/src/hyperlight_wasm/src/lib.rs @@ -17,6 +17,10 @@ limitations under the License. #![deny(dead_code, missing_docs, unused_mut)] //! This crate provides a Hyperlight implementation for WebAssembly (Wasm) guest code. +// Generated component bindings in tests use `alloc` paths. +#[cfg(all(test, component))] +extern crate alloc; + /// provides details about the build pub mod build_info; mod sandbox; @@ -57,13 +61,13 @@ pub use hyperlight_host::is_hypervisor_present; pub use hyperlight_host::new_error; /// The lifecycle state of a sandbox. pub use hyperlight_host::sandbox::SandboxStatus; -/// A snapshot of the memory of a sandbox at a given point in time. -pub use hyperlight_host::sandbox::snapshot::Snapshot; /// OCI Image Layout reference types used by [`Snapshot::save`] / /// [`Snapshot::load`] to identify snapshots inside an OCI Image Layout /// directory. Re-exported here so downstream crates do not need a /// direct dependency on `hyperlight-host`. pub use hyperlight_host::sandbox::snapshot::{OciDigest, OciReference, OciTag}; +pub use sandbox::sandbox_builder::{NewSandbox, SnapshotSource}; +pub use sandbox::snapshot::{Snapshot, SnapshotTarget}; /// Get the build information for this version of hyperlight-wasm pub fn get_build_info() -> BuildInfo { diff --git a/src/hyperlight_wasm/src/sandbox/loaded_wasm_sandbox.rs b/src/hyperlight_wasm/src/sandbox/loaded_wasm_sandbox.rs index a12cc0be..0929e43a 100644 --- a/src/hyperlight_wasm/src/sandbox/loaded_wasm_sandbox.rs +++ b/src/hyperlight_wasm/src/sandbox/loaded_wasm_sandbox.rs @@ -19,13 +19,15 @@ use std::sync::Arc; use hyperlight_host::func::{ParameterTuple, SupportedReturnType}; use hyperlight_host::hypervisor::InterruptHandle; -use hyperlight_host::sandbox::snapshot::Snapshot; use hyperlight_host::sandbox::{Callable, SandboxStatus}; use hyperlight_host::{Result, Sandbox, log_then_return, new_error}; use super::metrics::METRIC_TOTAL_LOADED_WASM_SANDBOXES; +use super::snapshot::Snapshot; use super::wasm_sandbox::WasmSandbox; -use crate::sandbox::metrics::{METRIC_ACTIVE_LOADED_WASM_SANDBOXES, METRIC_SANDBOX_UNLOADS}; +use crate::sandbox::metrics::{ + METRIC_ACTIVE_LOADED_WASM_SANDBOXES, METRIC_SANDBOX_LOADS, METRIC_SANDBOX_UNLOADS, +}; /// A sandbox that has both a Wasm engine and an arbitrary Wasm module /// loaded into memory. @@ -42,7 +44,7 @@ pub struct LoadedWasmSandbox { // because of this we cannot implement drop without making inner an Option (alternatively we could make Sandbox Copy but that would introduce other issues) inner: Option, // The state the sandbox was in before loading a wasm module. Used for transitioning back to a `WasmSandbox` (unloading the wasm module). - runtime_snapshot: Option>, + runtime_snapshot: Option>, } impl LoadedWasmSandbox { @@ -86,15 +88,22 @@ impl LoadedWasmSandbox { /// Returns `Err(HyperlightError::PoisonedSandbox)` if the sandbox is in a /// poisoned state. Use [`restore()`](Self::restore) with a previously /// taken snapshot to recover before taking a new snapshot. - pub fn snapshot(&mut self) -> Result> { + pub fn snapshot(&mut self) -> Result> { + let runtime = self + .runtime_snapshot + .as_ref() + .ok_or_else(|| new_error!("No runtime snapshot"))?; match &mut self.inner { - Some(inner) => inner.snapshot(), + Some(inner) => Ok(Snapshot::loaded(inner.snapshot()?, runtime)), None => log_then_return!("No inner Sandbox to snapshot"), } } /// Restore the state of the sandbox to the state captured in the given snapshot. /// + /// A later [`unload_module()`](Self::unload_module) returns to the runtime + /// state from before the given snapshot's module was loaded. + /// /// This method clears the poisoned state if the sandbox was poisoned, making /// it usable again for guest function calls. /// @@ -105,9 +114,14 @@ impl LoadedWasmSandbox { /// 1. Clear the poisoned state /// 2. Reset memory to the snapshot state /// 3. Allow subsequent [`call_guest_function()`](Self::call_guest_function) calls to succeed - pub fn restore(&mut self, snapshot: Arc) -> Result<()> { + pub fn restore(&mut self, snapshot: Snapshot) -> Result<()> { + let runtime = snapshot.runtime_snapshot(); match &mut self.inner { - Some(inner) => inner.restore(snapshot), + Some(inner) => { + inner.restore(snapshot.inner())?; + self.runtime_snapshot = Some(runtime); + Ok(()) + } None => log_then_return!("No inner Sandbox to restore"), } } @@ -131,15 +145,15 @@ impl LoadedWasmSandbox { .take() .ok_or_else(|| new_error!("No snapshot of the WasmSandbox to unload"))?; - WasmSandbox::new_from_loaded(sandbox, snapshot).inspect(|_| { - metrics::counter!(METRIC_SANDBOX_UNLOADS).increment(1); - }) + metrics::counter!(METRIC_SANDBOX_UNLOADS).increment(1); + Ok(WasmSandbox::dirty(sandbox, snapshot)) } pub(super) fn new( inner: Sandbox, - runtime_snapshot: Arc, + runtime_snapshot: Snapshot, ) -> Result { + metrics::counter!(METRIC_SANDBOX_LOADS).increment(1); metrics::gauge!(METRIC_ACTIVE_LOADED_WASM_SANDBOXES).increment(1); metrics::counter!(METRIC_TOTAL_LOADED_WASM_SANDBOXES).increment(1); Ok(LoadedWasmSandbox { diff --git a/src/hyperlight_wasm/src/sandbox/mod.rs b/src/hyperlight_wasm/src/sandbox/mod.rs index 24b3cb80..bdb41219 100644 --- a/src/hyperlight_wasm/src/sandbox/mod.rs +++ b/src/hyperlight_wasm/src/sandbox/mod.rs @@ -20,6 +20,7 @@ pub(crate) mod loaded_wasm_sandbox; pub(crate) mod metrics; /// A builder for a WasmSandbox. pub(crate) mod sandbox_builder; +pub(crate) mod snapshot; /// A Wasm Sandbox that can load a module. pub(crate) mod wasm_sandbox; diff --git a/src/hyperlight_wasm/src/sandbox/sandbox_builder.rs b/src/hyperlight_wasm/src/sandbox/sandbox_builder.rs index 2a4e94b9..a3930253 100644 --- a/src/hyperlight_wasm/src/sandbox/sandbox_builder.rs +++ b/src/hyperlight_wasm/src/sandbox/sandbox_builder.rs @@ -14,11 +14,14 @@ See the License for the specific language governing permissions and limitations under the License. */ -use hyperlight_host::func::HostFunction; +use hyperlight_host::func::{HostFunction, ParameterTuple, Registerable, SupportedReturnType}; use hyperlight_host::sandbox::SandboxConfiguration; -use hyperlight_host::{GuestBinary, HyperlightError, Result, is_hypervisor_present}; +use hyperlight_host::{ + GuestBinary, HostFunctions, HyperlightError, Result, Sandbox, is_hypervisor_present, +}; use super::proto_wasm_sandbox::ProtoWasmSandbox; +use super::snapshot::{Snapshot, SnapshotTarget}; // use large minimum scratch/heap/input data sizes // to deal with the size of wasmtime/wasi-libc aot artifacts @@ -26,22 +29,28 @@ pub const MIN_SCRATCH_SIZE: usize = 2 * 1024 * 1024; pub const MIN_INPUT_DATA_SIZE: usize = 192 * 1024; pub const MIN_HEAP_SIZE: u64 = 1024 * 1024; -/// A builder for WasmSandbox +/// Builder state for a new sandbox. +#[derive(Clone, Debug)] +pub struct NewSandbox; + +/// Builder state for a sandbox built from a snapshot. +pub struct SnapshotSource { + snapshot: Snapshot, + host_functions: HostFunctions, +} + +/// Builds a new sandbox or one from a snapshot. #[derive(Clone)] -pub struct SandboxBuilder { +pub struct SandboxBuilder { + source: S, config: SandboxConfiguration, host_print_fn: Option>, } -impl SandboxBuilder { - /// Create a new SandboxBuilder - pub fn new() -> Self { - let mut config: SandboxConfiguration = Default::default(); - config.set_h2g_pool_pages(MIN_INPUT_DATA_SIZE / hyperlight_common::vmem::PAGE_SIZE); - config.set_heap_size(MIN_HEAP_SIZE); - config.set_scratch_size(MIN_SCRATCH_SIZE); - +impl SandboxBuilder { + fn with_source(source: S, config: SandboxConfiguration) -> Self { Self { + source, config, host_print_fn: None, } @@ -78,6 +87,26 @@ impl SandboxBuilder { self } + /// Enable or disable crashdump generation for the sandbox + /// When enabled, core dumps will be generated when the guest crashes + /// This requires the `crashdump` feature to be enabled + #[cfg(feature = "crashdump")] + pub fn with_crashdump_enabled(mut self, enabled: bool) -> Self { + self.config.set_guest_core_dump(enabled); + self + } +} + +impl SandboxBuilder { + /// Configure a sandbox using the embedded runtime. + pub fn new() -> Self { + let mut config = SandboxConfiguration::default(); + config.set_h2g_pool_pages(MIN_INPUT_DATA_SIZE / hyperlight_common::vmem::PAGE_SIZE); + config.set_heap_size(MIN_HEAP_SIZE); + config.set_scratch_size(MIN_SCRATCH_SIZE); + Self::with_source(NewSandbox, config) + } + /// Set the guest output buffer size #[deprecated( since = "0.16.0", @@ -165,15 +194,6 @@ impl SandboxBuilder { self } - /// Enable or disable crashdump generation for the sandbox - /// When enabled, core dumps will be generated when the guest crashes - /// This requires the `crashdump` feature to be enabled - #[cfg(feature = "crashdump")] - pub fn with_crashdump_enabled(mut self, enabled: bool) -> Self { - self.config.set_guest_core_dump(enabled); - self - } - /// Build the ProtoWasmSandbox pub fn build(self) -> Result { if !is_hypervisor_present() { @@ -190,7 +210,73 @@ impl SandboxBuilder { } } -impl Default for SandboxBuilder { +/// ```compile_fail +/// use hyperlight_wasm::{LoadedWasmSandbox, SandboxBuilder, Snapshot}; +/// +/// fn restore(snapshot: Snapshot) { +/// let _builder = SandboxBuilder::from_snapshot(snapshot).with_guest_heap_size(1 << 21); +/// } +/// ``` +#[cfg(doctest)] +struct SnapshotBuilderHasNoLayoutSetters; + +impl SandboxBuilder> { + /// Create a builder that builds a sandbox from `snapshot`. + /// + /// Register every host function the guest imports before calling + /// [`build`](Self::build). Layout setters are unavailable because the snapshot + /// fixes the memory layout. + pub fn from_snapshot(snapshot: Snapshot) -> Self { + Self::with_source( + SnapshotSource { + snapshot, + host_functions: HostFunctions::default(), + }, + SandboxConfiguration::default(), + ) + } + + /// Register a host function the guest imports. + pub fn register( + &mut self, + name: impl AsRef, + host_func: impl Into>, + ) -> Result<()> { + self.register_host_function(name.as_ref(), host_func) + } + + /// Build the sandbox from the snapshot. + /// + /// Fails if a host function the guest imports is missing or has a different + /// signature. + pub fn build(mut self) -> Result { + if let Some(host_print_fn) = self.host_print_fn { + self.source + .host_functions + .register_host_function("HostPrint", host_print_fn)?; + } + let inner = Sandbox::from_snapshot( + self.source.snapshot.inner(), + self.source.host_functions, + Some(self.config), + )?; + self.source.snapshot.build(inner) + } +} + +impl Registerable for SandboxBuilder> { + fn register_host_function( + &mut self, + name: &str, + host_func: impl Into>, + ) -> Result<()> { + self.source + .host_functions + .register_host_function(name, host_func) + } +} + +impl Default for SandboxBuilder { fn default() -> Self { Self::new() } diff --git a/src/hyperlight_wasm/src/sandbox/snapshot.rs b/src/hyperlight_wasm/src/sandbox/snapshot.rs new file mode 100644 index 00000000..fe6cb191 --- /dev/null +++ b/src/hyperlight_wasm/src/sandbox/snapshot.rs @@ -0,0 +1,918 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +use std::fmt; +use std::marker::PhantomData; +use std::path::Path; +use std::sync::Arc; + +use hyperlight_host::sandbox::snapshot::{ + OciDigest, OciReference, OciTag, Snapshot as CoreSnapshot, +}; +use hyperlight_host::{Result, Sandbox, new_error}; +use serde::{Deserialize, Serialize}; + +use super::loaded_wasm_sandbox::LoadedWasmSandbox; +use super::wasm_sandbox::WasmSandbox; + +// Bump the version suffix when the metadata format changes. +const METADATA_NAMESPACE: &str = "dev.hyperlight.hyperlight-wasm.snapshot.v1"; +// Bump when host calls or their argument and result encodings change. +const GUEST_ABI_VERSION: u32 = 1; +const BASE_TAG_SUFFIX: &str = ".hlwasm-base"; +const MAX_TAG_LENGTH: usize = 128; + +/// A saved state of a wasm sandbox. +/// +/// `Snapshot` holds the runtime with no module loaded. +/// `Snapshot` holds a loaded module and its state, plus +/// the runtime that [`LoadedWasmSandbox::unload_module`] returns to. Clones +/// share the same underlying state. +/// +/// [`crate::SandboxBuilder::from_snapshot`] builds a sandbox of type `T` from +/// it. It takes fresh implementations of the host functions the guest imports. +/// +/// # Host state +/// +/// Host callbacks, files, sockets, and other host resources are outside the +/// snapshot. Restoring guest state leaves changes made by callbacks in place. +/// +/// # Portability +/// +/// A snapshot loads only on the CPU architecture, CPU vendor, and hypervisor +/// it was saved on. A component snapshot also needs the `WIT_WORLD` file and +/// `WIT_WORLD_NAME` that hyperlight-wasm was built with. Loading rejects a +/// mismatch. +/// +/// # Examples +/// +/// ```no_run +/// # use hyperlight_wasm::{LoadedWasmSandbox, OciTag, Result, SandboxBuilder, Snapshot}; +/// # fn example(mut loaded: LoadedWasmSandbox) -> Result<()> { +/// let digest = loaded.snapshot()?.save("./snapshots", &OciTag::new("app")?)?; +/// +/// // In another process: +/// let snapshot = Snapshot::::checked_load("./snapshots", digest)?; +/// let mut builder = SandboxBuilder::from_snapshot(snapshot); +/// builder.register("Add", |a: i32, b: i32| a + b)?; +/// let mut loaded = builder.build()?; +/// let result: i32 = loaded.call_guest_function("Run", ())?; +/// # Ok(()) +/// # } +/// ``` +pub struct Snapshot { + inner: Arc, + base: T::Base, + metadata: Arc, + target: PhantomData T>, +} + +impl Clone for Snapshot { + fn clone(&self) -> Self { + Self { + inner: Arc::clone(&self.inner), + base: self.base.clone(), + metadata: Arc::clone(&self.metadata), + target: PhantomData, + } + } +} + +impl fmt::Debug for Snapshot { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Snapshot") + .field("metadata", &self.metadata) + .finish_non_exhaustive() + } +} + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum SnapshotKind { + Runtime, + Loaded, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +struct SnapshotMetadata { + kind: SnapshotKind, + guest_abi_version: u32, + hyperlight_wasm_version: String, + component_interface: Option, + base_snapshot: Option, +} + +impl SnapshotMetadata { + fn runtime() -> Self { + Self { + kind: SnapshotKind::Runtime, + guest_abi_version: GUEST_ABI_VERSION, + hyperlight_wasm_version: env!("CARGO_PKG_VERSION").to_owned(), + component_interface: super::WASM_RUNTIME_COMPONENT_INTERFACE.map(str::to_owned), + base_snapshot: None, + } + } + + fn validate(&self) -> Result<()> { + if self.guest_abi_version != GUEST_ABI_VERSION { + return Err(new_error!( + "Snapshot guest ABI {} from hyperlight-wasm {} is incompatible with ABI {}", + self.guest_abi_version, + self.hyperlight_wasm_version, + GUEST_ABI_VERSION + )); + } + if self.component_interface.as_deref() != super::WASM_RUNTIME_COMPONENT_INTERFACE { + return Err(new_error!( + "Snapshot module/component interface does not match this host" + )); + } + match (self.kind, self.base_snapshot.is_some()) { + (SnapshotKind::Runtime, true) => Err(new_error!( + "Runtime snapshot must not reference another snapshot" + )), + (SnapshotKind::Loaded, false) => Err(new_error!( + "Loaded snapshot is missing its runtime snapshot" + )), + _ => Ok(()), + } + } +} + +mod private { + use std::path::Path; + + use hyperlight_host::Result; + use hyperlight_host::sandbox::snapshot::OciTag; + + pub trait Sealed: Sized { + const LOADED: bool; + + /// The runtime state that unloading returns to. + type Base: Clone; + + fn load_base(path: &Path, digest: Option<&str>, checked: bool) -> Result; + + /// Saves `base` alongside a snapshot tagged `tag` and returns its digest. + fn save_base(base: &Self::Base, path: &Path, tag: &OciTag) -> Result>; + + fn restore( + inner: hyperlight_host::Sandbox, + snapshot: super::Snapshot, + ) -> hyperlight_host::Result + where + Self: super::SnapshotTarget; + } +} + +/// A sandbox type that can be snapshotted. +#[allow(private_bounds)] +pub trait SnapshotTarget: private::Sealed {} + +impl private::Sealed for WasmSandbox { + const LOADED: bool = false; + + type Base = (); + + fn load_base(_: &Path, _: Option<&str>, _: bool) -> Result<()> { + Ok(()) + } + + fn save_base(_: &(), _: &Path, _: &OciTag) -> Result> { + Ok(None) + } + + fn restore(inner: Sandbox, snapshot: Snapshot) -> Result { + Ok(Self::clean(inner, snapshot)) + } +} + +impl SnapshotTarget for WasmSandbox {} + +impl private::Sealed for LoadedWasmSandbox { + const LOADED: bool = true; + + type Base = Snapshot; + + fn load_base(path: &Path, digest: Option<&str>, checked: bool) -> Result { + let digest = digest + .ok_or_else(|| new_error!("Loaded snapshot is missing its runtime snapshot"))? + .parse::()?; + Snapshot::::load_inner(path, OciReference::from(digest), checked) + } + + fn save_base(base: &Self::Base, path: &Path, tag: &OciTag) -> Result> { + let base_tag = baseline_tag(tag)?; + Ok(Some( + base.save_with_base(path, &base_tag, None)?.to_string(), + )) + } + + fn restore(inner: Sandbox, snapshot: Snapshot) -> Result { + Self::new(inner, snapshot.base) + } +} + +impl SnapshotTarget for LoadedWasmSandbox {} + +impl Snapshot { + fn new(inner: Arc, base: T::Base, metadata: SnapshotMetadata) -> Self { + Self { + inner, + base, + metadata: Arc::new(metadata), + target: PhantomData, + } + } + + /// Save to an OCI image layout and return the manifest digest. + /// + /// A `Snapshot` also saves its runtime snapshot under + /// `.hlwasm-base`. Copy both tags when distributing the layout. Tags + /// ending in `.hlwasm-base` are reserved. A failed save can leave a loaded + /// snapshot previously saved under `tag` unable to unload its module. + pub fn save(&self, path: impl AsRef, tag: &OciTag) -> Result { + let path = path.as_ref(); + if tag.as_str().ends_with(BASE_TAG_SUFFIX) { + return Err(new_error!( + "Snapshot tags ending in '{}' are reserved by hyperlight-wasm", + BASE_TAG_SUFFIX + )); + } + let base_snapshot = T::save_base(&self.base, path, tag)?; + self.save_with_base(path, tag, base_snapshot) + } + + fn save_with_base( + &self, + path: &Path, + tag: &OciTag, + base_snapshot: Option, + ) -> Result { + let mut metadata = self.metadata.as_ref().clone(); + metadata.base_snapshot = base_snapshot; + self.inner + .with_metadata(METADATA_NAMESPACE, &metadata)? + .save(path, tag) + } + + pub(super) fn inner(&self) -> Arc { + Arc::clone(&self.inner) + } + + /// Load from a trusted OCI layout, skipping blob digest checks. + /// + /// Fails if the snapshot holds the other sandbox type or is incompatible + /// with this host. + /// + /// Blob files stay mapped while this snapshot or a sandbox using it is + /// alive. Overwriting, truncating, or deleting those files can corrupt + /// guest memory and cause undefined behavior. + pub fn load(path: impl AsRef, reference: impl Into) -> Result { + Self::load_inner(path.as_ref(), reference.into(), false) + } + + /// Load and verify the digest of every blob, including the runtime snapshot + /// saved with a loaded snapshot. + /// + /// Verification checks integrity, not authenticity. Use a trusted source. + /// The file lifetime requirements of [`Self::load`] also apply. + pub fn checked_load( + path: impl AsRef, + reference: impl Into, + ) -> Result { + Self::load_inner(path.as_ref(), reference.into(), true) + } + + fn load_inner(path: &Path, reference: OciReference, checked: bool) -> Result { + let load = |reference| { + if checked { + CoreSnapshot::checked_load(path, reference) + } else { + CoreSnapshot::load(path, reference) + } + }; + let inner = Arc::new(load(reference)?); + let metadata = inner + .metadata::(METADATA_NAMESPACE)? + .ok_or_else(|| new_error!("Snapshot is missing hyperlight-wasm metadata"))?; + metadata.validate()?; + let expected = if T::LOADED { + SnapshotKind::Loaded + } else { + SnapshotKind::Runtime + }; + if metadata.kind != expected { + return Err(new_error!( + "Snapshot holds {:?} state, expected {:?}", + metadata.kind, + expected + )); + } + let base = T::load_base(path, metadata.base_snapshot.as_deref(), checked)?; + Ok(Self::new(inner, base, metadata)) + } + + pub(super) fn build(self, inner: Sandbox) -> Result { + T::restore(inner, self) + } +} + +impl Snapshot { + pub(super) fn runtime(inner: Arc) -> Self { + Self::new(inner, (), SnapshotMetadata::runtime()) + } +} + +impl Snapshot { + pub(super) fn loaded(inner: Arc, runtime: &Snapshot) -> Self { + let mut metadata = runtime.metadata.as_ref().clone(); + metadata.kind = SnapshotKind::Loaded; + Self::new(inner, runtime.clone(), metadata) + } + + pub(super) fn runtime_snapshot(&self) -> Snapshot { + self.base.clone() + } +} + +fn baseline_tag(tag: &OciTag) -> Result { + let max = MAX_TAG_LENGTH - BASE_TAG_SUFFIX.len(); + if tag.as_str().len() > max { + return Err(new_error!( + "Loaded snapshot tags must be at most {} bytes", + max + )); + } + OciTag::new(format!("{}{BASE_TAG_SUFFIX}", tag.as_str())) +} + +#[cfg(test)] +mod tests { + use examples_common::get_wasm_module_path; + use hyperlight_host::HyperlightError; + + use super::*; + use crate::SandboxBuilder; + + fn runtime_sandbox() -> Result { + let mut proto = SandboxBuilder::new() + .with_guest_heap_size(1280 * 1024) + .build()?; + proto.register("TestHostFunc", |value: i32| value + 1)?; + proto.load_runtime() + } + + fn restore(snapshot: Snapshot) -> Result { + let mut builder = SandboxBuilder::from_snapshot(snapshot); + builder.register("TestHostFunc", |value: i32| value + 1)?; + builder.build() + } + + fn loaded_sandbox() -> Result { + runtime_sandbox()?.load_module(get_wasm_module_path("rust_wasm_samples.aot")?) + } + + #[test] + fn metadata_checks_abi_and_component_contract() { + let metadata = SnapshotMetadata::runtime(); + metadata.validate().unwrap(); + + let mut invalid = metadata.clone(); + invalid.guest_abi_version += 1; + assert!( + invalid + .validate() + .unwrap_err() + .to_string() + .contains("guest ABI") + ); + + let mut invalid = metadata; + invalid.component_interface = Some("different-interface".to_owned()); + assert!( + invalid + .validate() + .unwrap_err() + .to_string() + .contains("interface") + ); + } + + #[test] + fn metadata_checks_baseline_relationship() { + let mut metadata = SnapshotMetadata::runtime(); + metadata.base_snapshot = Some("sha256:unused".to_owned()); + assert!(metadata.validate().is_err()); + metadata.kind = SnapshotKind::Loaded; + metadata.validate().unwrap(); + metadata.base_snapshot = None; + assert!(metadata.validate().is_err()); + } + + #[test] + fn baseline_tag_obeys_oci_limit() { + let max = MAX_TAG_LENGTH - BASE_TAG_SUFFIX.len(); + assert!(baseline_tag(&OciTag::new("a".repeat(max)).unwrap()).is_ok()); + assert!(baseline_tag(&OciTag::new("a".repeat(max + 1)).unwrap()).is_err()); + } + + #[test] + fn rejects_missing_or_incompatible_persisted_metadata() -> Result<()> { + let directory = tempfile::tempdir()?; + let snapshot = SandboxBuilder::new().build()?.load_runtime()?.snapshot()?; + let tag = OciTag::new("runtime")?; + + snapshot.inner.save(directory.path(), &tag)?; + assert!( + Snapshot::::checked_load(directory.path(), tag.clone()) + .unwrap_err() + .to_string() + .contains("missing hyperlight-wasm metadata") + ); + + let mut metadata = snapshot.metadata.as_ref().clone(); + metadata.guest_abi_version += 1; + snapshot + .inner + .with_metadata(METADATA_NAMESPACE, &metadata)? + .save(directory.path(), &tag)?; + assert!( + Snapshot::::checked_load(directory.path(), tag) + .unwrap_err() + .to_string() + .contains("guest ABI") + ); + Ok(()) + } + + #[test] + fn tag_validation_precedes_writes() -> Result<()> { + let directory = tempfile::tempdir()?; + let snapshot = SandboxBuilder::new().build()?.load_runtime()?.snapshot()?; + let loaded = Snapshot::::loaded(snapshot.inner(), &snapshot); + + assert!( + loaded + .save(directory.path(), &OciTag::new("a".repeat(117))?) + .is_err() + ); + assert!( + snapshot + .save(directory.path(), &OciTag::new("app.hlwasm-base")?) + .is_err() + ); + assert!(!directory.path().join("index.json").exists()); + Ok(()) + } + + #[test] + fn checked_load_verifies_the_baseline() -> Result<()> { + let directory = tempfile::tempdir()?; + let runtime = SandboxBuilder::new().build()?.load_runtime()?.snapshot()?; + let loaded = Snapshot::::loaded(runtime.inner(), &runtime); + let tag = OciTag::new("application")?; + loaded.save(directory.path(), &tag)?; + let restored = Snapshot::::load(directory.path(), tag.clone())?; + let base_digest = restored.metadata.base_snapshot.as_ref().unwrap(); + let blob = |digest: &str| { + directory + .path() + .join("blobs") + .join(digest.replace(':', "/")) + }; + let manifest: serde_json::Value = + serde_json::from_slice(&std::fs::read(blob(base_digest))?)?; + let config_path = blob(manifest["config"]["digest"].as_str().unwrap()); + let mut config = std::fs::read(&config_path)?; + *config.iter_mut().find(|byte| **byte == b'\n').unwrap() = b' '; + std::fs::write(config_path, config)?; + + Snapshot::::load(directory.path(), tag.clone())?; + assert!( + Snapshot::::checked_load(directory.path(), tag.clone()).is_err() + ); + std::fs::remove_file(blob(base_digest))?; + assert!(Snapshot::::load(directory.path(), tag).is_err()); + Ok(()) + } + + #[test] + fn baseline_is_the_runtime_checkpoint() -> Result<()> { + let directory = tempfile::tempdir()?; + let module = get_wasm_module_path("rust_wasm_samples.aot")?; + let runtime = runtime_sandbox()?; + let runtime_digest = runtime + .snapshot()? + .save(directory.path(), &OciTag::new("runtime")?)?; + let in_process = runtime + .load_module(&module)? + .snapshot()? + .save(directory.path(), &OciTag::new("in-process")?)?; + let runtime = + Snapshot::::checked_load(directory.path(), runtime_digest.clone())?; + let from_disk = restore(runtime)? + .load_module(&module)? + .snapshot()? + .save(directory.path(), &OciTag::new("from-disk")?)?; + + let base = |digest: OciDigest| -> Result> { + let snapshot = Snapshot::::load(directory.path(), digest)?; + Ok(snapshot.metadata.base_snapshot.clone()) + }; + assert_eq!(base(in_process)?, Some(runtime_digest.to_string())); + assert_eq!(base(from_disk)?, Some(runtime_digest.to_string())); + + let baseline = Snapshot::::checked_load( + directory.path(), + OciTag::new("in-process.hlwasm-base")?, + )?; + let mut loaded = restore(baseline)?.load_module(&module)?; + assert_eq!(loaded.call_guest_function::("counter", ())?, 1); + Ok(()) + } + + #[test] + fn rejects_a_baseline_of_the_wrong_kind() -> Result<()> { + let directory = tempfile::tempdir()?; + let loaded = loaded_sandbox()?.snapshot()?; + let digest = loaded.save(directory.path(), &OciTag::new("application")?)?; + let mut metadata = loaded.metadata.as_ref().clone(); + metadata.base_snapshot = Some(digest.to_string()); + let tag = OciTag::new("wrong-baseline")?; + loaded + .inner + .with_metadata(METADATA_NAMESPACE, &metadata)? + .save(directory.path(), &tag)?; + + let error = Snapshot::::load(directory.path(), tag).unwrap_err(); + assert!(error.to_string().contains("expected Runtime"), "{error}"); + Ok(()) + } + + #[test] + fn persists_runtime_and_checks_lifecycle() -> Result<()> { + let directory = tempfile::tempdir()?; + let snapshot = runtime_sandbox()?.snapshot()?; + let tag = OciTag::new("runtime")?; + let digest = snapshot.save(directory.path(), &tag)?; + + assert!(Snapshot::::checked_load(directory.path(), tag).is_err()); + let snapshot = Snapshot::::checked_load(directory.path(), digest)?; + let mut loaded = + restore(snapshot)?.load_module(get_wasm_module_path("rust_wasm_samples.aot")?)?; + assert_eq!(loaded.call_guest_function::("counter", ())?, 1); + Ok(()) + } + + #[test] + fn persists_mutable_state_and_preserves_unload() -> Result<()> { + let directory = tempfile::tempdir()?; + let mut loaded = loaded_sandbox()?; + assert_eq!(loaded.call_guest_function::("counter", ())?, 1); + let snapshot = loaded.snapshot()?; + let digest = snapshot.save(directory.path(), &OciTag::new("application")?)?; + drop(loaded); + drop(snapshot); + + assert!(Snapshot::::checked_load(directory.path(), digest.clone()).is_err()); + let snapshot = Snapshot::::checked_load(directory.path(), digest)?; + let mut loaded = restore(snapshot)?; + assert_eq!(loaded.call_guest_function::("counter", ())?, 2); + assert_eq!( + loaded.call_guest_function::("call_host_function", 4i32)?, + 5 + ); + + let runtime = loaded.unload_module()?; + let mut loaded = runtime.load_module(get_wasm_module_path("rust_wasm_samples.aot")?)?; + assert_eq!(loaded.call_guest_function::("counter", ())?, 1); + Ok(()) + } + + #[test] + fn unloaded_snapshot_captures_clean_runtime() -> Result<()> { + let runtime = runtime_sandbox()?; + let baseline = runtime.snapshot()?; + let mut loaded = runtime.load_module(get_wasm_module_path("rust_wasm_samples.aot")?)?; + loaded.call_guest_function::("counter", ())?; + let mut runtime = loaded.unload_module()?; + let snapshot = runtime.snapshot()?; + // Module memory after unload must not reach a runtime checkpoint. + assert!(Arc::ptr_eq(&snapshot.inner, &baseline.inner)); + runtime.restore(snapshot.clone())?; + let mut loaded = + restore(snapshot)?.load_module(get_wasm_module_path("rust_wasm_samples.aot")?)?; + assert_eq!(loaded.call_guest_function::("counter", ())?, 1); + Ok(()) + } + + #[test] + fn failed_runtime_restore_keeps_the_sandbox() -> Result<()> { + let mut proto = SandboxBuilder::new() + .with_guest_heap_size(1280 * 1024) + .build()?; + proto.register("TestHostFunc", |value: i32| value + 1)?; + proto.register("Extra", || 0i32)?; + let foreign = proto.load_runtime()?.snapshot()?; + + let mut runtime = loaded_sandbox()?.unload_module()?; + assert!(matches!( + runtime.restore(foreign), + Err(HyperlightError::SnapshotHostFunctionMismatch { .. }) + )); + let mut loaded = runtime.load_module(get_wasm_module_path("rust_wasm_samples.aot")?)?; + assert_eq!(loaded.call_guest_function::("counter", ())?, 1); + Ok(()) + } + + #[test] + fn restoring_loaded_state_replaces_the_unload_baseline() -> Result<()> { + let snapshot = loaded_sandbox()?.snapshot()?; + for load_from_runtime in [false, true] { + let mut proto = SandboxBuilder::new() + .with_guest_heap_size(1280 * 1024) + .build()?; + proto.register("TestHostFunc", |value: i32| value + 1)?; + proto.register("DestinationOnly", || 0i32)?; + let runtime = proto.load_runtime()?; + let loaded = if load_from_runtime { + runtime.load_from_snapshot(snapshot.clone())? + } else { + let mut loaded = + runtime.load_module(get_wasm_module_path("rust_wasm_samples.aot")?)?; + loaded.restore(snapshot.clone())?; + loaded + }; + // A destination-only baseline would require DestinationOnly here. + restore(loaded.unload_module()?.snapshot()?)?; + } + Ok(()) + } + + #[test] + fn validates_callback_replacements() -> Result<()> { + let snapshot = loaded_sandbox()?.snapshot()?; + let error = SandboxBuilder::from_snapshot(snapshot.clone()) + .build() + .unwrap_err(); + assert!(matches!( + error, + HyperlightError::SnapshotHostFunctionMismatch { missing, signature_mismatches } + if missing == ["TestHostFunc"] && signature_mismatches.is_empty() + )); + let mut builder = SandboxBuilder::from_snapshot(snapshot); + builder.register("TestHostFunc", |_: String| 0i32)?; + assert!(matches!( + builder.build().unwrap_err(), + HyperlightError::SnapshotHostFunctionMismatch { missing, signature_mismatches } + if missing.is_empty() && signature_mismatches.len() == 1 + && signature_mismatches[0].contains("TestHostFunc") + )); + Ok(()) + } + + #[test] + fn restores_custom_print_routing() -> Result<()> { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let directory = tempfile::tempdir()?; + let mut original = loaded_sandbox()?; + let digest = original + .snapshot()? + .save(directory.path(), &OciTag::new("application")?)?; + drop(original); + let snapshot = Snapshot::::checked_load(directory.path(), digest)?; + let calls = Arc::new(AtomicUsize::new(0)); + let print_calls = Arc::clone(&calls); + let mut builder = + SandboxBuilder::from_snapshot(snapshot).with_host_print_fn(move |message: String| { + print_calls.fetch_add(1, Ordering::Relaxed); + i32::try_from(message.len()).unwrap() + }); + builder.register("TestHostFunc", |value: i32| value + 1)?; + let mut loaded = builder.build()?; + loaded.call_guest_function::("hello_world", ())?; + assert_eq!(calls.load(Ordering::Relaxed), 1); + Ok(()) + } + + #[test] + fn restores_after_original_aot_file_is_removed() -> Result<()> { + let directory = tempfile::tempdir()?; + let module = directory.path().join("original.aot"); + std::fs::copy(get_wasm_module_path("rust_wasm_samples.aot")?, &module)?; + let mut loaded = runtime_sandbox()?.load_module(&module)?; + loaded.call_guest_function::("counter", ())?; + let snapshot = loaded.snapshot()?; + let digest = snapshot.save(directory.path(), &OciTag::new("application")?)?; + drop(loaded); + drop(snapshot); + std::fs::remove_file(module)?; + + let snapshot = Snapshot::::checked_load(directory.path(), digest)?; + assert_eq!( + restore(snapshot)?.call_guest_function::("counter", ())?, + 2 + ); + Ok(()) + } + + #[test] + fn buffer_loading_and_resaving_preserve_state() -> Result<()> { + let first = tempfile::tempdir()?; + let second = tempfile::tempdir()?; + let bytes = std::fs::read(get_wasm_module_path("rust_wasm_samples.aot")?)?; + // The module is sent through the transport, so the pool must hold it. + let mut proto = SandboxBuilder::new() + .with_guest_heap_size(1280 * 1024) + .with_h2g_pool_pages(128) + .with_h2g_queue_size(128) + .build()?; + proto.register("TestHostFunc", |value: i32| value + 1)?; + let mut loaded = proto.load_runtime()?.load_module_from_buffer(&bytes)?; + loaded.call_guest_function::("counter", ())?; + let digest = loaded + .snapshot()? + .save(first.path(), &OciTag::new("application")?)?; + let snapshot = Snapshot::::checked_load(first.path(), digest.clone())?; + assert_eq!(snapshot.save(second.path(), &OciTag::new("copy")?)?, digest); + let snapshot = Snapshot::::checked_load(second.path(), digest)?; + assert_eq!( + restore(snapshot)?.call_guest_function::("counter", ())?, + 2 + ); + Ok(()) + } + + #[test] + fn sandboxes_from_one_checkpoint_are_isolated() -> Result<()> { + let directory = tempfile::tempdir()?; + let mut loaded = loaded_sandbox()?; + loaded.call_guest_function::("counter", ())?; + let digest = loaded + .snapshot()? + .save(directory.path(), &OciTag::new("application")?)?; + drop(loaded); + + let snapshot = Snapshot::::load(directory.path(), digest.clone())?; + let mut first = restore(snapshot.clone())?; + assert_eq!(first.call_guest_function::("counter", ())?, 2); + assert_eq!(first.call_guest_function::("counter", ())?, 3); + let mut second = restore(snapshot)?; + assert_eq!(second.call_guest_function::("counter", ())?, 2); + + // Guest writes are copy-on-write, so the mapped blob still holds the saved state. + let snapshot = Snapshot::::checked_load(directory.path(), digest)?; + assert_eq!( + restore(snapshot)?.call_guest_function::("counter", ())?, + 2 + ); + Ok(()) + } + + #[test] + fn recaptures_restored_state_with_its_baseline() -> Result<()> { + let directory = tempfile::tempdir()?; + let mut loaded = loaded_sandbox()?; + loaded.call_guest_function::("counter", ())?; + let digest = loaded + .snapshot()? + .save(directory.path(), &OciTag::new("first")?)?; + drop(loaded); + + let snapshot = Snapshot::::checked_load(directory.path(), digest)?; + let mut loaded = restore(snapshot)?; + assert_eq!(loaded.call_guest_function::("counter", ())?, 2); + let digest = loaded + .snapshot()? + .save(directory.path(), &OciTag::new("second")?)?; + drop(loaded); + + let snapshot = Snapshot::::checked_load(directory.path(), digest)?; + let mut loaded = restore(snapshot)?; + assert_eq!(loaded.call_guest_function::("counter", ())?, 3); + let mut loaded = loaded + .unload_module()? + .load_module(get_wasm_module_path("rust_wasm_samples.aot")?)?; + assert_eq!(loaded.call_guest_function::("counter", ())?, 1); + Ok(()) + } + + #[test] + fn restores_in_another_process() -> Result<()> { + let directory = tempfile::tempdir()?; + let mut loaded = loaded_sandbox()?; + loaded.call_guest_function::("counter", ())?; + loaded + .snapshot()? + .save(directory.path(), &OciTag::new("application")?)?; + drop(loaded); + + let status = std::process::Command::new(std::env::current_exe()?) + .args([ + "--exact", + "sandbox::snapshot::tests::snapshot_reader", + "--ignored", + "--nocapture", + ]) + .env("HYPERLIGHT_WASM_SNAPSHOT_PATH", directory.path()) + .status()?; + assert!(status.success(), "snapshot reader exited with {status}"); + // libtest succeeds without running anything when the filter matches no test. + assert!(directory.path().join("reader-ran").exists()); + Ok(()) + } + + #[test] + #[ignore = "Invoked by restores_in_another_process"] + fn snapshot_reader() -> Result<()> { + let Some(path) = std::env::var_os("HYPERLIGHT_WASM_SNAPSHOT_PATH") else { + return Ok(()); + }; + let snapshot = + Snapshot::::checked_load(&path, OciTag::new("application")?)?; + let mut loaded = restore(snapshot)?; + assert_eq!(loaded.call_guest_function::("counter", ())?, 2); + let mut loaded = loaded + .unload_module()? + .load_module(get_wasm_module_path("rust_wasm_samples.aot")?)?; + assert_eq!(loaded.call_guest_function::("counter", ())?, 1); + std::fs::write(Path::new(&path).join("reader-ran"), [])?; + Ok(()) + } + + // Runs only in builds with `WIT_WORLD` set, through `just examples-components`. + #[cfg(component)] + mod component { + use super::*; + + mod bindings { + hyperlight_component_macro::host_bindgen!({ + path: "../tests/rust_guests/component_sample/wit/component-world.wasm", + }); + } + + use bindings::component_sample::example::{Adder, ExampleExports}; + + struct State; + + impl bindings::component_sample::example::Host for State { + fn r#print(&mut self, _message: String) {} + + fn r#host_function(&mut self, input: String) -> String { + format!("{input} and the host!") + } + } + + #[allow(refining_impl_trait)] + impl + bindings::component_sample::example::ExampleImports< + hyperlight_common::component::Negative, + > for State + { + type Host = State; + + fn r#host(&mut self) -> &mut Self { + self + } + } + + #[test] + fn persists_component_and_rebinds_callbacks() -> Result<()> { + let directory = tempfile::tempdir()?; + let module = get_wasm_module_path("component_sample.aot")?; + let mut proto = SandboxBuilder::new() + .with_guest_heap_size(200_000_000) + .with_guest_scratch_size(100 * 1024 * 1024) + .build()?; + let rt = bindings::register_host_functions(&mut proto, State)?; + let sb = proto.load_runtime()?.load_module(&module)?; + let mut wrapped = bindings::ExampleSandbox { sb, rt }; + assert_eq!(ExampleExports::adder(&mut wrapped).add(1, 2)?, 3); + let digest = wrapped + .sb + .snapshot()? + .save(directory.path(), &OciTag::new("component")?)?; + drop(wrapped); + + let snapshot = Snapshot::::checked_load(directory.path(), digest)?; + let mut builder = SandboxBuilder::from_snapshot(snapshot); + let rt = bindings::register_host_functions(&mut builder, State)?; + let mut wrapped = bindings::ExampleSandbox { + sb: builder.build()?, + rt, + }; + assert_eq!( + ExampleExports::adder(&mut wrapped).call_host("Restored".to_owned())?, + "Restored from component and the host!" + ); + + let sb = wrapped.sb.unload_module()?.load_module(&module)?; + let mut wrapped = bindings::ExampleSandbox { sb, rt: wrapped.rt }; + assert_eq!(ExampleExports::adder(&mut wrapped).add(4, 3)?, 7); + Ok(()) + } + } +} diff --git a/src/hyperlight_wasm/src/sandbox/wasm_sandbox.rs b/src/hyperlight_wasm/src/sandbox/wasm_sandbox.rs index 324f89fc..9146b4ea 100644 --- a/src/hyperlight_wasm/src/sandbox/wasm_sandbox.rs +++ b/src/hyperlight_wasm/src/sandbox/wasm_sandbox.rs @@ -19,13 +19,12 @@ use std::sync::Arc; #[cfg(target_os = "linux")] use hyperlight_host::mem::memory_region::{MemoryRegion, MemoryRegionFlags, MemoryRegionType}; -use hyperlight_host::sandbox::snapshot::Snapshot; +use hyperlight_host::sandbox::snapshot::Snapshot as CoreSnapshot; use hyperlight_host::{Result, Sandbox, new_error}; use super::loaded_wasm_sandbox::LoadedWasmSandbox; -use crate::sandbox::metrics::{ - METRIC_ACTIVE_WASM_SANDBOXES, METRIC_SANDBOX_LOADS, METRIC_TOTAL_WASM_SANDBOXES, -}; +use super::snapshot::Snapshot; +use crate::sandbox::metrics::{METRIC_ACTIVE_WASM_SANDBOXES, METRIC_TOTAL_WASM_SANDBOXES}; // All the logic around when to restore is nicely encapsulated here, // so that it would be harder for a `WasmSandbox` to end up in an @@ -47,7 +46,7 @@ mod backing_sandbox { Missing, } impl BackingSandbox { - pub(super) fn clean(&mut self, snapshot: Arc) -> Result<()> { + pub(super) fn clean(&mut self, snapshot: Arc) -> Result<()> { *self = match std::mem::replace(self, BackingSandbox::Missing) { BackingSandbox::Clean(x) => BackingSandbox::Clean(x), BackingSandbox::Loaded(_) => { @@ -56,7 +55,10 @@ mod backing_sandbox { )); } BackingSandbox::Dirty(mut x) => { - x.restore(snapshot)?; + if let Err(error) = x.restore(snapshot) { + *self = BackingSandbox::Dirty(x); + return Err(error); + } BackingSandbox::Clean(x) } BackingSandbox::Missing => { @@ -67,10 +69,14 @@ mod backing_sandbox { }; Ok(()) } - pub(super) fn load_via_restore(&mut self, snapshot: Arc) -> Result<()> { + pub(super) fn load_via_restore(&mut self, snapshot: Arc) -> Result<()> { *self = match std::mem::replace(self, BackingSandbox::Missing) { BackingSandbox::Clean(mut x) | BackingSandbox::Dirty(mut x) => { - x.restore(snapshot)?; + if let Err(error) = x.restore(snapshot) { + // A failed restore may have changed memory, so loading must restore again. + *self = BackingSandbox::Dirty(x); + return Err(error); + } BackingSandbox::Loaded(x) } BackingSandbox::Loaded(_) => { @@ -169,49 +175,52 @@ use backing_sandbox::*; /// function to load a Wasm module into memory. That function will return a /// `LoadedWasmSandbox` able to execute code in the loaded Wasm Module. pub struct WasmSandbox { - // inner is an Option as we need to take ownership of it - // We implement drop on the WasmSandbox 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 Sandbox Copy but that would introduce other issues) inner: BackingSandbox, - // Snapshot of state of an initial WasmSandbox (runtime loaded, but no guest module code loaded). - // Used for LoadedWasmSandbox to be able restore state back to WasmSandbox - snapshot: Option>, + // Runtime state, with no module loaded, that `load_module` starts from. + snapshot: Snapshot, } const MAPPED_BINARY_VA: u64 = 0x1_0000_0000u64; impl WasmSandbox { - /// Create a new WasmSandBox from a `Sandbox`. - /// This function should be used to create a new `WasmSandbox` from a ProtoWasmSandbox. - /// The difference between this function and creating a `WasmSandbox` directly is that - /// this function will increment the metrics for the number of `WasmSandbox`es in the system. + /// Snapshot a freshly initialized runtime and wrap it. pub(super) fn new(mut inner: Sandbox) -> Result { - let snapshot = inner.snapshot()?; - metrics::gauge!(METRIC_ACTIVE_WASM_SANDBOXES).increment(1); - metrics::counter!(METRIC_TOTAL_WASM_SANDBOXES).increment(1); - Ok(WasmSandbox { - inner: BackingSandbox::Clean(inner), - snapshot: Some(snapshot), - }) + let snapshot = Snapshot::runtime(inner.snapshot()?); + Ok(Self::clean(inner, snapshot)) + } + + /// Wrap a sandbox whose state matches `snapshot`. + pub(super) fn clean(inner: Sandbox, snapshot: Snapshot) -> Self { + Self::with_backing(BackingSandbox::Clean(inner), snapshot) } - /// Same as new, but doesn't take a new snapshot. Useful if `new` has already been called, - /// for example when creating a `WasmSandbox` from a `LoadedWasmSandbox`, since - /// the snapshot has already been created in that case. - /// Expects a snapshot of the state where wasm runtime is loaded, but no guest module code is loaded. - pub(super) fn new_from_loaded(loaded: Sandbox, snapshot: Arc) -> Result { + /// Wrap a sandbox that must be restored to `snapshot` before loading a module. + pub(super) fn dirty(inner: Sandbox, snapshot: Snapshot) -> Self { + Self::with_backing(BackingSandbox::Dirty(inner), snapshot) + } + + fn with_backing(inner: BackingSandbox, snapshot: Snapshot) -> Self { metrics::gauge!(METRIC_ACTIVE_WASM_SANDBOXES).increment(1); metrics::counter!(METRIC_TOTAL_WASM_SANDBOXES).increment(1); - Ok(WasmSandbox { - inner: BackingSandbox::Dirty(loaded), - snapshot: Some(snapshot), - }) + Self { inner, snapshot } } fn clean_inner(&mut self) -> Result<()> { - let snapshot = self.snapshot.as_ref().ok_or(new_error!( - "internal invariant violation: Snapshot is missing" - ))?; - self.inner.clean(snapshot.clone()) + self.inner.clean(self.snapshot.inner()) + } + + /// Return the runtime snapshot that `load_module` starts from. + /// + /// It excludes state left by an unloaded module. + pub fn snapshot(&self) -> Result> { + Ok(self.snapshot.clone()) + } + + /// Restore the sandbox to `snapshot`. Later module loads start from it. + pub fn restore(&mut self, snapshot: Snapshot) -> Result<()> { + self.inner.load_via_restore(snapshot.inner())?; + self.inner = BackingSandbox::Clean(self.inner.get_loaded()?); + self.snapshot = snapshot; + Ok(()) } /// Load a Wasm module at the given path into the sandbox and return a `LoadedWasmSandbox` @@ -235,10 +244,17 @@ impl WasmSandbox { self.finalize_module_load() } - /// Load a Wasm module by restoring a Hyperlight snapshot taken - /// from a `LoadedWasmSandbox`. - pub fn load_from_snapshot(mut self, snapshot: Arc) -> Result { - self.inner.load_via_restore(snapshot)?; + /// Load a Wasm module by restoring a snapshot taken from a `LoadedWasmSandbox`. + /// + /// A later `unload_module()` returns to the runtime state from before the + /// given snapshot's module was loaded. + pub fn load_from_snapshot( + mut self, + snapshot: Snapshot, + ) -> Result { + let runtime = snapshot.runtime_snapshot(); + self.inner.load_via_restore(snapshot.inner())?; + self.snapshot = runtime; self.finalize_module_load() } @@ -300,15 +316,8 @@ impl WasmSandbox { /// Helper function to finalize module loading and create LoadedWasmSandbox fn finalize_module_load(mut self) -> Result { - metrics::counter!(METRIC_SANDBOX_LOADS).increment(1); - let sandbox = self.inner.get_loaded()?; - - let snapshot = self.snapshot.take().ok_or(new_error!( - "internal invariant violation: Snapshot is missing" - ))?; - - LoadedWasmSandbox::new(sandbox, snapshot) + LoadedWasmSandbox::new(sandbox, self.snapshot.clone()) } } diff --git a/src/tests/rust_guests/rust_wasm_samples/src/lib.rs b/src/tests/rust_guests/rust_wasm_samples/src/lib.rs index 7a2e57e6..d66d6c63 100644 --- a/src/tests/rust_guests/rust_wasm_samples/src/lib.rs +++ b/src/tests/rust_guests/rust_wasm_samples/src/lib.rs @@ -70,3 +70,11 @@ pub extern "C" fn add(left: usize, right: usize) -> usize { pub extern "C" fn call_host_function(a: i32) -> i32 { hostfuncs::test_host_func(a) } + +#[no_mangle] +pub extern "C" fn counter() -> i32 { + use std::sync::atomic::{AtomicI32, Ordering}; + + static COUNT: AtomicI32 = AtomicI32::new(0); + COUNT.fetch_add(1, Ordering::Relaxed) + 1 +}