466 lines
19 KiB
Rust
466 lines
19 KiB
Rust
use tokio::sync::mpsc;
|
|
|
|
use arbiter_proto::{
|
|
proto::{
|
|
client::ClientInfo as ProtoClientMetadata,
|
|
evm::{
|
|
EvmError as ProtoEvmError, EvmGrantCreateRequest, EvmGrantCreateResponse,
|
|
EvmGrantDeleteRequest, EvmGrantDeleteResponse, EvmGrantList, EvmGrantListResponse,
|
|
GrantEntry, WalletCreateResponse, WalletEntry, WalletList, WalletListResponse,
|
|
evm_grant_create_response::Result as EvmGrantCreateResult,
|
|
evm_grant_delete_response::Result as EvmGrantDeleteResult,
|
|
evm_grant_list_response::Result as EvmGrantListResult,
|
|
wallet_create_response::Result as WalletCreateResult,
|
|
wallet_list_response::Result as WalletListResult,
|
|
},
|
|
user_agent::{
|
|
BootstrapEncryptedKey as ProtoBootstrapEncryptedKey,
|
|
BootstrapResult as ProtoBootstrapResult, ListWalletAccessResponse,
|
|
SdkClientConnectionCancel as ProtoSdkClientConnectionCancel,
|
|
SdkClientConnectionRequest as ProtoSdkClientConnectionRequest,
|
|
SdkClientEntry as ProtoSdkClientEntry, SdkClientError as ProtoSdkClientError,
|
|
SdkClientGrantWalletAccess, SdkClientList as ProtoSdkClientList,
|
|
SdkClientListResponse as ProtoSdkClientListResponse, SdkClientRevokeWalletAccess,
|
|
SdkClientWalletAccess, UnsealEncryptedKey as ProtoUnsealEncryptedKey,
|
|
UnsealResult as ProtoUnsealResult, UnsealStart, UserAgentRequest, UserAgentResponse,
|
|
VaultState as ProtoVaultState,
|
|
sdk_client_list_response::Result as ProtoSdkClientListResult,
|
|
user_agent_request::Payload as UserAgentRequestPayload,
|
|
user_agent_response::Payload as UserAgentResponsePayload,
|
|
},
|
|
},
|
|
transport::{Error as TransportError, Receiver, Sender, grpc::GrpcBi},
|
|
};
|
|
use async_trait::async_trait;
|
|
use kameo::{
|
|
actor::{ActorRef, Spawn as _},
|
|
error::SendError,
|
|
};
|
|
use tonic::Status;
|
|
use tracing::{error, info, warn};
|
|
|
|
use crate::{
|
|
actors::{
|
|
keyholder::KeyHolderState,
|
|
user_agent::{
|
|
OutOfBand, UserAgentConnection, UserAgentSession,
|
|
session::connection::{
|
|
BootstrapError, HandleBootstrapEncryptedKey, HandleEvmWalletCreate,
|
|
HandleEvmWalletList, HandleGrantCreate, HandleGrantDelete,
|
|
HandleGrantEvmWalletAccess, HandleGrantList, HandleListWalletAccess,
|
|
HandleNewClientApprove, HandleQueryVaultState, HandleRevokeEvmWalletAccess,
|
|
HandleSdkClientList, HandleUnsealEncryptedKey, HandleUnsealRequest, UnsealError,
|
|
},
|
|
},
|
|
},
|
|
db::models::{CoreEvmWalletAccess, NewEvmWalletAccess},
|
|
grpc::{Convert, TryConvert, request_tracker::RequestTracker},
|
|
};
|
|
mod auth;
|
|
mod inbound;
|
|
mod outbound;
|
|
|
|
pub struct OutOfBandAdapter(mpsc::Sender<OutOfBand>);
|
|
|
|
#[async_trait]
|
|
impl Sender<OutOfBand> for OutOfBandAdapter {
|
|
async fn send(&mut self, item: OutOfBand) -> Result<(), TransportError> {
|
|
self.0.send(item).await.map_err(|e| {
|
|
warn!(error = ?e, "Failed to send out-of-band message");
|
|
TransportError::ChannelClosed
|
|
})
|
|
}
|
|
}
|
|
|
|
async fn dispatch_loop(
|
|
mut bi: GrpcBi<UserAgentRequest, UserAgentResponse>,
|
|
actor: ActorRef<UserAgentSession>,
|
|
mut receiver: mpsc::Receiver<OutOfBand>,
|
|
mut request_tracker: RequestTracker,
|
|
) {
|
|
loop {
|
|
tokio::select! {
|
|
oob = receiver.recv() => {
|
|
let Some(oob) = oob else {
|
|
warn!("Out-of-band message channel closed");
|
|
return;
|
|
};
|
|
|
|
let payload = match oob {
|
|
OutOfBand::ClientConnectionRequest { profile } => {
|
|
UserAgentResponsePayload::SdkClientConnectionRequest(ProtoSdkClientConnectionRequest {
|
|
pubkey: profile.pubkey.to_bytes().to_vec(),
|
|
info: Some(ProtoClientMetadata {
|
|
name: profile.metadata.name,
|
|
description: profile.metadata.description,
|
|
version: profile.metadata.version,
|
|
}),
|
|
})
|
|
}
|
|
OutOfBand::ClientConnectionCancel { pubkey } => {
|
|
UserAgentResponsePayload::SdkClientConnectionCancel(ProtoSdkClientConnectionCancel {
|
|
pubkey: pubkey.to_bytes().to_vec(),
|
|
})
|
|
}
|
|
};
|
|
|
|
if bi.send(Ok(UserAgentResponse { id: None, payload: Some(payload) })).await.is_err() {
|
|
return;
|
|
}
|
|
}
|
|
|
|
message = bi.recv() => {
|
|
let Some(message) = message else { return; };
|
|
|
|
let conn = match message {
|
|
Ok(conn) => conn,
|
|
Err(err) => {
|
|
warn!(error = ?err, "Failed to receive user agent request");
|
|
return;
|
|
}
|
|
};
|
|
|
|
let request_id = match request_tracker.request(conn.id) {
|
|
Ok(id) => id,
|
|
Err(err) => {
|
|
let _ = bi.send(Err(err)).await;
|
|
return;
|
|
}
|
|
};
|
|
|
|
let Some(payload) = conn.payload else {
|
|
let _ = bi.send(Err(Status::invalid_argument("Missing user-agent request payload"))).await;
|
|
return;
|
|
};
|
|
|
|
match dispatch_inner(&actor, payload).await {
|
|
Ok(Some(response)) => {
|
|
if bi.send(Ok(UserAgentResponse {
|
|
id: Some(request_id),
|
|
payload: Some(response),
|
|
})).await.is_err() {
|
|
return;
|
|
}
|
|
}
|
|
Ok(None) => {}
|
|
Err(status) => {
|
|
error!(?status, "Failed to process user agent request");
|
|
let _ = bi.send(Err(status)).await;
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn dispatch_inner(
|
|
actor: &ActorRef<UserAgentSession>,
|
|
payload: UserAgentRequestPayload,
|
|
) -> Result<Option<UserAgentResponsePayload>, Status> {
|
|
let response = match payload {
|
|
UserAgentRequestPayload::UnsealStart(UnsealStart { client_pubkey }) => {
|
|
let client_pubkey = <[u8; 32]>::try_from(client_pubkey)
|
|
.map(x25519_dalek::PublicKey::from)
|
|
.map_err(|_| Status::invalid_argument("Invalid X25519 public key"))?;
|
|
|
|
let response = actor
|
|
.ask(HandleUnsealRequest { client_pubkey })
|
|
.await
|
|
.map_err(|err| {
|
|
warn!(error = ?err, "Failed to handle unseal start request");
|
|
Status::internal("Failed to start unseal flow")
|
|
})?;
|
|
|
|
UserAgentResponsePayload::UnsealStartResponse(
|
|
arbiter_proto::proto::user_agent::UnsealStartResponse {
|
|
server_pubkey: response.server_pubkey.as_bytes().to_vec(),
|
|
},
|
|
)
|
|
}
|
|
|
|
UserAgentRequestPayload::UnsealEncryptedKey(ProtoUnsealEncryptedKey {
|
|
nonce,
|
|
ciphertext,
|
|
associated_data,
|
|
}) => {
|
|
let result = match actor
|
|
.ask(HandleUnsealEncryptedKey {
|
|
nonce,
|
|
ciphertext,
|
|
associated_data,
|
|
})
|
|
.await
|
|
{
|
|
Ok(()) => ProtoUnsealResult::Success,
|
|
Err(SendError::HandlerError(UnsealError::InvalidKey)) => {
|
|
ProtoUnsealResult::InvalidKey
|
|
}
|
|
Err(err) => {
|
|
warn!(error = ?err, "Failed to handle unseal request");
|
|
return Err(Status::internal("Failed to unseal vault"));
|
|
}
|
|
};
|
|
UserAgentResponsePayload::UnsealResult(result.into())
|
|
}
|
|
|
|
UserAgentRequestPayload::BootstrapEncryptedKey(ProtoBootstrapEncryptedKey {
|
|
nonce,
|
|
ciphertext,
|
|
associated_data,
|
|
}) => {
|
|
let result = match actor
|
|
.ask(HandleBootstrapEncryptedKey {
|
|
nonce,
|
|
ciphertext,
|
|
associated_data,
|
|
})
|
|
.await
|
|
{
|
|
Ok(()) => ProtoBootstrapResult::Success,
|
|
Err(SendError::HandlerError(BootstrapError::InvalidKey)) => {
|
|
ProtoBootstrapResult::InvalidKey
|
|
}
|
|
Err(SendError::HandlerError(BootstrapError::AlreadyBootstrapped)) => {
|
|
ProtoBootstrapResult::AlreadyBootstrapped
|
|
}
|
|
Err(err) => {
|
|
warn!(error = ?err, "Failed to handle bootstrap request");
|
|
return Err(Status::internal("Failed to bootstrap vault"));
|
|
}
|
|
};
|
|
UserAgentResponsePayload::BootstrapResult(result.into())
|
|
}
|
|
|
|
UserAgentRequestPayload::QueryVaultState(_) => {
|
|
let state = match actor.ask(HandleQueryVaultState {}).await {
|
|
Ok(KeyHolderState::Unbootstrapped) => ProtoVaultState::Unbootstrapped,
|
|
Ok(KeyHolderState::Sealed) => ProtoVaultState::Sealed,
|
|
Ok(KeyHolderState::Unsealed) => ProtoVaultState::Unsealed,
|
|
Err(err) => {
|
|
warn!(error = ?err, "Failed to query vault state");
|
|
ProtoVaultState::Error
|
|
}
|
|
};
|
|
UserAgentResponsePayload::VaultState(state.into())
|
|
}
|
|
|
|
UserAgentRequestPayload::EvmWalletCreate(_) => {
|
|
let result = match actor.ask(HandleEvmWalletCreate {}).await {
|
|
Ok((wallet_id, address)) => WalletCreateResult::Wallet(WalletEntry {
|
|
id: wallet_id,
|
|
address: address.to_vec(),
|
|
}),
|
|
Err(err) => {
|
|
warn!(error = ?err, "Failed to create EVM wallet");
|
|
WalletCreateResult::Error(ProtoEvmError::Internal.into())
|
|
}
|
|
};
|
|
UserAgentResponsePayload::EvmWalletCreate(WalletCreateResponse {
|
|
result: Some(result),
|
|
})
|
|
}
|
|
|
|
UserAgentRequestPayload::EvmWalletList(_) => {
|
|
let result = match actor.ask(HandleEvmWalletList {}).await {
|
|
Ok(wallets) => WalletListResult::Wallets(WalletList {
|
|
wallets: wallets
|
|
.into_iter()
|
|
.map(|(id, address)| WalletEntry {
|
|
address: address.to_vec(),
|
|
id,
|
|
})
|
|
.collect(),
|
|
}),
|
|
Err(err) => {
|
|
warn!(error = ?err, "Failed to list EVM wallets");
|
|
WalletListResult::Error(ProtoEvmError::Internal.into())
|
|
}
|
|
};
|
|
UserAgentResponsePayload::EvmWalletList(WalletListResponse {
|
|
result: Some(result),
|
|
})
|
|
}
|
|
|
|
UserAgentRequestPayload::EvmGrantList(_) => {
|
|
let result = match actor.ask(HandleGrantList {}).await {
|
|
Ok(grants) => EvmGrantListResult::Grants(EvmGrantList {
|
|
grants: grants
|
|
.into_iter()
|
|
.map(|grant| GrantEntry {
|
|
id: grant.id,
|
|
wallet_access_id: grant.shared.wallet_access_id,
|
|
shared: Some(grant.shared.convert()),
|
|
specific: Some(grant.settings.convert()),
|
|
})
|
|
.collect(),
|
|
}),
|
|
Err(err) => {
|
|
warn!(error = ?err, "Failed to list EVM grants");
|
|
EvmGrantListResult::Error(ProtoEvmError::Internal.into())
|
|
}
|
|
};
|
|
UserAgentResponsePayload::EvmGrantList(EvmGrantListResponse {
|
|
result: Some(result),
|
|
})
|
|
}
|
|
|
|
UserAgentRequestPayload::EvmGrantCreate(EvmGrantCreateRequest { shared, specific }) => {
|
|
let basic = shared
|
|
.ok_or_else(|| Status::invalid_argument("Missing shared grant settings"))?
|
|
.try_convert()?;
|
|
let grant = specific
|
|
.ok_or_else(|| Status::invalid_argument("Missing specific grant settings"))?
|
|
.try_convert()?;
|
|
|
|
let result = match actor.ask(HandleGrantCreate { basic, grant }).await {
|
|
Ok(grant_id) => EvmGrantCreateResult::GrantId(grant_id),
|
|
Err(err) => {
|
|
warn!(error = ?err, "Failed to create EVM grant");
|
|
EvmGrantCreateResult::Error(ProtoEvmError::Internal.into())
|
|
}
|
|
};
|
|
UserAgentResponsePayload::EvmGrantCreate(EvmGrantCreateResponse {
|
|
result: Some(result),
|
|
})
|
|
}
|
|
|
|
UserAgentRequestPayload::EvmGrantDelete(EvmGrantDeleteRequest { grant_id }) => {
|
|
let result = match actor.ask(HandleGrantDelete { grant_id }).await {
|
|
Ok(()) => EvmGrantDeleteResult::Ok(()),
|
|
Err(err) => {
|
|
warn!(error = ?err, "Failed to delete EVM grant");
|
|
EvmGrantDeleteResult::Error(ProtoEvmError::Internal.into())
|
|
}
|
|
};
|
|
UserAgentResponsePayload::EvmGrantDelete(EvmGrantDeleteResponse {
|
|
result: Some(result),
|
|
})
|
|
}
|
|
|
|
UserAgentRequestPayload::SdkClientConnectionResponse(resp) => {
|
|
let pubkey_bytes = <[u8; 32]>::try_from(resp.pubkey)
|
|
.map_err(|_| Status::invalid_argument("Invalid Ed25519 public key length"))?;
|
|
let pubkey = ed25519_dalek::VerifyingKey::from_bytes(&pubkey_bytes)
|
|
.map_err(|_| Status::invalid_argument("Invalid Ed25519 public key"))?;
|
|
|
|
actor
|
|
.ask(HandleNewClientApprove {
|
|
approved: resp.approved,
|
|
pubkey,
|
|
})
|
|
.await
|
|
.map_err(|err| {
|
|
warn!(?err, "Failed to process client connection response");
|
|
Status::internal("Failed to process response")
|
|
})?;
|
|
|
|
return Ok(None);
|
|
}
|
|
|
|
UserAgentRequestPayload::SdkClientRevoke(_) => todo!(),
|
|
|
|
UserAgentRequestPayload::SdkClientList(_) => {
|
|
let result = match actor.ask(HandleSdkClientList {}).await {
|
|
Ok(clients) => ProtoSdkClientListResult::Clients(ProtoSdkClientList {
|
|
clients: clients
|
|
.into_iter()
|
|
.map(|(client, metadata)| ProtoSdkClientEntry {
|
|
id: client.id,
|
|
pubkey: client.public_key,
|
|
info: Some(ProtoClientMetadata {
|
|
name: metadata.name,
|
|
description: metadata.description,
|
|
version: metadata.version,
|
|
}),
|
|
created_at: client.created_at.0.timestamp() as i32,
|
|
})
|
|
.collect(),
|
|
}),
|
|
Err(err) => {
|
|
warn!(error = ?err, "Failed to list SDK clients");
|
|
ProtoSdkClientListResult::Error(ProtoSdkClientError::Internal.into())
|
|
}
|
|
};
|
|
UserAgentResponsePayload::SdkClientListResponse(ProtoSdkClientListResponse {
|
|
result: Some(result),
|
|
})
|
|
}
|
|
|
|
UserAgentRequestPayload::GrantWalletAccess(SdkClientGrantWalletAccess { accesses }) => {
|
|
let entries: Vec<NewEvmWalletAccess> =
|
|
accesses.into_iter().map(|a| a.convert()).collect();
|
|
|
|
match actor.ask(HandleGrantEvmWalletAccess { entries }).await {
|
|
Ok(()) => {
|
|
info!("Successfully granted wallet access");
|
|
return Ok(None);
|
|
}
|
|
Err(err) => {
|
|
warn!(error = ?err, "Failed to grant wallet access");
|
|
return Err(Status::internal("Failed to grant wallet access"));
|
|
}
|
|
}
|
|
}
|
|
|
|
UserAgentRequestPayload::RevokeWalletAccess(SdkClientRevokeWalletAccess { accesses }) => {
|
|
match actor.ask(HandleRevokeEvmWalletAccess { entries: accesses }).await {
|
|
Ok(()) => {
|
|
info!("Successfully revoked wallet access");
|
|
return Ok(None);
|
|
}
|
|
Err(err) => {
|
|
warn!(error = ?err, "Failed to revoke wallet access");
|
|
return Err(Status::internal("Failed to revoke wallet access"));
|
|
}
|
|
}
|
|
}
|
|
|
|
UserAgentRequestPayload::ListWalletAccess(_) => {
|
|
let result = match actor.ask(HandleListWalletAccess {}).await {
|
|
Ok(accesses) => ListWalletAccessResponse {
|
|
accesses: accesses.into_iter().map(|a| a.convert()).collect(),
|
|
},
|
|
Err(err) => {
|
|
warn!(error = ?err, "Failed to list wallet access");
|
|
return Err(Status::internal("Failed to list wallet access"));
|
|
}
|
|
};
|
|
UserAgentResponsePayload::ListWalletAccessResponse(result)
|
|
}
|
|
|
|
UserAgentRequestPayload::AuthChallengeRequest(..)
|
|
| UserAgentRequestPayload::AuthChallengeSolution(..) => {
|
|
warn!(?payload, "Unsupported post-auth user agent request");
|
|
return Err(Status::invalid_argument("Unsupported user-agent request"));
|
|
}
|
|
};
|
|
|
|
Ok(Some(response))
|
|
}
|
|
|
|
pub async fn start(
|
|
mut conn: UserAgentConnection,
|
|
mut bi: GrpcBi<UserAgentRequest, UserAgentResponse>,
|
|
) {
|
|
let mut request_tracker = RequestTracker::default();
|
|
|
|
let pubkey = match auth::start(&mut conn, &mut bi, &mut request_tracker).await {
|
|
Ok(pubkey) => pubkey,
|
|
Err(e) => {
|
|
warn!(error = ?e, "Authentication failed");
|
|
return;
|
|
}
|
|
};
|
|
|
|
let (oob_sender, oob_receiver) = mpsc::channel(16);
|
|
let oob_adapter = OutOfBandAdapter(oob_sender);
|
|
|
|
let actor = UserAgentSession::spawn(UserAgentSession::new(conn, Box::new(oob_adapter)));
|
|
let actor_for_cleanup = actor.clone();
|
|
|
|
info!(?pubkey, "User authenticated successfully");
|
|
dispatch_loop(bi, actor, oob_receiver, request_tracker).await;
|
|
actor_for_cleanup.kill();
|
|
}
|