refactor(server): added SafeCell abstraction for easier protected memory swap

server/crates/arbiter-server/src/actors/evm/mod.rs

@@ -4,14 +4,13 @@ use diesel::{
 };
 use diesel_async::RunQueryDsl;
 use kameo::{Actor, actor::ActorRef, messages};
-use memsafe::MemSafe;
 use rand::{SeedableRng, rng, rngs::StdRng};
 
 use crate::{
     actors::keyholder::{CreateNew, Decrypt, KeyHolder},
     db::{
         self, DatabasePool,
-        models::{self, EvmBasicGrant, SqliteTimestamp},
+        models::{self, SqliteTimestamp},
         schema,
     },
     evm::{
@@ -21,6 +20,7 @@ use crate::{
             ether_transfer::EtherTransfer, token_transfers::TokenTransfer,
         },
     },
+    safe_cell::{SafeCell, SafeCellHandle as _},
 };
 
 pub use crate::evm::safe_signer;
@@ -110,8 +110,8 @@ impl EvmActor {
 
         // Move raw key bytes into a Vec<u8> MemSafe for KeyHolder
         let plaintext = {
-            let reader = key_cell.read().expect("MemSafe read");
+            let reader = key_cell.read();
-            MemSafe::new(reader.to_vec()).expect("MemSafe allocation")
+            SafeCell::new(reader.to_vec())
         };
 
         let aead_id: i32 = self
@@ -249,7 +249,7 @@ impl EvmActor {
             .ok_or(SignTransactionError::WalletNotFound)?;
         drop(conn);
 
-        let raw_key: MemSafe<Vec<u8>> = self
+        let raw_key: SafeCell<Vec<u8>> = self
             .keyholder
             .ask(Decrypt {
                 aead_id: wallet.aead_encrypted_id,
@@ -257,7 +257,7 @@ impl EvmActor {
             .await
             .map_err(|_| SignTransactionError::KeyholderSend)?;
 
-        let signer = safe_signer::SafeSigner::from_memsafe(raw_key)?;
+        let signer = safe_signer::SafeSigner::from_cell(raw_key)?;
 
         self.engine
             .evaluate_transaction(

server/crates/arbiter-server/src/actors/keyholder/encryption/v1.rs

@@ -5,12 +5,13 @@ use chacha20poly1305::{
     AeadInPlace, Key, KeyInit as _, XChaCha20Poly1305, XNonce,
     aead::{AeadMut, Error, Payload},
 };
-use memsafe::MemSafe;
 use rand::{
     Rng as _, SeedableRng,
     rngs::{StdRng, SysRng},
 };
 
+use crate::safe_cell::{SafeCell, SafeCellHandle as _};
+
 pub const ROOT_KEY_TAG: &[u8] = "arbiter/seal/v1".as_bytes();
 pub const TAG: &[u8] = "arbiter/private-key/v1".as_bytes();
 
@@ -47,23 +48,23 @@ impl<'a> TryFrom<&'a [u8]> for Nonce {
     }
 }
 
-pub struct KeyCell(pub MemSafe<Key>);
+pub struct KeyCell(pub SafeCell<Key>);
-impl From<MemSafe<Key>> for KeyCell {
+impl From<SafeCell<Key>> for KeyCell {
-    fn from(value: MemSafe<Key>) -> Self {
+    fn from(value: SafeCell<Key>) -> Self {
         Self(value)
     }
 }
-impl TryFrom<MemSafe<Vec<u8>>> for KeyCell {
+impl TryFrom<SafeCell<Vec<u8>>> for KeyCell {
     type Error = ();
 
-    fn try_from(mut value: MemSafe<Vec<u8>>) -> Result<Self, Self::Error> {
+    fn try_from(mut value: SafeCell<Vec<u8>>) -> Result<Self, Self::Error> {
-        let value = value.read().unwrap();
+        let value = value.read();
         if value.len() != size_of::<Key>() {
             return Err(());
         }
-        let mut cell = MemSafe::new(Key::default()).unwrap();
+        let mut cell = SafeCell::new(Key::default());
         {
-            let mut cell_write = cell.write().unwrap();
+            let mut cell_write = cell.write();
             let cell_slice: &mut [u8] = cell_write.as_mut();
             cell_slice.copy_from_slice(&value);
         }
@@ -73,9 +74,9 @@ impl TryFrom<MemSafe<Vec<u8>>> for KeyCell {
 
 impl KeyCell {
     pub fn new_secure_random() -> Self {
-        let mut key = MemSafe::new(Key::default()).unwrap();
+        let mut key = SafeCell::new(Key::default());
         {
-            let mut key_buffer = key.write().unwrap();
+            let mut key_buffer = key.write();
             let key_buffer: &mut [u8] = key_buffer.as_mut();
 
             let mut rng = StdRng::try_from_rng(&mut SysRng).unwrap();
@@ -91,7 +92,7 @@ impl KeyCell {
         associated_data: &[u8],
         mut buffer: impl AsMut<Vec<u8>>,
     ) -> Result<(), Error> {
-        let key_reader = self.0.read().unwrap();
+        let key_reader = self.0.read();
         let key_ref = key_reader.deref();
         let cipher = XChaCha20Poly1305::new(key_ref);
         let nonce = XNonce::from_slice(nonce.0.as_ref());
@@ -102,13 +103,13 @@ impl KeyCell {
         &mut self,
         nonce: &Nonce,
         associated_data: &[u8],
-        buffer: &mut MemSafe<Vec<u8>>,
+        buffer: &mut SafeCell<Vec<u8>>,
     ) -> Result<(), Error> {
-        let key_reader = self.0.read().unwrap();
+        let key_reader = self.0.read();
         let key_ref = key_reader.deref();
         let cipher = XChaCha20Poly1305::new(key_ref);
         let nonce = XNonce::from_slice(nonce.0.as_ref());
-        let mut buffer = buffer.write().unwrap();
+        let mut buffer = buffer.write();
         let buffer: &mut Vec<u8> = buffer.as_mut();
         cipher.decrypt_in_place(nonce, associated_data, buffer)
     }
@@ -119,7 +120,7 @@ impl KeyCell {
         associated_data: &[u8],
         plaintext: impl AsRef<[u8]>,
     ) -> Result<Vec<u8>, Error> {
-        let key_reader = self.0.read().unwrap();
+        let key_reader = self.0.read();
         let key_ref = key_reader.deref();
         let mut cipher = XChaCha20Poly1305::new(key_ref);
         let nonce = XNonce::from_slice(nonce.0.as_ref());
@@ -146,13 +147,13 @@ pub fn generate_salt() -> Salt {
 
 /// User password might be of different length, have not enough entropy, etc...
 /// Derive a fixed-length key from the password using Argon2id, which is designed for password hashing and key derivation.
-pub fn derive_seal_key(mut password: MemSafe<Vec<u8>>, salt: &Salt) -> KeyCell {
+pub fn derive_seal_key(mut password: SafeCell<Vec<u8>>, salt: &Salt) -> KeyCell {
     let params = argon2::Params::new(262_144, 3, 4, None).unwrap();
     let hasher = Argon2::new(Algorithm::Argon2id, argon2::Version::V0x13, params);
-    let mut key = MemSafe::new(Key::default()).unwrap();
+    let mut key = SafeCell::new(Key::default());
     {
-        let password_source = password.read().unwrap();
+        let password_source = password.read();
-        let mut key_buffer = key.write().unwrap();
+        let mut key_buffer = key.write();
         let key_buffer: &mut [u8] = key_buffer.as_mut();
 
         hasher
@@ -166,20 +167,20 @@ pub fn derive_seal_key(mut password: MemSafe<Vec<u8>>, salt: &Salt) -> KeyCell {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use memsafe::MemSafe;
+    use crate::safe_cell::SafeCell;
 
     #[test]
     pub fn derive_seal_key_deterministic() {
         static PASSWORD: &[u8] = b"password";
-        let password = MemSafe::new(PASSWORD.to_vec()).unwrap();
+        let password = SafeCell::new(PASSWORD.to_vec());
-        let password2 = MemSafe::new(PASSWORD.to_vec()).unwrap();
+        let password2 = SafeCell::new(PASSWORD.to_vec());
         let salt = generate_salt();
 
         let mut key1 = derive_seal_key(password, &salt);
         let mut key2 = derive_seal_key(password2, &salt);
 
-        let key1_reader = key1.0.read().unwrap();
+        let key1_reader = key1.0.read();
-        let key2_reader = key2.0.read().unwrap();
+        let key2_reader = key2.0.read();
 
         assert_eq!(key1_reader.deref(), key2_reader.deref());
     }
@@ -187,11 +188,11 @@ mod tests {
     #[test]
     pub fn successful_derive() {
         static PASSWORD: &[u8] = b"password";
-        let password = MemSafe::new(PASSWORD.to_vec()).unwrap();
+        let password = SafeCell::new(PASSWORD.to_vec());
         let salt = generate_salt();
 
         let mut key = derive_seal_key(password, &salt);
-        let key_reader = key.0.read().unwrap();
+        let key_reader = key.0.read();
         let key_ref = key_reader.deref();
 
         assert_ne!(key_ref.as_slice(), &[0u8; 32][..]);
@@ -200,7 +201,7 @@ mod tests {
     #[test]
     pub fn encrypt_decrypt() {
         static PASSWORD: &[u8] = b"password";
-        let password = MemSafe::new(PASSWORD.to_vec()).unwrap();
+        let password = SafeCell::new(PASSWORD.to_vec());
         let salt = generate_salt();
 
         let mut key = derive_seal_key(password, &salt);
@@ -212,12 +213,12 @@ mod tests {
             .unwrap();
         assert_ne!(buffer, b"secret data");
 
-        let mut buffer = MemSafe::new(buffer).unwrap();
+        let mut buffer = SafeCell::new(buffer);
 
         key.decrypt_in_place(&nonce, associated_data, &mut buffer)
             .unwrap();
 
-        let buffer = buffer.read().unwrap();
+        let buffer = buffer.read();
         assert_eq!(*buffer, b"secret data");
     }
 

server/crates/arbiter-server/src/actors/keyholder/mod.rs

@@ -5,15 +5,15 @@ use diesel::{
 };
 use diesel_async::{AsyncConnection, RunQueryDsl};
 use kameo::{Actor, Reply, messages};
-use memsafe::MemSafe;
 use strum::{EnumDiscriminants, IntoDiscriminant};
 use tracing::{error, info};
 
-use crate::db::{
+use crate::{db::{
     self,
     models::{self, RootKeyHistory},
     schema::{self},
-};
+}, safe_cell::SafeCellHandle as _};
+use crate::safe_cell::SafeCell;
 use encryption::v1::{self, KeyCell, Nonce};
 
 pub mod encryption;
@@ -136,7 +136,7 @@ impl KeyHolder {
     }
 
     #[message]
-    pub async fn bootstrap(&mut self, seal_key_raw: MemSafe<Vec<u8>>) -> Result<(), Error> {
+    pub async fn bootstrap(&mut self, seal_key_raw: SafeCell<Vec<u8>>) -> Result<(), Error> {
         if !matches!(self.state, State::Unbootstrapped) {
             return Err(Error::AlreadyBootstrapped);
         }
@@ -149,7 +149,7 @@ impl KeyHolder {
         let data_encryption_nonce = v1::Nonce::default();
 
         let root_key_ciphertext: Vec<u8> = {
-            let root_key_reader = root_key.0.read().unwrap();
+            let root_key_reader = root_key.0.read();
             let root_key_reader = root_key_reader.as_slice();
             seal_key
                 .encrypt(&root_key_nonce, v1::ROOT_KEY_TAG, root_key_reader)
@@ -199,7 +199,7 @@ impl KeyHolder {
     }
 
     #[message]
-    pub async fn try_unseal(&mut self, seal_key_raw: MemSafe<Vec<u8>>) -> Result<(), Error> {
+    pub async fn try_unseal(&mut self, seal_key_raw: SafeCell<Vec<u8>>) -> Result<(), Error> {
         let State::Sealed {
             root_key_history_id,
         } = &self.state
@@ -225,7 +225,7 @@ impl KeyHolder {
         })?;
         let mut seal_key = v1::derive_seal_key(seal_key_raw, &salt);
 
-        let mut root_key = MemSafe::new(current_key.ciphertext.clone()).unwrap();
+        let mut root_key = SafeCell::new(current_key.ciphertext.clone());
 
         let nonce = v1::Nonce::try_from(current_key.root_key_encryption_nonce.as_slice()).map_err(
             |_| {
@@ -256,7 +256,7 @@ impl KeyHolder {
 
     // Decrypts the `aead_encrypted` entry with the given ID and returns the plaintext
     #[message]
-    pub async fn decrypt(&mut self, aead_id: i32) -> Result<MemSafe<Vec<u8>>, Error> {
+    pub async fn decrypt(&mut self, aead_id: i32) -> Result<SafeCell<Vec<u8>>, Error> {
         let State::Unsealed { root_key, .. } = &mut self.state else {
             return Err(Error::NotBootstrapped);
         };
@@ -279,14 +279,14 @@ impl KeyHolder {
             );
             Error::BrokenDatabase
         })?;
-        let mut output = MemSafe::new(row.ciphertext).unwrap();
+        let mut output = SafeCell::new(row.ciphertext);
         root_key.decrypt_in_place(&nonce, v1::TAG, &mut output)?;
         Ok(output)
     }
 
     // Creates new `aead_encrypted` entry in the database and returns it's ID
     #[message]
-    pub async fn create_new(&mut self, mut plaintext: MemSafe<Vec<u8>>) -> Result<i32, Error> {
+    pub async fn create_new(&mut self, mut plaintext: SafeCell<Vec<u8>>) -> Result<i32, Error> {
         let State::Unsealed {
             root_key,
             root_key_history_id,
@@ -299,7 +299,7 @@ impl KeyHolder {
         // Borrow checker note: &mut borrow a few lines above is disjoint from this field
         let nonce = Self::get_new_nonce(&self.db, *root_key_history_id).await?;
 
-        let mut ciphertext_buffer = plaintext.write().unwrap();
+        let mut ciphertext_buffer = plaintext.write();
         let ciphertext_buffer: &mut Vec<u8> = ciphertext_buffer.as_mut();
         root_key.encrypt_in_place(&nonce, v1::TAG, &mut *ciphertext_buffer)?;
 
@@ -348,15 +348,14 @@ mod tests {
     use diesel::SelectableHelper;
 
     use diesel_async::RunQueryDsl;
-    use memsafe::MemSafe;
 
-    use crate::db::{self};
+    use crate::{db::{self}, safe_cell::SafeCell};
 
     use super::*;
 
     async fn bootstrapped_actor(db: &db::DatabasePool) -> KeyHolder {
         let mut actor = KeyHolder::new(db.clone()).await.unwrap();
-        let seal_key = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
+        let seal_key = SafeCell::new(b"test-seal-key".to_vec());
         actor.bootstrap(seal_key).await.unwrap();
         actor
     }
@@ -391,7 +390,7 @@ mod tests {
         assert_eq!(root_row.data_encryption_nonce, n2.to_vec());
 
         let id = actor
-            .create_new(MemSafe::new(b"post-interleave".to_vec()).unwrap())
+            .create_new(SafeCell::new(b"post-interleave".to_vec()))
             .await
             .unwrap();
         let row: models::AeadEncrypted = schema::aead_encrypted::table

server/crates/arbiter-server/src/actors/user_agent/session.rs

@@ -1,5 +1,3 @@
-
-use chacha20poly1305::aead::KeyInit;
 use ed25519_dalek::VerifyingKey;
 use kameo::{Actor, messages, prelude::Context};
 use tokio::{select, sync::watch};

server/crates/arbiter-server/src/actors/user_agent/session/connection.rs

@@ -2,13 +2,11 @@ use std::{ops::DerefMut, sync::Mutex};
 
 use chacha20poly1305::{AeadInPlace, XChaCha20Poly1305, XNonce, aead::KeyInit};
 use kameo::error::SendError;
-use memsafe::MemSafe;
 use tracing::{error, info};
 use x25519_dalek::{EphemeralSecret, PublicKey};
 
-use crate::actors::{
+use crate::{actors::{
-    evm::{Generate, ListWallets, UseragentListGrants},
+    evm::{Generate, ListWallets, UseragentCreateGrant, UseragentDeleteGrant, UseragentListGrants},
-    evm::{UseragentCreateGrant, UseragentDeleteGrant},
     keyholder::{self, Bootstrap, TryUnseal},
     user_agent::{
         BootstrapError, Request, Response, TransportResponseError, UnsealError, VaultState,
@@ -17,7 +15,8 @@ use crate::actors::{
             state::{UnsealContext, UserAgentEvents, UserAgentStates},
         },
     },
-};
+}, safe_cell::SafeCellHandle as _};
+use crate::safe_cell::SafeCell;
 
 impl UserAgentSession {
     pub async fn process_transport_inbound(&mut self, req: Request) -> Output {
@@ -93,16 +92,16 @@ impl UserAgentSession {
         nonce: &[u8],
         ciphertext: &[u8],
         associated_data: &[u8],
-    ) -> Result<MemSafe<Vec<u8>>, ()> {
+    ) -> Result<SafeCell<Vec<u8>>, ()> {
         let nonce = XNonce::from_slice(nonce);
 
         let shared_secret = ephemeral_secret.diffie_hellman(&client_public_key);
         let cipher = XChaCha20Poly1305::new(shared_secret.as_bytes().into());
 
-        let mut key_buffer = MemSafe::new(ciphertext.to_vec()).unwrap();
+        let mut key_buffer = SafeCell::new(ciphertext.to_vec());
 
         let decryption_result = {
-            let mut write_handle = key_buffer.write().unwrap();
+            let mut write_handle = key_buffer.write();
             let write_handle = write_handle.deref_mut();
             cipher.decrypt_in_place(nonce, associated_data, write_handle)
         };

server/crates/arbiter-server/src/evm/safe_signer.rs

@@ -8,7 +8,7 @@ use alloy::{
 };
 use async_trait::async_trait;
 use k256::ecdsa::{self, RecoveryId, SigningKey, signature::hazmat::PrehashSigner};
-use memsafe::MemSafe;
+use crate::safe_cell::{SafeCell, SafeCellHandle as _};
 
 /// An Ethereum signer that stores its secp256k1 secret key inside a
 /// hardware-protected [`MemSafe`] cell.
@@ -20,7 +20,7 @@ use memsafe::MemSafe;
 /// Because [`MemSafe::read`] requires `&mut self` while the [`Signer`] trait
 /// requires `&self`, the cell is wrapped in a [`Mutex`].
 pub struct SafeSigner {
-    key: Mutex<MemSafe<SigningKey>>,
+    key: Mutex<SafeCell<SigningKey>>,
     address: Address,
     chain_id: Option<ChainId>,
 }
@@ -42,14 +42,14 @@ impl std::fmt::Debug for SafeSigner {
 /// rejection, but we retry to be correct).
 ///
 /// Returns the protected key bytes and the derived Ethereum address.
-pub fn generate(rng: &mut impl rand::Rng) -> (MemSafe<[u8; 32]>, Address) {
+pub fn generate(rng: &mut impl rand::Rng) -> (SafeCell<[u8; 32]>, Address) {
     loop {
-        let mut cell = MemSafe::new([0u8; 32]).expect("MemSafe allocation");
+        let mut cell = SafeCell::new([0u8; 32]);
         {
-            let mut w = cell.write().expect("MemSafe write");
+            let mut w = cell.write();
             rng.fill_bytes(w.as_mut());
         }
-        let reader = cell.read().expect("MemSafe read");
+        let reader = cell.read();
         if let Ok(sk) = SigningKey::from_slice(reader.as_ref()) {
             let address = secret_key_to_address(&sk);
             drop(reader);
@@ -64,8 +64,8 @@ impl SafeSigner {
     /// The key bytes are read from protected memory, parsed as a secp256k1
     /// scalar, and immediately moved into a new [`MemSafe`] cell. The raw
     /// bytes are never exposed outside this function.
-    pub fn from_memsafe(mut cell: MemSafe<Vec<u8>>) -> Result<Self> {
+    pub fn from_cell(mut cell: SafeCell<Vec<u8>>) -> Result<Self> {
-        let reader = cell.read().map_err(Error::other)?;
+        let reader = cell.read();
         let sk = SigningKey::from_slice(reader.as_slice()).map_err(Error::other)?;
         drop(reader);
         Self::new(sk)
@@ -75,7 +75,7 @@ impl SafeSigner {
     /// memory region.
     pub fn new(key: SigningKey) -> Result<Self> {
         let address = secret_key_to_address(&key);
-        let cell = MemSafe::new(key).map_err(Error::other)?;
+        let cell = SafeCell::new(key);
         Ok(Self {
             key: Mutex::new(cell),
             address,
@@ -85,7 +85,7 @@ impl SafeSigner {
 
     fn sign_hash_inner(&self, hash: &B256) -> Result<Signature> {
         let mut cell = self.key.lock().expect("SafeSigner mutex poisoned");
-        let reader = cell.read().map_err(Error::other)?;
+        let reader = cell.read();
         let sig: (ecdsa::Signature, RecoveryId) = reader.sign_prehash(hash.as_ref())?;
         Ok(sig.into())
     }

server/crates/arbiter-server/src/grpc/user_agent.rs

@@ -6,7 +6,6 @@ use arbiter_proto::{
             EvmGrantCreateRequest, EvmGrantCreateResponse, EvmGrantDeleteRequest,
             EvmGrantDeleteResponse, EvmGrantList, EvmGrantListResponse, GrantEntry,
             SharedSettings as ProtoSharedSettings, SpecificGrant as ProtoSpecificGrant,
-            SpecificGrant as ProtoGrantSpecificGrant,
             TokenTransferSettings as ProtoTokenTransferSettings,
             VolumeRateLimit as ProtoVolumeRateLimit, WalletCreateResponse, WalletEntry, WalletList,
             WalletListResponse, evm_grant_create_response::Result as EvmGrantCreateResult,
@@ -42,7 +41,6 @@ use crate::{
         TransportResponseError, UnsealError, VaultState,
     },
     evm::{
-        self,
         policies::{Grant, SpecificGrant},
         policies::{
             SharedGrantSettings, TransactionRateLimit, VolumeRateLimit, ether_transfer,

server/crates/arbiter-server/src/lib.rs

@@ -12,6 +12,7 @@ pub mod context;
 pub mod db;
 pub mod evm;
 pub mod grpc;
+pub mod safe_cell;
 
 const DEFAULT_CHANNEL_SIZE: usize = 1000;
 
@@ -25,4 +26,3 @@ impl Server {
     }
 }
 
-

server/crates/arbiter-server/src/safe_cell.rs

@@ -0,0 +1,79 @@
+use std::ops::{Deref, DerefMut};
+use std::{any::type_name, fmt};
+
+use memsafe::MemSafe;
+
+pub trait SafeCellHandle<T> {
+    type CellRead<'a>: Deref<Target = T>
+    where
+        Self: 'a,
+        T: 'a;
+    type CellWrite<'a>: Deref<Target = T> + DerefMut<Target = T>
+    where
+        Self: 'a,
+        T: 'a;
+
+    fn new(value: T) -> Self
+    where
+        Self: Sized;
+
+    fn read(&mut self) -> Self::CellRead<'_>;
+    fn write(&mut self) -> Self::CellWrite<'_>;
+}
+
+pub struct MemSafeCell<T>(MemSafe<T>);
+
+impl<T> fmt::Debug for MemSafeCell<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("MemSafeCell")
+            .field("inner", &format_args!("<protected {}>", type_name::<T>()))
+            .finish()
+    }
+}
+
+impl<T> SafeCellHandle<T> for MemSafeCell<T> {
+    type CellRead<'a>
+        = memsafe::MemSafeRead<'a, T>
+    where
+        Self: 'a,
+        T: 'a;
+    type CellWrite<'a>
+        = memsafe::MemSafeWrite<'a, T>
+    where
+        Self: 'a,
+        T: 'a;
+
+    fn new(value: T) -> Self {
+        match MemSafe::new(value) {
+            Ok(inner) => Self(inner),
+            Err(err) => {
+                // If protected memory cannot be allocated, process integrity is compromised.
+                abort_memory_breach("safe cell allocation", &err)
+            }
+        }
+    }
+
+    fn read(&mut self) -> Self::CellRead<'_> {
+        match self.0.read() {
+            Ok(inner) => inner,
+            Err(err) => abort_memory_breach("safe cell read", &err),
+        }
+    }
+
+    fn write(&mut self) -> Self::CellWrite<'_> {
+        match self.0.write() {
+            Ok(inner) => inner,
+            Err(err) => {
+                // If protected memory becomes unwritable here, treat it as a fatal memory breach.
+                abort_memory_breach("safe cell write", &err)
+            }
+        }
+    }
+}
+
+fn abort_memory_breach(action: &str, err: &memsafe::error::MemoryError) -> ! {
+    eprintln!("fatal {action}: {err}");
+    std::process::abort();
+}
+
+pub type SafeCell<T> = MemSafeCell<T>;

server/crates/arbiter-server/tests/common/mod.rs

@@ -1,19 +1,18 @@
 use arbiter_proto::transport::{Bi, Error};
 use arbiter_server::{
     actors::keyholder::KeyHolder,
-    db::{self, schema},
+    db::{self, schema}, safe_cell::{SafeCell, SafeCellHandle as _},
 };
 use async_trait::async_trait;
 use diesel::QueryDsl;
 use diesel_async::RunQueryDsl;
-use memsafe::MemSafe;
 use tokio::sync::mpsc;
 
 #[allow(dead_code)]
 pub async fn bootstrapped_keyholder(db: &db::DatabasePool) -> KeyHolder {
     let mut actor = KeyHolder::new(db.clone()).await.unwrap();
     actor
-        .bootstrap(MemSafe::new(b"test-seal-key".to_vec()).unwrap())
+        .bootstrap(SafeCell::new(b"test-seal-key".to_vec()))
         .await
         .unwrap();
     actor

server/crates/arbiter-server/tests/keyholder/concurrency.rs

@@ -3,11 +3,11 @@ use std::collections::{HashMap, HashSet};
 use arbiter_server::{
     actors::keyholder::{CreateNew, Error, KeyHolder},
     db::{self, models, schema},
+    safe_cell::{SafeCell, SafeCellHandle as _},
 };
 use diesel::{ExpressionMethods as _, QueryDsl, SelectableHelper, dsl::sql_query};
 use diesel_async::RunQueryDsl;
 use kameo::actor::{ActorRef, Spawn as _};
-use memsafe::MemSafe;
 use tokio::task::JoinSet;
 
 use crate::common;
@@ -24,7 +24,7 @@ async fn write_concurrently(
             let plaintext = format!("{prefix}-{i}").into_bytes();
             let id = actor
                 .ask(CreateNew {
-                    plaintext: MemSafe::new(plaintext.clone()).unwrap(),
+                    plaintext: SafeCell::new(plaintext.clone()),
                 })
                 .await
                 .unwrap();
@@ -118,7 +118,7 @@ async fn insert_failure_does_not_create_partial_row() {
     drop(conn);
 
     let err = actor
-        .create_new(MemSafe::new(b"should fail".to_vec()).unwrap())
+        .create_new(SafeCell::new(b"should fail".to_vec()))
         .await
         .unwrap_err();
     assert!(matches!(err, Error::DatabaseTransaction(_)));
@@ -162,12 +162,12 @@ async fn decrypt_roundtrip_after_high_concurrency() {
 
     let mut decryptor = KeyHolder::new(db.clone()).await.unwrap();
     decryptor
-        .try_unseal(MemSafe::new(b"test-seal-key".to_vec()).unwrap())
+        .try_unseal(SafeCell::new(b"test-seal-key".to_vec()))
         .await
         .unwrap();
 
     for (id, plaintext) in expected {
         let mut decrypted = decryptor.decrypt(id).await.unwrap();
-        assert_eq!(*decrypted.read().unwrap(), plaintext);
+        assert_eq!(*decrypted.read(), plaintext);
     }
 }

server/crates/arbiter-server/tests/keyholder/lifecycle.rs

@@ -1,10 +1,10 @@
 use arbiter_server::{
     actors::keyholder::{Error, KeyHolder},
     db::{self, models, schema},
+    safe_cell::{SafeCell, SafeCellHandle as _},
 };
 use diesel::{QueryDsl, SelectableHelper};
 use diesel_async::RunQueryDsl;
-use memsafe::MemSafe;
 
 use crate::common;
 
@@ -14,7 +14,7 @@ async fn test_bootstrap() {
     let db = db::create_test_pool().await;
     let mut actor = KeyHolder::new(db.clone()).await.unwrap();
 
-    let seal_key = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
+    let seal_key = SafeCell::new(b"test-seal-key".to_vec());
     actor.bootstrap(seal_key).await.unwrap();
 
     let mut conn = db.get().await.unwrap();
@@ -43,7 +43,7 @@ async fn test_bootstrap_rejects_double() {
     let db = db::create_test_pool().await;
     let mut actor = common::bootstrapped_keyholder(&db).await;
 
-    let seal_key2 = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
+    let seal_key2 = SafeCell::new(b"test-seal-key".to_vec());
     let err = actor.bootstrap(seal_key2).await.unwrap_err();
     assert!(matches!(err, Error::AlreadyBootstrapped));
 }
@@ -55,7 +55,7 @@ async fn test_create_new_before_bootstrap_fails() {
     let mut actor = KeyHolder::new(db).await.unwrap();
 
     let err = actor
-        .create_new(MemSafe::new(b"data".to_vec()).unwrap())
+        .create_new(SafeCell::new(b"data".to_vec()))
         .await
         .unwrap_err();
     assert!(matches!(err, Error::NotBootstrapped));
@@ -91,17 +91,17 @@ async fn test_unseal_correct_password() {
 
     let plaintext = b"survive a restart";
     let aead_id = actor
-        .create_new(MemSafe::new(plaintext.to_vec()).unwrap())
+        .create_new(SafeCell::new(plaintext.to_vec()))
         .await
         .unwrap();
     drop(actor);
 
     let mut actor = KeyHolder::new(db.clone()).await.unwrap();
-    let seal_key = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
+    let seal_key = SafeCell::new(b"test-seal-key".to_vec());
     actor.try_unseal(seal_key).await.unwrap();
 
     let mut decrypted = actor.decrypt(aead_id).await.unwrap();
-    assert_eq!(*decrypted.read().unwrap(), plaintext);
+    assert_eq!(*decrypted.read(), plaintext);
 }
 
 #[tokio::test]
@@ -112,20 +112,20 @@ async fn test_unseal_wrong_then_correct_password() {
 
     let plaintext = b"important data";
     let aead_id = actor
-        .create_new(MemSafe::new(plaintext.to_vec()).unwrap())
+        .create_new(SafeCell::new(plaintext.to_vec()))
         .await
         .unwrap();
     drop(actor);
 
     let mut actor = KeyHolder::new(db.clone()).await.unwrap();
 
-    let bad_key = MemSafe::new(b"wrong-password".to_vec()).unwrap();
+    let bad_key = SafeCell::new(b"wrong-password".to_vec());
     let err = actor.try_unseal(bad_key).await.unwrap_err();
     assert!(matches!(err, Error::InvalidKey));
 
-    let good_key = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
+    let good_key = SafeCell::new(b"test-seal-key".to_vec());
     actor.try_unseal(good_key).await.unwrap();
 
     let mut decrypted = actor.decrypt(aead_id).await.unwrap();
-    assert_eq!(*decrypted.read().unwrap(), plaintext);
+    assert_eq!(*decrypted.read(), plaintext);
 }

server/crates/arbiter-server/tests/keyholder/storage.rs

@@ -3,10 +3,10 @@ use std::collections::HashSet;
 use arbiter_server::{
     actors::keyholder::{Error, encryption::v1},
     db::{self, models, schema},
+    safe_cell::{SafeCell, SafeCellHandle as _},
 };
 use diesel::{ExpressionMethods as _, QueryDsl, SelectableHelper, dsl::update};
 use diesel_async::RunQueryDsl;
-use memsafe::MemSafe;
 
 use crate::common;
 
@@ -18,12 +18,12 @@ async fn test_create_decrypt_roundtrip() {
 
     let plaintext = b"hello arbiter";
     let aead_id = actor
-        .create_new(MemSafe::new(plaintext.to_vec()).unwrap())
+        .create_new(SafeCell::new(plaintext.to_vec()))
         .await
         .unwrap();
 
     let mut decrypted = actor.decrypt(aead_id).await.unwrap();
-    assert_eq!(*decrypted.read().unwrap(), plaintext);
+    assert_eq!(*decrypted.read(), plaintext);
 }
 
 #[tokio::test]
@@ -44,11 +44,11 @@ async fn test_ciphertext_differs_across_entries() {
 
     let plaintext = b"same content";
     let id1 = actor
-        .create_new(MemSafe::new(plaintext.to_vec()).unwrap())
+        .create_new(SafeCell::new(plaintext.to_vec()))
         .await
         .unwrap();
     let id2 = actor
-        .create_new(MemSafe::new(plaintext.to_vec()).unwrap())
+        .create_new(SafeCell::new(plaintext.to_vec()))
         .await
         .unwrap();
 
@@ -70,8 +70,8 @@ async fn test_ciphertext_differs_across_entries() {
 
     let mut d1 = actor.decrypt(id1).await.unwrap();
     let mut d2 = actor.decrypt(id2).await.unwrap();
-    assert_eq!(*d1.read().unwrap(), plaintext);
+    assert_eq!(*d1.read(), plaintext);
-    assert_eq!(*d2.read().unwrap(), plaintext);
+    assert_eq!(*d2.read(), plaintext);
 }
 
 #[tokio::test]
@@ -83,7 +83,7 @@ async fn test_nonce_never_reused() {
     let n = 5;
     for i in 0..n {
         actor
-            .create_new(MemSafe::new(format!("secret {i}").into_bytes()).unwrap())
+            .create_new(SafeCell::new(format!("secret {i}").into_bytes()))
             .await
             .unwrap();
     }
@@ -137,7 +137,7 @@ async fn broken_db_nonce_format_fails_closed() {
     drop(conn);
 
     let err = actor
-        .create_new(MemSafe::new(b"must fail".to_vec()).unwrap())
+        .create_new(SafeCell::new(b"must fail".to_vec()))
         .await
         .unwrap_err();
     assert!(matches!(err, Error::BrokenDatabase));
@@ -145,7 +145,7 @@ async fn broken_db_nonce_format_fails_closed() {
     let db = db::create_test_pool().await;
     let mut actor = common::bootstrapped_keyholder(&db).await;
     let id = actor
-        .create_new(MemSafe::new(b"decrypt target".to_vec()).unwrap())
+        .create_new(SafeCell::new(b"decrypt target".to_vec()))
         .await
         .unwrap();
     let mut conn = db.get().await.unwrap();

server/crates/arbiter-server/tests/user_agent/unseal.rs

@@ -5,9 +5,9 @@ use arbiter_server::{
         user_agent::{Request, Response, UnsealError, session::UserAgentSession},
     },
     db,
+    safe_cell::{SafeCell, SafeCellHandle as _},
 };
 use chacha20poly1305::{AeadInPlace, XChaCha20Poly1305, XNonce, aead::KeyInit};
-use memsafe::MemSafe;
 use x25519_dalek::{EphemeralSecret, PublicKey};
 
 async fn setup_sealed_user_agent(seal_key: &[u8]) -> (db::DatabasePool, UserAgentSession) {
@@ -17,7 +17,7 @@ async fn setup_sealed_user_agent(seal_key: &[u8]) -> (db::DatabasePool, UserAgen
     actors
         .key_holder
         .ask(Bootstrap {
-            seal_key_raw: MemSafe::new(seal_key.to_vec()).unwrap(),
+            seal_key_raw: SafeCell::new(seal_key.to_vec()),
         })
         .await
         .unwrap();