temp: fwd tcp

Author: dlon

test/socks-server/src/lib.rs

@@ -4,6 +4,8 @@ use std::net::IpAddr;
 use std::net::SocketAddr;
 use fast_socks5::util::target_addr::TargetAddr;
 use fast_socks5::client::{Config, Socks5Stream};
+use tokio::net::TcpStream;
+use tokio::net::TcpListener;
 
 #[derive(err_derive::Error, Debug)]
 #[error(no_from)]
@@ -14,6 +16,12 @@ pub enum Error {
     BindTempSocket(#[error(source)] io::Error),
     #[error(display = "Failed to find free port")]
     GetTempAddress(#[error(source)] io::Error),
+    #[error(display = "Failed to bind TCP listener")]
+    BindTcpListener(#[error(source)] io::Error),
+    #[error(display = "Failed to get TCP listener local addr")]
+    GetTcpListenerAddress(#[error(source)] io::Error),
+    #[error(display = "Failed to connect client to remote")]
+    ConnectToRemote(#[error(source)] io::Error),
 }
 
 pub struct Handle {
@@ -57,54 +65,35 @@ 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)?;
+/// Forward TCP traffic via `proxy_addr`
+pub async fn forward(bind_addr: SocketAddr, proxy_addr: SocketAddr)-> Result<Handle, Error> {
+    let listener = TcpListener::bind(&bind_addr).await.map_err(Error::BindTcpListener)?;
+    let bind_addr = listener.local_addr().map_err(Error::GetTcpListenerAddress)?;
 
     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;
-                        }
-                    };
-
+        loop {
+            match listener.accept().await {
+                Ok((mut client, _addr)) => {
+                    // FIXME: close these tasks
                     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}"),
+                        let mut proxy = match TcpStream::connect(proxy_addr).await {
+                            Ok(proxy) => proxy,
+                            Err(error) => {
+                                log::error!("failed to connect to TCP proxy: {error}");
+                                return;
+                            }
+                        };
+
+                        match tokio::io::copy_bidirectional(&mut client, &mut proxy).await {
+                            Ok((a_bytes, b_bytes)) => {
+                                log::info!("wrote {a_bytes} <-> {b_bytes} bytes, client <-> proxy");
+                            }
+                            Err(error) => log::error!("copy_directional failed: {error}"),
                         }
                     });
                 }
                 Err(error) => {
-                    log::error!("failed to accept socks client: {error}");
+                    log::error!("failed to accept TCP client: {error}");
                 }
             }
         }

test/test-runner/src/socks.rs

@@ -8,7 +8,7 @@ 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
+/// Spawn a SOCKS server that optionally forwards through a SOCKS server at `via_socks_server`
 pub async fn start_server(
     bind_addr: SocketAddr,
     via_socks_server: Option<SocketAddr>,
@@ -20,7 +20,7 @@ pub async fn start_server(
     let id = SocksHandleId(next_nonce());
 
     let handle = match via_socks_server {
-        Some(proxy) => socks_server::spawn_via(bind_addr, proxy).await,
+        Some(proxy) => socks_server::forward(bind_addr, proxy).await,
         None => socks_server::spawn(bind_addr).await,
     }.map_err(|error| {
         log::error!("Failed to spawn SOCKS server: {error}");