[PM-18100] Add mlock and memfd_secret implementations

crates/bitwarden-crypto/Cargo.toml

@@ -23,6 +23,7 @@ no-memory-hardening = [] # Disable memory hardening features
 
 [dependencies]
 aes = { version = ">=0.8.2, <0.9", features = ["zeroize"] }
+allocator-api2 = ">=0.2.21, <0.3"
 argon2 = { version = ">=0.5.0, <0.6", features = [
     "std",
     "zeroize",
@@ -35,6 +36,7 @@ ciborium = { version = ">=0.2.2, <0.3" }
 coset = { version = ">=0.3.8, <0.4" }
 ed25519-dalek = { version = ">=2.1.1, <=2.2.0", features = ["rand_core"] }
 generic-array = { version = ">=0.14.7, <1.0", features = ["zeroize"] }
+hashbrown = ">=0.15.3, <0.16"
 hkdf = ">=0.12.3, <0.13"
 hmac = ">=0.12.1, <0.13"
 num-bigint = ">=0.4, <0.5"
@@ -60,6 +62,9 @@ wasm-bindgen = { workspace = true, optional = true }
 zeroize = { version = ">=1.7.0, <2.0", features = ["derive", "aarch64"] }
 zeroizing-alloc = ">=0.1.0, <0.2"
 
+[target.'cfg(all(not(target_arch = "wasm32"), not(windows)))'.dependencies]
+memsec = { version = ">=0.7.0, <0.8", features = ["alloc_ext"] }
+
 [dev-dependencies]
 criterion = "0.6.0"
 rand_chacha = "0.3.1"
@@ -77,3 +82,7 @@ required-features = ["no-memory-hardening"]
 
 [lints]
 workspace = true
+
+[package.metadata.cargo-udeps.ignore]
+# This is unused when using --all-features, as that disables memory-hardening
+normal = ["memsec"]

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

@@ -47,3 +47,10 @@ impl<Key: KeyId> Drop for BasicBackend<Key> {
         self.clear();
     }
 }
+
+#[cfg(test)]
+impl<Key: KeyId> super::super::StoreBackendDebug<Key> for BasicBackend<Key> {
+    fn elements(&self) -> Vec<(Key, &Key::KeyValue)> {
+        self.keys.iter().map(|(k, v)| (*k, v)).collect()
+    }
+}

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

@@ -0,0 +1,78 @@
+use std::{alloc::Layout, ptr::NonNull, sync::LazyLock};
+
+use allocator_api2::alloc::{AllocError, Allocator};
+
+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.
+        static IS_SUPPORTED: LazyLock<bool> = LazyLock::new(|| {
+            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) };
+            result
+        });
+
+        (*IS_SUPPORTED).then_some(Self)
+    }
+}
+
+unsafe impl Allocator for LinuxMemfdSecretAlloc {
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        // Note: The allocator_api2 Allocator traits requires us to handle zero-sized allocations.
+        // We return an invalid pointer as you cannot allocate a zero-sized slice in most
+        // allocators. This is what allocator_api2::Global does as well:
+        // https://github.com/zakarumych/allocator-api2/blob/2dde97af85f3559619689cef152e90e6d8a0cee3/src/alloc/global.rs#L24-L29
+        if layout.size() == 0 {
+            return Ok(unsafe {
+                NonNull::new_unchecked(core::ptr::slice_from_raw_parts_mut(
+                    layout.align() as *mut u8,
+                    0,
+                ))
+            });
+        }
+
+        // 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);
+        };
+
+        // The pointer that we return needs to be aligned to the requested alignment.
+        // If that's not the case, we should free the memory and return an allocation error.
+        //  While we check for this error condition, this should never happen in practice, as the
+        // pointer returned by `memfd_secret_sized` should be page-aligned (typically 4KB)
+        // which should be larger than any possible alignment value.
+        if (ptr.as_ptr() as *mut u8).align_offset(layout.align()) != 0 {
+            unsafe { memsec::free_memfd_secret(ptr) };
+            return Err(AllocError);
+        }
+
+        Ok(ptr)
+    }
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        if layout.size() == 0 {
+            return;
+        }
+
+        memsec::free_memfd_secret(ptr);
+    }
+}

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

@@ -0,0 +1,51 @@
+use std::{
+    alloc::{GlobalAlloc, Layout},
+    ptr::NonNull,
+};
+
+use allocator_api2::alloc::{AllocError, Allocator};
+
+pub(crate) struct MlockAlloc(crate::ZeroizingAllocator<std::alloc::System>);
+
+impl MlockAlloc {
+    pub fn new() -> Self {
+        Self(crate::ZeroizingAllocator(std::alloc::System))
+    }
+}
+
+unsafe impl Allocator for MlockAlloc {
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        // Note: The allocator_api2 Allocator traits requires us to handle zero-sized allocations.
+        // We return an invalid pointer as you cannot allocate a zero-sized slice in most
+        // allocators. This is what allocator_api2::Global does as well:
+        // https://github.com/zakarumych/allocator-api2/blob/2dde97af85f3559619689cef152e90e6d8a0cee3/src/alloc/global.rs#L24-L29
+        if layout.size() == 0 {
+            return Ok(unsafe {
+                NonNull::new_unchecked(core::ptr::slice_from_raw_parts_mut(
+                    layout.align() as *mut u8,
+                    0,
+                ))
+            });
+        }
+
+        let ptr = unsafe { self.0.alloc(layout) };
+
+        if ptr.is_null() {
+            return Err(AllocError);
+        }
+        unsafe { memsec::mlock(ptr, layout.size()) };
+        Ok(unsafe {
+            let slice = std::slice::from_raw_parts_mut(ptr, layout.size());
+            NonNull::new(slice).expect("slice is never null")
+        })
+    }
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        if layout.size() == 0 {
+            return;
+        }
+
+        memsec::munlock(ptr.as_ptr(), layout.size());
+        self.0.dealloc(ptr.as_ptr(), layout);
+    }
+}

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

@@ -0,0 +1,57 @@
+use allocator_api2::alloc::Allocator;
+use zeroize::ZeroizeOnDrop;
+
+use crate::{store::backend::StoreBackend, KeyId};
+
+#[cfg(all(not(target_arch = "wasm32"), not(windows)))]
+pub(super) mod malloc;
+
+#[cfg(target_os = "linux")]
+pub(super) mod linux_memfd_secret;
+
+pub(super) struct CustomAllocBackend<Key: KeyId, Alloc: Allocator + Send + Sync> {
+    map: hashbrown::HashMap<Key, Key::KeyValue, hashbrown::DefaultHashBuilder, Alloc>,
+}
+
+impl<Key: KeyId, Alloc: Allocator + Send + Sync> CustomAllocBackend<Key, Alloc> {
+    pub(super) fn new(alloc: Alloc) -> Self {
+        Self {
+            map: hashbrown::HashMap::new_in(alloc),
+        }
+    }
+}
+
+impl<Key: KeyId, Alloc: Allocator + Send + Sync> ZeroizeOnDrop for CustomAllocBackend<Key, Alloc> {}
+
+impl<Key: KeyId, Alloc: Allocator + Send + Sync> StoreBackend<Key>
+    for CustomAllocBackend<Key, Alloc>
+{
+    fn upsert(&mut self, key_id: Key, key: <Key as KeyId>::KeyValue) {
+        self.map.insert(key_id, key);
+    }
+
+    fn get(&self, key_id: Key) -> Option<&<Key as KeyId>::KeyValue> {
+        self.map.get(&key_id)
+    }
+
+    fn remove(&mut self, key_id: Key) {
+        self.map.remove(&key_id);
+    }
+
+    fn clear(&mut self) {
+        self.map.clear();
+    }
+
+    fn retain(&mut self, f: fn(Key) -> bool) {
+        self.map.retain(|key, _| f(*key));
+    }
+}
+
+#[cfg(test)]
+impl<Key: KeyId, Alloc: Allocator + Send + Sync> super::super::StoreBackendDebug<Key>
+    for CustomAllocBackend<Key, Alloc>
+{
+    fn elements(&self) -> Vec<(Key, &Key::KeyValue)> {
+        self.map.iter().map(|(k, v)| (*k, v)).collect()
+    }
+}