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::{
+    db,
-        self,
-        custody::{CustodyStore, DieselCustodyStore},
-    },
 };
 
-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<Self, SpawnError> {
         let events = Self::spawn_message_bus();
-        let custody: Arc<dyn CustodyStore> = Arc::new(DieselCustodyStore);
+        let vault = Vault::spawn(Vault::new(db.clone(), events.clone()).await?);
-        let vault =
-            Vault::spawn(Vault::new(db.clone(), events.clone(), Arc::clone(&custody)).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,

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<MessageBus>,
-    custody: Arc<dyn CustodyStore>,
     unseal_failures: u32,
 }
 
 #[messages]
 impl Vault {
-    pub async fn new(
+    pub async fn new(db: db::DatabasePool, events: ActorRef<MessageBus>) -> Result<Self, Error> {
-        db: db::DatabasePool,
-        events: ActorRef<MessageBus>,
-        custody: Arc<dyn CustodyStore>,
-    ) -> Result<Self, Error> {
         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,16 +449,12 @@ 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(
+        let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus())
-            db.clone(),
-            GlobalActors::spawn_message_bus(),
-            Arc::new(DieselCustodyStore),
-        )
             .await
             .unwrap();
         let seal_key = KeyCell::from([0u8; 32]);

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`].
+//! everything it reads or writes goes through [`db::custody`].
-
-use std::sync::Arc;
 
 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<Vault>,
-    custody: Arc<dyn CustodyStore>,
     state: CoordinatorState,
 }
 
 impl VaultCoordinator {
-    pub fn new(
+    pub const fn new(db: db::DatabasePool, vault: ActorRef<Vault>) -> Self {
-        db: db::DatabasePool,
-        vault: ActorRef<Vault>,
-        custody: Arc<dyn CustodyStore>,
-    ) -> 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<dyn CustodyStore>,
     vault: &ActorRef<Vault>,
     threshold: usize,
     contributions: &mut Contributions,
 ) -> Result<(), Error> {
     let stored = {
         let mut conn = db.get().await?;
-        custody
+        custody::shares(&mut conn, &contributions.operators()).await?
-            .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(
+        match finalize_unseal(&self.db, &self.vault, threshold, &mut contributions).await {
-            &self.db,
-            &self.custody,
-            &self.vault,
-            threshold,
-            &mut contributions,
-        )
-        .await
-        {
             Ok(()) => Ok(true),
             Err(error) => {
                 self.state = CoordinatorState::Unsealing {

server/crates/arbiter-server/src/crypto/integrity/v1.rs

@@ -206,9 +206,6 @@ pub async fn is_signing_available(vault: &ActorRef<Vault>) -> Result<bool, Error
 
 #[cfg(test)]
 mod tests {
-    use std::sync::Arc;
-
-    use crate::db::custody::DieselCustodyStore;
     use diesel::{ExpressionMethods as _, QueryDsl};
     use diesel_async::RunQueryDsl;
     use kameo::{actor::ActorRef, prelude::Spawn};
@@ -234,11 +231,7 @@ mod tests {
 
     async fn bootstrapped_vault(db: &db::DatabasePool) -> ActorRef<Vault> {
         let actor = Vault::spawn(
-            Vault::new(
+            Vault::new(db.clone(), GlobalActors::spawn_message_bus())
-                db.clone(),
-                GlobalActors::spawn_message_bus(),
-                Arc::new(DieselCustodyStore),
-            )
                 .await
                 .unwrap(),
         );

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
+//! The queries live here so that the actors above hold no Diesel code of their
-//! no Diesel code of their own. [`CustodyStore::write_record`] borrows the
+//! own. Every one of them borrows the caller's connection instead of taking one
-//! caller's connection instead of taking one from the pool, which lets the
+//! from the pool, which is what lets the vault write custody material inside
-//! vault write custody material inside the same transaction that stores the
+//! the same transaction that stores the root key.
-//! root key.
 
 use std::collections::HashMap;
 
-use async_trait::async_trait;
 use diesel::{ExpressionMethods as _, QueryDsl};
 use diesel_async::RunQueryDsl;
 
@@ -46,35 +44,9 @@ 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(
+pub async fn write_record(
-        &self,
-        conn: &mut db::DatabaseConnection,
-        record: &CustodyRecord,
-    ) -> Result<(), Error>;
-
-    /// Number of shares required to reconstruct the seal key.
-    async fn threshold(&self, conn: &mut db::DatabaseConnection) -> Result<usize, Error>;
-
-    /// Load the shares of `operators` in one query, in the order requested.
-    async fn shares(
-        &self,
-        conn: &mut db::DatabaseConnection,
-        operators: &[OperatorId],
-    ) -> Result<Vec<EncryptedShare>, 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> {
@@ -108,7 +80,8 @@ impl CustodyStore for DieselCustodyStore {
     Ok(())
 }
 
-    async fn threshold(&self, conn: &mut db::DatabaseConnection) -> Result<usize, Error> {
+/// Number of shares required to reconstruct the seal key.
+pub async fn threshold(conn: &mut db::DatabaseConnection) -> Result<usize, Error> {
     let stored: Option<i32> = schema::arbiter_settings::table
         .select(schema::arbiter_settings::shamir_threshold)
         .first(conn)
@@ -120,14 +93,17 @@ impl CustodyStore for DieselCustodyStore {
         .ok_or(Error::BrokenThreshold)
 }
 
-    async fn shares(
+/// Load the shares of `operators` in one query, in the order requested.
-        &self,
+pub async fn shares(
     conn: &mut db::DatabaseConnection,
     operators: &[OperatorId],
 ) -> Result<Vec<EncryptedShare>, Error> {
+    /// (id, then the three columns that make up [`EncryptedShare`])
+    type ShareRow = (Option<OperatorId>, Vec<u8>, Vec<u8>, Vec<u8>);
+
     let wanted: Vec<Option<OperatorId>> = operators.iter().copied().map(Some).collect();
 
-        let rows: Vec<(Option<OperatorId>, Vec<u8>, Vec<u8>, Vec<u8>)> = schema::operator::table
+    let rows: Vec<ShareRow> = schema::operator::table
         .filter(schema::operator::id.eq_any(wanted))
         .select((
             schema::operator::id,
@@ -159,4 +135,3 @@ impl CustodyStore for DieselCustodyStore {
         .map(|id| found.remove(id).ok_or(Error::MissingShare(*id)))
         .collect()
 }
-}

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,11 +766,7 @@ mod tests {
 
     async fn bootstrapped_vault(db: &db::DatabasePool) -> ActorRef<Vault> {
         let actor = Vault::spawn(
-            Vault::new(
+            Vault::new(db.clone(), GlobalActors::spawn_message_bus())
-                db.clone(),
-                GlobalActors::spawn_message_bus(),
-                Arc::new(DieselCustodyStore),
-            )
                 .await
                 .unwrap(),
         );

server/crates/arbiter-server/tests/common/mod.rs

@@ -7,21 +7,16 @@ 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(
+    let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus())
-        db.clone(),
-        GlobalActors::spawn_message_bus(),
-        Arc::new(DieselCustodyStore),
-    )
         .await
         .unwrap();
     actor

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::{
+use std::collections::{HashMap, HashSet};
-    collections::{HashMap, HashSet},
-    sync::Arc,
-};
 use tokio::task::JoinSet;
 
 const TEST_AAD: &[u8] = b"test-aad";
@@ -169,11 +166,7 @@ async fn decrypt_roundtrip_after_high_concurrency() {
     let writes = write_concurrently(actor, "roundtrip", 40).await;
     let expected: HashMap<i32, Vec<u8>> = writes.into_iter().collect();
 
-    let mut decryptor = Vault::new(
+    let mut decryptor = Vault::new(db.clone(), GlobalActors::spawn_message_bus())
-        db.clone(),
-        GlobalActors::spawn_message_bus(),
-        Arc::new(DieselCustodyStore),
-    )
         .await
         .unwrap();
     decryptor

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,11 +21,7 @@ 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(
+    let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus())
-        db.clone(),
-        GlobalActors::spawn_message_bus(),
-        Arc::new(DieselCustodyStore),
-    )
         .await
         .unwrap();
 
@@ -62,11 +57,7 @@ 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(
+    let mut actor = Vault::new(db, GlobalActors::spawn_message_bus())
-        db,
-        GlobalActors::spawn_message_bus(),
-        Arc::new(DieselCustodyStore),
-    )
         .await
         .unwrap();
 
@@ -81,11 +72,7 @@ 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(
+    let mut actor = Vault::new(db, GlobalActors::spawn_message_bus())
-        db,
-        GlobalActors::spawn_message_bus(),
-        Arc::new(DieselCustodyStore),
-    )
         .await
         .unwrap();
 
@@ -100,11 +87,7 @@ async fn new_restores_sealed_state() {
     let actor = common::bootstrapped_vault(&db).await;
     drop(actor);
 
-    let mut actor2 = Vault::new(
+    let mut actor2 = Vault::new(db, GlobalActors::spawn_message_bus())
-        db,
-        GlobalActors::spawn_message_bus(),
-        Arc::new(DieselCustodyStore),
-    )
         .await
         .unwrap();
     let err = actor2.decrypt(1, TEST_AAD.to_vec()).await.unwrap_err();
@@ -124,11 +107,7 @@ async fn unseal_correct_password() {
         .unwrap();
     drop(actor);
 
-    let mut actor = Vault::new(
+    let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus())
-        db.clone(),
-        GlobalActors::spawn_message_bus(),
-        Arc::new(DieselCustodyStore),
-    )
         .await
         .unwrap();
     let seal_key = KeyCell::from([0u8; 32]);
@@ -151,11 +130,7 @@ async fn unseal_wrong_then_correct_password() {
         .unwrap();
     drop(actor);
 
-    let mut actor = Vault::new(
+    let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus())
-        db.clone(),
-        GlobalActors::spawn_message_bus(),
-        Arc::new(DieselCustodyStore),
-    )
         .await
         .unwrap();