tests(server): UserAgent bootstrap token auth flow test

server/Cargo.lock

@@ -103,6 +103,7 @@ dependencies = [
  "secrecy",
  "smlang",
  "statig",
+ "test-log",
  "thiserror",
  "tokio",
  "tokio-stream",
@@ -743,6 +744,7 @@ checksum = "053618a4c3d3bc24f188aa660ae75a46eeab74ef07fb415c61431e5e7cd4749b"
 dependencies = [
  "curve25519-dalek",
  "ed25519",
+ "rand_core 0.10.0",
  "sha2",
  "subtle",
  "zeroize",
@@ -1309,6 +1311,15 @@ version = "0.4.29"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
 
+[[package]]
+name = "matchers"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9"
+dependencies = [
+ "regex-automata",
+]
+
 [[package]]
 name = "matchit"
 version = "0.8.4"
@@ -2152,6 +2163,15 @@ dependencies = [
  "digest",
 ]
 
+[[package]]
+name = "sharded-slab"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6"
+dependencies = [
+ "lazy_static",
+]
+
 [[package]]
 name = "shlex"
 version = "1.3.0"
@@ -2373,6 +2393,27 @@ dependencies = [
  "windows-sys 0.60.2",
 ]
 
+[[package]]
+name = "test-log"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "37d53ac171c92a39e4769491c4b4dde7022c60042254b5fc044ae409d34a24d4"
+dependencies = [
+ "test-log-macros",
+ "tracing-subscriber",
+]
+
+[[package]]
+name = "test-log-macros"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "be35209fd0781c5401458ab66e4f98accf63553e8fae7425503e92fdd319783b"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.114",
+]
+
 [[package]]
 name = "textwrap"
 version = "0.16.2"
@@ -2403,6 +2444,15 @@ dependencies = [
  "syn 2.0.114",
 ]
 
+[[package]]
+name = "thread_local"
+version = "1.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185"
+dependencies = [
+ "cfg-if",
+]
+
 [[package]]
 name = "time"
 version = "0.3.47"
@@ -2677,6 +2727,21 @@ dependencies = [
  "once_cell",
 ]
 
+[[package]]
+name = "tracing-subscriber"
+version = "0.3.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2f30143827ddab0d256fd843b7a66d164e9f271cfa0dde49142c5ca0ca291f1e"
+dependencies = [
+ "matchers",
+ "once_cell",
+ "regex-automata",
+ "sharded-slab",
+ "thread_local",
+ "tracing",
+ "tracing-core",
+]
+
 [[package]]
 name = "try-lock"
 version = "0.2.5"

server/Cargo.toml

@@ -14,7 +14,7 @@ tonic = { version = "0.14.3", features = ["deflate", "gzip", "tls-connect-info",
 tracing = "0.1.44"
 tokio = { version = "1.49.0", features = ["full"] }
 ed25519 = "3.0.0-rc.4"
-ed25519-dalek = "3.0.0-pre.6"
+ed25519-dalek = { version = "3.0.0-pre.6", features = ["rand_core"] }
 chrono = { version = "0.4.43", features = ["serde"] }
 rand = "0.10.0"
 uuid = "1.20.0"
@@ -26,4 +26,4 @@ async-trait = "0.1.89"
 futures = "0.3.31"
 tokio-stream = { version = "0.1.18", features = ["full"] }
 kameo = "0.19.2"
-prost-types = { version = "0.14.3", features = ["chrono"] }
+prost-types = { version = "0.14.3", features = ["chrono"] }

server/crates/arbiter-proto/Cargo.toml

@@ -17,9 +17,8 @@ futures.workspace = true
 kameo.workspace = true
 hex = "0.4.3"
 
-
-
 [build-dependencies]
 prost-build = "0.14.3"
 serde_json = "1"
 tonic-prost-build = "0.14.3"
+

server/crates/arbiter-server/Cargo.toml

@@ -5,8 +5,19 @@ edition = "2024"
 repository = "https://git.markettakers.org/MarketTakers/arbiter"
 
 [dependencies]
-diesel = { version = "2.3.6", features = ["sqlite", "uuid", "time", "chrono", "serde_json"] }
-diesel-async = { version = "0.7.4", features = ["bb8", "migrations", "sqlite", "tokio"] }
+diesel = { version = "2.3.6", features = [
+    "sqlite",
+    "uuid",
+    "time",
+    "chrono",
+    "serde_json",
+] }
+diesel-async = { version = "0.7.4", features = [
+    "bb8",
+    "migrations",
+    "sqlite",
+    "tokio",
+] }
 ed25519.workspace = true
 ed25519-dalek.workspace = true
 arbiter-proto.path = "../arbiter-proto"
@@ -25,8 +36,17 @@ futures.workspace = true
 tokio-stream.workspace = true
 dashmap = "6.1.0"
 rand.workspace = true
-rcgen = { version = "0.14.7", features = ["aws_lc_rs", "pem", "x509-parser", "zeroize"], default-features = false }
-rkyv = { version = "0.8.15", features = ["aligned", "little_endian", "pointer_width_64"] }
+rcgen = { version = "0.14.7", features = [
+    "aws_lc_rs",
+    "pem",
+    "x509-parser",
+    "zeroize",
+], default-features = false }
+rkyv = { version = "0.8.15", features = [
+    "aligned",
+    "little_endian",
+    "pointer_width_64",
+] }
 restructed = "0.2.2"
 chrono.workspace = true
 bytes = "1.11.1"
@@ -34,4 +54,7 @@ memsafe = "0.4.0"
 chacha20poly1305 = { version = "0.10.1", features = ["std"] }
 zeroize = { version = "1.8.2", features = ["std", "simd"] }
 kameo.workspace = true
-prost-types.workspace = true
+prost-types.workspace = true
+
+[dev-dependencies]
+test-log = { version = "0.2", default-features = false, features = ["trace"] }

server/crates/arbiter-server/src/actors/user_agent.rs

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

server/crates/arbiter-server/src/context/bootstrap.rs

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

server/crates/arbiter-server/src/db.rs

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