Add forwarding SOCKS server

dlon · dlon · commit 4e5d673858d9 · 2024-02-07T10:55:07.000+01:00
diff --git a/test/socks-server/src/lib.rs b/test/socks-server/src/lib.rs
@@ -2,6 +2,8 @@ use futures::StreamExt;
 use std::io;
 use std::net::IpAddr;
 use std::net::SocketAddr;
+use fast_socks5::util::target_addr::TargetAddr;
+use fast_socks5::client::{Config, Socks5Stream};
 
 #[derive(err_derive::Error, Debug)]
 #[error(no_from)]
@@ -19,6 +21,7 @@ pub struct Handle {
     bind_addr: SocketAddr,
 }
 
+/// Spawn a SOCKS server bound to `bind_addr`
 pub async fn spawn(bind_addr: SocketAddr) -> Result<Handle, Error> {
     let bind_addr = match bind_addr.port() {
         0 => SocketAddr::new(bind_addr.ip(), find_free_port(bind_addr.ip())?),
@@ -36,6 +39,8 @@ pub async fn spawn(bind_addr: SocketAddr) -> Result<Handle, Error> {
             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"),
@@ -52,6 +57,61 @@ pub async fn spawn(bind_addr: SocketAddr) -> Result<Handle, Error> {
     Ok(Handle { handle, bind_addr })
 }
 
+/// Spawn a SOCKS server that forwards everything via another SOCKS server at `via_socks_server`
+pub async fn spawn_via(bind_addr: SocketAddr, via_socks_server: SocketAddr)-> Result<Handle, Error> {
+    let bind_addr = match bind_addr.port() {
+        0 => SocketAddr::new(bind_addr.ip(), find_free_port(bind_addr.ip())?),
+        _ => bind_addr,
+    };
+    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(mut socket) => {
+                    let (target_host, target_port) = match socket.target_addr().cloned() {
+                        Some(TargetAddr::Domain(host, port)) => (host, port),
+                        Some(TargetAddr::Ip(addr)) => (addr.ip().to_string(), addr.port()),
+                        None => {
+                            log::error!("missing target address");
+                            continue;
+                        }
+                    };
+
+                    let mut client = match Socks5Stream::connect(
+                        via_socks_server,
+                        target_host,
+                        target_port,
+                        Config::default(),
+                    ).await {
+                        Ok(client) => client,
+                        Err(error) => {
+                            log::error!("failed to connect client to proxy: {error}");
+                            continue;
+                        }
+                    };
+
+                    tokio::spawn(async move {
+                        match tokio::io::copy_bidirectional(&mut client, &mut socket).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, bind_addr })
+}
+
 impl Handle {
     pub fn bind_addr(&self) -> SocketAddr {
         self.bind_addr
diff --git a/test/test-rpc/src/client.rs b/test/test-rpc/src/client.rs
@@ -218,8 +218,9 @@ impl ServiceClient {
     pub async fn start_socks_server(
         &self,
         bind_addr: SocketAddr,
+        via_socks_server: Option<SocketAddr>,
     ) -> Result<crate::net::SocksHandle, Error> {
-        crate::net::SocksHandle::start_server(self.client.clone(), bind_addr).await
+        crate::net::SocksHandle::start_server(self.client.clone(), bind_addr, via_socks_server).await
     }
 
     /// Restarts the app.
diff --git a/test/test-rpc/src/lib.rs b/test/test-rpc/src/lib.rs
@@ -150,10 +150,12 @@ mod service {
         /// Perform DNS resolution.
         async fn resolve_hostname(hostname: String) -> Result<Vec<SocketAddr>, Error>;
 
-        /// Start a SOCKS5 server bound to the given address. Return an ID that can be used with
-        /// `stop_socks_server`, and the address that the listening socket was actually bound to.
+    /// Start a SOCKS5 server bound to the given address. This optionally proxies everything via a
+    /// SOCKS server at `via_socks_server`. Return an ID that can be used with `stop_socks_server`,
+    /// and the address that the listening socket was actually bound to.
         async fn start_socks_server(
             bind_addr: SocketAddr,
+            via_socks_server: Option<SocketAddr>,
         ) -> Result<(net::SocksHandleId, SocketAddr), Error>;
 
         /// Stop a SOCKS5 server that was previously started with `start_socks_server`.
diff --git a/test/test-rpc/src/net.rs b/test/test-rpc/src/net.rs
@@ -31,11 +31,12 @@ impl SocksHandle {
     pub(crate) async fn start_server(
         client: crate::service::ServiceClient,
         bind_addr: SocketAddr,
+        via_socks_server: Option<SocketAddr>,
     ) -> Result<Self, Error> {
         let (stop_tx, stop_rx) = oneshot::channel();
 
         let (id, bind_addr) = client
-            .start_socks_server(tarpc::context::current(), bind_addr)
+            .start_socks_server(tarpc::context::current(), bind_addr, via_socks_server)
             .await??;
 
         tokio::spawn(async move {
diff --git a/test/test-runner/src/main.rs b/test/test-runner/src/main.rs
@@ -169,14 +169,16 @@ impl Service for TestServer {
             .collect())
     }
 
-    /// Start a SOCKS5 server bound to the given address. Return an ID that can be used with
-    /// `stop_socks_server`, and the address that the listening socket was actually bound to.
+    /// Start a SOCKS5 server bound to the given address. This optionally proxies everything via a
+    /// SOCKS server at `via_socks_server`. Return an ID that can be used with `stop_socks_server`,
+    /// and the address that the listening socket was actually bound to.
     async fn start_socks_server(
         self,
         _: context::Context,
         bind_addr: SocketAddr,
+        via_socks_server: Option<SocketAddr>,
     ) -> Result<(SocksHandleId, SocketAddr), test_rpc::Error> {
-        socks::start_server(bind_addr).await
+        socks::start_server(bind_addr, via_socks_server).await
     }
 
     /// Stop a SOCKS5 server that was previously started with `start_socks_server`.
diff --git a/test/test-runner/src/socks.rs b/test/test-runner/src/socks.rs
@@ -8,16 +8,21 @@ use test_rpc::net::SocksHandleId;
 static SERVERS: Lazy<Mutex<HashMap<SocksHandleId, socks_server::Handle>>> =
     Lazy::new(|| Mutex::new(HashMap::new()));
 
+/// Spawn a SOCKS server that sends everything through the test-manager SOCKS server
 pub async fn start_server(
     bind_addr: SocketAddr,
+    via_socks_server: Option<SocketAddr>,
 ) -> Result<(SocksHandleId, SocketAddr), test_rpc::Error> {
     let next_nonce = {
         static NONCE: AtomicUsize = AtomicUsize::new(0);
         || NONCE.fetch_add(1, Ordering::Relaxed)
     };
     let id = SocksHandleId(next_nonce());
 
-    let handle = socks_server::spawn(bind_addr).await.map_err(|error| {
+    let handle = match via_socks_server {
+        Some(proxy) => socks_server::spawn_via(bind_addr, proxy).await,
+        None => socks_server::spawn(bind_addr).await,
+    }.map_err(|error| {
         log::error!("Failed to spawn SOCKS server: {error}");
         test_rpc::Error::SocksServer
     })?;
@@ -30,6 +35,7 @@ pub async fn start_server(
     Ok((id, bind_addr))
 }
 
+/// Stop SOCKS server given some ID returned by `start_server`
 pub async fn stop_server(id: SocksHandleId) -> Result<(), test_rpc::Error> {
     let handle = {
         let mut servers = SERVERS.lock().unwrap();
