Revise MPK support check to allow non-Intel CPUs.

crates/core/Cargo.toml

@@ -15,9 +15,19 @@ hashbrown = { workspace = true, features = ["default-hasher"] }
 libm = { workspace = true }
 serde = { workspace = true, optional = true }
 
+[dev-dependencies]
+libtest-mimic = { workspace = true }
+
+[target.'cfg(unix)'.dev-dependencies]
+libc = { workspace = true }
+
 [lints]
 workspace = true
 
+[[test]]
+name = "mpk"
+harness = false
+
 [features]
 # Enable the use of the `std` crate.
 std = []
@@ -27,3 +37,5 @@ backtrace = ["std"]
 # `Error::from_anyhow` constructor.
 anyhow = ["dep:anyhow"]
 serde = ["dep:serde"]
+
+

crates/core/src/mpk.rs

@@ -8,37 +8,45 @@
 
 /// Returns `true` if Wasmtime's MPK support can be used on this host.
 pub fn is_supported() -> bool {
-    cfg!(target_os = "linux") && cpuid_pku_bit_set()
+    cfg!(target_os = "linux") && cpuid_pku_enabled()
 }
 
-/// Check the `ECX.PKU` flag (bit 3, zero-based) of the `07h` `CPUID` leaf; see
-/// the Intel Software Development Manual, vol 3a, section 2.7. This flag is
-/// only set on Intel CPUs, so this function also checks the `CPUID` vendor
-/// string.
+/// Check whether protection keys (pkus) are supported by the CPU and enabled by the OS.
+///
+/// This is exposed in the ECX result of cpuid leaf 07h:
+///
+/// - bit 3 (`PKU`) is set if the processor supports protection keys;
+/// - bit 4 (`OSPKE`) is set if the OS has also set `CR4.PKE`, which is what makes
+///   `RDPKRU`/`WRPKRU` usable. Without it these instructions raise `#UD` (`SIGILL`), for example in
+///   virtual machines whose hypervisor reports `PKU` but doesn't enable it.
+///
+/// This is consistent between Intel and AMD CPUs.
+/// Sources:
+/// - Intel SDM Volume 1, Chapter 21, Section 21.3, Table 21-20
+///   (PDF at https://cdrdv2.intel.com/v1/dl/getContent/671436).
+/// - AMD64 Architecture Programmer's Manual, Volume 3, Appendix E, section E.3.6
+///   (PDF at https://docs.amd.com/api/khub/documents/BqEoC2FadIAgbrUdEzDJIg/content).
 #[cfg(target_arch = "x86_64")]
-fn cpuid_pku_bit_set() -> bool {
-    is_intel_cpu() && {
-        #[allow(
-            unused_unsafe,
-            reason = "rust is transitioning to `__cpuid` being a safe function"
-        )]
-        let result = unsafe { core::arch::x86_64::__cpuid(0x07) };
-        (result.ecx & 0b1000) != 0
+fn cpuid_pku_enabled() -> bool {
+    use core::arch::x86_64::{__cpuid, __cpuid_count};
+
+    const PKU: u32 = 1 << 3;
+    const OSPKE: u32 = 1 << 4;
+
+    #[allow(
+        unused_unsafe,
+        reason = "rust is transitioning to `__cpuid` being a safe function"
+    )]
+    let max_leaf = unsafe { __cpuid(0) }.eax;
+    if max_leaf < 0x07 {
+        return false;
     }
+    #[allow(unused_unsafe, reason = "as above")]
+    let ecx = unsafe { __cpuid_count(0x07, 0) }.ecx;
+    ecx & (PKU | OSPKE) == PKU | OSPKE
 }
 
 #[cfg(not(target_arch = "x86_64"))]
-fn cpuid_pku_bit_set() -> bool {
+fn cpuid_pku_enabled() -> bool {
     false
 }
-
-/// Check the `CPUID` vendor string for `GenuineIntel`; see the Intel Software
-/// Development Manual, vol 2a, `CPUID` description.
-#[cfg(target_arch = "x86_64")]
-fn is_intel_cpu() -> bool {
-    #[allow(unused_unsafe, reason = "see above about __cpuid")]
-    let result = unsafe { core::arch::x86_64::__cpuid(0) };
-    result.ebx == u32::from_le_bytes(*b"Genu")
-        && result.edx == u32::from_le_bytes(*b"ineI")
-        && result.ecx == u32::from_le_bytes(*b"ntel")
-}

crates/core/tests/mpk.rs

@@ -0,0 +1,128 @@
+//! Checks that `mpk::is_supported` agrees with what the hardware and kernel
+//! actually allow.
+//!
+//! We test that `RDPKRU`/`WRPKRU` work only if `mpk::is_supported` returns true.
+//! This is required to avoid hard to debug SIGILLs if our is_supported logic was wrong.
+//!
+//! On unsupported platforms we just test that is_supported is false.
+//! There are platforms where MPK is supported by hardware and the OS (32-bit x86, FreeBSD, ARM targets),
+//! but Wasmtime has no backend for that.
+
+use libtest_mimic::{Arguments, Failed, Trial};
+use std::process::ExitStatus;
+use std::process::{Command, exit};
+
+/// Set in the child process to make it run the probe instead of the test.
+const PROBE_ENV_VAR: &str = "WASMTIME_TEST_MPK_PROBE";
+
+// Exit codes of the probe, besides success.
+pub const EXIT_PKEY_ALLOC_FAILED: i32 = 2;
+pub const EXIT_PKEY_FREE_FAILED: i32 = 3;
+pub const EXIT_MPK_NOT_SUPPORTED_BY_WASMTIME: i32 = 4;
+
+fn main() {
+    // We are the child that is supposed to perform the probe.
+    if std::env::var_os(PROBE_ENV_VAR).is_some() {
+        probe();
+        // probe always exits, this is unreachable
+    }
+    // We are the test entrypoint.
+    let run_test = || {
+        let output = Command::new(std::env::current_exe().unwrap())
+            .env(PROBE_ENV_VAR, "1")
+            .output()
+            .expect("failed to spawn the probe process");
+        let usable = if output.status.success() {
+            true
+        } else if is_illegal_instruction(&output.status)
+            || is_not_supported_by_wasmtime(&output.status)
+        {
+            false
+        } else {
+            return Err(Failed::from(format!(
+                "inconclusive probe for protection key support: {}\n\
+                 stdout:\n{}\nstderr:\n{}",
+                output.status,
+                String::from_utf8_lossy(&output.stdout),
+                String::from_utf8_lossy(&output.stderr),
+            )));
+        };
+
+        let is_supported = wasmtime_internal_core::mpk::is_supported();
+        if is_supported != usable {
+            Err(Failed::from(format!(
+                "`mpk::is_supported()` disagrees with whether protection keys can \
+                 actually be used on this host: is_supported() = {is_supported}, usable = {usable}",
+            )))
+        } else {
+            Ok(())
+        }
+    };
+    let trial = Trial::test("is_supported_matches_hardware_and_os", run_test)
+        // Has to be excluded on miri as we can't even spawn the subcommand there.
+        .with_ignored_flag(cfg!(miri));
+    libtest_mimic::run(&Arguments::from_args(), vec![trial]).exit();
+}
+
+fn is_not_supported_by_wasmtime(status: &ExitStatus) -> bool {
+    status
+        .code()
+        .is_some_and(|x| x == EXIT_MPK_NOT_SUPPORTED_BY_WASMTIME)
+}
+
+#[cfg(unix)]
+fn is_illegal_instruction(status: &ExitStatus) -> bool {
+    use std::os::unix::process::ExitStatusExt;
+    status.signal() == Some(libc::SIGILL)
+}
+
+#[cfg(not(unix))]
+fn is_illegal_instruction(_status: &ExitStatus) -> bool {
+    false
+}
+
+/// Checks `RDPKRU`, `WRPKRU`, and tries to allocate and free a pkey.
+/// Exits the process with 0 if everything worked, raises `SIGILL` otherwise.
+#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
+fn probe() -> ! {
+    use std::arch::asm;
+    use std::io;
+    // No crate offers rdpkru or wrpkru as a Rust function (neither rustix nor libc), so we need to
+    // rawdog asm.
+    let pkru: u32;
+    // SAFETY: reading PKRU has no side effects; if protection keys are not
+    // enabled this raises `SIGILL`, which is what the parent checks for.
+    unsafe {
+        asm!("rdpkru", in("ecx") 0u32, out("eax") pkru, out("edx") _,
+            options(nomem, nostack, preserves_flags));
+    }
+    // SAFETY: writing back the same value doesn't change any access rights.
+    unsafe {
+        asm!("wrpkru", in("eax") pkru, in("ecx") 0u32, in("edx") 0u32,
+            options(nomem, nostack, preserves_flags));
+    }
+
+    // SAFETY: `pkey_alloc(flags = 0, access_rights = 0)` only allocates a key.
+    let key = unsafe { libc::syscall(libc::SYS_pkey_alloc, 0, 0) };
+    if key < 0 {
+        eprintln!("pkey_alloc failed: {}", io::Error::last_os_error());
+        exit(EXIT_PKEY_ALLOC_FAILED);
+    }
+    // SAFETY: frees the key just allocated, which nothing uses.
+    let ret = unsafe { libc::syscall(libc::SYS_pkey_free, key as usize) };
+    if ret != 0 {
+        eprintln!("pkey_free failed: {}", io::Error::last_os_error());
+        exit(EXIT_PKEY_FREE_FAILED);
+    }
+    exit(0)
+}
+
+/// Some platforms like ARM or FreeBSD support MPKs, but with different APIs that wasmtime
+/// currently does not hook into.
+/// We just return a special error on those.
+/// If wasmtime gets MPK support on those platforms, this should be updated to check the
+/// capabilities of the relevant API.
+#[cfg(not(all(target_os = "linux", target_arch = "x86_64")))]
+fn probe() -> ! {
+    exit(EXIT_MPK_NOT_SUPPORTED_BY_WASMTIME)
+}

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

@@ -60,15 +60,6 @@ mod tests {
     use super::*;
     use crate::runtime::vm::mpk::enabled::skip_if_mpk_unavailable;
 
-    #[test]
-    #[ignore = "cannot be run with other tests that munge the PKRU register"]
-    fn check_read() {
-        skip_if_mpk_unavailable!();
-        assert_eq!(read(), DISABLE_ACCESS ^ 1);
-        // By default, the Linux kernel only allows a process to access key 0,
-        // the default kernel key.
-    }
-
     #[test]
     fn check_roundtrip() {
         skip_if_mpk_unavailable!();