From ecd37470fca55afcbd9f1cf99f581fe878da7afb Mon Sep 17 00:00:00 2001
From: Andrew Brown <andrew.brown@intel.com>
Date: Fri, 27 Oct 2023 16:53:24 -0700
Subject: [PATCH] mpk: turn on `memory_protection_keys` during fuzzing (#7393)

This also twists the `max_memory_protection_keys` knob.
---
 crates/fuzzing/src/generators/pooling_config.rs | 9 +++++++++
 crates/runtime/src/mpk/mod.rs                   | 2 +-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/crates/fuzzing/src/generators/pooling_config.rs b/crates/fuzzing/src/generators/pooling_config.rs
index 9cdbd1d1319d..0997597f0350 100644
--- a/crates/fuzzing/src/generators/pooling_config.rs
+++ b/crates/fuzzing/src/generators/pooling_config.rs
@@ -1,6 +1,7 @@
 //! Generate instance limits for the pooling allocation strategy.
 
 use arbitrary::{Arbitrary, Unstructured};
+use wasmtime::MpkEnabled;
 
 /// Configuration for `wasmtime::PoolingAllocationStrategy`.
 #[derive(Debug, Clone, Eq, PartialEq, Hash)]
@@ -30,6 +31,9 @@ pub struct PoolingAllocationConfig {
 
     pub async_stack_zeroing: bool,
     pub async_stack_keep_resident: usize,
+
+    pub memory_protection_keys: MpkEnabled,
+    pub max_memory_protection_keys: usize,
 }
 
 impl PoolingAllocationConfig {
@@ -62,6 +66,8 @@ impl PoolingAllocationConfig {
         cfg.async_stack_zeroing(self.async_stack_zeroing);
         cfg.async_stack_keep_resident(self.async_stack_keep_resident);
 
+        cfg.memory_protection_keys(self.memory_protection_keys);
+
         cfg
     }
 }
@@ -104,6 +110,9 @@ impl<'a> Arbitrary<'a> for PoolingAllocationConfig {
 
             async_stack_zeroing: u.arbitrary()?,
             async_stack_keep_resident: u.int_in_range(0..=1 << 20)?,
+
+            memory_protection_keys: u.choose(&[MpkEnabled::Auto, MpkEnabled::Disable])?.clone(),
+            max_memory_protection_keys: u.int_in_range(0..=20)?,
         })
     }
 }
diff --git a/crates/runtime/src/mpk/mod.rs b/crates/runtime/src/mpk/mod.rs
index bfffe24118f7..f60b3880b09b 100644
--- a/crates/runtime/src/mpk/mod.rs
+++ b/crates/runtime/src/mpk/mod.rs
@@ -42,7 +42,7 @@ cfg_if::cfg_if! {
 }
 
 /// Describe the tri-state configuration of memory protection keys (MPK).
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
 pub enum MpkEnabled {
     /// Use MPK if supported by the current system; fall back to guard regions
     /// otherwise.