From c7fa910516ed877d0c900dd6ef368d2b30becab7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bogdan-=C8=98tefan=20Neac=C5=9Fu?= <bogdan@nymtech.net>
Date: Fri, 5 Jul 2024 10:36:18 +0000
Subject: [PATCH] Fix add of req id

---
 .../authenticator-requests/src/v1/request.rs   |  3 +++
 .../authenticator-requests/src/v1/response.rs  | 18 +++++++++++++-----
 .../authenticator/src/mixnet_listener.rs       | 10 +++++++---
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/common/authenticator-requests/src/v1/request.rs b/common/authenticator-requests/src/v1/request.rs
index d028856373..1538999828 100644
--- a/common/authenticator-requests/src/v1/request.rs
+++ b/common/authenticator-requests/src/v1/request.rs
@@ -20,6 +20,7 @@ pub struct AuthenticatorRequest {
     pub version: u8,
     pub data: AuthenticatorRequestData,
     pub reply_to: Recipient,
+    pub request_id: u64,
 }
 
 impl AuthenticatorRequest {
@@ -37,6 +38,7 @@ impl AuthenticatorRequest {
                 version: VERSION,
                 data: AuthenticatorRequestData::Initial(init_message),
                 reply_to,
+                request_id,
             },
             request_id,
         )
@@ -49,6 +51,7 @@ impl AuthenticatorRequest {
                 version: VERSION,
                 data: AuthenticatorRequestData::Final(gateway_client),
                 reply_to,
+                request_id,
             },
             request_id,
         )
diff --git a/common/authenticator-requests/src/v1/response.rs b/common/authenticator-requests/src/v1/response.rs
index 41e7f2473d..925ece872e 100644
--- a/common/authenticator-requests/src/v1/response.rs
+++ b/common/authenticator-requests/src/v1/response.rs
@@ -19,19 +19,27 @@ pub struct AuthenticatorResponse {
 impl AuthenticatorResponse {
     pub fn new_pending_registration_success(
         registration_data: RegistrationData,
+        request_id: u64,
         reply_to: Recipient,
     ) -> Self {
         Self {
             version: VERSION,
-            data: AuthenticatorResponseData::PendingRegistration(registration_data),
+            data: AuthenticatorResponseData::PendingRegistration(PendingRegistrationResponse {
+                reply: registration_data,
+                reply_to,
+                request_id,
+            }),
             reply_to,
         }
     }
 
-    pub fn new_registered(reply_to: Recipient) -> Self {
+    pub fn new_registered(reply_to: Recipient, request_id: u64) -> Self {
         Self {
             version: VERSION,
-            data: AuthenticatorResponseData::Registered,
+            data: AuthenticatorResponseData::Registered(RegisteredResponse {
+                reply_to,
+                request_id,
+            }),
             reply_to,
         }
     }
@@ -54,8 +62,8 @@ impl AuthenticatorResponse {
 
     pub fn id(&self) -> Option<u64> {
         match &self.data {
-            AuthenticatorResponseData::PendingRegistration(registration_data) => {}
-            AuthenticatorResponseData::Registered => None,
+            AuthenticatorResponseData::PendingRegistration(response) => Some(response.request_id),
+            AuthenticatorResponseData::Registered(response) => Some(response.request_id),
         }
     }
 }
diff --git a/service-providers/authenticator/src/mixnet_listener.rs b/service-providers/authenticator/src/mixnet_listener.rs
index 434180c6a1..1b43844018 100644
--- a/service-providers/authenticator/src/mixnet_listener.rs
+++ b/service-providers/authenticator/src/mixnet_listener.rs
@@ -124,6 +124,7 @@ impl MixnetListener {
     fn on_initial_request(
         &mut self,
         init_message: InitMessage,
+        request_id: u64,
         reply_to: Recipient,
     ) -> AuthenticatorHandleResult {
         let remote_public = init_message.pub_key();
@@ -131,6 +132,7 @@ impl MixnetListener {
         if let Some(registration_data) = self.registration_in_progres.get(&remote_public) {
             return Ok(AuthenticatorResponse::new_pending_registration_success(
                 registration_data.value().clone(),
+                request_id,
                 reply_to,
             ));
         }
@@ -177,6 +179,7 @@ impl MixnetListener {
 
         Ok(AuthenticatorResponse::new_pending_registration_success(
             registration_data,
+            request_id,
             reply_to,
         ))
     }
@@ -184,6 +187,7 @@ impl MixnetListener {
     fn on_final_request(
         &mut self,
         gateway_client: GatewayClient,
+        request_id: u64,
         reply_to: Recipient,
     ) -> AuthenticatorHandleResult {
         let registration_data = self
@@ -211,7 +215,7 @@ impl MixnetListener {
                 .client_registry()
                 .insert(gateway_client.pub_key(), gateway_client);
 
-            Ok(AuthenticatorResponse::new_registered(reply_to))
+            Ok(AuthenticatorResponse::new_registered(reply_to, request_id))
         } else {
             Err(AuthenticatorError::MacVerificationFailure)
         }
@@ -235,10 +239,10 @@ impl MixnetListener {
 
         match request.data {
             AuthenticatorRequestData::Initial(init_msg) => {
-                self.on_initial_request(init_msg, request.reply_to)
+                self.on_initial_request(init_msg, request.request_id, request.reply_to)
             }
             AuthenticatorRequestData::Final(client) => {
-                self.on_final_request(client, request.reply_to)
+                self.on_final_request(client, request.request_id, request.reply_to)
             }
         }
     }