Fix memfd alignment

Author: dani-garcia

crates/bitwarden-crypto/src/store/backend/implementation/custom_alloc/linux_memfd_secret.rs

@@ -6,13 +6,24 @@ pub(crate) struct LinuxMemfdSecretAlloc;
 
 impl LinuxMemfdSecretAlloc {
     pub fn new() -> Option<Self> {
-        // To test if memfd_secret is supported, we try to allocate a 1 byte and see if that succeeds.
+        // To test if memfd_secret is supported, we try to allocate a 1 byte and see if that
+        // succeeds.
         static IS_SUPPORTED: LazyLock<bool> = LazyLock::new(|| {
-            let Some(ptr) = (unsafe { memsec::memfd_secret_sized(1) }) else {
+            let Some(ptr): Option<NonNull<[u8]>> = (unsafe { memsec::memfd_secret_sized(1) })
+            else {
                 return false;
             };
+
+            // Check that the pointer is readable and writable
+            let result = unsafe {
+                let ptr = ptr.as_ptr() as *mut u8;
+                *ptr = 30;
+                *ptr += 107;
+                *ptr == 137
+            };
+
             unsafe { memsec::free_memfd_secret(ptr) };
-            true
+            result
         });
 
         (*IS_SUPPORTED).then_some(Self)
@@ -34,8 +45,23 @@ unsafe impl Allocator for LinuxMemfdSecretAlloc {
             });
         }
 
-        let ptr: NonNull<[u8]> = unsafe { memsec::memfd_secret_sized(layout.size()) }
-            .expect("memfd_secret_sized failed");
+        // Ensure the size we want to allocate is a multiple of the alignment,
+        // so that both the start and end of the allocation are aligned.
+        let layout = layout.pad_to_align();
+
+        let Some(ptr): Option<NonNull<[u8]>> =
+            (unsafe { memsec::memfd_secret_sized(layout.size()) })
+        else {
+            return Err(AllocError);
+        };
+
+        // Check that the pointer is aligned to the requested alignment.
+        // This should never happen, but just in case, we free the memory and return an allocation
+        // error.
+        if (ptr.as_ptr() as *mut u8).align_offset(layout.align()) != 0 {
+            unsafe { memsec::free_memfd_secret(ptr.as_ptr() as *mut u8) };
+            return Err(AllocError);
+        }
 
         Ok(ptr)
     }

crates/bitwarden-crypto/src/store/backend/implementation/mod.rs

@@ -59,11 +59,11 @@ mod tests {
         let basic_store: Box<dyn StoreBackendDebug<TestSymmKey>> =
             Box::new(basic::BasicBackend::new());
 
-        let mlock_store: Box<dyn StoreBackendDebug<TestSymmKey>> =
+        let memfd_store: Box<dyn StoreBackendDebug<TestSymmKey>> =
             Box::new(custom_alloc::CustomAllocBackend::new(
                 custom_alloc::LinuxMemfdSecretAlloc::new().unwrap(),
             ));
-        compare_stores(100_000, basic_store, mlock_store);
+        compare_stores(100_000, basic_store, memfd_store);
     }
 
     /// This function will perform a number of random operations on two stores,
@@ -88,11 +88,11 @@ mod tests {
             fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Operation {
                 match rng.gen_range(0..=4) {
                     0 => Operation::Upsert(
-                        TestSymmKey::A(rng.gen_range(0..=10)),
+                        TestSymmKey::A(rng.gen_range(0..=1000)),
                         SymmetricCryptoKey::generate(rng),
                     ),
-                    1 => Operation::Get(TestSymmKey::A(rng.gen_range(0..=10))),
-                    2 => Operation::Remove(TestSymmKey::A(rng.gen_range(0..=10))),
+                    1 => Operation::Get(TestSymmKey::A(rng.gen_range(0..=1000))),
+                    2 => Operation::Remove(TestSymmKey::A(rng.gen_range(0..=1000))),
                     3 => Operation::Clear,
                     // This one has to be constant as we can't capture variables
                     4 => Operation::Retain(|key| match key {

crates/bitwarden-crypto/src/store/mod.rs

@@ -360,7 +360,7 @@ pub(crate) mod tests {
 
         // Create some test data
         let data: Vec<_> = (0..300usize)
-            .map(|n| DataView(format!("Test {}", n), TestSymmKey::A((n % 15) as u8)))
+            .map(|n| DataView(format!("Test {}", n), TestSymmKey::A((n % 15) as u32)))
             .collect();
 
         // Encrypt the data

crates/bitwarden-crypto/src/traits/mod.rs

@@ -18,19 +18,19 @@ pub(crate) mod tests {
     key_ids! {
         #[symmetric]
         pub enum TestSymmKey {
-            A(u8),
+            A(u32),
 
             // We only support one variant value,
             // but that value can be a tuple
             B((u8, u8)),
 
             #[local]
-            C(u8),
+            C(u16),
         }
 
         #[asymmetric]
         pub enum TestAsymmKey {
-            A(u8),
+            A(u16),
             B,
             #[local]
             C(&'static str),