1 //!
2 
3 use super::{pkru, sys};
4 use crate::prelude::*;
5 use std::sync::OnceLock;
6 
7 /// Check if the MPK feature is supported.
8 pub 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`.
24 pub 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 }
46 static 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.
52 pub 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 pub fn current_mask() -> ProtectionMask {
64     ProtectionMask(pkru::read())
65 }
66 
67 /// An MPK protection key.
68 ///
69 /// The expected usage is:
70 /// - receive system-allocated keys from [`keys`]
71 /// - mark some regions of memory as accessible with [`ProtectionKey::protect`]
72 /// - [`allow`] or disallow access to the memory regions using a
73 ///   [`ProtectionMask`]; any accesses to unmarked pages result in a fault
74 /// - drop the key
75 #[derive(Clone, Copy, Debug)]
76 pub struct ProtectionKey {
77     id: u32,
78     stripe: u32,
79 }
80 
81 impl ProtectionKey {
82     /// Mark a page as protected by this [`ProtectionKey`].
83     ///
84     /// This "colors" the pages of `region` via a kernel `pkey_mprotect` call to
85     /// only allow reads and writes when this [`ProtectionKey`] is activated
86     /// (see [`allow`]).
87     ///
88     /// # Errors
89     ///
90     /// This will fail if the region is not page aligned or for some unknown
91     /// kernel reason.
92     pub fn protect(&self, region: &mut [u8]) -> Result<()> {
93         let addr = region.as_mut_ptr() as usize;
94         let len = region.len();
95         let prot = sys::PROT_NONE;
96         sys::pkey_mprotect(addr, len, prot, self.id).with_context(|| {
97             format!(
98                 "failed to mark region with pkey (addr = {addr:#x}, len = {len}, prot = {prot:#b})"
99             )
100         })
101     }
102 
103     /// Convert the [`ProtectionKey`] to its 0-based index; this is useful for
104     /// determining which allocation "stripe" a key belongs to.
105     ///
106     /// This function assumes that the kernel has allocated key 0 for itself.
107     pub fn as_stripe(&self) -> usize {
108         self.stripe as usize
109     }
110 }
111 
112 /// A bit field indicating which protection keys should be allowed and disabled.
113 ///
114 /// The internal representation makes it easy to use [`ProtectionMask`] directly
115 /// with the PKRU register. When bits `n` and `n+1` are set, it means the
116 /// protection key is *not* allowed (see the PKRU write and access disabled
117 /// bits).
118 pub struct ProtectionMask(u32);
119 impl ProtectionMask {
120     /// Allow access from all protection keys.
121     #[inline]
122     pub fn all() -> Self {
123         Self(pkru::ALLOW_ACCESS)
124     }
125 
126     /// Only allow access to memory protected with protection key 0; note that
127     /// this does not mean "none" but rather allows access from the default
128     /// kernel protection key.
129     #[inline]
130     pub fn zero() -> Self {
131         Self(pkru::DISABLE_ACCESS ^ 0b11)
132     }
133 
134     /// Include `pkey` as another allowed protection key in the mask.
135     #[inline]
136     pub fn or(self, pkey: ProtectionKey) -> Self {
137         let mask = pkru::DISABLE_ACCESS ^ 0b11 << (pkey.id * 2);
138         Self(self.0 & mask)
139     }
140 }
141 
142 /// Helper macro for skipping tests on systems that do not have MPK enabled
143 /// (e.g., older architecture, disabled by kernel, etc.)
144 #[cfg(test)]
145 macro_rules! skip_if_mpk_unavailable {
146     () => {
147         if !crate::runtime::vm::mpk::is_supported() {
148             println!("> mpk is not supported: ignoring test");
149             return;
150         }
151     };
152 }
153 /// Necessary for inter-module access.
154 #[cfg(test)]
155 pub(crate) use skip_if_mpk_unavailable;
156 
157 #[cfg(test)]
158 mod tests {
159     use super::*;
160 
161     #[test]
162     fn check_is_supported() {
163         println!("is pku supported = {}", is_supported());
164         if std::env::var("WASMTIME_TEST_FORCE_MPK").is_ok() {
165             assert!(is_supported());
166         }
167     }
168 
169     #[test]
170     fn check_initialized_keys() {
171         if is_supported() {
172             assert!(!keys(15).is_empty())
173         }
174     }
175 
176     #[test]
177     fn check_invalid_mark() {
178         skip_if_mpk_unavailable!();
179         let pkey = keys(15)[0];
180         let unaligned_region = unsafe {
181             let addr = 1 as *mut u8; // this is not page-aligned!
182             let len = 1;
183             std::slice::from_raw_parts_mut(addr, len)
184         };
185         let result = pkey.protect(unaligned_region);
186         assert!(result.is_err());
187         assert_eq!(
188             result.unwrap_err().to_string(),
189             "failed to mark region with pkey (addr = 0x1, len = 1, prot = 0b0)"
190         );
191     }
192 
193     #[test]
194     fn check_masking() {
195         skip_if_mpk_unavailable!();
196         let original = pkru::read();
197 
198         allow(ProtectionMask::all());
199         assert_eq!(0, pkru::read());
200 
201         allow(ProtectionMask::all().or(ProtectionKey { id: 5, stripe: 0 }));
202         assert_eq!(0, pkru::read());
203 
204         allow(ProtectionMask::zero());
205         assert_eq!(0b11111111_11111111_11111111_11111100, pkru::read());
206 
207         allow(ProtectionMask::zero().or(ProtectionKey { id: 5, stripe: 0 }));
208         assert_eq!(0b11111111_11111111_11110011_11111100, pkru::read());
209 
210         // Reset the PKRU state to what we originally observed.
211         pkru::write(original);
212     }
213 }
214