Cranelift: consider heap's guard pages when legalizing heap_addr (#5335)

* Cranelift: consider heap's guard pages when legalizing `heap_addr`

Fixes #5328

* Update comment to align more directly with implementation

* Add legalization tests for `heap_addr` and offset guard pages

cranelift/codegen/src/legalizer/heap.rs

@@ -28,7 +28,12 @@ pub fn expand_heap_addr(
         func.dfg.display_inst(inst)
     );
 
-    match func.heaps[heap].style {
+    let ir::HeapData {
+        offset_guard_size,
+        style,
+        ..
+    } = &func.heaps[heap];
+    match *style {
         ir::HeapStyle::Dynamic { bound_gv } => dynamic_addr(
             isa,
             inst,
@@ -47,6 +52,7 @@ pub fn expand_heap_addr(
             u32::from(offset_immediate),
             u8::from(access_size),
             bound.into(),
+            (*offset_guard_size).into(),
             func,
             cfg,
         ),
@@ -144,6 +150,7 @@ fn static_addr(
     offset: u32,
     access_size: u8,
     bound: u64,
+    guard_size: u64,
     func: &mut ir::Function,
     cfg: &mut ControlFlowGraph,
 ) {
@@ -172,23 +179,45 @@ fn static_addr(
     }
 
     // After the trivial case is done we're now mostly interested in trapping if
-    // `index > bound - offset - access_size`. We know `bound - offset -
-    // access_size` here is non-negative from the above comparison.
     //
-    // If we can know `bound - offset - access_size >= 4GB` then with a 32-bit
-    // offset we're guaranteed:
+    //     index + offset + size > bound
     //
-    //      bound - offset - access_size >= 4GB > index
+    // We know `bound - offset - access_size` here is non-negative from the
+    // above comparison, so we can rewrite that as
     //
-    // or, in other words, `index < bound - offset - access_size`, meaning we
-    // can't trap for any value of `index`.
+    //     index > bound - offset - size
     //
-    // With that we have an optimization here where with 32-bit offsets and
-    // `bound - access_size >= 4GB` we can omit a bounds check.
-    let limit = bound - offset as u64 - access_size as u64;
+    // Additionally, we add our guard pages (if any) to the right-hand side,
+    // since we can rely on the virtual memory subsystem at runtime to catch
+    // out-of-bound accesses within the range `bound .. bound + guard_size`. So
+    // now we are dealing with
+    //
+    //     index > bound + guard_size - offset - size
+    //
+    // (Note that `bound + guard_size` cannot overflow for correctly-configured
+    // heaps, as otherwise the heap wouldn't fit in a 64-bit memory space.)
+    //
+    // If we know the right-hand side is greater than or equal to 4GiB - 1, aka
+    // 0xffff_ffff, then with a 32-bit index we're guaranteed:
+    //
+    //     index <= 0xffff_ffff <= bound + guard_size - offset - access_size
+    //
+    // meaning that `index` is always either in bounds or within the guard page
+    // region, neither of which require emitting an explicit bounds check.
+    assert!(
+        bound.checked_add(guard_size).is_some(),
+        "heap's configuration doesn't fit in a 64-bit memory space"
+    );
     let mut spectre_oob_comparison = None;
     let index = cast_index_to_pointer_ty(index, index_ty, addr_ty, &mut pos);
-    if index_ty != ir::types::I32 || limit < 0xffff_ffff {
+    if index_ty == ir::types::I32
+        && bound + guard_size - offset_plus_size(offset, access_size) >= 0xffff_ffff
+    {
+        // Happy path! No bounds checks necessary!
+    } else {
+        // Since we have to emit explicit bounds checks anyways, ignore the
+        // guard pages and test against the precise limit.
+        let limit = bound - offset_plus_size(offset, access_size);
         // Here we want to test the condition `index > limit` and if that's true
         // then this is an out-of-bounds access and needs to trap.
         let oob = pos

cranelift/filetests/filetests/legalizer/static-heap-with-guard-pages.clif

@@ -0,0 +1,21 @@
+test legalizer
+set enable_heap_access_spectre_mitigation=true
+target x86_64
+
+;; The offset guard is large enough that we don't need explicit bounds checks.
+
+function %test(i64 vmctx, i32) -> i64 {
+    gv0 = vmctx
+    gv1 = load.i64 notrap aligned gv0+0
+    heap0 = static gv1, min 0x1000, bound 0x1000, offset_guard 0xffff_ffff, index_type i32
+
+block0(v0: i64, v1: i32):
+    v2 = heap_addr.i64 heap0, v1, 0, 4
+    return v2
+}
+
+; check:  block0(v0: i64, v1: i32):
+; nextln:     v3 = uextend.i64 v1
+; nextln:     v4 = load.i64 notrap aligned v0
+; nextln:     v2 = iadd v4, v3
+; nextln:     return v2

cranelift/filetests/filetests/legalizer/static-heap-without-guard-pages.clif

@@ -0,0 +1,34 @@
+test legalizer
+set enable_heap_access_spectre_mitigation=true
+target x86_64
+
+;; The offset guard is not large enough to avoid explicit bounds checks.
+
+function %test(i64 vmctx, i32) -> i64 {
+    gv0 = vmctx
+    gv1 = load.i64 notrap aligned gv0+0
+    heap0 = static gv1, min 0x1000, bound 0x1000, offset_guard 0xffff_0000, index_type i32
+
+block0(v0: i64, v1: i32):
+    v2 = heap_addr.i64 heap0, v1, 0, 4
+    return v2
+}
+
+; check:  block0(v0: i64, v1: i32):
+; nextln:     v3 = uextend.i64 v1
+; nextln:     v10 = iconst.i64 4092
+; nextln:     v4 = icmp ugt v3, v10  ; v10 = 4092
+; nextln:     brz v4, block2
+; nextln:     jump block1
+; nextln: 
+; nextln: block1:
+; nextln:     trap heap_oob
+; nextln: 
+; nextln: block2:
+; nextln:     v5 = iconst.i64 4092
+; nextln:     v6 = load.i64 notrap aligned v0
+; nextln:     v7 = iadd v6, v3
+; nextln:     v8 = iconst.i64 0
+; nextln:     v9 = icmp.i64 ugt v3, v5  ; v5 = 4092
+; nextln:     v2 = select_spectre_guard v9, v8, v7  ; v8 = 0
+; nextln:     return v2