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 e5be55e141
commit 805d5691d6

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,9 +158,47 @@ 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) => {
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(x25519_dalek::PublicKey::from)
.map_err(|_| Status::invalid_argument("Invalid X25519 public key"))?; .map_err(|_| Status::invalid_argument("Invalid X25519 public key"))?;
@@ -172,55 +210,49 @@ async fn dispatch_inner(
Status::internal("Failed to start unseal flow") Status::internal("Failed to start unseal flow")
})?; })?;
UserAgentResponsePayload::UnsealStartResponse( Ok(Some(UserAgentResponsePayload::UnsealStartResponse(
arbiter_proto::proto::user_agent::UnsealStartResponse { arbiter_proto::proto::user_agent::UnsealStartResponse {
server_pubkey: response.server_pubkey.as_bytes().to_vec(), server_pubkey: response.server_pubkey.as_bytes().to_vec(),
}, },
) )))
} }
UserAgentRequestPayload::UnsealEncryptedKey(ProtoUnsealEncryptedKey { async fn handle_unseal_encrypted_key(
nonce, actor: &ActorRef<UserAgentSession>,
ciphertext, req: ProtoUnsealEncryptedKey,
associated_data, ) -> Result<Option<UserAgentResponsePayload>, Status> {
}) => {
let result = match actor let result = match actor
.ask(HandleUnsealEncryptedKey { .ask(HandleUnsealEncryptedKey {
nonce, nonce: req.nonce,
ciphertext, ciphertext: req.ciphertext,
associated_data, associated_data: req.associated_data,
}) })
.await .await
{ {
Ok(()) => ProtoUnsealResult::Success, Ok(()) => ProtoUnsealResult::Success,
Err(SendError::HandlerError(UnsealError::InvalidKey)) => { Err(SendError::HandlerError(UnsealError::InvalidKey)) => ProtoUnsealResult::InvalidKey,
ProtoUnsealResult::InvalidKey
}
Err(err) => { Err(err) => {
warn!(error = ?err, "Failed to handle unseal request"); warn!(error = ?err, "Failed to handle unseal request");
return Err(Status::internal("Failed to unseal vault")); return Err(Status::internal("Failed to unseal vault"));
} }
}; };
UserAgentResponsePayload::UnsealResult(result.into()) Ok(Some(UserAgentResponsePayload::UnsealResult(result.into())))
} }
UserAgentRequestPayload::BootstrapEncryptedKey(ProtoBootstrapEncryptedKey { async fn handle_bootstrap_encrypted_key(
nonce, actor: &ActorRef<UserAgentSession>,
ciphertext, req: ProtoBootstrapEncryptedKey,
associated_data, ) -> Result<Option<UserAgentResponsePayload>, Status> {
}) => {
let result = match actor let result = match actor
.ask(HandleBootstrapEncryptedKey { .ask(HandleBootstrapEncryptedKey {
nonce, nonce: req.nonce,
ciphertext, ciphertext: req.ciphertext,
associated_data, associated_data: req.associated_data,
}) })
.await .await
{ {
Ok(()) => ProtoBootstrapResult::Success, Ok(()) => ProtoBootstrapResult::Success,
Err(SendError::HandlerError(BootstrapError::InvalidKey)) => { Err(SendError::HandlerError(BootstrapError::InvalidKey)) => ProtoBootstrapResult::InvalidKey,
ProtoBootstrapResult::InvalidKey
}
Err(SendError::HandlerError(BootstrapError::AlreadyBootstrapped)) => { Err(SendError::HandlerError(BootstrapError::AlreadyBootstrapped)) => {
ProtoBootstrapResult::AlreadyBootstrapped ProtoBootstrapResult::AlreadyBootstrapped
} }
@@ -229,10 +261,12 @@ async fn dispatch_inner(
return Err(Status::internal("Failed to bootstrap vault")); 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 { let state = match actor.ask(HandleQueryVaultState {}).await {
Ok(KeyHolderState::Unbootstrapped) => ProtoVaultState::Unbootstrapped, Ok(KeyHolderState::Unbootstrapped) => ProtoVaultState::Unbootstrapped,
Ok(KeyHolderState::Sealed) => ProtoVaultState::Sealed, Ok(KeyHolderState::Sealed) => ProtoVaultState::Sealed,
@@ -242,10 +276,12 @@ async fn dispatch_inner(
ProtoVaultState::Error 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 { let result = match actor.ask(HandleEvmWalletCreate {}).await {
Ok((wallet_id, address)) => WalletCreateResult::Wallet(WalletEntry { Ok((wallet_id, address)) => WalletCreateResult::Wallet(WalletEntry {
id: wallet_id, id: wallet_id,
@@ -256,12 +292,14 @@ async fn dispatch_inner(
WalletCreateResult::Error(ProtoEvmError::Internal.into()) WalletCreateResult::Error(ProtoEvmError::Internal.into())
} }
}; };
UserAgentResponsePayload::EvmWalletCreate(WalletCreateResponse { Ok(Some(UserAgentResponsePayload::EvmWalletCreate(
result: Some(result), 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 { let result = match actor.ask(HandleEvmWalletList {}).await {
Ok(wallets) => WalletListResult::Wallets(WalletList { Ok(wallets) => WalletListResult::Wallets(WalletList {
wallets: wallets wallets: wallets
@@ -277,12 +315,14 @@ async fn dispatch_inner(
WalletListResult::Error(ProtoEvmError::Internal.into()) WalletListResult::Error(ProtoEvmError::Internal.into())
} }
}; };
UserAgentResponsePayload::EvmWalletList(WalletListResponse { Ok(Some(UserAgentResponsePayload::EvmWalletList(
result: Some(result), 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 { let result = match actor.ask(HandleGrantList {}).await {
Ok(grants) => EvmGrantListResult::Grants(EvmGrantList { Ok(grants) => EvmGrantListResult::Grants(EvmGrantList {
grants: grants grants: grants
@@ -300,16 +340,21 @@ async fn dispatch_inner(
EvmGrantListResult::Error(ProtoEvmError::Internal.into()) EvmGrantListResult::Error(ProtoEvmError::Internal.into())
} }
}; };
UserAgentResponsePayload::EvmGrantList(EvmGrantListResponse { Ok(Some(UserAgentResponsePayload::EvmGrantList(
result: Some(result), EvmGrantListResponse { result: Some(result) },
}) )))
} }
UserAgentRequestPayload::EvmGrantCreate(EvmGrantCreateRequest { shared, specific }) => { async fn handle_evm_grant_create(
let basic = shared actor: &ActorRef<UserAgentSession>,
req: EvmGrantCreateRequest,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let basic = req
.shared
.ok_or_else(|| Status::invalid_argument("Missing shared grant settings"))? .ok_or_else(|| Status::invalid_argument("Missing shared grant settings"))?
.try_convert()?; .try_convert()?;
let grant = specific let grant = req
.specific
.ok_or_else(|| Status::invalid_argument("Missing specific grant settings"))? .ok_or_else(|| Status::invalid_argument("Missing specific grant settings"))?
.try_convert()?; .try_convert()?;
@@ -320,25 +365,31 @@ async fn dispatch_inner(
EvmGrantCreateResult::Error(ProtoEvmError::Internal.into()) EvmGrantCreateResult::Error(ProtoEvmError::Internal.into())
} }
}; };
UserAgentResponsePayload::EvmGrantCreate(EvmGrantCreateResponse { Ok(Some(UserAgentResponsePayload::EvmGrantCreate(
result: Some(result), EvmGrantCreateResponse { result: Some(result) },
}) )))
} }
UserAgentRequestPayload::EvmGrantDelete(EvmGrantDeleteRequest { grant_id }) => { async fn handle_evm_grant_delete(
let result = match actor.ask(HandleGrantDelete { grant_id }).await { actor: &ActorRef<UserAgentSession>,
req: EvmGrantDeleteRequest,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match actor.ask(HandleGrantDelete { grant_id: req.grant_id }).await {
Ok(()) => EvmGrantDeleteResult::Ok(()), Ok(()) => EvmGrantDeleteResult::Ok(()),
Err(err) => { Err(err) => {
warn!(error = ?err, "Failed to delete EVM grant"); warn!(error = ?err, "Failed to delete EVM grant");
EvmGrantDeleteResult::Error(ProtoEvmError::Internal.into()) EvmGrantDeleteResult::Error(ProtoEvmError::Internal.into())
} }
}; };
UserAgentResponsePayload::EvmGrantDelete(EvmGrantDeleteResponse { Ok(Some(UserAgentResponsePayload::EvmGrantDelete(
result: Some(result), 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) let pubkey_bytes = <[u8; 32]>::try_from(resp.pubkey)
.map_err(|_| Status::invalid_argument("Invalid Ed25519 public key length"))?; .map_err(|_| Status::invalid_argument("Invalid Ed25519 public key length"))?;
let pubkey = ed25519_dalek::VerifyingKey::from_bytes(&pubkey_bytes) let pubkey = ed25519_dalek::VerifyingKey::from_bytes(&pubkey_bytes)
@@ -355,12 +406,12 @@ async fn dispatch_inner(
Status::internal("Failed to process response") Status::internal("Failed to process response")
})?; })?;
return Ok(None); Ok(None)
} }
UserAgentRequestPayload::SdkClientRevoke(_) => todo!(), async fn handle_sdk_client_list(
actor: &ActorRef<UserAgentSession>,
UserAgentRequestPayload::SdkClientList(_) => { ) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match actor.ask(HandleSdkClientList {}).await { let result = match actor.ask(HandleSdkClientList {}).await {
Ok(clients) => ProtoSdkClientListResult::Clients(ProtoSdkClientList { Ok(clients) => ProtoSdkClientListResult::Clients(ProtoSdkClientList {
clients: clients clients: clients
@@ -382,61 +433,61 @@ async fn dispatch_inner(
ProtoSdkClientListResult::Error(ProtoSdkClientError::Internal.into()) ProtoSdkClientListResult::Error(ProtoSdkClientError::Internal.into())
} }
}; };
UserAgentResponsePayload::SdkClientListResponse(ProtoSdkClientListResponse { Ok(Some(UserAgentResponsePayload::SdkClientListResponse(
result: Some(result), ProtoSdkClientListResponse { result: Some(result) },
}) )))
} }
UserAgentRequestPayload::GrantWalletAccess(SdkClientGrantWalletAccess { accesses }) => {
let entries: Vec<NewEvmWalletAccess> =
accesses.into_iter().map(|a| a.convert()).collect();
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 { match actor.ask(HandleGrantEvmWalletAccess { entries }).await {
Ok(()) => { Ok(()) => {
info!("Successfully granted wallet access"); info!("Successfully granted wallet access");
return Ok(None); Ok(None)
} }
Err(err) => { Err(err) => {
warn!(error = ?err, "Failed to grant wallet access"); 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 }) => { async fn handle_revoke_wallet_access(
match actor.ask(HandleRevokeEvmWalletAccess { entries: accesses }).await { actor: &ActorRef<UserAgentSession>,
req: SdkClientRevokeWalletAccess,
) -> Result<Option<UserAgentResponsePayload>, Status> {
match actor
.ask(HandleRevokeEvmWalletAccess { entries: req.accesses })
.await
{
Ok(()) => { Ok(()) => {
info!("Successfully revoked wallet access"); info!("Successfully revoked wallet access");
return Ok(None); Ok(None)
} }
Err(err) => { Err(err) => {
warn!(error = ?err, "Failed to revoke wallet access"); 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(_) => { async fn handle_list_wallet_access(
let result = match actor.ask(HandleListWalletAccess {}).await { actor: &ActorRef<UserAgentSession>,
Ok(accesses) => ListWalletAccessResponse { ) -> 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(), accesses: accesses.into_iter().map(|a| a.convert()).collect(),
}, },
))),
Err(err) => { Err(err) => {
warn!(error = ?err, "Failed to list wallet access"); 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( pub async fn start(