tests(server): UserAgent bootstrap token auth flow test

This commit is contained in:
hdbg
2026-02-14 17:59:06 +01:00
parent 345a967c13
commit 69dd8f57ca
9 changed files with 181 additions and 20 deletions

View File

@@ -63,7 +63,8 @@ smlang::statemachine!(
}
);
impl UserAgentStateMachineContext for ServerContext {
pub struct DummyContext;
impl UserAgentStateMachineContext for DummyContext {
#[allow(missing_docs)]
#[allow(clippy::unused_unit)]
fn move_challenge(
@@ -88,7 +89,7 @@ impl UserAgentStateMachineContext for ServerContext {
pub struct UserAgentActor {
db: db::DatabasePool,
bootstapper: ActorRef<BootstrapActor>,
state: UserAgentStateMachine<ServerContext>,
state: UserAgentStateMachine<DummyContext>,
tx: Sender<Result<UserAgentResponse, Status>>,
}
@@ -100,7 +101,7 @@ impl UserAgentActor {
Self {
db: context.db.clone(),
bootstapper: context.bootstrapper.clone(),
state: UserAgentStateMachine::new(context),
state: UserAgentStateMachine::new(DummyContext),
tx,
}
}
@@ -108,13 +109,12 @@ impl UserAgentActor {
pub(crate) fn new_manual(
db: db::DatabasePool,
bootstapper: ActorRef<BootstrapActor>,
state: UserAgentStateMachine<ServerContext>,
tx: Sender<Result<UserAgentResponse, Status>>,
) -> Self {
Self {
db,
bootstapper,
state,
state: UserAgentStateMachine::new(DummyContext),
tx,
}
}
@@ -305,5 +305,67 @@ impl UserAgentActor {
}
}
#[cfg(test)]
mod tests {
use arbiter_proto::proto::{
UserAgentResponse, auth::{AuthChallengeRequest, AuthOk},
user_agent_response::Payload as UserAgentResponsePayload,
};
use kameo::actor::Spawn;
use crate::{
actors::user_agent::HandleAuthChallengeRequest, context::bootstrap::BootstrapActor, db,
};
use super::UserAgentActor;
#[tokio::test]
#[test_log::test]
pub async fn test_bootstrap_token_auth() {
let db = db::create_pool(Some("sqlite://:memory:"))
.await
.expect("Failed to create database pool");
// explicitly not installing any user_agent pubkeys
let bootstrapper = BootstrapActor::new(&db).await.unwrap(); // this will create bootstrap token
let token = bootstrapper.get_token().unwrap();
let bootstrapper_ref = BootstrapActor::spawn(bootstrapper);
let user_agent = UserAgentActor::new_manual(
db.clone(),
bootstrapper_ref,
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);
// simulate client sending auth request with bootstrap token
let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
let result = user_agent_ref
.ask(HandleAuthChallengeRequest {
req: AuthChallengeRequest {
pubkey: pubkey_bytes,
bootstrap_token: Some(token),
},
})
.await
.expect("Shouldn't fail to send message");
// auth succeeded
assert_eq!(
result,
UserAgentResponse {
payload: Some(UserAgentResponsePayload::AuthMessage(
arbiter_proto::proto::auth::ServerMessage {
payload: Some(arbiter_proto::proto::auth::server_message::Payload::AuthOk(
AuthOk {},
)),
},
)),
}
);
}
}
mod transport;
pub(crate) use transport::handle_user_agent;

View File

@@ -81,6 +81,11 @@ impl BootstrapActor {
Ok(Self { token })
}
#[cfg(test)]
pub fn get_token(&self) -> Option<String> {
self.token.clone()
}
}
#[messages]

View File

@@ -12,6 +12,7 @@ use diesel_async::{
use diesel_migrations::{EmbeddedMigrations, MigrationHarness, embed_migrations};
use miette::Diagnostic;
use thiserror::Error;
use tracing::info;
pub mod models;
pub mod schema;
@@ -23,7 +24,7 @@ pub type PoolError = diesel_async::pooled_connection::bb8::RunError;
static DB_FILE: &'static str = "arbiter.sqlite";
const MIGRATIONS: EmbeddedMigrations = embed_migrations!();
const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations");
#[derive(Error, Diagnostic, Debug)]
pub enum DatabaseSetupError {
@@ -48,6 +49,7 @@ pub enum DatabaseSetupError {
Pool(#[from] PoolInitError),
}
#[tracing::instrument(level = "info")]
fn database_path() -> Result<std::path::PathBuf, DatabaseSetupError> {
let arbiter_home = arbiter_proto::home_path().map_err(DatabaseSetupError::HomeDir)?;
@@ -56,6 +58,7 @@ fn database_path() -> Result<std::path::PathBuf, DatabaseSetupError> {
Ok(db_path)
}
#[tracing::instrument(level = "info", skip(conn))]
fn db_config(conn: &mut SqliteConnection) -> Result<(), diesel::result::Error> {
// fsync only in critical moments
conn.batch_execute("PRAGMA synchronous = NORMAL;")?;
@@ -77,6 +80,7 @@ fn db_config(conn: &mut SqliteConnection) -> Result<(), diesel::result::Error> {
Ok(())
}
#[tracing::instrument(level = "info", skip(url))]
fn initialize_database(url: &str) -> Result<(), DatabaseSetupError> {
let mut conn = SqliteConnection::establish(url).map_err(DatabaseSetupError::Connection)?;
@@ -85,16 +89,19 @@ fn initialize_database(url: &str) -> Result<(), DatabaseSetupError> {
conn.run_pending_migrations(MIGRATIONS)
.map_err(DatabaseSetupError::Migration)?;
info!(%url, "Database initialized successfully");
Ok(())
}
pub async fn create_pool() -> Result<DatabasePool, DatabaseSetupError> {
let database_url = format!(
#[tracing::instrument(level = "info")]
pub async fn create_pool(url: Option<&str>) -> Result<DatabasePool, DatabaseSetupError> {
let database_url = url.map(String::from).unwrap_or(format!(
"{}?mode=rwc",
database_path()?
(database_path()?
.to_str()
.expect("database path is not valid UTF-8")
);
.expect("database path is not valid UTF-8"))
));
initialize_database(&database_url)?;