use anyhow::{Context as AnyhowContext, Result};
use wasmtime::{
component::{Component, Linker},
error::Context,
Store,
};
use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiCtxView, WasiView};
#[cfg(feature = "dynamic")]
use crate::dynamic::host_imports::HostImportRegistry;
use crate::Image;
pub trait HostStateImpl: WasiView + Send + 'static {
fn as_any_mut(&mut self) -> &mut dyn std::any::Any;
}
pub struct HostState {
wasi: WasiCtx,
table: ResourceTable,
}
impl HostState {
pub fn new() -> Result<Self> {
let wasi = WasiCtxBuilder::new().build();
let table = ResourceTable::new();
Ok(Self { wasi, table })
}
pub fn with_wasi<F>(f: F) -> Result<Self>
where
F: FnOnce(&mut WasiCtxBuilder) -> &mut WasiCtxBuilder,
{
let mut builder = WasiCtxBuilder::new();
f(&mut builder);
let wasi = builder.build();
let table = ResourceTable::new();
Ok(Self { wasi, table })
}
}
impl Default for HostState {
fn default() -> Self {
let wasi = WasiCtxBuilder::new().build();
let table = ResourceTable::new();
Self { wasi, table }
}
}
impl WasiView for HostState {
fn ctx(&mut self) -> WasiCtxView<'_> {
WasiCtxView {
ctx: &mut self.wasi,
table: &mut self.table,
}
}
}
impl HostStateImpl for HostState {
fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
self
}
}
pub struct GuestInstance {
inner: Box<dyn std::any::Any + Send + Sync>,
#[cfg(feature = "dynamic")]
dynamic_instance: Option<wasmtime::component::Instance>,
}
impl GuestInstance {
pub fn new<T: 'static + Send + Sync>(instance: T) -> Self {
Self {
inner: Box::new(instance),
#[cfg(feature = "dynamic")]
dynamic_instance: None,
}
}
#[cfg(feature = "dynamic")]
pub fn new_dynamic(instance: wasmtime::component::Instance) -> Self {
Self {
inner: Box::new(instance),
dynamic_instance: Some(instance),
}
}
pub fn downcast_ref<T: 'static>(&self) -> Option<&T> {
self.inner.downcast_ref::<T>()
}
pub fn downcast_mut<T: 'static>(&mut self) -> Option<&mut T> {
self.inner.downcast_mut::<T>()
}
#[cfg(feature = "dynamic")]
pub fn get_export_func<T: HostStateImpl>(
&self,
_store: &mut Store<T>,
_func_name: &str,
) -> Result<Option<wasmtime::component::Func>> {
if let Some(ref instance) = self.dynamic_instance {
Ok(instance.get_func(_store, _func_name))
} else {
Ok(None)
}
}
#[cfg(feature = "dynamic")]
pub(crate) fn get_dynamic_instance_ref(&self) -> Option<&wasmtime::component::Instance> {
self.dynamic_instance.as_ref()
}
}
pub struct GuestHandlerContext<'a, T: HostStateImpl> {
pub linker: &'a mut Linker<T>,
pub store: &'a mut Store<T>,
pub component: &'a Component,
}
impl<'a, T: HostStateImpl> GuestHandlerContext<'a, T> {
pub fn new(
linker: &'a mut Linker<T>,
store: &'a mut Store<T>,
component: &'a Component,
) -> Self {
Self {
linker,
store,
component,
}
}
}
pub struct ContainerBuilder<T: HostStateImpl> {
image: Image,
host_state: T,
fuel_limit: Option<u64>,
epoch_deadline: Option<u64>,
#[allow(clippy::type_complexity)]
host_linker_init: Option<Box<dyn FnOnce(&mut Linker<T>) -> Result<(), anyhow::Error> + Send>>,
#[allow(clippy::type_complexity)]
guest_initializer: Option<
Box<
dyn for<'a> FnOnce(GuestHandlerContext<'a, T>) -> Result<GuestInstance, anyhow::Error>
+ Send,
>,
>,
}
impl<T: HostStateImpl> ContainerBuilder<T>
where
T: Default,
{
pub fn new(image: Image) -> Self {
Self {
image,
host_state: T::default(),
fuel_limit: None,
epoch_deadline: None,
host_linker_init: None,
guest_initializer: None,
}
}
}
impl<T: HostStateImpl> ContainerBuilder<T> {
pub fn with_host_state(mut self, state: T) -> Self {
self.host_state = state;
self
}
pub fn with_fuel_limit(mut self, limit: u64) -> Self {
self.fuel_limit = Some(limit);
self
}
pub fn with_epoch_deadline(mut self, deadline: u64) -> Self {
self.epoch_deadline = Some(deadline);
self
}
pub fn with_host_linker<F>(mut self, func: F) -> Self
where
F: FnOnce(&mut Linker<T>) -> Result<(), anyhow::Error> + Send + 'static,
{
self.host_linker_init = Some(Box::new(func));
self
}
pub fn with_guest_initializer<F>(mut self, f: F) -> Self
where
F: for<'a> FnOnce(GuestHandlerContext<'a, T>) -> Result<GuestInstance, anyhow::Error>
+ Send
+ 'static,
{
self.guest_initializer = Some(Box::new(f));
self
}
pub fn build(self) -> Result<Container<T>> {
let mut store = Store::new(self.image.engine(), self.host_state);
if let Some(fuel) = self.fuel_limit {
store
.set_fuel(fuel)
.context("Failed to set fuel limit (ensure Image was created with consume_fuel(true) in Config)")?;
}
if let Some(deadline) = self.epoch_deadline {
store.set_epoch_deadline(deadline);
}
let mut linker = Linker::new(self.image.engine());
wasmtime_wasi::p2::add_to_linker_sync(&mut linker)
.context("Failed to add WASI to linker")?;
if let Some(linker_init) = self.host_linker_init {
linker_init(&mut linker).context("Failed to configure host linker")?;
}
let component = self.image.component().clone();
let guest_instance = if let Some(initializer) = self.guest_initializer {
let ctx = GuestHandlerContext::new(&mut linker, &mut store, &component);
initializer(ctx)?
} else {
return Err(anyhow::anyhow!(
"Guest initializer is required. Use with_guest_initializer() to set it."
));
};
#[cfg(feature = "dynamic")]
let dynamic_instance = guest_instance.get_dynamic_instance_ref().cloned();
Ok(Container {
store,
guest: guest_instance,
state: ContainerState::Created,
#[cfg(feature = "dynamic")]
dynamic_instance,
#[cfg(feature = "dynamic")]
host_imports: None,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ContainerState {
Created,
Running,
Stopped,
Error(String),
}
pub struct Container<T: HostStateImpl = HostState> {
store: Store<T>,
guest: GuestInstance,
state: ContainerState,
#[cfg(feature = "dynamic")]
dynamic_instance: Option<wasmtime::component::Instance>,
#[cfg(feature = "dynamic")]
host_imports: Option<HostImportRegistry>,
}
impl Container {
pub fn builder(image: Image) -> ContainerBuilder<HostState> {
ContainerBuilder::new(image)
}
}
impl<T: HostStateImpl> Container<T> {
pub fn store_mut(&mut self) -> &mut Store<T> {
&mut self.store
}
pub fn store(&self) -> &Store<T> {
&self.store
}
pub fn guest(&self) -> &GuestInstance {
&self.guest
}
pub fn guest_mut(&mut self) -> &mut GuestInstance {
&mut self.guest
}
pub fn host_state_mut(&mut self) -> &mut T {
self.store.data_mut()
}
pub fn host_state(&self) -> &T {
self.store.data()
}
pub fn state(&self) -> &ContainerState {
&self.state
}
pub fn stop(&mut self) {
self.state = ContainerState::Stopped;
}
#[deprecated(
since = "0.6.0",
note = "use call_guest_raw_desc() with RON format instead"
)]
pub fn call_guest_json(&mut self, function_name: &str, _json_payload: &str) -> Result<String> {
anyhow::bail!(
"JSON invocation is not supported. Use call_guest_raw_desc() instead with RON format for better Rust type compatibility. Function: {}",
function_name,
)
}
#[cfg(feature = "dynamic")]
pub fn call_guest_raw_desc(
&mut self,
function_name: &str,
raw_desc_payload: &str,
) -> Result<String> {
match &self.state {
ContainerState::Stopped => {
anyhow::bail!("Container is stopped");
}
ContainerState::Error(e) => {
anyhow::bail!("Container is in error state: {}", e);
}
_ => {}
}
let result = self.call_guest_raw_desc_inner(function_name, raw_desc_payload);
match &result {
Ok(_) => self.state = ContainerState::Running,
Err(e) => {
let is_wasm_fatal = e.downcast_ref::<wasmtime::Trap>().is_some();
if is_wasm_fatal {
self.state = ContainerState::Error(e.to_string());
}
}
}
result
}
#[cfg(feature = "dynamic")]
fn call_guest_raw_desc_inner(
&mut self,
function_name: &str,
raw_desc_payload: &str,
) -> Result<String> {
use wasmtime::component::Val;
use crate::dynamic::{ron_to_val, val_to_ron};
let instance = self
.dynamic_instance
.as_ref()
.ok_or_else(|| anyhow::anyhow!("Dynamic instance not available"))?;
let func = instance
.get_func(&mut self.store, function_name)
.ok_or_else(|| anyhow::anyhow!("Export function not found: {}", function_name))?;
let func_ty = func.ty(&self.store);
let param_types: Vec<_> = func_ty.params().collect();
let result_types: Vec<_> = func_ty.results().collect();
let mut args = Vec::new();
if param_types.len() == 1 {
let param_type = ¶m_types[0].1;
args.push(ron_to_val(raw_desc_payload, param_type)?);
} else {
let ron_array = if raw_desc_payload.trim().starts_with('[') {
raw_desc_payload.to_string()
} else if raw_desc_payload.trim().starts_with('(') {
raw_desc_payload.to_string()
} else {
format!("[{}]", raw_desc_payload)
};
use ron::Value as RonValue;
let ron_value: RonValue = ron::from_str(&ron_array)?;
if let RonValue::Seq(items) = ron_value {
if items.len() != param_types.len() {
anyhow::bail!(
"Parameter count mismatch: expected {}, got {}",
param_types.len(),
items.len()
);
}
for (ron_val, (_param_name, param_type)) in
items.into_iter().zip(param_types.iter())
{
args.push(ron_value_to_val(ron_val, param_type)?);
}
} else {
anyhow::bail!(
"Invalid raw descriptor payload for function with multiple parameters"
);
}
}
let mut results = vec![Val::Bool(false); result_types.len()];
func.call(&mut self.store, &args, &mut results)
.context("Function call failed")?;
let output_ron: Result<Vec<_>> = results.iter().map(val_to_ron).collect();
let output_ron = output_ron.context("Failed to convert result to RON")?;
let output = match output_ron.len() {
0 => "()".to_string(),
1 => output_ron[0].clone(),
_ => format!("({})", output_ron.join(", ")),
};
Ok(output)
}
#[cfg(feature = "dynamic")]
pub fn call_guest_binary(
&mut self,
function_name: &str,
args: &[wasmtime::component::Val],
) -> Result<Vec<wasmtime::component::Val>> {
match &self.state {
ContainerState::Stopped => {
anyhow::bail!("Container is stopped");
}
ContainerState::Error(e) => {
anyhow::bail!("Container is in error state: {}", e);
}
_ => {}
}
let result = self.call_guest_binary_inner(function_name, args);
match &result {
Ok(_) => self.state = ContainerState::Running,
Err(e) => {
let msg = e.to_string();
let is_wasm_fatal = msg.contains("trap")
|| msg.contains("out of memory")
|| msg.contains("fuel")
|| e.downcast_ref::<wasmtime::Error>().is_some()
|| e.downcast_ref::<wasmtime::Trap>().is_some();
if is_wasm_fatal {
self.state = ContainerState::Error(msg);
}
}
}
result
}
#[cfg(feature = "dynamic")]
fn call_guest_binary_inner(
&mut self,
function_name: &str,
args: &[wasmtime::component::Val],
) -> Result<Vec<wasmtime::component::Val>> {
use wasmtime::component::Val;
let instance = self
.dynamic_instance
.as_ref()
.ok_or_else(|| anyhow::anyhow!("Dynamic instance not available"))?;
let func = instance
.get_func(&mut self.store, function_name)
.ok_or_else(|| anyhow::anyhow!("Export function not found: {}", function_name))?;
let func_ty = func.ty(&self.store);
let num_results = func_ty.results().count();
let mut results = vec![Val::Bool(false); num_results];
func.call(&mut self.store, args, &mut results)
.context("Function call failed")?;
Ok(results)
}
#[cfg(feature = "dynamic")]
pub fn with_host_import_registry(&mut self, registry: HostImportRegistry) {
self.host_imports = Some(registry);
}
#[cfg(feature = "dynamic")]
pub fn host_imports_mut(&mut self) -> Option<&mut HostImportRegistry> {
self.host_imports.as_mut()
}
#[cfg(feature = "dynamic")]
pub fn call_host_import_raw_desc(
&mut self,
function_name: &str,
raw_desc_payload: &str,
) -> Result<String> {
use crate::dynamic::{ron_to_val, val_to_ron};
let registry = self
.host_imports
.as_ref()
.ok_or_else(|| anyhow::anyhow!("Host import registry not initialized"))?;
let (params, _results) = registry
.get_signature(function_name)
.ok_or_else(|| anyhow::anyhow!("Host import not found: {}", function_name))?;
let mut args = Vec::new();
if params.len() == 1 {
args.push(ron_to_val(raw_desc_payload, ¶ms[0])?);
} else {
let payload = raw_desc_payload.trim();
let ron_array = if payload.starts_with('[') || payload.starts_with('(') {
raw_desc_payload.to_string()
} else {
format!("[{}]", raw_desc_payload)
};
use ron::Value as RonValue;
let ron_value: RonValue = ron::from_str(&ron_array)?;
if let RonValue::Seq(items) = ron_value {
if items.len() != params.len() {
anyhow::bail!(
"Parameter count mismatch: expected {}, got {}",
params.len(),
items.len()
);
}
for (ron_val, param_type) in items.into_iter().zip(params.iter()) {
args.push(ron_value_to_val(ron_val, param_type)?);
}
} else {
anyhow::bail!(
"Invalid raw descriptor payload for function with multiple parameters"
);
}
}
let result_vals = registry.call(function_name, &args)?;
let output_ron: Result<Vec<_>> = result_vals.iter().map(val_to_ron).collect();
let output_ron = output_ron.context("Failed to convert result to RON")?;
let output = match output_ron.len() {
0 => "()".to_string(),
1 => output_ron[0].clone(),
_ => format!("({})", output_ron.join(", ")),
};
Ok(output)
}
#[cfg(feature = "dynamic")]
pub fn list_guest_exports(&mut self) -> Result<Vec<ExportInfo>> {
if let Some(instance) = &self.dynamic_instance {
let mut exports = Vec::new();
let potential_exports = vec![
"init",
"process",
"getname",
"getversion",
"getfeatures",
"shutdown",
"notify",
"add",
"sub",
"mul",
"div",
"to-upper",
"to-lower",
"reverse",
"length",
"process-numbers",
"transform",
];
for export_name in potential_exports {
if let Some(func) = instance.get_func(&mut self.store, export_name) {
let func_ty = func.ty(&self.store);
let mut params = Vec::new();
let mut results = Vec::new();
for (param_name, param_type) in func_ty.params() {
params.push((param_name.to_string(), param_type.clone()));
}
for result_type in func_ty.results() {
results.push(result_type.clone());
}
exports.push(ExportInfo {
name: export_name.to_string(),
params,
results,
});
}
}
Ok(exports)
} else {
Ok(Vec::new())
}
}
#[cfg(feature = "dynamic")]
pub fn list_host_imports(&self) -> Result<Vec<ImportInfo>> {
if let Some(ref registry) = self.host_imports {
let imports = registry.list_imports();
imports
.into_iter()
.map(|name| {
let (params, results) =
AnyhowContext::with_context(registry.get_signature(name), || {
format!("missing signature for import: {name}")
})?;
Ok(ImportInfo {
name: name.to_string(),
params,
results,
})
})
.collect()
} else {
Ok(Vec::new())
}
}
}
#[derive(Debug, Clone)]
pub struct ExportInfo {
pub name: String,
pub params: Vec<(String, wasmtime::component::Type)>,
pub results: Vec<wasmtime::component::Type>,
}
#[derive(Debug, Clone)]
pub struct ImportInfo {
pub name: String,
pub params: Vec<wasmtime::component::Type>,
pub results: Vec<wasmtime::component::Type>,
}
#[cfg(feature = "dynamic")]
fn ron_value_to_val(
ron_value: ron::Value,
target_type: &wasmtime::component::Type,
) -> Result<wasmtime::component::Val> {
use crate::dynamic::ron_to_val;
let ron_str = ron::to_string(&ron_value)?;
ron_to_val(&ron_str, target_type)
}
impl<T: HostStateImpl> std::fmt::Debug for Container<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Container")
.field("state", &self.state)
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_host_state_new_succeeds() {
let state = HostState::new();
assert!(state.is_ok(), "HostState::new() should succeed");
}
#[test]
fn test_host_state_new_does_not_inherit_stdio_or_network() {
let state = HostState::new().expect("should succeed");
let _ = state;
}
#[test]
fn test_host_state_default_does_not_panic() {
let state = HostState::default();
let _ = state;
}
#[test]
fn test_host_state_new_equals_default() {
let new_state = HostState::new().expect("new should succeed");
let default_state = HostState::default();
let _ = (new_state, default_state);
}
#[test]
fn test_guest_instance_new_and_downcast_ref() {
let instance: GuestInstance = GuestInstance::new(42i32);
let val = instance.downcast_ref::<i32>();
assert!(val.is_some());
assert_eq!(*val.unwrap(), 42);
let wrong = instance.downcast_ref::<String>();
assert!(wrong.is_none());
}
#[test]
fn test_guest_instance_downcast_mut() {
let mut instance: GuestInstance = GuestInstance::new(vec![1, 2, 3]);
{
let v = instance.downcast_mut::<Vec<i32>>();
assert!(v.is_some());
v.unwrap().push(4);
}
let v = instance.downcast_ref::<Vec<i32>>();
assert_eq!(v.unwrap(), &vec![1, 2, 3, 4]);
}
#[test]
fn test_guest_instance_with_string_type() {
let instance: GuestInstance = GuestInstance::new(String::from("hello"));
let val = instance.downcast_ref::<String>();
assert_eq!(val.unwrap(), "hello");
let mut instance = instance;
instance
.downcast_mut::<String>()
.unwrap()
.push_str(" world");
assert_eq!(instance.downcast_ref::<String>().unwrap(), "hello world");
}
#[test]
fn test_host_state_with_wasi_custom() {
let state = HostState::with_wasi(|builder| {
builder.arg("test").env("KEY", "VALUE");
builder
});
assert!(state.is_ok());
}
#[test]
fn test_container_state_transitions() {
let created = ContainerState::Created;
let running = ContainerState::Running;
let stopped = ContainerState::Stopped;
let error = ContainerState::Error("trap".to_string());
assert_eq!(created, ContainerState::Created);
assert_eq!(running, ContainerState::Running);
assert_eq!(stopped, ContainerState::Stopped);
assert!(matches!(error, ContainerState::Error(msg) if msg == "trap"));
}
#[test]
fn test_container_state_equality() {
assert_eq!(ContainerState::Created, ContainerState::Created);
assert_ne!(ContainerState::Created, ContainerState::Running);
assert_ne!(ContainerState::Running, ContainerState::Stopped);
assert_ne!(
ContainerState::Error("a".into()),
ContainerState::Error("b".into())
);
}
#[test]
fn test_container_state_clone_debug() {
let state = ContainerState::Error("oom".to_string());
let cloned = state.clone();
assert_eq!(state, cloned);
let debug = format!("{:?}", state);
assert!(debug.contains("oom"));
}
#[test]
#[cfg(feature = "dynamic")]
fn test_call_guest_raw_desc_rejects_invalid_payload() {
use crate::Image;
let wasm = bytes::Bytes::from_static(b"\x00asm\x01\x00\x00\x00");
let img = match Image::new(wasm) {
Ok(img) => img,
Err(_) => return,
};
let mut container = Container::builder(img)
.with_guest_initializer(|_ctx| Ok(GuestInstance::new(())))
.build()
.expect("build should succeed");
let sensitive = r#"("SECRET_API_KEY_12345",)"#;
let result = container.call_guest_raw_desc("test_fn", sensitive);
assert!(result.is_err());
let err_msg = format!("{}", result.unwrap_err());
assert!(
!err_msg.contains("SECRET_API_KEY_12345"),
"Error message should not contain the payload, but got: {}",
err_msg
);
}
}