Skip to content

Commit 47fa1ad

Browse files
authored
Rework bounds checking for atomic operations (#5239)
Before, we would do a `heap_addr` to translate the given Wasm memory address into a native memory address and pass it into the libcall that implemented the atomic operation, which would then treat the address as a Wasm memory address and pass it to `validate_atomic_addr` to be bounds checked a second time. This is a bit nonsensical, as we are validating a native memory address as if it were a Wasm memory address. Now, we no longer do a `heap_addr` to translate the Wasm memory address to a native memory address. Instead, we pass the Wasm memory address to the libcall, and the libcall is responsible for doing the bounds check (by calling `validate_atomic_addr` with the correct type of memory address now).
1 parent 8667948 commit 47fa1ad

File tree

4 files changed

+74
-52
lines changed

4 files changed

+74
-52
lines changed

cranelift/wasm/src/code_translator.rs

+39-17
Original file line numberDiff line numberDiff line change
@@ -1066,21 +1066,24 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
10661066
let heap = state.get_heap(builder.func, memarg.memory, environ)?;
10671067
let timeout = state.pop1(); // 64 (fixed)
10681068
let expected = state.pop1(); // 32 or 64 (per the `Ixx` in `IxxAtomicWait`)
1069-
let (_flags, addr) = prepare_atomic_addr(
1070-
memarg,
1071-
u8::try_from(implied_ty.bytes()).unwrap(),
1072-
builder,
1073-
state,
1074-
environ,
1075-
)?;
10761069
assert!(builder.func.dfg.value_type(expected) == implied_ty);
1070+
let addr = state.pop1();
1071+
let effective_addr = if memarg.offset == 0 {
1072+
addr
1073+
} else {
1074+
let index_type = builder.func.heaps[heap].index_type;
1075+
let offset = builder.ins().iconst(index_type, memarg.offset as i64);
1076+
builder
1077+
.ins()
1078+
.uadd_overflow_trap(addr, offset, ir::TrapCode::HeapOutOfBounds)
1079+
};
10771080
// `fn translate_atomic_wait` can inspect the type of `expected` to figure out what
10781081
// code it needs to generate, if it wants.
10791082
let res = environ.translate_atomic_wait(
10801083
builder.cursor(),
10811084
heap_index,
10821085
heap,
1083-
addr,
1086+
effective_addr,
10841087
expected,
10851088
timeout,
10861089
)?;
@@ -1090,12 +1093,23 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
10901093
let heap_index = MemoryIndex::from_u32(memarg.memory);
10911094
let heap = state.get_heap(builder.func, memarg.memory, environ)?;
10921095
let count = state.pop1(); // 32 (fixed)
1093-
1094-
// `memory.atomic.notify` is defined to have an access size of 4
1095-
// bytes in the spec, even though it doesn't necessarily access memory.
1096-
let (_flags, addr) = prepare_atomic_addr(memarg, 4, builder, state, environ)?;
1097-
let res =
1098-
environ.translate_atomic_notify(builder.cursor(), heap_index, heap, addr, count)?;
1096+
let addr = state.pop1();
1097+
let effective_addr = if memarg.offset == 0 {
1098+
addr
1099+
} else {
1100+
let index_type = builder.func.heaps[heap].index_type;
1101+
let offset = builder.ins().iconst(index_type, memarg.offset as i64);
1102+
builder
1103+
.ins()
1104+
.uadd_overflow_trap(addr, offset, ir::TrapCode::HeapOutOfBounds)
1105+
};
1106+
let res = environ.translate_atomic_notify(
1107+
builder.cursor(),
1108+
heap_index,
1109+
heap,
1110+
effective_addr,
1111+
count,
1112+
)?;
10991113
state.push1(res);
11001114
}
11011115
Operator::I32AtomicLoad { memarg } => {
@@ -2324,13 +2338,12 @@ fn prepare_addr<FE: FuncEnvironment + ?Sized>(
23242338
Ok((flags, addr))
23252339
}
23262340

2327-
fn prepare_atomic_addr<FE: FuncEnvironment + ?Sized>(
2341+
fn align_atomic_addr(
23282342
memarg: &MemArg,
23292343
loaded_bytes: u8,
23302344
builder: &mut FunctionBuilder,
23312345
state: &mut FuncTranslationState,
2332-
environ: &mut FE,
2333-
) -> WasmResult<(MemFlags, Value)> {
2346+
) {
23342347
// Atomic addresses must all be aligned correctly, and for now we check
23352348
// alignment before we check out-of-bounds-ness. The order of this check may
23362349
// need to be updated depending on the outcome of the official threads
@@ -2358,7 +2371,16 @@ fn prepare_atomic_addr<FE: FuncEnvironment + ?Sized>(
23582371
let f = builder.ins().icmp_imm(IntCC::NotEqual, misalignment, 0);
23592372
builder.ins().trapnz(f, ir::TrapCode::HeapMisaligned);
23602373
}
2374+
}
23612375

2376+
fn prepare_atomic_addr<FE: FuncEnvironment + ?Sized>(
2377+
memarg: &MemArg,
2378+
loaded_bytes: u8,
2379+
builder: &mut FunctionBuilder,
2380+
state: &mut FuncTranslationState,
2381+
environ: &mut FE,
2382+
) -> WasmResult<(MemFlags, Value)> {
2383+
align_atomic_addr(memarg, loaded_bytes, builder, state);
23622384
prepare_addr(memarg, loaded_bytes, builder, state, environ)
23632385
}
23642386

crates/cranelift/src/func_environ.rs

+2
Original file line numberDiff line numberDiff line change
@@ -1981,6 +1981,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
19811981
expected: ir::Value,
19821982
timeout: ir::Value,
19831983
) -> WasmResult<ir::Value> {
1984+
let addr = self.cast_memory_index_to_i64(&mut pos, addr, memory_index);
19841985
let implied_ty = pos.func.dfg.value_type(expected);
19851986
let (func_sig, memory_index, func_idx) =
19861987
self.get_memory_atomic_wait(&mut pos.func, memory_index, implied_ty);
@@ -2006,6 +2007,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
20062007
addr: ir::Value,
20072008
count: ir::Value,
20082009
) -> WasmResult<ir::Value> {
2010+
let addr = self.cast_memory_index_to_i64(&mut pos, addr, memory_index);
20092011
let func_sig = self
20102012
.builtin_function_signatures
20112013
.memory_atomic_notify(&mut pos.func);

crates/environ/src/builtin.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -42,11 +42,11 @@ macro_rules! foreach_builtin_function {
4242
/// Returns an index for Wasm's `global.get` instruction for `externref`s.
4343
externref_global_set(vmctx: vmctx, global: i32, val: reference);
4444
/// Returns an index for wasm's `memory.atomic.notify` instruction.
45-
memory_atomic_notify(vmctx: vmctx, memory: i32, addr: pointer, count: i32) -> i32;
45+
memory_atomic_notify(vmctx: vmctx, memory: i32, addr: i64, count: i32) -> i32;
4646
/// Returns an index for wasm's `memory.atomic.wait32` instruction.
47-
memory_atomic_wait32(vmctx: vmctx, memory: i32, addr: pointer, expected: i32, timeout: i64) -> i32;
47+
memory_atomic_wait32(vmctx: vmctx, memory: i32, addr: i64, expected: i32, timeout: i64) -> i32;
4848
/// Returns an index for wasm's `memory.atomic.wait64` instruction.
49-
memory_atomic_wait64(vmctx: vmctx, memory: i32, addr: pointer, expected: i64, timeout: i64) -> i32;
49+
memory_atomic_wait64(vmctx: vmctx, memory: i32, addr: i64, expected: i64, timeout: i64) -> i32;
5050
/// Invoked when fuel has run out while executing a function.
5151
out_of_gas(vmctx: vmctx);
5252
/// Invoked when we reach a new epoch.

crates/runtime/src/libcalls.rs

+30-32
Original file line numberDiff line numberDiff line change
@@ -434,17 +434,12 @@ unsafe fn externref_global_set(vmctx: *mut VMContext, index: u32, externref: *mu
434434
unsafe fn memory_atomic_notify(
435435
vmctx: *mut VMContext,
436436
memory_index: u32,
437-
addr: *mut u8,
437+
addr: u64,
438438
_count: u32,
439439
) -> Result<u32, TrapReason> {
440-
let addr = addr as usize;
441440
let memory = MemoryIndex::from_u32(memory_index);
442441
let instance = (*vmctx).instance();
443-
// this should never overflow since addr + 4 either hits a guard page
444-
// or it's been validated to be in-bounds already. Double-check for now
445-
// just to be sure.
446-
let addr_to_check = addr.checked_add(4).unwrap();
447-
validate_atomic_addr(instance, memory, addr_to_check)?;
442+
validate_atomic_addr(instance, memory, addr, 4, 4)?;
448443
Err(
449444
anyhow::anyhow!("unimplemented: wasm atomics (fn memory_atomic_notify) unsupported",)
450445
.into(),
@@ -455,17 +450,13 @@ unsafe fn memory_atomic_notify(
455450
unsafe fn memory_atomic_wait32(
456451
vmctx: *mut VMContext,
457452
memory_index: u32,
458-
addr: *mut u8,
453+
addr: u64,
459454
_expected: u32,
460455
_timeout: u64,
461456
) -> Result<u32, TrapReason> {
462-
let addr = addr as usize;
463457
let memory = MemoryIndex::from_u32(memory_index);
464458
let instance = (*vmctx).instance();
465-
// see wasmtime_memory_atomic_notify for why this shouldn't overflow
466-
// but we still double-check
467-
let addr_to_check = addr.checked_add(4).unwrap();
468-
validate_atomic_addr(instance, memory, addr_to_check)?;
459+
validate_atomic_addr(instance, memory, addr, 4, 4)?;
469460
Err(
470461
anyhow::anyhow!("unimplemented: wasm atomics (fn memory_atomic_wait32) unsupported",)
471462
.into(),
@@ -476,40 +467,47 @@ unsafe fn memory_atomic_wait32(
476467
unsafe fn memory_atomic_wait64(
477468
vmctx: *mut VMContext,
478469
memory_index: u32,
479-
addr: *mut u8,
470+
addr: u64,
480471
_expected: u64,
481472
_timeout: u64,
482473
) -> Result<u32, TrapReason> {
483-
let addr = addr as usize;
484474
let memory = MemoryIndex::from_u32(memory_index);
485475
let instance = (*vmctx).instance();
486-
// see wasmtime_memory_atomic_notify for why this shouldn't overflow
487-
// but we still double-check
488-
let addr_to_check = addr.checked_add(8).unwrap();
489-
validate_atomic_addr(instance, memory, addr_to_check)?;
476+
validate_atomic_addr(instance, memory, addr, 8, 8)?;
490477
Err(
491478
anyhow::anyhow!("unimplemented: wasm atomics (fn memory_atomic_wait64) unsupported",)
492479
.into(),
493480
)
494481
}
495482

496-
/// For atomic operations we still check the actual address despite this also
497-
/// being checked via the `heap_addr` instruction in cranelift. The reason for
498-
/// that is because the `heap_addr` instruction can defer to a later segfault to
499-
/// actually recognize the out-of-bounds whereas once we're running Rust code
500-
/// here we don't want to segfault.
501-
///
502-
/// In the situations where bounds checks were elided in JIT code (because oob
503-
/// would then be later guaranteed to segfault) this manual check is here
504-
/// so we don't segfault from Rust.
483+
macro_rules! ensure {
484+
($cond:expr, $trap:expr) => {
485+
if !($cond) {
486+
return Err($trap);
487+
}
488+
};
489+
}
490+
491+
/// In the configurations where bounds checks were elided in JIT code (because
492+
/// we are using static memories with virtual memory guard pages) this manual
493+
/// check is here so we don't segfault from Rust. For other configurations,
494+
/// these checks are required anyways.
505495
unsafe fn validate_atomic_addr(
506496
instance: &Instance,
507497
memory: MemoryIndex,
508-
addr: usize,
498+
addr: u64,
499+
access_size: u64,
500+
access_alignment: u64,
509501
) -> Result<(), TrapCode> {
510-
if addr > instance.get_memory(memory).current_length() {
511-
return Err(TrapCode::HeapOutOfBounds);
512-
}
502+
debug_assert!(access_alignment.is_power_of_two());
503+
ensure!(addr % access_alignment == 0, TrapCode::HeapMisaligned);
504+
505+
let length = u64::try_from(instance.get_memory(memory).current_length()).unwrap();
506+
ensure!(
507+
addr.saturating_add(access_size) < length,
508+
TrapCode::HeapOutOfBounds
509+
);
510+
513511
Ok(())
514512
}
515513

0 commit comments

Comments
 (0)