diff --git a/server/crates/arbiter-server/src/actors/mod.rs b/server/crates/arbiter-server/src/actors/mod.rs index 292c1ae..9a0e197 100644 --- a/server/crates/arbiter-server/src/actors/mod.rs +++ b/server/crates/arbiter-server/src/actors/mod.rs @@ -3,14 +3,9 @@ use crate::{ bootstrap::Bootstrapper, evm::EvmActor, flow_coordinator::FlowCoordinator, operator_registry::OperatorRegistry, vault::Vault, vault_coordinator::VaultCoordinator, }, - db::{ - self, - custody::{CustodyStore, DieselCustodyStore}, - }, + db, }; -use std::sync::Arc; - use kameo::actor::{ActorRef, Spawn}; use kameo_actors::{DeliveryStrategy, message_bus::MessageBus}; use thiserror::Error; @@ -50,12 +45,10 @@ impl GlobalActors { pub async fn spawn(db: db::DatabasePool) -> Result { let events = Self::spawn_message_bus(); - let custody: Arc = Arc::new(DieselCustodyStore); - let vault = - Vault::spawn(Vault::new(db.clone(), events.clone(), Arc::clone(&custody)).await?); + let vault = Vault::spawn(Vault::new(db.clone(), events.clone()).await?); let bootstrapper = Bootstrapper::spawn(Bootstrapper::new(&db, events.clone()).await?); let vault_coordinator = - VaultCoordinator::spawn(VaultCoordinator::new(db.clone(), vault.clone(), custody)); + VaultCoordinator::spawn(VaultCoordinator::new(db.clone(), vault.clone())); let operator_registry = OperatorRegistry::spawn(OperatorRegistry::default()); Ok(Self { bootstrapper, diff --git a/server/crates/arbiter-server/src/actors/vault/mod.rs b/server/crates/arbiter-server/src/actors/vault/mod.rs index d92eb70..0954511 100644 --- a/server/crates/arbiter-server/src/actors/vault/mod.rs +++ b/server/crates/arbiter-server/src/actors/vault/mod.rs @@ -6,14 +6,13 @@ use crate::{ }, db::{ self, - custody::{CustodyRecord, CustodyStore}, + custody::{self, CustodyRecord}, models::{self, RootKeyHistory, RootKeyHistoryId}, schema::{self}, }, }; use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _}; -use std::sync::Arc; use chrono::Utc; use diesel::{ @@ -64,7 +63,7 @@ pub enum Error { DatabaseTransaction(#[from] diesel::result::Error), #[error("Custody storage error: {0}")] - Custody(#[from] db::custody::Error), + Custody(#[from] custody::Error), #[error("Broken database")] BrokenDatabase, @@ -101,17 +100,12 @@ pub struct Vault { db: db::DatabasePool, state: State, events: ActorRef, - custody: Arc, unseal_failures: u32, } #[messages] impl Vault { - pub async fn new( - db: db::DatabasePool, - events: ActorRef, - custody: Arc, - ) -> Result { + pub async fn new(db: db::DatabasePool, events: ActorRef) -> Result { let state = { let mut conn = db.get().await?; @@ -133,7 +127,6 @@ impl Vault { db, state, events, - custody, unseal_failures: 0, }) } @@ -213,7 +206,6 @@ impl Vault { let mut conn = self.db.get().await?; let data_encryption_nonce_bytes = data_encryption_nonce.to_vec(); - let custody_store = Arc::clone(&self.custody); let root_key_history_id = conn .transaction(async |conn| { let root_key_history_id = insert_into(schema::root_key_history::table) @@ -235,7 +227,7 @@ impl Vault { .await?; if let Some(record) = custody.as_ref() { - custody_store.write_record(&mut *conn, record).await?; + custody::write_record(&mut *conn, record).await?; } Result::<_, Error>::Ok(RootKeyHistoryId::from_raw(root_key_history_id)) @@ -457,18 +449,14 @@ impl Vault { #[cfg(test)] mod tests { - use crate::{actors::GlobalActors, db::custody::DieselCustodyStore}; + use crate::actors::GlobalActors; use super::*; async fn bootstrapped_actor(db: &db::DatabasePool) -> Vault { - let mut actor = Vault::new( - db.clone(), - GlobalActors::spawn_message_bus(), - Arc::new(DieselCustodyStore), - ) - .await - .unwrap(); + let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus()) + .await + .unwrap(); let seal_key = KeyCell::from([0u8; 32]); actor.bootstrap(seal_key, None).await.unwrap(); actor diff --git a/server/crates/arbiter-server/src/actors/vault_coordinator/mod.rs b/server/crates/arbiter-server/src/actors/vault_coordinator/mod.rs index a3f6aa3..d4ed285 100644 --- a/server/crates/arbiter-server/src/actors/vault_coordinator/mod.rs +++ b/server/crates/arbiter-server/src/actors/vault_coordinator/mod.rs @@ -2,9 +2,7 @@ //! //! The coordinator collects one passphrase per committee member, then hands the //! assembled material to [`Vault`] in a single message. It owns no Diesel code: -//! everything it reads or writes goes through [`CustodyStore`]. - -use std::sync::Arc; +//! everything it reads or writes goes through [`db::custody`]. use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _}; use argon2::RECOMMENDED_SALT_LEN; @@ -17,7 +15,7 @@ use crate::{ crypto::{KeyCell, derive_key, encryption::v1::Nonce, shamir}, db::{ self, - custody::{CustodyRecord, CustodyStore, EncryptedShare}, + custody::{self, CustodyRecord, EncryptedShare}, models::OperatorId, }, }; @@ -52,7 +50,7 @@ pub enum Error { #[error("Database connection error: {0}")] DatabaseConnection(#[from] db::PoolError), #[error("Custody storage error: {0}")] - Custody(#[from] db::custody::Error), + Custody(#[from] custody::Error), #[error("Encryption error")] Encryption, #[error("The vault is already bootstrapped")] @@ -111,20 +109,14 @@ enum CoordinatorState { pub struct VaultCoordinator { db: db::DatabasePool, vault: ActorRef, - custody: Arc, state: CoordinatorState, } impl VaultCoordinator { - pub fn new( - db: db::DatabasePool, - vault: ActorRef, - custody: Arc, - ) -> Self { + pub const fn new(db: db::DatabasePool, vault: ActorRef) -> Self { Self { db, vault, - custody, state: CoordinatorState::Idle, } } @@ -217,16 +209,13 @@ async fn finalize_bootstrap( /// Reconstruct the seal key from the contributed passphrases and unseal. async fn finalize_unseal( db: &db::DatabasePool, - custody: &Arc, vault: &ActorRef, threshold: usize, contributions: &mut Contributions, ) -> Result<(), Error> { let stored = { let mut conn = db.get().await?; - custody - .shares(&mut conn, &contributions.operators()) - .await? + custody::shares(&mut conn, &contributions.operators()).await? }; let mut plaintext = Vec::with_capacity(stored.len()); @@ -336,7 +325,7 @@ impl VaultCoordinator { if matches!(self.state, CoordinatorState::Idle) { let threshold = { let mut conn = self.db.get().await?; - self.custody.threshold(&mut conn).await? + custody::threshold(&mut conn).await? }; self.state = CoordinatorState::Unsealing { threshold, @@ -374,15 +363,7 @@ impl VaultCoordinator { unreachable!("state was matched as Unsealing above") }; - match finalize_unseal( - &self.db, - &self.custody, - &self.vault, - threshold, - &mut contributions, - ) - .await - { + match finalize_unseal(&self.db, &self.vault, threshold, &mut contributions).await { Ok(()) => Ok(true), Err(error) => { self.state = CoordinatorState::Unsealing { diff --git a/server/crates/arbiter-server/src/crypto/integrity/v1.rs b/server/crates/arbiter-server/src/crypto/integrity/v1.rs index 66c1f01..d92e3b4 100644 --- a/server/crates/arbiter-server/src/crypto/integrity/v1.rs +++ b/server/crates/arbiter-server/src/crypto/integrity/v1.rs @@ -206,9 +206,6 @@ pub async fn is_signing_available(vault: &ActorRef) -> Result ActorRef { let actor = Vault::spawn( - Vault::new( - db.clone(), - GlobalActors::spawn_message_bus(), - Arc::new(DieselCustodyStore), - ) - .await - .unwrap(), + Vault::new(db.clone(), GlobalActors::spawn_message_bus()) + .await + .unwrap(), ); actor .ask(Bootstrap { diff --git a/server/crates/arbiter-server/src/db/custody.rs b/server/crates/arbiter-server/src/db/custody.rs index 01da33d..d790780 100644 --- a/server/crates/arbiter-server/src/db/custody.rs +++ b/server/crates/arbiter-server/src/db/custody.rs @@ -1,15 +1,13 @@ //! Storage for Shamir custody material: the reconstruction threshold and the //! per-operator encrypted shares of the vault seal key. //! -//! Every query lives behind [`CustodyStore`] so that the actors above it hold -//! no Diesel code of their own. [`CustodyStore::write_record`] borrows the -//! caller's connection instead of taking one from the pool, which lets the -//! vault write custody material inside the same transaction that stores the -//! root key. +//! The queries live here so that the actors above hold no Diesel code of their +//! own. Every one of them borrows the caller's connection instead of taking one +//! from the pool, which is what lets the vault write custody material inside +//! the same transaction that stores the root key. use std::collections::HashMap; -use async_trait::async_trait; use diesel::{ExpressionMethods as _, QueryDsl}; use diesel_async::RunQueryDsl; @@ -46,117 +44,94 @@ pub enum Error { MissingShare(OperatorId), } -#[async_trait] -pub trait CustodyStore: std::fmt::Debug + Send + Sync { - /// Persist threshold and shares on the caller's connection, joining any - /// transaction the caller has already opened. - async fn write_record( - &self, - conn: &mut db::DatabaseConnection, - record: &CustodyRecord, - ) -> Result<(), Error>; +/// Persist threshold and shares on the caller's connection, joining any +/// transaction the caller has already opened. +pub async fn write_record( + conn: &mut db::DatabaseConnection, + record: &CustodyRecord, +) -> Result<(), Error> { + let threshold = i32::try_from(record.threshold).map_err(|_| Error::BrokenThreshold)?; - /// Number of shares required to reconstruct the seal key. - async fn threshold(&self, conn: &mut db::DatabaseConnection) -> Result; - - /// Load the shares of `operators` in one query, in the order requested. - async fn shares( - &self, - conn: &mut db::DatabaseConnection, - operators: &[OperatorId], - ) -> Result, Error>; -} - -/// The production [`CustodyStore`], backed by the `SQLite` schema. -#[derive(Debug, Clone, Copy, Default)] -pub struct DieselCustodyStore; - -#[async_trait] -impl CustodyStore for DieselCustodyStore { - async fn write_record( - &self, - conn: &mut db::DatabaseConnection, - record: &CustodyRecord, - ) -> Result<(), Error> { - let threshold = i32::try_from(record.threshold).map_err(|_| Error::BrokenThreshold)?; - - // SQLite has no batch form for REPLACE INTO in diesel, so the rows go - // in one at a time. The caller's transaction still makes them atomic. - let now = SqliteTimestamp::now(); - for (operator_id, share) in &record.shares { - diesel::replace_into(schema::operator::table) - .values(( - schema::operator::id.eq(Some(*operator_id)), - schema::operator::share.eq(&share.ciphertext), - schema::operator::share_nonce.eq(&share.nonce), - schema::operator::share_salt.eq(&share.salt), - schema::operator::created_at.eq(now.clone()), - schema::operator::updated_at.eq(now.clone()), - )) - .execute(&mut *conn) - .await?; - } - - let updated = diesel::update(schema::arbiter_settings::table) - .set(schema::arbiter_settings::shamir_threshold.eq(Some(threshold))) + // SQLite has no batch form for REPLACE INTO in diesel, so the rows go + // in one at a time. The caller's transaction still makes them atomic. + let now = SqliteTimestamp::now(); + for (operator_id, share) in &record.shares { + diesel::replace_into(schema::operator::table) + .values(( + schema::operator::id.eq(Some(*operator_id)), + schema::operator::share.eq(&share.ciphertext), + schema::operator::share_nonce.eq(&share.nonce), + schema::operator::share_salt.eq(&share.salt), + schema::operator::created_at.eq(now.clone()), + schema::operator::updated_at.eq(now.clone()), + )) .execute(&mut *conn) .await?; - if updated != 1 { - return Err(Error::MissingSettings); - } - - Ok(()) } - async fn threshold(&self, conn: &mut db::DatabaseConnection) -> Result { - let stored: Option = schema::arbiter_settings::table - .select(schema::arbiter_settings::shamir_threshold) - .first(conn) - .await?; - - stored - .and_then(|value| usize::try_from(value).ok()) - .filter(|threshold| *threshold > 0) - .ok_or(Error::BrokenThreshold) + let updated = diesel::update(schema::arbiter_settings::table) + .set(schema::arbiter_settings::shamir_threshold.eq(Some(threshold))) + .execute(&mut *conn) + .await?; + if updated != 1 { + return Err(Error::MissingSettings); } - async fn shares( - &self, - conn: &mut db::DatabaseConnection, - operators: &[OperatorId], - ) -> Result, Error> { - let wanted: Vec> = operators.iter().copied().map(Some).collect(); - - let rows: Vec<(Option, Vec, Vec, Vec)> = schema::operator::table - .filter(schema::operator::id.eq_any(wanted)) - .select(( - schema::operator::id, - schema::operator::share, - schema::operator::share_nonce, - schema::operator::share_salt, - )) - .load(conn) - .await?; - - let mut found: HashMap = rows - .into_iter() - .filter_map(|(id, ciphertext, nonce, salt)| { - id.map(|id| { - ( - id, - EncryptedShare { - ciphertext, - nonce, - salt, - }, - ) - }) - }) - .collect(); - - operators - .iter() - .map(|id| found.remove(id).ok_or(Error::MissingShare(*id))) - .collect() - } + Ok(()) +} + +/// Number of shares required to reconstruct the seal key. +pub async fn threshold(conn: &mut db::DatabaseConnection) -> Result { + let stored: Option = schema::arbiter_settings::table + .select(schema::arbiter_settings::shamir_threshold) + .first(conn) + .await?; + + stored + .and_then(|value| usize::try_from(value).ok()) + .filter(|threshold| *threshold > 0) + .ok_or(Error::BrokenThreshold) +} + +/// Load the shares of `operators` in one query, in the order requested. +pub async fn shares( + conn: &mut db::DatabaseConnection, + operators: &[OperatorId], +) -> Result, Error> { + /// (id, then the three columns that make up [`EncryptedShare`]) + type ShareRow = (Option, Vec, Vec, Vec); + + let wanted: Vec> = operators.iter().copied().map(Some).collect(); + + let rows: Vec = schema::operator::table + .filter(schema::operator::id.eq_any(wanted)) + .select(( + schema::operator::id, + schema::operator::share, + schema::operator::share_nonce, + schema::operator::share_salt, + )) + .load(conn) + .await?; + + let mut found: HashMap = rows + .into_iter() + .filter_map(|(id, ciphertext, nonce, salt)| { + id.map(|id| { + ( + id, + EncryptedShare { + ciphertext, + nonce, + salt, + }, + ) + }) + }) + .collect(); + + operators + .iter() + .map(|id| found.remove(id).ok_or(Error::MissingShare(*id))) + .collect() } diff --git a/server/crates/arbiter-server/src/evm/mod.rs b/server/crates/arbiter-server/src/evm/mod.rs index 930fa04..ef4e0de 100644 --- a/server/crates/arbiter-server/src/evm/mod.rs +++ b/server/crates/arbiter-server/src/evm/mod.rs @@ -501,9 +501,6 @@ impl Engine { #[cfg(test)] mod tests { - use std::sync::Arc; - - use crate::db::custody::DieselCustodyStore; use alloy::primitives::{Address, Bytes, U256, address}; use chrono::{Duration, Utc}; use diesel::{SelectableHelper, insert_into}; @@ -769,13 +766,9 @@ mod tests { async fn bootstrapped_vault(db: &db::DatabasePool) -> ActorRef { let actor = Vault::spawn( - Vault::new( - db.clone(), - GlobalActors::spawn_message_bus(), - Arc::new(DieselCustodyStore), - ) - .await - .unwrap(), + Vault::new(db.clone(), GlobalActors::spawn_message_bus()) + .await + .unwrap(), ); actor .ask(Bootstrap { diff --git a/server/crates/arbiter-server/tests/common/mod.rs b/server/crates/arbiter-server/tests/common/mod.rs index 5878a97..8eeaba0 100644 --- a/server/crates/arbiter-server/tests/common/mod.rs +++ b/server/crates/arbiter-server/tests/common/mod.rs @@ -7,23 +7,18 @@ use arbiter_proto::transport::{Bi, Error, Receiver, Sender}; use arbiter_server::{ actors::{GlobalActors, vault::Vault}, crypto::KeyCell, - db::{self, custody::DieselCustodyStore, schema}, + db::{self, schema}, }; use async_trait::async_trait; use diesel::QueryDsl; use diesel_async::RunQueryDsl; -use std::sync::Arc; use tokio::sync::mpsc; pub(crate) async fn bootstrapped_vault(db: &db::DatabasePool) -> Vault { - let mut actor = Vault::new( - db.clone(), - GlobalActors::spawn_message_bus(), - Arc::new(DieselCustodyStore), - ) - .await - .unwrap(); + let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus()) + .await + .unwrap(); actor .bootstrap(KeyCell::from([0u8; 32]), None) .await diff --git a/server/crates/arbiter-server/tests/vault/concurrency.rs b/server/crates/arbiter-server/tests/vault/concurrency.rs index 0cd156c..0659031 100644 --- a/server/crates/arbiter-server/tests/vault/concurrency.rs +++ b/server/crates/arbiter-server/tests/vault/concurrency.rs @@ -6,16 +6,13 @@ use arbiter_server::{ vault::{CreateNew, Error, Vault}, }, crypto::KeyCell, - db::{self, custody::DieselCustodyStore, models, schema}, + db::{self, models, schema}, }; use diesel::{ExpressionMethods as _, QueryDsl, SelectableHelper, dsl::sql_query}; use diesel_async::RunQueryDsl; use kameo::actor::{ActorRef, Spawn as _}; -use std::{ - collections::{HashMap, HashSet}, - sync::Arc, -}; +use std::collections::{HashMap, HashSet}; use tokio::task::JoinSet; const TEST_AAD: &[u8] = b"test-aad"; @@ -169,13 +166,9 @@ async fn decrypt_roundtrip_after_high_concurrency() { let writes = write_concurrently(actor, "roundtrip", 40).await; let expected: HashMap> = writes.into_iter().collect(); - let mut decryptor = Vault::new( - db.clone(), - GlobalActors::spawn_message_bus(), - Arc::new(DieselCustodyStore), - ) - .await - .unwrap(); + let mut decryptor = Vault::new(db.clone(), GlobalActors::spawn_message_bus()) + .await + .unwrap(); decryptor .try_unseal(KeyCell::from([0u8; 32])) .await diff --git a/server/crates/arbiter-server/tests/vault/lifecycle.rs b/server/crates/arbiter-server/tests/vault/lifecycle.rs index 085bb36..bb0bde5 100644 --- a/server/crates/arbiter-server/tests/vault/lifecycle.rs +++ b/server/crates/arbiter-server/tests/vault/lifecycle.rs @@ -9,12 +9,11 @@ use arbiter_server::{ KeyCell, encryption::v1::{Nonce, ROOT_KEY_TAG}, }, - db::{self, custody::DieselCustodyStore, models, schema}, + db::{self, models, schema}, }; use diesel::{QueryDsl, SelectableHelper}; use diesel_async::RunQueryDsl; -use std::sync::Arc; const TEST_AAD: &[u8] = b"test-aad"; @@ -22,13 +21,9 @@ const TEST_AAD: &[u8] = b"test-aad"; #[test_log::test] async fn bootstrap() { let db = db::create_test_pool().await; - let mut actor = Vault::new( - db.clone(), - GlobalActors::spawn_message_bus(), - Arc::new(DieselCustodyStore), - ) - .await - .unwrap(); + let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus()) + .await + .unwrap(); let seal_key = KeyCell::from([0u8; 32]); actor.bootstrap(seal_key, None).await.unwrap(); @@ -62,13 +57,9 @@ async fn bootstrap_rejects_double() { #[test_log::test] async fn create_new_before_bootstrap_fails() { let db = db::create_test_pool().await; - let mut actor = Vault::new( - db, - GlobalActors::spawn_message_bus(), - Arc::new(DieselCustodyStore), - ) - .await - .unwrap(); + let mut actor = Vault::new(db, GlobalActors::spawn_message_bus()) + .await + .unwrap(); let err = actor .create_new(SafeCell::new(b"data".to_vec()), TEST_AAD.to_vec()) @@ -81,13 +72,9 @@ async fn create_new_before_bootstrap_fails() { #[test_log::test] async fn decrypt_before_bootstrap_fails() { let db = db::create_test_pool().await; - let mut actor = Vault::new( - db, - GlobalActors::spawn_message_bus(), - Arc::new(DieselCustodyStore), - ) - .await - .unwrap(); + let mut actor = Vault::new(db, GlobalActors::spawn_message_bus()) + .await + .unwrap(); let err = actor.decrypt(1, TEST_AAD.to_vec()).await.unwrap_err(); assert!(matches!(err, Error::NotBootstrapped)); @@ -100,13 +87,9 @@ async fn new_restores_sealed_state() { let actor = common::bootstrapped_vault(&db).await; drop(actor); - let mut actor2 = Vault::new( - db, - GlobalActors::spawn_message_bus(), - Arc::new(DieselCustodyStore), - ) - .await - .unwrap(); + let mut actor2 = Vault::new(db, GlobalActors::spawn_message_bus()) + .await + .unwrap(); let err = actor2.decrypt(1, TEST_AAD.to_vec()).await.unwrap_err(); assert!(matches!(err, Error::Sealed)); } @@ -124,13 +107,9 @@ async fn unseal_correct_password() { .unwrap(); drop(actor); - let mut actor = Vault::new( - db.clone(), - GlobalActors::spawn_message_bus(), - Arc::new(DieselCustodyStore), - ) - .await - .unwrap(); + let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus()) + .await + .unwrap(); let seal_key = KeyCell::from([0u8; 32]); actor.try_unseal(seal_key).await.unwrap(); @@ -151,13 +130,9 @@ async fn unseal_wrong_then_correct_password() { .unwrap(); drop(actor); - let mut actor = Vault::new( - db.clone(), - GlobalActors::spawn_message_bus(), - Arc::new(DieselCustodyStore), - ) - .await - .unwrap(); + let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus()) + .await + .unwrap(); let bad_key = KeyCell::from([1u8; 32]); let err = actor.try_unseal(bad_key).await.unwrap_err();