Move testing SOCKS server to own crate

mullvad · Feb 6, 2024 · d6f28b6 · d6f28b6
1 parent 6f33280
commit d6f28b6
Show file tree

Hide file tree

Showing 6 changed files with 77 additions and 31 deletions.
diff --git a/test/Cargo.lock b/test/Cargo.lock
diff --git a/test/Cargo.toml b/test/Cargo.toml
@@ -11,6 +11,7 @@ members = [
     "test-manager",
     "test-runner",
     "test-rpc",
+    "socks-server",
 ]
 
 [workspace.lints.rust]

diff --git a/test/socks-server/Cargo.toml b/test/socks-server/Cargo.toml
@@ -0,0 +1,18 @@
+[package]
+name = "socks-server"
+description = "Contains a simple SOCKS5 server"
+authors.workspace = true
+repository.workspace = true
+license.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+fast-socks5 = "0.9.5"
+err-derive = { workspace = true }
+tokio = { workspace = true }
+log = { workspace = true }
+futures = { workspace = true }
diff --git a/test/socks-server/src/lib.rs b/test/socks-server/src/lib.rs
@@ -0,0 +1,38 @@
+use futures::StreamExt;
+use std::io;
+use std::net::SocketAddr;
+
+#[derive(err_derive::Error, Debug)]
+pub enum Error {
+    #[error(display = "Failed to start SOCKS5 server")]
+    StartSocksServer(#[error(source)] io::Error),
+}
+
+pub async fn spawn(bind_addr: SocketAddr) -> Result<tokio::task::JoinHandle<()>, Error> {
+    let socks_server: fast_socks5::server::Socks5Server =
+        fast_socks5::server::Socks5Server::bind(bind_addr)
+            .await
+            .map_err(Error::StartSocksServer)?;
+
+    let handle = tokio::spawn(async move {
+        let mut incoming = socks_server.incoming();
+
+        while let Some(new_client) = incoming.next().await {
+            match new_client {
+                Ok(socket) => {
+                    let fut = socket.upgrade_to_socks5();
+                    tokio::spawn(async move {
+                        match fut.await {
+                            Ok(_socket) => log::info!("socks client disconnected"),
+                            Err(error) => log::error!("socks client failed: {error}"),
+                        }
+                    });
+                }
+                Err(error) => {
+                    log::error!("failed to accept socks client: {error}");
+                }
+            }
+        }
+    });
+    Ok(handle)
+}
diff --git a/test/test-manager/Cargo.toml b/test/test-manager/Cargo.toml
@@ -11,7 +11,6 @@ rust-version.workspace = true
 workspace = true
 
 [dependencies]
-fast-socks5 = "0.9.5"
 anyhow = { version = "1", features = ["backtrace"] }
 futures = { workspace = true }
 regex = "1"
@@ -43,6 +42,7 @@ pcap = { version = "0.10.1", features = ["capture-stream"] }
 pnet_packet = "0.31.0"
 
 test-rpc = { path = "../test-rpc" }
+socks-server = { path = "../socks-server" }
 
 env_logger = { workspace = true }
 

diff --git a/test/test-manager/src/main.rs b/test/test-manager/src/main.rs
@@ -14,7 +14,6 @@ use std::path::PathBuf;
 use anyhow::Context;
 use anyhow::Result;
 use clap::Parser;
-use futures::StreamExt;
 use std::net::SocketAddr;
 use tests::config::DEFAULT_MULLVAD_HOST;
 
@@ -250,33 +249,12 @@ async fn main() -> Result<()> {
                 .await
                 .context("Failed to run provisioning for VM")?;
 
-            let socks_server: fast_socks5::server::Socks5Server =
-                fast_socks5::server::Socks5Server::bind(SocketAddr::new(
-                    crate::vm::network::NON_TUN_GATEWAY.into(),
-                    crate::vm::network::SOCKS5_PORT,
-                ))
-                .await
-                .context("Failed to start SOCKS5 server")?;
-            let socks_server = tokio::spawn(async move {
-                let mut incoming = socks_server.incoming();
-
-                while let Some(new_client) = incoming.next().await {
-                    match new_client {
-                        Ok(socket) => {
-                            let fut = socket.upgrade_to_socks5();
-                            tokio::spawn(async move {
-                                match fut.await {
-                                    Ok(_socket) => log::info!("socks client disconnected"),
-                                    Err(error) => log::error!("socks client failed: {error}"),
-                                }
-                            });
-                        }
-                        Err(error) => {
-                            log::error!("failed to accept socks client: {error}");
-                        }
-                    }
-                }
-            });
+            // For convenience, spawn a SOCKS5 server that is reachable for tests that need it
+            let socks = socks_server::spawn(SocketAddr::new(
+                crate::vm::network::NON_TUN_GATEWAY.into(),
+                crate::vm::network::SOCKS5_PORT,
+            ))
+            .await?;
 
             let skip_wait = vm_config.provisioner != config::Provisioner::Noop;
 
@@ -321,7 +299,7 @@ async fn main() -> Result<()> {
             if display {
                 instance.wait().await;
             }
-            socks_server.abort();
+            socks.abort();
             result
         }
         Commands::FormatTestReports { reports } => {