Add tproxy udp

erebe · Dec 3, 2023 · 592fab9 · 592fab9
1 parent 430f19a
commit 592fab9
Show file tree

Hide file tree

Showing 6 changed files with 223 additions and 71 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -7,7 +7,7 @@ repository = "https://github.com/erebe/wstunnel.git"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-clap = { version = "4.4.8", features = ["derive"]}
+clap = { version = "4.4.10", features = ["derive"]}
 url = "2.5.0"
 anyhow = "1.0.75"
 
@@ -24,6 +24,7 @@ rustls-pemfile = { version = "1.0.4", features = [] }
 bytes = { version = "1.5.0", features = [] }
 parking_lot = "0.12.1"
 urlencoding = "2.1.3"
+nix = { version = "0.27.1", features = ["socket", "net", "uio"] }
 
 rustls-native-certs = { version = "0.6.3", features = [] }
 tokio = { version = "1.34.0", features = ["full"] }

diff --git a/src/main.rs b/src/main.rs
@@ -51,17 +51,18 @@ enum Commands {
 struct Client {
     /// Listen on local and forwards traffic from remote. Can be specified multiple times
     /// examples:
-    /// 'tcp://1212:google.com:443'      =>     listen locally on tcp on port 1212 and forward to google.com on port 443
+    /// 'tcp://1212:google.com:443'      =>       listen locally on tcp on port 1212 and forward to google.com on port 443
     ///
-    /// 'udp://1212:1.1.1.1:53'          =>     listen locally on udp on port 1212 and forward to cloudflare dns 1.1.1.1 on port 53
-    /// 'udp://1212:1.1.1.1:53?timeout_sec=10'  timeout_sec on udp force close the tunnel after 10sec. Set it to 0 to disable the timeout [default: 30]
+    /// 'udp://1212:1.1.1.1:53'          =>       listen locally on udp on port 1212 and forward to cloudflare dns 1.1.1.1 on port 53
+    /// 'udp://1212:1.1.1.1:53?timeout_sec=10'    timeout_sec on udp force close the tunnel after 10sec. Set it to 0 to disable the timeout [default: 30]
     ///
-    /// 'socks5://[::1]:1212'            =>     listen locally with socks5 on port 1212 and forward dynamically requested tunnel
+    /// 'socks5://[::1]:1212'            =>       listen locally with socks5 on port 1212 and forward dynamically requested tunnel
     ///
-    /// 'tproxy+tcp://[::1]:1212'        =>     listen locally on tcp on port 1212 as a *transparent proxy* and forward dynamically requested tunnel
-    ///                                         linux only and requires sudo/CAP_NET_ADMIN
+    /// 'tproxy+tcp://[::1]:1212'        =>       listen locally on tcp on port 1212 as a *transparent proxy* and forward dynamically requested tunnel
+    /// 'tproxy+udp://[::1]:1212?timeout_sec=10'  listen locally on udp on port 1212 as a *transparent proxy* and forward dynamically requested tunnel
+    ///                                           linux only and requires sudo/CAP_NET_ADMIN
     ///
-    /// 'stdio://google.com:443'         =>     listen for data from stdio, mainly for `ssh -o ProxyCommand="wstunnel client -L stdio://%h:%p ws://localhost:8080" my-server`
+    /// 'stdio://google.com:443'         =>       listen for data from stdio, mainly for `ssh -o ProxyCommand="wstunnel client -L stdio://%h:%p ws://localhost:8080" my-server`
     #[arg(short='L', long, value_name = "{tcp,udp,socks5,stdio}://[BIND:]PORT:HOST:PORT", value_parser = parse_tunnel_arg, verbatim_doc_comment)]
     local_to_remote: Vec<LocalToRemote>,
 
@@ -180,6 +181,7 @@ enum LocalProtocol {
     Stdio,
     Socks5,
     TProxyTcp,
+    TProxyUdp { timeout: Option<Duration> },
     ReverseTcp,
     ReverseUdp { timeout: Option<Duration> },
     ReverseSocks5,
@@ -334,6 +336,21 @@ fn parse_tunnel_arg(arg: &str) -> Result<LocalToRemote, io::Error> {
                     remote: (dest_host, dest_port),
                 })
             }
+            "tproxy+u" => {
+                let (local_bind, remaining) = parse_local_bind(&arg["tproxy+udp://".len()..])?;
+                let x = format!("0.0.0.0:0?{}", remaining);
+                let (dest_host, dest_port, options) = parse_tunnel_dest(&x)?;
+                let timeout = options
+                    .get("timeout_sec")
+                    .and_then(|x| x.parse::<u64>().ok())
+                    .map(|d| if d == 0 { None } else { Some(Duration::from_secs(d)) })
+                    .unwrap_or(Some(Duration::from_secs(30)));
+                Ok(LocalToRemote {
+                    local_protocol: LocalProtocol::TProxyUdp { timeout },
+                    local: local_bind,
+                    remote: (dest_host, dest_port),
+                })
+            }
             _ => Err(Error::new(
                 ErrorKind::InvalidInput,
                 format!("Invalid local protocol for tunnel {}", arg),
@@ -644,7 +661,7 @@ async fn main() {
                     LocalProtocol::TProxyTcp => {
                         let server = tcp::run_server(tunnel.local, true)
                             .await
-                            .unwrap_or_else(|err| panic!("Cannot start TProxy server on {}: {}", tunnel.local, err))
+                            .unwrap_or_else(|err| panic!("Cannot start TProxy TCP server on {}: {}", tunnel.local, err))
                             .map_err(anyhow::Error::new)
                             .map_ok(move |stream| {
                                 // In TProxy mode local destination is the final ip:port destination
@@ -658,13 +675,34 @@ async fn main() {
                             }
                         });
                     }
+                    #[cfg(target_os = "linux")]
+                    LocalProtocol::TProxyUdp { timeout } => {
+                        let server =
+                            udp::run_server(tunnel.local, *timeout, udp::configure_tproxy, udp::mk_send_socket_tproxy)
+                                .await
+                                .unwrap_or_else(|err| {
+                                    panic!("Cannot start TProxy UDP server on {}: {}", tunnel.local, err)
+                                })
+                                .map_err(anyhow::Error::new)
+                                .map_ok(move |stream| {
+                                    // In TProxy mode local destination is the final ip:port destination
+                                    let dest = to_host_port(stream.local_addr().unwrap());
+                                    (tokio::io::split(stream), dest)
+                                });
+
+                        tokio::spawn(async move {
+                            if let Err(err) = tunnel::client::run_tunnel(client_config, tunnel, server).await {
+                                error!("{:?}", err);
+                            }
+                        });
+                    }
                     #[cfg(not(target_os = "linux"))]
-                    LocalProtocol::TProxyTcp => {
+                    LocalProtocol::TProxyTcp | LocalProtocol::TProxyUdp { .. } => {
                         panic!("Transparent proxy is not available for non Linux platform")
                     }
                     LocalProtocol::Udp { timeout } => {
                         let remote = tunnel.remote.clone();
-                        let server = udp::run_server(tunnel.local, *timeout)
+                        let server = udp::run_server(tunnel.local, *timeout, |_| Ok(()), |s| Ok(s.clone()))
                             .await
                             .unwrap_or_else(|err| panic!("Cannot start UDP server on {}: {}", tunnel.local, err))
                             .map_err(anyhow::Error::new)

diff --git a/src/tunnel/mod.rs b/src/tunnel/mod.rs
@@ -40,6 +40,7 @@ impl JwtTunnelConfig {
                 LocalProtocol::ReverseUdp { .. } => tunnel.local_protocol,
                 LocalProtocol::ReverseSocks5 => LocalProtocol::ReverseSocks5,
                 LocalProtocol::TProxyTcp => LocalProtocol::Tcp,
+                LocalProtocol::TProxyUdp { timeout } => LocalProtocol::Udp { timeout },
             },
             r: tunnel.remote.0.to_string(),
             rp: tunnel.remote.1,

diff --git a/src/tunnel/server.rs b/src/tunnel/server.rs
@@ -103,7 +103,8 @@ async fn from_query(
 
             let local_srv = (Host::parse(&jwt.claims.r)?, jwt.claims.rp);
             let bind = format!("{}:{}", local_srv.0, local_srv.1);
-            let listening_server = udp::run_server(bind.parse()?, timeout);
+            let listening_server =
+                udp::run_server(bind.parse()?, timeout, |_| Ok(()), |send_socket| Ok(send_socket.clone()));
             let udp = run_listening_server(&local_srv, SERVERS.deref(), listening_server).await?;
             let (local_rx, local_tx) = tokio::io::split(udp);