From 33bba9e46bcc5edcb9724fefb2db9141d95b2f97 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20L=C3=B6nnhager?= <david.l@mullvad.net>
Date: Sun, 19 Nov 2023 14:03:42 +0100
Subject: [PATCH] Set sock nonblocking after connect

---
 test/test-manager/src/tests/settings.rs     |  2 +-
 test/test-manager/src/tests/tunnel_state.rs |  4 +-
 test/test-runner/src/net.rs                 | 50 +++++++++++----------
 3 files changed, 30 insertions(+), 26 deletions(-)

diff --git a/test/test-manager/src/tests/settings.rs b/test/test-manager/src/tests/settings.rs
index 75fd3979c696..b00192efa072 100644
--- a/test/test-manager/src/tests/settings.rs
+++ b/test/test-manager/src/tests/settings.rs
@@ -1,5 +1,5 @@
 use super::helpers;
-use super::helpers::{connect_and_wait, disconnect_and_wait, get_tunnel_state, send_guest_probes};
+use super::helpers::{connect_and_wait, get_tunnel_state, send_guest_probes};
 use super::{Error, TestContext};
 use crate::assert_tunnel_state;
 use crate::vm::network::DUMMY_LAN_INTERFACE_IP;
diff --git a/test/test-manager/src/tests/tunnel_state.rs b/test/test-manager/src/tests/tunnel_state.rs
index c82022c09dad..7b8cacf00d88 100644
--- a/test/test-manager/src/tests/tunnel_state.rs
+++ b/test/test-manager/src/tests/tunnel_state.rs
@@ -1,6 +1,6 @@
 use super::helpers::{
-    self, connect_and_wait, disconnect_and_wait, get_tunnel_state, send_guest_probes,
-    set_relay_settings, unreachable_wireguard_tunnel, wait_for_tunnel_state,
+    self, connect_and_wait, get_tunnel_state, send_guest_probes, set_relay_settings,
+    unreachable_wireguard_tunnel, wait_for_tunnel_state,
 };
 use super::{ui, Error, TestContext};
 use crate::assert_tunnel_state;
diff --git a/test/test-runner/src/net.rs b/test/test-runner/src/net.rs
index 79e8e34cce34..90d691f99537 100644
--- a/test/test-runner/src/net.rs
+++ b/test/test-runner/src/net.rs
@@ -26,11 +26,6 @@ pub async fn send_tcp(
             test_rpc::Error::SendTcp
         })?;
 
-    sock.set_nonblocking(true).map_err(|error| {
-        log::error!("Failed to set non-blocking TCP socket: {error}");
-        test_rpc::Error::SendTcp
-    })?;
-
     if let Some(iface) = bind_interface {
         #[cfg(target_os = "macos")]
         let interface_index = unsafe {
@@ -59,24 +54,33 @@ pub async fn send_tcp(
         log::trace!("Bind interface {iface} is ignored on Windows")
     }
 
-    sock.bind(&SockAddr::from(bind_addr)).map_err(|error| {
-        log::error!("Failed to bind TCP socket to {bind_addr}: {error}");
-        test_rpc::Error::SendTcp
-    })?;
-
     log::debug!("Connecting from {bind_addr} to {destination}/TCP");
 
-    sock.connect(&SockAddr::from(destination))
-        .map_err(|error| {
-            log::error!("Failed to connect to {destination}: {error}");
+    let mut stream = tokio::task::spawn_blocking(move || {
+        sock.bind(&SockAddr::from(bind_addr)).map_err(|error| {
+            log::error!("Failed to bind TCP socket to {bind_addr}: {error}");
             test_rpc::Error::SendTcp
         })?;
 
-    let std_stream = std::net::TcpStream::from(sock);
-    let mut stream = TcpStream::from_std(std_stream).map_err(|error| {
-        log::error!("Failed to convert to TCP stream to tokio stream: {error}");
-        test_rpc::Error::SendTcp
-    })?;
+        sock.connect(&SockAddr::from(destination))
+            .map_err(|error| {
+                log::error!("Failed to connect to {destination}: {error}");
+                test_rpc::Error::SendTcp
+            })?;
+
+        sock.set_nonblocking(true).map_err(|error| {
+            log::error!("Failed to set non-blocking TCP socket: {error}");
+            test_rpc::Error::SendTcp
+        })?;
+
+        let std_stream = std::net::TcpStream::from(sock);
+        TcpStream::from_std(std_stream).map_err(|error| {
+            log::error!("Failed to convert to TCP stream to tokio stream: {error}");
+            test_rpc::Error::SendTcp
+        })
+    })
+    .await
+    .unwrap()?;
 
     stream.write_all(b"hello").await.map_err(|error| {
         log::error!("Failed to send message to {destination}: {error}");
@@ -101,11 +105,6 @@ pub async fn send_udp(
             test_rpc::Error::SendUdp
         })?;
 
-    sock.set_nonblocking(true).map_err(|error| {
-        log::error!("Failed to set non-blocking UDP socket: {error}");
-        test_rpc::Error::SendUdp
-    })?;
-
     if let Some(iface) = bind_interface {
         #[cfg(target_os = "macos")]
         let interface_index = unsafe {
@@ -139,6 +138,11 @@ pub async fn send_udp(
         test_rpc::Error::SendUdp
     })?;
 
+    sock.set_nonblocking(true).map_err(|error| {
+        log::error!("Failed to set non-blocking UDP socket: {error}");
+        test_rpc::Error::SendUdp
+    })?;
+
     log::debug!("Send message from {bind_addr} to {destination}/UDP");
 
     let std_socket = std::net::UdpSocket::from(sock);