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),