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

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

Browse Source
This commit is contained in:
hdbg
2026-03-18 23:37:40 +01:00
committed by Stas
parent 2ff4d0961c
commit 60ce1cc110
4 changed files with 107 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
server/crates/arbiter-server/src/actors/user_agent/session.rs
View File

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

4
server/crates/arbiter-server/src/actors/user_agent/session/connection.rs
View File

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

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

@@ -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 task = tokio::spawn(connect_user_agent(props));
let db_for_task = db.clone();
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 task = tokio::spawn(connect_user_agent(props));
let db_for_task = db.clone();
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
task.await.unwrap();
assert!(matches!(
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 task = tokio::spawn(connect_user_agent(props));
let db_for_task = db.clone();
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
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();
}

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

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