diff --git a/crates/wasmi/src/engine/code_map.rs b/crates/wasmi/src/engine/code_map.rs index e790d44456..c6eb5d1f52 100644 --- a/crates/wasmi/src/engine/code_map.rs +++ b/crates/wasmi/src/engine/code_map.rs @@ -10,6 +10,7 @@ use crate::{ core::UntypedValue, engine::bytecode::Instruction, module::{FuncIdx, ModuleHeader}, + store::StoreInner, Error, }; use alloc::boxed::Box; @@ -83,6 +84,22 @@ impl InternalFuncEntity { Self::from(CompiledFuncEntity::uninit()) } + /// Charge fuel for compiling the given `bytes` representing the Wasm function body. + /// + /// # Note + /// + /// This only charges fuel if `ctx` is `Some` and `fuel_metering` is enabled. + fn charge_compilation_fuel(ctx: Option<&mut StoreInner>, bytes: &[u8]) -> Result<(), Error> { + let Some(ctx) = ctx else { return Ok(()) }; + if !ctx.engine().config().get_consume_fuel() { + return Ok(()); + } + let fuel_costs = ctx.engine().config().fuel_costs(); + let delta = fuel_costs.fuel_for_bytes(bytes.len() as u64); + ctx.fuel_mut().consume_fuel(delta)?; + Ok(()) + } + /// Compile the uncompiled [`FuncEntity`]. /// /// # Panics @@ -92,8 +109,9 @@ impl InternalFuncEntity { /// /// # Errors /// - /// If function translation failed. - fn compile(&mut self) -> Result<(), Error> { + /// - If function translation failed. + /// - If `ctx` ran out of fuel in case fuel consumption is enabled. + fn compile(&mut self, ctx: Option<&mut StoreInner>) -> Result<(), Error> { let uncompiled = match self { InternalFuncEntity::Uncompiled(func) => func, InternalFuncEntity::Compiled(func) => { @@ -102,6 +120,7 @@ impl InternalFuncEntity { }; let func_idx = uncompiled.func_idx; let bytes = mem::take(&mut uncompiled.bytes); + Self::charge_compilation_fuel(ctx, bytes.as_slice())?; let module = uncompiled.module.clone(); let Some(engine) = module.engine().upgrade() else { panic!( @@ -609,9 +628,13 @@ impl FuncEntity { /// /// # Errors /// - /// If translation or Wasm validation of the [`FuncEntity`] failed. + /// - If translation or Wasm validation of the [`FuncEntity`] failed. + /// - If `ctx` ran out of fuel in case fuel consumption is enabled. #[cold] - pub fn compile_and_get(&self) -> Result<&CompiledFuncEntity, Error> { + pub fn compile_and_get( + &self, + mut ctx: Option<&mut StoreInner>, + ) -> Result<&CompiledFuncEntity, Error> { loop { if let Some(func) = self.get_compiled() { // Case: The function has been compiled and can be returned. @@ -630,7 +653,10 @@ impl FuncEntity { // SAFETY: This method is only called after a lock has been acquired // to take responsibility for driving the function translation. let func = unsafe { &mut *self.func.get() }; - match func.compile() { + // Note: We need to use `take` because Rust doesn't know that this part of + // the loop is only executed once. + let ctx = ctx.take(); + match func.compile(ctx) { Ok(()) => { self.phase .set_compiled() @@ -692,14 +718,23 @@ impl CodeMap { } /// Returns the [`InternalFuncEntity`] of the [`CompiledFunc`]. + /// + /// # Errors + /// + /// - If translation or Wasm validation of the [`FuncEntity`] failed. + /// - If `ctx` ran out of fuel in case fuel consumption is enabled. #[track_caller] - pub fn get(&self, compiled_func: CompiledFunc) -> Result<&CompiledFuncEntity, Error> { + pub fn get( + &self, + ctx: Option<&mut StoreInner>, + compiled_func: CompiledFunc, + ) -> Result<&CompiledFuncEntity, Error> { let Some(func) = self.funcs.get(compiled_func) else { panic!("invalid compiled func: {compiled_func:?}") }; match func.get_compiled() { Some(func) => Ok(func), - None => func.compile_and_get(), + None => func.compile_and_get(ctx), } } } diff --git a/crates/wasmi/src/engine/executor/instrs/call.rs b/crates/wasmi/src/engine/executor/instrs/call.rs index c38b0e91af..8292f31615 100644 --- a/crates/wasmi/src/engine/executor/instrs/call.rs +++ b/crates/wasmi/src/engine/executor/instrs/call.rs @@ -172,7 +172,7 @@ impl<'ctx, 'engine> Executor<'ctx, 'engine> { params: CallParams, call_kind: CallKind, ) -> Result<(), Error> { - let func = self.code_map.get(func)?; + let func = self.code_map.get(Some(self.ctx), func)?; let mut called = self.dispatch_compiled_func(results, func)?; if let CallParams::Some = params { let called_sp = self.frame_stack_ptr(&called); diff --git a/crates/wasmi/src/engine/executor/mod.rs b/crates/wasmi/src/engine/executor/mod.rs index 8f86fd078b..0a2114cb3f 100644 --- a/crates/wasmi/src/engine/executor/mod.rs +++ b/crates/wasmi/src/engine/executor/mod.rs @@ -201,8 +201,13 @@ impl<'engine> EngineExecutor<'engine> { let len_results = results.len_results(); self.stack.values.reserve(len_results)?; self.stack.values.extend_zeros(len_results); - let instance = wasm_func.instance(); - let compiled_func = self.res.code_map.get(wasm_func.func_body())?; + let instance = *wasm_func.instance(); + let compiled_func = wasm_func.func_body(); + let ctx = ctx.as_context_mut(); + let compiled_func = self + .res + .code_map + .get(Some(&mut ctx.store.inner), compiled_func)?; let (base_ptr, frame_ptr) = self.stack.values.alloc_call_frame(compiled_func)?; // Safety: We use the `base_ptr` that we just received upon allocating the new // call frame which is guaranteed to be valid for this particular operation @@ -215,9 +220,9 @@ impl<'engine> EngineExecutor<'engine> { frame_ptr, base_ptr, RegisterSpan::new(Register::from_i16(0)), - *instance, + instance, ))?; - self.execute_func(ctx.as_context_mut())?; + self.execute_func(ctx)?; } FuncEntity::Host(host_func) => { // The host function signature is required for properly diff --git a/crates/wasmi/src/engine/mod.rs b/crates/wasmi/src/engine/mod.rs index 85ff29bd8d..b76b18c7e7 100644 --- a/crates/wasmi/src/engine/mod.rs +++ b/crates/wasmi/src/engine/mod.rs @@ -739,7 +739,8 @@ impl EngineInner { where F: FnOnce(&CompiledFuncEntity) -> R, { - Ok(f(self.res.read().code_map.get(func)?)) + // Note: We use `None` so this test-only function will never charge for compilation fuel. + Ok(f(self.res.read().code_map.get(None, func)?)) } /// Returns the [`Instruction`] of `func` at `index`.