From 60ce1cc11021406ac2f0e6e881bc2e749c46e112 Mon Sep 17 00:00:00 2001 From: hdbg Date: Wed, 18 Mar 2026 23:37:40 +0100 Subject: [PATCH] test(user-agent): add test helpers and update actor integration tests --- .../src/actors/user_agent/session.rs | 19 +++- .../actors/user_agent/session/connection.rs | 4 +- .../arbiter-server/tests/user_agent/auth.rs | 64 ++++++++++---- .../arbiter-server/tests/user_agent/unseal.rs | 88 +++++++++---------- 4 files changed, 107 insertions(+), 68 deletions(-) diff --git a/server/crates/arbiter-server/src/actors/user_agent/session.rs b/server/crates/arbiter-server/src/actors/user_agent/session.rs index 382165a..398b09f 100644 --- a/server/crates/arbiter-server/src/actors/user_agent/session.rs +++ b/server/crates/arbiter-server/src/actors/user_agent/session.rs @@ -1,6 +1,7 @@ use std::{borrow::Cow, convert::Infallible}; use arbiter_proto::transport::Sender; +use async_trait::async_trait; use ed25519_dalek::VerifyingKey; use kameo::{Actor, messages, prelude::Context}; use thiserror::Error; @@ -42,8 +43,8 @@ mod connection; pub(crate) use connection::{ BootstrapError, HandleBootstrapEncryptedKey, HandleEvmWalletCreate, HandleEvmWalletList, HandleGrantCreate, HandleGrantDelete, HandleGrantList, HandleQueryVaultState, - HandleUnsealEncryptedKey, HandleUnsealRequest, UnsealError, }; +pub use connection::{HandleUnsealEncryptedKey, HandleUnsealRequest, UnsealError}; impl UserAgentSession { pub(crate) fn new(props: UserAgentConnection, sender: Box>) -> Self { @@ -54,6 +55,22 @@ impl UserAgentSession { } } + pub fn new_test(db: crate::db::DatabasePool, actors: crate::actors::GlobalActors) -> Self { + struct DummySender; + + #[async_trait] + impl Sender for DummySender { + async fn send( + &mut self, + _item: OutOfBand, + ) -> Result<(), arbiter_proto::transport::Error> { + Ok(()) + } + } + + Self::new(UserAgentConnection::new(db, actors), Box::new(DummySender)) + } + fn transition(&mut self, event: UserAgentEvents) -> Result<(), Error> { self.state.process_event(event).map_err(|e| { error!(?e, "State transition failed"); diff --git a/server/crates/arbiter-server/src/actors/user_agent/session/connection.rs b/server/crates/arbiter-server/src/actors/user_agent/session/connection.rs index ed9a107..364dbf4 100644 --- a/server/crates/arbiter-server/src/actors/user_agent/session/connection.rs +++ b/server/crates/arbiter-server/src/actors/user_agent/session/connection.rs @@ -109,7 +109,7 @@ pub enum BootstrapError { #[messages] impl UserAgentSession { #[message] - pub(crate) async fn handle_unseal_request( + pub async fn handle_unseal_request( &mut self, client_pubkey: x25519_dalek::PublicKey, ) -> Result { @@ -127,7 +127,7 @@ impl UserAgentSession { } #[message] - pub(crate) async fn handle_unseal_encrypted_key( + pub async fn handle_unseal_encrypted_key( &mut self, nonce: Vec, ciphertext: Vec, diff --git a/server/crates/arbiter-server/tests/user_agent/auth.rs b/server/crates/arbiter-server/tests/user_agent/auth.rs index bfe308a..285ddcf 100644 --- a/server/crates/arbiter-server/tests/user_agent/auth.rs +++ b/server/crates/arbiter-server/tests/user_agent/auth.rs @@ -1,9 +1,9 @@ -use arbiter_proto::transport::Bi; +use arbiter_proto::transport::{Receiver, Sender}; use arbiter_server::{ actors::{ GlobalActors, bootstrap::GetToken, - user_agent::{AuthPublicKey, OutOfBand, Request, UserAgentConnection, connect_user_agent}, + user_agent::{AuthPublicKey, UserAgentConnection, auth}, }, db::{self, schema}, }; @@ -21,19 +21,31 @@ pub async fn test_bootstrap_token_auth() { let token = actors.bootstrapper.ask(GetToken).await.unwrap().unwrap(); let (server_transport, mut test_transport) = ChannelTransport::new(); - let props = UserAgentConnection::new(db.clone(), actors, Box::new(server_transport)); - let task = tokio::spawn(connect_user_agent(props)); + let db_for_task = db.clone(); + let task = tokio::spawn(async move { + let mut props = UserAgentConnection::new(db_for_task, actors); + auth::authenticate(&mut props, server_transport).await + }); let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng()); test_transport - .send(Request::AuthChallengeRequest { + .send(auth::Inbound::AuthChallengeRequest { pubkey: AuthPublicKey::Ed25519(new_key.verifying_key()), bootstrap_token: Some(token), }) .await .unwrap(); - task.await.unwrap(); + let response = test_transport + .recv() + .await + .expect("should receive auth result"); + match response { + Ok(auth::Outbound::AuthSuccess) => {} + other => panic!("Expected AuthSuccess, got {other:?}"), + } + + task.await.unwrap().unwrap(); let mut conn = db.get().await.unwrap(); let stored_pubkey: Vec = schema::useragent_client::table @@ -51,20 +63,25 @@ pub async fn test_bootstrap_invalid_token_auth() { let actors = GlobalActors::spawn(db.clone()).await.unwrap(); let (server_transport, mut test_transport) = ChannelTransport::new(); - let props = UserAgentConnection::new(db.clone(), actors, Box::new(server_transport)); - let task = tokio::spawn(connect_user_agent(props)); + let db_for_task = db.clone(); + let task = tokio::spawn(async move { + let mut props = UserAgentConnection::new(db_for_task, actors); + auth::authenticate(&mut props, server_transport).await + }); let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng()); test_transport - .send(Request::AuthChallengeRequest { + .send(auth::Inbound::AuthChallengeRequest { pubkey: AuthPublicKey::Ed25519(new_key.verifying_key()), bootstrap_token: Some("invalid_token".to_string()), }) .await .unwrap(); - // Auth fails, connect_user_agent returns, transport drops - task.await.unwrap(); + assert!(matches!( + task.await.unwrap(), + Err(auth::Error::InvalidBootstrapToken) + )); // Verify no key was registered let mut conn = db.get().await.unwrap(); @@ -99,12 +116,15 @@ pub async fn test_challenge_auth() { } let (server_transport, mut test_transport) = ChannelTransport::new(); - let props = UserAgentConnection::new(db.clone(), actors, Box::new(server_transport)); - let task = tokio::spawn(connect_user_agent(props)); + let db_for_task = db.clone(); + let task = tokio::spawn(async move { + let mut props = UserAgentConnection::new(db_for_task, actors); + auth::authenticate(&mut props, server_transport).await + }); // Send challenge request test_transport - .send(Request::AuthChallengeRequest { + .send(auth::Inbound::AuthChallengeRequest { pubkey: AuthPublicKey::Ed25519(new_key.verifying_key()), bootstrap_token: None, }) @@ -118,7 +138,7 @@ pub async fn test_challenge_auth() { .expect("should receive challenge"); let challenge = match response { Ok(resp) => match resp { - OutOfBand::AuthChallenge { nonce } => nonce, + auth::Outbound::AuthChallenge { nonce } => nonce, other => panic!("Expected AuthChallenge, got {other:?}"), }, Err(err) => panic!("Expected Ok response, got Err({err:?})"), @@ -128,12 +148,20 @@ pub async fn test_challenge_auth() { let signature = new_key.sign(&formatted_challenge); test_transport - .send(Request::AuthChallengeSolution { + .send(auth::Inbound::AuthChallengeSolution { signature: signature.to_bytes().to_vec(), }) .await .unwrap(); - // Auth completes, session spawned - task.await.unwrap(); + let response = test_transport + .recv() + .await + .expect("should receive auth result"); + match response { + Ok(auth::Outbound::AuthSuccess) => {} + other => panic!("Expected AuthSuccess, got {other:?}"), + } + + task.await.unwrap().unwrap(); } diff --git a/server/crates/arbiter-server/tests/user_agent/unseal.rs b/server/crates/arbiter-server/tests/user_agent/unseal.rs index 0b2eea6..cd59b01 100644 --- a/server/crates/arbiter-server/tests/user_agent/unseal.rs +++ b/server/crates/arbiter-server/tests/user_agent/unseal.rs @@ -2,15 +2,20 @@ use arbiter_server::{ actors::{ GlobalActors, keyholder::{Bootstrap, Seal}, - user_agent::{OutOfBand, Request, UnsealError, session::UserAgentSession}, + user_agent::session::{ + HandleUnsealEncryptedKey, HandleUnsealRequest, UnsealError, UserAgentSession, + }, }, db, safe_cell::{SafeCell, SafeCellHandle as _}, }; use chacha20poly1305::{AeadInPlace, XChaCha20Poly1305, XNonce, aead::KeyInit}; +use kameo::actor::Spawn as _; use x25519_dalek::{EphemeralSecret, PublicKey}; -async fn setup_sealed_user_agent(seal_key: &[u8]) -> (db::DatabasePool, UserAgentSession) { +async fn setup_sealed_user_agent( + seal_key: &[u8], +) -> (db::DatabasePool, kameo::actor::ActorRef) { let db = db::create_test_pool().await; let actors = GlobalActors::spawn(db.clone()).await.unwrap(); @@ -23,26 +28,26 @@ async fn setup_sealed_user_agent(seal_key: &[u8]) -> (db::DatabasePool, UserAgen .unwrap(); actors.key_holder.ask(Seal).await.unwrap(); - let session = UserAgentSession::new_test(db.clone(), actors); + let session = UserAgentSession::spawn(UserAgentSession::new_test(db.clone(), actors)); (db, session) } -async fn client_dh_encrypt(user_agent: &mut UserAgentSession, key_to_send: &[u8]) -> Request { +async fn client_dh_encrypt( + user_agent: &kameo::actor::ActorRef, + key_to_send: &[u8], +) -> HandleUnsealEncryptedKey { let client_secret = EphemeralSecret::random(); let client_public = PublicKey::from(&client_secret); let response = user_agent - .process_transport_inbound(Request::UnsealStart { + .ask(HandleUnsealRequest { client_pubkey: client_public, }) .await .unwrap(); - let server_pubkey = match response { - OutOfBand::UnsealStartResponse { server_pubkey } => server_pubkey, - other => panic!("Expected UnsealStartResponse, got {other:?}"), - }; + let server_pubkey = response.server_pubkey; let shared_secret = client_secret.diffie_hellman(&server_pubkey); let cipher = XChaCha20Poly1305::new(shared_secret.as_bytes().into()); @@ -53,7 +58,7 @@ async fn client_dh_encrypt(user_agent: &mut UserAgentSession, key_to_send: &[u8] .encrypt_in_place(&nonce, associated_data, &mut ciphertext) .unwrap(); - Request::UnsealEncryptedKey { + HandleUnsealEncryptedKey { nonce: nonce.to_vec(), ciphertext, associated_data: associated_data.to_vec(), @@ -64,63 +69,58 @@ async fn client_dh_encrypt(user_agent: &mut UserAgentSession, key_to_send: &[u8] #[test_log::test] pub async fn test_unseal_success() { let seal_key = b"test-seal-key"; - let (_db, mut user_agent) = setup_sealed_user_agent(seal_key).await; + let (_db, user_agent) = setup_sealed_user_agent(seal_key).await; - let encrypted_key = client_dh_encrypt(&mut user_agent, seal_key).await; + let encrypted_key = client_dh_encrypt(&user_agent, seal_key).await; - let response = user_agent - .process_transport_inbound(encrypted_key) - .await - .unwrap(); - - assert!(matches!(response, OutOfBand::UnsealResult(Ok(())))); + let response = user_agent.ask(encrypted_key).await; + assert!(matches!(response, Ok(()))); } #[tokio::test] #[test_log::test] pub async fn test_unseal_wrong_seal_key() { - let (_db, mut user_agent) = setup_sealed_user_agent(b"correct-key").await; + let (_db, user_agent) = setup_sealed_user_agent(b"correct-key").await; - let encrypted_key = client_dh_encrypt(&mut user_agent, b"wrong-key").await; - - let response = user_agent - .process_transport_inbound(encrypted_key) - .await - .unwrap(); + let encrypted_key = client_dh_encrypt(&user_agent, b"wrong-key").await; + let response = user_agent.ask(encrypted_key).await; assert!(matches!( response, - OutOfBand::UnsealResult(Err(UnsealError::InvalidKey)) + Err(kameo::error::SendError::HandlerError( + UnsealError::InvalidKey + )) )); } #[tokio::test] #[test_log::test] pub async fn test_unseal_corrupted_ciphertext() { - let (_db, mut user_agent) = setup_sealed_user_agent(b"test-key").await; + let (_db, user_agent) = setup_sealed_user_agent(b"test-key").await; let client_secret = EphemeralSecret::random(); let client_public = PublicKey::from(&client_secret); user_agent - .process_transport_inbound(Request::UnsealStart { + .ask(HandleUnsealRequest { client_pubkey: client_public, }) .await .unwrap(); let response = user_agent - .process_transport_inbound(Request::UnsealEncryptedKey { + .ask(HandleUnsealEncryptedKey { nonce: vec![0u8; 24], ciphertext: vec![0u8; 32], associated_data: vec![], }) - .await - .unwrap(); + .await; assert!(matches!( response, - OutOfBand::UnsealResult(Err(UnsealError::InvalidKey)) + Err(kameo::error::SendError::HandlerError( + UnsealError::InvalidKey + )) )); } @@ -128,30 +128,24 @@ pub async fn test_unseal_corrupted_ciphertext() { #[test_log::test] pub async fn test_unseal_retry_after_invalid_key() { let seal_key = b"real-seal-key"; - let (_db, mut user_agent) = setup_sealed_user_agent(seal_key).await; + let (_db, user_agent) = setup_sealed_user_agent(seal_key).await; { - let encrypted_key = client_dh_encrypt(&mut user_agent, b"wrong-key").await; - - let response = user_agent - .process_transport_inbound(encrypted_key) - .await - .unwrap(); + let encrypted_key = client_dh_encrypt(&user_agent, b"wrong-key").await; + let response = user_agent.ask(encrypted_key).await; assert!(matches!( response, - OutOfBand::UnsealResult(Err(UnsealError::InvalidKey)) + Err(kameo::error::SendError::HandlerError( + UnsealError::InvalidKey + )) )); } { - let encrypted_key = client_dh_encrypt(&mut user_agent, seal_key).await; + let encrypted_key = client_dh_encrypt(&user_agent, seal_key).await; - let response = user_agent - .process_transport_inbound(encrypted_key) - .await - .unwrap(); - - assert!(matches!(response, OutOfBand::UnsealResult(Ok(())))); + let response = user_agent.ask(encrypted_key).await; + assert!(matches!(response, Ok(()))); } }