mpk: enable more tests (#7376)

* mpk: add Wasm test

This adds a Wasmtime-level test checking that, when MPK is available on
the system, an MPK-configured memory pool obeys the same invariants as
other memory pool configurations (e.g., `guards_present` for normal
allocation, `guards_present_pooling` for non-MPK pooling allocation).

It also fixes a bug. A previous commit had updated the validation logic
of `MemoryPool` to check that the memory plan's bound would be reached
before the next slot of the same stripe. But `MemoryPool::allocate` has
some "double-check" logic that was triggered by this test. It now
matches the logic in `MemoryPool::validate`.

* mpk: explicitly disable MPK on certain tests

Some Wasmtime tests rely on specific limits to the memory pool. When MPK
is enabled, these tests fail because MPK splits access to the pool
slices among different stores. This change does not yet enable MPK,
though locally I run with MPK enabled. With this commit and MPK enabled
locally, all Wasmtime tests now pass.

crates/runtime/src/instance/allocator/pooling/memory_pool.rs

@@ -341,8 +341,9 @@ impl MemoryPool {
             // but double-check here to be sure.
             match memory_plan.style {
                 MemoryStyle::Static { bound } => {
-                    let bound = bound * u64::from(WASM_PAGE_SIZE);
-                    assert!(bound <= u64::try_from(self.layout.slot_bytes).unwrap());
+                    let pages_to_next_stripe_slot =
+                        u64::try_from(self.layout.pages_to_next_stripe_slot()).unwrap();
+                    assert!(bound <= pages_to_next_stripe_slot);
                 }
                 MemoryStyle::Dynamic { .. } => {}
             }

tests/all/limits.rs

@@ -1,5 +1,6 @@
 use anyhow::Result;
 use wasmtime::*;
+use wasmtime_runtime::MpkEnabled;
 
 const WASM_PAGE_SIZE: usize = wasmtime_environ::WASM_PAGE_SIZE as usize;
 
@@ -355,7 +356,8 @@ fn test_pooling_allocator_initial_limits_exceeded() -> Result<()> {
     let mut pool = crate::small_pool_config();
     pool.total_memories(2)
         .max_memories_per_module(2)
-        .memory_pages(5);
+        .memory_pages(5)
+        .memory_protection_keys(MpkEnabled::Disable);
     let mut config = Config::new();
     config.wasm_multi_memory(true);
     config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool));

tests/all/memory.rs

@@ -3,6 +3,7 @@ use rayon::prelude::*;
 use std::sync::atomic::{AtomicU32, Ordering::SeqCst};
 use std::time::{Duration, Instant};
 use wasmtime::*;
+use wasmtime_runtime::{mpk, MpkEnabled};
 
 fn module(engine: &Engine) -> Result<Module> {
     let mut wat = format!("(module\n");
@@ -192,7 +193,72 @@ fn guards_present_pooling() -> Result<()> {
     const GUARD_SIZE: u64 = 65536;
 
     let mut pool = crate::small_pool_config();
-    pool.total_memories(2).memory_pages(10);
+    pool.total_memories(2)
+        .memory_pages(10)
+        .memory_protection_keys(MpkEnabled::Disable);
+    let mut config = Config::new();
+    config.static_memory_maximum_size(1 << 20);
+    config.dynamic_memory_guard_size(GUARD_SIZE);
+    config.static_memory_guard_size(GUARD_SIZE);
+    config.guard_before_linear_memory(true);
+    config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool));
+    let engine = Engine::new(&config)?;
+
+    let mut store = Store::new(&engine, ());
+
+    let mem1 = {
+        let m = Module::new(&engine, "(module (memory (export \"\") 1 2))")?;
+        Instance::new(&mut store, &m, &[])?
+            .get_memory(&mut store, "")
+            .unwrap()
+    };
+    let mem2 = {
+        let m = Module::new(&engine, "(module (memory (export \"\") 1))")?;
+        Instance::new(&mut store, &m, &[])?
+            .get_memory(&mut store, "")
+            .unwrap()
+    };
+
+    unsafe fn assert_guards(store: &Store<()>, mem: &Memory) {
+        // guards before
+        println!("check pre-mem");
+        assert_faults(mem.data_ptr(&store).offset(-(GUARD_SIZE as isize)));
+
+        // unmapped just after memory
+        println!("check mem");
+        assert_faults(mem.data_ptr(&store).add(mem.data_size(&store)));
+
+        // guards after memory
+        println!("check post-mem");
+        assert_faults(mem.data_ptr(&store).add(1 << 20));
+    }
+    unsafe {
+        assert_guards(&store, &mem1);
+        assert_guards(&store, &mem2);
+        println!("growing");
+        mem1.grow(&mut store, 1).unwrap();
+        mem2.grow(&mut store, 1).unwrap();
+        assert_guards(&store, &mem1);
+        assert_guards(&store, &mem2);
+    }
+
+    Ok(())
+}
+
+#[test]
+#[cfg_attr(miri, ignore)]
+fn guards_present_pooling_mpk() -> Result<()> {
+    if !mpk::is_supported() {
+        println!("skipping `guards_present_pooling_mpk` test; mpk is not supported");
+        return Ok(());
+    }
+
+    const GUARD_SIZE: u64 = 65536;
+    let mut pool = crate::small_pool_config();
+    pool.total_memories(4)
+        .memory_pages(10)
+        .memory_protection_keys(MpkEnabled::Enable)
+        .max_memory_protection_keys(2);
     let mut config = Config::new();
     config.static_memory_maximum_size(1 << 20);
     config.dynamic_memory_guard_size(GUARD_SIZE);

tests/all/pooling_allocator.rs

@@ -1,6 +1,7 @@
 use super::skip_pooling_allocator_tests;
 use anyhow::Result;
 use wasmtime::*;
+use wasmtime_runtime::MpkEnabled;
 
 #[test]
 fn successful_instantiation() -> Result<()> {
@@ -1151,7 +1152,8 @@ fn total_memories_limit() -> Result<()> {
 
     let mut pool = crate::small_pool_config();
     pool.total_memories(TOTAL_MEMORIES)
-        .total_core_instances(TOTAL_MEMORIES + 1);
+        .total_core_instances(TOTAL_MEMORIES + 1)
+        .memory_protection_keys(MpkEnabled::Disable);
     let mut config = Config::new();
     config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool));