From d5116116414aa2985fff48fdb3fcdda488944978 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bogdan-=C8=98tefan=20Neac=C5=9Fu?= <bogdan@nymtech.net>
Date: Mon, 24 Feb 2025 14:23:43 +0200
Subject: [PATCH] Connection fd callback before actual connection (#5494)

---
 .../gateway-client/src/client/mod.rs          | 12 ++---
 .../gateway-client/src/client/websockets.rs   | 47 ++++++++++++++++---
 .../client-libs/gateway-client/src/error.rs   |  3 ++
 3 files changed, 50 insertions(+), 12 deletions(-)

diff --git a/common/client-libs/gateway-client/src/client/mod.rs b/common/client-libs/gateway-client/src/client/mod.rs
index 84cb571c72..b72c7a033c 100644
--- a/common/client-libs/gateway-client/src/client/mod.rs
+++ b/common/client-libs/gateway-client/src/client/mod.rs
@@ -204,15 +204,15 @@ impl<C, St> GatewayClient<C, St> {
             "Attemting to establish connection to gateway at: {}",
             self.gateway_address
         );
-        let (ws_stream, _) = connect_async(&self.gateway_address).await?;
+        let (ws_stream, _) = connect_async(
+            &self.gateway_address,
+            #[cfg(unix)]
+            self.connection_fd_callback.clone(),
+        )
+        .await?;
 
         self.connection = SocketState::Available(Box::new(ws_stream));
 
-        #[cfg(unix)]
-        if let (Some(callback), Some(fd)) = (self.connection_fd_callback.as_ref(), self.ws_fd()) {
-            callback.as_ref()(fd);
-        }
-
         Ok(())
     }
 
diff --git a/common/client-libs/gateway-client/src/client/websockets.rs b/common/client-libs/gateway-client/src/client/websockets.rs
index 32527571ec..293435b1cf 100644
--- a/common/client-libs/gateway-client/src/client/websockets.rs
+++ b/common/client-libs/gateway-client/src/client/websockets.rs
@@ -1,6 +1,11 @@
 use crate::error::GatewayClientError;
 
 use nym_http_api_client::HickoryDnsResolver;
+#[cfg(unix)]
+use std::{
+    os::fd::{AsRawFd, RawFd},
+    sync::Arc,
+};
 use tokio::net::TcpStream;
 use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
 use tungstenite::handshake::client::Response;
@@ -11,7 +16,10 @@ use std::net::SocketAddr;
 #[cfg(not(target_arch = "wasm32"))]
 pub(crate) async fn connect_async(
     endpoint: &str,
+    #[cfg(unix)] connection_fd_callback: Option<Arc<dyn Fn(RawFd) + Send + Sync>>,
 ) -> Result<(WebSocketStream<MaybeTlsStream<TcpStream>>, Response), GatewayClientError> {
+    use tokio::net::TcpSocket;
+
     let resolver = HickoryDnsResolver::default();
     let uri =
         Url::parse(endpoint).map_err(|_| GatewayClientError::InvalidUrl(endpoint.to_owned()))?;
@@ -37,14 +45,41 @@ pub(crate) async fn connect_async(
         }
     };
 
-    let stream = TcpStream::connect(&sock_addrs[..]).await.map_err(|error| {
-        GatewayClientError::NetworkConnectionFailed {
-            address: endpoint.to_owned(),
-            source: error.into(),
+    let mut stream = Err(GatewayClientError::NoEndpointForConnection {
+        address: endpoint.to_owned(),
+    });
+    for sock_addr in sock_addrs {
+        let socket = if sock_addr.is_ipv4() {
+            TcpSocket::new_v4()
+        } else {
+            TcpSocket::new_v6()
         }
-    })?;
+        .map_err(|err| GatewayClientError::NetworkConnectionFailed {
+            address: endpoint.to_owned(),
+            source: err.into(),
+        })?;
 
-    tokio_tungstenite::client_async_tls(endpoint, stream)
+        #[cfg(unix)]
+        if let Some(callback) = connection_fd_callback.as_ref() {
+            callback.as_ref()(socket.as_raw_fd());
+        }
+
+        match socket.connect(sock_addr).await {
+            Ok(s) => {
+                stream = Ok(s);
+                break;
+            }
+            Err(err) => {
+                stream = Err(GatewayClientError::NetworkConnectionFailed {
+                    address: endpoint.to_owned(),
+                    source: err.into(),
+                });
+                continue;
+            }
+        }
+    }
+
+    tokio_tungstenite::client_async_tls(endpoint, stream?)
         .await
         .map_err(|error| GatewayClientError::NetworkConnectionFailed {
             address: endpoint.to_owned(),
diff --git a/common/client-libs/gateway-client/src/error.rs b/common/client-libs/gateway-client/src/error.rs
index 3432f43131..a9ce175504 100644
--- a/common/client-libs/gateway-client/src/error.rs
+++ b/common/client-libs/gateway-client/src/error.rs
@@ -43,6 +43,9 @@ pub enum GatewayClientError {
     #[error("connection failed: {address}: {source}")]
     NetworkConnectionFailed { address: String, source: WsError },
 
+    #[error("no socket address for endpoint: {address}")]
+    NoEndpointForConnection { address: String },
+
     #[error("Invalid URL: {0}")]
     InvalidUrl(String),