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 #[cfg(feature = "async")]
64 pub 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)]
77 pub struct ProtectionKey {
78     id: u32,
79     stripe: u32,
80 }
81 
82 impl 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).
119 pub struct ProtectionMask(u32);
120 impl 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)]
146 macro_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)]
156 pub(crate) use skip_if_mpk_unavailable;
157 
158 #[cfg(test)]
159 mod 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 }
215