From c12c12d73eaaee39d6df984ad0ac4251d74afe6a Mon Sep 17 00:00:00 2001
From: CleverWild <cleverwilddev@gmail.com>
Date: Fri, 19 Jun 2026 22:21:48 +0200
Subject: [PATCH] fix(UA): signing endpoint accepts arbitrary `client_id`

---
 .../src/actors/flow_coordinator/mod.rs        | 13 ++++++++++-
 .../arbiter-server/src/grpc/operator/evm.rs   |  5 +++++
 .../src/peers/client/session.rs               |  2 +-
 .../src/peers/operator/session/handlers.rs    | 22 +++++++++++++++++--
 4 files changed, 38 insertions(+), 4 deletions(-)

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 fa31334..fba78c8 100644
--- a/server/crates/arbiter-server/src/actors/flow_coordinator/mod.rs
+++ b/server/crates/arbiter-server/src/actors/flow_coordinator/mod.rs
@@ -20,6 +20,8 @@ pub mod client_connect_approval;
 
 pub struct FlowCoordinator {
     pub clients: HashMap<ActorId, ActorRef<ClientSession>>,
+    /// Maps DB client_id → ActorId for fast connected-client lookup.
+    client_ids: HashMap<i32, ActorId>,
     operator_registry: ActorRef<OperatorRegistry>,
 }
 
@@ -27,6 +29,7 @@ impl FlowCoordinator {
     pub fn new(operator_registry: ActorRef<OperatorRegistry>) -> Self {
         Self {
             clients: HashMap::default(),
+            client_ids: HashMap::default(),
             operator_registry,
         }
     }
@@ -48,6 +51,7 @@ impl Actor for FlowCoordinator {
         _: ActorStopReason,
     ) -> Result<ControlFlow<ActorStopReason>, Self::Error> {
         if self.clients.remove(&id).is_some() {
+            self.client_ids.retain(|_, actor_id| *actor_id != id);
             info!(
                 ?id,
                 actor = "FlowCoordinator",
@@ -75,14 +79,21 @@ impl FlowCoordinator {
     #[message(ctx)]
     pub async fn register_client(
         &mut self,
+        client_id: i32,
         actor: ActorRef<ClientSession>,
         ctx: &mut Context<Self, ()>,
     ) {
-        info!(id = %actor.id(), actor = "FlowCoordinator", event = "client.connected");
+        info!(id = %actor.id(), client_id, actor = "FlowCoordinator", event = "client.connected");
         ctx.actor_ref().link(&actor).await;
+        self.client_ids.insert(client_id, actor.id());
         self.clients.insert(actor.id(), actor);
     }
 
+    #[message]
+    pub fn is_client_connected(&self, client_id: i32) -> bool {
+        self.client_ids.contains_key(&client_id)
+    }
+
     #[message(ctx)]
     pub async fn request_client_approval(
         &mut self,
diff --git a/server/crates/arbiter-server/src/grpc/operator/evm.rs b/server/crates/arbiter-server/src/grpc/operator/evm.rs
index 0b3ac2c..2a70abb 100644
--- a/server/crates/arbiter-server/src/grpc/operator/evm.rs
+++ b/server/crates/arbiter-server/src/grpc/operator/evm.rs
@@ -217,6 +217,11 @@ async fn handle_sign_transaction(
                 result: Some(vet_error.convert()),
             }
         }
+        Err(kameo::error::SendError::HandlerError(
+            SessionSignTransactionError::ClientNotConnected,
+        )) => {
+            return Err(Status::permission_denied("client not connected"));
+        }
         Err(kameo::error::SendError::HandlerError(SessionSignTransactionError::Internal)) => {
             EvmSignTransactionResponse {
                 result: Some(EvmSignTransactionResult::Error(
diff --git a/server/crates/arbiter-server/src/peers/client/session.rs b/server/crates/arbiter-server/src/peers/client/session.rs
index 23ebf3c..6106a6d 100644
--- a/server/crates/arbiter-server/src/peers/client/session.rs
+++ b/server/crates/arbiter-server/src/peers/client/session.rs
@@ -83,7 +83,7 @@ impl Actor for ClientSession {
         args.props
             .actors
             .flow_coordinator
-            .ask(RegisterClient { actor: this })
+            .ask(RegisterClient { client_id: args.client_id, actor: this })
             .await
             .map_err(|_| Error::ConnectionRegistrationFailed)?;
         Ok(args)
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 5ac6cb4..034fff4 100644
--- a/server/crates/arbiter-server/src/peers/operator/session/handlers.rs
+++ b/server/crates/arbiter-server/src/peers/operator/session/handlers.rs
@@ -4,7 +4,7 @@ use crate::{
         ClientSignTransaction, Generate, ListWallets, OperatorCreateGrant, OperatorListGrants,
         SignTransactionError as EvmSignError,
     },
-    actors::flow_coordinator::client_connect_approval::ClientApprovalAnswer,
+    actors::flow_coordinator::{IsClientConnected, client_connect_approval::ClientApprovalAnswer},
     actors::vault::VaultState,
     db::models::{EvmWalletAccess, NewEvmWalletAccess, ProgramClient, ProgramClientMetadata},
     evm::policies::{Grant, SpecificGrant},
@@ -15,13 +15,16 @@ use alloy::{consensus::TxEip1559, primitives::Address, signers::Signature};
 use diesel::{ExpressionMethods as _, QueryDsl as _, SelectableHelper};
 use diesel_async::{AsyncConnection, RunQueryDsl};
 use kameo::{error::SendError, messages, prelude::Context};
-use tracing::error;
+use tracing::{error, info, warn};
 
 #[derive(Debug, Error)]
 pub enum SignTransactionError {
     #[error("Policy evaluation failed")]
     Vet(#[from] crate::evm::VetError),
 
+    #[error("Client not connected")]
+    ClientNotConnected,
+
     #[error("Internal signing error")]
     Internal,
 }
@@ -141,6 +144,21 @@ impl OperatorSession {
         wallet_address: Address,
         transaction: TxEip1559,
     ) -> Result<Signature, SignTransactionError> {
+        let connected = self
+            .props
+            .actors
+            .flow_coordinator
+            .ask(IsClientConnected { client_id })
+            .await
+            .unwrap_or(false);
+
+        if !connected {
+            warn!(client_id, "operator attempted to sign for disconnected client");
+            return Err(SignTransactionError::ClientNotConnected);
+        }
+
+        info!(client_id, event = "sign_transaction", "operator.sign_transaction");
+
         match self
             .props
             .actors