Add support for user-registered host functions (#89)

* Add support for user-registered host functions

* format the code

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<AppState>,
     operation_tx: mpsc::Sender<ProgramOperation>,
+    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<Params, Args>(
+        &mut self,
+        module: &str,
+        name: &str,
+        func: impl IntoFunc<AppState, Params, Args>,
+    ) -> anyhow::Result<()> {
+        self.linker.func_wrap(module, name, func).map(|_| ())
+    }
 }

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
+}

runtime/wasm-bpf-rs/tests/custom_host_func/.gitignore

@@ -0,0 +1 @@
+custom_host_func.wasm

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 ..

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.

runtime/wasm-bpf-rs/tests/custom_host_func/custom_host_func.c

@@ -0,0 +1,10 @@
+#include <inttypes.h>
+#include <assert.h>
+__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;
+}