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

This commit is contained in:
CleverWild
2026-09-12 16:47:35 +02:00
parent d49c39130a
commit 3db7ece6c0
10 changed files with 209 additions and 256 deletions

View File

@@ -3,14 +3,9 @@ use crate::{
bootstrap::Bootstrapper, evm::EvmActor, flow_coordinator::FlowCoordinator, bootstrap::Bootstrapper, evm::EvmActor, flow_coordinator::FlowCoordinator,
operator_registry::OperatorRegistry, vault::Vault, vault_coordinator::VaultCoordinator, 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::actor::{ActorRef, Spawn};
use kameo_actors::{DeliveryStrategy, message_bus::MessageBus}; use kameo_actors::{DeliveryStrategy, message_bus::MessageBus};
use thiserror::Error; use thiserror::Error;
@@ -50,12 +45,10 @@ impl GlobalActors {
pub async fn spawn(db: db::DatabasePool) -> Result<Self, SpawnError> { pub async fn spawn(db: db::DatabasePool) -> Result<Self, SpawnError> {
let events = Self::spawn_message_bus(); 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 bootstrapper = Bootstrapper::spawn(Bootstrapper::new(&db, events.clone()).await?);
let vault_coordinator = 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()); let operator_registry = OperatorRegistry::spawn(OperatorRegistry::default());
Ok(Self { Ok(Self {
bootstrapper, bootstrapper,

View File

@@ -6,14 +6,13 @@ use crate::{
}, },
db::{ db::{
self, self,
custody::{CustodyRecord, CustodyStore}, custody::CustodyRecord,
models::{self, RootKeyHistory, RootKeyHistoryId}, models::{self, RootKeyHistory, RootKeyHistoryId},
schema::{self}, schema::{self},
}, },
}; };
use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _}; use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
use std::sync::Arc;
use chrono::Utc; use chrono::Utc;
use diesel::{ use diesel::{
@@ -101,17 +100,12 @@ pub struct Vault {
db: db::DatabasePool, db: db::DatabasePool,
state: State, state: State,
events: ActorRef<MessageBus>, events: ActorRef<MessageBus>,
custody: Arc<dyn CustodyStore>,
unseal_failures: u32, unseal_failures: u32,
} }
#[messages] #[messages]
impl Vault { 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 state = {
let mut conn = db.get().await?; let mut conn = db.get().await?;
@@ -133,7 +127,6 @@ impl Vault {
db, db,
state, state,
events, events,
custody,
unseal_failures: 0, unseal_failures: 0,
}) })
} }
@@ -213,7 +206,6 @@ impl Vault {
let mut conn = self.db.get().await?; let mut conn = self.db.get().await?;
let data_encryption_nonce_bytes = data_encryption_nonce.to_vec(); let data_encryption_nonce_bytes = data_encryption_nonce.to_vec();
let custody_store = Arc::clone(&self.custody);
let root_key_history_id = conn let root_key_history_id = conn
.transaction(async |conn| { .transaction(async |conn| {
let root_key_history_id = insert_into(schema::root_key_history::table) let root_key_history_id = insert_into(schema::root_key_history::table)
@@ -235,7 +227,7 @@ impl Vault {
.await?; .await?;
if let Some(record) = custody.as_ref() { 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)) Result::<_, Error>::Ok(RootKeyHistoryId::from_raw(root_key_history_id))
@@ -457,18 +449,14 @@ impl Vault {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use crate::{actors::GlobalActors, db::custody::DieselCustodyStore}; use crate::actors::GlobalActors;
use super::*; use super::*;
async fn bootstrapped_actor(db: &db::DatabasePool) -> Vault { 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(), .await
GlobalActors::spawn_message_bus(), .unwrap();
Arc::new(DieselCustodyStore),
)
.await
.unwrap();
let seal_key = KeyCell::from([0u8; 32]); let seal_key = KeyCell::from([0u8; 32]);
actor.bootstrap(seal_key, None).await.unwrap(); actor.bootstrap(seal_key, None).await.unwrap();
actor actor

View File

@@ -2,9 +2,7 @@
//! //!
//! The coordinator collects one passphrase per committee member, then hands the //! The coordinator collects one passphrase per committee member, then hands the
//! assembled material to [`Vault`] in a single message. It owns no Diesel code: //! 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 the [`db::custody`] functions.
use std::sync::Arc;
use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _}; use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
use argon2::RECOMMENDED_SALT_LEN; use argon2::RECOMMENDED_SALT_LEN;
@@ -17,7 +15,7 @@ use crate::{
crypto::{KeyCell, derive_key, encryption::v1::Nonce, shamir}, crypto::{KeyCell, derive_key, encryption::v1::Nonce, shamir},
db::{ db::{
self, self,
custody::{CustodyRecord, CustodyStore, EncryptedShare}, custody::{CustodyRecord, EncryptedShare},
models::OperatorId, models::OperatorId,
}, },
}; };
@@ -111,20 +109,14 @@ enum CoordinatorState {
pub struct VaultCoordinator { pub struct VaultCoordinator {
db: db::DatabasePool, db: db::DatabasePool,
vault: ActorRef<Vault>, vault: ActorRef<Vault>,
custody: Arc<dyn CustodyStore>,
state: CoordinatorState, state: CoordinatorState,
} }
impl VaultCoordinator { 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 { Self {
db, db,
vault, vault,
custody,
state: CoordinatorState::Idle, state: CoordinatorState::Idle,
} }
} }
@@ -217,16 +209,13 @@ async fn finalize_bootstrap(
/// Reconstruct the seal key from the contributed passphrases and unseal. /// Reconstruct the seal key from the contributed passphrases and unseal.
async fn finalize_unseal( async fn finalize_unseal(
db: &db::DatabasePool, db: &db::DatabasePool,
custody: &Arc<dyn CustodyStore>,
vault: &ActorRef<Vault>, vault: &ActorRef<Vault>,
threshold: usize, threshold: usize,
contributions: &mut Contributions, contributions: &mut Contributions,
) -> Result<(), Error> { ) -> Result<(), Error> {
let stored = { let stored = {
let mut conn = db.get().await?; let mut conn = db.get().await?;
custody db::custody::shares(&mut conn, &contributions.operators()).await?
.shares(&mut conn, &contributions.operators())
.await?
}; };
let mut plaintext = Vec::with_capacity(stored.len()); let mut plaintext = Vec::with_capacity(stored.len());
@@ -336,7 +325,7 @@ impl VaultCoordinator {
if matches!(self.state, CoordinatorState::Idle) { if matches!(self.state, CoordinatorState::Idle) {
let threshold = { let threshold = {
let mut conn = self.db.get().await?; let mut conn = self.db.get().await?;
self.custody.threshold(&mut conn).await? db::custody::threshold(&mut conn).await?
}; };
self.state = CoordinatorState::Unsealing { self.state = CoordinatorState::Unsealing {
threshold, threshold,
@@ -374,15 +363,7 @@ impl VaultCoordinator {
unreachable!("state was matched as Unsealing above") 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), Ok(()) => Ok(true),
Err(error) => { Err(error) => {
self.state = CoordinatorState::Unsealing { self.state = CoordinatorState::Unsealing {

View File

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

View File

@@ -1,20 +1,17 @@
//! Storage for Shamir custody material: the reconstruction threshold and the //! Storage for Shamir custody material: the reconstruction threshold and the
//! per-operator encrypted shares of the vault seal key. //! per-operator encrypted shares of the vault seal key.
//! //!
//! Every query lives behind [`CustodyStore`] so that the actors above it hold //! Every query lives here so that the actors above hold no Diesel code of their
//! no Diesel code of their own. [`CustodyStore::write_record`] borrows the //! own. The functions borrow the caller's connection instead of taking one from
//! caller's connection instead of taking one from the pool, which lets the //! the pool, which lets the vault write custody material inside the same
//! vault write custody material inside the same transaction that stores the //! transaction that stores the root key.
//! root key.
use std::collections::HashMap; use std::collections::HashMap;
use async_trait::async_trait; use diesel::{ExpressionMethods as _, QueryDsl, sqlite::Sqlite};
use diesel::{ExpressionMethods as _, QueryDsl}; use diesel_async::{AsyncConnection, RunQueryDsl};
use diesel_async::RunQueryDsl;
use crate::db::{ use crate::db::{
self,
models::{OperatorId, SqliteTimestamp}, models::{OperatorId, SqliteTimestamp},
schema, schema,
}; };
@@ -46,117 +43,92 @@ pub enum Error {
MissingShare(OperatorId), MissingShare(OperatorId),
} }
#[async_trait] /// Persist threshold and shares on the caller's connection, joining any
pub trait CustodyStore: std::fmt::Debug + Send + Sync { /// transaction the caller has already opened.
/// Persist threshold and shares on the caller's connection, joining any pub async fn write_record(
/// transaction the caller has already opened. conn: &mut impl AsyncConnection<Backend = Sqlite>,
async fn write_record( record: &CustodyRecord,
&self, ) -> Result<(), Error> {
conn: &mut db::DatabaseConnection, let threshold = i32::try_from(record.threshold).map_err(|_| Error::BrokenThreshold)?;
record: &CustodyRecord,
) -> Result<(), Error>;
/// Number of shares required to reconstruct the seal key. let now = SqliteTimestamp::now();
async fn threshold(&self, conn: &mut db::DatabaseConnection) -> Result<usize, Error>; for (operator_id, share) in &record.shares {
diesel::replace_into(schema::operator::table)
/// Load the shares of `operators` in one query, in the order requested. .values((
async fn shares( schema::operator::id.eq(Some(*operator_id)),
&self, schema::operator::share.eq(&share.ciphertext),
conn: &mut db::DatabaseConnection, schema::operator::share_nonce.eq(&share.nonce),
operators: &[OperatorId], schema::operator::share_salt.eq(&share.salt),
) -> Result<Vec<EncryptedShare>, Error>; schema::operator::created_at.eq(now.clone()),
} schema::operator::updated_at.eq(now.clone()),
))
/// 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)))
.execute(&mut *conn) .execute(&mut *conn)
.await?; .await?;
if updated != 1 {
return Err(Error::MissingSettings);
}
Ok(())
} }
async fn threshold(&self, conn: &mut db::DatabaseConnection) -> Result<usize, Error> { let updated = diesel::update(schema::arbiter_settings::table)
let stored: Option<i32> = schema::arbiter_settings::table .set(schema::arbiter_settings::shamir_threshold.eq(Some(threshold)))
.select(schema::arbiter_settings::shamir_threshold) .execute(&mut *conn)
.first(conn) .await?;
.await?; if updated != 1 {
return Err(Error::MissingSettings);
stored
.and_then(|value| usize::try_from(value).ok())
.filter(|threshold| *threshold > 0)
.ok_or(Error::BrokenThreshold)
} }
async fn shares( Ok(())
&self, }
conn: &mut db::DatabaseConnection,
operators: &[OperatorId], /// Number of shares required to reconstruct the seal key.
) -> Result<Vec<EncryptedShare>, Error> { pub async fn threshold(conn: &mut impl AsyncConnection<Backend = Sqlite>) -> Result<usize, Error> {
let wanted: Vec<Option<OperatorId>> = operators.iter().copied().map(Some).collect(); let stored: Option<i32> = schema::arbiter_settings::table
.select(schema::arbiter_settings::shamir_threshold)
let rows: Vec<(Option<OperatorId>, Vec<u8>, Vec<u8>, Vec<u8>)> = schema::operator::table .first(conn)
.filter(schema::operator::id.eq_any(wanted)) .await?;
.select((
schema::operator::id, stored
schema::operator::share, .and_then(|value| usize::try_from(value).ok())
schema::operator::share_nonce, .filter(|threshold| *threshold > 0)
schema::operator::share_salt, .ok_or(Error::BrokenThreshold)
)) }
.load(conn)
.await?; /// One row of the share query: operator id, ciphertext, nonce, salt.
type ShareRow = (Option<OperatorId>, Vec<u8>, Vec<u8>, Vec<u8>);
let mut found: HashMap<OperatorId, EncryptedShare> = rows
.into_iter() /// Load the shares of `operators` in one query, in the order requested.
.filter_map(|(id, ciphertext, nonce, salt)| { pub async fn shares(
id.map(|id| { conn: &mut impl AsyncConnection<Backend = Sqlite>,
( operators: &[OperatorId],
id, ) -> Result<Vec<EncryptedShare>, Error> {
EncryptedShare { let wanted: Vec<Option<OperatorId>> = operators.iter().copied().map(Some).collect();
ciphertext,
nonce, let rows: Vec<ShareRow> = schema::operator::table
salt, .filter(schema::operator::id.eq_any(wanted))
}, .select((
) schema::operator::id,
}) schema::operator::share,
}) schema::operator::share_nonce,
.collect(); schema::operator::share_salt,
))
operators .load(conn)
.iter() .await?;
.map(|id| found.remove(id).ok_or(Error::MissingShare(*id)))
.collect() 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()
} }

View File

@@ -501,9 +501,6 @@ impl Engine {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use std::sync::Arc;
use crate::db::custody::DieselCustodyStore;
use alloy::primitives::{Address, Bytes, U256, address}; use alloy::primitives::{Address, Bytes, U256, address};
use chrono::{Duration, Utc}; use chrono::{Duration, Utc};
use diesel::{SelectableHelper, insert_into}; use diesel::{SelectableHelper, insert_into};
@@ -769,13 +766,9 @@ mod tests {
async fn bootstrapped_vault(db: &db::DatabasePool) -> ActorRef<Vault> { async fn bootstrapped_vault(db: &db::DatabasePool) -> ActorRef<Vault> {
let actor = Vault::spawn( let actor = Vault::spawn(
Vault::new( Vault::new(db.clone(), GlobalActors::spawn_message_bus())
db.clone(), .await
GlobalActors::spawn_message_bus(), .unwrap(),
Arc::new(DieselCustodyStore),
)
.await
.unwrap(),
); );
actor actor
.ask(Bootstrap { .ask(Bootstrap {

View File

@@ -7,23 +7,18 @@ use arbiter_proto::transport::{Bi, Error, Receiver, Sender};
use arbiter_server::{ use arbiter_server::{
actors::{GlobalActors, vault::Vault}, actors::{GlobalActors, vault::Vault},
crypto::KeyCell, crypto::KeyCell,
db::{self, custody::DieselCustodyStore, schema}, db::{self, schema},
}; };
use async_trait::async_trait; use async_trait::async_trait;
use diesel::QueryDsl; use diesel::QueryDsl;
use diesel_async::RunQueryDsl; use diesel_async::RunQueryDsl;
use std::sync::Arc;
use tokio::sync::mpsc; use tokio::sync::mpsc;
pub(crate) async fn bootstrapped_vault(db: &db::DatabasePool) -> Vault { 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(), .await
GlobalActors::spawn_message_bus(), .unwrap();
Arc::new(DieselCustodyStore),
)
.await
.unwrap();
actor actor
.bootstrap(KeyCell::from([0u8; 32]), None) .bootstrap(KeyCell::from([0u8; 32]), None)
.await .await

View File

@@ -6,16 +6,13 @@ use arbiter_server::{
vault::{CreateNew, Error, Vault}, vault::{CreateNew, Error, Vault},
}, },
crypto::KeyCell, crypto::KeyCell,
db::{self, custody::DieselCustodyStore, models, schema}, db::{self, models, schema},
}; };
use diesel::{ExpressionMethods as _, QueryDsl, SelectableHelper, dsl::sql_query}; use diesel::{ExpressionMethods as _, QueryDsl, SelectableHelper, dsl::sql_query};
use diesel_async::RunQueryDsl; use diesel_async::RunQueryDsl;
use kameo::actor::{ActorRef, Spawn as _}; use kameo::actor::{ActorRef, Spawn as _};
use std::{ use std::collections::{HashMap, HashSet};
collections::{HashMap, HashSet},
sync::Arc,
};
use tokio::task::JoinSet; use tokio::task::JoinSet;
const TEST_AAD: &[u8] = b"test-aad"; 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 writes = write_concurrently(actor, "roundtrip", 40).await;
let expected: HashMap<i32, Vec<u8>> = writes.into_iter().collect(); 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(), .await
GlobalActors::spawn_message_bus(), .unwrap();
Arc::new(DieselCustodyStore),
)
.await
.unwrap();
decryptor decryptor
.try_unseal(KeyCell::from([0u8; 32])) .try_unseal(KeyCell::from([0u8; 32]))
.await .await

View File

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

View File

@@ -9,12 +9,11 @@ use arbiter_server::{
KeyCell, KeyCell,
encryption::v1::{Nonce, ROOT_KEY_TAG}, encryption::v1::{Nonce, ROOT_KEY_TAG},
}, },
db::{self, custody::DieselCustodyStore, models, schema}, db::{self, models, schema},
}; };
use diesel::{QueryDsl, SelectableHelper}; use diesel::{QueryDsl, SelectableHelper};
use diesel_async::RunQueryDsl; use diesel_async::RunQueryDsl;
use std::sync::Arc;
const TEST_AAD: &[u8] = b"test-aad"; const TEST_AAD: &[u8] = b"test-aad";
@@ -22,13 +21,9 @@ const TEST_AAD: &[u8] = b"test-aad";
#[test_log::test] #[test_log::test]
async fn bootstrap() { async fn bootstrap() {
let db = db::create_test_pool().await; 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(), .await
GlobalActors::spawn_message_bus(), .unwrap();
Arc::new(DieselCustodyStore),
)
.await
.unwrap();
let seal_key = KeyCell::from([0u8; 32]); let seal_key = KeyCell::from([0u8; 32]);
actor.bootstrap(seal_key, None).await.unwrap(); actor.bootstrap(seal_key, None).await.unwrap();
@@ -62,13 +57,9 @@ async fn bootstrap_rejects_double() {
#[test_log::test] #[test_log::test]
async fn create_new_before_bootstrap_fails() { async fn create_new_before_bootstrap_fails() {
let db = db::create_test_pool().await; let db = db::create_test_pool().await;
let mut actor = Vault::new( let mut actor = Vault::new(db, GlobalActors::spawn_message_bus())
db, .await
GlobalActors::spawn_message_bus(), .unwrap();
Arc::new(DieselCustodyStore),
)
.await
.unwrap();
let err = actor let err = actor
.create_new(SafeCell::new(b"data".to_vec()), TEST_AAD.to_vec()) .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] #[test_log::test]
async fn decrypt_before_bootstrap_fails() { async fn decrypt_before_bootstrap_fails() {
let db = db::create_test_pool().await; let db = db::create_test_pool().await;
let mut actor = Vault::new( let mut actor = Vault::new(db, GlobalActors::spawn_message_bus())
db, .await
GlobalActors::spawn_message_bus(), .unwrap();
Arc::new(DieselCustodyStore),
)
.await
.unwrap();
let err = actor.decrypt(1, TEST_AAD.to_vec()).await.unwrap_err(); let err = actor.decrypt(1, TEST_AAD.to_vec()).await.unwrap_err();
assert!(matches!(err, Error::NotBootstrapped)); assert!(matches!(err, Error::NotBootstrapped));
@@ -100,13 +87,9 @@ async fn new_restores_sealed_state() {
let actor = common::bootstrapped_vault(&db).await; let actor = common::bootstrapped_vault(&db).await;
drop(actor); drop(actor);
let mut actor2 = Vault::new( let mut actor2 = Vault::new(db, GlobalActors::spawn_message_bus())
db, .await
GlobalActors::spawn_message_bus(), .unwrap();
Arc::new(DieselCustodyStore),
)
.await
.unwrap();
let err = actor2.decrypt(1, TEST_AAD.to_vec()).await.unwrap_err(); let err = actor2.decrypt(1, TEST_AAD.to_vec()).await.unwrap_err();
assert!(matches!(err, Error::Sealed)); assert!(matches!(err, Error::Sealed));
} }
@@ -124,13 +107,9 @@ async fn unseal_correct_password() {
.unwrap(); .unwrap();
drop(actor); drop(actor);
let mut actor = Vault::new( let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus())
db.clone(), .await
GlobalActors::spawn_message_bus(), .unwrap();
Arc::new(DieselCustodyStore),
)
.await
.unwrap();
let seal_key = KeyCell::from([0u8; 32]); let seal_key = KeyCell::from([0u8; 32]);
actor.try_unseal(seal_key).await.unwrap(); actor.try_unseal(seal_key).await.unwrap();
@@ -151,13 +130,9 @@ async fn unseal_wrong_then_correct_password() {
.unwrap(); .unwrap();
drop(actor); drop(actor);
let mut actor = Vault::new( let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus())
db.clone(), .await
GlobalActors::spawn_message_bus(), .unwrap();
Arc::new(DieselCustodyStore),
)
.await
.unwrap();
let bad_key = KeyCell::from([1u8; 32]); let bad_key = KeyCell::from([1u8; 32]);
let err = actor.try_unseal(bad_key).await.unwrap_err(); let err = actor.try_unseal(bad_key).await.unwrap_err();