From b3566c8af64b1d60efd51648af2706d9d8262b00 Mon Sep 17 00:00:00 2001 From: hdbg Date: Mon, 16 Feb 2026 21:44:11 +0100 Subject: [PATCH] refactor(server): separated global actors into their own handle --- server/crates/arbiter-server/src/actors.rs | 42 ++++++++++++- .../arbiter-server/src/actors/bootstrap.rs | 18 +++--- .../src/actors/user_agent/mod.rs | 31 +++++----- .../src/actors/user_agent/tests.rs | 59 +++++++------------ server/crates/arbiter-server/src/context.rs | 22 ++----- 5 files changed, 92 insertions(+), 80 deletions(-) diff --git a/server/crates/arbiter-server/src/actors.rs b/server/crates/arbiter-server/src/actors.rs index a987633..2a6567c 100644 --- a/server/crates/arbiter-server/src/actors.rs +++ b/server/crates/arbiter-server/src/actors.rs @@ -1,4 +1,40 @@ -pub mod user_agent; -pub mod client; +use kameo::actor::{ActorRef, Spawn}; +use miette::Diagnostic; +use thiserror::Error; + +use crate::{ + actors::{bootstrap::Bootstrapper, keyholder::KeyHolder}, + db, +}; + pub(crate) mod bootstrap; -pub(crate) mod keyholder; \ No newline at end of file +pub mod client; +pub(crate) mod keyholder; +pub mod user_agent; + +#[derive(Error, Debug, Diagnostic)] +pub enum SpawnError { + #[error("Failed to spawn Bootstrapper actor")] + #[diagnostic(code(SpawnError::Bootstrapper))] + Bootstrapper(#[from] bootstrap::Error), + + #[error("Failed to spawn KeyHolder actor")] + #[diagnostic(code(SpawnError::KeyHolder))] + KeyHolder(#[from] keyholder::Error), +} + +/// Long-lived actors that are shared across all connections and handle global state and operations +#[derive(Clone)] +pub struct GlobalActors { + pub key_holder: ActorRef, + pub bootstrapper: ActorRef, +} + +impl GlobalActors { + pub async fn spawn(db: db::DatabasePool) -> Result { + Ok(Self { + bootstrapper: Bootstrapper::spawn(Bootstrapper::new(&db).await?), + key_holder: KeyHolder::spawn(KeyHolder::new(db.clone()).await?), + }) + } +} diff --git a/server/crates/arbiter-server/src/actors/bootstrap.rs b/server/crates/arbiter-server/src/actors/bootstrap.rs index fd69e3d..063545c 100644 --- a/server/crates/arbiter-server/src/actors/bootstrap.rs +++ b/server/crates/arbiter-server/src/actors/bootstrap.rs @@ -28,7 +28,7 @@ pub async fn generate_token() -> Result { } #[derive(Error, Debug, Diagnostic)] -pub enum BootstrapError { +pub enum Error { #[error("Database error: {0}")] #[diagnostic(code(arbiter_server::bootstrap::database))] Database(#[from] db::PoolError), @@ -48,7 +48,7 @@ pub struct Bootstrapper { } impl Bootstrapper { - pub async fn new(db: &DatabasePool) -> Result { + pub async fn new(db: &DatabasePool) -> Result { let mut conn = db.get().await?; let row_count: i64 = schema::useragent_client::table @@ -69,11 +69,6 @@ impl Bootstrapper { Ok(Self { token }) } - - #[cfg(test)] - pub fn get_token(&self) -> Option { - self.token.clone() - } } #[messages] @@ -96,3 +91,12 @@ impl Bootstrapper { } } } + +#[cfg(test)] +#[messages] +impl Bootstrapper { + #[message] + pub fn get_token(&self) -> Option { + self.token.clone() + } +} diff --git a/server/crates/arbiter-server/src/actors/user_agent/mod.rs b/server/crates/arbiter-server/src/actors/user_agent/mod.rs index 098333a..82f6722 100644 --- a/server/crates/arbiter-server/src/actors/user_agent/mod.rs +++ b/server/crates/arbiter-server/src/actors/user_agent/mod.rs @@ -10,9 +10,9 @@ use arbiter_proto::proto::{ }; use chacha20poly1305::{AeadInPlace, XChaCha20Poly1305, XNonce, aead::KeyInit}; use diesel::{ExpressionMethods as _, OptionalExtension as _, QueryDsl, dsl::update}; -use diesel_async::{RunQueryDsl}; +use diesel_async::RunQueryDsl; use ed25519_dalek::VerifyingKey; -use kameo::{Actor, actor::ActorRef, error::SendError, messages}; +use kameo::{Actor, error::SendError, messages}; use memsafe::MemSafe; use tokio::sync::mpsc::Sender; use tonic::Status; @@ -22,11 +22,12 @@ use x25519_dalek::{EphemeralSecret, PublicKey}; use crate::{ ServerContext, actors::{ - bootstrap::{Bootstrapper, ConsumeToken}, - keyholder::{self, KeyHolder, TryUnseal}, + GlobalActors, + bootstrap::ConsumeToken, + keyholder::{self, TryUnseal}, user_agent::state::{ - ChallengeContext, DummyContext, UnsealContext, UserAgentEvents, - UserAgentStateMachine, UserAgentStates, + ChallengeContext, DummyContext, UnsealContext, UserAgentEvents, UserAgentStateMachine, + UserAgentStates, }, }, db::{self, schema}, @@ -43,8 +44,7 @@ pub(crate) use transport::handle_user_agent; #[derive(Actor)] pub struct UserAgentActor { db: db::DatabasePool, - bootstapper: ActorRef, - keyholder: ActorRef, + actors: GlobalActors, state: UserAgentStateMachine, // will be used in future _tx: Sender>, @@ -57,8 +57,7 @@ impl UserAgentActor { ) -> Self { Self { db: context.db.clone(), - bootstapper: context.bootstrapper.clone(), - keyholder: context.keyholder.clone(), + actors: context.actors.clone(), state: UserAgentStateMachine::new(DummyContext), _tx: tx, } @@ -67,14 +66,12 @@ impl UserAgentActor { #[cfg(test)] pub(crate) fn new_manual( db: db::DatabasePool, - bootstapper: ActorRef, - keyholder: ActorRef, + actors: GlobalActors, tx: Sender>, ) -> Self { Self { db, - bootstapper, - keyholder, + actors, state: UserAgentStateMachine::new(DummyContext), _tx: tx, } @@ -94,7 +91,8 @@ impl UserAgentActor { token: String, ) -> Result { let token_ok: bool = self - .bootstapper + .actors + .bootstrapper .ask(ConsumeToken { token }) .await .map_err(|e| { @@ -288,7 +286,8 @@ impl UserAgentActor { match decryption_result { Ok(_) => { match self - .keyholder + .actors + .key_holder .ask(TryUnseal { seal_key_raw: seal_key_buffer, }) diff --git a/server/crates/arbiter-server/src/actors/user_agent/tests.rs b/server/crates/arbiter-server/src/actors/user_agent/tests.rs index 2dbcdd3..20f95e4 100644 --- a/server/crates/arbiter-server/src/actors/user_agent/tests.rs +++ b/server/crates/arbiter-server/src/actors/user_agent/tests.rs @@ -13,8 +13,9 @@ use x25519_dalek::{EphemeralSecret, PublicKey}; use crate::{ actors::{ - bootstrap::Bootstrapper, - keyholder::KeyHolder, + GlobalActors, + bootstrap::GetToken, + keyholder::{Bootstrap, Seal}, user_agent::{ HandleAuthChallengeRequest, HandleAuthChallengeSolution, HandleUnsealEncryptedKey, HandleUnsealRequest, @@ -47,25 +48,23 @@ async fn setup_authenticated_user_agent( let db = db::create_test_pool().await; seed_settings(&db).await; - let mut keyholder = KeyHolder::new(db.clone()).await.unwrap(); - keyholder - .bootstrap(MemSafe::new(seal_key.to_vec()).unwrap()) + let actors = GlobalActors::spawn(db.clone()).await.unwrap(); + actors + .key_holder + .ask(Bootstrap { + seal_key_raw: MemSafe::new(seal_key.to_vec()).unwrap(), + }) .await .unwrap(); - keyholder.seal().unwrap(); - let keyholder_ref = KeyHolder::spawn(keyholder); - - let bootstrapper = Bootstrapper::new(&db).await.unwrap(); - let token = bootstrapper.get_token().unwrap(); - let bootstrapper_ref = Bootstrapper::spawn(bootstrapper); + actors.key_holder.ask(Seal).await.unwrap(); let user_agent = UserAgentActor::new_manual( db.clone(), - bootstrapper_ref, - keyholder_ref, + actors.clone(), tokio::sync::mpsc::channel(1).0, ); let user_agent_ref = UserAgentActor::spawn(user_agent); + let token = actors.bootstrapper.ask(GetToken).await.unwrap().unwrap(); let auth_key = ed25519_dalek::SigningKey::generate(&mut rand::rng()); user_agent_ref @@ -128,17 +127,11 @@ async fn client_dh_encrypt( pub async fn test_bootstrap_token_auth() { let db = db::create_test_pool().await; seed_settings(&db).await; - // explicitly not installing any user_agent pubkeys - let bootstrapper = Bootstrapper::new(&db).await.unwrap(); // this will create bootstrap token - let keyholder = KeyHolder::new(db.clone()).await.unwrap(); - let token = bootstrapper.get_token().unwrap(); - - let bootstrapper_ref = Bootstrapper::spawn(bootstrapper); - let keyholder_ref = KeyHolder::spawn(keyholder); + let actors = GlobalActors::spawn(db.clone()).await.unwrap(); + let token = actors.bootstrapper.ask(GetToken).await.unwrap().unwrap(); let user_agent = UserAgentActor::new_manual( db.clone(), - bootstrapper_ref, - keyholder_ref, + actors.clone(), tokio::sync::mpsc::channel(1).0, // dummy channel, we won't actually send responses in this test ); let user_agent_ref = UserAgentActor::spawn(user_agent); @@ -186,17 +179,11 @@ pub async fn test_bootstrap_token_auth() { pub async fn test_bootstrap_invalid_token_auth() { let db = db::create_test_pool().await; seed_settings(&db).await; - // explicitly not installing any user_agent pubkeys - let bootstrapper = Bootstrapper::new(&db).await.unwrap(); // this will create bootstrap token - let keyholder = KeyHolder::new(db.clone()).await.unwrap(); - - let bootstrapper_ref = Bootstrapper::spawn(bootstrapper); - let keyholder_ref = KeyHolder::spawn(keyholder); + let actors = GlobalActors::spawn(db.clone()).await.unwrap(); let user_agent = UserAgentActor::new_manual( db.clone(), - bootstrapper_ref, - keyholder_ref, + actors, tokio::sync::mpsc::channel(1).0, // dummy channel, we won't actually send responses in this test ); let user_agent_ref = UserAgentActor::spawn(user_agent); @@ -240,12 +227,10 @@ pub async fn test_challenge_auth() { let db = db::create_test_pool().await; seed_settings(&db).await; - let bootstrapper_ref = Bootstrapper::spawn(Bootstrapper::new(&db).await.unwrap()); - let keyholder_ref = KeyHolder::spawn(KeyHolder::new(db.clone()).await.unwrap()); + let actors = GlobalActors::spawn(db.clone()).await.unwrap(); let user_agent = UserAgentActor::new_manual( db.clone(), - bootstrapper_ref, - keyholder_ref, + actors, tokio::sync::mpsc::channel(1).0, // dummy channel, we won't actually send responses in this test ); let user_agent_ref = UserAgentActor::spawn(user_agent); @@ -394,13 +379,11 @@ pub async fn test_unseal_start_without_auth_fails() { let db = db::create_test_pool().await; seed_settings(&db).await; - let keyholder_ref = KeyHolder::spawn( KeyHolder::new(db.clone()).await.unwrap()); - let bootstrapper_ref = Bootstrapper::spawn(Bootstrapper::new(&db).await.unwrap()); + let actors = GlobalActors::spawn(db.clone()).await.unwrap(); let user_agent = UserAgentActor::new_manual( db.clone(), - bootstrapper_ref, - keyholder_ref, + actors, tokio::sync::mpsc::channel(1).0, ); let user_agent_ref = UserAgentActor::spawn(user_agent); diff --git a/server/crates/arbiter-server/src/context.rs b/server/crates/arbiter-server/src/context.rs index c8c40b7..01cd2ed 100644 --- a/server/crates/arbiter-server/src/context.rs +++ b/server/crates/arbiter-server/src/context.rs @@ -2,15 +2,11 @@ use std::sync::Arc; use diesel::OptionalExtension as _; use diesel_async::RunQueryDsl as _; -use kameo::actor::{ActorRef, Spawn}; use miette::Diagnostic; use thiserror::Error; use crate::{ - actors::{ - bootstrap::{self, Bootstrapper}, - keyholder::KeyHolder, - }, + actors::GlobalActors, context::tls::{TlsDataRaw, TlsManager}, db::{self, models::ArbiterSetting, schema::arbiter_settings}, }; @@ -35,13 +31,9 @@ pub enum InitError { #[diagnostic(code(arbiter_server::init::tls_init))] Tls(#[from] tls::TlsInitError), - #[error("Bootstrap token generation failed: {0}")] - #[diagnostic(code(arbiter_server::init::bootstrap_token))] - BootstrapToken(#[from] bootstrap::BootstrapError), - - #[error("KeyHolder initialization failed: {0}")] - #[diagnostic(code(arbiter_server::init::keyholder_init))] - KeyHolder(#[from] crate::actors::keyholder::Error), + #[error("Actor spawn failed: {0}")] + #[diagnostic(code(arbiter_server::init::actor_spawn))] + ActorSpawn(#[from] crate::actors::SpawnError), #[error("I/O Error: {0}")] #[diagnostic(code(arbiter_server::init::io))] @@ -51,8 +43,7 @@ pub enum InitError { pub struct _ServerContextInner { pub db: db::DatabasePool, pub tls: TlsManager, - pub bootstrapper: ActorRef, - pub keyholder: ActorRef, + pub actors: GlobalActors, } #[derive(Clone)] pub struct ServerContext(Arc<_ServerContextInner>); @@ -111,8 +102,7 @@ impl ServerContext { drop(conn); Ok(Self(Arc::new(_ServerContextInner { - bootstrapper: Bootstrapper::spawn(Bootstrapper::new(&db).await?), - keyholder: KeyHolder::spawn(KeyHolder::new(db.clone()).await?), + actors: GlobalActors::spawn(db.clone()).await?, db, tls, })))