Add query request id

Cargo.lock

@@ -4162,6 +4162,7 @@ dependencies = [
  "nym-sdk",
  "nym-service-providers-common",
  "nym-socks5-requests",
+ "rand 0.8.5",
  "serde",
  "tokio",
 ]

common/socks5/requests/src/request.rs

@@ -84,6 +84,8 @@ pub struct SendRequest {
     pub local_closed: bool,
 }
 
+pub type QueryRequestId = u64;
+
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub enum QueryRequest {
     OpenProxy,
@@ -192,11 +194,12 @@ impl Socks5Request {
 
     pub fn new_query(
         protocol_version: Socks5ProtocolVersion,
+        req_id: u64,
         query: QueryRequest,
     ) -> Socks5Request {
         Socks5Request {
             protocol_version,
-            content: Socks5RequestContent::Query(query),
+            content: Socks5RequestContent::Query((req_id, query)),
         }
     }
 }
@@ -213,7 +216,7 @@ pub enum Socks5RequestContent {
     /// Re-use an existing TCP connection, sending more request data up it.
     Send(SendRequest),
 
-    Query(QueryRequest),
+    Query((QueryRequestId, QueryRequest)),
 }
 
 impl Socks5RequestContent {
@@ -677,13 +680,13 @@ mod request_deserialization_tests {
 
         #[test]
         fn serialize_there_and_back() {
-            let open_proxy = Socks5RequestContent::Query(QueryRequest::OpenProxy);
+            let open_proxy = Socks5RequestContent::Query((42, QueryRequest::OpenProxy));
             let bytes_open_proxy = open_proxy.clone().into_bytes();
-            assert_eq!(bytes_open_proxy, vec![2, 0]);
+            assert_eq!(bytes_open_proxy, vec![42, 2, 0]);
 
-            let description = Socks5RequestContent::Query(QueryRequest::Description);
+            let description = Socks5RequestContent::Query((8723647823, QueryRequest::Description));
             let bytes_description = description.clone().into_bytes();
-            assert_eq!(bytes_description, vec![2, 1]);
+            assert_eq!(bytes_description, vec![23, 2, 1]);
 
             let open_proxy2 = Socks5RequestContent::try_from_bytes(&bytes_open_proxy).unwrap();
             let description2 = Socks5RequestContent::try_from_bytes(&bytes_description).unwrap();

common/socks5/requests/src/response.rs

@@ -1,7 +1,7 @@
 // Copyright 2020-2023 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::{make_bincode_serializer, ConnectionId, Socks5ProtocolVersion, Socks5RequestError};
+use crate::{make_bincode_serializer, ConnectionId, Socks5ProtocolVersion, Socks5RequestError, QueryRequestId};
 use nym_service_providers_common::interface::{Serializable, ServiceProviderResponse};
 use serde::{Deserialize, Serialize};
 use tap::TapFallible;
@@ -148,11 +148,12 @@ impl Socks5Response {
 
     pub fn new_query(
         protocol_version: Socks5ProtocolVersion,
+        query_request_id: QueryRequestId,
         query_response: QueryResponse,
     ) -> Socks5Response {
         Socks5Response {
             protocol_version,
-            content: Socks5ResponseContent::Query(query_response),
+            content: Socks5ResponseContent::Query((query_request_id, query_response)),
         }
     }
 }
@@ -161,7 +162,7 @@ impl Socks5Response {
 pub enum Socks5ResponseContent {
     NetworkData(NetworkData),
     ConnectionError(ConnectionError),
-    Query(QueryResponse),
+    Query((QueryRequestId, QueryResponse)),
 }
 
 impl Socks5ResponseContent {
@@ -234,9 +235,9 @@ impl Socks5ResponseContent {
         }
     }
 
-    pub fn as_query(&self) -> Option<&QueryResponse> {
+    pub fn as_query(&self) -> Option<(&QueryRequestId, &QueryResponse)> {
         match self {
-            Socks5ResponseContent::Query(query) => Some(query),
+            Socks5ResponseContent::Query((query_id, query)) => Some((query_id, query)),
             _ => None,
         }
     }

service-providers/network-requester/src/core.rs

@@ -25,8 +25,8 @@ use nym_socks5_proxy_helpers::connection_controller::{
 };
 use nym_socks5_proxy_helpers::proxy_runner::{MixProxyReader, MixProxySender};
 use nym_socks5_requests::{
-    ConnectRequest, ConnectionId, NetworkData, QueryRequest, QueryResponse, SendRequest,
-    Socks5ProtocolVersion, Socks5ProviderRequest, Socks5Request, Socks5RequestContent,
+    ConnectRequest, ConnectionId, NetworkData, QueryRequest, QueryRequestId, QueryResponse,
+    SendRequest, Socks5ProtocolVersion, Socks5ProviderRequest, Socks5Request, Socks5RequestContent,
     Socks5Response,
 };
 use nym_sphinx::addressing::clients::Recipient;
@@ -149,7 +149,9 @@ impl ServiceProvider<Socks5Request> for NRServiceProvider {
                 }
                 self.handle_proxy_send(req)
             }
-            Socks5RequestContent::Query(query) => return self.handle_query(query),
+            Socks5RequestContent::Query((query_id, query)) => {
+                return self.handle_query(query_id, query)
+            }
         }
 
         Ok(None)
@@ -480,16 +482,19 @@ impl NRServiceProvider {
 
     fn handle_query(
         &self,
+        query_id: QueryRequestId,
         query: QueryRequest,
     ) -> Result<Option<Socks5Response>, NetworkRequesterError> {
         let protocol_version = Socks5ProtocolVersion::default();
         let response = match query {
             QueryRequest::OpenProxy => Socks5Response::new_query(
                 protocol_version,
+                query_id,
                 QueryResponse::OpenProxy(self.open_proxy),
             ),
             QueryRequest::Description => Socks5Response::new_query(
                 protocol_version,
+                query_id,
                 QueryResponse::Description("Description (placeholder)".to_string()),
             ),
         };

tools/nym-nr-query/Cargo.toml

@@ -14,5 +14,6 @@ nym-network-defaults = { path = "../../common/network-defaults" }
 nym-sdk = { path = "../../sdk/rust/nym-sdk" }
 nym-service-providers-common = { path = "../../service-providers/common" }
 nym-socks5-requests = { path = "../../common/socks5/requests" }
+rand = "0.8"
 serde = { workspace = true }
 tokio = { workspace = true, features = [ "net", "rt-multi-thread", "macros" ] }

tools/nym-nr-query/src/main.rs

@@ -7,13 +7,20 @@ use nym_service_providers_common::interface::{
     ControlRequest, ControlResponse, ProviderInterfaceVersion, Request, Response, ResponseContent,
 };
 use nym_socks5_requests::{
-    QueryRequest, QueryResponse, Socks5ProtocolVersion, Socks5Request, Socks5Response,
+    QueryRequest, QueryRequestId, QueryResponse, Socks5ProtocolVersion, Socks5Request,
+    Socks5Response,
 };
+use rand::RngCore;
 use serde::Serialize;
 use tokio::time::{timeout, Duration};
 
 const RESPONSE_TIMEOUT: Duration = Duration::from_secs(10);
 
+fn generate_random() -> u64 {
+    let mut rng = rand::rngs::OsRng;
+    rng.next_u64()
+}
+
 #[derive(Parser)]
 #[command(version, about, long_about = None)]
 struct Cli {
@@ -128,9 +135,10 @@ fn new_supported_request_versions_request() -> Request {
     Request::new_control(ProviderInterfaceVersion::new_current(), request_versions)
 }
 
-fn new_open_proxy_request() -> Request<Socks5Request> {
+fn new_open_proxy_request(query_id: QueryRequestId) -> Request<Socks5Request> {
     let request_open_proxy = Socks5Request::new_query(
         Socks5ProtocolVersion::new_current(),
+        query_id,
         QueryRequest::OpenProxy,
     );
     Request::new_provider_data(ProviderInterfaceVersion::new_current(), request_open_proxy)
@@ -175,19 +183,22 @@ impl QueryClient {
     }
 
     async fn query_open_proxy(&mut self) -> QueryResponse {
+        let query_id: QueryRequestId = generate_random();
         self.client
             .send_bytes(
                 self.provider,
-                new_open_proxy_request().into_bytes(),
+                new_open_proxy_request(query_id).into_bytes(),
                 IncludedSurbs::new(10),
             )
             .await;
         let response = wait_for_socks5_response(&mut self.client).await;
-        response
+        let query_response = response
             .content
             .as_query()
-            .expect("Unexpected response type!")
-            .clone()
+            .expect("Unexpected response type!");
+        // TODO: we should just drop unexpected messages
+        assert_eq!(&query_id, query_response.0);
+        query_response.1.clone()
     }
 
     async fn ping(&mut self) -> PingResponse {