Merge remote-tracking branch 'origin/main' into SDK-client-UA-registration
Some checks failed
ci/woodpecker/pr/server-audit Pipeline failed
ci/woodpecker/pr/server-vet Pipeline failed
ci/woodpecker/pr/server-lint Pipeline was successful
ci/woodpecker/pr/server-test Pipeline was successful

This commit is contained in:
CleverWild
2026-03-21 21:14:41 +01:00
173 changed files with 16890 additions and 3578 deletions

View File

@@ -1,20 +1,7 @@
use alloy::{
consensus::TxEip1559,
primitives::{Address, Bytes, TxKind, U256},
rlp::Encodable,
};
use arbiter_proto::proto::{
client::{
AuthChallengeRequest, AuthChallengeSolution, ClientRequest,
client_request::Payload as ClientRequestPayload,
client_response::Payload as ClientResponsePayload,
},
evm::EvmSignTransactionRequest,
};
use arbiter_proto::transport::Bi;
use arbiter_proto::transport::{Receiver, Sender};
use arbiter_server::actors::GlobalActors;
use arbiter_server::{
actors::client::{ClientConnection, connect_client},
actors::client::{ClientConnection, auth, connect_client},
db::{self, schema},
};
use diesel::{ExpressionMethods as _, insert_into};
@@ -30,19 +17,17 @@ pub async fn test_unregistered_pubkey_rejected() {
let (server_transport, mut test_transport) = ChannelTransport::new();
let actors = GlobalActors::spawn(db.clone()).await.unwrap();
let props = ClientConnection::new(db.clone(), Box::new(server_transport), actors);
let task = tokio::spawn(connect_client(props));
let props = ClientConnection::new(db.clone(), actors);
let task = tokio::spawn(async move {
let mut server_transport = server_transport;
connect_client(props, &mut server_transport).await;
});
let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
test_transport
.send(ClientRequest {
payload: Some(ClientRequestPayload::AuthChallengeRequest(
AuthChallengeRequest {
pubkey: pubkey_bytes,
},
)),
.send(auth::Inbound::AuthChallengeRequest {
pubkey: new_key.verifying_key(),
})
.await
.unwrap();
@@ -71,17 +56,16 @@ pub async fn test_challenge_auth() {
let (server_transport, mut test_transport) = ChannelTransport::new();
let actors = GlobalActors::spawn(db.clone()).await.unwrap();
let props = ClientConnection::new(db.clone(), Box::new(server_transport), actors);
let task = tokio::spawn(connect_client(props));
let props = ClientConnection::new(db.clone(), actors);
let task = tokio::spawn(async move {
let mut server_transport = server_transport;
connect_client(props, &mut server_transport).await;
});
// Send challenge request
test_transport
.send(ClientRequest {
payload: Some(ClientRequestPayload::AuthChallengeRequest(
AuthChallengeRequest {
pubkey: pubkey_bytes,
},
)),
.send(auth::Inbound::AuthChallengeRequest {
pubkey: new_key.verifying_key(),
})
.await
.unwrap();
@@ -92,34 +76,29 @@ pub async fn test_challenge_auth() {
.await
.expect("should receive challenge");
let challenge = match response {
Ok(resp) => match resp.payload {
Some(ClientResponsePayload::AuthChallenge(c)) => c,
Ok(resp) => match resp {
auth::Outbound::AuthChallenge { pubkey, nonce } => (pubkey, nonce),
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 formatted_challenge = arbiter_proto::format_challenge(challenge.1, challenge.0.as_bytes());
let signature = new_key.sign(&formatted_challenge);
test_transport
.send(ClientRequest {
payload: Some(ClientRequestPayload::AuthChallengeSolution(
AuthChallengeSolution {
signature: signature.to_bytes().to_vec(),
},
)),
})
.send(auth::Inbound::AuthChallengeSolution { signature })
.await
.unwrap();
let response = test_transport.recv().await.expect("should receive auth ok");
let response = test_transport
.recv()
.await
.expect("should receive auth success");
match response {
Ok(resp) => match resp.payload {
Some(ClientResponsePayload::AuthOk(_)) => {}
other => panic!("Expected AuthOk, got {other:?}"),
},
Ok(auth::Outbound::AuthSuccess) => {}
Ok(other) => panic!("Expected AuthSuccess, got {other:?}"),
Err(err) => panic!("Expected Ok response, got Err({err:?})"),
}
@@ -127,114 +106,3 @@ pub async fn test_challenge_auth() {
task.await.unwrap();
}
#[tokio::test]
#[test_log::test]
pub async fn test_evm_sign_request_payload_is_handled() {
let db = db::create_test_pool().await;
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 (server_transport, mut test_transport) = ChannelTransport::new();
let actors = GlobalActors::spawn(db.clone()).await.unwrap();
let props = ClientConnection::new(db.clone(), Box::new(server_transport), actors);
let task = tokio::spawn(connect_client(props));
test_transport
.send(ClientRequest {
payload: Some(ClientRequestPayload::AuthChallengeRequest(
AuthChallengeRequest {
pubkey: pubkey_bytes,
},
)),
})
.await
.unwrap();
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:?})"),
};
let formatted_challenge = arbiter_proto::format_challenge(challenge.nonce, &challenge.pubkey);
let signature = new_key.sign(&formatted_challenge);
test_transport
.send(ClientRequest {
payload: Some(ClientRequestPayload::AuthChallengeSolution(
AuthChallengeSolution {
signature: signature.to_bytes().to_vec(),
},
)),
})
.await
.unwrap();
let response = test_transport.recv().await.expect("should receive auth ok");
match response {
Ok(resp) => match resp.payload {
Some(ClientResponsePayload::AuthOk(_)) => {}
other => panic!("Expected AuthOk, got {other:?}"),
},
Err(err) => panic!("Expected Ok response, got Err({err:?})"),
}
task.await.unwrap();
let tx = TxEip1559 {
chain_id: 1,
nonce: 0,
gas_limit: 21_000,
max_fee_per_gas: 1,
max_priority_fee_per_gas: 1,
to: TxKind::Call(Address::from_slice(&[0x11; 20])),
value: U256::ZERO,
input: Bytes::new(),
access_list: Default::default(),
};
let mut rlp_transaction = Vec::new();
tx.encode(&mut rlp_transaction);
test_transport
.send(ClientRequest {
payload: Some(ClientRequestPayload::EvmSignTransaction(
EvmSignTransactionRequest {
wallet_address: [0x22; 20].to_vec(),
rlp_transaction,
},
)),
})
.await
.unwrap();
let response = test_transport
.recv()
.await
.expect("should receive sign response");
match response {
Ok(resp) => match resp.payload {
Some(ClientResponsePayload::EvmSignTransaction(_)) => {}
other => panic!("Expected EvmSignTransaction response, got {other:?}"),
},
Err(err) => panic!("Expected Ok response, got Err({err:?})"),
}
}

View File

@@ -1,19 +1,19 @@
use arbiter_proto::transport::{Bi, Error};
use arbiter_proto::transport::{Bi, Error, Receiver, Sender};
use arbiter_server::{
actors::keyholder::KeyHolder,
db::{self, schema},
safe_cell::{SafeCell, SafeCellHandle as _},
};
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 {
let mut actor = KeyHolder::new(db.clone()).await.unwrap();
actor
.bootstrap(MemSafe::new(b"test-seal-key".to_vec()).unwrap())
.bootstrap(SafeCell::new(b"test-seal-key".to_vec()))
.await
.unwrap();
actor
@@ -55,10 +55,10 @@ impl<T, Y> ChannelTransport<T, Y> {
}
#[async_trait]
impl<T, Y> Bi<T, Y> for ChannelTransport<T, Y>
impl<T, Y> Sender<Y> for ChannelTransport<T, Y>
where
T: Send + 'static,
Y: Send + 'static,
T: Send + Sync + 'static,
Y: Send + Sync + 'static,
{
async fn send(&mut self, item: Y) -> Result<(), Error> {
self.sender
@@ -66,8 +66,22 @@ where
.await
.map_err(|_| Error::ChannelClosed)
}
}
#[async_trait]
impl<T, Y> Receiver<T> for ChannelTransport<T, Y>
where
T: Send + Sync + 'static,
Y: Send + Sync + 'static,
{
async fn recv(&mut self) -> Option<T> {
self.receiver.recv().await
}
}
impl<T, Y> Bi<T, Y> for ChannelTransport<T, Y>
where
T: Send + Sync + 'static,
Y: Send + Sync + 'static,
{
}

View File

@@ -3,11 +3,11 @@ use std::collections::{HashMap, HashSet};
use arbiter_server::{
actors::keyholder::{CreateNew, Error, KeyHolder},
db::{self, models, schema},
safe_cell::{SafeCell, SafeCellHandle as _},
};
use diesel::{ExpressionMethods as _, QueryDsl, SelectableHelper, dsl::sql_query};
use diesel_async::RunQueryDsl;
use kameo::actor::{ActorRef, Spawn as _};
use memsafe::MemSafe;
use tokio::task::JoinSet;
use crate::common;
@@ -24,7 +24,7 @@ async fn write_concurrently(
let plaintext = format!("{prefix}-{i}").into_bytes();
let id = actor
.ask(CreateNew {
plaintext: MemSafe::new(plaintext.clone()).unwrap(),
plaintext: SafeCell::new(plaintext.clone()),
})
.await
.unwrap();
@@ -118,7 +118,7 @@ async fn insert_failure_does_not_create_partial_row() {
drop(conn);
let err = actor
.create_new(MemSafe::new(b"should fail".to_vec()).unwrap())
.create_new(SafeCell::new(b"should fail".to_vec()))
.await
.unwrap_err();
assert!(matches!(err, Error::DatabaseTransaction(_)));
@@ -162,12 +162,12 @@ async fn decrypt_roundtrip_after_high_concurrency() {
let mut decryptor = KeyHolder::new(db.clone()).await.unwrap();
decryptor
.try_unseal(MemSafe::new(b"test-seal-key".to_vec()).unwrap())
.try_unseal(SafeCell::new(b"test-seal-key".to_vec()))
.await
.unwrap();
for (id, plaintext) in expected {
let mut decrypted = decryptor.decrypt(id).await.unwrap();
assert_eq!(*decrypted.read().unwrap(), plaintext);
assert_eq!(*decrypted.read(), plaintext);
}
}

View File

@@ -1,10 +1,10 @@
use arbiter_server::{
actors::keyholder::{Error, KeyHolder},
db::{self, models, schema},
safe_cell::{SafeCell, SafeCellHandle as _},
};
use diesel::{QueryDsl, SelectableHelper};
use diesel_async::RunQueryDsl;
use memsafe::MemSafe;
use crate::common;
@@ -14,7 +14,7 @@ async fn test_bootstrap() {
let db = db::create_test_pool().await;
let mut actor = KeyHolder::new(db.clone()).await.unwrap();
let seal_key = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
let seal_key = SafeCell::new(b"test-seal-key".to_vec());
actor.bootstrap(seal_key).await.unwrap();
let mut conn = db.get().await.unwrap();
@@ -43,7 +43,7 @@ async fn test_bootstrap_rejects_double() {
let db = db::create_test_pool().await;
let mut actor = common::bootstrapped_keyholder(&db).await;
let seal_key2 = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
let seal_key2 = SafeCell::new(b"test-seal-key".to_vec());
let err = actor.bootstrap(seal_key2).await.unwrap_err();
assert!(matches!(err, Error::AlreadyBootstrapped));
}
@@ -55,7 +55,7 @@ async fn test_create_new_before_bootstrap_fails() {
let mut actor = KeyHolder::new(db).await.unwrap();
let err = actor
.create_new(MemSafe::new(b"data".to_vec()).unwrap())
.create_new(SafeCell::new(b"data".to_vec()))
.await
.unwrap_err();
assert!(matches!(err, Error::NotBootstrapped));
@@ -91,17 +91,17 @@ async fn test_unseal_correct_password() {
let plaintext = b"survive a restart";
let aead_id = actor
.create_new(MemSafe::new(plaintext.to_vec()).unwrap())
.create_new(SafeCell::new(plaintext.to_vec()))
.await
.unwrap();
drop(actor);
let mut actor = KeyHolder::new(db.clone()).await.unwrap();
let seal_key = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
let seal_key = SafeCell::new(b"test-seal-key".to_vec());
actor.try_unseal(seal_key).await.unwrap();
let mut decrypted = actor.decrypt(aead_id).await.unwrap();
assert_eq!(*decrypted.read().unwrap(), plaintext);
assert_eq!(*decrypted.read(), plaintext);
}
#[tokio::test]
@@ -112,20 +112,20 @@ async fn test_unseal_wrong_then_correct_password() {
let plaintext = b"important data";
let aead_id = actor
.create_new(MemSafe::new(plaintext.to_vec()).unwrap())
.create_new(SafeCell::new(plaintext.to_vec()))
.await
.unwrap();
drop(actor);
let mut actor = KeyHolder::new(db.clone()).await.unwrap();
let bad_key = MemSafe::new(b"wrong-password".to_vec()).unwrap();
let bad_key = SafeCell::new(b"wrong-password".to_vec());
let err = actor.try_unseal(bad_key).await.unwrap_err();
assert!(matches!(err, Error::InvalidKey));
let good_key = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
let good_key = SafeCell::new(b"test-seal-key".to_vec());
actor.try_unseal(good_key).await.unwrap();
let mut decrypted = actor.decrypt(aead_id).await.unwrap();
assert_eq!(*decrypted.read().unwrap(), plaintext);
assert_eq!(*decrypted.read(), plaintext);
}

View File

@@ -3,10 +3,10 @@ use std::collections::HashSet;
use arbiter_server::{
actors::keyholder::{Error, encryption::v1},
db::{self, models, schema},
safe_cell::{SafeCell, SafeCellHandle as _},
};
use diesel::{ExpressionMethods as _, QueryDsl, SelectableHelper, dsl::update};
use diesel_async::RunQueryDsl;
use memsafe::MemSafe;
use crate::common;
@@ -18,12 +18,12 @@ async fn test_create_decrypt_roundtrip() {
let plaintext = b"hello arbiter";
let aead_id = actor
.create_new(MemSafe::new(plaintext.to_vec()).unwrap())
.create_new(SafeCell::new(plaintext.to_vec()))
.await
.unwrap();
let mut decrypted = actor.decrypt(aead_id).await.unwrap();
assert_eq!(*decrypted.read().unwrap(), plaintext);
assert_eq!(*decrypted.read(), plaintext);
}
#[tokio::test]
@@ -44,11 +44,11 @@ async fn test_ciphertext_differs_across_entries() {
let plaintext = b"same content";
let id1 = actor
.create_new(MemSafe::new(plaintext.to_vec()).unwrap())
.create_new(SafeCell::new(plaintext.to_vec()))
.await
.unwrap();
let id2 = actor
.create_new(MemSafe::new(plaintext.to_vec()).unwrap())
.create_new(SafeCell::new(plaintext.to_vec()))
.await
.unwrap();
@@ -70,8 +70,8 @@ async fn test_ciphertext_differs_across_entries() {
let mut d1 = actor.decrypt(id1).await.unwrap();
let mut d2 = actor.decrypt(id2).await.unwrap();
assert_eq!(*d1.read().unwrap(), plaintext);
assert_eq!(*d2.read().unwrap(), plaintext);
assert_eq!(*d1.read(), plaintext);
assert_eq!(*d2.read(), plaintext);
}
#[tokio::test]
@@ -83,7 +83,7 @@ async fn test_nonce_never_reused() {
let n = 5;
for i in 0..n {
actor
.create_new(MemSafe::new(format!("secret {i}").into_bytes()).unwrap())
.create_new(SafeCell::new(format!("secret {i}").into_bytes()))
.await
.unwrap();
}
@@ -137,7 +137,7 @@ async fn broken_db_nonce_format_fails_closed() {
drop(conn);
let err = actor
.create_new(MemSafe::new(b"must fail".to_vec()).unwrap())
.create_new(SafeCell::new(b"must fail".to_vec()))
.await
.unwrap_err();
assert!(matches!(err, Error::BrokenDatabase));
@@ -145,7 +145,7 @@ async fn broken_db_nonce_format_fails_closed() {
let db = db::create_test_pool().await;
let mut actor = common::bootstrapped_keyholder(&db).await;
let id = actor
.create_new(MemSafe::new(b"decrypt target".to_vec()).unwrap())
.create_new(SafeCell::new(b"decrypt target".to_vec()))
.await
.unwrap();
let mut conn = db.get().await.unwrap();

View File

@@ -2,7 +2,5 @@ mod common;
#[path = "user_agent/auth.rs"]
mod auth;
#[path = "user_agent/sdk_client.rs"]
mod sdk_client;
#[path = "user_agent/unseal.rs"]
mod unseal;

View File

@@ -1,14 +1,9 @@
use arbiter_proto::proto::user_agent::{
AuthChallengeRequest, AuthChallengeSolution, KeyType as ProtoKeyType, UserAgentRequest,
user_agent_request::Payload as UserAgentRequestPayload,
user_agent_response::Payload as UserAgentResponsePayload,
};
use arbiter_proto::transport::Bi;
use arbiter_proto::transport::{Receiver, Sender};
use arbiter_server::{
actors::{
GlobalActors,
bootstrap::GetToken,
user_agent::{UserAgentConnection, connect_user_agent},
user_agent::{AuthPublicKey, UserAgentConnection, auth},
},
db::{self, schema},
};
@@ -26,26 +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());
let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
test_transport
.send(UserAgentRequest {
payload: Some(UserAgentRequestPayload::AuthChallengeRequest(
AuthChallengeRequest {
pubkey: pubkey_bytes,
bootstrap_token: Some(token),
key_type: ProtoKeyType::Ed25519.into(),
},
)),
.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
@@ -63,27 +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());
let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
test_transport
.send(UserAgentRequest {
payload: Some(UserAgentRequestPayload::AuthChallengeRequest(
AuthChallengeRequest {
pubkey: pubkey_bytes,
bootstrap_token: Some("invalid_token".to_string()),
key_type: ProtoKeyType::Ed25519.into(),
},
)),
.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();
@@ -118,19 +116,17 @@ 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(UserAgentRequest {
payload: Some(UserAgentRequestPayload::AuthChallengeRequest(
AuthChallengeRequest {
pubkey: pubkey_bytes,
bootstrap_token: None,
key_type: ProtoKeyType::Ed25519.into(),
},
)),
.send(auth::Inbound::AuthChallengeRequest {
pubkey: AuthPublicKey::Ed25519(new_key.verifying_key()),
bootstrap_token: None,
})
.await
.unwrap();
@@ -141,28 +137,31 @@ pub async fn test_challenge_auth() {
.await
.expect("should receive challenge");
let challenge = match response {
Ok(resp) => match resp.payload {
Some(UserAgentResponsePayload::AuthChallenge(c)) => c,
Ok(resp) => match resp {
auth::Outbound::AuthChallenge { nonce } => nonce,
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 formatted_challenge = arbiter_proto::format_challenge(challenge, &pubkey_bytes);
let signature = new_key.sign(&formatted_challenge);
test_transport
.send(UserAgentRequest {
payload: Some(UserAgentRequestPayload::AuthChallengeSolution(
AuthChallengeSolution {
signature: signature.to_bytes().to_vec(),
},
)),
.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();
}

View File

@@ -1,259 +0,0 @@
use arbiter_proto::proto::user_agent::{
SdkClientApproveRequest, SdkClientError as ProtoSdkClientError, SdkClientRevokeRequest,
UserAgentRequest, sdk_client_approve_response, sdk_client_list_response,
sdk_client_revoke_response, user_agent_request::Payload as UserAgentRequestPayload,
user_agent_response::Payload as UserAgentResponsePayload,
};
use arbiter_server::{
actors::{
GlobalActors,
user_agent::{TransportResponseError, session::UserAgentSession},
},
db,
};
/// Shared helper: create a session and register a client pubkey via sdk_client_approve.
async fn make_session(db: &db::DatabasePool) -> UserAgentSession {
let actors = GlobalActors::spawn(db.clone()).await.unwrap();
UserAgentSession::new_test(db.clone(), actors)
}
#[tokio::test]
#[test_log::test]
async fn test_sdk_client_approve_registers_client() {
let db = db::create_test_pool().await;
let mut session = make_session(&db).await;
let pubkey = [0x42u8; 32];
let response = session
.process_transport_inbound(UserAgentRequest {
payload: Some(UserAgentRequestPayload::SdkClientApprove(
SdkClientApproveRequest {
pubkey: pubkey.to_vec(),
},
)),
})
.await
.expect("handler should succeed");
let entry = match response.payload.unwrap() {
UserAgentResponsePayload::SdkClientApprove(resp) => match resp.result.unwrap() {
sdk_client_approve_response::Result::Client(e) => e,
sdk_client_approve_response::Result::Error(e) => {
panic!("Expected Client, got error {:?}", e)
}
},
other => panic!("Expected SdkClientApprove, got {other:?}"),
};
assert_eq!(entry.pubkey, pubkey.to_vec());
assert!(entry.id > 0);
}
#[tokio::test]
#[test_log::test]
async fn test_sdk_client_approve_duplicate_returns_already_exists() {
let db = db::create_test_pool().await;
let mut session = make_session(&db).await;
let pubkey = [0x11u8; 32];
let req = UserAgentRequest {
payload: Some(UserAgentRequestPayload::SdkClientApprove(
SdkClientApproveRequest {
pubkey: pubkey.to_vec(),
},
)),
};
session
.process_transport_inbound(req.clone())
.await
.unwrap();
let err = session
.process_transport_inbound(req)
.await
.expect_err("second insert should return typed TransportResponseError");
assert_eq!(
err,
TransportResponseError::SdkClientApprove(ProtoSdkClientError::AlreadyExists)
);
}
#[tokio::test]
#[test_log::test]
async fn test_sdk_client_list_shows_registered_clients() {
let db = db::create_test_pool().await;
let mut session = make_session(&db).await;
let pubkey_a = [0x0Au8; 32];
let pubkey_b = [0x0Bu8; 32];
for pubkey in [pubkey_a, pubkey_b] {
session
.process_transport_inbound(UserAgentRequest {
payload: Some(UserAgentRequestPayload::SdkClientApprove(
SdkClientApproveRequest {
pubkey: pubkey.to_vec(),
},
)),
})
.await
.unwrap();
}
let response = session
.process_transport_inbound(UserAgentRequest {
payload: Some(UserAgentRequestPayload::SdkClientList(())),
})
.await
.expect("list should succeed");
let clients = match response.payload.unwrap() {
UserAgentResponsePayload::SdkClientList(resp) => match resp.result.unwrap() {
sdk_client_list_response::Result::Clients(list) => list.clients,
sdk_client_list_response::Result::Error(e) => {
panic!("Expected Clients, got error {:?}", e)
}
},
other => panic!("Expected SdkClientList, got {other:?}"),
};
assert_eq!(clients.len(), 2);
let pubkeys: Vec<Vec<u8>> = clients.into_iter().map(|e| e.pubkey).collect();
assert!(pubkeys.contains(&pubkey_a.to_vec()));
assert!(pubkeys.contains(&pubkey_b.to_vec()));
}
#[tokio::test]
#[test_log::test]
async fn test_sdk_client_revoke_removes_client() {
let db = db::create_test_pool().await;
let mut session = make_session(&db).await;
let pubkey = [0xBBu8; 32];
// Register a client and get its id
let approve_response = session
.process_transport_inbound(UserAgentRequest {
payload: Some(UserAgentRequestPayload::SdkClientApprove(
SdkClientApproveRequest {
pubkey: pubkey.to_vec(),
},
)),
})
.await
.unwrap();
let client_id = match approve_response.payload.unwrap() {
UserAgentResponsePayload::SdkClientApprove(resp) => match resp.result.unwrap() {
sdk_client_approve_response::Result::Client(e) => e.id,
sdk_client_approve_response::Result::Error(e) => panic!("approve failed: {:?}", e),
},
other => panic!("{other:?}"),
};
// Revoke the client
let revoke_response = session
.process_transport_inbound(UserAgentRequest {
payload: Some(UserAgentRequestPayload::SdkClientRevoke(
SdkClientRevokeRequest { client_id },
)),
})
.await
.expect("revoke should succeed");
match revoke_response.payload.unwrap() {
UserAgentResponsePayload::SdkClientRevoke(resp) => match resp.result.unwrap() {
sdk_client_revoke_response::Result::Ok(_) => {}
sdk_client_revoke_response::Result::Error(e) => {
panic!("Expected Ok, got error {:?}", e)
}
},
other => panic!("Expected SdkClientRevoke, got {other:?}"),
}
// List should now be empty
let list_response = session
.process_transport_inbound(UserAgentRequest {
payload: Some(UserAgentRequestPayload::SdkClientList(())),
})
.await
.unwrap();
let clients = match list_response.payload.unwrap() {
UserAgentResponsePayload::SdkClientList(resp) => match resp.result.unwrap() {
sdk_client_list_response::Result::Clients(list) => list.clients,
sdk_client_list_response::Result::Error(e) => panic!("list error: {:?}", e),
},
other => panic!("{other:?}"),
};
assert!(clients.is_empty(), "client should be removed after revoke");
}
#[tokio::test]
#[test_log::test]
async fn test_sdk_client_revoke_not_found_returns_error() {
let db = db::create_test_pool().await;
let mut session = make_session(&db).await;
let err = session
.process_transport_inbound(UserAgentRequest {
payload: Some(UserAgentRequestPayload::SdkClientRevoke(
SdkClientRevokeRequest { client_id: 9999 },
)),
})
.await
.expect_err("missing client should return typed TransportResponseError");
assert_eq!(
err,
TransportResponseError::SdkClientRevoke(ProtoSdkClientError::NotFound)
);
}
#[tokio::test]
#[test_log::test]
async fn test_sdk_client_approve_rejected_client_cannot_auth() {
// Verify the core flow: only pre-approved clients can authenticate
use arbiter_proto::proto::client::{
AuthChallengeRequest, ClientRequest, client_request::Payload as ClientRequestPayload,
client_response::Payload as ClientResponsePayload,
};
use arbiter_proto::transport::Bi as _;
use arbiter_server::actors::client::{ClientConnection, connect_client};
let db = db::create_test_pool().await;
let actors = GlobalActors::spawn(db.clone()).await.unwrap();
let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
let (server_transport, mut test_transport) = super::common::ChannelTransport::<_, _>::new();
let props = ClientConnection::new(db.clone(), Box::new(server_transport), actors.clone());
let task = tokio::spawn(connect_client(props));
test_transport
.send(ClientRequest {
payload: Some(ClientRequestPayload::AuthChallengeRequest(
AuthChallengeRequest {
pubkey: pubkey_bytes.clone(),
},
)),
})
.await
.unwrap();
let response = test_transport.recv().await.unwrap().unwrap();
assert!(
matches!(
response.payload.unwrap(),
ClientResponsePayload::ClientConnectError(_)
),
"unregistered client should be rejected"
);
task.await.unwrap();
}

View File

@@ -1,63 +1,55 @@
use arbiter_proto::proto::user_agent::{
UnsealEncryptedKey, UnsealResult, UnsealStart, UserAgentRequest,
user_agent_request::Payload as UserAgentRequestPayload,
user_agent_response::Payload as UserAgentResponsePayload,
};
use arbiter_server::{
actors::{
GlobalActors,
keyholder::{Bootstrap, Seal},
user_agent::session::UserAgentSession,
user_agent::session::{
HandleUnsealEncryptedKey, HandleUnsealRequest, UnsealError, UserAgentSession,
},
},
db,
safe_cell::{SafeCell, SafeCellHandle as _},
};
use chacha20poly1305::{AeadInPlace, XChaCha20Poly1305, XNonce, aead::KeyInit};
use memsafe::MemSafe;
use kameo::actor::Spawn as _;
use x25519_dalek::{EphemeralSecret, PublicKey};
async fn setup_sealed_user_agent(
seal_key: &[u8],
) -> (db::DatabasePool, UserAgentSession) {
) -> (db::DatabasePool, kameo::actor::ActorRef<UserAgentSession>) {
let db = db::create_test_pool().await;
let actors = GlobalActors::spawn(db.clone()).await.unwrap();
actors
.key_holder
.ask(Bootstrap {
seal_key_raw: MemSafe::new(seal_key.to_vec()).unwrap(),
seal_key_raw: SafeCell::new(seal_key.to_vec()),
})
.await
.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,
user_agent: &kameo::actor::ActorRef<UserAgentSession>,
key_to_send: &[u8],
) -> UnsealEncryptedKey {
) -> HandleUnsealEncryptedKey {
let client_secret = EphemeralSecret::random();
let client_public = PublicKey::from(&client_secret);
let response = user_agent
.process_transport_inbound(UserAgentRequest {
payload: Some(UserAgentRequestPayload::UnsealStart(UnsealStart {
client_pubkey: client_public.as_bytes().to_vec(),
})),
.ask(HandleUnsealRequest {
client_pubkey: client_public,
})
.await
.unwrap();
let server_pubkey = match response.payload.unwrap() {
UserAgentResponsePayload::UnsealStartResponse(resp) => resp.server_pubkey,
other => panic!("Expected UnsealStartResponse, got {other:?}"),
};
let server_public = PublicKey::from(<[u8; 32]>::try_from(server_pubkey.as_slice()).unwrap());
let server_pubkey = response.server_pubkey;
let shared_secret = client_secret.diffie_hellman(&server_public);
let shared_secret = client_secret.diffie_hellman(&server_pubkey);
let cipher = XChaCha20Poly1305::new(shared_secret.as_bytes().into());
let nonce = XNonce::from([0u8; 24]);
let associated_data = b"unseal";
@@ -66,119 +58,94 @@ async fn client_dh_encrypt(
.encrypt_in_place(&nonce, associated_data, &mut ciphertext)
.unwrap();
UnsealEncryptedKey {
HandleUnsealEncryptedKey {
nonce: nonce.to_vec(),
ciphertext,
associated_data: associated_data.to_vec(),
}
}
fn unseal_key_request(req: UnsealEncryptedKey) -> UserAgentRequest {
UserAgentRequest {
payload: Some(UserAgentRequestPayload::UnsealEncryptedKey(req)),
}
}
#[tokio::test]
#[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(unseal_key_request(encrypted_key))
.await
.unwrap();
assert_eq!(
response.payload.unwrap(),
UserAgentResponsePayload::UnsealResult(UnsealResult::Success.into()),
);
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 encrypted_key = client_dh_encrypt(&user_agent, b"wrong-key").await;
let response = user_agent
.process_transport_inbound(unseal_key_request(encrypted_key))
.await
.unwrap();
assert_eq!(
response.payload.unwrap(),
UserAgentResponsePayload::UnsealResult(UnsealResult::InvalidKey.into()),
);
let response = user_agent.ask(encrypted_key).await;
assert!(matches!(
response,
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(UserAgentRequest {
payload: Some(UserAgentRequestPayload::UnsealStart(UnsealStart {
client_pubkey: client_public.as_bytes().to_vec(),
})),
.ask(HandleUnsealRequest {
client_pubkey: client_public,
})
.await
.unwrap();
let response = user_agent
.process_transport_inbound(unseal_key_request(UnsealEncryptedKey {
.ask(HandleUnsealEncryptedKey {
nonce: vec![0u8; 24],
ciphertext: vec![0u8; 32],
associated_data: vec![],
}))
.await
.unwrap();
})
.await;
assert_eq!(
response.payload.unwrap(),
UserAgentResponsePayload::UnsealResult(UnsealResult::InvalidKey.into()),
);
assert!(matches!(
response,
Err(kameo::error::SendError::HandlerError(
UnsealError::InvalidKey
))
));
}
#[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_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 encrypted_key = client_dh_encrypt(&user_agent, b"wrong-key").await;
let response = user_agent
.process_transport_inbound(unseal_key_request(encrypted_key))
.await
.unwrap();
assert_eq!(
response.payload.unwrap(),
UserAgentResponsePayload::UnsealResult(UnsealResult::InvalidKey.into()),
);
let response = user_agent.ask(encrypted_key).await;
assert!(matches!(
response,
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(unseal_key_request(encrypted_key))
.await
.unwrap();
assert_eq!(
response.payload.unwrap(),
UserAgentResponsePayload::UnsealResult(UnsealResult::Success.into()),
);
let response = user_agent.ask(encrypted_key).await;
assert!(matches!(response, Ok(())));
}
}