MarketTakers/arbiter
6
0
Fork 0
Code Issues 26 Pull Requests 1 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

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