MarketTakers/arbiter
5
0
Fork 0
Code Issues 27 Pull Requests 5 Packages Projects Releases Wiki Activity

refactor: consolidate auth messages into client and user_agent packages

Browse Source
This commit is contained in:
hdbg
2026-03-01 11:44:34 +01:00
parent 06f4d628db
commit 4b4a8f4489
19 changed files with 686 additions and 264 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
server/crates/arbiter-server/tests/client.rs Normal file
View File

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

102
server/crates/arbiter-server/tests/client/auth.rs Normal file
View File

@@ -0,0 +1,102 @@
use arbiter_proto::proto::client::{
AuthChallengeRequest, AuthChallengeSolution, AuthOk, ClientRequest, ClientResponse,
client_request::Payload as ClientRequestPayload,
client_response::Payload as ClientResponsePayload,
};
use arbiter_server::{
actors::client::{ClientActor, ClientError},
db::{self, schema},
};
use diesel::{ExpressionMethods as _, insert_into};
use diesel_async::RunQueryDsl;
use ed25519_dalek::Signer as _;
#[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 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 {
payload: Some(ClientRequestPayload::AuthChallengeRequest(
AuthChallengeRequest {
pubkey: pubkey_bytes,
},
)),
})
.await;
match result {
Err(err) => {
assert_eq!(err, ClientError::PublicKeyNotRegistered);
}
Ok(_) => {
panic!("Expected error due to unregistered pubkey, but got success");
}
}
}
#[tokio::test]
#[test_log::test]
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();
{
let mut conn = db.get().await.unwrap();
insert_into(schema::program_client::table)
.values(schema::program_client::public_key.eq(pubkey_bytes.clone()))
.execute(&mut conn)
.await
.unwrap();
}
let result = client
.process_transport_inbound(ClientRequest {
payload: Some(ClientRequestPayload::AuthChallengeRequest(
AuthChallengeRequest {
pubkey: pubkey_bytes,
},
)),
})
.await
.expect("Shouldn't fail to process message");
let ClientResponse {
payload: Some(ClientResponsePayload::AuthChallenge(challenge)),
} = result
else {
panic!("Expected auth challenge response, got {result:?}");
};
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 {
payload: Some(ClientRequestPayload::AuthChallengeSolution(
AuthChallengeSolution {
signature: serialized_signature,
},
)),
})
.await
.expect("Shouldn't fail to process message");
assert_eq!(
result,
ClientResponse {
payload: Some(ClientResponsePayload::AuthOk(AuthOk {})),
}
);
}

94
server/crates/arbiter-server/tests/user_agent/auth.rs
View File

@@ -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 {})),
}
);
}

70
server/crates/arbiter-server/tests/user_agent/unseal.rs
View File

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