Merge remote-tracking branch 'origin/main' into SDK-client-UA-registration

server/crates/arbiter-server/src/actors/client/auth.rs

@@ -1,18 +1,12 @@
 use arbiter_proto::{
     format_challenge,
-    proto::client::{
-        AuthChallenge, AuthChallengeSolution, AuthOk, ClientConnectError, ClientRequest,
-        ClientResponse, client_connect_error::Code as ConnectErrorCode,
-        client_request::Payload as ClientRequestPayload,
-        client_response::Payload as ClientResponsePayload,
-    },
-    transport::expect_message,
+    transport::{Bi, expect_message},
 };
 use diesel::{
     ExpressionMethods as _, OptionalExtension as _, QueryDsl as _, dsl::insert_into, update,
 };
 use diesel_async::RunQueryDsl as _;
-use ed25519_dalek::VerifyingKey;
+use ed25519_dalek::{Signature, VerifyingKey};
 use kameo::error::SendError;
 use tracing::error;
 
@@ -24,35 +18,8 @@ use crate::{
     db::{self, schema::program_client},
 };
 
-use super::session::ClientSession;
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub struct ClientId(i32);
-
-impl ClientId {
-    pub fn new(raw: i32) -> Self {
-        Self(raw)
-    }
-
-    pub fn as_i32(self) -> i32 {
-        self.0
-    }
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-struct ClientNonceState {
-    client_id: ClientId,
-    nonce: i32,
-}
-
 #[derive(thiserror::Error, Debug, Clone, PartialEq, Eq)]
 pub enum Error {
-    #[error("Unexpected message payload")]
-    UnexpectedMessagePayload,
-    #[error("Invalid client public key length")]
-    InvalidClientPubkeyLength,
-    #[error("Invalid client public key encoding")]
-    InvalidAuthPubkeyEncoding,
     #[error("Database pool unavailable")]
     DatabasePoolUnavailable,
     #[error("Database operation failed")]
@@ -61,8 +28,6 @@ pub enum Error {
     InvalidChallengeSolution,
     #[error("Client approval request failed")]
     ApproveError(#[from] ApproveError),
-    #[error("Internal error")]
-    InternalError,
     #[error("Transport error")]
     Transport,
 }
@@ -77,12 +42,21 @@ pub enum ApproveError {
     Upstream(router::ApprovalError),
 }
 
+#[derive(Debug, Clone)]
+pub enum Inbound {
+    AuthChallengeRequest { pubkey: VerifyingKey },
+    AuthChallengeSolution { signature: Signature },
+}
+
+#[derive(Debug, Clone)]
+pub enum Outbound {
+    AuthChallenge { pubkey: VerifyingKey, nonce: i32 },
+    AuthSuccess,
+}
+
 /// Atomically reads and increments the nonce for a known client.
 /// Returns `None` if the pubkey is not registered.
-async fn get_nonce(
-    db: &db::DatabasePool,
-    pubkey: &VerifyingKey,
-) -> Result<Option<ClientNonceState>, Error> {
+async fn get_nonce(db: &db::DatabasePool, pubkey: &VerifyingKey) -> Result<Option<i32>, Error> {
     let pubkey_bytes = pubkey.as_bytes().to_vec();
 
     let mut conn = db.get().await.map_err(|e| {
@@ -109,10 +83,8 @@ async fn get_nonce(
                 .execute(conn)
                 .await?;
 
-            Ok(Some(ClientNonceState {
-                client_id: ClientId::new(client_id),
-                nonce: current_nonce,
-            }))
+            let _ = client_id;
+            Ok(Some(current_nonce))
         })
     })
     .await
@@ -148,7 +120,7 @@ async fn approve_new_client(
 }
 
 enum InsertClientResult {
-    Inserted(ClientId),
+    Inserted,
     AlreadyExists,
 }
 
@@ -198,126 +170,80 @@ async fn insert_client(
             Error::DatabaseOperationFailed
         })?;
 
-    Ok(InsertClientResult::Inserted(ClientId::new(client_id)))
+    let _ = client_id;
+    Ok(InsertClientResult::Inserted)
 }
 
-async fn challenge_client(
-    props: &mut ClientConnection,
+async fn challenge_client<T>(
+    transport: &mut T,
     pubkey: VerifyingKey,
     nonce: i32,
-) -> Result<(), Error> {
-    let challenge = AuthChallenge {
-        pubkey: pubkey.as_bytes().to_vec(),
-        nonce,
-    };
-
-    props
-        .transport
-        .send(Ok(ClientResponse {
-            payload: Some(ClientResponsePayload::AuthChallenge(challenge.clone())),
-        }))
+) -> Result<(), Error>
+where
+    T: Bi<Inbound, Result<Outbound, Error>> + ?Sized,
+{
+    transport
+        .send(Ok(Outbound::AuthChallenge { pubkey, nonce }))
         .await
         .map_err(|e| {
             error!(error = ?e, "Failed to send auth challenge");
             Error::Transport
         })?;
 
-    let AuthChallengeSolution { signature } =
-        expect_message(&mut *props.transport, |req: ClientRequest| {
-            match req.payload? {
-                ClientRequestPayload::AuthChallengeSolution(s) => Some(s),
-                _ => None,
-            }
-        })
-        .await
-        .map_err(|e| {
-            error!(error = ?e, "Failed to receive challenge solution");
-            Error::Transport
-        })?;
-
-    let formatted = format_challenge(nonce, &challenge.pubkey);
-    let sig = signature.as_slice().try_into().map_err(|_| {
-        error!("Invalid signature length");
-        Error::InvalidChallengeSolution
+    let signature = expect_message(transport, |req: Inbound| match req {
+        Inbound::AuthChallengeSolution { signature } => Some(signature),
+        _ => None,
+    })
+    .await
+    .map_err(|e| {
+        error!(error = ?e, "Failed to receive challenge solution");
+        Error::Transport
     })?;
 
-    pubkey.verify_strict(&formatted, &sig).map_err(|_| {
+    let formatted = format_challenge(nonce, pubkey.as_bytes());
+
+    pubkey.verify_strict(&formatted, &signature).map_err(|_| {
         error!("Challenge solution verification failed");
         Error::InvalidChallengeSolution
     })?;
 
-    props
-        .transport
-        .send(Ok(ClientResponse {
-            payload: Some(ClientResponsePayload::AuthOk(AuthOk {})),
-        }))
-        .await
-        .map_err(|e| {
-            error!(error = ?e, "Failed to send auth ok");
-            Error::Transport
-        })?;
-
     Ok(())
 }
 
-fn connect_error_code(err: &Error) -> ConnectErrorCode {
-    match err {
-        Error::ApproveError(ApproveError::Denied) => ConnectErrorCode::ApprovalDenied,
-        Error::ApproveError(ApproveError::Upstream(
-            router::ApprovalError::NoUserAgentsConnected,
-        )) => ConnectErrorCode::NoUserAgentsOnline,
-        _ => ConnectErrorCode::Unknown,
-    }
-}
-
-async fn authenticate(props: &mut ClientConnection) -> Result<(VerifyingKey, ClientId), Error> {
-    let Some(ClientRequest {
-        payload: Some(ClientRequestPayload::AuthChallengeRequest(challenge)),
-    }) = props.transport.recv().await
+pub async fn authenticate<T>(
+    props: &mut ClientConnection,
+    transport: &mut T,
+) -> Result<VerifyingKey, Error>
+where
+    T: Bi<Inbound, Result<Outbound, Error>> + Send + ?Sized,
+{
+    let Some(Inbound::AuthChallengeRequest { pubkey }) = transport.recv().await
     else {
         return Err(Error::Transport);
     };
 
-    let pubkey_bytes = challenge
-        .pubkey
-        .as_array()
-        .ok_or(Error::InvalidClientPubkeyLength)?;
-    let pubkey =
-        VerifyingKey::from_bytes(pubkey_bytes).map_err(|_| Error::InvalidAuthPubkeyEncoding)?;
-
-    let (client_id, nonce) = match get_nonce(&props.db, &pubkey).await? {
-        Some(state) => (state.client_id, state.nonce),
+    let nonce = match get_nonce(&props.db, &pubkey).await? {
+        Some(nonce) => nonce,
         None => {
             approve_new_client(&props.actors, pubkey).await?;
             match insert_client(&props.db, &pubkey).await? {
-                InsertClientResult::Inserted(client_id) => (client_id, 0),
+                InsertClientResult::Inserted => 0,
                 InsertClientResult::AlreadyExists => match get_nonce(&props.db, &pubkey).await? {
-                    Some(state) => (state.client_id, state.nonce),
-                    None => return Err(Error::InternalError),
+                    Some(nonce) => nonce,
+                    None => return Err(Error::DatabaseOperationFailed),
                 },
             }
         }
     };
 
-    challenge_client(props, pubkey, nonce).await?;
+    challenge_client(transport, pubkey, nonce).await?;
+    transport
+        .send(Ok(Outbound::AuthSuccess))
+        .await
+        .map_err(|e| {
+            error!(error = ?e, "Failed to send auth success");
+            Error::Transport
+        })?;
 
-    Ok((pubkey, client_id))
-}
-
-pub async fn authenticate_and_create(mut props: ClientConnection) -> Result<ClientSession, Error> {
-    match authenticate(&mut props).await {
-        Ok((_pubkey, client_id)) => Ok(ClientSession::new(props, client_id)),
-        Err(err) => {
-            let code = connect_error_code(&err);
-            let _ = props
-                .transport
-                .send(Ok(ClientResponse {
-                    payload: Some(ClientResponsePayload::ClientConnectError(
-                        ClientConnectError { code: code.into() },
-                    )),
-                }))
-                .await;
-            Err(err)
-        }
-    }
+    Ok(pubkey)
 }

server/crates/arbiter-server/src/actors/client/mod.rs

@@ -1,7 +1,4 @@
-use arbiter_proto::{
-    proto::client::{ClientRequest, ClientResponse},
-    transport::Bi,
-};
+use arbiter_proto::transport::Bi;
 use kameo::actor::Spawn;
 use tracing::{error, info};
 
@@ -10,48 +7,31 @@ use crate::{
     db,
 };
 
-#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
-pub enum ClientError {
-    #[error("Expected message with payload")]
-    MissingRequestPayload,
-    #[error("Unexpected request payload")]
-    UnexpectedRequestPayload,
-    #[error("State machine error")]
-    StateTransitionFailed,
-    #[error("Connection registration failed")]
-    ConnectionRegistrationFailed,
-    #[error(transparent)]
-    Auth(#[from] auth::Error),
-}
-
-pub type Transport = Box<dyn Bi<ClientRequest, Result<ClientResponse, ClientError>> + Send>;
-
 pub struct ClientConnection {
     pub(crate) db: db::DatabasePool,
-    pub(crate) transport: Transport,
     pub(crate) actors: GlobalActors,
 }
 
 impl ClientConnection {
-    pub fn new(db: db::DatabasePool, transport: Transport, actors: GlobalActors) -> Self {
-        Self {
-            db,
-            transport,
-            actors,
-        }
+    pub fn new(db: db::DatabasePool, actors: GlobalActors) -> Self {
+        Self { db, actors }
     }
 }
 
 pub mod auth;
 pub mod session;
 
-pub async fn connect_client(props: ClientConnection) {
-    match auth::authenticate_and_create(props).await {
-        Ok(session) => {
-            ClientSession::spawn(session);
+pub async fn connect_client<T>(mut props: ClientConnection, transport: &mut T)
+where
+    T: Bi<auth::Inbound, Result<auth::Outbound, auth::Error>> + Send + ?Sized,
+{
+    match auth::authenticate(&mut props, transport).await {
+        Ok(_pubkey) => {
+            ClientSession::spawn(ClientSession::new(props));
             info!("Client authenticated, session started");
         }
         Err(err) => {
+            let _ = transport.send(Err(err.clone())).await;
             error!(?err, "Authentication failed, closing connection");
         }
     }

server/crates/arbiter-server/src/actors/client/session.rs

@@ -1,92 +1,45 @@
-use alloy::{consensus::TxEip1559, primitives::Address, rlp::Decodable};
-use arbiter_proto::proto::{
-    client::{
-        ClientRequest, ClientResponse, client_request::Payload as ClientRequestPayload,
-        client_response::Payload as ClientResponsePayload,
-    },
-    evm::{
-        EvmError, EvmSignTransactionResponse, evm_sign_transaction_response::Result as SignResult,
-    },
-};
-use kameo::Actor;
-use tokio::select;
-use tracing::{error, info};
+use kameo::{Actor, messages};
+use tracing::error;
 
 use crate::{
     actors::{
-        GlobalActors,
-        client::{ClientConnection, ClientError, auth::ClientId},
-        evm::ClientSignTransaction,
-        router::RegisterClient,
+        GlobalActors, client::ClientConnection, keyholder::KeyHolderState, router::RegisterClient,
     },
     db,
 };
 
 pub struct ClientSession {
     props: ClientConnection,
-    client_id: ClientId,
 }
 
 impl ClientSession {
-    pub(crate) fn new(props: ClientConnection, client_id: ClientId) -> Self {
-        Self { props, client_id }
-    }
-
-    pub async fn process_transport_inbound(&mut self, req: ClientRequest) -> Output {
-        let msg = req.payload.ok_or_else(|| {
-            error!(actor = "client", "Received message with no payload");
-            ClientError::MissingRequestPayload
-        })?;
-
-        match msg {
-            ClientRequestPayload::EvmSignTransaction(sign_req) => {
-                let wallet_address: [u8; 20] = sign_req
-                    .wallet_address
-                    .try_into()
-                    .map_err(|_| ClientError::UnexpectedRequestPayload)?;
-
-                let mut rlp_bytes: &[u8] = &sign_req.rlp_transaction;
-                let tx = TxEip1559::decode(&mut rlp_bytes)
-                    .map_err(|_| ClientError::UnexpectedRequestPayload)?;
-
-                let result = self
-                    .props
-                    .actors
-                    .evm
-                    .ask(ClientSignTransaction {
-                        client_id: self.client_id.as_i32(),
-                        wallet_address: Address::from_slice(&wallet_address),
-                        transaction: tx,
-                    })
-                    .await;
-
-                let response_result = match result {
-                    Ok(signature) => SignResult::Signature(signature.as_bytes().to_vec()),
-                    Err(err) => {
-                        error!(?err, "client sign transaction failed");
-                        SignResult::Error(EvmError::Internal.into())
-                    }
-                };
-
-                Ok(ClientResponse {
-                    payload: Some(ClientResponsePayload::EvmSignTransaction(
-                        EvmSignTransactionResponse {
-                            result: Some(response_result),
-                        },
-                    )),
-                })
-            }
-            _ => Err(ClientError::UnexpectedRequestPayload),
-        }
+    pub(crate) fn new(props: ClientConnection) -> Self {
+        Self { props }
     }
 }
 
-type Output = Result<ClientResponse, ClientError>;
+#[messages]
+impl ClientSession {
+    #[message]
+    pub(crate) async fn handle_query_vault_state(&mut self) -> Result<KeyHolderState, Error> {
+        use crate::actors::keyholder::GetState;
+
+        let vault_state = match self.props.actors.key_holder.ask(GetState {}).await {
+            Ok(state) => state,
+            Err(err) => {
+                error!(?err, actor = "client", "keyholder.query.failed");
+                return Err(Error::Internal);
+            }
+        };
+
+        Ok(vault_state)
+    }
+}
 
 impl Actor for ClientSession {
     type Args = Self;
 
-    type Error = ClientError;
+    type Error = Error;
 
     async fn on_start(
         args: Self::Args,
@@ -97,55 +50,22 @@ impl Actor for ClientSession {
             .router
             .ask(RegisterClient { actor: this })
             .await
-            .map_err(|_| ClientError::ConnectionRegistrationFailed)?;
+            .map_err(|_| Error::ConnectionRegistrationFailed)?;
         Ok(args)
     }
-
-    async fn next(
-        &mut self,
-        _actor_ref: kameo::prelude::WeakActorRef<Self>,
-        mailbox_rx: &mut kameo::prelude::MailboxReceiver<Self>,
-    ) -> Option<kameo::mailbox::Signal<Self>> {
-        loop {
-            select! {
-                signal = mailbox_rx.recv() => {
-                    return signal;
-                }
-                msg = self.props.transport.recv() => {
-                    match msg {
-                        Some(request) => {
-                            match self.process_transport_inbound(request).await {
-                                Ok(resp) => {
-                                    if self.props.transport.send(Ok(resp)).await.is_err() {
-                                        error!(actor = "client", reason = "channel closed", "send.failed");
-                                        return Some(kameo::mailbox::Signal::Stop);
-                                    }
-                                }
-                                Err(err) => {
-                                    let _ = self.props.transport.send(Err(err)).await;
-                                    return Some(kameo::mailbox::Signal::Stop);
-                                }
-                            }
-                        }
-                        None => {
-                            info!(actor = "client", "transport.closed");
-                            return Some(kameo::mailbox::Signal::Stop);
-                        }
-                    }
-                }
-            }
-        }
-    }
 }
 
 impl ClientSession {
     pub fn new_test(db: db::DatabasePool, actors: GlobalActors) -> Self {
-        use arbiter_proto::transport::DummyTransport;
-        let transport: super::Transport = Box::new(DummyTransport::new());
-        let props = ClientConnection::new(db, transport, actors);
-        Self {
-            props,
-            client_id: ClientId::new(0),
-        }
+        let props = ClientConnection::new(db, actors);
+        Self { props }
     }
 }
+
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("Connection registration failed")]
+    ConnectionRegistrationFailed,
+    #[error("Internal error")]
+    Internal,
+}