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, SdkClientEntry as ProtoSdkClientEntry, SdkClientError as ProtoSdkClientError,
SdkClientGrantWalletAccess, SdkClientList as ProtoSdkClientList, SdkClientGrantWalletAccess, SdkClientList as ProtoSdkClientList,
SdkClientListResponse as ProtoSdkClientListResponse, SdkClientRevokeWalletAccess, SdkClientListResponse as ProtoSdkClientListResponse, SdkClientRevokeWalletAccess,
SdkClientWalletAccess, UnsealEncryptedKey as ProtoUnsealEncryptedKey, UnsealEncryptedKey as ProtoUnsealEncryptedKey,
UnsealResult as ProtoUnsealResult, UnsealStart, UserAgentRequest, UserAgentResponse, UnsealResult as ProtoUnsealResult, UnsealStart, UserAgentRequest, UserAgentResponse,
VaultState as ProtoVaultState, VaultState as ProtoVaultState,
sdk_client_list_response::Result as ProtoSdkClientListResult, 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}, grpc::{Convert, TryConvert, request_tracker::RequestTracker},
}; };
mod auth; mod auth;
@@ -158,285 +158,336 @@ async fn dispatch_inner(
actor: &ActorRef<UserAgentSession>, actor: &ActorRef<UserAgentSession>,
payload: UserAgentRequestPayload, payload: UserAgentRequestPayload,
) -> Result<Option<UserAgentResponsePayload>, Status> { ) -> Result<Option<UserAgentResponsePayload>, Status> {
let response = match payload { match payload {
UserAgentRequestPayload::UnsealStart(UnsealStart { client_pubkey }) => { UserAgentRequestPayload::UnsealStart(req) => handle_unseal_start(actor, req).await,
let client_pubkey = <[u8; 32]>::try_from(client_pubkey) UserAgentRequestPayload::UnsealEncryptedKey(req) => {
.map(x25519_dalek::PublicKey::from) handle_unseal_encrypted_key(actor, req).await
.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::BootstrapEncryptedKey(req) => {
UserAgentRequestPayload::UnsealEncryptedKey(ProtoUnsealEncryptedKey { handle_bootstrap_encrypted_key(actor, req).await
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::QueryVaultState(_) => handle_query_vault_state(actor).await,
UserAgentRequestPayload::BootstrapEncryptedKey(ProtoBootstrapEncryptedKey { UserAgentRequestPayload::EvmWalletCreate(_) => handle_evm_wallet_create(actor).await,
nonce, UserAgentRequestPayload::EvmWalletList(_) => handle_evm_wallet_list(actor).await,
ciphertext, UserAgentRequestPayload::EvmGrantList(_) => handle_evm_grant_list(actor).await,
associated_data, UserAgentRequestPayload::EvmGrantCreate(req) => handle_evm_grant_create(actor, req).await,
}) => { UserAgentRequestPayload::EvmGrantDelete(req) => handle_evm_grant_delete(actor, req).await,
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) => { UserAgentRequestPayload::SdkClientConnectionResponse(resp) => {
let pubkey_bytes = <[u8; 32]>::try_from(resp.pubkey) handle_sdk_client_connection_response(actor, resp).await
.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(_) => {
UserAgentRequestPayload::SdkClientRevoke(_) => todo!(), Err(Status::unimplemented("SdkClientRevoke is not yet implemented"))
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::SdkClientList(_) => handle_sdk_client_list(actor).await,
UserAgentRequestPayload::GrantWalletAccess(SdkClientGrantWalletAccess { accesses }) => { UserAgentRequestPayload::GrantWalletAccess(req) => {
let entries: Vec<NewEvmWalletAccess> = handle_grant_wallet_access(actor, req).await
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(req) => {
UserAgentRequestPayload::RevokeWalletAccess(SdkClientRevokeWalletAccess { accesses }) => { handle_revoke_wallet_access(actor, req).await
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(_) => handle_list_wallet_access(actor).await,
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::AuthChallengeRequest(..)
| UserAgentRequestPayload::AuthChallengeSolution(..) => { | UserAgentRequestPayload::AuthChallengeSolution(..) => {
warn!(?payload, "Unsupported post-auth user agent request"); warn!(?payload, "Unsupported post-auth user agent request");
return Err(Status::invalid_argument("Unsupported 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"))?;
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")
})?;
Ok(Some(UserAgentResponsePayload::UnsealStartResponse(
arbiter_proto::proto::user_agent::UnsealStartResponse {
server_pubkey: response.server_pubkey.as_bytes().to_vec(),
},
)))
}
async fn handle_unseal_encrypted_key(
actor: &ActorRef<UserAgentSession>,
req: ProtoUnsealEncryptedKey,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match actor
.ask(HandleUnsealEncryptedKey {
nonce: req.nonce,
ciphertext: req.ciphertext,
associated_data: req.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"));
} }
}; };
Ok(Some(UserAgentResponsePayload::UnsealResult(result.into())))
}
Ok(Some(response)) async fn handle_bootstrap_encrypted_key(
actor: &ActorRef<UserAgentSession>,
req: ProtoBootstrapEncryptedKey,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match actor
.ask(HandleBootstrapEncryptedKey {
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::AlreadyBootstrapped)) => {
ProtoBootstrapResult::AlreadyBootstrapped
}
Err(err) => {
warn!(error = ?err, "Failed to handle bootstrap request");
return Err(Status::internal("Failed to bootstrap vault"));
}
};
Ok(Some(UserAgentResponsePayload::BootstrapResult(result.into())))
}
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,
Ok(KeyHolderState::Unsealed) => ProtoVaultState::Unsealed,
Err(err) => {
warn!(error = ?err, "Failed to query vault state");
ProtoVaultState::Error
}
};
Ok(Some(UserAgentResponsePayload::VaultState(state.into())))
}
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,
address: address.to_vec(),
}),
Err(err) => {
warn!(error = ?err, "Failed to create EVM wallet");
WalletCreateResult::Error(ProtoEvmError::Internal.into())
}
};
Ok(Some(UserAgentResponsePayload::EvmWalletCreate(
WalletCreateResponse { result: Some(result) },
)))
}
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
.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())
}
};
Ok(Some(UserAgentResponsePayload::EvmWalletList(
WalletListResponse { result: Some(result) },
)))
}
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
.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())
}
};
Ok(Some(UserAgentResponsePayload::EvmGrantList(
EvmGrantListResponse { result: Some(result) },
)))
}
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 = req
.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())
}
};
Ok(Some(UserAgentResponsePayload::EvmGrantCreate(
EvmGrantCreateResponse { result: Some(result) },
)))
}
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())
}
};
Ok(Some(UserAgentResponsePayload::EvmGrantDelete(
EvmGrantDeleteResponse { result: Some(result) },
)))
}
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)
.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")
})?;
Ok(None)
}
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
.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())
}
};
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");
Ok(None)
}
Err(err) => {
warn!(error = ?err, "Failed to grant wallet access");
Err(Status::internal("Failed to grant wallet access"))
}
}
}
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");
Ok(None)
}
Err(err) => {
warn!(error = ?err, "Failed to revoke wallet access");
Err(Status::internal("Failed to revoke wallet access"))
}
}
}
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");
Err(Status::internal("Failed to list wallet access"))
}
}
} }
pub async fn start( pub async fn start(