diff --git a/mise.lock b/mise.lock index e3ad0a5..2037239 100644 --- a/mise.lock +++ b/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" diff --git a/mise.toml b/mise.toml index ea6265d..9e1e802 100644 --- a/mise.toml +++ b/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'] @@ -17,4 +18,4 @@ outputs = ['useragent/lib/proto/*'] run = ''' dart pub global activate protoc_plugin && \ protoc --dart_out=grpc:useragent/lib/proto --proto_path=protobufs/ protobufs/*.proto -''' \ No newline at end of file +''' diff --git a/server/crates/arbiter-server/src/actors/evm/mod.rs b/server/crates/arbiter-server/src/actors/evm/mod.rs index 05a9d94..16e4200 100644 --- a/server/crates/arbiter-server/src/actors/evm/mod.rs +++ b/server/crates/arbiter-server/src/actors/evm/mod.rs @@ -108,11 +108,7 @@ impl EvmActor { pub async fn generate(&mut self) -> Result { let (mut key_cell, address) = safe_signer::generate(&mut self.rng); - // Move raw key bytes into a Vec 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 diff --git a/server/crates/arbiter-server/src/actors/keyholder/encryption/v1.rs b/server/crates/arbiter-server/src/actors/keyholder/encryption/v1.rs index 9ca5a45..99348f4 100644 --- a/server/crates/arbiter-server/src/actors/keyholder/encryption/v1.rs +++ b/server/crates/arbiter-server/src/actors/keyholder/encryption/v1.rs @@ -62,26 +62,19 @@ impl TryFrom>> for KeyCell { if value.len() != size_of::() { 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>, 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() } diff --git a/server/crates/arbiter-server/src/actors/keyholder/mod.rs b/server/crates/arbiter-server/src/actors/keyholder/mod.rs index e53f6a3..f37284a 100644 --- a/server/crates/arbiter-server/src/actors/keyholder/mod.rs +++ b/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 = { - let root_key_reader = root_key.0.read(); - let root_key_reader = root_key_reader.as_slice(); + let root_key_ciphertext: Vec = 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::*; diff --git a/server/crates/arbiter-server/src/actors/user_agent/session/connection.rs b/server/crates/arbiter-server/src/actors/user_agent/session/connection.rs index 984f464..627e5f1 100644 --- a/server/crates/arbiter-server/src/actors/user_agent/session/connection.rs +++ b/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), diff --git a/server/crates/arbiter-server/src/evm/safe_signer.rs b/server/crates/arbiter-server/src/evm/safe_signer.rs index 124a248..f1f5bcd 100644 --- a/server/crates/arbiter-server/src/evm/safe_signer.rs +++ b/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); diff --git a/server/crates/arbiter-server/src/safe_cell.rs b/server/crates/arbiter-server/src/safe_cell.rs index 056e131..dc44065 100644 --- a/server/crates/arbiter-server/src/safe_cell.rs +++ b/server/crates/arbiter-server/src/safe_cell.rs @@ -19,6 +19,36 @@ pub trait SafeCellHandle { fn read(&mut self) -> Self::CellRead<'_>; fn write(&mut self) -> Self::CellWrite<'_>; + + fn new_inline(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(&mut self, f: F) -> R + where + F: FnOnce(&T) -> R, + { + f(&*self.read()) + } + + #[inline(always)] + fn write_inline(&mut self, f: F) -> R + where + F: FnOnce(&mut T) -> R, + { + f(&mut *self.write()) + } } pub struct MemSafeCell(MemSafe); @@ -53,6 +83,7 @@ impl SafeCellHandle for MemSafeCell { } } + #[inline(always)] fn read(&mut self) -> Self::CellRead<'_> { match self.0.read() { Ok(inner) => inner, @@ -60,6 +91,7 @@ impl SafeCellHandle for MemSafeCell { } } + #[inline(always)] fn write(&mut self) -> Self::CellWrite<'_> { match self.0.write() { Ok(inner) => inner, diff --git a/server/rules/.gitkeep b/server/rules/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/server/rules/safecell/new-inline.yaml b/server/rules/safecell/new-inline.yaml new file mode 100644 index 0000000..48c2e4e --- /dev/null +++ b/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) \ No newline at end of file diff --git a/server/rules/safecell/read-inline.yaml b/server/rules/safecell/read-inline.yaml new file mode 100644 index 0000000..e5bbee6 --- /dev/null +++ b/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 + }); \ No newline at end of file diff --git a/server/rules/safecell/write-inline.yaml b/server/rules/safecell/write-inline.yaml new file mode 100644 index 0000000..cfa13fe --- /dev/null +++ b/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 + }); \ No newline at end of file diff --git a/server/sgconfig.yml b/server/sgconfig.yml new file mode 100644 index 0000000..b0f5823 --- /dev/null +++ b/server/sgconfig.yml @@ -0,0 +1,2 @@ +ruleDirs: +- ./rules