test(user-agent): add test helpers and update actor integration tests

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<dyn Sender<OutOfBand>>) -> 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<OutOfBand> 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");

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<UnsealStartResponse, Error> {
@@ -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<u8>,
         ciphertext: Vec<u8>,

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 db_for_task = db.clone();
-    let task = tokio::spawn(connect_user_agent(props));
+    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<u8> = 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 db_for_task = db.clone();
-    let task = tokio::spawn(connect_user_agent(props));
+    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
+    assert!(matches!(
-    task.await.unwrap();
+        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 db_for_task = db.clone();
-    let task = tokio::spawn(connect_user_agent(props));
+    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
+    let response = test_transport
-    task.await.unwrap();
+        .recv()
+        .await
+        .expect("should receive auth result");
+    match response {
+        Ok(auth::Outbound::AuthSuccess) => {}
+        other => panic!("Expected AuthSuccess, got {other:?}"),
+    }
+
+    task.await.unwrap().unwrap();
 }

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<UserAgentSession>) {
     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<UserAgentSession>,
+    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 {
+    let server_pubkey = response.server_pubkey;
-        OutOfBand::UnsealStartResponse { server_pubkey } => server_pubkey,
-        other => panic!("Expected UnsealStartResponse, got {other:?}"),
-    };
 
     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
+    let response = user_agent.ask(encrypted_key).await;
-        .process_transport_inbound(encrypted_key)
+    assert!(matches!(response, Ok(())));
-        .await
-        .unwrap();
-
-    assert!(matches!(response, OutOfBand::UnsealResult(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 encrypted_key = client_dh_encrypt(&user_agent, b"wrong-key").await;
-
-    let response = user_agent
-        .process_transport_inbound(encrypted_key)
-        .await
-        .unwrap();
 
+    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
+        .await;
-        .unwrap();
 
     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 encrypted_key = client_dh_encrypt(&user_agent, b"wrong-key").await;
-
-        let response = user_agent
-            .process_transport_inbound(encrypted_key)
-            .await
-            .unwrap();
 
+        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
+        let response = user_agent.ask(encrypted_key).await;
-            .process_transport_inbound(encrypted_key)
+        assert!(matches!(response, Ok(())));
-            .await
-            .unwrap();
-
-        assert!(matches!(response, OutOfBand::UnsealResult(Ok(()))));
     }
 }