refactor(server): rewrote cell access using new helpers and added ast-grep rules for it

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

@@ -108,11 +108,7 @@ impl EvmActor {
     pub async fn generate(&mut self) -> Result<Address, Error> {
         let (mut key_cell, address) = safe_signer::generate(&mut self.rng);
 
-        // Move raw key bytes into a Vec<u8> MemSafe for KeyHolder
-        let plaintext = {
-            let reader = key_cell.read();
-            SafeCell::new(reader.to_vec())
-        };
+        let plaintext = key_cell.read_inline(|reader| SafeCell::new(reader.to_vec()));
 
         let aead_id: i32 = self
             .keyholder

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

@@ -62,26 +62,19 @@ impl TryFrom<SafeCell<Vec<u8>>> for KeyCell {
         if value.len() != size_of::<Key>() {
             return Err(());
         }
-        let mut cell = SafeCell::new(Key::default());
-        {
-            let mut cell_write = cell.write();
-            let cell_slice: &mut [u8] = cell_write.as_mut();
-            cell_slice.copy_from_slice(&value);
-        }
+        let cell = SafeCell::new_inline(|cell_write: &mut Key| {
+            cell_write.copy_from_slice(&value);
+        });
         Ok(Self(cell))
     }
 }
 
 impl KeyCell {
     pub fn new_secure_random() -> Self {
-        let mut key = SafeCell::new(Key::default());
-        {
-            let mut key_buffer = key.write();
-            let key_buffer: &mut [u8] = key_buffer.as_mut();
-
+        let key = SafeCell::new_inline(|key_buffer: &mut Key| {
             let mut rng = StdRng::try_from_rng(&mut SysRng).unwrap();
             rng.fill_bytes(key_buffer);
-        }
+        });
 
         key.into()
     }
@@ -151,15 +144,14 @@ 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 = SafeCell::new(Key::default());
-    {
-        let password_source = password.read();
+    password.read_inline(|password_source| {
         let mut key_buffer = key.write();
         let key_buffer: &mut [u8] = key_buffer.as_mut();
 
         hasher
             .hash_password_into(password_source.deref(), salt, key_buffer)
             .unwrap();
-    }
+    });
 
     key.into()
 }

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

@@ -8,12 +8,15 @@ use kameo::{Actor, Reply, messages};
 use strum::{EnumDiscriminants, IntoDiscriminant};
 use tracing::{error, info};
 
-use crate::{db::{
-    self,
-    models::{self, RootKeyHistory},
-    schema::{self},
-}, safe_cell::SafeCellHandle as _};
 use crate::safe_cell::SafeCell;
+use crate::{
+    db::{
+        self,
+        models::{self, RootKeyHistory},
+        schema::{self},
+    },
+    safe_cell::SafeCellHandle as _,
+};
 use encryption::v1::{self, KeyCell, Nonce};
 
 pub mod encryption;
@@ -148,16 +151,15 @@ impl KeyHolder {
         let root_key_nonce = v1::Nonce::default();
         let data_encryption_nonce = v1::Nonce::default();
 
-        let root_key_ciphertext: Vec<u8> = {
-            let root_key_reader = root_key.0.read();
-            let root_key_reader = root_key_reader.as_slice();
+        let root_key_ciphertext: Vec<u8> = root_key.0.read_inline(|reader| {
+            let root_key_reader = reader.as_slice();
             seal_key
                 .encrypt(&root_key_nonce, v1::ROOT_KEY_TAG, root_key_reader)
                 .map_err(|err| {
                     error!(?err, "Fatal bootstrap error");
                     Error::Encryption(err)
-                })?
-        };
+                })
+        })?;
 
         let mut conn = self.db.get().await?;
 
@@ -349,7 +351,10 @@ mod tests {
 
     use diesel_async::RunQueryDsl;
 
-    use crate::{db::{self}, safe_cell::SafeCell};
+    use crate::{
+        db::{self},
+        safe_cell::SafeCell,
+    };
 
     use super::*;
 

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

@@ -5,18 +5,23 @@ use kameo::error::SendError;
 use tracing::{error, info};
 use x25519_dalek::{EphemeralSecret, PublicKey};
 
-use crate::{actors::{
-    evm::{Generate, ListWallets, UseragentCreateGrant, UseragentDeleteGrant, UseragentListGrants},
-    keyholder::{self, Bootstrap, TryUnseal},
-    user_agent::{
-        BootstrapError, Request, Response, TransportResponseError, UnsealError, VaultState,
-        session::{
-            UserAgentSession,
-            state::{UnsealContext, UserAgentEvents, UserAgentStates},
+use crate::safe_cell::SafeCell;
+use crate::{
+    actors::{
+        evm::{
+            Generate, ListWallets, UseragentCreateGrant, UseragentDeleteGrant, UseragentListGrants,
+        },
+        keyholder::{self, Bootstrap, TryUnseal},
+        user_agent::{
+            BootstrapError, Request, Response, TransportResponseError, UnsealError, VaultState,
+            session::{
+                UserAgentSession,
+                state::{UnsealContext, UserAgentEvents, UserAgentStates},
+            },
         },
     },
-}, safe_cell::SafeCellHandle as _};
-use crate::safe_cell::SafeCell;
+    safe_cell::SafeCellHandle as _,
+};
 
 impl UserAgentSession {
     pub async fn process_transport_inbound(&mut self, req: Request) -> Output {
@@ -100,11 +105,9 @@ impl UserAgentSession {
 
         let mut key_buffer = SafeCell::new(ciphertext.to_vec());
 
-        let decryption_result = {
-            let mut write_handle = key_buffer.write();
-            let write_handle = write_handle.deref_mut();
+        let decryption_result = key_buffer.write_inline(|write_handle| {
             cipher.decrypt_in_place(nonce, associated_data, write_handle)
-        };
+        });
 
         match decryption_result {
             Ok(_) => Ok(key_buffer),

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

@@ -1,5 +1,6 @@
 use std::sync::Mutex;
 
+use crate::safe_cell::{SafeCell, SafeCellHandle as _};
 use alloy::{
     consensus::SignableTransaction,
     network::{TxSigner, TxSignerSync},
@@ -8,7 +9,6 @@ use alloy::{
 };
 use async_trait::async_trait;
 use k256::ecdsa::{self, RecoveryId, SigningKey, signature::hazmat::PrehashSigner};
-use crate::safe_cell::{SafeCell, SafeCellHandle as _};
 
 /// An Ethereum signer that stores its secp256k1 secret key inside a
 /// hardware-protected [`MemSafe`] cell.
@@ -44,11 +44,10 @@ impl std::fmt::Debug for SafeSigner {
 /// Returns the protected key bytes and the derived Ethereum address.
 pub fn generate(rng: &mut impl rand::Rng) -> (SafeCell<[u8; 32]>, Address) {
     loop {
-        let mut cell = SafeCell::new([0u8; 32]);
-        {
-            let mut w = cell.write();
-            rng.fill_bytes(w.as_mut());
-        }
+        let mut cell = SafeCell::new_inline(|w: &mut [u8; 32]| {
+            rng.fill_bytes(w);
+        });
+
         let reader = cell.read();
         if let Ok(sk) = SigningKey::from_slice(reader.as_ref()) {
             let address = secret_key_to_address(&sk);

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

@@ -19,6 +19,36 @@ pub trait SafeCellHandle<T> {
 
     fn read(&mut self) -> Self::CellRead<'_>;
     fn write(&mut self) -> Self::CellWrite<'_>;
+
+    fn new_inline<F>(f: F) -> Self
+    where
+        Self: Sized,
+        T: Default,
+        F: for<'a> FnOnce(&'a mut T),
+    {
+        let mut cell = Self::new(T::default());
+        {
+            let mut handle = cell.write();
+            f(handle.deref_mut());
+        }
+        cell
+    }
+
+    #[inline(always)]
+    fn read_inline<F, R>(&mut self, f: F) -> R
+    where
+        F: FnOnce(&T) -> R,
+    {
+        f(&*self.read())
+    }
+
+    #[inline(always)]
+    fn write_inline<F, R>(&mut self, f: F) -> R
+    where
+        F: FnOnce(&mut T) -> R,
+    {
+        f(&mut *self.write())
+    }
 }
 
 pub struct MemSafeCell<T>(MemSafe<T>);
@@ -53,6 +83,7 @@ impl<T> SafeCellHandle<T> for MemSafeCell<T> {
         }
     }
 
+    #[inline(always)]
     fn read(&mut self) -> Self::CellRead<'_> {
         match self.0.read() {
             Ok(inner) => inner,
@@ -60,6 +91,7 @@ impl<T> SafeCellHandle<T> for MemSafeCell<T> {
         }
     }
 
+    #[inline(always)]
     fn write(&mut self) -> Self::CellWrite<'_> {
         match self.0.write() {
             Ok(inner) => inner,