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

refactor(server::{user_agent, client}): move auth part to separate function to not to pollute actor session with one-time concerns
Some checks failed
ci/woodpecker/pr/server-lint Pipeline failed
Details
ci/woodpecker/pr/server-audit Pipeline was successful
Details
ci/woodpecker/pr/server-vet Pipeline failed
Details
ci/woodpecker/pr/server-test Pipeline was successful
Details

Browse Source
This commit is contained in:
hdbg
2026-03-01 19:59:42 +01:00
parent 004b14a168
commit f709650bb1
18 changed files with 1176 additions and 975 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

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