refactor(proto): restructure wallet access messages for improved data organization

This commit is contained in:
hdbg
2026-03-28 12:49:47 +01:00
parent 27428f709a
commit ca35b9fed7
10 changed files with 212 additions and 113 deletions

View File

@@ -3,11 +3,6 @@ use crate::{
db::{self, models::KeyType},
};
pub struct EvmAccessEntry {
pub wallet_id: i32,
pub sdk_client_id: i32,
}
/// Abstraction over Ed25519 / ECDSA-secp256k1 / RSA public keys used during the auth handshake.
#[derive(Clone, Debug)]
pub enum AuthPublicKey {

View File

@@ -13,9 +13,10 @@ use x25519_dalek::{EphemeralSecret, PublicKey};
use crate::actors::flow_coordinator::client_connect_approval::ClientApprovalAnswer;
use crate::actors::keyholder::KeyHolderState;
use crate::actors::user_agent::EvmAccessEntry;
use crate::actors::user_agent::session::Error;
use crate::db::models::{ProgramClient, ProgramClientMetadata};
use crate::db::models::{
CoreEvmWalletAccess, EvmWalletAccess, NewEvmWalletAccess, ProgramClient, ProgramClientMetadata,
};
use crate::db::schema::evm_wallet_access;
use crate::evm::policies::{Grant, SpecificGrant};
use crate::safe_cell::SafeCell;
@@ -304,8 +305,6 @@ impl UserAgentSession {
}
}
#[messages]
impl UserAgentSession {
#[message]
@@ -360,20 +359,16 @@ impl UserAgentSession {
#[message]
pub(crate) async fn handle_grant_evm_wallet_access(
&mut self,
entries: Vec<EvmAccessEntry>,
entries: Vec<NewEvmWalletAccess>,
) -> Result<(), Error> {
let mut conn = self.props.db.get().await?;
conn.transaction(|conn| {
Box::pin(async move {
use crate::db::models::NewEvmWalletAccess;
use crate::db::schema::evm_wallet_access;
for entry in entries {
diesel::insert_into(evm_wallet_access::table)
.values(&NewEvmWalletAccess {
wallet_id: entry.wallet_id,
client_id: entry.sdk_client_id,
})
.values(&entry)
.on_conflict_do_nothing()
.execute(conn)
.await?;
@@ -389,7 +384,7 @@ impl UserAgentSession {
#[message]
pub(crate) async fn handle_revoke_evm_wallet_access(
&mut self,
entries: Vec<EvmAccessEntry>,
entries: Vec<i32>,
) -> Result<(), Error> {
let mut conn = self.props.db.get().await?;
conn.transaction(|conn| {
@@ -397,11 +392,7 @@ impl UserAgentSession {
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.wallet_id)
.and(evm_wallet_access::client_id.eq(entry.sdk_client_id)),
)
.filter(evm_wallet_access::wallet_id.eq(entry))
.execute(conn)
.await?;
}
@@ -414,19 +405,15 @@ impl UserAgentSession {
}
#[message]
pub(crate) async fn handle_list_wallet_access(&mut self) -> Result<Vec<EvmAccessEntry>, Error> {
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((evm_wallet_access::wallet_id, evm_wallet_access::client_id))
.load::<(i32, i32)>(&mut conn)
.await?
.into_iter()
.map(|(wallet_id, sdk_client_id)| EvmAccessEntry {
wallet_id,
sdk_client_id,
})
.collect();
.select(EvmWalletAccess::as_select())
.load::<_>(&mut conn)
.await?;
Ok(access_entries)
}
}

View File

@@ -193,6 +193,12 @@ pub struct EvmWallet {
omit(id, created_at),
attributes_with = "deriveless"
)]
#[view(
CoreEvmWalletAccess,
derive(Insertable),
omit(created_at),
attributes_with = "deriveless"
)]
pub struct EvmWalletAccess {
pub id: i32,
pub wallet_id: i32,

View File

@@ -45,10 +45,15 @@ use crate::{
user_agent::{
OutOfBand, UserAgentConnection, UserAgentSession,
session::connection::{
BootstrapError, HandleBootstrapEncryptedKey, HandleEvmWalletCreate, HandleEvmWalletList, HandleGrantCreate, HandleGrantDelete, HandleGrantEvmWalletAccess, HandleGrantList, HandleListWalletAccess, HandleNewClientApprove, HandleQueryVaultState, HandleRevokeEvmWalletAccess, HandleSdkClientList, HandleUnsealEncryptedKey, HandleUnsealRequest, UnsealError
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;
@@ -383,7 +388,8 @@ async fn dispatch_inner(
}
UserAgentRequestPayload::GrantWalletAccess(SdkClientGrantWalletAccess { accesses }) => {
let entries = accesses.try_convert()?;
let entries: Vec<NewEvmWalletAccess> =
accesses.into_iter().map(|a| a.convert()).collect();
match actor.ask(HandleGrantEvmWalletAccess { entries }).await {
Ok(()) => {
@@ -398,9 +404,7 @@ async fn dispatch_inner(
}
UserAgentRequestPayload::RevokeWalletAccess(SdkClientRevokeWalletAccess { accesses }) => {
let entries = accesses.try_convert()?;
match actor.ask(HandleRevokeEvmWalletAccess { entries }).await {
match actor.ask(HandleRevokeEvmWalletAccess { entries: accesses }).await {
Ok(()) => {
info!("Successfully revoked wallet access");
return Ok(None);

View File

@@ -1,23 +1,21 @@
use alloy::primitives::{Address, U256};
use arbiter_proto::proto::evm::{
EtherTransferSettings as ProtoEtherTransferSettings,
SharedSettings as ProtoSharedSettings,
SpecificGrant as ProtoSpecificGrant,
TokenTransferSettings as ProtoTokenTransferSettings,
TransactionRateLimit as ProtoTransactionRateLimit,
VolumeRateLimit as ProtoVolumeRateLimit,
EtherTransferSettings as ProtoEtherTransferSettings, SharedSettings as ProtoSharedSettings,
SpecificGrant as ProtoSpecificGrant, TokenTransferSettings as ProtoTokenTransferSettings,
TransactionRateLimit as ProtoTransactionRateLimit, VolumeRateLimit as ProtoVolumeRateLimit,
specific_grant::Grant as ProtoSpecificGrantType,
};
use arbiter_proto::proto::user_agent::SdkClientWalletAccess;
use alloy::primitives::{Address, U256};
use arbiter_proto::proto::user_agent::{SdkClientWalletAccess, WalletAccess};
use chrono::{DateTime, TimeZone, Utc};
use prost_types::Timestamp as ProtoTimestamp;
use tonic::Status;
use crate::actors::user_agent::EvmAccessEntry;
use crate::db::models::{CoreEvmWalletAccess, NewEvmWallet, NewEvmWalletAccess};
use crate::grpc::Convert;
use crate::{
evm::policies::{
SharedGrantSettings, SpecificGrant, TransactionRateLimit, VolumeRateLimit,
ether_transfer, token_transfers,
SharedGrantSettings, SpecificGrant, TransactionRateLimit, VolumeRateLimit, ether_transfer,
token_transfers,
},
grpc::TryConvert,
};
@@ -79,8 +77,14 @@ impl TryConvert for ProtoSharedSettings {
Ok(SharedGrantSettings {
wallet_access_id: self.wallet_access_id,
chain: self.chain_id,
valid_from: self.valid_from.map(ProtoTimestamp::try_convert).transpose()?,
valid_until: self.valid_until.map(ProtoTimestamp::try_convert).transpose()?,
valid_from: self
.valid_from
.map(ProtoTimestamp::try_convert)
.transpose()?,
valid_until: self
.valid_until
.map(ProtoTimestamp::try_convert)
.transpose()?,
max_gas_fee_per_gas: self
.max_gas_fee_per_gas
.as_deref()
@@ -136,17 +140,29 @@ impl TryConvert for ProtoSpecificGrant {
}
}
impl TryConvert for Vec<SdkClientWalletAccess> {
type Output = Vec<EvmAccessEntry>;
type Error = Status;
impl Convert for WalletAccess {
type Output = NewEvmWalletAccess;
fn try_convert(self) -> Result<Vec<EvmAccessEntry>, Status> {
Ok(self
.into_iter()
.map(|SdkClientWalletAccess { client_id, wallet_id }| EvmAccessEntry {
wallet_id,
sdk_client_id: client_id,
})
.collect())
fn convert(self) -> Self::Output {
NewEvmWalletAccess {
wallet_id: self.wallet_id,
client_id: self.sdk_client_id,
}
}
}
impl TryConvert for SdkClientWalletAccess {
type Output = CoreEvmWalletAccess;
type Error = Status;
fn try_convert(self) -> Result<CoreEvmWalletAccess, Status> {
let Some(access) = self.access else {
return Err(Status::invalid_argument("Missing wallet access entry"));
};
Ok(CoreEvmWalletAccess {
wallet_id: access.wallet_id,
client_id: access.sdk_client_id,
id: self.id,
})
}
}

View File

@@ -5,13 +5,13 @@ use arbiter_proto::proto::{
TransactionRateLimit as ProtoTransactionRateLimit, VolumeRateLimit as ProtoVolumeRateLimit,
specific_grant::Grant as ProtoSpecificGrantType,
},
user_agent::SdkClientWalletAccess as ProtoSdkClientWalletAccess,
user_agent::{SdkClientWalletAccess as ProtoSdkClientWalletAccess, WalletAccess},
};
use chrono::{DateTime, Utc};
use prost_types::Timestamp as ProtoTimestamp;
use crate::{
actors::user_agent::EvmAccessEntry,
db::models::EvmWalletAccess,
evm::policies::{SharedGrantSettings, SpecificGrant, TransactionRateLimit, VolumeRateLimit},
grpc::Convert,
};
@@ -96,13 +96,16 @@ impl Convert for SpecificGrant {
}
}
impl Convert for EvmAccessEntry {
impl Convert for EvmWalletAccess {
type Output = ProtoSdkClientWalletAccess;
fn convert(self) -> Self::Output {
ProtoSdkClientWalletAccess {
client_id: self.sdk_client_id,
wallet_id: self.wallet_id,
Self::Output {
id: self.id,
access: Some(WalletAccess {
wallet_id: self.wallet_id,
sdk_client_id: self.client_id,
}),
}
}
}