From f7da166abdc1e7c6a266897566482d88da3d6974 Mon Sep 17 00:00:00 2001
From: Stephen Touset <stouset@crusoe.ai>
Date: Thu, 14 Aug 2025 13:33:45 -0700
Subject: [PATCH] Align `Secret` to the page size of the current system.

Closes #109.

Multiple `Secret<T>` instances could previously share the same memory
page. This would cause `munlock` to unlock multiple secrets at a time.
With secrets aligned to page boundaries, we ensure that no two secrets
can share the same page.

Additionally we had detected the multiple-unlock case in Windows, which
returns an error when `munlock` is called on an already-unlocked page.
We swallowed errors in this case, but with this change we no longer need
to do so.
---
 Cargo.toml    |  1 +
 build.rs      |  2 +-
 src/secret.rs | 53 ++++++++++++++++++++++++++++++++++++++++++++-------
 3 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index f712fbc..5bb7ac6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,6 +19,7 @@ keywords = [ "crypto", "cryptography", "allocator" ]
 [dependencies]
 libc          = '0.2'
 libsodium-sys = { version = '0.2', optional = true }
+page_size     = '0.6'
 
 [target.'cfg(target_family = "unix")'.build-dependencies]
 pkg-config = '0.3'
diff --git a/build.rs b/build.rs
index e2233a0..d5b7510 100644
--- a/build.rs
+++ b/build.rs
@@ -116,4 +116,4 @@ fn link(name: &str, _version: &str) -> Option<()> {
             None
         }
     }
-}
\ No newline at end of file
+}
diff --git a/src/secret.rs b/src/secret.rs
index 7455119..55224dc 100644
--- a/src/secret.rs
+++ b/src/secret.rs
@@ -78,6 +78,22 @@ use std::thread;
 /// ```
 ///
 /// [mlock]: http://man7.org/linux/man-pages/man2/mlock.2.html
+//
+// Aligned to the target's memory page size so that no two `Secret`s
+// can share a page; `munlock` operates on whole pages and would
+// otherwise unlock a sibling secret's memory while it is still live.
+// `repr(align)` requires a literal, so we dispatch on the target.
+// `Secret::new` asserts at runtime that the chosen alignment is at
+// least the page size reported by the OS, catching any target whose
+// real page size exceeds what this table promises.
+#[cfg_attr(any(target_arch = "x86",         target_arch   = "x86_64"), repr(align(4096)))]
+#[cfg_attr(all(target_arch = "aarch64",     target_vendor = "apple"),  repr(align(16384)))]
+#[cfg_attr(all(target_arch = "aarch64", not(target_vendor = "apple")), repr(align(65536)))]
+#[cfg_attr(not(any(
+    target_arch = "x86",
+    target_arch = "x86_64",
+    target_arch = "aarch64",
+)),                                                                    repr(align(65536)))]
 pub struct Secret<T: Bytes> {
     /// The internal protected memory for the [`Secret`].
     data: T,
@@ -119,6 +135,14 @@ impl<T: Bytes> Secret<T> {
     {
         tested!(size_of::<T>() == 0);
 
+        assert!(
+            align_of::<Self>() >= page_size::get(),
+            "secrets: Secret alignment ({}) is smaller than the system page size ({}); \
+             this target is not yet supported",
+            align_of::<Self>(),
+            page_size::get(),
+        );
+
         let mut secret = Self {
             data: T::uninitialized(),
         };
@@ -203,13 +227,7 @@ impl<T: Bytes> Drop for Secret<T> {
     /// Ensures that the [`Secret`]'s underlying memory is `munlock`ed
     /// and zeroed when it leaves scope.
     fn drop(&mut self) {
-        // When we call sodium_munlock on some data, it actually unlocks the entire page that
-        // contains the memory. If two locked items were on the same page, then the second one
-        // fails because it was already unlocked. On Linux, this does now throw an error. On
-        // Windows, it does. We'll ignore it for now, and provide a better fix later.
-        if unsafe { !sodium::munlock(&raw mut self.data) }
-            && !(cfg!(target_family = "windows")
-                && (std::io::Error::last_os_error().raw_os_error() == Some(158))) {
+        if unsafe { !sodium::munlock(&raw mut self.data) } {
             // [`Drop::drop`] is called during stack unwinding, so we
             // may be in a panic already.
             assert!(
@@ -272,6 +290,27 @@ mod tests {
         Secret::<u16>::new(|s| assert_eq!(*s, 0xdbdb));
     }
 
+    #[test]
+    fn it_aligns_to_at_least_the_page_size() {
+        let page = page_size::get();
+        assert!(align_of::<Secret<u8>>()       >= page);
+        assert!(align_of::<Secret<[u8; 32]>>() >= page);
+        assert!(align_of::<Secret<[u64; 4]>>() >= page);
+    }
+
+    #[test]
+    fn it_does_not_share_a_page_between_secrets() {
+        // Regression test for #109: two stack-allocated `Secret`s
+        // could land on the same page, so dropping one would `munlock`
+        // the other's still-live memory.
+        Secret::<u64>::zero(|a| {
+            let addr_a = &raw const *a as usize;
+            let addr_b = Secret::<u64>::zero(|b| &raw const *b as usize);
+            let page   = page_size::get();
+            assert_ne!(addr_a / page, addr_b / page);
+        });
+    }
+
     #[test]
     fn it_zeroes_when_leaving_scope() {
         unsafe {