feat(client-auth): emit and require AuthOk for SDK client challenge flow

This commit is contained in:
CleverWild
2026-03-19 19:06:27 +01:00
parent ec70561c93
commit e1a8553142
3 changed files with 63 additions and 10 deletions

View File

@@ -367,7 +367,9 @@ async fn receive_auth_confirmation(
.await .await
.map_err(|_| ConnectError::UnexpectedAuthResponse)?; .map_err(|_| ConnectError::UnexpectedAuthResponse)?;
let payload = response.payload.ok_or(ConnectError::UnexpectedAuthResponse)?; let payload = response
.payload
.ok_or(ConnectError::UnexpectedAuthResponse)?;
match payload { match payload {
ClientResponsePayload::AuthOk(_) => Ok(()), ClientResponsePayload::AuthOk(_) => Ok(()),
ClientResponsePayload::ClientConnectError(err) => Err(map_connect_error(err.code)), ClientResponsePayload::ClientConnectError(err) => Err(map_connect_error(err.code)),

View File

@@ -1,8 +1,8 @@
use arbiter_proto::{ use arbiter_proto::{
format_challenge, format_challenge,
proto::client::{ proto::client::{
AuthChallenge, AuthChallengeSolution, ClientConnectError, ClientRequest, ClientResponse, AuthChallenge, AuthChallengeSolution, AuthOk, ClientConnectError, ClientRequest,
client_connect_error::Code as ConnectErrorCode, ClientResponse, client_connect_error::Code as ConnectErrorCode,
client_request::Payload as ClientRequestPayload, client_request::Payload as ClientRequestPayload,
client_response::Payload as ClientResponsePayload, client_response::Payload as ClientResponsePayload,
}, },
@@ -26,6 +26,25 @@ use crate::{
use super::session::ClientSession; 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)] #[derive(thiserror::Error, Debug, Clone, PartialEq, Eq)]
pub enum Error { pub enum Error {
#[error("Unexpected message payload")] #[error("Unexpected message payload")]
@@ -63,7 +82,7 @@ pub enum ApproveError {
async fn get_nonce( async fn get_nonce(
db: &db::DatabasePool, db: &db::DatabasePool,
pubkey: &VerifyingKey, pubkey: &VerifyingKey,
) -> Result<Option<(i32, i32)>, Error> { ) -> Result<Option<ClientNonceState>, Error> {
let pubkey_bytes = pubkey.as_bytes().to_vec(); let pubkey_bytes = pubkey.as_bytes().to_vec();
let mut conn = db.get().await.map_err(|e| { let mut conn = db.get().await.map_err(|e| {
@@ -90,7 +109,10 @@ async fn get_nonce(
.execute(conn) .execute(conn)
.await?; .await?;
Ok(Some((client_id, current_nonce))) Ok(Some(ClientNonceState {
client_id: ClientId::new(client_id),
nonce: current_nonce,
}))
}) })
}) })
.await .await
@@ -126,7 +148,7 @@ async fn approve_new_client(
} }
enum InsertClientResult { enum InsertClientResult {
Inserted(i32), Inserted(ClientId),
AlreadyExists, AlreadyExists,
} }
@@ -176,7 +198,7 @@ async fn insert_client(
Error::DatabaseOperationFailed Error::DatabaseOperationFailed
})?; })?;
Ok(InsertClientResult::Inserted(client_id)) Ok(InsertClientResult::Inserted(ClientId::new(client_id)))
} }
async fn challenge_client( async fn challenge_client(
@@ -224,6 +246,17 @@ async fn challenge_client(
Error::InvalidChallengeSolution 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(()) Ok(())
} }
@@ -237,7 +270,7 @@ fn connect_error_code(err: &Error) -> ConnectErrorCode {
} }
} }
async fn authenticate(props: &mut ClientConnection) -> Result<(VerifyingKey, i32), Error> { async fn authenticate(props: &mut ClientConnection) -> Result<(VerifyingKey, ClientId), Error> {
let Some(ClientRequest { let Some(ClientRequest {
payload: Some(ClientRequestPayload::AuthChallengeRequest(challenge)), payload: Some(ClientRequestPayload::AuthChallengeRequest(challenge)),
}) = props.transport.recv().await }) = props.transport.recv().await
@@ -253,13 +286,13 @@ async fn authenticate(props: &mut ClientConnection) -> Result<(VerifyingKey, i32
VerifyingKey::from_bytes(pubkey_bytes).map_err(|_| Error::InvalidAuthPubkeyEncoding)?; VerifyingKey::from_bytes(pubkey_bytes).map_err(|_| Error::InvalidAuthPubkeyEncoding)?;
let (client_id, nonce) = match get_nonce(&props.db, &pubkey).await? { let (client_id, nonce) = match get_nonce(&props.db, &pubkey).await? {
Some((client_id, nonce)) => (client_id, nonce), Some(state) => (state.client_id, state.nonce),
None => { None => {
approve_new_client(&props.actors, pubkey).await?; approve_new_client(&props.actors, pubkey).await?;
match insert_client(&props.db, &pubkey).await? { match insert_client(&props.db, &pubkey).await? {
InsertClientResult::Inserted(client_id) => (client_id, 0), InsertClientResult::Inserted(client_id) => (client_id, 0),
InsertClientResult::AlreadyExists => match get_nonce(&props.db, &pubkey).await? { InsertClientResult::AlreadyExists => match get_nonce(&props.db, &pubkey).await? {
Some((client_id, nonce)) => (client_id, nonce), Some(state) => (state.client_id, state.nonce),
None => return Err(Error::InternalError), None => return Err(Error::InternalError),
}, },
} }

View File

@@ -114,6 +114,15 @@ pub async fn test_challenge_auth() {
.await .await
.unwrap(); .unwrap();
let response = test_transport.recv().await.expect("should receive auth ok");
match response {
Ok(resp) => match resp.payload {
Some(ClientResponsePayload::AuthOk(_)) => {}
other => panic!("Expected AuthOk, got {other:?}"),
},
Err(err) => panic!("Expected Ok response, got Err({err:?})"),
}
// Auth completes, session spawned // Auth completes, session spawned
task.await.unwrap(); task.await.unwrap();
} }
@@ -178,6 +187,15 @@ pub async fn test_evm_sign_request_payload_is_handled() {
.await .await
.unwrap(); .unwrap();
let response = test_transport.recv().await.expect("should receive auth ok");
match response {
Ok(resp) => match resp.payload {
Some(ClientResponsePayload::AuthOk(_)) => {}
other => panic!("Expected AuthOk, got {other:?}"),
},
Err(err) => panic!("Expected Ok response, got Err({err:?})"),
}
task.await.unwrap(); task.await.unwrap();
let tx = TxEip1559 { let tx = TxEip1559 {