From 654aee18dfe2337e58997fff868ef97479b23fcf Mon Sep 17 00:00:00 2001
From: Andrew Brown <andrew.brown@intel.com>
Date: Wed, 29 Nov 2023 14:44:30 -0800
Subject: [PATCH] Revert "mpk: optimize layout of protected stripes (#7603)"
 (#7611)

This reverts commit 043e4cef69a4b22996c54c5d844ca236a69a0e5d. The
following fails when running locally on an MPK-enabled machine:

```
WASMTIME_TEST_FORCE_MPK=1 cargo test
```
---
 .../allocator/pooling/memory_pool.txt         |  1 -
 .../instance/allocator/pooling/memory_pool.rs | 40 +++++++++----------
 2 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/crates/runtime/proptest-regressions/instance/allocator/pooling/memory_pool.txt b/crates/runtime/proptest-regressions/instance/allocator/pooling/memory_pool.txt
index c1d3c19df5..f95bfcd387 100644
--- a/crates/runtime/proptest-regressions/instance/allocator/pooling/memory_pool.txt
+++ b/crates/runtime/proptest-regressions/instance/allocator/pooling/memory_pool.txt
@@ -6,4 +6,3 @@
 # everyone who runs the test benefits from these saved cases.
 cc 696808084287d5d58b85c60c4720227ab4dd83ada7be6841a67162023aaf4914 # shrinks to c = SlabConstraints { max_memory_bytes: 0, num_memory_slots: 1, num_pkeys_available: 0, guard_bytes: 9223372036854775808 }
 cc cf9f6c36659f7f56ed8ea646e8c699cbf46708cef6911cdd376418ad69ea1388 # shrinks to c = SlabConstraints { max_memory_bytes: 14161452635954640438, num_memory_slots: 0, num_pkeys_available: 0, guard_bytes: 4285291437754911178 }
-cc 58f42405c4fbfb4b464950f372995d4d08a77d6884335e38c2e68d590cacf7d8 # shrinks to c = SlabConstraints { expected_slot_bytes: 0, max_memory_bytes: 8483846582232735745, num_slots: 0, num_pkeys_available: 0, guard_bytes: 0, guard_before_slots: false }
diff --git a/crates/runtime/src/instance/allocator/pooling/memory_pool.rs b/crates/runtime/src/instance/allocator/pooling/memory_pool.rs
index bfeaf26caa..2b25a6d076 100644
--- a/crates/runtime/src/instance/allocator/pooling/memory_pool.rs
+++ b/crates/runtime/src/instance/allocator/pooling/memory_pool.rs
@@ -673,7 +673,8 @@ fn calculate(constraints: &SlabConstraints) -> Result<SlabLayout> {
     // to define a slot as "all of the memory and guard region."
     let slot_bytes = expected_slot_bytes
         .max(max_memory_bytes)
-        .saturating_add(guard_bytes);
+        .checked_add(guard_bytes)
+        .unwrap_or(usize::MAX);
 
     let (num_stripes, slot_bytes) = if guard_bytes == 0 || max_memory_bytes == 0 || num_slots == 0 {
         // In the uncommon case where the memory/guard regions are empty or we don't need any slots , we
@@ -702,22 +703,20 @@ fn calculate(constraints: &SlabConstraints) -> Result<SlabLayout> {
         // 3`), we will run into failures if we attempt to set up more than
         // three stripes.
         let needed_num_stripes =
-            slot_bytes / max_memory_bytes + usize::from(slot_bytes % max_memory_bytes != 0);
-        assert!(needed_num_stripes > 0);
+            slot_bytes / max_memory_bytes + usize::from(slot_bytes % max_memory_bytes != 0) + 1;
         let num_stripes = num_pkeys_available.min(needed_num_stripes).min(num_slots);
 
-        // Next, we try to reduce the slot size by "overlapping" the stripes: we
-        // must respect the `expected_slot_bytes + guard bytes` expectations
-        // that codegen constrains us to, but, since a stripe of a different
-        // color is just as inaccessible as `PROT_NONE`, we can squeeze them
-        // together, effectively reducing the size of `slot_bytes` but still
-        // guaranteeing the codegen constraints.
+        // Next, we try to reduce the slot size by "overlapping" the
+        // stripes: we can make slot `n` smaller since we know that slot
+        // `n+1` and following are in different stripes and will look just
+        // like `PROT_NONE` memory.
+        let next_slots_overlapping_bytes = max_memory_bytes
+            .checked_mul(num_stripes - 1)
+            .unwrap_or(usize::MAX);
         let needed_slot_bytes = slot_bytes
-            .checked_div(num_stripes)
-            .unwrap_or(slot_bytes)
+            .checked_sub(next_slots_overlapping_bytes)
+            .unwrap_or(0)
             .max(max_memory_bytes);
-        assert!(needed_slot_bytes >= max_memory_bytes);
-
         (num_stripes, needed_slot_bytes)
     };
 
@@ -729,14 +728,13 @@ fn calculate(constraints: &SlabConstraints) -> Result<SlabLayout> {
         .ok_or_else(|| anyhow!("slot size is too large"))?;
 
     // We may need another guard region (like `pre_slab_guard_bytes`) at the end
-    // of our slab. If we have overlapped stripes of different colors to reduce
-    // the slot size, the last stripe will be vulnerable--it is not followed by
-    // other inaccessible stripes. To fix this, we ensure that the
-    // `post_slab_guard_bytes` is large enough to ensure the last slot meets the
-    // `expected_slot_bytes + guard_bytes` guarantees.
-    let post_slab_guard_bytes = expected_slot_bytes
-        .saturating_add(guard_bytes)
-        .saturating_sub(slot_bytes);
+    // of our slab. We could be conservative and just create it as large as
+    // `guard_bytes`, but because we know that the last slot already has
+    // `guard_bytes` factored in to its guard region, we can reduce the final
+    // guard region by that much.
+    let post_slab_guard_bytes = guard_bytes
+        .checked_sub(slot_bytes - max_memory_bytes)
+        .unwrap_or(0);
 
     // Check that we haven't exceeded the slab we can calculate given the limits
     // of `usize`.