refactor(custody): drop the single-implementation CustodyStore trait

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

@@ -6,7 +6,7 @@ use crate::{
     },
     db::{
         self,
-        custody::CustodyRecord,
+        custody::{self, CustodyRecord},
         models::{self, RootKeyHistory, RootKeyHistoryId},
         schema::{self},
     },
@@ -63,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,
@@ -227,7 +227,7 @@ impl Vault {
                     .await?;
 
                 if let Some(record) = custody.as_ref() {
-                    db::custody::write_record(&mut *conn, record).await?;
+                    custody::write_record(&mut *conn, record).await?;
                 }
 
                 Result::<_, Error>::Ok(RootKeyHistoryId::from_raw(root_key_history_id))

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

@@ -2,7 +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 the [`db::custody`] functions.
+//! everything it reads or writes goes through [`db::custody`].
 
 use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
 use argon2::RECOMMENDED_SALT_LEN;
@@ -15,7 +15,7 @@ use crate::{
     crypto::{KeyCell, derive_key, encryption::v1::Nonce, shamir},
     db::{
         self,
-        custody::{CustodyRecord, EncryptedShare},
+        custody::{self, CustodyRecord, EncryptedShare},
         models::OperatorId,
     },
 };
@@ -50,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")]
@@ -215,7 +215,7 @@ async fn finalize_unseal(
 ) -> Result<(), Error> {
     let stored = {
         let mut conn = db.get().await?;
-        db::custody::shares(&mut conn, &contributions.operators()).await?
+        custody::shares(&mut conn, &contributions.operators()).await?
     };
 
     let mut plaintext = Vec::with_capacity(stored.len());
@@ -325,7 +325,7 @@ impl VaultCoordinator {
         if matches!(self.state, CoordinatorState::Idle) {
             let threshold = {
                 let mut conn = self.db.get().await?;
-                db::custody::threshold(&mut conn).await?
+                custody::threshold(&mut conn).await?
             };
             self.state = CoordinatorState::Unsealing {
                 threshold,

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

@@ -1,17 +1,18 @@
 //! Storage for Shamir custody material: the reconstruction threshold and the
 //! per-operator encrypted shares of the vault seal 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.
+//! 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 diesel::{ExpressionMethods as _, QueryDsl, sqlite::Sqlite};
-use diesel_async::{AsyncConnection, RunQueryDsl};
+use diesel::{ExpressionMethods as _, QueryDsl};
+use diesel_async::RunQueryDsl;
 
 use crate::db::{
+    self,
     models::{OperatorId, SqliteTimestamp},
     schema,
 };
@@ -46,11 +47,13 @@ pub enum 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>,
+    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)
@@ -78,7 +81,7 @@ pub async fn write_record(
 }
 
 /// Number of shares required to reconstruct the seal key.
-pub async fn threshold(conn: &mut impl AsyncConnection<Backend = Sqlite>) -> Result<usize, Error> {
+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)
@@ -90,14 +93,14 @@ pub async fn threshold(conn: &mut impl AsyncConnection<Backend = Sqlite>) -> Res
         .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>,
+    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<ShareRow> = schema::operator::table

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

@@ -4,16 +4,11 @@ use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
 use arbiter_server::{
     actors::{
         GlobalActors,
-        vault::{Bootstrap, Error as VaultError, GetState, Seal, Vault, VaultState},
+        vault::{Bootstrap, GetState, Seal, VaultState},
         vault_coordinator::{ContributeBootstrap, ContributeUnseal, StartBootstrap},
     },
     crypto::{KeyCell, shamir},
-    db::{
-        self,
-        custody::{CustodyRecord, EncryptedShare},
-        models::OperatorId,
-        schema,
-    },
+    db::{self, models::OperatorId, schema},
 };
 
 use diesel::{ExpressionMethods as _, QueryDsl};
@@ -180,71 +175,6 @@ 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() {