Skip to content

Commit

Permalink
cranelift-wasm: Attach table OOB traps to loads/stores (#8171)
Browse files Browse the repository at this point in the history 
Currently, every access to a table element does a bounds-check with a
conditional branch to a block that explicitly traps.

Instead, when Spectre mitigations are enabled, let's change the address
computation to return a null pointer for out-of-bounds accesses, and
then allow the subsequent load or store to trap.

This is less code in that case since we can reuse instructions we needed
anyway.

We return the MemFlags that the memory access should use, in addition to
the address it should access. That way we don't record trap metadata on
memory access instructions which can't actually trap due to being
preceded by a `trapnz`-based bounds check, when Spectre mitigations are
disabled.

In addition, when the table has constant size and the element index is a
constant and mid-end optimization is enabled, this allows the
bounds-check to be constant folded away. Later, #8139 will let us
optimize away the select_spectre_guard instruction in this case too.

Once we also implement #8160, `tests/disas/typed-funcrefs.wat` should be
almost as fast as native indirect function calls.
  • Loading branch information
@jameysharp
jameysharp authored Mar 19, 2024
1 parent 3deaa87 commit 310e97e
Show file tree
Hide file tree
Showing 10 changed files with 97 additions and 172 deletions.
8 changes: 4 additions & 4 deletions cranelift/wasm/src/environ/dummy.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -530,8 +530,8 @@ impl<'dummy_environment> FuncEnvironment for DummyFuncEnvironment<'dummy_environ
let pointer_type = self.pointer_type();
self.ensure_table_exists(builder.func, table_index);
let table = self.tables[table_index].as_ref().unwrap();
let table_entry_addr = table.prepare_table_addr(builder, index, pointer_type, true);
let flags = ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
let (table_entry_addr, flags) =
table.prepare_table_addr(builder, index, pointer_type, true);
let value = builder
.ins()
.load(self.reference_type(), flags, table_entry_addr, 0);
Expand All @@ -548,8 +548,8 @@ impl<'dummy_environment> FuncEnvironment for DummyFuncEnvironment<'dummy_environ
let pointer_type = self.pointer_type();
self.ensure_table_exists(builder.func, table_index);
let table = self.tables[table_index].as_ref().unwrap();
let table_entry_addr = table.prepare_table_addr(builder, index, pointer_type, true);
let flags = ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
let (table_entry_addr, flags) =
table.prepare_table_addr(builder, index, pointer_type, true);
builder.ins().store(flags, value, table_entry_addr, 0);
Ok(())
}
Expand Down
34 changes: 17 additions & 17 deletions cranelift/wasm/src/table.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ impl TableData {
mut index: ir::Value,
addr_ty: ir::Type,
enable_table_access_spectre_mitigation: bool,
) -> ir::Value {
) -> (ir::Value, ir::MemFlags) {
let index_ty = pos.func.dfg.value_type(index);

// Start with the bounds check. Trap if `index + 1 > bound`.
Expand All @@ -60,16 +60,10 @@ impl TableData {
let oob = pos
.ins()
.icmp(IntCC::UnsignedGreaterThanOrEqual, index, bound);
pos.ins().trapnz(oob, ir::TrapCode::TableOutOfBounds);

// If Spectre mitigations are enabled, we will use a comparison to
// short-circuit the computed table element address to the start
// of the table on the misspeculation path when out-of-bounds.
let spectre_oob_cmp = if enable_table_access_spectre_mitigation {
Some((index, bound))
} else {
None
};
if !enable_table_access_spectre_mitigation {
pos.ins().trapnz(oob, ir::TrapCode::TableOutOfBounds);
}

// Convert `index` to `addr_ty`.
if index_ty != addr_ty {
Expand All @@ -91,14 +85,20 @@ impl TableData {

let element_addr = pos.ins().iadd(base, offset);

if let Some((index, bound)) = spectre_oob_cmp {
let cond = pos
.ins()
.icmp(IntCC::UnsignedGreaterThanOrEqual, index, bound);
// If out-of-bounds, choose the table base on the misspeculation path.
pos.ins().select_spectre_guard(cond, base, element_addr)
let base_flags = ir::MemFlags::new()
.with_aligned()
.with_alias_region(Some(ir::AliasRegion::Table));
if enable_table_access_spectre_mitigation {
// Short-circuit the computed table element address to a null pointer
// when out-of-bounds. The consumer of this address will trap when
// trying to access it.
let zero = pos.ins().iconst(addr_ty, 0);
(
pos.ins().select_spectre_guard(oob, zero, element_addr),
base_flags.with_trap_code(Some(ir::TrapCode::TableOutOfBounds)),
)
} else {
element_addr
(element_addr, base_flags.with_trap_code(None))
}
}
}
16 changes: 7 additions & 9 deletions crates/cranelift/src/func_environ.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -955,13 +955,12 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
// contents, we check for a null entry here, and
// if null, we take a slow-path that invokes a
// libcall.
let table_entry_addr = table_data.prepare_table_addr(
let (table_entry_addr, flags) = table_data.prepare_table_addr(
builder,
index,
pointer_type,
self.isa.flags().enable_table_access_spectre_mitigation(),
);
let flags = ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
let value = builder.ins().load(pointer_type, flags, table_entry_addr, 0);
// Mask off the "initialized bit". See documentation on
// FUNCREF_INIT_BIT in crates/environ/src/ref_bits.rs for more
Expand Down Expand Up @@ -1548,13 +1547,12 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
// Load the table element.
self.ensure_table_exists(builder.func, table_index);
let table_data = self.tables[table_index].as_ref().unwrap();
let elem_addr = table_data.prepare_table_addr(
let (elem_addr, flags) = table_data.prepare_table_addr(
builder,
index,
pointer_type,
self.isa.flags().enable_table_access_spectre_mitigation(),
);
let flags = ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
let elem = builder.ins().load(reference_type, flags, elem_addr, 0);

let elem_is_null = builder.ins().is_null(elem);
Expand Down Expand Up @@ -1659,7 +1657,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
WasmHeapType::Func | WasmHeapType::Concrete(_) | WasmHeapType::NoFunc => {
match plan.style {
TableStyle::CallerChecksSignature => {
let table_entry_addr = table_data.prepare_table_addr(
let (table_entry_addr, flags) = table_data.prepare_table_addr(
builder,
index,
pointer_type,
Expand All @@ -1671,8 +1669,6 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
let value_with_init_bit = builder
.ins()
.bor_imm(value, Imm64::from(FUNCREF_INIT_BIT as i64));
let flags =
ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
builder
.ins()
.store(flags, value_with_init_bit, table_entry_addr, 0);
Expand Down Expand Up @@ -1736,14 +1732,16 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
// not recorded in the stack map (which would lead to the GC
// saving a reference to a deallocated object, and then using it
// after its been freed).
let table_entry_addr = table_data.prepare_table_addr(
let (table_entry_addr, flags) = table_data.prepare_table_addr(
builder,
index,
pointer_type,
self.isa.flags().enable_table_access_spectre_mitigation(),
);
let flags = ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
let current_elem = builder.ins().load(pointer_type, flags, table_entry_addr, 0);

// After the load, a store to the same address can't trap.
let flags = ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
builder.ins().store(flags, value, table_entry_addr, 0);

// If value is not null, increment `value`'s ref count.
Expand Down
7 changes: 3 additions & 4 deletions tests/disas/icall-simd.wat
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,13 @@
;; block0(v0: i64, v1: i64, v2: i32, v3: i8x16):
;; @0033 v5 = iconst.i32 23
;; @0033 v6 = icmp uge v2, v5 ; v5 = 23
;; @0033 trapnz v6, table_oob
;; @0033 v7 = uextend.i64 v2
;; @0033 v8 = global_value.i64 gv4
;; @0033 v9 = ishl_imm v7, 3
;; @0033 v10 = iadd v8, v9
;; @0033 v11 = icmp uge v2, v5 ; v5 = 23
;; @0033 v12 = select_spectre_guard v11, v8, v10
;; @0033 v13 = load.i64 notrap aligned table v12
;; @0033 v11 = iconst.i64 0
;; @0033 v12 = select_spectre_guard v6, v11, v10 ; v11 = 0
;; @0033 v13 = load.i64 table_oob aligned table v12
;; @0033 v14 = band_imm v13, -2
;; @0033 brif v13, block3(v14), block2
;;
Expand Down
7 changes: 3 additions & 4 deletions tests/disas/icall.wat
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,13 @@
;; block0(v0: i64, v1: i64, v2: i32, v3: f32):
;; @0033 v5 = iconst.i32 23
;; @0033 v6 = icmp uge v2, v5 ; v5 = 23
;; @0033 trapnz v6, table_oob
;; @0033 v7 = uextend.i64 v2
;; @0033 v8 = global_value.i64 gv4
;; @0033 v9 = ishl_imm v7, 3
;; @0033 v10 = iadd v8, v9
;; @0033 v11 = icmp uge v2, v5 ; v5 = 23
;; @0033 v12 = select_spectre_guard v11, v8, v10
;; @0033 v13 = load.i64 notrap aligned table v12
;; @0033 v11 = iconst.i64 0
;; @0033 v12 = select_spectre_guard v6, v11, v10 ; v11 = 0
;; @0033 v13 = load.i64 table_oob aligned table v12
;; @0033 v14 = band_imm v13, -2
;; @0033 brif v13, block3(v14), block2
;;
Expand Down
24 changes: 6 additions & 18 deletions tests/disas/table-get-fixed-size.wat
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,20 +33,14 @@
;; @0052 v3 = iconst.i32 0
;; @0054 v4 = iconst.i32 7
;; @0054 v5 = icmp uge v3, v4 ; v3 = 0, v4 = 7
;; @0054 brif v5, block6, block7
;;
;; block6 cold:
;; @0054 trap table_oob
;;
;; block7:
;; @0054 v6 = uextend.i64 v3 ; v3 = 0
;; @0054 v7 = load.i64 notrap aligned v25+72
;; v26 = iconst.i64 3
;; @0054 v8 = ishl v6, v26 ; v26 = 3
;; @0054 v9 = iadd v7, v8
;; @0054 v10 = icmp.i32 uge v3, v4 ; v3 = 0, v4 = 7
;; @0054 v11 = select_spectre_guard v10, v7, v9
;; @0054 v12 = load.r64 notrap aligned table v11
;; @0054 v10 = iconst.i64 0
;; @0054 v11 = select_spectre_guard v5, v10, v9 ; v10 = 0
;; @0054 v12 = load.r64 table_oob aligned table v11
;; v2 -> v12
;; @0054 v13 = is_null v12
;; @0054 brif v13, block2, block3
Expand Down Expand Up @@ -99,20 +93,14 @@
;; v25 -> v0
;; @005b v4 = iconst.i32 7
;; @005b v5 = icmp uge v2, v4 ; v4 = 7
;; @005b brif v5, block6, block7
;;
;; block6 cold:
;; @005b trap table_oob
;;
;; block7:
;; @005b v6 = uextend.i64 v2
;; @005b v7 = load.i64 notrap aligned v25+72
;; v26 = iconst.i64 3
;; @005b v8 = ishl v6, v26 ; v26 = 3
;; @005b v9 = iadd v7, v8
;; @005b v10 = icmp.i32 uge v2, v4 ; v4 = 7
;; @005b v11 = select_spectre_guard v10, v7, v9
;; @005b v12 = load.r64 notrap aligned table v11
;; @005b v10 = iconst.i64 0
;; @005b v11 = select_spectre_guard v5, v10, v9 ; v10 = 0
;; @005b v12 = load.r64 table_oob aligned table v11
;; v3 -> v12
;; @005b v13 = is_null v12
;; @005b brif v13, block2, block3
Expand Down
24 changes: 6 additions & 18 deletions tests/disas/table-get.wat
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,20 +34,14 @@
;; @0051 v3 = iconst.i32 0
;; @0053 v4 = load.i32 notrap aligned v25+80
;; @0053 v5 = icmp uge v3, v4 ; v3 = 0
;; @0053 brif v5, block6, block7
;;
;; block6 cold:
;; @0053 trap table_oob
;;
;; block7:
;; @0053 v6 = uextend.i64 v3 ; v3 = 0
;; @0053 v7 = load.i64 notrap aligned v26+72
;; v27 = iconst.i64 3
;; @0053 v8 = ishl v6, v27 ; v27 = 3
;; @0053 v9 = iadd v7, v8
;; @0053 v10 = icmp.i32 uge v3, v4 ; v3 = 0
;; @0053 v11 = select_spectre_guard v10, v7, v9
;; @0053 v12 = load.r64 notrap aligned table v11
;; @0053 v10 = iconst.i64 0
;; @0053 v11 = select_spectre_guard v5, v10, v9 ; v10 = 0
;; @0053 v12 = load.r64 table_oob aligned table v11
;; v2 -> v12
;; @0053 v13 = is_null v12
;; @0053 brif v13, block2, block3
Expand Down Expand Up @@ -102,20 +96,14 @@
;; v26 -> v0
;; @005a v4 = load.i32 notrap aligned v25+80
;; @005a v5 = icmp uge v2, v4
;; @005a brif v5, block6, block7
;;
;; block6 cold:
;; @005a trap table_oob
;;
;; block7:
;; @005a v6 = uextend.i64 v2
;; @005a v7 = load.i64 notrap aligned v26+72
;; v27 = iconst.i64 3
;; @005a v8 = ishl v6, v27 ; v27 = 3
;; @005a v9 = iadd v7, v8
;; @005a v10 = icmp.i32 uge v2, v4
;; @005a v11 = select_spectre_guard v10, v7, v9
;; @005a v12 = load.r64 notrap aligned table v11
;; @005a v10 = iconst.i64 0
;; @005a v11 = select_spectre_guard v5, v10, v9 ; v10 = 0
;; @005a v12 = load.r64 table_oob aligned table v11
;; v3 -> v12
;; @005a v13 = is_null v12
;; @005a brif v13, block2, block3
Expand Down
32 changes: 10 additions & 22 deletions tests/disas/table-set-fixed-size.wat
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,22 +34,16 @@
;; @0052 v3 = iconst.i32 0
;; @0056 v4 = iconst.i32 7
;; @0056 v5 = icmp uge v3, v4 ; v3 = 0, v4 = 7
;; @0056 brif v5, block7, block8
;;
;; block7 cold:
;; @0056 trap table_oob
;;
;; block8:
;; @0056 v6 = uextend.i64 v3 ; v3 = 0
;; @0056 v7 = load.i64 notrap aligned v23+72
;; v24 = iconst.i64 3
;; @0056 v8 = ishl v6, v24 ; v24 = 3
;; @0056 v9 = iadd v7, v8
;; @0056 v10 = icmp.i32 uge v3, v4 ; v3 = 0, v4 = 7
;; @0056 v11 = select_spectre_guard v10, v7, v9
;; @0056 v12 = load.i64 notrap aligned table v11
;; @0056 store.r64 notrap aligned table v2, v11
;; @0056 v13 = is_null.r64 v2
;; @0056 v10 = iconst.i64 0
;; @0056 v11 = select_spectre_guard v5, v10, v9 ; v10 = 0
;; @0056 v12 = load.i64 table_oob aligned table v11
;; @0056 store notrap aligned table v2, v11
;; @0056 v13 = is_null v2
;; @0056 brif v13, block3, block2
;;
;; block2:
Expand Down Expand Up @@ -102,22 +96,16 @@
;; v23 -> v0
;; @005f v4 = iconst.i32 7
;; @005f v5 = icmp uge v2, v4 ; v4 = 7
;; @005f brif v5, block7, block8
;;
;; block7 cold:
;; @005f trap table_oob
;;
;; block8:
;; @005f v6 = uextend.i64 v2
;; @005f v7 = load.i64 notrap aligned v23+72
;; v24 = iconst.i64 3
;; @005f v8 = ishl v6, v24 ; v24 = 3
;; @005f v9 = iadd v7, v8
;; @005f v10 = icmp.i32 uge v2, v4 ; v4 = 7
;; @005f v11 = select_spectre_guard v10, v7, v9
;; @005f v12 = load.i64 notrap aligned table v11
;; @005f store.r64 notrap aligned table v3, v11
;; @005f v13 = is_null.r64 v3
;; @005f v10 = iconst.i64 0
;; @005f v11 = select_spectre_guard v5, v10, v9 ; v10 = 0
;; @005f v12 = load.i64 table_oob aligned table v11
;; @005f store notrap aligned table v3, v11
;; @005f v13 = is_null v3
;; @005f brif v13, block3, block2
;;
;; block2:
Expand Down
Loading

0 comments on commit 310e97e

Please sign in to comment.