wasmer middleware import injection MWE

lib.rs

use std::sync::Mutex;

use loupe::MemoryUsage;
use wasmer::{FunctionMiddleware, FunctionType, ModuleMiddleware, Type};
use wasmer_types::{FunctionIndex, ImportIndex};
use wasmer_vm::ModuleInfo;

#[derive(Debug, MemoryUsage)]
struct MyMiddleware {
    // This would then be passed to FunctionMiddleware things, which would use
    // this to inject calls of the imported function into the Wasm code and so on...
    fn_index: Mutex<Option<FunctionIndex>>,
}

impl MyMiddleware {
    fn new() -> Self {
        Self {
            fn_index: Mutex::new(None),
        }
    }
}

#[derive(Debug, MemoryUsage)]
struct MyFunctionMiddleware;

impl ModuleMiddleware for MyMiddleware {
    fn transform_module_info(&self, mod_info: &mut ModuleInfo) {
        let mut fn_index = self.fn_index.lock().unwrap();

        if fn_index.is_some() {
            panic!("MyModuleMiddleware::transform_module_info: Attempting to use the middleware from multiple modules.");
        }

        // Essentially we're trying to register an imported function in the
        // module here.

        let sig = mod_info.signatures.push(FunctionType::new([Type::I32], []));
        *fn_index = Some(mod_info.functions.push(sig));
        let import_index = mod_info.imports().len();
        mod_info.imports.insert(
            (
                "profiling".to_string(),
                "start_measurement".to_string(),
                import_index as u32,
            ),
            ImportIndex::Function(fn_index.unwrap()),
        );
        mod_info.num_imported_functions += 1;
    }

    fn generate_function_middleware(
        &self,
        _local_function_index: wasmer::LocalFunctionIndex,
    ) -> Box<dyn FunctionMiddleware> {
        Box::new(MyFunctionMiddleware)
    }
}

impl FunctionMiddleware for MyFunctionMiddleware {}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use wasmer::{CompilerConfig, Cranelift, Function, Instance, Module, Store, Universal, imports, wat2wasm};
    use wasmer_types::Value;

    use crate::MyMiddleware;

    const WAT: &[u8] = br#"
    (module
    (type $t0 (func (param i32) (result i32)))
    (func $add_one (export "add_one") (type $t0) (param $p0 i32) (result i32)
        get_local $p0
        i32.const 1
        i32.add)
    (func $multisub (export "multisub") (type $t0) (param $p0 i32) (result i32)
        get_local $p0
        i32.const 2
        i32.mul
        call $sub_one
        i32.const 1
        i32.sub)
    (func $sub_one (type $t0) (param $p0 i32) (result i32)
        get_local $p0
        i32.const 1
        i32.sub))
    "#;

    #[test]
    fn middleware_test() {
        let profiling = Arc::new(MyMiddleware::new());

        // Create the module with our middleware.
        let mut compiler_config = Cranelift::default();
        compiler_config.push_middleware(profiling.clone());
        let store = Store::new(&Universal::new(compiler_config).engine());
        let wasm = wat2wasm(WAT).unwrap();
        let module = Module::new(&store, wasm).unwrap();

        let imports = imports! {
            "profiling" => {
                "start_measurement" => Function::new_native(&store, |_: u32| println!("start measuring")),
            }
        };
        let instance = Instance::new(&module, &imports).unwrap();

        let add_one_fn = instance.exports.get_function("add_one").unwrap();
        let result = add_one_fn.call(&[Value::I32(42)]).unwrap();
        assert_eq!(result[0], Value::I32(43));
    }
}

Quick Cargo.toml for this:

[package]
name = "profiler-mwe"
version = "0.1.0"
edition = "2018"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
loupe = "0.1.3"
wasmer = "2.0.0"
wasmer-types = "2.0.0"
wasmer-vm = "2.0.0"

Output:

failures:

---- tests::middleware_test stdout ----
start measuring
thread 'tests::middleware_test' panicked at 'assertion failed: `(left == right)`
  left: `I32(0)`,
 right: `I32(43)`', src/lib.rs:110:9
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace