From b7c4f2e735e6f84316d3e0368960eec2d3eb9150 Mon Sep 17 00:00:00 2001 From: hdbg Date: Tue, 10 Mar 2026 18:56:31 +0100 Subject: [PATCH] feat(evm): add find_all_grants to Policy trait with shared auto_type queries --- server/crates/arbiter-server/src/evm/mod.rs | 34 ++++++- .../crates/arbiter-server/src/evm/policies.rs | 18 ++++ .../src/evm/policies/ether_transfer.rs | 96 +++++++++++++++--- .../src/evm/policies/token_transfers.rs | 99 +++++++++++++++++-- 4 files changed, 224 insertions(+), 23 deletions(-) diff --git a/server/crates/arbiter-server/src/evm/mod.rs b/server/crates/arbiter-server/src/evm/mod.rs index 37bbae4..b0262b9 100644 --- a/server/crates/arbiter-server/src/evm/mod.rs +++ b/server/crates/arbiter-server/src/evm/mod.rs @@ -16,7 +16,7 @@ use crate::{ schema::{self, evm_transaction_log}, }, evm::policies::{ - EvalContext, EvalViolation, FullGrant, Policy, SpecificMeaning, + EvalContext, EvalViolation, FullGrant, Grant, Policy, SpecificGrant, SpecificMeaning, ether_transfer::EtherTransfer, token_transfers::TokenTransfer, }, }; @@ -87,6 +87,17 @@ pub enum CreationError { Database(#[from] diesel::result::Error), } +#[derive(Debug, thiserror::Error, miette::Diagnostic)] +pub enum ListGrantsError { + #[error("Database connection pool error")] + #[diagnostic(code(arbiter_server::evm::list_grants_error::pool))] + Pool(#[from] db::PoolError), + + #[error("Database returned error")] + #[diagnostic(code(arbiter_server::evm::list_grants_error::database))] + Database(#[from] diesel::result::Error), +} + /// Controls whether a transaction should be executed or only validated #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RunKind { @@ -201,6 +212,27 @@ impl Engine { Ok(id) } + pub async fn list_all_grants(&self) -> Result>, ListGrantsError> { + let mut conn = self.db.get().await?; + + let mut grants: Vec> = Vec::new(); + + grants.extend( + EtherTransfer::find_all_grants(&mut conn) + .await? + .into_iter() + .map(Grant::from), + ); + grants.extend( + TokenTransfer::find_all_grants(&mut conn) + .await? + .into_iter() + .map(Grant::from), + ); + + Ok(grants) + } + pub async fn evaluate_transaction( &self, wallet_id: i32, diff --git a/server/crates/arbiter-server/src/evm/policies.rs b/server/crates/arbiter-server/src/evm/policies.rs index 5650f0f..086de24 100644 --- a/server/crates/arbiter-server/src/evm/policies.rs +++ b/server/crates/arbiter-server/src/evm/policies.rs @@ -97,6 +97,11 @@ pub trait Policy: Sized { conn: &mut impl AsyncConnection, ) -> impl Future>>> + Send; + // Return all non-revoked grants, eagerly loading policy-specific settings + fn find_all_grants( + conn: &mut impl AsyncConnection, + ) -> impl Future>>> + Send; + // Records, updates or deletes rate limits // In other words, records grant-specific things after transaction is executed fn record_transaction( @@ -192,6 +197,19 @@ pub enum SpecificGrant { TokenTransfer(token_transfers::Settings), } +/// Blanket conversion from a typed `Grant` into `Grant`. +/// Lets the engine collect across all policies into one `Vec>`. +impl> From> for Grant { + fn from(g: Grant) -> Self { + Grant { + id: g.id, + shared_grant_id: g.shared_grant_id, + shared: g.shared, + settings: g.settings.into(), + } + } +} + pub struct FullGrant { pub basic: SharedGrantSettings, pub specific: PolicyGrant, diff --git a/server/crates/arbiter-server/src/evm/policies/ether_transfer.rs b/server/crates/arbiter-server/src/evm/policies/ether_transfer.rs index bc67c33..1333932 100644 --- a/server/crates/arbiter-server/src/evm/policies/ether_transfer.rs +++ b/server/crates/arbiter-server/src/evm/policies/ether_transfer.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::fmt::Display; use alloy::primitives::{Address, U256}; @@ -11,7 +12,7 @@ use crate::db::models::{ EvmBasicGrant, EvmEtherTransferGrant, EvmEtherTransferGrantTarget, EvmEtherTransferLimit, NewEvmEtherTransferLimit, SqliteTimestamp, }; -use crate::db::schema::{evm_ether_transfer_limit, evm_transaction_log}; +use crate::db::schema::{evm_basic_grant, evm_ether_transfer_limit, evm_transaction_log}; use crate::evm::policies::{ Grant, SharedGrantSettings, SpecificGrant, SpecificMeaning, VolumeRateLimit, }; @@ -23,6 +24,14 @@ use crate::{ evm::{policies::Policy, utils}, }; +#[diesel::auto_type] +fn grant_join() -> _ { + evm_ether_transfer_grant::table.inner_join( + evm_basic_grant::table + .on(evm_ether_transfer_grant::basic_grant_id.eq(evm_basic_grant::id)), + ) +} + use super::{DatabaseID, EvalContext, EvalViolation}; // Plain ether transfer @@ -183,24 +192,12 @@ impl Policy for EtherTransfer { context: &EvalContext, conn: &mut impl AsyncConnection, ) -> diesel::result::QueryResult>> { - use crate::db::schema::{ - evm_basic_grant, evm_ether_transfer_grant, evm_ether_transfer_grant_target, - }; - let target_bytes = context.to.to_vec(); // Find a grant where: // 1. The basic grant's wallet_id and client_id match the context // 2. Any of the grant's targets match the context's `to` address - let grant: Option<(EvmBasicGrant, EvmEtherTransferGrant)> = evm_ether_transfer_grant::table - .inner_join( - evm_basic_grant::table - .on(evm_ether_transfer_grant::basic_grant_id.eq(evm_basic_grant::id)), - ) - .inner_join( - evm_ether_transfer_grant_target::table - .on(evm_ether_transfer_grant::id.eq(evm_ether_transfer_grant_target::grant_id)), - ) + let grant: Option<(EvmBasicGrant, EvmEtherTransferGrant)> = grant_join() .filter(evm_basic_grant::wallet_id.eq(context.wallet_id)) .filter(evm_basic_grant::client_id.eq(context.client_id)) .filter(evm_ether_transfer_grant_target::address.eq(&target_bytes)) @@ -266,4 +263,75 @@ impl Policy for EtherTransfer { Ok(()) } + + async fn find_all_grants( + conn: &mut impl AsyncConnection, + ) -> QueryResult>> { + let grants: Vec<(EvmBasicGrant, EvmEtherTransferGrant)> = grant_join() + .filter(evm_basic_grant::revoked_at.is_null()) + .select((EvmBasicGrant::as_select(), EvmEtherTransferGrant::as_select())) + .load(conn) + .await?; + + if grants.is_empty() { + return Ok(Vec::new()); + } + + let grant_ids: Vec = grants.iter().map(|(_, g)| g.id).collect(); + let limit_ids: Vec = grants.iter().map(|(_, g)| g.limit_id).collect(); + + let all_targets: Vec = evm_ether_transfer_grant_target::table + .filter(evm_ether_transfer_grant_target::grant_id.eq_any(&grant_ids)) + .select(EvmEtherTransferGrantTarget::as_select()) + .load(conn) + .await?; + + let all_limits: Vec = evm_ether_transfer_limit::table + .filter(evm_ether_transfer_limit::id.eq_any(&limit_ids)) + .select(EvmEtherTransferLimit::as_select()) + .load(conn) + .await?; + + let mut targets_by_grant: HashMap> = HashMap::new(); + for target in all_targets { + targets_by_grant.entry(target.grant_id).or_default().push(target); + } + + let limits_by_id: HashMap = + all_limits.into_iter().map(|l| (l.id, l)).collect(); + + grants + .into_iter() + .map(|(basic, specific)| { + let targets: Vec
= targets_by_grant + .get(&specific.id) + .map(|v| v.as_slice()) + .unwrap_or_default() + .iter() + .filter_map(|t| { + let arr: [u8; 20] = t.address.clone().try_into().ok()?; + Some(Address::from(arr)) + }) + .collect(); + + let limit = limits_by_id + .get(&specific.limit_id) + .ok_or(diesel::result::Error::NotFound)?; + + Ok(Grant { + id: specific.id, + shared_grant_id: specific.basic_grant_id, + shared: SharedGrantSettings::try_from_model(basic)?, + settings: Settings { + target: targets, + limit: VolumeRateLimit { + max_volume: utils::try_bytes_to_u256(&limit.max_volume) + .map_err(|e| diesel::result::Error::DeserializationError(Box::new(e)))?, + window: Duration::seconds(limit.window_secs as i64), + }, + }, + }) + }) + .collect() + } } diff --git a/server/crates/arbiter-server/src/evm/policies/token_transfers.rs b/server/crates/arbiter-server/src/evm/policies/token_transfers.rs index f21fc0c..9ef3dc0 100644 --- a/server/crates/arbiter-server/src/evm/policies/token_transfers.rs +++ b/server/crates/arbiter-server/src/evm/policies/token_transfers.rs @@ -1,3 +1,5 @@ +use std::collections::HashMap; + use alloy::{ primitives::{Address, U256}, sol_types::SolCall, @@ -14,7 +16,8 @@ use crate::db::models::{ NewEvmTokenTransferLog, NewEvmTokenTransferVolumeLimit, SqliteTimestamp, }; use crate::db::schema::{ - evm_token_transfer_grant, evm_token_transfer_log, evm_token_transfer_volume_limit, + evm_basic_grant, evm_token_transfer_grant, evm_token_transfer_log, + evm_token_transfer_volume_limit, }; use crate::evm::{ abi::IERC20::transferCall, @@ -24,6 +27,14 @@ use crate::evm::{ use super::{DatabaseID, EvalContext, EvalViolation}; +#[diesel::auto_type] +fn grant_join() -> _ { + evm_token_transfer_grant::table.inner_join( + evm_basic_grant::table + .on(evm_token_transfer_grant::basic_grant_id.eq(evm_basic_grant::id)), + ) +} + #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct Meaning { token: &'static TokenInfo, @@ -192,15 +203,9 @@ impl Policy for TokenTransfer { context: &EvalContext, conn: &mut impl AsyncConnection, ) -> QueryResult>> { - use crate::db::schema::{evm_basic_grant, evm_token_transfer_grant}; - let token_contract_bytes = context.to.to_vec(); - let grant: Option<(EvmBasicGrant, EvmTokenTransferGrant)> = evm_token_transfer_grant::table - .inner_join( - evm_basic_grant::table - .on(evm_token_transfer_grant::basic_grant_id.eq(evm_basic_grant::id)), - ) + let grant: Option<(EvmBasicGrant, EvmTokenTransferGrant)> = grant_join() .filter(evm_basic_grant::wallet_id.eq(context.wallet_id)) .filter(evm_basic_grant::client_id.eq(context.client_id)) .filter(evm_token_transfer_grant::token_contract.eq(&token_contract_bytes)) @@ -288,4 +293,82 @@ impl Policy for TokenTransfer { Ok(()) } + + async fn find_all_grants( + conn: &mut impl AsyncConnection, + ) -> QueryResult>> { + let grants: Vec<(EvmBasicGrant, EvmTokenTransferGrant)> = grant_join() + .filter(evm_basic_grant::revoked_at.is_null()) + .select((EvmBasicGrant::as_select(), EvmTokenTransferGrant::as_select())) + .load(conn) + .await?; + + if grants.is_empty() { + return Ok(Vec::new()); + } + + let grant_ids: Vec = grants.iter().map(|(_, g)| g.id).collect(); + + let all_volume_limits: Vec = + evm_token_transfer_volume_limit::table + .filter(evm_token_transfer_volume_limit::grant_id.eq_any(&grant_ids)) + .select(EvmTokenTransferVolumeLimit::as_select()) + .load(conn) + .await?; + + let mut limits_by_grant: HashMap> = HashMap::new(); + for limit in all_volume_limits { + limits_by_grant.entry(limit.grant_id).or_default().push(limit); + } + + grants + .into_iter() + .map(|(basic, specific)| { + let volume_limits: Vec = limits_by_grant + .get(&specific.id) + .map(|v| v.as_slice()) + .unwrap_or_default() + .iter() + .map(|row| { + Ok(VolumeRateLimit { + max_volume: utils::try_bytes_to_u256(&row.max_volume) + .map_err(|e| diesel::result::Error::DeserializationError(Box::new(e)))?, + window: Duration::seconds(row.window_secs as i64), + }) + }) + .collect::>>()?; + + let token_contract: [u8; 20] = specific + .token_contract + .clone() + .try_into() + .map_err(|_| diesel::result::Error::DeserializationError( + "Invalid token contract address length".into(), + ))?; + + let target: Option
= match &specific.receiver { + None => None, + Some(bytes) => { + let arr: [u8; 20] = bytes.clone().try_into().map_err(|_| { + diesel::result::Error::DeserializationError( + "Invalid receiver address length".into(), + ) + })?; + Some(Address::from(arr)) + } + }; + + Ok(Grant { + id: specific.id, + shared_grant_id: specific.basic_grant_id, + shared: SharedGrantSettings::try_from_model(basic)?, + settings: Settings { + token_contract: Address::from(token_contract), + target, + volume_limits, + }, + }) + }) + .collect() + } }