Merge branch 'testing-add-socks-server'

3 files changed, 78 insertions, 1 deletions

diff --git a/test/test-rpc/src/client.rs b/test/test-rpc/src/client.rs
index 2c47328e00..4d103ed44e 100644
--- a/test/test-rpc/src/client.rs
+++ b/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
diff --git a/test/test-rpc/src/lib.rs b/test/test-rpc/src/lib.rs
index 5919a894d1..d2bee40dbb 100644
--- a/test/test-rpc/src/lib.rs
+++ b/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>;
 
diff --git a/test/test-rpc/src/net.rs b/test/test-rpc/src/net.rs
index b4e114ea47..77aa5c938a 100644
--- a/test/test-rpc/src/net.rs
+++ b/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)?;