diff --git a/cranelift/wasm/src/code_translator.rs b/cranelift/wasm/src/code_translator.rs index 05f7d0b991..2250fb3cff 100644 --- a/cranelift/wasm/src/code_translator.rs +++ b/cranelift/wasm/src/code_translator.rs @@ -2316,17 +2316,18 @@ where // offsets in `memarg` are <=2gb, which means we get the fast path of one // `heap_addr` instruction plus a hardcoded i32-offset in memory-related // instructions. + let heap = environ.heaps()[heap].clone(); let addr = match u32::try_from(memarg.offset) { // If our offset fits within a u32, then we can place the it into the // offset immediate of the `heap_addr` instruction. Ok(offset) => bounds_checks::bounds_check_and_compute_addr( builder, - &*environ, - &environ.heaps()[heap], + environ, + &heap, index, offset, access_size, - ), + )?, // If the offset doesn't fit within a u32, then we can't pass it // directly into `heap_addr`. @@ -2355,20 +2356,19 @@ where // relatively odd/rare. In the future if needed we can look into // optimizing this more. Err(_) => { - let index_type = environ.heaps()[heap].index_type; - let offset = builder.ins().iconst(index_type, memarg.offset as i64); + let offset = builder.ins().iconst(heap.index_type, memarg.offset as i64); let adjusted_index = builder .ins() .uadd_overflow_trap(index, offset, ir::TrapCode::HeapOutOfBounds); bounds_checks::bounds_check_and_compute_addr( builder, - &*environ, - &environ.heaps()[heap], + environ, + &heap, adjusted_index, 0, access_size, - ) + )? } }; let addr = match addr { diff --git a/cranelift/wasm/src/code_translator/bounds_checks.rs b/cranelift/wasm/src/code_translator/bounds_checks.rs index 406e60863e..fa49b08b99 100644 --- a/cranelift/wasm/src/code_translator/bounds_checks.rs +++ b/cranelift/wasm/src/code_translator/bounds_checks.rs @@ -8,21 +8,22 @@ //! //! bounds check the memory access and translate it into a native memory access. -use crate::{HeapData, HeapStyle, TargetEnvironment}; +use crate::{FuncEnvironment, HeapData, HeapStyle}; use cranelift_codegen::{ cursor::{Cursor, FuncCursor}, ir::{self, condcodes::IntCC, InstBuilder, RelSourceLoc}, }; use cranelift_frontend::FunctionBuilder; +use wasmtime_types::WasmResult; /// Helper used to emit bounds checks (as necessary) and compute the native /// address of a heap access. /// /// Returns the `ir::Value` holding the native address of the heap access, or /// `None` if the heap access will unconditionally trap. -pub fn bounds_check_and_compute_addr( +pub fn bounds_check_and_compute_addr( builder: &mut FunctionBuilder, - env: &TE, + env: &mut Env, heap: &HeapData, // Dynamic operand indexing into the heap. index: ir::Value, @@ -30,9 +31,9 @@ pub fn bounds_check_and_compute_addr( offset: u32, // Static size of the heap access. access_size: u8, -) -> Option +) -> WasmResult> where - TE: TargetEnvironment + ?Sized, + Env: FuncEnvironment + ?Sized, { let index = cast_index_to_pointer_ty( index, @@ -62,7 +63,7 @@ where // different bounds checks and optimizations of those bounds checks. It is // intentionally written in a straightforward case-matching style that will // hopefully make it easy to port to ISLE one day. - match heap.style { + Ok(match heap.style { // ====== Dynamic Memories ====== // // 1. First special case for when `offset + access_size == 1`: @@ -216,6 +217,7 @@ where // bound`, since we will end up being out-of-bounds regardless of the // given `index`. HeapStyle::Static { bound } if offset_and_size > bound.into() => { + env.before_unconditionally_trapping_memory_access(builder)?; builder.ins().trap(ir::TrapCode::HeapOutOfBounds); None } @@ -325,7 +327,7 @@ where None, )) } - } + }) } fn cast_index_to_pointer_ty( diff --git a/cranelift/wasm/src/environ/spec.rs b/cranelift/wasm/src/environ/spec.rs index 08b27597ab..1b64ec811a 100644 --- a/cranelift/wasm/src/environ/spec.rs +++ b/cranelift/wasm/src/environ/spec.rs @@ -490,6 +490,18 @@ pub trait FuncEnvironment: TargetEnvironment { Ok(()) } + /// Optional callback for the `FuncEnvironment` performing this translation + /// to maintain, prepare, or finalize custom, internal state when we + /// statically determine that a Wasm memory access will unconditionally + /// trap, rendering the rest of the block unreachable. Called just before + /// the unconditional trap is emitted. + fn before_unconditionally_trapping_memory_access( + &mut self, + _builder: &mut FunctionBuilder, + ) -> WasmResult<()> { + Ok(()) + } + /// Optional callback for the `FunctionEnvironment` performing this translation to perform work /// before the function body is translated. fn before_translate_function( diff --git a/crates/cranelift/src/func_environ.rs b/crates/cranelift/src/func_environ.rs index 6a3473ca6d..f938d7dab7 100644 --- a/crates/cranelift/src/func_environ.rs +++ b/crates/cranelift/src/func_environ.rs @@ -2078,6 +2078,17 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m Ok(()) } + fn before_unconditionally_trapping_memory_access( + &mut self, + builder: &mut FunctionBuilder, + ) -> WasmResult<()> { + if self.tunables.consume_fuel { + self.fuel_increment_var(builder); + self.fuel_save_from_var(builder); + } + Ok(()) + } + fn before_translate_function( &mut self, builder: &mut FunctionBuilder, diff --git a/tests/all/fuel.rs b/tests/all/fuel.rs index 9b924f7284..730d56b44b 100644 --- a/tests/all/fuel.rs +++ b/tests/all/fuel.rs @@ -183,3 +183,43 @@ fn manual_edge_cases() { assert!(store.consume_fuel(i64::MAX as u64 + 1).is_err()); assert_eq!(store.consume_fuel(i64::MAX as u64).unwrap(), 0); } + +#[test] +fn unconditionally_trapping_memory_accesses_save_fuel_before_trapping() { + let mut config = Config::new(); + config.consume_fuel(true); + config.static_memory_maximum_size(0x1_0000); + + let engine = Engine::new(&config).unwrap(); + + let module = Module::new( + &engine, + r#" + (module + (memory 1 1) + (func (export "f") (param i32) (result i32) + local.get 0 + local.get 0 + i32.add + ;; This offset is larger than our memory max size and therefore + ;; will unconditionally trap. + i32.load8_s offset=0xffffffff)) + "#, + ) + .unwrap(); + + let mut store = Store::new(&engine, ()); + store.add_fuel(1_000).unwrap(); + + let instance = Instance::new(&mut store, &module, &[]).unwrap(); + let f = instance + .get_typed_func::(&mut store, "f") + .unwrap(); + + let trap = f.call(&mut store, 0).unwrap_err(); + assert_eq!(trap.downcast::().unwrap(), Trap::MemoryOutOfBounds); + + // The `i32.add` consumed some fuel before the unconditionally trapping + // memory access. + assert!(store.fuel_consumed().unwrap() > 0); +}