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    cfg!(target_os = "linux") && cfg!(target_arch = "x86_64") && pkru::has_cpuid_bit_set()
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
113/// A bit field indicating which protection keys should be allowed and disabled.
114///
115/// The internal representation makes it easy to use [`ProtectionMask`] directly
116/// with the PKRU register. When bits `n` and `n+1` are set, it means the
117/// protection key is *not* allowed (see the PKRU write and access disabled
118/// bits).
119pub struct ProtectionMask(u32);
120impl ProtectionMask {
121    /// Allow access from all protection keys.
122    #[inline]
123    pub fn all() -> Self {
124        Self(pkru::ALLOW_ACCESS)
125    }
126
127    /// Only allow access to memory protected with protection key 0; note that
128    /// this does not mean "none" but rather allows access from the default
129    /// kernel protection key.
130    #[inline]
131    pub fn zero() -> Self {
132        Self(pkru::DISABLE_ACCESS ^ 0b11)
133    }
134
135    /// Include `pkey` as another allowed protection key in the mask.
136    #[inline]
137    pub fn or(self, pkey: ProtectionKey) -> Self {
138        let mask = pkru::DISABLE_ACCESS ^ 0b11 << (pkey.id * 2);
139        Self(self.0 & mask)
140    }
141}
142
143/// Helper macro for skipping tests on systems that do not have MPK enabled
144/// (e.g., older architecture, disabled by kernel, etc.)
145#[cfg(test)]
146macro_rules! skip_if_mpk_unavailable {
147    () => {
148        if !crate::runtime::vm::mpk::is_supported() {
149            println!("> mpk is not supported: ignoring test");
150            return;
151        }
152    };
153}
154/// Necessary for inter-module access.
155#[cfg(test)]
156pub(crate) use skip_if_mpk_unavailable;
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    #[test]
163    fn check_is_supported() {
164        println!("is pku supported = {}", is_supported());
165        if std::env::var("WASMTIME_TEST_FORCE_MPK").is_ok() {
166            assert!(is_supported());
167        }
168    }
169
170    #[test]
171    fn check_initialized_keys() {
172        if is_supported() {
173            assert!(!keys(15).is_empty())
174        }
175    }
176
177    #[test]
178    fn check_invalid_mark() {
179        skip_if_mpk_unavailable!();
180        let pkey = keys(15)[0];
181        let unaligned_region = unsafe {
182            let addr = 1 as *mut u8; // this is not page-aligned!
183            let len = 1;
184            std::slice::from_raw_parts_mut(addr, len)
185        };
186        let result = pkey.protect(unaligned_region);
187        assert!(result.is_err());
188        assert_eq!(
189            result.unwrap_err().to_string(),
190            "failed to mark region with pkey (addr = 0x1, len = 1, prot = 0b0)"
191        );
192    }
193
194    #[test]
195    fn check_masking() {
196        skip_if_mpk_unavailable!();
197        let original = pkru::read();
198
199        allow(ProtectionMask::all());
200        assert_eq!(0, pkru::read());
201
202        allow(ProtectionMask::all().or(ProtectionKey { id: 5, stripe: 0 }));
203        assert_eq!(0, pkru::read());
204
205        allow(ProtectionMask::zero());
206        assert_eq!(0b11111111_11111111_11111111_11111100, pkru::read());
207
208        allow(ProtectionMask::zero().or(ProtectionKey { id: 5, stripe: 0 }));
209        assert_eq!(0b11111111_11111111_11110011_11111100, pkru::read());
210
211        // Reset the PKRU state to what we originally observed.
212        pkru::write(original);
213    }
214}
wasmtime/runtime/vm/mpk/enabled.rs

wasmtime/runtime/vm/mpk/
enabled.rs
