From 442e752cba6f33a040bb196cbd707c9e1c7732ab Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Mon, 21 Oct 2024 14:59:26 +0800
Subject: [PATCH] udpgw feature

---
 Cargo.toml  |  5 +++++
 src/args.rs | 17 ++++++++++++++---
 src/lib.rs  | 13 ++++++++++---
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index c5aaf56..3ad63ce 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -13,6 +13,10 @@ rust-version = "1.80"
 [lib]
 crate-type = ["staticlib", "cdylib", "lib"]
 
+[features]
+default = ["udpgw"]
+udpgw = []
+
 [dependencies]
 async-trait = "0.1"
 base64 = { version = "0.22" }
@@ -68,6 +72,7 @@ path = "src/bin/main.rs"
 [[bin]]
 name = "udpgw-server"
 path = "src/bin/udpgw_server.rs"
+required-features = ["udpgw"]
 
 [profile.release]
 strip = "symbols"
diff --git a/src/args.rs b/src/args.rs
index 5eaa052..efcd2e7 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -113,10 +113,12 @@ pub struct Args {
     pub max_sessions: usize,
 
     /// UDP gateway server address, similar to badvpn-udpgw
+    #[cfg(feature = "udpgw")]
     #[arg(long, value_name = "IP:PORT")]
     pub udpgw_server: Option<SocketAddr>,
 
     /// Max udpgw connections
+    #[cfg(feature = "udpgw")]
     #[arg(long, value_name = "number")]
     pub udpgw_max_connections: Option<u16>,
 }
@@ -152,8 +154,6 @@ impl Default for Args {
             admin_command: Vec::new(),
             ipv6_enabled: false,
             setup,
-            udpgw_server: None,
-            udpgw_max_connections: None,
             dns: ArgDns::default(),
             dns_addr: "8.8.8.8".parse().unwrap(),
             bypass: vec![],
@@ -164,6 +164,10 @@ impl Default for Args {
             daemonize: false,
             exit_on_fatal_error: false,
             max_sessions: 200,
+            #[cfg(feature = "udpgw")]
+            udpgw_server: None,
+            #[cfg(feature = "udpgw")]
+            udpgw_max_connections: None,
         }
     }
 }
@@ -191,11 +195,18 @@ impl Args {
         self
     }
 
-    pub fn udpgw(&mut self, udpgw: SocketAddr) -> &mut Self {
+    #[cfg(feature = "udpgw")]
+    pub fn udpgw_server(&mut self, udpgw: SocketAddr) -> &mut Self {
         self.udpgw_server = Some(udpgw);
         self
     }
 
+    #[cfg(feature = "udpgw")]
+    pub fn udpgw_max_connections(&mut self, udpgw_max_connections: u16) -> &mut Self {
+        self.udpgw_max_connections = Some(udpgw_max_connections);
+        self
+    }
+
     #[cfg(unix)]
     pub fn tun_fd(&mut self, tun_fd: Option<i32>) -> &mut Self {
         self.tun_fd = tun_fd;
diff --git a/src/lib.rs b/src/lib.rs
index f373c41..7206478 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,19 +1,22 @@
+#[cfg(feature = "udpgw")]
+use crate::udpgw::UdpGwClient;
 use crate::{
     directions::{IncomingDataEvent, IncomingDirection, OutgoingDirection},
     http::HttpManager,
     no_proxy::NoProxyManager,
     session_info::{IpProtocol, SessionInfo},
-    udpgw::UdpGwClient,
     virtual_dns::VirtualDns,
 };
 use ipstack::stream::{IpStackStream, IpStackTcpStream, IpStackUdpStream};
 use proxy_handler::{ProxyHandler, ProxyHandlerManager};
 use socks::SocksProxyManager;
 pub use socks5_impl::protocol::UserKey;
+#[cfg(feature = "udpgw")]
+use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
 use std::{
     collections::VecDeque,
     io::ErrorKind,
-    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
+    net::{IpAddr, SocketAddr},
     sync::Arc,
 };
 use tokio::{
@@ -24,6 +27,7 @@ use tokio::{
 pub use tokio_util::sync::CancellationToken;
 use tproxy_config::is_private_ip;
 use udp_stream::UdpStream;
+#[cfg(feature = "udpgw")]
 use udpgw::{UdpGwClientStream, UdpGwResponse, UDPGW_KEEPALIVE_TIME};
 
 pub use {
@@ -61,6 +65,7 @@ mod session_info;
 pub mod socket_transfer;
 mod socks;
 mod traffic_status;
+#[cfg(feature = "udpgw")]
 pub mod udpgw;
 mod virtual_dns;
 #[doc(hidden)]
@@ -234,6 +239,7 @@ where
 
     let mut ip_stack = ipstack::IpStack::new(ipstack_config, device);
 
+    #[cfg(feature = "udpgw")]
     let udpgw_client = match args.udpgw_server {
         None => None,
         Some(addr) => {
@@ -290,7 +296,6 @@ where
                     if let Err(err) = handle_tcp_session(tcp, proxy_handler, socket_queue).await {
                         log::error!("{} error \"{}\"", info, err);
                     }
-
                     log::trace!("Session count {}", TASK_COUNT.fetch_sub(1, Relaxed) - 1);
                 });
             }
@@ -341,6 +346,7 @@ where
                 } else {
                     None
                 };
+                #[cfg(feature = "udpgw")]
                 if let Some(udpgw) = udpgw_client.clone() {
                     let tcp_src = match udp.peer_addr() {
                         SocketAddr::V4(_) => SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0)),
@@ -477,6 +483,7 @@ async fn handle_tcp_session(
     Ok(())
 }
 
+#[cfg(feature = "udpgw")]
 async fn handle_udp_gateway_session(
     mut udp_stack: IpStackUdpStream,
     udpgw_client: Arc<UdpGwClient>,