diff --git a/runtime/wasm-bpf-rs/src/runner.rs b/runtime/wasm-bpf-rs/src/runner.rs index 365f8f7..03a6e96 100644 --- a/runtime/wasm-bpf-rs/src/runner.rs +++ b/runtime/wasm-bpf-rs/src/runner.rs @@ -1,7 +1,7 @@ use std::sync::mpsc; use anyhow::{anyhow, Context}; -use wasmtime::{Engine, Linker, Module, Store, TypedFunc}; +use wasmtime::{Engine, IntoFunc, Linker, Module, Store, TypedFunc}; use wasmtime_wasi::WasiCtxBuilder; use crate::add_bind_function_with_module; @@ -39,6 +39,7 @@ pub struct WasmBpfModuleRunner { /// The linker which will be used pub linker: Linker, operation_tx: mpsc::Sender, + main_module: Module, } impl WasmBpfModuleRunner { @@ -88,14 +89,12 @@ impl WasmBpfModuleRunner { wrapper_poll::bpf_buffer_poll_wrapper, POLL_WRAPPER_FUNCTION_NAME )?; - linker - .module(&mut store, MAIN_MODULE_NAME, &main_module) - .with_context(|| anyhow!("Failed to link main module"))?; Ok(Self { engine, store, linker, operation_tx: tx, + main_module, }) } /// Consume this runner, return a handle to the wasm program, which can control the pause/resume/terminate of the program @@ -104,6 +103,10 @@ impl WasmBpfModuleRunner { pub fn into_engine_and_entry_func( mut self, ) -> anyhow::Result<(WasmProgramHandle, WasmBpfEntryFuncWrapper)> { + self.linker + .module(&mut self.store, MAIN_MODULE_NAME, &self.main_module) + .with_context(|| anyhow!("Failed to link main module"))?; + let func = self .linker .get(&mut self.store, MAIN_MODULE_NAME, "_start") @@ -119,4 +122,13 @@ impl WasmBpfModuleRunner { }, )) } + /// Register a custom host function. It has the similar signature as `wasmtime::linker::Linker::func_wrap` + pub fn register_host_function( + &mut self, + module: &str, + name: &str, + func: impl IntoFunc, + ) -> anyhow::Result<()> { + self.linker.func_wrap(module, name, func).map(|_| ()) + } } diff --git a/runtime/wasm-bpf-rs/src/tests/mod.rs b/runtime/wasm-bpf-rs/src/tests/mod.rs index c5f4ea7..a1f2e5d 100644 --- a/runtime/wasm-bpf-rs/src/tests/mod.rs +++ b/runtime/wasm-bpf-rs/src/tests/mod.rs @@ -1,5 +1,6 @@ use crate::handle::WasmProgramHandle; use crate::pipe::ReadableWritePipe; +use crate::state::CallerType; use super::*; use std::fs::File; @@ -154,3 +155,20 @@ fn test_pause_and_resume_wasm_program() { assert!(tick_count_3 - tick_count_2 >= 2); handle.as_mut().unwrap().terminate().unwrap(); } + +#[test] +fn test_custom_host_function() { + let module_binary = std::fs::read(get_test_file_path("custom_host_func.wasm")).unwrap(); + let args = vec!["test".to_string()]; + let mut runner = + WasmBpfModuleRunner::new(&module_binary[..], &args[..], Config::default()).unwrap(); + runner + .register_host_function("host_func_test", "plus_i32", host_func_plus_i32) + .unwrap(); + let (_, wrapper) = runner.into_engine_and_entry_func().unwrap(); + wrapper.run().unwrap(); +} + +fn host_func_plus_i32(_caller: CallerType, a: i32, b: i32) -> i32 { + a + b +} diff --git a/runtime/wasm-bpf-rs/tests/custom_host_func.wasm b/runtime/wasm-bpf-rs/tests/custom_host_func.wasm new file mode 100755 index 0000000..3165402 Binary files /dev/null and b/runtime/wasm-bpf-rs/tests/custom_host_func.wasm differ diff --git a/runtime/wasm-bpf-rs/tests/custom_host_func/.gitignore b/runtime/wasm-bpf-rs/tests/custom_host_func/.gitignore new file mode 100644 index 0000000..f00b29a --- /dev/null +++ b/runtime/wasm-bpf-rs/tests/custom_host_func/.gitignore @@ -0,0 +1 @@ +custom_host_func.wasm diff --git a/runtime/wasm-bpf-rs/tests/custom_host_func/Makefile b/runtime/wasm-bpf-rs/tests/custom_host_func/Makefile new file mode 100644 index 0000000..88352a6 --- /dev/null +++ b/runtime/wasm-bpf-rs/tests/custom_host_func/Makefile @@ -0,0 +1,6 @@ +WASI_CLANG = /opt/wasi-sdk/bin/clang +WASI_CFLAGS = -O2 --sysroot=/opt/wasi-sdk/share/wasi-sysroot -Wl,--allow-undefined,--export-table + +custom_host_func.wasm: custom_host_func.c + $(WASI_CLANG) $(WASI_CFLAGS) -o $@ $< + cp custom_host_func.wasm .. diff --git a/runtime/wasm-bpf-rs/tests/custom_host_func/README.md b/runtime/wasm-bpf-rs/tests/custom_host_func/README.md new file mode 100644 index 0000000..e3175f0 --- /dev/null +++ b/runtime/wasm-bpf-rs/tests/custom_host_func/README.md @@ -0,0 +1,3 @@ +# Test program for custom host func + +It invokes a host function to calculate `a+b`, and verifies the result. diff --git a/runtime/wasm-bpf-rs/tests/custom_host_func/custom_host_func.c b/runtime/wasm-bpf-rs/tests/custom_host_func/custom_host_func.c new file mode 100644 index 0000000..1814539 --- /dev/null +++ b/runtime/wasm-bpf-rs/tests/custom_host_func/custom_host_func.c @@ -0,0 +1,10 @@ +#include +#include +__attribute__((import_module("host_func_test"), import_name("plus_i32"))) +int32_t +plus_i32(int32_t a, int32_t b); +int main() { + int32_t c = plus_i32(0xABCD, 0x1234); + assert(c == 0xABCD + 0x1234); + return 0; +}