fix(stub): resolve TCP socket leak on target server disconnect and hanging connection attempt

stub/src/main.rs

@@ -74,8 +74,19 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                 }
             };
 
-            match tokio_tungstenite::connect_async(request).await {
-                Ok((ws_stream, _)) => {
+            let connect_result = tokio::select! {
+                _ = shutdown_rx.recv() => {
+                    tracing::info!("Shutdown received during connect");
+                    break;
+                }
+                res = tokio::time::timeout(
+                    tokio::time::Duration::from_secs(5),
+                    tokio_tungstenite::connect_async(request)
+                ) => res,
+            };
+
+            match connect_result {
+                Ok(Ok((ws_stream, _))) => {
                     tracing::info!("Successfully connected to target server");
                     let (mut write, mut read) = ws_stream.split();
 
@@ -138,7 +149,6 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                             tracing::error!("Send task exited");
                             tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
                             recv_task.abort();
-                            break;
                         }
                         _ = &mut recv_task => {
                             tracing::error!("Recv task exited");
@@ -147,12 +157,18 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                     }
                     tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
                 }
-                Err(e) => {
+                Ok(Err(e)) => {
                     tracing::error!("Failed to connect via WSS: {}", e);
                     tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
                 }
+                Err(_) => {
+                    tracing::error!("Connection attempt timed out");
+                    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+                }
             }
         }
         Ok(())
     })
 }
+
+