diff --git a/server/crates/arbiter-server/src/actors/flow_coordinator/mod.rs b/server/crates/arbiter-server/src/actors/flow_coordinator/mod.rs index fba78c8..c9f8f5b 100644 --- a/server/crates/arbiter-server/src/actors/flow_coordinator/mod.rs +++ b/server/crates/arbiter-server/src/actors/flow_coordinator/mod.rs @@ -94,6 +94,13 @@ impl FlowCoordinator { self.client_ids.contains_key(&client_id) } + /// Returns the DB client_ids of all currently connected SDK clients. + /// Used by operator sessions on startup to seed their approved-client set. + #[message] + pub fn get_connected_client_ids(&self) -> Vec { + self.client_ids.keys().copied().collect() + } + #[message(ctx)] pub async fn request_client_approval( &mut self, diff --git a/server/crates/arbiter-server/src/peers/operator/session/handlers.rs b/server/crates/arbiter-server/src/peers/operator/session/handlers.rs index 9565146..4bad681 100644 --- a/server/crates/arbiter-server/src/peers/operator/session/handlers.rs +++ b/server/crates/arbiter-server/src/peers/operator/session/handlers.rs @@ -6,7 +6,10 @@ use crate::{ }, actors::flow_coordinator::{IsClientConnected, client_connect_approval::ClientApprovalAnswer}, actors::vault::VaultState, - db::models::{EvmWalletAccess, NewEvmWalletAccess, ProgramClient, ProgramClientMetadata}, + db::{ + models::{EvmWalletAccess, NewEvmWalletAccess, ProgramClient, ProgramClientMetadata}, + schema::program_client, + }, evm::policies::{Grant, SpecificGrant}, }; use arbiter_crypto::authn; @@ -144,6 +147,14 @@ impl OperatorSession { wallet_address: Address, transaction: TxEip1559, ) -> Result { + if !self.approved_client_ids.contains(&client_id) { + warn!( + client_id, + "operator attempted to sign for client not in its approved set" + ); + return Err(SignTransactionError::ClientNotConnected); + } + let connected = self .props .actors @@ -153,6 +164,7 @@ impl OperatorSession { .unwrap_or(false); if !connected { + self.approved_client_ids.remove(&client_id); warn!(client_id, "operator attempted to sign for disconnected client"); return Err(SignTransactionError::ClientNotConnected); } @@ -267,6 +279,30 @@ impl OperatorSession { ctx.actor_ref().unlink(&pending_approval.controller).await; + if approved { + let pubkey_bytes = pending_approval.pubkey.to_bytes(); + match self.props.db.get().await { + Ok(mut conn) => { + match program_client::table + .filter(program_client::public_key.eq(pubkey_bytes.as_slice())) + .select(program_client::id) + .first::(&mut conn) + .await + { + Ok(client_id) => { + self.approved_client_ids.insert(client_id); + } + Err(err) => { + error!(?err, "Failed to look up client_id for approved pubkey"); + } + } + } + Err(err) => { + error!(?err, "DB pool error after client approval"); + } + } + } + Ok(()) } diff --git a/server/crates/arbiter-server/src/peers/operator/session/mod.rs b/server/crates/arbiter-server/src/peers/operator/session/mod.rs index 0fe2c84..eab628a 100644 --- a/server/crates/arbiter-server/src/peers/operator/session/mod.rs +++ b/server/crates/arbiter-server/src/peers/operator/session/mod.rs @@ -1,7 +1,7 @@ use super::{OutOfBand, OperatorConnection}; use crate::{ actors::{ - flow_coordinator::client_connect_approval::ClientApprovalController, + flow_coordinator::{GetConnectedClientIds, client_connect_approval::ClientApprovalController}, operator_registry::ConnectOperator, }, peers::client::ClientProfile, @@ -10,7 +10,7 @@ use arbiter_crypto::authn; use arbiter_proto::transport::Sender; use kameo::{Actor, actor::ActorRef, messages}; -use std::{borrow::Cow, collections::HashMap}; +use std::{borrow::Cow, collections::{HashMap, HashSet}}; use thiserror::Error; use tracing::error; @@ -54,6 +54,10 @@ pub struct OperatorSession { sender: Box>, pending_client_approvals: HashMap, PendingClientApproval>, + /// DB client_ids this operator session is allowed to sign for. + /// Seeded from currently-connected clients on start, then updated as + /// approvals are granted or denied during the session lifetime. + approved_client_ids: HashSet, } pub mod handlers; @@ -64,6 +68,7 @@ impl OperatorSession { props, sender, pending_client_approvals: HashMap::default(), + approved_client_ids: HashSet::default(), } } } @@ -106,7 +111,7 @@ impl Actor for OperatorSession { type Error = Error; - async fn on_start(args: Self::Args, this: ActorRef) -> Result { + async fn on_start(mut args: Self::Args, this: ActorRef) -> Result { args.props .actors .operator_registry @@ -121,6 +126,16 @@ impl Actor for OperatorSession { ); Error::internal("Failed to register operator connection with operator registry") })?; + + // Seed approved set with clients already connected when this session starts. + // New clients will be added via handle_new_client_approve as they are approved. + match args.props.actors.flow_coordinator.ask(GetConnectedClientIds {}).await { + Ok(ids) => args.approved_client_ids.extend(ids), + Err(err) => { + error!(?err, "Failed to fetch connected client IDs on operator session start"); + } + } + Ok(args) }