use crate::context::ExecutionContext;
use crate::error::WasmError;
use crate::executor::CommandHandler;
use crate::plugin::Plugin;
use crate::Result;
use std::collections::HashMap;
use std::path::Path;
use wasmtime::{Engine, Instance, Module, Store};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum WasmSerializationFormat {
#[default]
Yaml,
Json,
}
impl WasmSerializationFormat {
fn serialize(self, args: &HashMap<String, String>) -> Result<Vec<u8>> {
let bytes = match self {
Self::Yaml => serde_yaml::to_string(args)
.map_err(|e| WasmError::SerializationFailed(e.to_string()))?
.into_bytes(),
Self::Json => serde_json::to_vec(args)
.map_err(|e| WasmError::SerializationFailed(e.to_string()))?,
};
Ok(bytes)
}
}
const EXPORT_MEMORY: &str = "memory";
const EXPORT_ALLOC: &str = "dcli_alloc";
const EXPORT_DEALLOC: &str = "dcli_dealloc";
const EXPORT_LAST_ERROR: &str = "dcli_last_error_message";
pub struct WasmPlugin {
name: String,
version: String,
description: String,
engine: Engine,
module: Module,
function_map: HashMap<String, String>,
format: WasmSerializationFormat,
}
impl WasmPlugin {
pub fn load(path: &Path) -> Result<Self> {
let bytes = std::fs::read(path).map_err(|e| WasmError::LoadFailed {
path: path.to_path_buf(),
source: anyhow::Error::new(e),
suggestion: Some("Verify the path is correct and the file is readable.".to_string()),
})?;
Self::from_bytes(&bytes, path)
}
pub fn from_bytes(bytes: &[u8], origin: &Path) -> Result<Self> {
let engine = Engine::default();
let module = Module::new(&engine, bytes).map_err(|e| WasmError::LoadFailed {
path: origin.to_path_buf(),
source: e.into(),
suggestion: Some("Verify the file is a valid WASM binary or WAT module.".to_string()),
})?;
let module_label = origin.display().to_string();
Self::validate_mandatory_exports(&module, &module_label)?;
let default_name = origin
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("wasm-plugin")
.to_string();
Ok(Self {
name: default_name,
version: "0.0.0".to_string(),
description: "WASM plugin (no metadata provided)".to_string(),
engine,
module,
function_map: HashMap::new(),
format: WasmSerializationFormat::default(),
})
}
fn validate_mandatory_exports(module: &Module, module_label: &str) -> Result<()> {
let exported_names: Vec<&str> = module.exports().map(|e| e.name()).collect();
for required in [EXPORT_MEMORY, EXPORT_ALLOC, EXPORT_DEALLOC] {
if !exported_names.contains(&required) {
return Err(WasmError::missing_mandatory_export(required, module_label).into());
}
}
Ok(())
}
pub fn with_function_map(mut self, impl_name: &str, wasm_fn_name: &str) -> Self {
self.function_map
.insert(impl_name.to_string(), wasm_fn_name.to_string());
self
}
pub fn with_format(mut self, format: WasmSerializationFormat) -> Self {
self.format = format;
self
}
pub fn with_metadata(mut self, name: &str, version: &str, description: &str) -> Self {
self.name = name.to_string();
self.version = version.to_string();
self.description = description.to_string();
self
}
}
impl Plugin for WasmPlugin {
fn name(&self) -> &str {
&self.name
}
fn version(&self) -> &str {
&self.version
}
fn description(&self) -> &str {
&self.description
}
fn handlers(&self) -> Vec<(String, Box<dyn CommandHandler>)> {
self.function_map
.iter()
.map(|(impl_name, wasm_fn_name)| {
let handler: Box<dyn CommandHandler> = Box::new(WasmHandler {
engine: self.engine.clone(),
module: self.module.clone(),
module_label: self.name.clone(),
wasm_fn_name: wasm_fn_name.clone(),
format: self.format,
});
(impl_name.clone(), handler)
})
.collect()
}
}
struct WasmHandler {
engine: Engine,
module: Module,
module_label: String,
wasm_fn_name: String,
format: WasmSerializationFormat,
}
impl WasmHandler {
fn call_guest(&self, args: &HashMap<String, String>) -> Result<()> {
let mut store = Store::new(&self.engine, ());
let instance =
Instance::new(&mut store, &self.module, &[]).map_err(|e| WasmError::LoadFailed {
path: std::path::PathBuf::from(&self.module_label),
source: e.into(),
suggestion: Some("Failed to instantiate the WASM module.".to_string()),
})?;
let memory = instance
.get_memory(&mut store, EXPORT_MEMORY)
.ok_or_else(|| {
WasmError::missing_mandatory_export(EXPORT_MEMORY, &self.module_label)
})?;
let alloc = instance
.get_typed_func::<i32, i32>(&mut store, EXPORT_ALLOC)
.map_err(|_| WasmError::missing_mandatory_export(EXPORT_ALLOC, &self.module_label))?;
let dealloc = instance
.get_typed_func::<(i32, i32), ()>(&mut store, EXPORT_DEALLOC)
.map_err(|_| WasmError::missing_mandatory_export(EXPORT_DEALLOC, &self.module_label))?;
let business_fn = instance
.get_typed_func::<(i32, i32), i32>(&mut store, self.wasm_fn_name.as_str())
.map_err(|_| WasmError::FunctionNotFound {
function: self.wasm_fn_name.clone(),
module: self.module_label.clone(),
suggestion: Some(format!(
"Export `fn {}(ptr: i32, len: i32) -> i32` from the WASM module, \
or check the name passed to `with_function_map`.",
self.wasm_fn_name
)),
})?;
let payload = self.format.serialize(args)?;
let len = payload.len() as i32;
let ptr = alloc
.call(&mut store, len)
.map_err(|e| WasmError::MemoryAccessFailed {
reason: e.to_string(),
})?;
memory
.write(&mut store, ptr as usize, &payload)
.map_err(|e| WasmError::MemoryAccessFailed {
reason: e.to_string(),
})?;
let call_result = business_fn.call(&mut store, (ptr, len));
let dealloc_result = dealloc.call(&mut store, (ptr, len));
let code = call_result.map_err(|e| WasmError::MemoryAccessFailed {
reason: format!("business function trapped: {e}"),
})?;
if code == 0 {
dealloc_result.map_err(|e| WasmError::MemoryAccessFailed {
reason: format!("dcli_dealloc failed: {e}"),
})?;
return Ok(());
}
let message = self.read_last_error_message(&mut store, &instance);
Err(WasmError::GuestError { code, message }.into())
}
fn read_last_error_message(
&self,
store: &mut Store<()>,
instance: &Instance,
) -> Option<String> {
let last_error_fn = instance
.get_typed_func::<(), (i32, i32)>(&mut *store, EXPORT_LAST_ERROR)
.ok()?;
let (ptr, len) = last_error_fn.call(&mut *store, ()).ok()?;
if len <= 0 {
return None;
}
let memory = instance.get_memory(&mut *store, EXPORT_MEMORY)?;
let mut buf = vec![0u8; len as usize];
memory.read(&mut *store, ptr as usize, &mut buf).ok()?;
String::from_utf8(buf).ok()
}
}
impl CommandHandler for WasmHandler {
fn execute(
&self,
_ctx: &mut dyn ExecutionContext,
args: &HashMap<String, String>,
) -> Result<()> {
self.call_guest(args)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::any::Any;
use std::path::PathBuf;
#[derive(Default)]
struct TestContext;
impl ExecutionContext for TestContext {
fn as_any(&self) -> &dyn Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
}
const WAT_MINIMAL_OK: &str = r#"
(module
(memory (export "memory") 1)
(func (export "dcli_alloc") (param i32) (result i32)
i32.const 1024)
(func (export "dcli_dealloc") (param i32 i32))
(func (export "ok_handler") (param i32 i32) (result i32)
i32.const 0)
)
"#;
const WAT_MINIMAL_ERROR: &str = r#"
(module
(memory (export "memory") 1)
(func (export "dcli_alloc") (param i32) (result i32)
i32.const 1024)
(func (export "dcli_dealloc") (param i32 i32))
(func (export "err_handler") (param i32 i32) (result i32)
i32.const 1)
)
"#;
const WAT_WITH_ERROR_MESSAGE: &str = r#"
(module
(memory (export "memory") 1)
(data (i32.const 2048) "boom")
(func (export "dcli_alloc") (param i32) (result i32)
i32.const 1024)
(func (export "dcli_dealloc") (param i32 i32))
(func (export "err_handler") (param i32 i32) (result i32)
i32.const 1)
(func (export "dcli_last_error_message") (result i32 i32)
i32.const 2048
i32.const 4)
)
"#;
const WAT_MISSING_DEALLOC: &str = r#"
(module
(memory (export "memory") 1)
(func (export "dcli_alloc") (param i32) (result i32)
i32.const 1024)
(func (export "ok_handler") (param i32 i32) (result i32)
i32.const 0)
)
"#;
const WAT_MISSING_MEMORY: &str = r#"
(module
(func (export "dcli_alloc") (param i32) (result i32)
i32.const 1024)
(func (export "dcli_dealloc") (param i32 i32))
(func (export "ok_handler") (param i32 i32) (result i32)
i32.const 0)
)
"#;
fn load_wat(wat: &str) -> Result<WasmPlugin> {
WasmPlugin::from_bytes(wat.as_bytes(), &PathBuf::from("test.wat"))
}
#[test]
fn test_load_minimal_valid_module_succeeds() {
let plugin = load_wat(WAT_MINIMAL_OK);
assert!(plugin.is_ok());
}
#[test]
fn test_load_missing_dealloc_fails() {
let result = load_wat(WAT_MISSING_DEALLOC);
assert!(result.is_err());
}
#[test]
fn test_load_missing_memory_fails() {
let result = load_wat(WAT_MISSING_MEMORY);
assert!(result.is_err());
}
#[test]
fn test_load_invalid_bytes_fails() {
let result = WasmPlugin::from_bytes(b"not a wasm module", &PathBuf::from("bad.wasm"));
assert!(result.is_err());
}
#[test]
fn test_default_metadata_derived_from_filename() {
let plugin = WasmPlugin::from_bytes(
WAT_MINIMAL_OK.as_bytes(),
&PathBuf::from("plugins/greet.wat"),
)
.unwrap();
assert_eq!(plugin.name(), "greet");
}
#[test]
fn test_with_metadata_overrides_defaults() {
let plugin =
load_wat(WAT_MINIMAL_OK)
.unwrap()
.with_metadata("custom", "2.5.0", "A custom plugin");
assert_eq!(plugin.name(), "custom");
assert_eq!(plugin.version(), "2.5.0");
assert_eq!(plugin.description(), "A custom plugin");
}
#[test]
fn test_with_format_defaults_to_yaml() {
let plugin = load_wat(WAT_MINIMAL_OK).unwrap();
assert_eq!(plugin.format, WasmSerializationFormat::Yaml);
}
#[test]
fn test_with_format_can_be_set_to_json() {
let plugin = load_wat(WAT_MINIMAL_OK)
.unwrap()
.with_format(WasmSerializationFormat::Json);
assert_eq!(plugin.format, WasmSerializationFormat::Json);
}
#[test]
fn test_handlers_reflects_function_map() {
let plugin = load_wat(WAT_MINIMAL_OK)
.unwrap()
.with_function_map("my_command", "ok_handler");
let handlers = plugin.handlers();
assert_eq!(handlers.len(), 1);
assert_eq!(handlers[0].0, "my_command");
}
#[test]
fn test_handlers_empty_without_function_map() {
let plugin = load_wat(WAT_MINIMAL_OK).unwrap();
assert_eq!(plugin.handlers().len(), 0);
}
#[test]
fn test_execute_success_returns_ok() {
let plugin = load_wat(WAT_MINIMAL_OK)
.unwrap()
.with_function_map("cmd", "ok_handler");
let handlers = plugin.handlers();
let (_, handler) = &handlers[0];
let mut ctx = TestContext;
let args = HashMap::new();
assert!(handler.execute(&mut ctx, &args).is_ok());
}
#[test]
fn test_execute_with_args_serializes_without_error() {
let plugin = load_wat(WAT_MINIMAL_OK)
.unwrap()
.with_function_map("cmd", "ok_handler");
let handlers = plugin.handlers();
let (_, handler) = &handlers[0];
let mut ctx = TestContext;
let mut args = HashMap::new();
args.insert("name".to_string(), "World".to_string());
args.insert("count".to_string(), "3".to_string());
assert!(handler.execute(&mut ctx, &args).is_ok());
}
#[test]
fn test_execute_repeated_calls_do_not_exhaust_guest() {
let plugin = load_wat(WAT_MINIMAL_OK)
.unwrap()
.with_function_map("cmd", "ok_handler");
let handlers = plugin.handlers();
let (_, handler) = &handlers[0];
let mut ctx = TestContext;
for _ in 0..50 {
let args = HashMap::new();
assert!(handler.execute(&mut ctx, &args).is_ok());
}
}
#[test]
fn test_execute_guest_error_without_message() {
let plugin = load_wat(WAT_MINIMAL_ERROR)
.unwrap()
.with_function_map("cmd", "err_handler");
let handlers = plugin.handlers();
let (_, handler) = &handlers[0];
let mut ctx = TestContext;
let args = HashMap::new();
let result = handler.execute(&mut ctx, &args);
assert!(result.is_err());
match result.unwrap_err() {
crate::error::DynamicCliError::Wasm(WasmError::GuestError { code, message }) => {
assert_eq!(code, 1);
assert!(message.is_none());
}
other => panic!("unexpected error variant: {other:?}"),
}
}
#[test]
fn test_execute_guest_error_with_message() {
let plugin = load_wat(WAT_WITH_ERROR_MESSAGE)
.unwrap()
.with_function_map("cmd", "err_handler");
let handlers = plugin.handlers();
let (_, handler) = &handlers[0];
let mut ctx = TestContext;
let args = HashMap::new();
let result = handler.execute(&mut ctx, &args);
assert!(result.is_err());
match result.unwrap_err() {
crate::error::DynamicCliError::Wasm(WasmError::GuestError { code, message }) => {
assert_eq!(code, 1);
assert_eq!(message, Some("boom".to_string()));
}
other => panic!("unexpected error variant: {other:?}"),
}
}
#[test]
fn test_execute_unmapped_function_name_fails() {
let plugin = load_wat(WAT_MINIMAL_OK)
.unwrap()
.with_function_map("cmd", "does_not_exist");
let handlers = plugin.handlers();
let (_, handler) = &handlers[0];
let mut ctx = TestContext;
let args = HashMap::new();
assert!(handler.execute(&mut ctx, &args).is_err());
}
#[test]
fn test_wasm_plugin_is_send_sync() {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<WasmPlugin>();
}
}