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

mise.lock

@@ -1,3 +1,12 @@
+[[tools.ast-grep]]
+version = "0.42.0"
+backend = "aqua:ast-grep/ast-grep"
+"platforms.linux-arm64" = { checksum = "sha256:5c830eae8456569e2f7212434ed9c238f58dca412d76045418ed6d394a755836", url = "https://github.com/ast-grep/ast-grep/releases/download/0.42.0/app-aarch64-unknown-linux-gnu.zip"}
+"platforms.linux-x64" = { checksum = "sha256:e825a05603f0bcc4cd9076c4cc8c9abd6d008b7cd07d9aa3cc323ba4b8606651", url = "https://github.com/ast-grep/ast-grep/releases/download/0.42.0/app-x86_64-unknown-linux-gnu.zip"}
+"platforms.macos-arm64" = { checksum = "sha256:fc300d5293b1c770a5aece03a8a193b92e71e87cec726c28096990691a582620", url = "https://github.com/ast-grep/ast-grep/releases/download/0.42.0/app-aarch64-apple-darwin.zip"}
+"platforms.macos-x64" = { checksum = "sha256:979ffe611327056f4730a1ae71b0209b3b830f58b22c6ed194cda34f55400db2", url = "https://github.com/ast-grep/ast-grep/releases/download/0.42.0/app-x86_64-apple-darwin.zip"}
+"platforms.windows-x64" = { checksum = "sha256:55836fa1b2c65dc7d61615a4d9368622a0d2371a76d28b9a165e5a3ab6ae32a4", url = "https://github.com/ast-grep/ast-grep/releases/download/0.42.0/app-x86_64-pc-windows-msvc.zip"}
+
 [[tools."cargo:cargo-audit"]]
 version = "0.22.1"
 backend = "cargo:cargo-audit"

mise.toml

@@ -10,6 +10,7 @@ protoc = "29.6"
 "cargo:cargo-shear" = "latest"
 "cargo:cargo-insta" = "1.46.3"
 python = "3.14.3"
+ast-grep = "0.42.0"
 
 [tasks.codegen]
 sources = ['protobufs/*.proto']

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::{
+use crate::safe_cell::SafeCell;
+use crate::{
+    db::{
         self,
         models::{self, RootKeyHistory},
         schema::{self},
-}, safe_cell::SafeCellHandle as _};
-use crate::safe_cell::SafeCell;
+    },
+    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,8 +5,12 @@ use kameo::error::SendError;
 use tracing::{error, info};
 use x25519_dalek::{EphemeralSecret, PublicKey};
 
-use crate::{actors::{
-    evm::{Generate, ListWallets, UseragentCreateGrant, UseragentDeleteGrant, UseragentListGrants},
+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,
@@ -15,8 +19,9 @@ use crate::{actors::{
                 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,

server/rules/safecell/new-inline.yaml

@@ -0,0 +1,10 @@
+id: safecell-new-inline
+language: Rust
+rule:
+  pattern: $CELL.write_inline(|$W| $BODY);
+  follows:
+    pattern: let mut $CELL = SafeCell::new($INIT);
+fix: 
+  template: let mut $CELL = SafeCell::new_inline(|$W| $BODY);
+  expandStart:
+    pattern: let mut $CELL = SafeCell::new($INIT)

server/rules/safecell/read-inline.yaml

@@ -0,0 +1,17 @@
+id: safecell-read-inline
+language: Rust
+rule:
+  pattern: 
+    context: |
+      {
+          let $READ = $CELL.read();
+          $$$BODY
+      }
+    selector: block
+  inside:
+    kind: block
+fix: 
+  template: | 
+    $CELL.read_inline(|$READ| { 
+      $$$BODY 
+    });

server/rules/safecell/write-inline.yaml

@@ -0,0 +1,13 @@
+id: safecell-write-inline
+language: Rust
+rule:
+  pattern: |
+    {
+        let mut $WRITE = $CELL.write();
+        $$$BODY
+    }
+fix: 
+  template: | 
+    $CELL.write_inline(|$WRITE| { 
+      $$$BODY 
+    });

server/sgconfig.yml

@@ -0,0 +1,2 @@
+ruleDirs:
+- ./rules