refactor: consolidate auth messages into client and user_agent packages

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

@@ -1,7 +1,5 @@
-use arbiter_proto::proto::{
-    UserAgentResponse,
-    UserAgentRequest,
-    auth::{self, AuthChallengeRequest, AuthOk, ClientMessage, client_message::Payload as ClientAuthPayload},
+use arbiter_proto::proto::user_agent::{
+    AuthChallengeRequest, AuthChallengeSolution, AuthOk, UserAgentRequest, UserAgentResponse,
     user_agent_request::Payload as UserAgentRequestPayload,
     user_agent_response::Payload as UserAgentResponsePayload,
 };
@@ -17,18 +15,10 @@ use diesel::{ExpressionMethods as _, QueryDsl, insert_into};
 use diesel_async::RunQueryDsl;
 use ed25519_dalek::Signer as _;
 
-fn auth_request(payload: ClientAuthPayload) -> UserAgentRequest {
-    UserAgentRequest {
-        payload: Some(UserAgentRequestPayload::AuthMessage(ClientMessage {
-            payload: Some(payload),
-        })),
-    }
-}
-
 #[tokio::test]
 #[test_log::test]
 pub async fn test_bootstrap_token_auth() {
-    let db =db::create_test_pool().await;
+    let db = db::create_test_pool().await;
 
     let actors = GlobalActors::spawn(db.clone()).await.unwrap();
     let token = actors.bootstrapper.ask(GetToken).await.unwrap().unwrap();
@@ -38,25 +28,21 @@ pub async fn test_bootstrap_token_auth() {
     let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
 
     let result = user_agent
-        .process_transport_inbound(auth_request(ClientAuthPayload::AuthChallengeRequest(
-            AuthChallengeRequest {
-                pubkey: pubkey_bytes,
-                bootstrap_token: Some(token),
-            },
-        )))
+        .process_transport_inbound(UserAgentRequest {
+            payload: Some(UserAgentRequestPayload::AuthChallengeRequest(
+                AuthChallengeRequest {
+                    pubkey: pubkey_bytes,
+                    bootstrap_token: Some(token),
+                },
+            )),
+        })
         .await
         .expect("Shouldn't fail to process message");
 
     assert_eq!(
         result,
         UserAgentResponse {
-            payload: Some(UserAgentResponsePayload::AuthMessage(
-                arbiter_proto::proto::auth::ServerMessage {
-                    payload: Some(arbiter_proto::proto::auth::server_message::Payload::AuthOk(
-                        AuthOk {},
-                    )),
-                },
-            )),
+            payload: Some(UserAgentResponsePayload::AuthOk(AuthOk {})),
         }
     );
 
@@ -81,12 +67,14 @@ pub async fn test_bootstrap_invalid_token_auth() {
     let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
 
     let result = user_agent
-        .process_transport_inbound(auth_request(ClientAuthPayload::AuthChallengeRequest(
-            AuthChallengeRequest {
-                pubkey: pubkey_bytes,
-                bootstrap_token: Some("invalid_token".to_string()),
-            },
-        )))
+        .process_transport_inbound(UserAgentRequest {
+            payload: Some(UserAgentRequestPayload::AuthChallengeRequest(
+                AuthChallengeRequest {
+                    pubkey: pubkey_bytes,
+                    bootstrap_token: Some("invalid_token".to_string()),
+                },
+            )),
+        })
         .await;
 
     match result {
@@ -120,49 +108,43 @@ pub async fn test_challenge_auth() {
     }
 
     let result = user_agent
-        .process_transport_inbound(auth_request(ClientAuthPayload::AuthChallengeRequest(
-            AuthChallengeRequest {
-                pubkey: pubkey_bytes,
-                bootstrap_token: None,
-            },
-        )))
+        .process_transport_inbound(UserAgentRequest {
+            payload: Some(UserAgentRequestPayload::AuthChallengeRequest(
+                AuthChallengeRequest {
+                    pubkey: pubkey_bytes,
+                    bootstrap_token: None,
+                },
+            )),
+        })
         .await
         .expect("Shouldn't fail to process message");
 
     let UserAgentResponse {
-        payload:
-            Some(UserAgentResponsePayload::AuthMessage(arbiter_proto::proto::auth::ServerMessage {
-                payload:
-                    Some(arbiter_proto::proto::auth::server_message::Payload::AuthChallenge(challenge)),
-            })),
+        payload: Some(UserAgentResponsePayload::AuthChallenge(challenge)),
     } = result
     else {
         panic!("Expected auth challenge response, got {result:?}");
     };
 
-    let formatted_challenge = arbiter_proto::format_challenge(&challenge);
+    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(auth_request(ClientAuthPayload::AuthChallengeSolution(
-            auth::AuthChallengeSolution {
-                signature: serialized_signature,
-            },
-        )))
+        .process_transport_inbound(UserAgentRequest {
+            payload: Some(UserAgentRequestPayload::AuthChallengeSolution(
+                AuthChallengeSolution {
+                    signature: serialized_signature,
+                },
+            )),
+        })
         .await
         .expect("Shouldn't fail to process message");
 
     assert_eq!(
         result,
         UserAgentResponse {
-            payload: Some(UserAgentResponsePayload::AuthMessage(
-                arbiter_proto::proto::auth::ServerMessage {
-                    payload: Some(arbiter_proto::proto::auth::server_message::Payload::AuthOk(
-                        AuthOk {},
-                    )),
-                },
-            )),
+            payload: Some(UserAgentResponsePayload::AuthOk(AuthOk {})),
         }
     );
 }

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

@@ -1,6 +1,6 @@
-use arbiter_proto::proto::{
-    UnsealEncryptedKey, UnsealResult, UnsealStart, UserAgentRequest, UserAgentResponse,
-    auth::{AuthChallengeRequest, ClientMessage, client_message::Payload as ClientAuthPayload},
+use arbiter_proto::proto::user_agent::{
+    AuthChallengeRequest, UnsealEncryptedKey, UnsealResult, UnsealStart,
+    UserAgentRequest, UserAgentResponse,
     user_agent_request::Payload as UserAgentRequestPayload,
     user_agent_response::Payload as UserAgentResponsePayload,
 };
@@ -21,26 +21,6 @@ use x25519_dalek::{EphemeralSecret, PublicKey};
 type TestUserAgent =
     UserAgentActor<DummyTransport<UserAgentRequest, Result<UserAgentResponse, UserAgentError>>>;
 
-fn auth_request(payload: ClientAuthPayload) -> UserAgentRequest {
-    UserAgentRequest {
-        payload: Some(UserAgentRequestPayload::AuthMessage(ClientMessage {
-            payload: Some(payload),
-        })),
-    }
-}
-
-fn unseal_start_request(req: UnsealStart) -> UserAgentRequest {
-    UserAgentRequest {
-        payload: Some(UserAgentRequestPayload::UnsealStart(req)),
-    }
-}
-
-fn unseal_key_request(req: UnsealEncryptedKey) -> UserAgentRequest {
-    UserAgentRequest {
-        payload: Some(UserAgentRequestPayload::UnsealEncryptedKey(req)),
-    }
-}
-
 async fn setup_authenticated_user_agent(
     seal_key: &[u8],
 ) -> (
@@ -64,12 +44,14 @@ async fn setup_authenticated_user_agent(
     let token = actors.bootstrapper.ask(GetToken).await.unwrap().unwrap();
     let auth_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
     user_agent
-        .process_transport_inbound(auth_request(ClientAuthPayload::AuthChallengeRequest(
-            AuthChallengeRequest {
-                pubkey: auth_key.verifying_key().to_bytes().to_vec(),
-                bootstrap_token: Some(token),
-            },
-        )))
+        .process_transport_inbound(UserAgentRequest {
+            payload: Some(UserAgentRequestPayload::AuthChallengeRequest(
+                AuthChallengeRequest {
+                    pubkey: auth_key.verifying_key().to_bytes().to_vec(),
+                    bootstrap_token: Some(token),
+                },
+            )),
+        })
         .await
         .unwrap();
 
@@ -84,9 +66,11 @@ async fn client_dh_encrypt(
     let client_public = PublicKey::from(&client_secret);
 
     let response = user_agent
-        .process_transport_inbound(unseal_start_request(UnsealStart {
-            client_pubkey: client_public.as_bytes().to_vec(),
-        }))
+        .process_transport_inbound(UserAgentRequest {
+            payload: Some(UserAgentRequestPayload::UnsealStart(UnsealStart {
+                client_pubkey: client_public.as_bytes().to_vec(),
+            })),
+        })
         .await
         .unwrap();
 
@@ -112,6 +96,12 @@ async fn client_dh_encrypt(
     }
 }
 
+fn unseal_key_request(req: UnsealEncryptedKey) -> UserAgentRequest {
+    UserAgentRequest {
+        payload: Some(UserAgentRequestPayload::UnsealEncryptedKey(req)),
+    }
+}
+
 #[tokio::test]
 #[test_log::test]
 pub async fn test_unseal_success() {
@@ -158,9 +148,11 @@ pub async fn test_unseal_corrupted_ciphertext() {
     let client_public = PublicKey::from(&client_secret);
 
     user_agent
-        .process_transport_inbound(unseal_start_request(UnsealStart {
-            client_pubkey: client_public.as_bytes().to_vec(),
-        }))
+        .process_transport_inbound(UserAgentRequest {
+            payload: Some(UserAgentRequestPayload::UnsealStart(UnsealStart {
+                client_pubkey: client_public.as_bytes().to_vec(),
+            })),
+        })
         .await
         .unwrap();
 
@@ -191,9 +183,11 @@ pub async fn test_unseal_start_without_auth_fails() {
     let client_public = PublicKey::from(&client_secret);
 
     let result = user_agent
-        .process_transport_inbound(unseal_start_request(UnsealStart {
-            client_pubkey: client_public.as_bytes().to_vec(),
-        }))
+        .process_transport_inbound(UserAgentRequest {
+            payload: Some(UserAgentRequestPayload::UnsealStart(UnsealStart {
+                client_pubkey: client_public.as_bytes().to_vec(),
+            })),
+        })
         .await;
 
     match result {