refactor(server::{user_agent, client}): move auth part to separate function to not to pollute actor session with one-time concerns

server/crates/arbiter-server/tests/client.rs

@@ -1,2 +1,4 @@
+mod common;
+
 #[path = "client/auth.rs"]
 mod auth;

server/crates/arbiter-server/tests/client/auth.rs

@@ -1,44 +1,44 @@
+use arbiter_proto::transport::Bi;
 use arbiter_proto::proto::client::{
-    AuthChallengeRequest, AuthChallengeSolution, AuthOk, ClientRequest, ClientResponse,
+    AuthChallengeRequest, AuthChallengeSolution, ClientRequest,
     client_request::Payload as ClientRequestPayload,
     client_response::Payload as ClientResponsePayload,
 };
 use arbiter_server::{
-    actors::client::{ClientActor, ClientError},
+    actors::client::{ConnectionProps, connect_client},
     db::{self, schema},
 };
 use diesel::{ExpressionMethods as _, insert_into};
 use diesel_async::RunQueryDsl;
 use ed25519_dalek::Signer as _;
 
+use super::common::ChannelTransport;
+
 #[tokio::test]
 #[test_log::test]
 pub async fn test_unregistered_pubkey_rejected() {
     let db = db::create_test_pool().await;
 
-    let mut client = ClientActor::new_manual(db.clone());
+    let (server_transport, mut test_transport) = ChannelTransport::new();
+    let props = ConnectionProps::new(db.clone(), Box::new(server_transport));
+    let task = tokio::spawn(connect_client(props));
 
     let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
     let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
 
-    let result = client
-        .process_transport_inbound(ClientRequest {
+    test_transport
+        .send(ClientRequest {
             payload: Some(ClientRequestPayload::AuthChallengeRequest(
                 AuthChallengeRequest {
                     pubkey: pubkey_bytes,
                 },
             )),
         })
-        .await;
+        .await
+        .unwrap();
 
-    match result {
-        Err(err) => {
-            assert_eq!(err, ClientError::PublicKeyNotRegistered);
-        }
-        Ok(_) => {
-            panic!("Expected error due to unregistered pubkey, but got success");
-        }
-    }
+    // Auth fails, connect_client returns, transport drops
+    task.await.unwrap();
 }
 
 #[tokio::test]
@@ -46,8 +46,6 @@ pub async fn test_unregistered_pubkey_rejected() {
 pub async fn test_challenge_auth() {
     let db = db::create_test_pool().await;
 
-    let mut client = ClientActor::new_manual(db.clone());
-
     let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
     let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
 
@@ -60,8 +58,13 @@ pub async fn test_challenge_auth() {
             .unwrap();
     }
 
-    let result = client
-        .process_transport_inbound(ClientRequest {
+    let (server_transport, mut test_transport) = ChannelTransport::new();
+    let props = ConnectionProps::new(db.clone(), Box::new(server_transport));
+    let task = tokio::spawn(connect_client(props));
+
+    // Send challenge request
+    test_transport
+        .send(ClientRequest {
             payload: Some(ClientRequestPayload::AuthChallengeRequest(
                 AuthChallengeRequest {
                     pubkey: pubkey_bytes,
@@ -69,34 +72,36 @@ pub async fn test_challenge_auth() {
             )),
         })
         .await
-        .expect("Shouldn't fail to process message");
+        .unwrap();
 
-    let ClientResponse {
-        payload: Some(ClientResponsePayload::AuthChallenge(challenge)),
-    } = result
-    else {
-        panic!("Expected auth challenge response, got {result:?}");
+    // Read the challenge response
+    let response = test_transport
+        .recv()
+        .await
+        .expect("should receive challenge");
+    let challenge = match response {
+        Ok(resp) => match resp.payload {
+            Some(ClientResponsePayload::AuthChallenge(c)) => c,
+            other => panic!("Expected AuthChallenge, got {other:?}"),
+        },
+        Err(err) => panic!("Expected Ok response, got Err({err:?})"),
     };
 
+    // Sign the challenge and send solution
     let formatted_challenge = arbiter_proto::format_challenge(challenge.nonce, &challenge.pubkey);
     let signature = new_key.sign(&formatted_challenge);
-    let serialized_signature = signature.to_bytes().to_vec();
 
-    let result = client
-        .process_transport_inbound(ClientRequest {
+    test_transport
+        .send(ClientRequest {
             payload: Some(ClientRequestPayload::AuthChallengeSolution(
                 AuthChallengeSolution {
-                    signature: serialized_signature,
+                    signature: signature.to_bytes().to_vec(),
                 },
             )),
         })
         .await
-        .expect("Shouldn't fail to process message");
+        .unwrap();
 
-    assert_eq!(
-        result,
-        ClientResponse {
-            payload: Some(ClientResponsePayload::AuthOk(AuthOk {})),
-        }
-    );
+    // Auth completes, session spawned
+    task.await.unwrap();
 }

server/crates/arbiter-server/tests/common/mod.rs

@@ -1,10 +1,14 @@
+use arbiter_proto::transport::{Bi, Error};
 use arbiter_server::{
     actors::keyholder::KeyHolder,
     db::{self, schema},
 };
+use async_trait::async_trait;
 use diesel::QueryDsl;
 use diesel_async::RunQueryDsl;
 use memsafe::MemSafe;
+use tokio::sync::mpsc;
+
 
 #[allow(dead_code)]
 pub async fn bootstrapped_keyholder(db: &db::DatabasePool) -> KeyHolder {
@@ -26,3 +30,46 @@ pub async fn root_key_history_id(db: &db::DatabasePool) -> i32 {
         .unwrap();
     id.expect("root_key_id should be set after bootstrap")
 }
+
+
+pub struct ChannelTransport<T, Y> {
+    receiver: mpsc::Receiver<T>,
+    sender: mpsc::Sender<Y>,
+}
+
+impl<T, Y> ChannelTransport<T, Y> {
+    pub fn new() -> (Self, ChannelTransport<Y, T>) {
+        let (tx1, rx1) = mpsc::channel(10);
+        let (tx2, rx2) = mpsc::channel(10);
+        (
+            Self {
+                receiver: rx1,
+                sender: tx2,
+            },
+            ChannelTransport {
+                receiver: rx2,
+                sender: tx1,
+            },
+        )
+    }
+}
+
+
+
+#[async_trait]
+impl<T, Y> Bi<T, Y> for ChannelTransport<T, Y>
+where
+    T: Send + 'static,
+    Y: Send + 'static,
+{
+    async fn send(&mut self, item: Y) -> Result<(), Error> {
+        self.sender
+            .send(item)
+            .await
+            .map_err(|_| Error::ChannelClosed)
+    }
+
+    async fn recv(&mut self) -> Option<T> {
+        self.receiver.recv().await
+    }
+}

server/crates/arbiter-server/tests/user_agent/auth.rs

@@ -1,13 +1,14 @@
 use arbiter_proto::proto::user_agent::{
-    AuthChallengeRequest, AuthChallengeSolution, AuthOk, UserAgentRequest, UserAgentResponse,
+    AuthChallengeRequest, AuthChallengeSolution, UserAgentRequest,
     user_agent_request::Payload as UserAgentRequestPayload,
     user_agent_response::Payload as UserAgentResponsePayload,
 };
+use arbiter_proto::transport::Bi;
 use arbiter_server::{
     actors::{
         GlobalActors,
         bootstrap::GetToken,
-        user_agent::{UserAgentActor, UserAgentError},
+        user_agent::{ConnectionProps, connect_user_agent},
     },
     db::{self, schema},
 };
@@ -15,20 +16,24 @@ use diesel::{ExpressionMethods as _, QueryDsl, insert_into};
 use diesel_async::RunQueryDsl;
 use ed25519_dalek::Signer as _;
 
+use super::common::ChannelTransport;
+
 #[tokio::test]
 #[test_log::test]
 pub async fn test_bootstrap_token_auth() {
     let db = db::create_test_pool().await;
-
     let actors = GlobalActors::spawn(db.clone()).await.unwrap();
     let token = actors.bootstrapper.ask(GetToken).await.unwrap().unwrap();
-    let mut user_agent = UserAgentActor::new_manual(db.clone(), actors);
+
+    let (server_transport, mut test_transport) = ChannelTransport::new();
+    let props = ConnectionProps::new(db.clone(), actors, Box::new(server_transport));
+    let task = tokio::spawn(connect_user_agent(props));
 
     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
-        .process_transport_inbound(UserAgentRequest {
+    test_transport
+        .send(UserAgentRequest {
             payload: Some(UserAgentRequestPayload::AuthChallengeRequest(
                 AuthChallengeRequest {
                     pubkey: pubkey_bytes,
@@ -37,14 +42,9 @@ pub async fn test_bootstrap_token_auth() {
             )),
         })
         .await
-        .expect("Shouldn't fail to process message");
+        .unwrap();
 
-    assert_eq!(
-        result,
-        UserAgentResponse {
-            payload: Some(UserAgentResponsePayload::AuthOk(AuthOk {})),
-        }
-    );
+    task.await.unwrap();
 
     let mut conn = db.get().await.unwrap();
     let stored_pubkey: Vec<u8> = schema::useragent_client::table
@@ -59,15 +59,17 @@ pub async fn test_bootstrap_token_auth() {
 #[test_log::test]
 pub async fn test_bootstrap_invalid_token_auth() {
     let db = db::create_test_pool().await;
-
     let actors = GlobalActors::spawn(db.clone()).await.unwrap();
-    let mut user_agent = UserAgentActor::new_manual(db.clone(), actors);
+
+    let (server_transport, mut test_transport) = ChannelTransport::new();
+    let props = ConnectionProps::new(db.clone(), actors, Box::new(server_transport));
+    let task = tokio::spawn(connect_user_agent(props));
 
     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
-        .process_transport_inbound(UserAgentRequest {
+    test_transport
+        .send(UserAgentRequest {
             payload: Some(UserAgentRequestPayload::AuthChallengeRequest(
                 AuthChallengeRequest {
                     pubkey: pubkey_bytes,
@@ -75,25 +77,27 @@ pub async fn test_bootstrap_invalid_token_auth() {
                 },
             )),
         })
-        .await;
+        .await
+        .unwrap();
 
-    match result {
-        Err(err) => {
-            assert_eq!(err, UserAgentError::InvalidBootstrapToken);
-        }
-        Ok(_) => {
-            panic!("Expected error due to invalid bootstrap token, but got success");
-        }
-    }
+    // Auth fails, connect_user_agent returns, transport drops
+    task.await.unwrap();
+
+    // Verify no key was registered
+    let mut conn = db.get().await.unwrap();
+    let count: i64 = schema::useragent_client::table
+        .count()
+        .get_result::<i64>(&mut conn)
+        .await
+        .unwrap();
+    assert_eq!(count, 0);
 }
 
 #[tokio::test]
 #[test_log::test]
 pub async fn test_challenge_auth() {
     let db = db::create_test_pool().await;
-
     let actors = GlobalActors::spawn(db.clone()).await.unwrap();
-    let mut user_agent = UserAgentActor::new_manual(db.clone(), actors);
 
     let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
     let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
@@ -107,8 +111,13 @@ pub async fn test_challenge_auth() {
             .unwrap();
     }
 
-    let result = user_agent
-        .process_transport_inbound(UserAgentRequest {
+    let (server_transport, mut test_transport) = ChannelTransport::new();
+    let props = ConnectionProps::new(db.clone(), actors, Box::new(server_transport));
+    let task = tokio::spawn(connect_user_agent(props));
+
+    // Send challenge request
+    test_transport
+        .send(UserAgentRequest {
             payload: Some(UserAgentRequestPayload::AuthChallengeRequest(
                 AuthChallengeRequest {
                     pubkey: pubkey_bytes,
@@ -117,34 +126,36 @@ pub async fn test_challenge_auth() {
             )),
         })
         .await
-        .expect("Shouldn't fail to process message");
+        .unwrap();
 
-    let UserAgentResponse {
-        payload: Some(UserAgentResponsePayload::AuthChallenge(challenge)),
-    } = result
-    else {
-        panic!("Expected auth challenge response, got {result:?}");
+    // Read the challenge response
+    let response = test_transport
+        .recv()
+        .await
+        .expect("should receive challenge");
+    let challenge = match response {
+        Ok(resp) => match resp.payload {
+            Some(UserAgentResponsePayload::AuthChallenge(c)) => c,
+            other => panic!("Expected AuthChallenge, got {other:?}"),
+        },
+        Err(err) => panic!("Expected Ok response, got Err({err:?})"),
     };
 
+    // Sign the challenge and send solution
     let formatted_challenge = arbiter_proto::format_challenge(challenge.nonce, &challenge.pubkey);
     let signature = new_key.sign(&formatted_challenge);
-    let serialized_signature = signature.to_bytes().to_vec();
 
-    let result = user_agent
-        .process_transport_inbound(UserAgentRequest {
+    test_transport
+        .send(UserAgentRequest {
             payload: Some(UserAgentRequestPayload::AuthChallengeSolution(
                 AuthChallengeSolution {
-                    signature: serialized_signature,
+                    signature: signature.to_bytes().to_vec(),
                 },
             )),
         })
         .await
-        .expect("Shouldn't fail to process message");
+        .unwrap();
 
-    assert_eq!(
-        result,
-        UserAgentResponse {
-            payload: Some(UserAgentResponsePayload::AuthOk(AuthOk {})),
-        }
-    );
+    // Auth completes, session spawned
+    task.await.unwrap();
 }

server/crates/arbiter-server/tests/user_agent/unseal.rs

@@ -1,15 +1,13 @@
 use arbiter_proto::proto::user_agent::{
-    AuthChallengeRequest, UnsealEncryptedKey, UnsealResult, UnsealStart,
-    UserAgentRequest,
+    UnsealEncryptedKey, UnsealResult, UnsealStart, UserAgentRequest,
     user_agent_request::Payload as UserAgentRequestPayload,
     user_agent_response::Payload as UserAgentResponsePayload,
 };
 use arbiter_server::{
     actors::{
         GlobalActors,
-        bootstrap::GetToken,
         keyholder::{Bootstrap, Seal},
-        user_agent::{UserAgentActor, UserAgentError},
+        user_agent::session::UserAgentSession,
     },
     db,
 };
@@ -17,16 +15,12 @@ use chacha20poly1305::{AeadInPlace, XChaCha20Poly1305, XNonce, aead::KeyInit};
 use memsafe::MemSafe;
 use x25519_dalek::{EphemeralSecret, PublicKey};
 
-
-async fn setup_authenticated_user_agent(
+async fn setup_sealed_user_agent(
     seal_key: &[u8],
-) -> (
-    arbiter_server::db::DatabasePool,
-    UserAgentActor,
-) {
+) -> (db::DatabasePool, UserAgentSession) {
     let db = db::create_test_pool().await;
-
     let actors = GlobalActors::spawn(db.clone()).await.unwrap();
+
     actors
         .key_holder
         .ask(Bootstrap {
@@ -36,27 +30,13 @@ async fn setup_authenticated_user_agent(
         .unwrap();
     actors.key_holder.ask(Seal).await.unwrap();
 
-    let mut user_agent = UserAgentActor::new_manual(db.clone(), actors.clone());
+    let session = UserAgentSession::new_test(db.clone(), actors);
 
-    let token = actors.bootstrapper.ask(GetToken).await.unwrap().unwrap();
-    let auth_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
-    user_agent
-        .process_transport_inbound(UserAgentRequest {
-            payload: Some(UserAgentRequestPayload::AuthChallengeRequest(
-                AuthChallengeRequest {
-                    pubkey: auth_key.verifying_key().to_bytes().to_vec(),
-                    bootstrap_token: Some(token),
-                },
-            )),
-        })
-        .await
-        .unwrap();
-
-    (db, user_agent)
+    (db, session)
 }
 
 async fn client_dh_encrypt(
-    user_agent: &mut UserAgentActor,
+    user_agent: &mut UserAgentSession,
     key_to_send: &[u8],
 ) -> UnsealEncryptedKey {
     let client_secret = EphemeralSecret::random();
@@ -103,7 +83,7 @@ fn unseal_key_request(req: UnsealEncryptedKey) -> UserAgentRequest {
 #[test_log::test]
 pub async fn test_unseal_success() {
     let seal_key = b"test-seal-key";
-    let (_db, mut user_agent) = setup_authenticated_user_agent(seal_key).await;
+    let (_db, mut user_agent) = setup_sealed_user_agent(seal_key).await;
 
     let encrypted_key = client_dh_encrypt(&mut user_agent, seal_key).await;
 
@@ -121,7 +101,7 @@ pub async fn test_unseal_success() {
 #[tokio::test]
 #[test_log::test]
 pub async fn test_unseal_wrong_seal_key() {
-    let (_db, mut user_agent) = setup_authenticated_user_agent(b"correct-key").await;
+    let (_db, mut user_agent) = setup_sealed_user_agent(b"correct-key").await;
 
     let encrypted_key = client_dh_encrypt(&mut user_agent, b"wrong-key").await;
 
@@ -139,7 +119,7 @@ pub async fn test_unseal_wrong_seal_key() {
 #[tokio::test]
 #[test_log::test]
 pub async fn test_unseal_corrupted_ciphertext() {
-    let (_db, mut user_agent) = setup_authenticated_user_agent(b"test-key").await;
+    let (_db, mut user_agent) = setup_sealed_user_agent(b"test-key").await;
 
     let client_secret = EphemeralSecret::random();
     let client_public = PublicKey::from(&client_secret);
@@ -168,38 +148,11 @@ pub async fn test_unseal_corrupted_ciphertext() {
     );
 }
 
-#[tokio::test]
-#[test_log::test]
-pub async fn test_unseal_start_without_auth_fails() {
-    let db = db::create_test_pool().await;
-
-    let actors = GlobalActors::spawn(db.clone()).await.unwrap();
-    let mut user_agent = UserAgentActor::new_manual(db.clone(), actors);
-
-    let client_secret = EphemeralSecret::random();
-    let client_public = PublicKey::from(&client_secret);
-
-    let result = user_agent
-        .process_transport_inbound(UserAgentRequest {
-            payload: Some(UserAgentRequestPayload::UnsealStart(UnsealStart {
-                client_pubkey: client_public.as_bytes().to_vec(),
-            })),
-        })
-        .await;
-
-    match result {
-        Err(err) => {
-            assert_eq!(err, UserAgentError::StateTransitionFailed);
-        }
-        other => panic!("Expected state machine error, got {other:?}"),
-    }
-}
-
 #[tokio::test]
 #[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_authenticated_user_agent(seal_key).await;
+    let (_db, mut user_agent) = setup_sealed_user_agent(seal_key).await;
 
     {
         let encrypted_key = client_dh_encrypt(&mut user_agent, b"wrong-key").await;