refactor(server): separated global actors into their own handle

This commit is contained in:
hdbg
2026-02-16 21:44:11 +01:00
parent bdb9f01757
commit b3566c8af6
5 changed files with 92 additions and 80 deletions

View File

@@ -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 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<KeyHolder>,
pub bootstrapper: ActorRef<Bootstrapper>,
}
impl GlobalActors {
pub async fn spawn(db: db::DatabasePool) -> Result<Self, SpawnError> {
Ok(Self {
bootstrapper: Bootstrapper::spawn(Bootstrapper::new(&db).await?),
key_holder: KeyHolder::spawn(KeyHolder::new(db.clone()).await?),
})
}
}

View File

@@ -28,7 +28,7 @@ pub async fn generate_token() -> Result<String, std::io::Error> {
}
#[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<Self, BootstrapError> {
pub async fn new(db: &DatabasePool) -> Result<Self, Error> {
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<String> {
self.token.clone()
}
}
#[messages]
@@ -96,3 +91,12 @@ impl Bootstrapper {
}
}
}
#[cfg(test)]
#[messages]
impl Bootstrapper {
#[message]
pub fn get_token(&self) -> Option<String> {
self.token.clone()
}
}

View File

@@ -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<Bootstrapper>,
keyholder: ActorRef<KeyHolder>,
actors: GlobalActors,
state: UserAgentStateMachine<DummyContext>,
// will be used in future
_tx: Sender<Result<UserAgentResponse, Status>>,
@@ -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<Bootstrapper>,
keyholder: ActorRef<KeyHolder>,
actors: GlobalActors,
tx: Sender<Result<UserAgentResponse, Status>>,
) -> Self {
Self {
db,
bootstapper,
keyholder,
actors,
state: UserAgentStateMachine::new(DummyContext),
_tx: tx,
}
@@ -94,7 +91,8 @@ impl UserAgentActor {
token: String,
) -> Result<UserAgentResponse, Status> {
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,
})

View File

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

View File

@@ -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<Bootstrapper>,
pub keyholder: ActorRef<KeyHolder>,
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,
})))