mpk: limit the number of protection keys

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

@@ -219,6 +219,8 @@ pub struct PoolingInstanceAllocatorConfig {
     pub table_keep_resident: usize,
     /// Whether to enable memory protection keys.
     pub memory_protection_keys: MpkEnabled,
+    /// How many memory protection keys to allocate.
+    pub max_memory_protection_keys: usize,
 }
 
 impl Default for PoolingInstanceAllocatorConfig {
@@ -232,6 +234,7 @@ impl Default for PoolingInstanceAllocatorConfig {
             linear_memory_keep_resident: 0,
             table_keep_resident: 0,
             memory_protection_keys: MpkEnabled::Disable,
+            max_memory_protection_keys: 16,
         }
     }
 }

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

@@ -156,14 +156,14 @@ impl MemoryPool {
         let pkeys = match config.memory_protection_keys {
             MpkEnabled::Auto => {
                 if mpk::is_supported() {
-                    mpk::keys()
+                    mpk::keys(config.max_memory_protection_keys)
                 } else {
                     &[]
                 }
             }
             MpkEnabled::Enable => {
                 if mpk::is_supported() {
-                    mpk::keys()
+                    mpk::keys(config.max_memory_protection_keys)
                 } else {
                     bail!("mpk is disabled on this system")
                 }

crates/runtime/src/mpk/disabled.rs

@@ -8,7 +8,7 @@ use anyhow::Result;
 pub fn is_supported() -> bool {
     false
 }
-pub fn keys() -> &'static [ProtectionKey] {
+pub fn keys(_: usize) -> &'static [ProtectionKey] {
     &[]
 }
 pub fn allow(_: ProtectionMask) {}

crates/runtime/src/mpk/enabled.rs

@@ -9,34 +9,39 @@ pub fn is_supported() -> bool {
     cfg!(target_os = "linux") && cfg!(target_arch = "x86_64") && pkru::has_cpuid_bit_set()
 }
 
-/// Allocate all protection keys available to this process.
+/// Allocate up to `max` protection keys.
 ///
-/// This asks the kernel for all available keys (we expect 1-15; 0 is
-/// kernel-reserved) in a thread-safe way. This avoids interference when
+/// This asks the kernel for all available keys up to `max` in a thread-safe way
+/// (we can expect 1-15; 0 is kernel-reserved). This avoids interference when
 /// multiple threads try to allocate keys at the same time (e.g., during
-/// testing). It also ensures that a single copy of the keys are reserved for
-/// the lifetime of the process.
+/// testing). It also ensures that a single copy of the keys is reserved for the
+/// lifetime of the process. Because of this, `max` is only a hint to
+/// allocation: it only is effective on the first invocation of this function.
 ///
 /// TODO: this is not the best-possible design. This creates global state that
 /// would prevent any other code in the process from using protection keys; the
 /// `KEYS` are never deallocated from the system with `pkey_dealloc`.
-pub fn keys() -> &'static [ProtectionKey] {
+pub fn keys(max: usize) -> &'static [ProtectionKey] {
     let keys = KEYS.get_or_init(|| {
         let mut allocated = vec![];
         if is_supported() {
-            while let Ok(key_id) = sys::pkey_alloc(0, 0) {
-                debug_assert!(key_id < 16);
-                // UNSAFETY: here we unsafely assume that the system-allocated
-                // pkey will exist forever.
-                allocated.push(ProtectionKey {
-                    id: key_id,
-                    stripe: allocated.len().try_into().unwrap(),
-                });
+            while allocated.len() < max {
+                if let Ok(key_id) = sys::pkey_alloc(0, 0) {
+                    debug_assert!(key_id < 16);
+                    // UNSAFETY: here we unsafely assume that the
+                    // system-allocated pkey will exist forever.
+                    allocated.push(ProtectionKey {
+                        id: key_id,
+                        stripe: allocated.len().try_into().unwrap(),
+                    });
+                } else {
+                    break;
+                }
             }
         }
         allocated
     });
-    &keys
+    &keys[..keys.len().min(max)]
 }
 static KEYS: OnceLock<Vec<ProtectionKey>> = OnceLock::new();
 
@@ -152,14 +157,14 @@ mod tests {
     #[test]
     fn check_initialized_keys() {
         if is_supported() {
-            assert!(!keys().is_empty())
+            assert!(!keys(15).is_empty())
         }
     }
 
     #[test]
     fn check_invalid_mark() {
         skip_if_mpk_unavailable!();
-        let pkey = keys()[0];
+        let pkey = keys(15)[0];
         let unaligned_region = unsafe {
             let addr = 1 as *mut u8; // this is not page-aligned!
             let len = 1;

crates/wasmtime/src/config.rs

@@ -2386,6 +2386,23 @@ impl PoolingAllocationConfig {
         self
     }
 
+    /// Sets an upper limit on how many memory protection keys (MPK) Wasmtime
+    /// will use.
+    ///
+    /// This setting is only applicable when
+    /// [`PoolingAllocationConfig::memory_protection_keys`] is set to `enable`
+    /// or `auto`. Configuring this above the HW and OS limits (typically 15)
+    /// has no effect.
+    ///
+    /// If multiple Wasmtime engines are used in the same process, note that all
+    /// engines will share the same set of allocated keys; this setting will
+    /// limit how many keys are allocated initially and thus available to all
+    /// other engines.
+    pub fn max_memory_protection_keys(&mut self, max: usize) -> &mut Self {
+        self.config.max_memory_protection_keys = max;
+        self
+    }
+
     /// Check if memory protection keys (MPK) are available on the current host.
     ///
     /// This is a convenience method for determining MPK availability using the

tests/all/wast.rs

@@ -123,11 +123,13 @@ fn run_wast(wast: &str, strategy: Strategy, pooling: bool) -> anyhow::Result<()>
         }
 
         // The limits here are crafted such that the wast tests should pass.
-        // However, these limits may become insufficient in the future as the wast tests change.
-        // If a wast test fails because of a limit being "exceeded" or if memory/table
-        // fails to grow, the values here will need to be adjusted.
+        // However, these limits may become insufficient in the future as the
+        // wast tests change. If a wast test fails because of a limit being
+        // "exceeded" or if memory/table fails to grow, the values here will
+        // need to be adjusted.
         let mut pool = PoolingAllocationConfig::default();
-        pool.total_memories(450)
+        pool.total_memories(450 * 2)
+            .max_memory_protection_keys(2)
             .memory_pages(805)
             .max_memories_per_module(if multi_memory { 9 } else { 1 })
             .max_tables_per_module(4);
@@ -180,7 +182,7 @@ fn feature_found_src(bytes: &[u8], name: &str) -> bool {
 // specified maximum we can put a cap on the virtual address space reservations
 // made.
 fn lock_pooling() -> impl Drop {
-    const MAX_CONCURRENT_POOLING: u32 = 8;
+    const MAX_CONCURRENT_POOLING: u32 = 4;
 
     static ACTIVE: Lazy<MyState> = Lazy::new(MyState::default);