wasmtime/runtime/vm/mpk/enabled.rs
1//!
2
3use super::{pkru, sys};
4use crate::prelude::*;
5use std::sync::OnceLock;
6
7/// Check if the MPK feature is supported.
8pub fn is_supported() -> bool {
9 wasmtime_core::mpk::is_supported()
10}
11
12/// Allocate up to `max` protection keys.
13///
14/// This asks the kernel for all available keys up to `max` in a thread-safe way
15/// (we can expect 1-15; 0 is kernel-reserved). This avoids interference when
16/// multiple threads try to allocate keys at the same time (e.g., during
17/// testing). It also ensures that a single copy of the keys is reserved for the
18/// lifetime of the process. Because of this, `max` is only a hint to
19/// allocation: it only is effective on the first invocation of this function.
20///
21/// TODO: this is not the best-possible design. This creates global state that
22/// would prevent any other code in the process from using protection keys; the
23/// `KEYS` are never deallocated from the system with `pkey_dealloc`.
24pub fn keys(max: usize) -> &'static [ProtectionKey] {
25 let keys = KEYS.get_or_init(|| {
26 let mut allocated = vec![];
27 if is_supported() {
28 while allocated.len() < max {
29 if let Ok(key_id) = sys::pkey_alloc(0, 0) {
30 debug_assert!(key_id < 16);
31 // UNSAFETY: here we unsafely assume that the
32 // system-allocated pkey will exist forever.
33 allocated.push(ProtectionKey {
34 id: key_id,
35 stripe: allocated.len().try_into().unwrap(),
36 });
37 } else {
38 break;
39 }
40 }
41 }
42 allocated
43 });
44 &keys[..keys.len().min(max)]
45}
46static KEYS: OnceLock<Vec<ProtectionKey>> = OnceLock::new();
47
48/// Only allow access to pages marked by the keys set in `mask`.
49///
50/// Any accesses to pages marked by another key will result in a `SIGSEGV`
51/// fault.
52pub fn allow(mask: ProtectionMask) {
53 let previous = if log::log_enabled!(log::Level::Trace) {
54 pkru::read()
55 } else {
56 0
57 };
58 pkru::write(mask.0);
59 log::trace!("PKRU change: {:#034b} => {:#034b}", previous, pkru::read());
60}
61
62/// Retrieve the current protection mask.
63#[cfg(feature = "async")]
64pub fn current_mask() -> ProtectionMask {
65 ProtectionMask(pkru::read())
66}
67
68/// An MPK protection key.
69///
70/// The expected usage is:
71/// - receive system-allocated keys from [`keys`]
72/// - mark some regions of memory as accessible with [`ProtectionKey::protect`]
73/// - [`allow`] or disallow access to the memory regions using a
74/// [`ProtectionMask`]; any accesses to unmarked pages result in a fault
75/// - drop the key
76#[derive(Clone, Copy, Debug)]
77pub struct ProtectionKey {
78 id: u32,
79 stripe: u32,
80}
81
82impl ProtectionKey {
83 /// Mark a page as protected by this [`ProtectionKey`].
84 ///
85 /// This "colors" the pages of `region` via a kernel `pkey_mprotect` call to
86 /// only allow reads and writes when this [`ProtectionKey`] is activated
87 /// (see [`allow`]).
88 ///
89 /// # Errors
90 ///
91 /// This will fail if the region is not page aligned or for some unknown
92 /// kernel reason.
93 pub fn protect(&self, region: &mut [u8]) -> Result<()> {
94 let addr = region.as_mut_ptr() as usize;
95 let len = region.len();
96 let prot = sys::PROT_NONE;
97 sys::pkey_mprotect(addr, len, prot, self.id).with_context(|| {
98 format!(
99 "failed to mark region with pkey (addr = {addr:#x}, len = {len}, prot = {prot:#b})"
100 )
101 })
102 }
103
104 /// Convert the [`ProtectionKey`] to its 0-based index; this is useful for
105 /// determining which allocation "stripe" a key belongs to.
106 ///
107 /// This function assumes that the kernel has allocated key 0 for itself.
108 pub fn as_stripe(&self) -> usize {
109 self.stripe as usize
110 }
111
112 /// Re-apply this [`ProtectionKey`] to a region that has just been re-mapped.
113 ///
114 /// A fresh `mmap` over a region discards that region's protection key,
115 /// leaving it associated with the default key 0 which is always accessible.
116 /// Any code that maps over pkey-protected memory must therefore call this
117 /// afterwards to restore the key, otherwise the memory becomes readable and
118 /// writable from any stripe.
119 ///
120 /// Note that `mprotect` (unlike `mmap`) preserves the existing key, so only
121 /// `mmap` call sites need this.
122 ///
123 /// # Safety
124 ///
125 /// `addr` must be page-aligned and `addr..addr + len` must describe a mapped
126 /// region owned by the caller. `readwrite` must match the page protections
127 /// the region was just mapped with, since this overwrites them.
128 pub unsafe fn reprotect(&self, addr: usize, len: usize, readwrite: bool) -> Result<()> {
129 let prot = if readwrite {
130 sys::PROT_READ | sys::PROT_WRITE
131 } else {
132 sys::PROT_NONE
133 };
134 sys::pkey_mprotect(addr, len, prot, self.id).with_context(|| {
135 format!(
136 "failed to restore pkey on region (addr = {addr:#x}, len = {len}, prot = {prot:#b})"
137 )
138 })
139 }
140}
141
142/// A bit field indicating which protection keys should be allowed and disabled.
143///
144/// The internal representation makes it easy to use [`ProtectionMask`] directly
145/// with the PKRU register. When bits `n` and `n+1` are set, it means the
146/// protection key is *not* allowed (see the PKRU write and access disabled
147/// bits).
148pub struct ProtectionMask(u32);
149impl ProtectionMask {
150 /// Allow access from all protection keys.
151 #[inline]
152 pub fn all() -> Self {
153 Self(pkru::ALLOW_ACCESS)
154 }
155
156 /// Only allow access to memory protected with protection key 0; note that
157 /// this does not mean "none" but rather allows access from the default
158 /// kernel protection key.
159 #[inline]
160 pub fn zero() -> Self {
161 Self(pkru::DISABLE_ACCESS ^ 0b11)
162 }
163
164 /// Include `pkey` as another allowed protection key in the mask.
165 #[inline]
166 pub fn or(self, pkey: ProtectionKey) -> Self {
167 let mask = pkru::DISABLE_ACCESS ^ 0b11 << (pkey.id * 2);
168 Self(self.0 & mask)
169 }
170}
171
172/// Helper macro for skipping tests on systems that do not have MPK enabled
173/// (e.g., older architecture, disabled by kernel, etc.)
174#[cfg(test)]
175macro_rules! skip_if_mpk_unavailable {
176 () => {
177 if !crate::runtime::vm::mpk::is_supported() {
178 println!("> mpk is not supported: ignoring test");
179 return;
180 }
181 };
182}
183/// Necessary for inter-module access.
184#[cfg(test)]
185pub(crate) use skip_if_mpk_unavailable;
186
187#[cfg(test)]
188mod tests {
189 use super::*;
190
191 #[test]
192 fn check_is_supported() {
193 println!("is pku supported = {}", is_supported());
194 if std::env::var("WASMTIME_TEST_FORCE_MPK").is_ok() {
195 assert!(is_supported());
196 }
197 }
198
199 #[test]
200 fn check_initialized_keys() {
201 if is_supported() {
202 assert!(!keys(15).is_empty())
203 }
204 }
205
206 #[test]
207 fn check_invalid_mark() {
208 skip_if_mpk_unavailable!();
209 let pkey = keys(15)[0];
210 let unaligned_region = unsafe {
211 let addr = 1 as *mut u8; // this is not page-aligned!
212 let len = 1;
213 std::slice::from_raw_parts_mut(addr, len)
214 };
215 let result = pkey.protect(unaligned_region);
216 assert!(result.is_err());
217 assert_eq!(
218 result.unwrap_err().to_string(),
219 "failed to mark region with pkey (addr = 0x1, len = 1, prot = 0b0)"
220 );
221 }
222
223 #[test]
224 fn check_masking() {
225 skip_if_mpk_unavailable!();
226 let original = pkru::read();
227
228 allow(ProtectionMask::all());
229 assert_eq!(0, pkru::read());
230
231 allow(ProtectionMask::all().or(ProtectionKey { id: 5, stripe: 0 }));
232 assert_eq!(0, pkru::read());
233
234 allow(ProtectionMask::zero());
235 assert_eq!(0b11111111_11111111_11111111_11111100, pkru::read());
236
237 allow(ProtectionMask::zero().or(ProtectionKey { id: 5, stripe: 0 }));
238 assert_eq!(0b11111111_11111111_11110011_11111100, pkru::read());
239
240 // Reset the PKRU state to what we originally observed.
241 pkru::write(original);
242 }
243}