refactor(grpc): extract user agent request handlers into separate functions

This commit is contained in:
hdbg
2026-03-29 12:45:00 +02:00
parent e2d8b7841b
commit 523bf783ac

View File

@@ -21,7 +21,7 @@ use arbiter_proto::{
SdkClientEntry as ProtoSdkClientEntry, SdkClientError as ProtoSdkClientError,
SdkClientGrantWalletAccess, SdkClientList as ProtoSdkClientList,
SdkClientListResponse as ProtoSdkClientListResponse, SdkClientRevokeWalletAccess,
SdkClientWalletAccess, UnsealEncryptedKey as ProtoUnsealEncryptedKey,
UnsealEncryptedKey as ProtoUnsealEncryptedKey,
UnsealResult as ProtoUnsealResult, UnsealStart, UserAgentRequest, UserAgentResponse,
VaultState as ProtoVaultState,
sdk_client_list_response::Result as ProtoSdkClientListResult,
@@ -53,7 +53,7 @@ use crate::{
},
},
},
db::models::{CoreEvmWalletAccess, NewEvmWalletAccess},
db::models::NewEvmWalletAccess,
grpc::{Convert, TryConvert, request_tracker::RequestTracker},
};
mod auth;
@@ -158,9 +158,47 @@ 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)
match payload {
UserAgentRequestPayload::UnsealStart(req) => handle_unseal_start(actor, req).await,
UserAgentRequestPayload::UnsealEncryptedKey(req) => {
handle_unseal_encrypted_key(actor, req).await
}
UserAgentRequestPayload::BootstrapEncryptedKey(req) => {
handle_bootstrap_encrypted_key(actor, req).await
}
UserAgentRequestPayload::QueryVaultState(_) => handle_query_vault_state(actor).await,
UserAgentRequestPayload::EvmWalletCreate(_) => handle_evm_wallet_create(actor).await,
UserAgentRequestPayload::EvmWalletList(_) => handle_evm_wallet_list(actor).await,
UserAgentRequestPayload::EvmGrantList(_) => handle_evm_grant_list(actor).await,
UserAgentRequestPayload::EvmGrantCreate(req) => handle_evm_grant_create(actor, req).await,
UserAgentRequestPayload::EvmGrantDelete(req) => handle_evm_grant_delete(actor, req).await,
UserAgentRequestPayload::SdkClientConnectionResponse(resp) => {
handle_sdk_client_connection_response(actor, resp).await
}
UserAgentRequestPayload::SdkClientRevoke(_) => {
Err(Status::unimplemented("SdkClientRevoke is not yet implemented"))
}
UserAgentRequestPayload::SdkClientList(_) => handle_sdk_client_list(actor).await,
UserAgentRequestPayload::GrantWalletAccess(req) => {
handle_grant_wallet_access(actor, req).await
}
UserAgentRequestPayload::RevokeWalletAccess(req) => {
handle_revoke_wallet_access(actor, req).await
}
UserAgentRequestPayload::ListWalletAccess(_) => handle_list_wallet_access(actor).await,
UserAgentRequestPayload::AuthChallengeRequest(..)
| UserAgentRequestPayload::AuthChallengeSolution(..) => {
warn!(?payload, "Unsupported post-auth user agent request");
Err(Status::invalid_argument("Unsupported user-agent request"))
}
}
}
async fn handle_unseal_start(
actor: &ActorRef<UserAgentSession>,
req: UnsealStart,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let client_pubkey = <[u8; 32]>::try_from(req.client_pubkey)
.map(x25519_dalek::PublicKey::from)
.map_err(|_| Status::invalid_argument("Invalid X25519 public key"))?;
@@ -172,55 +210,49 @@ async fn dispatch_inner(
Status::internal("Failed to start unseal flow")
})?;
UserAgentResponsePayload::UnsealStartResponse(
Ok(Some(UserAgentResponsePayload::UnsealStartResponse(
arbiter_proto::proto::user_agent::UnsealStartResponse {
server_pubkey: response.server_pubkey.as_bytes().to_vec(),
},
)
}
)))
}
UserAgentRequestPayload::UnsealEncryptedKey(ProtoUnsealEncryptedKey {
nonce,
ciphertext,
associated_data,
}) => {
async fn handle_unseal_encrypted_key(
actor: &ActorRef<UserAgentSession>,
req: ProtoUnsealEncryptedKey,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match actor
.ask(HandleUnsealEncryptedKey {
nonce,
ciphertext,
associated_data,
nonce: req.nonce,
ciphertext: req.ciphertext,
associated_data: req.associated_data,
})
.await
{
Ok(()) => ProtoUnsealResult::Success,
Err(SendError::HandlerError(UnsealError::InvalidKey)) => {
ProtoUnsealResult::InvalidKey
}
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())
}
Ok(Some(UserAgentResponsePayload::UnsealResult(result.into())))
}
UserAgentRequestPayload::BootstrapEncryptedKey(ProtoBootstrapEncryptedKey {
nonce,
ciphertext,
associated_data,
}) => {
async fn handle_bootstrap_encrypted_key(
actor: &ActorRef<UserAgentSession>,
req: ProtoBootstrapEncryptedKey,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match actor
.ask(HandleBootstrapEncryptedKey {
nonce,
ciphertext,
associated_data,
nonce: req.nonce,
ciphertext: req.ciphertext,
associated_data: req.associated_data,
})
.await
{
Ok(()) => ProtoBootstrapResult::Success,
Err(SendError::HandlerError(BootstrapError::InvalidKey)) => {
ProtoBootstrapResult::InvalidKey
}
Err(SendError::HandlerError(BootstrapError::InvalidKey)) => ProtoBootstrapResult::InvalidKey,
Err(SendError::HandlerError(BootstrapError::AlreadyBootstrapped)) => {
ProtoBootstrapResult::AlreadyBootstrapped
}
@@ -229,10 +261,12 @@ async fn dispatch_inner(
return Err(Status::internal("Failed to bootstrap vault"));
}
};
UserAgentResponsePayload::BootstrapResult(result.into())
}
Ok(Some(UserAgentResponsePayload::BootstrapResult(result.into())))
}
UserAgentRequestPayload::QueryVaultState(_) => {
async fn handle_query_vault_state(
actor: &ActorRef<UserAgentSession>,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let state = match actor.ask(HandleQueryVaultState {}).await {
Ok(KeyHolderState::Unbootstrapped) => ProtoVaultState::Unbootstrapped,
Ok(KeyHolderState::Sealed) => ProtoVaultState::Sealed,
@@ -242,10 +276,12 @@ async fn dispatch_inner(
ProtoVaultState::Error
}
};
UserAgentResponsePayload::VaultState(state.into())
}
Ok(Some(UserAgentResponsePayload::VaultState(state.into())))
}
UserAgentRequestPayload::EvmWalletCreate(_) => {
async fn handle_evm_wallet_create(
actor: &ActorRef<UserAgentSession>,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match actor.ask(HandleEvmWalletCreate {}).await {
Ok((wallet_id, address)) => WalletCreateResult::Wallet(WalletEntry {
id: wallet_id,
@@ -256,12 +292,14 @@ async fn dispatch_inner(
WalletCreateResult::Error(ProtoEvmError::Internal.into())
}
};
UserAgentResponsePayload::EvmWalletCreate(WalletCreateResponse {
result: Some(result),
})
}
Ok(Some(UserAgentResponsePayload::EvmWalletCreate(
WalletCreateResponse { result: Some(result) },
)))
}
UserAgentRequestPayload::EvmWalletList(_) => {
async fn handle_evm_wallet_list(
actor: &ActorRef<UserAgentSession>,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match actor.ask(HandleEvmWalletList {}).await {
Ok(wallets) => WalletListResult::Wallets(WalletList {
wallets: wallets
@@ -277,12 +315,14 @@ async fn dispatch_inner(
WalletListResult::Error(ProtoEvmError::Internal.into())
}
};
UserAgentResponsePayload::EvmWalletList(WalletListResponse {
result: Some(result),
})
}
Ok(Some(UserAgentResponsePayload::EvmWalletList(
WalletListResponse { result: Some(result) },
)))
}
UserAgentRequestPayload::EvmGrantList(_) => {
async fn handle_evm_grant_list(
actor: &ActorRef<UserAgentSession>,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match actor.ask(HandleGrantList {}).await {
Ok(grants) => EvmGrantListResult::Grants(EvmGrantList {
grants: grants
@@ -300,16 +340,21 @@ async fn dispatch_inner(
EvmGrantListResult::Error(ProtoEvmError::Internal.into())
}
};
UserAgentResponsePayload::EvmGrantList(EvmGrantListResponse {
result: Some(result),
})
}
Ok(Some(UserAgentResponsePayload::EvmGrantList(
EvmGrantListResponse { result: Some(result) },
)))
}
UserAgentRequestPayload::EvmGrantCreate(EvmGrantCreateRequest { shared, specific }) => {
let basic = shared
async fn handle_evm_grant_create(
actor: &ActorRef<UserAgentSession>,
req: EvmGrantCreateRequest,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let basic = req
.shared
.ok_or_else(|| Status::invalid_argument("Missing shared grant settings"))?
.try_convert()?;
let grant = specific
let grant = req
.specific
.ok_or_else(|| Status::invalid_argument("Missing specific grant settings"))?
.try_convert()?;
@@ -320,25 +365,31 @@ async fn dispatch_inner(
EvmGrantCreateResult::Error(ProtoEvmError::Internal.into())
}
};
UserAgentResponsePayload::EvmGrantCreate(EvmGrantCreateResponse {
result: Some(result),
})
}
Ok(Some(UserAgentResponsePayload::EvmGrantCreate(
EvmGrantCreateResponse { result: Some(result) },
)))
}
UserAgentRequestPayload::EvmGrantDelete(EvmGrantDeleteRequest { grant_id }) => {
let result = match actor.ask(HandleGrantDelete { grant_id }).await {
async fn handle_evm_grant_delete(
actor: &ActorRef<UserAgentSession>,
req: EvmGrantDeleteRequest,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match actor.ask(HandleGrantDelete { grant_id: req.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),
})
}
Ok(Some(UserAgentResponsePayload::EvmGrantDelete(
EvmGrantDeleteResponse { result: Some(result) },
)))
}
UserAgentRequestPayload::SdkClientConnectionResponse(resp) => {
async fn handle_sdk_client_connection_response(
actor: &ActorRef<UserAgentSession>,
resp: arbiter_proto::proto::user_agent::SdkClientConnectionResponse,
) -> Result<Option<UserAgentResponsePayload>, Status> {
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)
@@ -355,12 +406,12 @@ async fn dispatch_inner(
Status::internal("Failed to process response")
})?;
return Ok(None);
}
Ok(None)
}
UserAgentRequestPayload::SdkClientRevoke(_) => todo!(),
UserAgentRequestPayload::SdkClientList(_) => {
async fn handle_sdk_client_list(
actor: &ActorRef<UserAgentSession>,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match actor.ask(HandleSdkClientList {}).await {
Ok(clients) => ProtoSdkClientListResult::Clients(ProtoSdkClientList {
clients: clients
@@ -382,61 +433,61 @@ async fn dispatch_inner(
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();
Ok(Some(UserAgentResponsePayload::SdkClientListResponse(
ProtoSdkClientListResponse { result: Some(result) },
)))
}
async fn handle_grant_wallet_access(
actor: &ActorRef<UserAgentSession>,
req: SdkClientGrantWalletAccess,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let entries: Vec<NewEvmWalletAccess> = req.accesses.into_iter().map(|a| a.convert()).collect();
match actor.ask(HandleGrantEvmWalletAccess { entries }).await {
Ok(()) => {
info!("Successfully granted wallet access");
return Ok(None);
Ok(None)
}
Err(err) => {
warn!(error = ?err, "Failed to grant wallet access");
return Err(Status::internal("Failed to grant wallet access"));
}
Err(Status::internal("Failed to grant wallet access"))
}
}
}
UserAgentRequestPayload::RevokeWalletAccess(SdkClientRevokeWalletAccess { accesses }) => {
match actor.ask(HandleRevokeEvmWalletAccess { entries: accesses }).await {
async fn handle_revoke_wallet_access(
actor: &ActorRef<UserAgentSession>,
req: SdkClientRevokeWalletAccess,
) -> Result<Option<UserAgentResponsePayload>, Status> {
match actor
.ask(HandleRevokeEvmWalletAccess { entries: req.accesses })
.await
{
Ok(()) => {
info!("Successfully revoked wallet access");
return Ok(None);
Ok(None)
}
Err(err) => {
warn!(error = ?err, "Failed to revoke wallet access");
return Err(Status::internal("Failed to revoke wallet access"));
}
Err(Status::internal("Failed to revoke wallet access"))
}
}
}
UserAgentRequestPayload::ListWalletAccess(_) => {
let result = match actor.ask(HandleListWalletAccess {}).await {
Ok(accesses) => ListWalletAccessResponse {
async fn handle_list_wallet_access(
actor: &ActorRef<UserAgentSession>,
) -> Result<Option<UserAgentResponsePayload>, Status> {
match actor.ask(HandleListWalletAccess {}).await {
Ok(accesses) => Ok(Some(UserAgentResponsePayload::ListWalletAccessResponse(
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"));
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(