MarketTakers/arbiter
5
0
Fork 0
Code Issues 27 Pull Requests 4 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
hdbg
2026-03-16 18:56:13 +01:00
parent 088fa6fe72
commit 9017ea4017
14 changed files with 178 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
server/crates/arbiter-server/src/actors/evm/mod.rs
View File

@@ -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");
MemSafe::new(reader.to_vec()).expect("MemSafe allocation")
let reader = key_cell.read();
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(

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

@@ -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>);
impl From<MemSafe<Key>> for KeyCell {
fn from(value: MemSafe<Key>) -> Self {
pub struct KeyCell(pub SafeCell<Key>);
impl From<SafeCell<Key>> for KeyCell {
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> {
let value = value.read().unwrap();
fn try_from(mut value: SafeCell<Vec<u8>>) -> Result<Self, Self::Error> {
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 mut key_buffer = key.write().unwrap();
let password_source = password.read();
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 password2 = MemSafe::new(PASSWORD.to_vec()).unwrap();
let password = SafeCell::new(PASSWORD.to_vec());
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 key2_reader = key2.0.read().unwrap();
let key1_reader = key1.0.read();
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");
}

29
server/crates/arbiter-server/src/actors/keyholder/mod.rs
View File

@@ -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

2
server/crates/arbiter-server/src/actors/user_agent/session.rs
View File

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

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

@@ -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::{
evm::{Generate, ListWallets, UseragentListGrants},
evm::{UseragentCreateGrant, UseragentDeleteGrant},
use crate::{actors::{
evm::{Generate, ListWallets, UseragentCreateGrant, UseragentDeleteGrant, UseragentListGrants},
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)
};

20
server/crates/arbiter-server/src/evm/safe_signer.rs
View File

@@ -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> {
let reader = cell.read().map_err(Error::other)?;
pub fn from_cell(mut cell: SafeCell<Vec<u8>>) -> Result<Self> {
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())
}

2
server/crates/arbiter-server/src/grpc/user_agent.rs
View File

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

2
server/crates/arbiter-server/src/lib.rs
View File

@@ -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 {
}
}

79
server/crates/arbiter-server/src/safe_cell.rs Normal file
View File

@@ -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>;

5
server/crates/arbiter-server/tests/common/mod.rs
View File

@@ -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

10
server/crates/arbiter-server/tests/keyholder/concurrency.rs
View File

@@ -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);
}
}

22
server/crates/arbiter-server/tests/keyholder/lifecycle.rs
View File

@@ -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);
}

20
server/crates/arbiter-server/tests/keyholder/storage.rs
View File

@@ -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!(*d2.read().unwrap(), plaintext);
assert_eq!(*d1.read(), 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();

4
server/crates/arbiter-server/tests/user_agent/unseal.rs
View File

@@ -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();