refactor(custody): load shares through the Operator model

server/crates/arbiter-server/migrations/2026-02-14-171124-0000_init/up.sql

@@ -53,7 +53,7 @@ create table if not exists operator_identity (
 create unique index if not exists uniq_operator_identity_public_key on operator_identity (public_key);
 
 create table if not exists operator (
-    id integer primary key references operator_identity(id) on delete restrict, -- same id as operator_identity
+    id integer not null primary key references operator_identity(id) on delete restrict, -- same id as operator_identity
 
     share blob not null,
     share_nonce blob not null,

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,19 +1,13 @@
 //! 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
-//! 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};
-use diesel_async::RunQueryDsl;
+use diesel::{ExpressionMethods as _, QueryDsl, SelectableHelper as _, sqlite::Sqlite};
+use diesel_async::{AsyncConnection, RunQueryDsl};
 
 use crate::db::{
-    self,
-    models::{OperatorId, SqliteTimestamp},
+    models::{Operator, OperatorId, SqliteTimestamp},
     schema,
 };
 
@@ -47,18 +41,16 @@ 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)
             .values((
-                schema::operator::id.eq(Some(*operator_id)),
+                schema::operator::id.eq(*operator_id),
                 schema::operator::share.eq(&share.ciphertext),
                 schema::operator::share_nonce.eq(&share.nonce),
                 schema::operator::share_salt.eq(&share.salt),
@@ -81,7 +73,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)
@@ -95,39 +87,27 @@ pub async fn threshold(conn: &mut db::DatabaseConnection) -> Result<usize, Error
 
 /// 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
-        .filter(schema::operator::id.eq_any(wanted))
-        .select((
-            schema::operator::id,
-            schema::operator::share,
-            schema::operator::share_nonce,
-            schema::operator::share_salt,
-        ))
+    let rows: Vec<Operator> = schema::operator::table
+        .filter(schema::operator::id.eq_any(operators))
+        .select(Operator::as_select())
         .load(conn)
         .await?;
 
     let mut found: HashMap<OperatorId, EncryptedShare> = rows
         .into_iter()
-        .filter_map(|(id, ciphertext, nonce, salt)| {
-            id.map(|id| {
+        .map(|row| {
             (
-                    id,
+                row.id,
                 EncryptedShare {
-                        ciphertext,
-                        nonce,
-                        salt,
+                    ciphertext: row.share,
+                    nonce: row.share_nonce,
+                    salt: row.share_salt,
                 },
             )
         })
-        })
         .collect();
 
     operators

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

@@ -292,7 +292,7 @@ pub struct OperatorClient {
     pub updated_at: SqliteTimestamp,
 }
 
-#[derive(Queryable, Debug)]
+#[derive(Queryable, Debug, Selectable)]
 #[diesel(table_name = schema::operator, check_for_backend(Sqlite))]
 pub struct Operator {
     pub id: OperatorId,

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

@@ -155,7 +155,7 @@ diesel::table! {
 
 diesel::table! {
     operator (id) {
-        id -> Nullable<Integer>,
+        id -> Integer,
         share -> Binary,
         share_nonce -> Binary,
         share_salt -> Binary,

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