mpk: restore protection keys after mmap'ing memory images

crates/wasmtime/src/runtime/vm/cow.rs

@@ -4,6 +4,7 @@
 use super::sys::DecommitBehavior;
 use crate::Engine;
 use crate::prelude::*;
+use crate::runtime::vm::mpk::ProtectionKey;
 use crate::runtime::vm::sys::vm::{self, MemoryImageSource, PageMap, reset_with_pagemap};
 use crate::runtime::vm::{
     HostAlignedByteCount, MmapOffset, ModuleMemoryImageSource, host_page_size,
@@ -337,6 +338,14 @@ pub struct MemoryImageSlot {
     /// initial image content, as appropriate. Everything between
     /// `self.accessible` and `self.static_size` is inaccessible.
     dirty: bool,
+
+    /// The MPK protection key that this slot's stripe was colored with, if the
+    /// pooling allocator is striping memory with protection keys.
+    ///
+    /// This must be re-applied after every `mmap` performed on this slot, since
+    /// `mmap` resets the affected pages back to the default key 0 which is
+    /// accessible from every stripe.
+    pkey: Option<ProtectionKey>,
 }
 
 impl fmt::Debug for MemoryImageSlot {
@@ -363,13 +372,15 @@ impl MemoryImageSlot {
         base: MmapOffset,
         accessible: HostAlignedByteCount,
         static_size: usize,
+        pkey: Option<ProtectionKey>,
     ) -> Self {
         MemoryImageSlot {
             base,
             static_size,
             accessible,
             image: None,
             dirty: false,
+            pkey,
         }
     }
 
@@ -489,6 +500,11 @@ impl MemoryImageSlot {
                     unsafe {
                         image.map_at(&self.base)?;
                     }
+                    // `map_at` above `mmap`'d over part of this slot, which
+                    // reset those pages to the default protection key. Restore
+                    // this slot's key so the image is not left accessible to
+                    // every other stripe.
+                    self.reapply_pkey(image.linear_memory_offset, image.len, true)?;
                 }
             }
             self.image = maybe_image.cloned();
@@ -506,6 +522,9 @@ impl MemoryImageSlot {
             unsafe {
                 image.remap_as_zeros_at(self.base.as_mut_ptr())?;
             }
+            // As in `instantiate`, the `mmap` above dropped this slot's
+            // protection key over the image's range, so restore it.
+            self.reapply_pkey(image.linear_memory_offset, image.len, true)?;
             self.image = None;
         }
         Ok(())
@@ -682,6 +701,43 @@ impl MemoryImageSlot {
         Ok(())
     }
 
+    /// Re-color `offset..offset + len` within this slot with this slot's MPK
+    /// protection key, if any.
+    ///
+    /// This is a no-op unless the pooling allocator is striping memory with
+    /// protection keys. It must be called after every `mmap` that lands inside
+    /// this slot: `mmap` associates the pages it replaces with the default key
+    /// 0, which is accessible regardless of which stripe is currently active,
+    /// so skipping this would let one instance read and write another
+    /// instance's memory.
+    ///
+    /// Note that `mprotect` preserves the existing key, so `set_protection`
+    /// does not need this treatment.
+    fn reapply_pkey(
+        &self,
+        offset: HostAlignedByteCount,
+        len: HostAlignedByteCount,
+        readwrite: bool,
+    ) -> Result<()> {
+        let Some(pkey) = self.pkey else {
+            return Ok(());
+        };
+        if len.is_zero() {
+            return Ok(());
+        }
+        // `mmap` rounds lengths up to a page boundary, so the restored range is
+        // allowed to extend to the end of the slot's final page.
+        debug_assert!(
+            offset.byte_count() + len.byte_count()
+                <= self.static_size.next_multiple_of(host_page_size())
+        );
+        unsafe {
+            let start = self.base.as_mut_ptr().add(offset.byte_count());
+            pkey.reprotect(start.addr(), len.byte_count(), readwrite)?;
+        }
+        Ok(())
+    }
+
     pub(crate) fn has_image(&self) -> bool {
         self.image.is_some()
     }
@@ -704,6 +760,11 @@ impl MemoryImageSlot {
             vm::erase_existing_mapping(self.base.as_mut_ptr(), self.static_size)?;
         }
 
+        // The `mmap` above covers the whole slot and left it inaccessible, so
+        // restore this slot's protection key across the same range.
+        let static_size = HostAlignedByteCount::new_rounded_up(self.static_size)?;
+        self.reapply_pkey(HostAlignedByteCount::ZERO, static_size, false)?;
+
         self.image = None;
         self.accessible = HostAlignedByteCount::ZERO;
 
@@ -813,8 +874,12 @@ mod test {
         // 4 MiB mmap'd area, not accessible
         let mmap = mmap_4mib_inaccessible();
         // Create a MemoryImageSlot on top of it
-        let mut memfd =
-            MemoryImageSlot::create(mmap.zero_offset(), HostAlignedByteCount::ZERO, 4 << 20);
+        let mut memfd = MemoryImageSlot::create(
+            mmap.zero_offset(),
+            HostAlignedByteCount::ZERO,
+            4 << 20,
+            None,
+        );
         assert!(!memfd.is_dirty());
         // instantiate with 64 KiB initial size
         memfd
@@ -872,8 +937,12 @@ mod test {
         // 4 MiB mmap'd area, not accessible
         let mmap = mmap_4mib_inaccessible();
         // Create a MemoryImageSlot on top of it
-        let mut memfd =
-            MemoryImageSlot::create(mmap.zero_offset(), HostAlignedByteCount::ZERO, 4 << 20);
+        let mut memfd = MemoryImageSlot::create(
+            mmap.zero_offset(),
+            HostAlignedByteCount::ZERO,
+            4 << 20,
+            None,
+        );
         // Create an image with some data.
         let image = Arc::new(create_memfd_with_data(page_size, &[1, 2, 3, 4]).unwrap());
         // Instantiate with this image
@@ -983,8 +1052,12 @@ mod test {
             ..Tunables::default_miri()
         };
         let mmap = mmap_4mib_inaccessible();
-        let mut memfd =
-            MemoryImageSlot::create(mmap.zero_offset(), HostAlignedByteCount::ZERO, 4 << 20);
+        let mut memfd = MemoryImageSlot::create(
+            mmap.zero_offset(),
+            HostAlignedByteCount::ZERO,
+            4 << 20,
+            None,
+        );
 
         // Test basics with the image
         for image_off in [0, page_size, page_size * 2] {
@@ -1057,8 +1130,12 @@ mod test {
         };
 
         let mmap = mmap_4mib_inaccessible();
-        let mut memfd =
-            MemoryImageSlot::create(mmap.zero_offset(), HostAlignedByteCount::ZERO, 4 << 20);
+        let mut memfd = MemoryImageSlot::create(
+            mmap.zero_offset(),
+            HostAlignedByteCount::ZERO,
+            4 << 20,
+            None,
+        );
         let image = Arc::new(create_memfd_with_data(page_size, &[1, 2, 3, 4]).unwrap());
         let initial = 64 << 10;
 
@@ -1166,8 +1243,12 @@ mod test {
         };
         let mmap = mmap_4mib_inaccessible();
         let mmap_len = page_size * 9;
-        let mut memfd =
-            MemoryImageSlot::create(mmap.zero_offset(), HostAlignedByteCount::ZERO, mmap_len);
+        let mut memfd = MemoryImageSlot::create(
+            mmap.zero_offset(),
+            HostAlignedByteCount::ZERO,
+            mmap_len,
+            None,
+        );
         let pagemap = PageMap::new();
         let pagemap = pagemap.as_ref();
 

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

@@ -589,6 +589,19 @@ impl MemoryPool {
         self.mapping.offset(offset).expect("offset is in bounds")
     }
 
+    /// Return the protection key that this slot's memory was striped with when
+    /// the pool was created, if any.
+    ///
+    /// This mirrors the striping performed in `new`: memory is only colored
+    /// when there are at least two stripes, and slot `i` is colored with the
+    /// `i % num_stripes`th key.
+    fn pkey_for_slot(&self, allocation_index: MemoryAllocationIndex) -> Option<ProtectionKey> {
+        if self.stripes.len() < 2 {
+            return None;
+        }
+        self.stripes[allocation_index.index() % self.stripes.len()].pkey
+    }
+
     /// Take ownership of the given image slot.
     ///
     /// This method is used when a `MemoryAllocationIndex` has been allocated
@@ -620,6 +633,7 @@ impl MemoryPool {
                 self.get_base(allocation_index),
                 HostAlignedByteCount::ZERO,
                 self.layout.max_memory_bytes.byte_count(),
+                self.pkey_for_slot(allocation_index),
             )
         });
 

crates/wasmtime/src/runtime/vm/memory.rs

@@ -570,8 +570,10 @@ impl LocalMemory {
                         }
                     };
 
+                    // Memories allocated on demand are never striped with MPK
+                    // protection keys, so this slot has no key to preserve.
                     let mut slot =
-                        MemoryImageSlot::create(mmap_base, byte_size, alloc.byte_capacity());
+                        MemoryImageSlot::create(mmap_base, byte_size, alloc.byte_capacity(), None);
                     slot.instantiate(alloc.byte_size(), Some(image), ty, memory_tunables)?;
                     Some(slot)
                 } else {

crates/wasmtime/src/runtime/vm/mpk/disabled.rs

@@ -1,7 +1,7 @@
 //! Noop implementations of MPK primitives for environments that do not support
 //! the feature.
 
-#[cfg(feature = "pooling-allocator")]
+#[cfg(any(feature = "pooling-allocator", has_virtual_memory))]
 use crate::prelude::*;
 
 #[cfg(feature = "pooling-allocator")]
@@ -34,6 +34,12 @@ impl ProtectionKey {
     pub fn as_stripe(&self) -> usize {
         match *self {}
     }
+    // Note: gated on `has_virtual_memory` rather than `pooling-allocator`
+    // because this is called from `cow.rs`, which is not pooling-specific.
+    #[cfg(has_virtual_memory)]
+    pub unsafe fn reprotect(&self, _: usize, _: usize, _: bool) -> Result<()> {
+        match *self {}
+    }
 }
 
 #[derive(Clone, Copy, Debug)]

crates/wasmtime/src/runtime/vm/mpk/enabled.rs

@@ -108,6 +108,35 @@ impl ProtectionKey {
     pub fn as_stripe(&self) -> usize {
         self.stripe as usize
     }
+
+    /// Re-apply this [`ProtectionKey`] to a region that has just been re-mapped.
+    ///
+    /// A fresh `mmap` over a region discards that region's protection key,
+    /// leaving it associated with the default key 0 which is always accessible.
+    /// Any code that maps over pkey-protected memory must therefore call this
+    /// afterwards to restore the key, otherwise the memory becomes readable and
+    /// writable from any stripe.
+    ///
+    /// Note that `mprotect` (unlike `mmap`) preserves the existing key, so only
+    /// `mmap` call sites need this.
+    ///
+    /// # Safety
+    ///
+    /// `addr` must be page-aligned and `addr..addr + len` must describe a mapped
+    /// region owned by the caller. `readwrite` must match the page protections
+    /// the region was just mapped with, since this overwrites them.
+    pub unsafe fn reprotect(&self, addr: usize, len: usize, readwrite: bool) -> Result<()> {
+        let prot = if readwrite {
+            sys::PROT_READ | sys::PROT_WRITE
+        } else {
+            sys::PROT_NONE
+        };
+        sys::pkey_mprotect(addr, len, prot, self.id).with_context(|| {
+            format!(
+                "failed to restore pkey on region (addr = {addr:#x}, len = {len}, prot = {prot:#b})"
+            )
+        })
+    }
 }
 
 /// A bit field indicating which protection keys should be allowed and disabled.

crates/wasmtime/src/runtime/vm/mpk/sys.rs

@@ -18,6 +18,14 @@ use std::io::Error;
 /// to start as `PROT_NONE`.
 pub const PROT_NONE: u32 = libc::PROT_NONE as u32; // == 0b0000;
 
+/// Protection mask allowing reads of pkey-protected memory (see `prot` in
+/// [`pkey_mprotect`]).
+pub const PROT_READ: u32 = libc::PROT_READ as u32; // == 0b0001;
+
+/// Protection mask allowing writes of pkey-protected memory (see `prot` in
+/// [`pkey_mprotect`]).
+pub const PROT_WRITE: u32 = libc::PROT_WRITE as u32; // == 0b0010;
+
 /// Allocate a new protection key in the Linux kernel ([docs]); returns the
 /// key ID.
 ///