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

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

@@ -6,7 +6,7 @@ use crate::{
     },
     db::{
         self,
-        custody::{self, CustodyRecord},
+        custody::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] custody::Error),
+    Custody(#[from] db::custody::Error),
 
     #[error("Broken database")]
     BrokenDatabase,
@@ -227,7 +227,7 @@ impl Vault {
                     .await?;
 
                 if let Some(record) = custody.as_ref() {
-                    custody::write_record(&mut *conn, record).await?;
+                    db::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 [`db::custody`].
+//! everything it reads or writes goes through the [`db::custody`] functions.
 
 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::{self, CustodyRecord, EncryptedShare},
+        custody::{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] custody::Error),
+    Custody(#[from] db::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?;
-        custody::shares(&mut conn, &contributions.operators()).await?
+        db::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?;
-                custody::threshold(&mut conn).await?
+                db::custody::threshold(&mut conn).await?
             };
             self.state = CoordinatorState::Unsealing {
                 threshold,

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

@@ -1,18 +1,17 @@
 //! Storage for Shamir custody material: the reconstruction threshold and the
 //! per-operator encrypted shares of the vault seal key.
 //!
-//! The queries live here so that the actors above hold no Diesel code of their
+//! Every query lives 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
+//! own. The functions borrow the caller's connection instead of taking one from
-//! from the pool, which is what lets the vault write custody material inside
+//! the pool, which lets the vault write custody material inside the same
-//! the same transaction that stores the root key.
+//! transaction that stores the root key.
 
 use std::collections::HashMap;
 
-use diesel::{ExpressionMethods as _, QueryDsl};
+use diesel::{ExpressionMethods as _, QueryDsl, sqlite::Sqlite};
-use diesel_async::RunQueryDsl;
+use diesel_async::{AsyncConnection, RunQueryDsl};
 
 use crate::db::{
-    self,
     models::{OperatorId, SqliteTimestamp},
     schema,
 };
@@ -47,13 +46,11 @@ 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 db::DatabaseConnection,
+    conn: &mut impl AsyncConnection<Backend = Sqlite>,
     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)
@@ -81,7 +78,7 @@ pub async fn write_record(
 }
 
 /// Number of shares required to reconstruct the seal key.
-pub async fn threshold(conn: &mut db::DatabaseConnection) -> Result<usize, Error> {
+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)
@@ -93,14 +90,14 @@ pub async fn threshold(conn: &mut db::DatabaseConnection) -> Result<usize, Error
         .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 db::DatabaseConnection,
+    conn: &mut impl AsyncConnection<Backend = Sqlite>,
     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,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() {