Add MPK-protected stripes to the pooling allocator

crates/runtime/Cargo.toml

@@ -50,6 +50,7 @@ features = [
 
 [dev-dependencies]
 once_cell = { workspace = true }
+proptest = "1.0.0"
 
 [build-dependencies]
 cc = "1.0"
@@ -61,9 +62,6 @@ async = ["wasmtime-fiber"]
 # Enables support for the pooling instance allocator
 pooling-allocator = []
 
-component-model = [
-  "wasmtime-environ/component-model",
-  "dep:encoding_rs",
-]
+component-model = ["wasmtime-environ/component-model", "dep:encoding_rs"]
 
 wmemcheck = []

crates/runtime/proptest-regressions/instance/allocator/pooling/memory_pool.txt

@@ -0,0 +1,8 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc 696808084287d5d58b85c60c4720227ab4dd83ada7be6841a67162023aaf4914 # shrinks to c = SlabConstraints { max_memory_bytes: 0, num_memory_slots: 1, num_pkeys_available: 0, guard_bytes: 9223372036854775808 }
+cc cf9f6c36659f7f56ed8ea646e8c699cbf46708cef6911cdd376418ad69ea1388 # shrinks to c = SlabConstraints { max_memory_bytes: 14161452635954640438, num_memory_slots: 0, num_pkeys_available: 0, guard_bytes: 4285291437754911178 }

crates/runtime/src/instance/allocator.rs

@@ -1,6 +1,7 @@
 use crate::imports::Imports;
 use crate::instance::{Instance, InstanceHandle};
 use crate::memory::Memory;
+use crate::mpk::ProtectionKey;
 use crate::table::Table;
 use crate::{CompiledModuleId, ModuleRuntimeInfo, Store};
 use anyhow::{anyhow, bail, Result};
@@ -59,6 +60,10 @@ pub struct InstanceAllocationRequest<'a> {
 
     /// Indicates '--wmemcheck' flag.
     pub wmemcheck: bool,
+
+    /// Request that the instance's memories be protected by a specific
+    /// protection key.
+    pub pkey: Option<ProtectionKey>,
 }
 
 /// A pointer to a Store. This Option<*mut dyn Store> is wrapped in a struct
@@ -267,6 +272,24 @@ pub unsafe trait InstanceAllocatorImpl {
     /// Primarily present for the pooling allocator to remove mappings of
     /// this module from slots in linear memory.
     fn purge_module(&self, module: CompiledModuleId);
+
+    /// Use the next available protection key.
+    ///
+    /// The pooling allocator can use memory protection keys (MPK) for
+    /// compressing the guard regions protecting against OOB. Each
+    /// pool-allocated store needs its own key.
+    fn next_available_pkey(&self) -> Option<ProtectionKey>;
+
+    /// Restrict access to memory regions protected by `pkey`.
+    ///
+    /// This is useful for the pooling allocator, which can use memory
+    /// protection keys (MPK). Note: this may still allow access to other
+    /// protection keys, such as the default kernel key; see implementations of
+    /// this.
+    fn restrict_to_pkey(&self, pkey: ProtectionKey);
+
+    /// Allow access to memory regions protected by any protection key.
+    fn allow_all_pkeys(&self);
 }
 
 /// A thing that can allocate instances.

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

@@ -3,6 +3,7 @@ use super::{
 };
 use crate::instance::RuntimeMemoryCreator;
 use crate::memory::{DefaultMemoryCreator, Memory};
+use crate::mpk::ProtectionKey;
 use crate::table::Table;
 use crate::CompiledModuleId;
 use anyhow::Result;
@@ -151,4 +152,25 @@ unsafe impl InstanceAllocatorImpl for OnDemandInstanceAllocator {
     }
 
     fn purge_module(&self, _: CompiledModuleId) {}
+
+    fn next_available_pkey(&self) -> Option<ProtectionKey> {
+        // The on-demand allocator cannot use protection keys--it requires
+        // back-to-back allocation of memory slots that this allocator cannot
+        // guarantee.
+        None
+    }
+
+    fn restrict_to_pkey(&self, _: ProtectionKey) {
+        // The on-demand allocator cannot use protection keys; an on-demand
+        // allocator will never hand out protection keys to the stores its
+        // engine creates.
+        unreachable!()
+    }
+
+    fn allow_all_pkeys(&self) {
+        // The on-demand allocator cannot use protection keys; an on-demand
+        // allocator will never hand out protection keys to the stores its
+        // engine creates.
+        unreachable!()
+    }
 }

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

@@ -1,11 +1,62 @@
 //! Implements the pooling instance allocator.
 //!
-//! The pooling instance allocator maps memory in advance
-//! and allocates instances, memories, tables, and stacks from
-//! a pool of available resources.
+//! The pooling instance allocator maps memory in advance and allocates
+//! instances, memories, tables, and stacks from a pool of available resources.
+//! Using the pooling instance allocator can speed up module instantiation when
+//! modules can be constrained based on configurable limits
+//! ([`InstanceLimits`]). Each new instance is stored in a "slot"; as instances
+//! are allocated and freed, these slots are either filled or emptied:
 //!
-//! Using the pooling instance allocator can speed up module instantiation
-//! when modules can be constrained based on configurable limits.
+//! ```text
+//! ┌──────┬──────┬──────┬──────┬──────┐
+//! │Slot 0│Slot 1│Slot 2│Slot 3│......│
+//! └──────┴──────┴──────┴──────┴──────┘
+//! ```
+//!
+//! Note that these slots are a useful abstraction but not exactly how this is
+//! mapped to memory in fact. Each new instance _does_ get associated with a
+//! slot number (see uses of `index` and [`SlotId`] in this module) but the
+//! parts of the instances are stored in separate pools: memories in the
+//! [`MemoryPool`], tables in the [`TablePool`], etc. What ties these various
+//! parts together is the slot number generated by an [`IndexAllocator`] .
+//!
+//! The [`MemoryPool`] protects Wasmtime from out-of-bounds memory accesses by
+//! inserting inaccessible guard regions between memory slots. The
+//! [`MemoryPool`] documentation has a more in-depth chart but one can think of
+//! memories being laid out like the following:
+//!
+//! ```text
+//! ┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐
+//! │Guard│Mem 0│Guard│Mem 1│Guard│Mem 2│.....│Guard│
+//! └─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘
+//! ```
+//!
+//! To complicate matters, each instance can have multiple memories, multiple
+//! tables, etc. You might think these would be stored consecutively in their
+//! respective pools (for instance `n`, table 0 is at table pool slot `n + 0`
+//! and table 1 is at `n + 1`), but for memories this is not the case. With
+//! protection keys enabled, memories do not need interleaved guard regions
+//! because the protection key will signal a fault if the wrong memory is
+//! accessed. Instead, the pooling allocator "stripes" the memories with
+//! different protection keys.
+//!
+//! This concept, dubbed [ColorGuard] in the original paper, relies on careful
+//! calculation of the memory sizes to prevent any "overlapping access": there
+//! are limited protection keys available (15) so the next memory using the same
+//! key must be at least as far away as the guard region we would insert
+//! otherwise. This ends up looking like the following, where a store for
+//! instance 0 (`I0`) "stripes" two memories (`M0` and `M1`) with the same
+//! protection key 1 and far enough apart to signal an OOB access:
+//!
+//! ```text
+//! ┌─────┬─────┬─────┬─────┬────────────────┬─────┬─────┬─────┐
+//! │.....│I0:M1│.....│.....│.<enough slots>.│I0:M2│.....│.....│
+//! ├─────┼─────┼─────┼─────┼────────────────┼─────┼─────┼─────┤
+//! │.....│key 1│key 2│key 3│..<more keys>...│key 1│key 2│.....│
+//! └─────┴─────┴─────┴─────┴────────────────┴─────┴─────┴─────┘
+//! ```
+//!
+//! [ColorGuard]: https://plas2022.github.io/files/pdf/SegueColorGuard.pdf
 
 mod index_allocator;
 mod memory_pool;
@@ -27,7 +78,11 @@ cfg_if::cfg_if! {
 use super::{
     InstanceAllocationRequest, InstanceAllocatorImpl, MemoryAllocationIndex, TableAllocationIndex,
 };
-use crate::{instance::Instance, CompiledModuleId, Memory, Table};
+use crate::{
+    instance::Instance,
+    mpk::{self, MpkEnabled, ProtectionKey, ProtectionMask},
+    CompiledModuleId, Memory, Table,
+};
 use anyhow::{bail, Result};
 use memory_pool::MemoryPool;
 use std::{
@@ -162,6 +217,8 @@ pub struct PoolingInstanceAllocatorConfig {
     pub linear_memory_keep_resident: usize,
     /// Same as `linear_memory_keep_resident` but for tables.
     pub table_keep_resident: usize,
+    /// Whether to enable memory protection keys.
+    pub memory_protection_keys: MpkEnabled,
 }
 
 impl Default for PoolingInstanceAllocatorConfig {
@@ -174,15 +231,18 @@ impl Default for PoolingInstanceAllocatorConfig {
             async_stack_keep_resident: 0,
             linear_memory_keep_resident: 0,
             table_keep_resident: 0,
+            memory_protection_keys: MpkEnabled::Disable,
         }
     }
 }
 
 /// Implements the pooling instance allocator.
 ///
-/// This allocator internally maintains pools of instances, memories, tables, and stacks.
+/// This allocator internally maintains pools of instances, memories, tables,
+/// and stacks.
 ///
-/// Note: the resource pools are manually dropped so that the fault handler terminates correctly.
+/// Note: the resource pools are manually dropped so that the fault handler
+/// terminates correctly.
 #[derive(Debug)]
 pub struct PoolingInstanceAllocator {
     limits: InstanceLimits,
@@ -533,6 +593,18 @@ unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator {
     fn purge_module(&self, module: CompiledModuleId) {
         self.memories.purge_module(module);
     }
+
+    fn next_available_pkey(&self) -> Option<ProtectionKey> {
+        self.memories.next_available_pkey()
+    }
+
+    fn restrict_to_pkey(&self, pkey: ProtectionKey) {
+        mpk::allow(ProtectionMask::zero().or(pkey));
+    }
+
+    fn allow_all_pkeys(&self) {
+        mpk::allow(ProtectionMask::all());
+    }
 }
 
 #[cfg(test)]

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

@@ -167,6 +167,12 @@ impl ModuleAffinityIndexAllocator {
         }))
     }
 
+    /// How many slots can this allocator allocate?
+    pub fn len(&self) -> usize {
+        let inner = self.0.lock().unwrap();
+        inner.slot_state.len()
+    }
+
     /// Are zero slots in use right now?
     pub fn is_empty(&self) -> bool {
         let inner = self.0.lock().unwrap();
@@ -299,8 +305,16 @@ impl ModuleAffinityIndexAllocator {
         });
     }
 
-    /// For testing only, we want to be able to assert what is on the
-    /// single freelist, for the policies that keep just one.
+    /// Return the number of empty slots available in this allocator.
+    #[cfg(test)]
+    pub fn num_empty_slots(&self) -> usize {
+        let inner = self.0.lock().unwrap();
+        let total_slots = inner.slot_state.len();
+        (total_slots - inner.last_cold as usize) + inner.unused_warm_slots as usize
+    }
+
+    /// For testing only, we want to be able to assert what is on the single
+    /// freelist, for the policies that keep just one.
     #[cfg(test)]
     #[allow(unused)]
     pub(crate) fn testing_freelist(&self) -> Vec<SlotId> {
@@ -311,8 +325,8 @@ impl ModuleAffinityIndexAllocator {
             .collect()
     }
 
-    /// For testing only, get the list of all modules with at least
-    /// one slot with affinity for that module.
+    /// For testing only, get the list of all modules with at least one slot
+    /// with affinity for that module.
     #[cfg(test)]
     pub(crate) fn testing_module_affinity_list(&self) -> Vec<MemoryInModule> {
         let inner = self.0.lock().unwrap();
@@ -475,7 +489,9 @@ mod test {
     fn test_next_available_allocation_strategy() {
         for size in 0..20 {
             let state = ModuleAffinityIndexAllocator::new(size, 0);
+            assert_eq!(state.num_empty_slots() as u32, size);
             for i in 0..size {
+                assert_eq!(state.num_empty_slots() as u32, size - i);
                 assert_eq!(state.alloc(None).unwrap().index(), i as usize);
             }
             assert!(state.alloc(None).is_none());
@@ -496,6 +512,9 @@ mod test {
         assert_ne!(index1, index2);
 
         state.free(index1);
+        assert_eq!(state.num_empty_slots(), 99);
+
+        // Allocate to the same `index1` slot again.
         let index3 = state.alloc(Some(id1)).unwrap();
         assert_eq!(index3, index1);
         state.free(index3);
@@ -512,13 +531,14 @@ mod test {
         // for id1, and 98 empty. Allocate 100 for id2. The first
         // should be equal to the one we know was previously used for
         // id2. The next 99 are arbitrary.
-
+        assert_eq!(state.num_empty_slots(), 100);
         let mut indices = vec![];
         for _ in 0..100 {
             indices.push(state.alloc(Some(id2)).unwrap());
         }
         assert!(state.alloc(None).is_none());
         assert_eq!(indices[0], index2);
+        assert_eq!(state.num_empty_slots(), 0);
 
         for i in indices {
             state.free(i);