MarketTakers/arbiter
6
0
Fork 0
Code Issues 14 Pull Requests 8 Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

13
server/crates/arbiter-server/tests/common/mod.rs
View File

@@ -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

17
server/crates/arbiter-server/tests/vault/concurrency.rs
View File

@@ -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

74
server/crates/arbiter-server/tests/vault/custody.rs
View File

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

63
server/crates/arbiter-server/tests/vault/lifecycle.rs
View File

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