Merge branch 'testing-add-socks-server'

test/Cargo.toml

@@ -11,6 +11,7 @@ members = [
     "test-manager",
     "test-runner",
     "test-rpc",
+    "socks-server",
 ]
 
 [workspace.lints.rust]

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 }

test/socks-server/src/lib.rs

@@ -0,0 +1,58 @@
+use futures::StreamExt;
+use std::io;
+use std::net::SocketAddr;
+
+#[derive(err_derive::Error, Debug)]
+#[error(no_from)]
+pub enum Error {
+    #[error(display = "Failed to start SOCKS5 server")]
+    StartSocksServer(#[error(source)] io::Error),
+}
+
+pub struct Handle {
+    handle: tokio::task::JoinHandle<()>,
+}
+
+/// Spawn a SOCKS server bound to `bind_addr`
+pub async fn spawn(bind_addr: SocketAddr) -> Result<Handle, 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();
+
+                    // Act as normal SOCKS server
+                    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 { handle })
+}
+
+impl Handle {
+    pub fn close(&self) {
+        self.handle.abort();
+    }
+}
+
+impl Drop for Handle {
+    fn drop(&mut self) {
+        self.close();
+    }
+}

test/test-manager/Cargo.toml

@@ -42,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 }
 

test/test-manager/src/main.rs

@@ -14,6 +14,7 @@ use std::path::PathBuf;
 use anyhow::Context;
 use anyhow::Result;
 use clap::Parser;
+use std::net::SocketAddr;
 use tests::config::DEFAULT_MULLVAD_HOST;
 
 /// Test manager for Mullvad VPN app
@@ -248,6 +249,13 @@ async fn main() -> Result<()> {
                 .await
                 .context("Failed to run provisioning for VM")?;
 
+            // 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;
 
             let result = run_tests::run(
@@ -291,6 +299,7 @@ async fn main() -> Result<()> {
             if display {
                 instance.wait().await;
             }
+            socks.close();
             result
         }
         Commands::FormatTestReports { reports } => {

test/test-manager/src/vm/network/mod.rs

@@ -15,3 +15,6 @@ pub use platform::{
     CUSTOM_TUN_REMOTE_REAL_PORT, CUSTOM_TUN_REMOTE_TUN_ADDR, DUMMY_LAN_INTERFACE_IP,
     NON_TUN_GATEWAY,
 };
+
+/// Port on NON_TUN_GATEWAY that hosts a SOCKS5 server
+pub const SOCKS5_PORT: u16 = 54321;

test/test-rpc/src/client.rs

@@ -213,6 +213,16 @@ impl ServiceClient {
             .await?
     }
 
+    /// Start forwarding TCP from a server listening on `bind_addr` to the given address, and return a handle that closes the
+    /// server when dropped
+    pub async fn start_tcp_forward(
+        &self,
+        bind_addr: SocketAddr,
+        via_addr: SocketAddr,
+    ) -> Result<crate::net::SockHandle, Error> {
+        crate::net::SockHandle::start_tcp_forward(self.client.clone(), bind_addr, via_addr).await
+    }
+
     /// Restarts the app.
     ///
     /// Shuts down a running app, making it disconnect from any current tunnel

test/test-rpc/src/lib.rs

@@ -53,6 +53,8 @@ pub enum Error {
     InvalidUrl,
     #[error(display = "Timeout")]
     Timeout,
+    #[error(display = "TCP forward error")]
+    TcpForward,
 }
 
 /// Response from am.i.mullvad.net
@@ -148,6 +150,16 @@ mod service {
         /// Perform DNS resolution.
         async fn resolve_hostname(hostname: String) -> Result<Vec<SocketAddr>, Error>;
 
+        /// Start forwarding TCP bound to the given address. Return an ID that can be used with
+        /// `stop_tcp_forward`, and the address that the listening socket was actually bound to.
+        async fn start_tcp_forward(
+            bind_addr: SocketAddr,
+            via_addr: SocketAddr,
+        ) -> Result<(net::SockHandleId, SocketAddr), Error>;
+
+        /// Stop forwarding TCP that was previously started with `start_tcp_forward`.
+        async fn stop_tcp_forward(id: net::SockHandleId) -> Result<(), Error>;
+
         /// Restart the Mullvad VPN application.
         async fn restart_mullvad_daemon() -> Result<(), Error>;
 

test/test-rpc/src/net.rs

@@ -1,6 +1,8 @@
+use futures::channel::oneshot;
 use hyper::{Client, Uri};
 use once_cell::sync::Lazy;
-use serde::de::DeserializeOwned;
+use serde::{de::DeserializeOwned, Deserialize, Serialize};
+use std::net::SocketAddr;
 use tokio_rustls::rustls::ClientConfig;
 
 use crate::{AmIMullvad, Error};
@@ -17,6 +19,59 @@ static CLIENT_CONFIG: Lazy<ClientConfig> = Lazy::new(|| {
         .with_no_client_auth()
 });
 
+#[derive(Debug, Serialize, Deserialize, Clone, Copy, Hash, PartialEq, Eq)]
+pub struct SockHandleId(pub usize);
+
+pub struct SockHandle {
+    stop_tx: Option<oneshot::Sender<()>>,
+    bind_addr: SocketAddr,
+}
+
+impl SockHandle {
+    pub(crate) async fn start_tcp_forward(
+        client: crate::service::ServiceClient,
+        bind_addr: SocketAddr,
+        via_addr: SocketAddr,
+    ) -> Result<Self, Error> {
+        let (stop_tx, stop_rx) = oneshot::channel();
+
+        let (id, bind_addr) = client
+            .start_tcp_forward(tarpc::context::current(), bind_addr, via_addr)
+            .await??;
+
+        tokio::spawn(async move {
+            let _ = stop_rx.await;
+
+            log::trace!("Stopping TCP forward");
+
+            if let Err(error) = client.stop_tcp_forward(tarpc::context::current(), id).await {
+                log::error!("Failed to stop TCP forward: {error}");
+            }
+        });
+
+        Ok(SockHandle {
+            stop_tx: Some(stop_tx),
+            bind_addr,
+        })
+    }
+
+    pub fn stop(&mut self) {
+        if let Some(stop_tx) = self.stop_tx.take() {
+            let _ = stop_tx.send(());
+        }
+    }
+
+    pub fn bind_addr(&self) -> SocketAddr {
+        self.bind_addr
+    }
+}
+
+impl Drop for SockHandle {
+    fn drop(&mut self) {
+        self.stop()
+    }
+}
+
 pub async fn geoip_lookup(mullvad_host: String) -> Result<AmIMullvad, Error> {
     let uri = Uri::try_from(format!("https://ipv4.am.i.{mullvad_host}/json"))
         .map_err(|_| Error::InvalidUrl)?;