From 0a6a083f96777c745da7bf91cc4bbc8d899fcc8d Mon Sep 17 00:00:00 2001 From: Emil Taylor Bye Date: Thu, 1 Aug 2024 09:53:12 +0200 Subject: [PATCH] Add `Store::call_hook` API (#1144) * Start implementing `Store::call_hook` (#1083) * Add the `Store::call_hook` function, which stores a call hook in `Store` * Create tests to verify call hook behavior * Let call hooks return any `wasmi::Error` * Changed signature of call hooks to return a `wasmi::Error` instead of a `TrapCode`. * Use `assert_eq!` instead of `assert!` for tests, cosmetic change * Invoke call hooks on function calls * An `invoke_call_hook` function in Store to avoid problems with the borrow checker - we need a reference to the underlying data and the stored call hook * Invoke call hook on host -> wasm and wasm -> host calls * Clarified documentation on error propagation for call hooks * Fixed attribute placement and clippy errors * Placed attributes after doc comments * Added hint to allow complex type for `generate_error_after_n_calls` * Fixed clippy warnings on implicit conversion of ! to () in call_hook tests * Help the compiler optimize function calls for the no call hook scenario * Inline `Store::invoke_call_hook` to avoid extra function call * Add a `Store::invoke_call_hook_impl` function that is `#[cold]` to hint that the compiler should optimize for the scenario that there are no call hooks * Update crates/wasmi/src/store.rs Co-authored-by: Robin Freyler --------- Co-authored-by: Robin Freyler --- .../wasmi/src/engine/executor/instrs/call.rs | 9 +- crates/wasmi/src/engine/executor/mod.rs | 3 + crates/wasmi/src/lib.rs | 2 +- crates/wasmi/src/store.rs | 81 ++++++ crates/wasmi/tests/e2e/v1/call_hook.rs | 262 ++++++++++++++++++ crates/wasmi/tests/e2e/v1/mod.rs | 1 + 6 files changed, 356 insertions(+), 2 deletions(-) create mode 100644 crates/wasmi/tests/e2e/v1/call_hook.rs diff --git a/crates/wasmi/src/engine/executor/instrs/call.rs b/crates/wasmi/src/engine/executor/instrs/call.rs index 96ada021f6..2ac1c09140 100644 --- a/crates/wasmi/src/engine/executor/instrs/call.rs +++ b/crates/wasmi/src/engine/executor/instrs/call.rs @@ -18,6 +18,7 @@ use crate::{ }, func::{FuncEntity, HostFuncEntity}, store::StoreInner, + CallHook, Error, Func, FuncRef, @@ -469,7 +470,13 @@ impl<'engine> Executor<'engine> { Ok(()) } FuncEntity::Host(host_func) => { - self.execute_host_func::(store, results, func, *host_func) + let host_func = *host_func; + + store.invoke_call_hook(CallHook::CallingHost)?; + self.execute_host_func::(store, results, func, host_func)?; + store.invoke_call_hook(CallHook::ReturningFromHost)?; + + Ok(()) } } } diff --git a/crates/wasmi/src/engine/executor/mod.rs b/crates/wasmi/src/engine/executor/mod.rs index 24e93f09a0..054a292c36 100644 --- a/crates/wasmi/src/engine/executor/mod.rs +++ b/crates/wasmi/src/engine/executor/mod.rs @@ -14,6 +14,7 @@ use crate::{ ResumableInvocation, }, func::HostFuncEntity, + CallHook, Error, Func, FuncEntity, @@ -219,7 +220,9 @@ impl<'engine> EngineExecutor<'engine> { ), Some(instance), )?; + store.invoke_call_hook(CallHook::CallingWasm)?; self.execute_func(store)?; + store.invoke_call_hook(CallHook::ReturningFromWasm)?; } FuncEntity::Host(host_func) => { // The host function signature is required for properly diff --git a/crates/wasmi/src/lib.rs b/crates/wasmi/src/lib.rs index ebefb6bf03..70c6a794e9 100644 --- a/crates/wasmi/src/lib.rs +++ b/crates/wasmi/src/lib.rs @@ -169,7 +169,7 @@ pub use self::{ ModuleImportsIter, Read, }, - store::{AsContext, AsContextMut, Store, StoreContext, StoreContextMut}, + store::{AsContext, AsContextMut, CallHook, Store, StoreContext, StoreContextMut}, table::{Table, TableType}, value::Val, }; diff --git a/crates/wasmi/src/store.rs b/crates/wasmi/src/store.rs index 44d62d0d66..b52d759826 100644 --- a/crates/wasmi/src/store.rs +++ b/crates/wasmi/src/store.rs @@ -103,6 +103,19 @@ impl Debug for ResourceLimiterQuery { } } +/// A wrapper used to store hooks added with [`Store::call_hook`], containing a +/// boxed `FnMut(&mut T, CallHook) -> Result<(), Error>`. +/// +/// This wrapper exists to provide a `Debug` impl so that `#[derive(Debug)]` +/// works for [`Store`]. +#[allow(clippy::type_complexity)] +struct CallHookWrapper(Box Result<(), Error> + Send + Sync>); +impl Debug for CallHookWrapper { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "CallHook(...)") + } +} + /// The store that owns all data associated to Wasm modules. #[derive(Debug)] pub struct Store { @@ -119,6 +132,10 @@ pub struct Store { data: T, /// User provided hook to retrieve a [`ResourceLimiter`]. limiter: Option>, + /// User provided callback called when a host calls a WebAssembly function + /// or a WebAssembly function calls a host function, or these functions + /// return. + call_hook: Option>, } /// The inner store that owns all data not associated to the host state. @@ -166,6 +183,20 @@ fn test_store_is_send_sync() { }; } +/// Argument to the callback set by [`Store::call_hook`] to indicate why the +/// callback was invoked. +#[derive(Debug)] +pub enum CallHook { + /// Indicates that a WebAssembly function is being called from the host. + CallingWasm, + /// Indicates that a WebAssembly function called from the host is returning. + ReturningFromWasm, + /// Indicates that a host function is being called from a WebAssembly function. + CallingHost, + /// Indicates that a host function called from a WebAssembly function is returning. + ReturningFromHost, +} + /// An error that may be encountered when operating on the [`Store`]. #[derive(Debug, Clone)] pub enum FuelError { @@ -821,6 +852,7 @@ where trampolines: Arena::new(), data: T::default(), limiter: None, + call_hook: None, } } } @@ -833,6 +865,7 @@ impl Store { trampolines: Arena::new(), data, limiter: None, + call_hook: None, } } @@ -978,6 +1011,54 @@ impl Store { .get(entity_index) .unwrap_or_else(|| panic!("failed to resolve stored host function: {entity_index:?}")) } + + /// Sets a callback function that is executed whenever a WebAssembly + /// function is called from the host or a host function is called from + /// WebAssembly, or these functions return. + /// + /// The function is passed a `&mut T` to the underlying store, and a + /// [`CallHook`]. [`CallHook`] can be used to find out what kind of function + /// is being called or returned from. + /// + /// The callback can either return `Ok(())` or an `Err` with an + /// [`Error`]. If an error is returned, it is returned to the host + /// caller. If there are nested calls, only the most recent host caller + /// receives the error and it is not propagated further automatically. The + /// hook may be invoked again as new functions are called and returned from. + pub fn call_hook( + &mut self, + hook: impl FnMut(&mut T, CallHook) -> Result<(), Error> + Send + Sync + 'static, + ) { + self.call_hook = Some(CallHookWrapper(Box::new(hook))); + } + + /// Executes the callback set by [`Store::call_hook`] if any has been set. + /// + /// # Note + /// + /// - Returns the value returned by the call hook. + /// - Returns `Ok(())` if no call hook exists. + #[inline] + pub(crate) fn invoke_call_hook(&mut self, call_type: CallHook) -> Result<(), Error> { + match self.call_hook.as_mut() { + None => Ok(()), + Some(call_hook) => Self::invoke_call_hook_impl(&mut self.data, call_type, call_hook), + } + } + + /// Utility function to invoke the [`Store::call_hook`] that is asserted to + /// be available in this case. + /// + /// This is kept as a separate `#[cold]` function to help the compiler speed + /// up the code path without any call hooks. + #[cold] + fn invoke_call_hook_impl( + data: &mut T, + call_type: CallHook, + call_hook: &mut CallHookWrapper, + ) -> Result<(), Error> { + call_hook.0(data, call_type) + } } /// A trait used to get shared access to a [`Store`] in Wasmi. diff --git a/crates/wasmi/tests/e2e/v1/call_hook.rs b/crates/wasmi/tests/e2e/v1/call_hook.rs new file mode 100644 index 0000000000..cc32ba5c1c --- /dev/null +++ b/crates/wasmi/tests/e2e/v1/call_hook.rs @@ -0,0 +1,262 @@ +//! Tests to check if `Store::call_hook` works as intended. + +use wasmi::{ + core::TrapCode, + AsContext, + AsContextMut, + CallHook, + Caller, + Error, + Extern, + Func, + Linker, + Module, + Store, +}; + +/// Number of times different callback events have fired. +#[derive(Default)] +struct CallHookTestState { + calling_wasm: u32, + returning_from_wasm: u32, + calling_host: u32, + returning_from_host: u32, + erroneous_callback_invocation: bool, +} + +fn test_setup() -> (Store, Linker) { + let store = Store::default(); + let linker = >::new(store.engine()); + (store, linker) +} + +/// Prepares the test WAT and executes it. The wat defines two functions, +/// `wasm_fn_a` and `wasm_fn_b` and two imports, `host_fn_a` and `host_fn_b`. +/// `wasm_fn_a` calls `host_fn_a`, and `wasm_fn_b` calls `host_fn_b`. +/// None of the functions accept any arguments or return any value. +fn execute_wasm_fn_a( + mut store: &mut Store, + linker: &mut Linker, +) -> Result<(), Error> { + const TEST_WAT: &str = r#" + (module + (import "env" "host_fn_a" (func $host_fn_a)) + (import "env" "host_fn_b" (func $host_fn_b)) + (func (export "wasm_fn_a") + (call $host_fn_a) + ) + (func (export "wasm_fn_b") + (call $host_fn_b) + ) + ) + "#; + + let wasm = wat::parse_str(TEST_WAT).unwrap(); + let module = Module::new(store.engine(), &wasm).unwrap(); + let instance = linker + .instantiate(&mut store, &module) + .unwrap() + .start(store.as_context_mut()) + .unwrap(); + let wasm_fn = instance + .get_export(store.as_context(), "wasm_fn_a") + .and_then(Extern::into_func) + .unwrap() + .typed::<(), ()>(&store) + .unwrap(); + + wasm_fn.call(store.as_context_mut(), ()) +} + +#[test] +fn call_hooks_get_called() { + let (mut store, mut linker) = test_setup(); + + store.call_hook( + |data: &mut CallHookTestState, hook_type: CallHook| -> Result<(), Error> { + match hook_type { + CallHook::CallingWasm => data.calling_wasm += 1, + CallHook::ReturningFromWasm => data.returning_from_wasm += 1, + CallHook::CallingHost => data.calling_host += 1, + CallHook::ReturningFromHost => data.returning_from_host += 1, + }; + + Ok(()) + }, + ); + + let host_fn_a = Func::wrap(&mut store, |mut caller: Caller| { + // Call wasm_fn_a + // Call host_fn_a + assert_eq!(caller.data().calling_wasm, 1); + assert_eq!(caller.data().returning_from_wasm, 0); + assert_eq!(caller.data().calling_host, 1); + assert_eq!(caller.data().returning_from_host, 0); + + caller + .get_export("wasm_fn_b") + .and_then(Extern::into_func) + .unwrap() + .typed::<(), ()>(&caller) + .unwrap() + .call(&mut caller, ()) + .unwrap(); + + // Call wasm_fn_a + // Call host_fn_a + // Call wasm_fn_b + // Call host_fn_b + // Return host_fn_b + // Return wasm_fn_b + assert_eq!(caller.data().calling_wasm, 2); + assert_eq!(caller.data().returning_from_wasm, 1); + assert_eq!(caller.data().calling_host, 2); + assert_eq!(caller.data().returning_from_host, 1); + }); + linker.define("env", "host_fn_a", host_fn_a).unwrap(); + + let host_fn_b = Func::wrap(&mut store, |caller: Caller| { + // Call wasm_fn_a + // Call host_fn_a + // Call wasm_fn_b + // Call host_fn_b + assert_eq!(caller.data().calling_wasm, 2); + assert_eq!(caller.data().returning_from_wasm, 0); + assert_eq!(caller.data().calling_host, 2); + assert_eq!(caller.data().returning_from_host, 0); + }); + linker.define("env", "host_fn_b", host_fn_b).unwrap(); + + execute_wasm_fn_a(&mut store, &mut linker).unwrap(); + + assert_eq!(store.data().calling_wasm, 2); + assert_eq!(store.data().returning_from_wasm, 2); + assert_eq!(store.data().calling_host, 2); + assert_eq!(store.data().returning_from_host, 2); +} + +/// Utility function to generate a callback that fails after is has been called +/// `n` times. +#[allow(clippy::type_complexity)] +fn generate_error_after_n_calls + Clone + Send + Sync + 'static>( + limit: u32, + error: E, +) -> Box Result<(), Error> + Send + Sync> { + Box::new(move |data, hook_type| -> Result<(), Error> { + if (data.calling_wasm + + data.returning_from_wasm + + data.calling_host + + data.returning_from_host) + >= limit + { + return Err(error.clone().into()); + } + + match hook_type { + CallHook::CallingWasm => data.calling_wasm += 1, + CallHook::ReturningFromWasm => data.returning_from_wasm += 1, + CallHook::CallingHost => data.calling_host += 1, + CallHook::ReturningFromHost => data.returning_from_host += 1, + }; + + Ok(()) + }) +} + +#[test] +fn call_hook_prevents_wasm_execution() { + let (mut store, mut linker) = test_setup(); + + store.call_hook(generate_error_after_n_calls( + 0, + wasmi_core::TrapCode::BadConversionToInteger, + )); + + let should_not_run = Func::wrap(&mut store, |mut caller: Caller| { + caller.data_mut().erroneous_callback_invocation = true; + }); + + linker.define("env", "host_fn_a", should_not_run).unwrap(); + linker.define("env", "host_fn_b", should_not_run).unwrap(); + + let result = execute_wasm_fn_a(&mut store, &mut linker).map_err(|err| { + err.as_trap_code() + .expect("The returned error is not a trap code") + }); + + assert!( + !store.data().erroneous_callback_invocation, + "A callback that should have been prevented was executed." + ); + assert_eq!(result, Err(TrapCode::BadConversionToInteger)); +} + +#[test] +fn call_hook_prevents_host_execution() { + let (mut store, mut linker) = test_setup(); + + store.call_hook(generate_error_after_n_calls(1, TrapCode::BadSignature)); + + let should_not_run = Func::wrap(&mut store, |mut caller: Caller| { + caller.data_mut().erroneous_callback_invocation = true; + }); + + linker.define("env", "host_fn_a", should_not_run).unwrap(); + linker.define("env", "host_fn_b", should_not_run).unwrap(); + + let result = execute_wasm_fn_a(&mut store, &mut linker).map_err(|err| { + err.as_trap_code() + .expect("The returned error is not a trap code") + }); + + assert!( + !store.data().erroneous_callback_invocation, + "A callback that should have been prevented was executed." + ); + assert_eq!(result, Err(TrapCode::BadSignature)); +} + +#[test] +fn call_hook_prevents_nested_wasm_execution() { + let (mut store, mut linker) = test_setup(); + + store.call_hook(generate_error_after_n_calls( + 2, + TrapCode::GrowthOperationLimited, + )); + + let host_fn_a = Func::wrap(&mut store, |mut caller: Caller| { + let result = caller + .get_export("wasm_fn_b") + .and_then(Extern::into_func) + .unwrap() + .typed::<(), ()>(&caller) + .unwrap() + .call(&mut caller, ()) + .map_err(|err| { + err.as_trap_code() + .expect("The returned error is not a trap code") + }); + + assert_eq!(result, Err(TrapCode::GrowthOperationLimited)); + }); + + let should_not_run = Func::wrap(&mut store, |mut caller: Caller| { + caller.data_mut().erroneous_callback_invocation = true; + }); + + linker.define("env", "host_fn_a", host_fn_a).unwrap(); + linker.define("env", "host_fn_b", should_not_run).unwrap(); + + // wasm_fn_a should also return a `TrapCode` from `CallHook::ReturningFromHost` hook. + let result = execute_wasm_fn_a(&mut store, &mut linker).map_err(|err| { + err.as_trap_code() + .expect("The returned error is not a trap code") + }); + + assert!( + !store.data().erroneous_callback_invocation, + "A callback that should have been prevented was executed." + ); + assert_eq!(result, Err(TrapCode::GrowthOperationLimited)); +} diff --git a/crates/wasmi/tests/e2e/v1/mod.rs b/crates/wasmi/tests/e2e/v1/mod.rs index 2132a71a9a..09ff5e5ed5 100644 --- a/crates/wasmi/tests/e2e/v1/mod.rs +++ b/crates/wasmi/tests/e2e/v1/mod.rs @@ -1,3 +1,4 @@ +mod call_hook; mod fuel_consumption; mod fuel_metering; mod func;