refactor(custody): replace the store abstraction with direct database functions

server/Cargo.toml

@@ -78,7 +78,7 @@ pub_underscore_fields = "allow"
 redundant_pub_crate = "allow" 
 uninhabited_references = "allow" # safe with unsafe_code = "forbid" and standard uninhabited pattern (match *self {})
 too-many-lines = "allow" # this is a very common pattern in server code, and it's not always possible to break it down into smaller modules without hurting readability
-unused_async_trait_impl = "allow" # to pedantic
+unused_async_trait_impl = "allow" # too pedantic
 
 # restriction lints
 alloc_instead_of_core = "warn"

server/crates/arbiter-server/src/actors/bootstrap.rs

@@ -76,6 +76,9 @@ pub enum Error {
 /// declared, so the token stays valid across several registrations and is
 /// retired by the `Bootstrapped` event rather than by first use, whichever
 /// bootstrap path fired it.
+///
+/// Every daemon start mints a fresh token and overwrites the file: a token
+/// handed out by an earlier run is dead.
 pub struct Bootstrapper {
     token: Option<SafeCell<[u8; TOKEN_LENGTH]>>,
     token_path: Option<PathBuf>,
@@ -116,41 +119,13 @@ impl Bootstrapper {
             });
         }
 
-        let registered = diesel::select(diesel::dsl::exists(
-            schema::operator_identity::table.select(schema::operator_identity::id),
-        ))
-        .get_result::<bool>(&mut conn)
-        .await?;
-
-        let token = if registered {
-            match tokio::fs::read_to_string(&path).await {
-                Ok(existing)
-                    if existing.len() == TOKEN_LENGTH
-                        && existing.chars().all(|c| c.is_ascii_alphanumeric()) =>
-                {
-                    let mut cell = SafeCell::new([0u8; TOKEN_LENGTH]);
-                    cell.write().copy_from_slice(existing.as_bytes());
-                    cell
-                }
-                Ok(_) | Err(_) => generate_token(&path).await?,
-            }
-        } else {
-            generate_token(&path).await?
-        };
-
         Ok(Self {
-            token: Some(token),
+            token: Some(generate_token(&path).await?),
             token_path: Some(path),
             events,
         })
     }
 
-    fn is_correct_token(&mut self, token: &[u8]) -> bool {
-        self.token.as_mut().is_some_and(|expected| {
-            expected.read_inline(|bytes| bool::from(bytes.as_ref().ct_eq(token)))
-        })
-    }
-
     async fn forget(&mut self) {
         self.token = None;
         if let Some(path) = self.token_path.take() {
@@ -175,7 +150,9 @@ impl Message<events::Bootstrapped> for Bootstrapper {
 impl Bootstrapper {
     #[message]
     pub fn verify_token(&mut self, token: Vec<u8>) -> bool {
-        self.is_correct_token(&token)
+        self.token.as_mut().is_some_and(|expected| {
+            expected.read_inline(|bytes| bool::from(bytes.as_ref().ct_eq(token.as_slice())))
+        })
     }
 
     #[message]

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<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(), 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,

server/crates/arbiter-server/src/actors/vault/mod.rs

@@ -6,14 +6,13 @@ use crate::{
     },
     db::{
         self,
-        custody::{CustodyRecord, CustodyStore},
+        custody::CustodyRecord,
         models::{self, RootKeyHistory, RootKeyHistoryId},
         schema::{self},
     },
 };
 
 use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
-use std::sync::Arc;
 
 use chrono::Utc;
 use diesel::{
@@ -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(
-        db: db::DatabasePool,
-        events: ActorRef<MessageBus>,
-        custody: Arc<dyn CustodyStore>,
-    ) -> Result<Self, Error> {
+    pub async fn new(db: db::DatabasePool, events: ActorRef<MessageBus>) -> 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?;
+                    db::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

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 the [`db::custody`] functions.
 
 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::{CustodyRecord, EncryptedShare},
         models::OperatorId,
     },
 };
@@ -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(
-        db: db::DatabasePool,
-        vault: ActorRef<Vault>,
-        custody: Arc<dyn CustodyStore>,
-    ) -> Self {
+    pub const fn new(db: db::DatabasePool, vault: ActorRef<Vault>) -> Self {
         Self {
             db,
             vault,
-            custody,
             state: CoordinatorState::Idle,
         }
     }
@@ -192,12 +184,14 @@ async fn finalize_bootstrap(
     let mut shares = shamir::split_key(threshold, total, &mut seal_key, UnwrapErr(SysRng))
         .map_err(|error| Error::Shamir(error.to_string()))?;
 
+    if shares.len() < total {
+        return Err(Error::Shamir("missing share for operator".to_owned()));
+    }
+
     let mut encrypted = Vec::with_capacity(total);
-    for (index, (operator_id, passphrase)) in contributions.0.iter_mut().enumerate() {
-        let share = shares
-            .read_inline(|shares| shares.get(index).cloned())
-            .ok_or_else(|| Error::Shamir("missing share for operator".to_owned()))?;
-        encrypted.push((*operator_id, encrypt_share(passphrase, &share)?));
+    for ((operator_id, passphrase), share) in contributions.0.iter_mut().zip(shares.iter_mut()) {
+        let share = share.read_inline(|share| encrypt_share(passphrase, share))?;
+        encrypted.push((*operator_id, share));
     }
 
     vault
@@ -215,24 +209,18 @@ 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
-            .shares(&mut conn, &contributions.operators())
-            .await?
+        db::custody::shares(&mut conn, &contributions.operators()).await?
     };
 
-    let mut plaintext = SafeCell::new(Vec::with_capacity(stored.len()));
+    let mut plaintext = Vec::with_capacity(stored.len());
     for ((_, passphrase), share) in contributions.0.iter_mut().zip(stored) {
-        let mut decrypted = decrypt_share(passphrase, share)?;
-        decrypted.read_inline(|share| {
-            plaintext.write_inline(|shares| shares.push(share.clone()));
-        });
+        plaintext.push(decrypt_share(passphrase, share)?);
     }
 
     let seal_key = shamir::combine_shares(threshold, &mut plaintext)
@@ -337,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?
+                db::custody::threshold(&mut conn).await?
             };
             self.state = CoordinatorState::Unsealing {
                 threshold,
@@ -375,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 {

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,13 +231,9 @@ mod tests {
 
     async fn bootstrapped_vault(db: &db::DatabasePool) -> ActorRef<Vault> {
         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 {

server/crates/arbiter-server/src/crypto/shamir.rs

@@ -16,10 +16,13 @@ pub enum ShamirError {
     Combine(String),
 }
 
-/// Return the required majority threshold for an ordinary operator committee.
+/// Return the required threshold for a Shamir share pool of `committee_size`.
 ///
-/// Committees of two are rejected: a majority of two is two, which gives each
-/// member a veto over every unseal without giving either one recovery.
+/// A pool of two is rejected: a majority of two is two, which gives each holder
+/// a veto over every unseal without giving either one recovery. That rejects no
+/// supported committee, because a two-operator vault must carry at least one
+/// recovery share and so never splits into a pool of two -- see
+/// `docs/ARCHITECTURE.md` 3.9.
 #[expect(
     clippy::integer_division,
     reason = "majority thresholds use integer arithmetic"
@@ -40,7 +43,7 @@ pub fn split_key(
     total: usize,
     key: &mut KeyCell,
     rng: impl CryptoRng,
-) -> Result<SafeCell<Vec<Vec<u8>>>, ShamirError> {
+) -> Result<Vec<SafeCell<Vec<u8>>>, ShamirError> {
     if total == 0 || threshold == 0 || threshold > total || total == 2 || total > MAX_COMMITTEE_SIZE
     {
         return Err(ShamirError::Split(
@@ -48,20 +51,23 @@ pub fn split_key(
         ));
     }
 
+    // Nothing to interpolate when one share suffices.
+    if threshold == 1 {
+        return Ok(key.0.read_inline(|key| {
+            std::iter::repeat_with(|| SafeCell::new(key.as_slice().to_vec()))
+                .take(total)
+                .collect()
+        }));
+    }
+
     key.0.read_inline(|key| {
         let key: &[u8; 32] = key
             .as_slice()
             .try_into()
             .map_err(|_| ShamirError::Split("unexpected seal key length".to_owned()))?;
 
-        if threshold == 1 {
-            return Ok(SafeCell::new(
-                std::iter::repeat_n(key.to_vec(), total).collect(),
-            ));
-        }
-
         Gf256::split_array(threshold, total, key, rng)
-            .map(SafeCell::new)
+            .map(|shares| shares.into_iter().map(SafeCell::new).collect())
             .map_err(|error| ShamirError::Split(format!("{error:?}")))
     })
 }
@@ -74,30 +80,43 @@ pub fn split_key(
 /// sized.
 pub fn combine_shares(
     threshold: usize,
-    shares: &mut SafeCell<Vec<Vec<u8>>>,
+    shares: &mut [SafeCell<Vec<u8>>],
 ) -> Result<KeyCell, ShamirError> {
     if threshold == 0 {
         return Err(ShamirError::Combine("threshold is zero".to_owned()));
     }
-    if shares.read().len() < threshold {
+    if shares.len() < threshold {
         return Err(ShamirError::Combine(
             "not enough shares supplied".to_owned(),
         ));
     }
 
-    let combined = shares.read_inline(|shares| {
-        if threshold == 1 {
-            let share = shares
-                .first()
-                .ok_or_else(|| ShamirError::Combine("no shares supplied".to_owned()))?;
-            return Ok(SafeCell::new(share.clone()));
-        }
-        Gf256::combine_array(shares)
+    // Mirror of the one-of-one case in [`split_key`]: the share is the key.
+    if threshold == 1 {
+        let share = shares
+            .first_mut()
+            .ok_or_else(|| ShamirError::Combine("no shares supplied".to_owned()))?;
+        return reconstructed_key(share.read_inline(|share| SafeCell::new(share.clone())));
+    }
+
+    let mut gathered = SafeCell::new(Vec::with_capacity(shares.len()));
+    for share in shares.iter_mut() {
+        share.read_inline(|share| {
+            gathered.write_inline(|gathered| gathered.push(share.clone()));
+        });
+    }
+
+    let combined = gathered.read_inline(|gathered| {
+        Gf256::combine_array(gathered.as_slice())
             .map(SafeCell::new)
             .map_err(|error| ShamirError::Combine(format!("{error:?}")))
     })?;
 
-    KeyCell::try_from(combined)
+    reconstructed_key(combined)
+}
+
+fn reconstructed_key(bytes: SafeCell<Vec<u8>>) -> Result<KeyCell, ShamirError> {
+    KeyCell::try_from(bytes)
         .map_err(|()| ShamirError::Combine("unexpected reconstructed key length".to_owned()))
 }
 
@@ -108,6 +127,7 @@ mod tests {
     use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
     use rand::rngs::SysRng;
     use rand_core::UnwrapErr;
+    use rstest::rstest;
 
     fn key_bytes(mut key: KeyCell) -> [u8; 32] {
         key.0.read_inline(|key| {
@@ -117,28 +137,29 @@ mod tests {
         })
     }
 
-    fn select(shares: &mut SafeCell<Vec<Vec<u8>>>, indexes: &[usize]) -> SafeCell<Vec<Vec<u8>>> {
-        shares.read_inline(|shares| {
-            SafeCell::new(
-                indexes
-                    .iter()
-                    .filter_map(|index| shares.get(*index).cloned())
-                    .collect(),
-            )
-        })
+    fn select(shares: &mut [SafeCell<Vec<u8>>], indexes: &[usize]) -> Vec<SafeCell<Vec<u8>>> {
+        indexes
+            .iter()
+            .filter_map(|index| {
+                shares
+                    .get_mut(*index)
+                    .map(|share| share.read_inline(|share| SafeCell::new(share.clone())))
+            })
+            .collect()
     }
 
-    #[test]
-    fn threshold_shares_reconstruct_fixed_key() {
+    #[rstest]
+    #[case(&[0, 1])]
+    #[case(&[0, 2])]
+    #[case(&[1, 2])]
+    fn threshold_shares_reconstruct_fixed_key(#[case] indexes: &[usize]) {
         let expected = [9_u8; 32];
         let mut key = KeyCell::from(expected);
         let rng = UnwrapErr(SysRng);
         let mut shares = split_key(2, 3, &mut key, rng).expect("split should succeed");
-        for indexes in [[0_usize, 1_usize], [0, 2], [1, 2]] {
-            let mut selected = select(&mut shares, &indexes);
-            let combined = combine_shares(2, &mut selected).expect("combine should succeed");
-            assert_eq!(key_bytes(combined), expected);
-        }
+        let mut selected = select(&mut shares, indexes);
+        let combined = combine_shares(2, &mut selected).expect("combine should succeed");
+        assert_eq!(key_bytes(combined), expected);
     }
 
     #[test]
@@ -163,22 +184,23 @@ mod tests {
         );
     }
 
-    #[test]
-    fn empty_committee_has_no_threshold() {
-        assert_eq!(shamir_threshold(0), None);
+    #[rstest]
+    #[case(0, None)]
+    #[case(1, Some(1))]
+    #[case(2, None)]
+    #[case(3, Some(2))]
+    #[case(4, Some(3))]
+    #[case(MAX_COMMITTEE_SIZE, Some(128))]
+    #[case(MAX_COMMITTEE_SIZE + 1, None)]
+    fn committee_threshold_is_a_majority(
+        #[case] committee_size: usize,
+        #[case] expected: Option<usize>,
+    ) {
+        assert_eq!(shamir_threshold(committee_size), expected);
     }
 
     #[test]
-    fn committee_threshold_is_majority_for_three_or_more() {
-        assert_eq!(shamir_threshold(1), Some(1));
-        assert_eq!(shamir_threshold(3), Some(2));
-        assert_eq!(shamir_threshold(4), Some(3));
-    }
-
-    #[test]
-    fn oversized_committee_has_no_threshold() {
-        assert_eq!(shamir_threshold(MAX_COMMITTEE_SIZE), Some(128));
-        assert_eq!(shamir_threshold(MAX_COMMITTEE_SIZE + 1), None);
+    fn oversized_committee_is_rejected_by_split() {
         let mut key = KeyCell::from([1_u8; 32]);
         let rng = UnwrapErr(SysRng);
         assert!(
@@ -189,7 +211,6 @@ mod tests {
 
     #[test]
     fn two_operator_committee_is_explicitly_unsupported() {
-        assert_eq!(shamir_threshold(2), None);
         let mut key = KeyCell::from([7_u8; 32]);
         let rng = UnwrapErr(SysRng);
         assert!(

server/crates/arbiter-server/src/db/custody.rs

@@ -1,20 +1,17 @@
 //! 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.
+//! Every query lives here so that the actors above hold no Diesel code of their
+//! own. The functions borrow 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.
 
 use std::collections::HashMap;
 
-use async_trait::async_trait;
-use diesel::{ExpressionMethods as _, QueryDsl};
-use diesel_async::RunQueryDsl;
+use diesel::{ExpressionMethods as _, QueryDsl, sqlite::Sqlite};
+use diesel_async::{AsyncConnection, RunQueryDsl};
 
 use crate::db::{
-    self,
     models::{OperatorId, SqliteTimestamp},
     schema,
 };
@@ -46,117 +43,92 @@ 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 impl AsyncConnection<Backend = Sqlite>,
+    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<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> {
-        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)))
+    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<usize, Error> {
-        let stored: Option<i32> = 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<Vec<EncryptedShare>, Error> {
-        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
-            .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<OperatorId, EncryptedShare> = 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 impl AsyncConnection<Backend = Sqlite>) -> Result<usize, Error> {
+    let stored: Option<i32> = 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)
+}
+
+/// One row of the share query: operator id, ciphertext, nonce, salt.
+type ShareRow = (Option<OperatorId>, Vec<u8>, Vec<u8>, Vec<u8>);
+
+/// Load the shares of `operators` in one query, in the order requested.
+pub async fn shares(
+    conn: &mut impl AsyncConnection<Backend = Sqlite>,
+    operators: &[OperatorId],
+) -> Result<Vec<EncryptedShare>, Error> {
+    let wanted: Vec<Option<OperatorId>> = operators.iter().copied().map(Some).collect();
+
+    let rows: Vec<ShareRow> = 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<OperatorId, EncryptedShare> = 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()
 }

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<Vault> {
         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 {

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

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<i32, Vec<u8>> = 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

server/crates/arbiter-server/tests/vault/custody.rs

@@ -4,11 +4,16 @@ use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
 use arbiter_server::{
     actors::{
         GlobalActors,
-        vault::{Bootstrap, GetState, Seal, VaultState},
+        vault::{Bootstrap, Error as VaultError, GetState, Seal, Vault, VaultState},
         vault_coordinator::{ContributeBootstrap, ContributeUnseal, StartBootstrap},
     },
     crypto::{KeyCell, shamir},
-    db::{self, models::OperatorId, schema},
+    db::{
+        self,
+        custody::{CustodyRecord, EncryptedShare},
+        models::OperatorId,
+        schema,
+    },
 };
 
 use diesel::{ExpressionMethods as _, QueryDsl};
@@ -175,6 +180,71 @@ async fn refused_bootstrap_stores_no_shares() {
     );
 }
 
+/// A failing custody write must take the whole bootstrap down with it: a vault
+/// that kept its root key but lost the shares could never be unsealed again.
+#[tokio::test]
+#[test_log::test]
+async fn custody_write_failure_rolls_back_bootstrap() {
+    let db = db::create_test_pool().await;
+    let operators = register_operators(&db, 1).await;
+    let record = CustodyRecord {
+        threshold: 1,
+        shares: operators
+            .into_iter()
+            .map(|operator_id| {
+                (
+                    operator_id,
+                    EncryptedShare {
+                        ciphertext: vec![1; 32],
+                        nonce: vec![2; 24],
+                        salt: vec![3; 16],
+                    },
+                )
+            })
+            .collect(),
+    };
+    let mut vault = Vault::new(db.clone(), GlobalActors::spawn_message_bus())
+        .await
+        .unwrap();
+
+    // The threshold update is the last statement of the custody write, so the
+    // trigger fails the transaction once the share row is already in place.
+    let mut conn = db.get().await.unwrap();
+    diesel::sql_query(
+        "CREATE TRIGGER fail_custody_threshold BEFORE UPDATE OF shamir_threshold ON arbiter_settings BEGIN SELECT RAISE(ABORT, 'forced custody failure'); END;",
+    )
+    .execute(&mut conn)
+    .await
+    .unwrap();
+    drop(conn);
+
+    let error = vault
+        .bootstrap(KeyCell::from([4u8; 32]), Some(record))
+        .await
+        .expect_err("a failing custody write must fail the bootstrap");
+    assert!(
+        matches!(error, VaultError::Custody(_)),
+        "expected a custody error, got {error:?}"
+    );
+    assert_eq!(vault.get_state(), VaultState::Unbootstrapped);
+    assert_eq!(stored_share_count(&db).await, 0);
+    assert_eq!(stored_threshold(&db).await, None);
+
+    let mut conn = db.get().await.unwrap();
+    let root_count: i64 = schema::root_key_history::table
+        .count()
+        .get_result(&mut conn)
+        .await
+        .unwrap();
+    let root_key_id: Option<i32> = schema::arbiter_settings::table
+        .select(schema::arbiter_settings::root_key_id)
+        .first(&mut conn)
+        .await
+        .unwrap();
+    assert_eq!(root_count, 0, "the root key write must roll back as well");
+    assert_eq!(root_key_id, None);
+}
+
 #[tokio::test]
 #[test_log::test]
 async fn oversized_and_degenerate_committees_are_rejected() {

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