feat(proto): request / response pair tracking by assigning id

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

@@ -8,15 +8,35 @@ use arbiter_proto::{
     transport::{Bi, Error as TransportError, Receiver, Sender, grpc::GrpcBi},
 };
 use async_trait::async_trait;
+use tonic::Status;
 use tracing::warn;
 
-use crate::actors::client::{self, ClientConnection, auth};
+use crate::{
+    actors::client::{self, ClientConnection, auth},
+    grpc::request_tracker::RequestTracker,
+};
 
-pub struct AuthTransportAdapter<'a>(pub(super) &'a mut GrpcBi<ClientRequest, ClientResponse>);
+pub struct AuthTransportAdapter<'a> {
+    bi: &'a mut GrpcBi<ClientRequest, ClientResponse>,
+    request_tracker: &'a mut RequestTracker,
+    response_id: &'a mut Option<i32>,
+}
 
-impl AuthTransportAdapter<'_> {
-    fn response_to_proto(response: auth::Outbound) -> ClientResponse {
-        let payload = match response {
+impl<'a> AuthTransportAdapter<'a> {
+    pub fn new(
+        bi: &'a mut GrpcBi<ClientRequest, ClientResponse>,
+        request_tracker: &'a mut RequestTracker,
+        response_id: &'a mut Option<i32>,
+    ) -> Self {
+        Self {
+            bi,
+            request_tracker,
+            response_id,
+        }
+    }
+
+    fn response_to_proto(response: auth::Outbound) -> ClientResponsePayload {
+        match response {
             auth::Outbound::AuthChallenge { pubkey, nonce } => {
                 ClientResponsePayload::AuthChallenge(ProtoAuthChallenge {
                     pubkey: pubkey.to_bytes().to_vec(),
@@ -26,39 +46,44 @@ impl AuthTransportAdapter<'_> {
             auth::Outbound::AuthSuccess => {
                 ClientResponsePayload::AuthResult(ProtoAuthResult::Success.into())
             }
-        };
-
-        ClientResponse {
-            payload: Some(payload),
         }
     }
 
-    fn error_to_proto(error: auth::Error) -> ClientResponse {
-        ClientResponse {
-            payload: Some(ClientResponsePayload::AuthResult(
-                match error {
-                    auth::Error::InvalidChallengeSolution => ProtoAuthResult::InvalidSignature,
-                    auth::Error::ApproveError(auth::ApproveError::Denied) => {
-                        ProtoAuthResult::ApprovalDenied
-                    }
-                    auth::Error::ApproveError(auth::ApproveError::Upstream(
-                        crate::actors::router::ApprovalError::NoUserAgentsConnected,
-                    )) => ProtoAuthResult::NoUserAgentsOnline,
-                    auth::Error::ApproveError(auth::ApproveError::Internal)
-                    | auth::Error::DatabasePoolUnavailable
-                    | auth::Error::DatabaseOperationFailed
-                    | auth::Error::Transport => ProtoAuthResult::Internal,
+    fn error_to_proto(error: auth::Error) -> ClientResponsePayload {
+        ClientResponsePayload::AuthResult(
+            match error {
+                auth::Error::InvalidChallengeSolution => ProtoAuthResult::InvalidSignature,
+                auth::Error::ApproveError(auth::ApproveError::Denied) => {
+                    ProtoAuthResult::ApprovalDenied
                 }
-                .into(),
-            )),
-        }
+                auth::Error::ApproveError(auth::ApproveError::Upstream(
+                    crate::actors::router::ApprovalError::NoUserAgentsConnected,
+                )) => ProtoAuthResult::NoUserAgentsOnline,
+                auth::Error::ApproveError(auth::ApproveError::Internal)
+                | auth::Error::DatabasePoolUnavailable
+                | auth::Error::DatabaseOperationFailed
+                | auth::Error::Transport => ProtoAuthResult::Internal,
+            }
+            .into(),
+        )
+    }
+
+    async fn send_client_response(
+        &mut self,
+        payload: ClientResponsePayload,
+    ) -> Result<(), TransportError> {
+        let request_id = self.response_id.take();
+
+        self.bi
+            .send(Ok(ClientResponse {
+                request_id,
+                payload: Some(payload),
+            }))
+            .await
     }
 
     async fn send_auth_result(&mut self, result: ProtoAuthResult) -> Result<(), TransportError> {
-        self.0
-            .send(Ok(ClientResponse {
-                payload: Some(ClientResponsePayload::AuthResult(result.into())),
-            }))
+        self.send_client_response(ClientResponsePayload::AuthResult(result.into()))
             .await
     }
 }
@@ -69,19 +94,19 @@ impl Sender<Result<auth::Outbound, auth::Error>> for AuthTransportAdapter<'_> {
         &mut self,
         item: Result<auth::Outbound, auth::Error>,
     ) -> Result<(), TransportError> {
-        let outbound = match item {
-            Ok(message) => Ok(AuthTransportAdapter::response_to_proto(message)),
-            Err(err) => Ok(AuthTransportAdapter::error_to_proto(err)),
+        let payload = match item {
+            Ok(message) => AuthTransportAdapter::response_to_proto(message),
+            Err(err) => AuthTransportAdapter::error_to_proto(err),
         };
 
-        self.0.send(outbound).await
+        self.send_client_response(payload).await
     }
 }
 
 #[async_trait]
 impl Receiver<auth::Inbound> for AuthTransportAdapter<'_> {
     async fn recv(&mut self) -> Option<auth::Inbound> {
-        let request = match self.0.recv().await? {
+        let request = match self.bi.recv().await? {
             Ok(request) => request,
             Err(error) => {
                 warn!(error = ?error, "grpc client recv failed; closing stream");
@@ -89,6 +114,15 @@ impl Receiver<auth::Inbound> for AuthTransportAdapter<'_> {
             }
         };
 
+        let request_id = match self.request_tracker.request(request.request_id) {
+            Ok(request_id) => request_id,
+            Err(error) => {
+                let _ = self.bi.send(Err(error)).await;
+                return None;
+            }
+        };
+        *self.response_id = Some(request_id);
+
         let payload = request.payload?;
 
         match payload {
@@ -114,7 +148,13 @@ impl Receiver<auth::Inbound> for AuthTransportAdapter<'_> {
                 };
                 Some(auth::Inbound::AuthChallengeSolution { signature })
             }
-            _ => None,
+            _ => {
+                let _ = self
+                    .bi
+                    .send(Err(Status::invalid_argument("Unsupported client auth request")))
+                    .await;
+                None
+            }
         }
     }
 }
@@ -124,8 +164,10 @@ impl Bi<auth::Inbound, Result<auth::Outbound, auth::Error>> for AuthTransportAda
 pub async fn start(
     conn: &mut ClientConnection,
     bi: &mut GrpcBi<ClientRequest, ClientResponse>,
+    request_tracker: &mut RequestTracker,
+    response_id: &mut Option<i32>,
 ) -> Result<(), auth::Error> {
-    let mut transport = AuthTransportAdapter(bi);
+    let mut transport = AuthTransportAdapter::new(bi, request_tracker, response_id);
     client::auth::authenticate(conn, &mut transport).await?;
     Ok(())
 }