refactor(server): reorganized client/user_agent actors into separate module peers and added event MessageBus

This commit is contained in:
hdbg
2026-04-07 23:54:29 +02:00
parent f3cf6a9438
commit 1585f90cae
42 changed files with 332 additions and 286 deletions

View File

@@ -0,0 +1,423 @@
use arbiter_crypto::authn::{self, CLIENT_CONTEXT};
use arbiter_proto::{
ClientMetadata,
transport::{Bi, expect_message},
};
use chrono::Utc;
use diesel::{
ExpressionMethods as _, OptionalExtension as _, QueryDsl as _, SelectableHelper as _,
dsl::insert_into, update,
};
use diesel_async::RunQueryDsl as _;
use kameo::{actor::ActorRef, error::SendError};
use tracing::error;
use crate::{
actors::{
GlobalActors, flow_coordinator::{self, RequestClientApproval}, vault::Vault
},
crypto::integrity::{self, AttestationStatus},
db::{
self,
models::{ProgramClientMetadata, SqliteTimestamp},
schema::program_client,
},
};
use super::{ClientConnection, ClientCredentials, ClientProfile};
#[derive(thiserror::Error, Debug, Clone, PartialEq, Eq)]
pub enum Error {
#[error("Database pool unavailable")]
DatabasePoolUnavailable,
#[error("Database operation failed")]
DatabaseOperationFailed,
#[error("Integrity check failed")]
IntegrityCheckFailed,
#[error("Invalid challenge solution")]
InvalidChallengeSolution,
#[error("Client approval request failed")]
ApproveError(#[from] ApproveError),
#[error("Transport error")]
Transport,
}
impl From<diesel::result::Error> for Error {
fn from(e: diesel::result::Error) -> Self {
error!(?e, "Database error");
Self::DatabaseOperationFailed
}
}
#[derive(thiserror::Error, Debug, Clone, PartialEq, Eq)]
pub enum ApproveError {
#[error("Internal error")]
Internal,
#[error("Client connection denied by user agents")]
Denied,
#[error("Upstream error: {0}")]
Upstream(flow_coordinator::ApprovalError),
}
#[derive(Debug, Clone)]
pub enum Inbound {
AuthChallengeRequest {
pubkey: authn::PublicKey,
metadata: ClientMetadata,
},
AuthChallengeSolution {
signature: authn::Signature,
},
}
#[derive(Debug, Clone)]
pub enum Outbound {
AuthChallenge {
pubkey: authn::PublicKey,
nonce: i32,
},
AuthSuccess,
}
/// Returns the current nonce and client ID for a registered client.
/// Returns `None` if the pubkey is not registered.
async fn get_current_nonce_and_id(
db: &db::DatabasePool,
pubkey: &authn::PublicKey,
) -> Result<Option<(i32, i32)>, Error> {
let pubkey_bytes = pubkey.to_bytes();
let mut conn = db.get().await.map_err(|e| {
error!(error = ?e, "Database pool error");
Error::DatabasePoolUnavailable
})?;
program_client::table
.filter(program_client::public_key.eq(&pubkey_bytes))
.select((program_client::id, program_client::nonce))
.first::<(i32, i32)>(&mut conn)
.await
.optional()
.map_err(|e| {
error!(error = ?e, "Database error");
Error::DatabaseOperationFailed
})
}
async fn verify_integrity(
db: &db::DatabasePool,
vault: &ActorRef<Vault>,
pubkey: &authn::PublicKey,
) -> Result<(), Error> {
let mut db_conn = db.get().await.map_err(|e| {
error!(error = ?e, "Database pool error");
Error::DatabasePoolUnavailable
})?;
let (id, nonce) = get_current_nonce_and_id(db, pubkey).await?.ok_or_else(|| {
error!("Client not found during integrity verification");
Error::DatabaseOperationFailed
})?;
let attestation = integrity::verify_entity(
&mut db_conn,
vault,
&ClientCredentials {
pubkey: pubkey.clone(),
nonce,
},
id,
)
.await
.map_err(|e| {
error!(?e, "Integrity verification failed");
Error::IntegrityCheckFailed
})?;
if attestation != AttestationStatus::Attested {
error!("Integrity attestation unavailable for client {id}");
return Err(Error::IntegrityCheckFailed);
}
Ok(())
}
/// Atomically increments the nonce and re-signs the integrity envelope.
/// Returns the new nonce, which is used as the challenge nonce.
async fn create_nonce(
db: &db::DatabasePool,
vault: &ActorRef<Vault>,
pubkey: &authn::PublicKey,
) -> Result<i32, Error> {
let pubkey_bytes = pubkey.to_bytes();
let pubkey = pubkey.clone();
let mut conn = db.get().await.map_err(|e| {
error!(error = ?e, "Database pool error");
Error::DatabasePoolUnavailable
})?;
conn.exclusive_transaction(|conn| {
let vault = vault.clone();
let pubkey = pubkey.clone();
Box::pin(async move {
let (id, new_nonce): (i32, i32) = update(program_client::table)
.filter(program_client::public_key.eq(&pubkey_bytes))
.set(program_client::nonce.eq(program_client::nonce + 1))
.returning((program_client::id, program_client::nonce))
.get_result(conn)
.await?;
integrity::sign_entity(
conn,
&vault,
&ClientCredentials {
pubkey: pubkey.clone(),
nonce: new_nonce,
},
id,
)
.await
.map_err(|e| {
error!(?e, "Integrity sign failed after nonce update");
Error::DatabaseOperationFailed
})?;
Ok(new_nonce)
})
})
.await
}
async fn approve_new_client(
actors: &GlobalActors,
profile: ClientProfile,
) -> Result<(), Error> {
let result = actors
.flow_coordinator
.ask(RequestClientApproval { client: profile })
.await;
match result {
Ok(true) => Ok(()),
Ok(false) => Err(Error::ApproveError(ApproveError::Denied)),
Err(SendError::HandlerError(e)) => {
error!(error = ?e, "Approval upstream error");
Err(Error::ApproveError(ApproveError::Upstream(e)))
}
Err(e) => {
error!(error = ?e, "Approval request to flow coordinator failed");
Err(Error::ApproveError(ApproveError::Internal))
}
}
}
async fn insert_client(
db: &db::DatabasePool,
vault: &ActorRef<Vault>,
pubkey: &authn::PublicKey,
metadata: &ClientMetadata,
) -> Result<i32, Error> {
use crate::db::schema::{client_metadata, program_client};
let pubkey = pubkey.clone();
let metadata = metadata.clone();
let mut conn = db.get().await.map_err(|e| {
error!(error = ?e, "Database pool error");
Error::DatabasePoolUnavailable
})?;
conn.exclusive_transaction(|conn| {
let vault = vault.clone();
let pubkey = pubkey.clone();
Box::pin(async move {
const NONCE_START: i32 = 1;
let metadata_id = insert_into(client_metadata::table)
.values((
client_metadata::name.eq(&metadata.name),
client_metadata::description.eq(&metadata.description),
client_metadata::version.eq(&metadata.version),
))
.returning(client_metadata::id)
.get_result::<i32>(conn)
.await?;
let client_id = insert_into(program_client::table)
.values((
program_client::public_key.eq(pubkey.to_bytes()),
program_client::metadata_id.eq(metadata_id),
program_client::nonce.eq(NONCE_START),
))
.on_conflict_do_nothing()
.returning(program_client::id)
.get_result::<i32>(conn)
.await?;
integrity::sign_entity(
conn,
&vault,
&ClientCredentials {
pubkey: pubkey.clone(),
nonce: NONCE_START,
},
client_id,
)
.await
.map_err(|e| {
error!(error = ?e, "Failed to sign integrity tag for new client key");
Error::DatabaseOperationFailed
})?;
Ok(client_id)
})
})
.await
}
async fn sync_client_metadata(
db: &db::DatabasePool,
client_id: i32,
metadata: &ClientMetadata,
) -> Result<(), Error> {
use crate::db::schema::{client_metadata, client_metadata_history};
let now = SqliteTimestamp(Utc::now());
let mut conn = db.get().await.map_err(|e| {
error!(error = ?e, "Database pool error");
Error::DatabasePoolUnavailable
})?;
conn.exclusive_transaction(|conn| {
let metadata = metadata.clone();
Box::pin(async move {
let (current_metadata_id, current): (i32, ProgramClientMetadata) =
program_client::table
.find(client_id)
.inner_join(client_metadata::table)
.select((
program_client::metadata_id,
ProgramClientMetadata::as_select(),
))
.first(conn)
.await?;
let unchanged = current.name == metadata.name
&& current.description == metadata.description
&& current.version == metadata.version;
if unchanged {
return Ok(());
}
insert_into(client_metadata_history::table)
.values((
client_metadata_history::metadata_id.eq(current_metadata_id),
client_metadata_history::client_id.eq(client_id),
))
.execute(conn)
.await?;
let metadata_id = insert_into(client_metadata::table)
.values((
client_metadata::name.eq(&metadata.name),
client_metadata::description.eq(&metadata.description),
client_metadata::version.eq(&metadata.version),
))
.returning(client_metadata::id)
.get_result::<i32>(conn)
.await?;
update(program_client::table.find(client_id))
.set((
program_client::metadata_id.eq(metadata_id),
program_client::updated_at.eq(now),
))
.execute(conn)
.await?;
Ok::<(), diesel::result::Error>(())
})
})
.await
.map_err(|e| {
error!(error = ?e, "Database error");
Error::DatabaseOperationFailed
})
}
async fn challenge_client<T>(
transport: &mut T,
pubkey: authn::PublicKey,
nonce: i32,
) -> Result<(), Error>
where
T: Bi<Inbound, Result<Outbound, Error>> + ?Sized,
{
transport
.send(Ok(Outbound::AuthChallenge {
pubkey: pubkey.clone(),
nonce,
}))
.await
.map_err(|e| {
error!(error = ?e, "Failed to send auth challenge");
Error::Transport
})?;
let signature = expect_message(transport, |req: Inbound| match req {
Inbound::AuthChallengeSolution { signature } => Some(signature),
_ => None,
})
.await
.map_err(|e| {
error!(error = ?e, "Failed to receive challenge solution");
Error::Transport
})?;
if !pubkey.verify(nonce, CLIENT_CONTEXT, &signature) {
error!("Challenge solution verification failed");
return Err(Error::InvalidChallengeSolution);
}
Ok(())
}
pub async fn authenticate<T>(props: &mut ClientConnection, transport: &mut T) -> Result<i32, Error>
where
T: Bi<Inbound, Result<Outbound, Error>> + Send + ?Sized,
{
let Some(Inbound::AuthChallengeRequest { pubkey, metadata }) = transport.recv().await else {
return Err(Error::Transport);
};
let client_id = match get_current_nonce_and_id(&props.db, &pubkey).await? {
Some((id, _)) => {
verify_integrity(&props.db, &props.actors.vault, &pubkey).await?;
id
}
None => {
approve_new_client(
&props.actors,
ClientProfile {
pubkey: pubkey.clone(),
metadata: metadata.clone(),
},
)
.await?;
insert_client(&props.db, &props.actors.vault, &pubkey, &metadata).await?
}
};
sync_client_metadata(&props.db, client_id, &metadata).await?;
let challenge_nonce = create_nonce(&props.db, &props.actors.vault, &pubkey).await?;
challenge_client(transport, pubkey, challenge_nonce).await?;
transport
.send(Ok(Outbound::AuthSuccess))
.await
.map_err(|e| {
error!(error = ?e, "Failed to send auth success");
Error::Transport
})?;
Ok(client_id)
}

View File

@@ -0,0 +1,64 @@
use arbiter_crypto::authn;
use arbiter_proto::{ClientMetadata, transport::Bi};
use kameo::actor::Spawn;
use tracing::{error, info};
use crate::{
actors::GlobalActors,
crypto::integrity::{Integrable, hashing::Hashable},
db, peers::client::session::ClientSession,
};
#[derive(Debug, Clone)]
pub struct ClientProfile {
pub pubkey: authn::PublicKey,
pub metadata: ClientMetadata,
}
pub struct ClientCredentials {
pub pubkey: authn::PublicKey,
pub nonce: i32,
}
impl Integrable for ClientCredentials {
const KIND: &'static str = "client_credentials";
}
impl Hashable for ClientCredentials {
fn hash<H: sha2::Digest>(&self, hasher: &mut H) {
hasher.update(self.pubkey.to_bytes());
self.nonce.hash(hasher);
}
}
pub struct ClientConnection {
pub(crate) db: db::DatabasePool,
pub(crate) actors: GlobalActors,
}
impl ClientConnection {
pub fn new(db: db::DatabasePool, actors: GlobalActors) -> Self {
Self { db, actors }
}
}
pub mod auth;
pub mod session;
pub async fn connect_client<T>(mut props: ClientConnection, transport: &mut T)
where
T: Bi<auth::Inbound, Result<auth::Outbound, auth::Error>> + Send + ?Sized,
{
let fut = auth::authenticate(&mut props, transport);
println!("authenticate future size: {}", std::mem::size_of_val(&fut));
match fut.await {
Ok(client_id) => {
ClientSession::spawn(ClientSession::new(props, client_id));
info!("Client authenticated, session started");
}
Err(err) => {
let _ = transport.send(Err(err.clone())).await;
error!(?err, "Authentication failed, closing connection");
}
}
}

View File

@@ -0,0 +1,120 @@
use kameo::{Actor, messages};
use tracing::error;
use alloy::{consensus::TxEip1559, primitives::Address, signers::Signature};
use crate::{
actors::{
GlobalActors,
evm::{ClientSignTransaction, SignTransactionError},
flow_coordinator::RegisterClient,
vault::VaultState,
},
db,
evm::VetError,
};
use super::ClientConnection;
pub struct ClientSession {
props: ClientConnection,
client_id: i32,
}
impl ClientSession {
pub(crate) fn new(props: ClientConnection, client_id: i32) -> Self {
Self { props, client_id }
}
}
#[messages]
impl ClientSession {
#[message]
pub(crate) async fn handle_query_vault_state(&mut self) -> Result<VaultState, Error> {
use crate::actors::vault::GetState;
let vault_state = match self.props.actors.vault.ask(GetState {}).await {
Ok(state) => state,
Err(err) => {
error!(?err, actor = "client", "vault.query.failed");
return Err(Error::Internal);
}
};
Ok(vault_state)
}
#[message]
pub(crate) async fn handle_sign_transaction(
&mut self,
wallet_address: Address,
transaction: TxEip1559,
) -> Result<Signature, SignTransactionRpcError> {
match self
.props
.actors
.evm
.ask(ClientSignTransaction {
client_id: self.client_id,
wallet_address,
transaction,
})
.await
{
Ok(signature) => Ok(signature),
Err(kameo::error::SendError::HandlerError(SignTransactionError::Vet(vet_error))) => {
Err(SignTransactionRpcError::Vet(vet_error))
}
Err(err) => {
error!(?err, "Failed to sign EVM transaction in client session");
Err(SignTransactionRpcError::Internal)
}
}
}
}
impl Actor for ClientSession {
type Args = Self;
type Error = Error;
async fn on_start(
args: Self::Args,
this: kameo::prelude::ActorRef<Self>,
) -> Result<Self, Self::Error> {
args.props
.actors
.flow_coordinator
.ask(RegisterClient { actor: this })
.await
.map_err(|_| Error::ConnectionRegistrationFailed)?;
Ok(args)
}
}
impl ClientSession {
pub fn new_test(db: db::DatabasePool, actors: GlobalActors) -> Self {
let props = ClientConnection::new(db, actors);
Self {
props,
client_id: 0,
}
}
}
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("Connection registration failed")]
ConnectionRegistrationFailed,
#[error("Internal error")]
Internal,
}
#[derive(Debug, thiserror::Error)]
pub enum SignTransactionRpcError {
#[error("Policy evaluation failed")]
Vet(#[from] VetError),
#[error("Internal error")]
Internal,
}

View File

@@ -0,0 +1,2 @@
pub mod user_agent;
pub mod client;

View File

@@ -0,0 +1,106 @@
use arbiter_crypto::authn;
use arbiter_proto::transport::Bi;
use tracing::error;
mod state;
use state::*;
use super::UserAgentConnection;
#[derive(Debug, Clone)]
pub enum Inbound {
AuthChallengeRequest {
pubkey: authn::PublicKey,
bootstrap_token: Option<String>,
},
AuthChallengeSolution {
signature: Vec<u8>,
},
}
#[derive(Debug)]
pub enum Error {
UnregisteredPublicKey,
InvalidChallengeSolution,
InvalidBootstrapToken,
Internal { details: String },
Transport,
}
impl Error {
fn internal(details: impl Into<String>) -> Self {
Self::Internal {
details: details.into(),
}
}
}
impl From<diesel::result::Error> for Error {
fn from(e: diesel::result::Error) -> Self {
error!(?e, "Database error");
Self::internal("Database error")
}
}
#[derive(Debug, Clone)]
pub enum Outbound {
AuthChallenge { nonce: i32 },
AuthSuccess,
}
fn parse_auth_event(payload: Inbound) -> AuthEvents {
match payload {
Inbound::AuthChallengeRequest {
pubkey,
bootstrap_token: None,
} => AuthEvents::AuthRequest(ChallengeRequest { pubkey }),
Inbound::AuthChallengeRequest {
pubkey,
bootstrap_token: Some(token),
} => AuthEvents::BootstrapAuthRequest(BootstrapAuthRequest { pubkey, token }),
Inbound::AuthChallengeSolution { signature } => {
AuthEvents::ReceivedSolution(ChallengeSolution {
solution: signature,
})
}
}
}
pub async fn authenticate<T>(
props: &mut UserAgentConnection,
transport: T,
) -> Result<authn::PublicKey, Error>
where
T: Bi<Inbound, Result<Outbound, Error>> + Send,
{
let mut state = AuthStateMachine::new(AuthContext::new(props, transport));
loop {
// `state` holds a mutable reference to `props` so we can't access it directly here
let Some(payload) = state.context_mut().transport.recv().await else {
return Err(Error::Transport);
};
match state.process_event(parse_auth_event(payload)).await {
Ok(AuthStates::AuthOk(key)) => return Ok(key.clone()),
Err(AuthError::ActionFailed(err)) => {
error!(?err, "State machine action failed");
return Err(err);
}
Err(AuthError::GuardFailed(err)) => {
error!(?err, "State machine guard failed");
return Err(err);
}
Err(AuthError::InvalidEvent) => {
error!("Invalid event for current state");
return Err(Error::InvalidChallengeSolution);
}
Err(AuthError::TransitionsFailed) => {
error!("Invalid state transition");
return Err(Error::InvalidChallengeSolution);
}
_ => (),
}
}
}

View File

@@ -0,0 +1,319 @@
use arbiter_crypto::authn::{self, USERAGENT_CONTEXT};
use arbiter_proto::transport::Bi;
use diesel::{ExpressionMethods as _, OptionalExtension as _, QueryDsl, update};
use diesel_async::{AsyncConnection, RunQueryDsl};
use kameo::actor::ActorRef;
use tracing::error;
use super::super::{UserAgentCredentials, UserAgentConnection};
use super::Error;
use crate::peers::user_agent::auth::Outbound;
use crate::{
actors::{
bootstrap::ConsumeToken,
vault::Vault,
},
crypto::integrity,
db::{DatabasePool, schema::useragent_client},
};
pub struct ChallengeRequest {
pub pubkey: authn::PublicKey,
}
pub struct BootstrapAuthRequest {
pub pubkey: authn::PublicKey,
pub token: String,
}
pub struct ChallengeContext {
pub challenge_nonce: i32,
pub key: authn::PublicKey,
}
pub struct ChallengeSolution {
pub solution: Vec<u8>,
}
smlang::statemachine!(
name: Auth,
custom_error: true,
transitions: {
*Init + AuthRequest(ChallengeRequest) / async prepare_challenge = SentChallenge(ChallengeContext),
Init + BootstrapAuthRequest(BootstrapAuthRequest) / async verify_bootstrap_token = AuthOk(authn::PublicKey),
SentChallenge(ChallengeContext) + ReceivedSolution(ChallengeSolution) / async verify_solution = AuthOk(authn::PublicKey),
}
);
/// Returns the current nonce, ready to use for the challenge nonce.
async fn get_current_nonce_and_id(
db: &DatabasePool,
key: &authn::PublicKey,
) -> Result<(i32, i32), Error> {
let mut db_conn = db.get().await.map_err(|e| {
error!(error = ?e, "Database pool error");
Error::internal("Database unavailable")
})?;
db_conn
.exclusive_transaction(|conn| {
Box::pin(async move {
useragent_client::table
.filter(useragent_client::public_key.eq(key.to_bytes()))
.select((useragent_client::id, useragent_client::nonce))
.first::<(i32, i32)>(conn)
.await
})
})
.await
.optional()
.map_err(|e| {
error!(error = ?e, "Database error");
Error::internal("Database operation failed")
})?
.ok_or_else(|| {
error!(?key, "Public key not found in database");
Error::UnregisteredPublicKey
})
}
async fn verify_integrity(
db: &DatabasePool,
vault: &ActorRef<Vault>,
pubkey: &authn::PublicKey,
) -> Result<(), Error> {
let mut db_conn = db.get().await.map_err(|e| {
error!(error = ?e, "Database pool error");
Error::internal("Database unavailable")
})?;
let (id, nonce) = get_current_nonce_and_id(db, pubkey).await?;
let _result = integrity::verify_entity(
&mut db_conn,
vault,
&UserAgentCredentials {
pubkey: pubkey.clone(),
nonce,
},
id,
)
.await
.map_err(|e| {
error!(?e, "Integrity verification failed");
Error::internal("Integrity verification failed")
})?;
Ok(())
}
async fn create_nonce(
db: &DatabasePool,
vault: &ActorRef<Vault>,
pubkey: &authn::PublicKey,
) -> Result<i32, Error> {
let mut db_conn = db.get().await.map_err(|e| {
error!(error = ?e, "Database pool error");
Error::internal("Database unavailable")
})?;
let new_nonce = db_conn
.exclusive_transaction(|conn| {
Box::pin(async move {
let (id, new_nonce): (i32, i32) = update(useragent_client::table)
.filter(useragent_client::public_key.eq(pubkey.to_bytes()))
.set(useragent_client::nonce.eq(useragent_client::nonce + 1))
.returning((useragent_client::id, useragent_client::nonce))
.get_result(conn)
.await
.map_err(|e| {
error!(error = ?e, "Database error");
Error::internal("Database operation failed")
})?;
integrity::sign_entity(
conn,
vault,
&UserAgentCredentials {
pubkey: pubkey.clone(),
nonce: new_nonce,
},
id,
)
.await
.map_err(|e| {
error!(?e, "Integrity signature update failed");
Error::internal("Database error")
})?;
Result::<_, Error>::Ok(new_nonce)
})
})
.await?;
Ok(new_nonce)
}
async fn register_key(
db: &DatabasePool,
vault: &ActorRef<Vault>,
pubkey: &authn::PublicKey,
) -> Result<(), Error> {
let pubkey_bytes = pubkey.to_bytes();
let mut conn = db.get().await.map_err(|e| {
error!(error = ?e, "Database pool error");
Error::internal("Database unavailable")
})?;
conn.transaction(|conn| {
Box::pin(async move {
const NONCE_START: i32 = 1;
let id: i32 = diesel::insert_into(useragent_client::table)
.values((
useragent_client::public_key.eq(pubkey_bytes),
useragent_client::nonce.eq(NONCE_START),
))
.returning(useragent_client::id)
.get_result(conn)
.await
.map_err(|e| {
error!(error = ?e, "Database error");
Error::internal("Database operation failed")
})?;
let entity = UserAgentCredentials {
pubkey: pubkey.clone(),
nonce: NONCE_START,
};
integrity::sign_entity(conn, vault, &entity, id)
.await
.map_err(|e| {
error!(error = ?e, "Failed to sign integrity tag for new user-agent key");
Error::internal("Failed to register public key")
})?;
Result::<_, Error>::Ok(())
})
})
.await?;
Ok(())
}
pub struct AuthContext<'a, T> {
pub(super) conn: &'a mut UserAgentConnection,
pub(super) transport: T,
}
impl<'a, T> AuthContext<'a, T> {
pub fn new(conn: &'a mut UserAgentConnection, transport: T) -> Self {
Self { conn, transport }
}
}
impl<T> AuthStateMachineContext for AuthContext<'_, T>
where
T: Bi<super::Inbound, Result<super::Outbound, Error>> + Send,
{
type Error = Error;
async fn prepare_challenge(
&mut self,
ChallengeRequest { pubkey }: ChallengeRequest,
) -> Result<ChallengeContext, Self::Error> {
verify_integrity(&self.conn.db, &self.conn.actors.vault, &pubkey).await?;
let nonce = create_nonce(&self.conn.db, &self.conn.actors.vault, &pubkey).await?;
self.transport
.send(Ok(Outbound::AuthChallenge { nonce }))
.await
.map_err(|e| {
error!(?e, "Failed to send auth challenge");
Error::Transport
})?;
Ok(ChallengeContext {
challenge_nonce: nonce,
key: pubkey,
})
}
#[allow(missing_docs)]
#[allow(clippy::result_unit_err)]
async fn verify_bootstrap_token(
&mut self,
BootstrapAuthRequest { pubkey, token }: BootstrapAuthRequest,
) -> Result<authn::PublicKey, Self::Error> {
let token_ok: bool = self
.conn
.actors
.bootstrapper
.ask(ConsumeToken {
token: token.clone(),
})
.await
.map_err(|e| {
error!(?e, "Failed to consume bootstrap token");
Error::internal("Failed to consume bootstrap token")
})?;
if !token_ok {
error!("Invalid bootstrap token provided");
return Err(Error::InvalidBootstrapToken);
}
match token_ok {
true => {
register_key(&self.conn.db, &self.conn.actors.vault, &pubkey).await?;
self.transport
.send(Ok(Outbound::AuthSuccess))
.await
.map_err(|_| Error::Transport)?;
Ok(pubkey)
}
false => {
error!("Invalid bootstrap token provided");
self.transport
.send(Err(Error::InvalidBootstrapToken))
.await
.map_err(|_| Error::Transport)?;
Err(Error::InvalidBootstrapToken)
}
}
}
#[allow(missing_docs)]
#[allow(clippy::unused_unit)]
async fn verify_solution(
&mut self,
ChallengeContext {
challenge_nonce,
key,
}: &ChallengeContext,
ChallengeSolution { solution }: ChallengeSolution,
) -> Result<authn::PublicKey, Self::Error> {
let signature = authn::Signature::try_from(solution.as_slice()).map_err(|_| {
error!("Failed to decode signature in challenge solution");
Error::InvalidChallengeSolution
})?;
let valid = key.verify(*challenge_nonce, USERAGENT_CONTEXT, &signature);
match valid {
true => {
self.transport
.send(Ok(Outbound::AuthSuccess))
.await
.map_err(|_| Error::Transport)?;
Ok(key.clone())
}
false => {
self.transport
.send(Err(Error::InvalidChallengeSolution))
.await
.map_err(|_| Error::Transport)?;
Err(Error::InvalidChallengeSolution)
}
}
}
}

View File

@@ -0,0 +1,55 @@
use crate::{
actors::GlobalActors,
crypto::integrity::Integrable,
db, peers::client::ClientProfile,
};
use arbiter_crypto::authn;
#[derive(Debug)]
pub struct UserAgentCredentials {
pub pubkey: authn::PublicKey,
pub nonce: i32,
}
impl Integrable for UserAgentCredentials {
const KIND: &'static str = "useragent_credentials";
}
// Messages, sent by user agent to connection client without having a request
#[derive(Debug)]
pub enum OutOfBand {
ClientConnectionRequest { profile: ClientProfile },
ClientConnectionCancel { pubkey: authn::PublicKey },
}
pub struct UserAgentConnection {
pub(crate) db: db::DatabasePool,
pub(crate) actors: GlobalActors,
}
impl UserAgentConnection {
pub fn new(db: db::DatabasePool, actors: GlobalActors) -> Self {
Self { db, actors }
}
}
pub mod auth;
pub mod session;
pub use auth::authenticate;
pub use session::UserAgentSession;
use crate::crypto::integrity::hashing::Hashable;
impl Hashable for authn::PublicKey {
fn hash<H: sha2::Digest>(&self, hasher: &mut H) {
hasher.update(self.to_bytes());
}
}
impl Hashable for UserAgentCredentials {
fn hash<H: sha2::Digest>(&self, hasher: &mut H) {
self.pubkey.hash(hasher);
self.nonce.hash(hasher);
}
}

View File

@@ -0,0 +1,189 @@
use arbiter_crypto::authn;
use std::{borrow::Cow, collections::HashMap};
use arbiter_proto::transport::Sender;
use async_trait::async_trait;
use kameo::{Actor, actor::ActorRef, messages};
use thiserror::Error;
use tracing::error;
use crate::{actors::flow_coordinator::{RegisterUserAgent, client_connect_approval::ClientApprovalController}, peers::client::ClientProfile};
mod state;
use state::{DummyContext, UserAgentEvents, UserAgentStateMachine};
use super::{OutOfBand, UserAgentConnection};
#[derive(Debug, Error)]
pub enum Error {
#[error("State transition failed")]
State,
#[error("Internal error: {message}")]
Internal { message: Cow<'static, str> },
}
impl From<crate::db::PoolError> for Error {
fn from(err: crate::db::PoolError) -> Self {
error!(?err, "Database pool error");
Self::internal("Database pool error")
}
}
impl From<diesel::result::Error> for Error {
fn from(err: diesel::result::Error) -> Self {
error!(?err, "Database error");
Self::internal("Database error")
}
}
impl Error {
pub fn internal(message: impl Into<Cow<'static, str>>) -> Self {
Self::Internal {
message: message.into(),
}
}
}
pub struct PendingClientApproval {
pubkey: authn::PublicKey,
controller: ActorRef<ClientApprovalController>,
}
pub struct UserAgentSession {
props: UserAgentConnection,
state: UserAgentStateMachine<DummyContext>,
sender: Box<dyn Sender<OutOfBand>>,
pending_client_approvals: HashMap<Vec<u8>, PendingClientApproval>,
}
pub mod connection;
impl UserAgentSession {
pub(crate) fn new(props: UserAgentConnection, sender: Box<dyn Sender<OutOfBand>>) -> Self {
Self {
props,
state: UserAgentStateMachine::new(DummyContext),
sender,
pending_client_approvals: Default::default(),
}
}
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");
Error::State
})?;
Ok(())
}
}
#[messages]
impl UserAgentSession {
#[message]
pub async fn begin_new_client_approval(
&mut self,
client: ClientProfile,
controller: ActorRef<ClientApprovalController>,
) {
if let Err(e) = self
.sender
.send(OutOfBand::ClientConnectionRequest {
profile: client.clone(),
})
.await
{
error!(
?e,
actor = "user_agent",
event = "failed to announce new client connection"
);
return;
}
self.pending_client_approvals.insert(
client.pubkey.to_bytes(),
PendingClientApproval {
pubkey: client.pubkey,
controller,
},
);
}
}
impl Actor for UserAgentSession {
type Args = Self;
type Error = Error;
async fn on_start(
args: Self::Args,
this: kameo::prelude::ActorRef<Self>,
) -> Result<Self, Self::Error> {
args.props
.actors
.flow_coordinator
.ask(RegisterUserAgent {
actor: this.clone(),
})
.await
.map_err(|err| {
error!(
?err,
"Failed to register user agent connection with flow coordinator"
);
Error::internal("Failed to register user agent connection with flow coordinator")
})?;
Ok(args)
}
async fn on_link_died(
&mut self,
_: kameo::prelude::WeakActorRef<Self>,
id: kameo::prelude::ActorId,
_: kameo::prelude::ActorStopReason,
) -> Result<std::ops::ControlFlow<kameo::prelude::ActorStopReason>, Self::Error> {
let cancelled_pubkey = self
.pending_client_approvals
.iter()
.find_map(|(k, v)| (v.controller.id() == id).then_some(k.clone()));
if let Some(pubkey_bytes) = cancelled_pubkey {
let Some(approval) = self.pending_client_approvals.remove(&pubkey_bytes) else {
return Ok(std::ops::ControlFlow::Continue(()));
};
if let Err(e) = self
.sender
.send(OutOfBand::ClientConnectionCancel {
pubkey: approval.pubkey,
})
.await
{
error!(
?e,
actor = "user_agent",
event = "failed to announce client connection cancellation"
);
}
}
Ok(std::ops::ControlFlow::Continue(()))
}
}

View File

@@ -0,0 +1,523 @@
use std::sync::Mutex;
use alloy::{consensus::TxEip1559, primitives::Address, signers::Signature};
use arbiter_crypto::{
authn,
safecell::{SafeCell, SafeCellHandle as _},
};
use chacha20poly1305::{AeadInPlace, XChaCha20Poly1305, XNonce, aead::KeyInit};
use diesel::{ExpressionMethods as _, QueryDsl as _, SelectableHelper};
use diesel_async::{AsyncConnection, RunQueryDsl};
use kameo::error::SendError;
use kameo::messages;
use kameo::prelude::Context;
use tracing::{error, info};
use x25519_dalek::{EphemeralSecret, PublicKey};
use crate::{actors::vault::VaultState, peers::user_agent::session::state::{UnsealContext, UserAgentEvents}};
use crate::actors::{
evm::{
ClientSignTransaction, Generate, ListWallets, SignTransactionError as EvmSignError,
UseragentCreateGrant, UseragentListGrants,
},
vault::{self, Bootstrap, TryUnseal},
};
use crate::db::models::{
EvmWalletAccess, NewEvmWalletAccess, ProgramClient, ProgramClientMetadata,
};
use crate::evm::policies::{Grant, SpecificGrant};
use crate::{
actors::flow_coordinator::client_connect_approval::ClientApprovalAnswer,
};
use super::{UserAgentSession, state, Error};
impl UserAgentSession {
fn take_unseal_secret(&mut self) -> Result<(EphemeralSecret, PublicKey), Error> {
let state::UserAgentStates::WaitingForUnsealKey(unseal_context) = self.state.state() else {
error!("Received encrypted key in invalid state");
return Err(Error::internal("Invalid state for unseal encrypted key"));
};
let ephemeral_secret = {
#[allow(
clippy::unwrap_used,
reason = "Mutex poison is unrecoverable and should panic"
)]
let mut secret_lock = unseal_context.secret.lock().unwrap();
let secret = secret_lock.take();
match secret {
Some(secret) => secret,
None => {
drop(secret_lock);
error!("Ephemeral secret already taken");
return Err(Error::internal("Ephemeral secret already taken"));
}
}
};
Ok((ephemeral_secret, unseal_context.client_public_key))
}
fn decrypt_client_key_material(
ephemeral_secret: EphemeralSecret,
client_public_key: PublicKey,
nonce: &[u8],
ciphertext: &[u8],
associated_data: &[u8],
) -> Result<SafeCell<Vec<u8>>, ()> {
let nonce = XNonce::from_slice(nonce);
let shared_secret = ephemeral_secret.diffie_hellman(&client_public_key);
let cipher = XChaCha20Poly1305::new(shared_secret.as_bytes().into());
let mut key_buffer = SafeCell::new(ciphertext.to_vec());
let decryption_result = key_buffer.write_inline(|write_handle| {
cipher.decrypt_in_place(nonce, associated_data, write_handle)
});
match decryption_result {
Ok(_) => Ok(key_buffer),
Err(err) => {
error!(?err, "Failed to decrypt encrypted key material");
Err(())
}
}
}
}
pub struct UnsealStartResponse {
pub server_pubkey: PublicKey,
}
#[derive(Debug, Error)]
pub enum UnsealError {
#[error("Invalid key provided for unsealing")]
InvalidKey,
#[error("Internal error during unsealing process")]
General(#[from] super::Error),
}
#[derive(Debug, Error)]
pub enum BootstrapError {
#[error("Invalid key provided for bootstrapping")]
InvalidKey,
#[error("Vault is already bootstrapped")]
AlreadyBootstrapped,
#[error("Internal error during bootstrapping process")]
General(#[from] super::Error),
}
#[derive(Debug, Error)]
pub enum SignTransactionError {
#[error("Policy evaluation failed")]
Vet(#[from] crate::evm::VetError),
#[error("Internal signing error")]
Internal,
}
#[derive(Debug, Error)]
pub enum GrantMutationError {
#[error("Vault is sealed")]
VaultSealed,
#[error("Internal grant mutation error")]
Internal,
}
#[messages]
impl UserAgentSession {
#[message]
pub async fn handle_unseal_request(
&mut self,
client_pubkey: x25519_dalek::PublicKey,
) -> Result<UnsealStartResponse, Error> {
let secret = EphemeralSecret::random();
let public_key = PublicKey::from(&secret);
self.transition(UserAgentEvents::UnsealRequest(UnsealContext {
secret: Mutex::new(Some(secret)),
client_public_key: client_pubkey,
}))?;
Ok(UnsealStartResponse {
server_pubkey: public_key,
})
}
#[message]
pub async fn handle_unseal_encrypted_key(
&mut self,
nonce: Vec<u8>,
ciphertext: Vec<u8>,
associated_data: Vec<u8>,
) -> Result<(), UnsealError> {
let (ephemeral_secret, client_public_key) = match self.take_unseal_secret() {
Ok(values) => values,
Err(Error::State) => {
self.transition(UserAgentEvents::ReceivedInvalidKey)?;
return Err(UnsealError::InvalidKey);
}
Err(_err) => {
return Err(Error::internal("Failed to take unseal secret").into());
}
};
let seal_key_buffer = match Self::decrypt_client_key_material(
ephemeral_secret,
client_public_key,
&nonce,
&ciphertext,
&associated_data,
) {
Ok(buffer) => buffer,
Err(()) => {
self.transition(UserAgentEvents::ReceivedInvalidKey)?;
return Err(UnsealError::InvalidKey);
}
};
match self
.props
.actors
.vault
.ask(TryUnseal {
seal_key_raw: seal_key_buffer,
})
.await
{
Ok(_) => {
info!("Successfully unsealed key with client-provided key");
self.transition(UserAgentEvents::ReceivedValidKey)?;
Ok(())
}
Err(SendError::HandlerError(vault::Error::InvalidKey)) => {
self.transition(UserAgentEvents::ReceivedInvalidKey)?;
Err(UnsealError::InvalidKey)
}
Err(SendError::HandlerError(err)) => {
error!(?err, "Vault failed to unseal key");
self.transition(UserAgentEvents::ReceivedInvalidKey)?;
Err(UnsealError::InvalidKey)
}
Err(err) => {
error!(?err, "Failed to send unseal request to vault");
self.transition(UserAgentEvents::ReceivedInvalidKey)?;
Err(Error::internal("Vault actor error").into())
}
}
}
#[message]
pub(crate) async fn handle_bootstrap_encrypted_key(
&mut self,
nonce: Vec<u8>,
ciphertext: Vec<u8>,
associated_data: Vec<u8>,
) -> Result<(), BootstrapError> {
let (ephemeral_secret, client_public_key) = match self.take_unseal_secret() {
Ok(values) => values,
Err(Error::State) => {
self.transition(UserAgentEvents::ReceivedInvalidKey)?;
return Err(BootstrapError::InvalidKey);
}
Err(err) => return Err(err.into()),
};
let seal_key_buffer = match Self::decrypt_client_key_material(
ephemeral_secret,
client_public_key,
&nonce,
&ciphertext,
&associated_data,
) {
Ok(buffer) => buffer,
Err(()) => {
self.transition(UserAgentEvents::ReceivedInvalidKey)?;
return Err(BootstrapError::InvalidKey);
}
};
match self
.props
.actors
.vault
.ask(Bootstrap {
seal_key_raw: seal_key_buffer,
})
.await
{
Ok(_) => {
info!("Successfully bootstrapped vault with client-provided key");
self.transition(UserAgentEvents::ReceivedValidKey)?;
Ok(())
}
Err(SendError::HandlerError(vault::Error::AlreadyBootstrapped)) => {
self.transition(UserAgentEvents::ReceivedInvalidKey)?;
Err(BootstrapError::AlreadyBootstrapped)
}
Err(SendError::HandlerError(err)) => {
error!(?err, "Vault failed to bootstrap vault");
self.transition(UserAgentEvents::ReceivedInvalidKey)?;
Err(BootstrapError::InvalidKey)
}
Err(err) => {
error!(?err, "Failed to send bootstrap request to vault");
self.transition(UserAgentEvents::ReceivedInvalidKey)?;
Err(BootstrapError::General(Error::internal(
"Vault actor error",
)))
}
}
}
}
#[messages]
impl UserAgentSession {
#[message]
pub(crate) async fn handle_query_vault_state(&mut self) -> Result<VaultState, Error> {
use crate::actors::vault::GetState;
let vault_state = match self.props.actors.vault.ask(GetState {}).await {
Ok(state) => state,
Err(err) => {
error!(?err, actor = "useragent", "vault.query.failed");
return Err(Error::internal("Vault is in broken state"));
}
};
Ok(vault_state)
}
}
#[messages]
impl UserAgentSession {
#[message]
pub(crate) async fn handle_evm_wallet_create(&mut self) -> Result<(i32, Address), Error> {
match self.props.actors.evm.ask(Generate {}).await {
Ok(address) => Ok(address),
Err(SendError::HandlerError(err)) => Err(Error::internal(format!(
"EVM wallet generation failed: {err}"
))),
Err(err) => {
error!(?err, "EVM actor unreachable during wallet create");
Err(Error::internal("EVM actor unreachable"))
}
}
}
#[message]
pub(crate) async fn handle_evm_wallet_list(&mut self) -> Result<Vec<(i32, Address)>, Error> {
match self.props.actors.evm.ask(ListWallets {}).await {
Ok(wallets) => Ok(wallets),
Err(err) => {
error!(?err, "EVM wallet list failed");
Err(Error::internal("Failed to list EVM wallets"))
}
}
}
}
#[messages]
impl UserAgentSession {
#[message]
pub(crate) async fn handle_grant_list(&mut self) -> Result<Vec<Grant<SpecificGrant>>, Error> {
match self.props.actors.evm.ask(UseragentListGrants {}).await {
Ok(grants) => Ok(grants),
Err(err) => {
error!(?err, "EVM grant list failed");
Err(Error::internal("Failed to list EVM grants"))
}
}
}
#[message]
pub(crate) async fn handle_grant_create(
&mut self,
basic: crate::evm::policies::SharedGrantSettings,
grant: crate::evm::policies::SpecificGrant,
) -> Result<i32, GrantMutationError> {
match self
.props
.actors
.evm
.ask(UseragentCreateGrant { basic, grant })
.await
{
Ok(grant_id) => Ok(grant_id),
Err(err) => {
error!(?err, "EVM grant create failed");
Err(GrantMutationError::Internal)
}
}
}
#[message]
pub(crate) async fn handle_grant_delete(
&mut self,
grant_id: i32,
) -> Result<(), GrantMutationError> {
// match self
// .props
// .actors
// .evm
// .ask(UseragentDeleteGrant { grant_id })
// .await
// {
// Ok(()) => Ok(()),
// Err(err) => {
// error!(?err, "EVM grant delete failed");
// Err(GrantMutationError::Internal)
// }
// }
let _ = grant_id;
todo!()
}
#[message]
pub(crate) async fn handle_sign_transaction(
&mut self,
client_id: i32,
wallet_address: Address,
transaction: TxEip1559,
) -> Result<Signature, SignTransactionError> {
match self
.props
.actors
.evm
.ask(ClientSignTransaction {
client_id,
wallet_address,
transaction,
})
.await
{
Ok(signature) => Ok(signature),
Err(SendError::HandlerError(EvmSignError::Vet(vet_error))) => {
Err(SignTransactionError::Vet(vet_error))
}
Err(err) => {
error!(?err, "EVM sign transaction failed in user-agent session");
Err(SignTransactionError::Internal)
}
}
}
#[message]
pub(crate) async fn handle_grant_evm_wallet_access(
&mut self,
entries: Vec<NewEvmWalletAccess>,
) -> Result<(), Error> {
let mut conn = self.props.db.get().await?;
conn.transaction(|conn| {
Box::pin(async move {
use crate::db::schema::evm_wallet_access;
for entry in entries {
diesel::insert_into(evm_wallet_access::table)
.values(&entry)
.on_conflict_do_nothing()
.execute(conn)
.await?;
}
Result::<_, Error>::Ok(())
})
})
.await?;
Ok(())
}
#[message]
pub(crate) async fn handle_revoke_evm_wallet_access(
&mut self,
entries: Vec<i32>,
) -> Result<(), Error> {
let mut conn = self.props.db.get().await?;
conn.transaction(|conn| {
Box::pin(async move {
use crate::db::schema::evm_wallet_access;
for entry in entries {
diesel::delete(evm_wallet_access::table)
.filter(evm_wallet_access::wallet_id.eq(entry))
.execute(conn)
.await?;
}
Result::<_, Error>::Ok(())
})
})
.await?;
Ok(())
}
#[message]
pub(crate) async fn handle_list_wallet_access(
&mut self,
) -> Result<Vec<EvmWalletAccess>, Error> {
let mut conn = self.props.db.get().await?;
use crate::db::schema::evm_wallet_access;
let access_entries = evm_wallet_access::table
.select(EvmWalletAccess::as_select())
.load::<_>(&mut conn)
.await?;
Ok(access_entries)
}
}
#[messages]
impl UserAgentSession {
#[message(ctx)]
pub(crate) async fn handle_new_client_approve(
&mut self,
approved: bool,
pubkey: authn::PublicKey,
ctx: &mut Context<Self, Result<(), Error>>,
) -> Result<(), Error> {
let pending_approval = match self.pending_client_approvals.remove(&pubkey.to_bytes()) {
Some(approval) => approval,
None => {
error!("Received client connection response for unknown client");
return Err(Error::internal("Unknown client in connection response"));
}
};
pending_approval
.controller
.tell(ClientApprovalAnswer { approved })
.await
.map_err(|err| {
error!(
?err,
"Failed to send client approval response to controller"
);
Error::internal("Failed to send client approval response to controller")
})?;
ctx.actor_ref().unlink(&pending_approval.controller).await;
Ok(())
}
#[message]
pub(crate) async fn handle_sdk_client_list(
&mut self,
) -> Result<Vec<(ProgramClient, ProgramClientMetadata)>, Error> {
use crate::db::schema::{client_metadata, program_client};
let mut conn = self.props.db.get().await?;
let clients = program_client::table
.inner_join(client_metadata::table)
.select((
ProgramClient::as_select(),
ProgramClientMetadata::as_select(),
))
.load::<(ProgramClient, ProgramClientMetadata)>(&mut conn)
.await?;
Ok(clients)
}
}

View File

@@ -0,0 +1,27 @@
use std::sync::Mutex;
use x25519_dalek::{EphemeralSecret, PublicKey};
pub struct UnsealContext {
pub client_public_key: PublicKey,
pub secret: Mutex<Option<EphemeralSecret>>,
}
smlang::statemachine!(
name: UserAgent,
custom_error: false,
transitions: {
*Idle + UnsealRequest(UnsealContext) / generate_temp_keypair = WaitingForUnsealKey(UnsealContext),
WaitingForUnsealKey(UnsealContext) + ReceivedValidKey = Unsealed,
WaitingForUnsealKey(UnsealContext) + ReceivedInvalidKey = Idle,
}
);
pub struct DummyContext;
impl UserAgentStateMachineContext for DummyContext {
#[allow(missing_docs)]
#[allow(clippy::unused_unit)]
fn generate_temp_keypair(&mut self, event_data: UnsealContext) -> Result<UnsealContext, ()> {
Ok(event_data)
}
}