From ab730bad36caaaba843a520faea4e3fc22ea2e0a Mon Sep 17 00:00:00 2001
From: zonyitoo <zonyitoo@gmail.com>
Date: Mon, 12 Oct 2020 00:31:24 +0800
Subject: [PATCH] sslocal support customizing udp assoc bind address

ref: shadowsocks/shadowsocks-android#2571
---
 src/bin/local.rs                   |  6 +++++
 src/config.rs                      | 41 ++++++++++++++++++++++++++++++
 src/relay/tcprelay/socks5_local.rs | 12 ++++++++-
 src/relay/udprelay/socks5_local.rs |  9 +++++--
 4 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/src/bin/local.rs b/src/bin/local.rs
index ecd489e26ca0..874a3d424959 100644
--- a/src/bin/local.rs
+++ b/src/bin/local.rs
@@ -86,6 +86,8 @@ fn main() {
 
         (@arg UDP_TIMEOUT: --("udp-timeout") +takes_value {validator::validate_u64} "Timeout seconds for UDP relay")
         (@arg UDP_MAX_ASSOCIATIONS: --("udp-max-associations") +takes_value {validator::validate_u64} "Maximum associations to be kept simultaneously for UDP relay")
+
+        (@arg UDP_BIND_ADDR: --("udp-bind-addr") +takes_value {validator::validate_server_addr} "UDP relay's bind address, default is the same as local-addr")
     );
 
     // FIXME: -6 is not a identifier, so we cannot build it with clap_app!
@@ -338,6 +340,10 @@ fn main() {
         config.udp_max_associations = Some(udp_max_assoc.parse::<usize>().expect("udp-max-associations"));
     }
 
+    if let Some(udp_bind_addr) = matches.value_of("UDP_BIND_ADDR") {
+        config.udp_bind_addr = Some(udp_bind_addr.parse::<ServerAddr>().expect("udp-bind-addr"));
+    }
+
     // DONE READING options
 
     if config.local_addr.is_none() {
diff --git a/src/config.rs b/src/config.rs
index bdb387c95313..090c553232a3 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -218,6 +218,42 @@ impl<I: Into<String>> From<(I, u16)> for ServerAddr {
     }
 }
 
+impl From<Address> for ServerAddr {
+    fn from(addr: Address) -> ServerAddr {
+        match addr {
+            Address::SocketAddress(sa) => ServerAddr::SocketAddr(sa),
+            Address::DomainNameAddress(dn, port) => ServerAddr::DomainName(dn, port),
+        }
+    }
+}
+
+impl From<&Address> for ServerAddr {
+    fn from(addr: &Address) -> ServerAddr {
+        match *addr {
+            Address::SocketAddress(sa) => ServerAddr::SocketAddr(sa),
+            Address::DomainNameAddress(ref dn, port) => ServerAddr::DomainName(dn.clone(), port),
+        }
+    }
+}
+
+impl From<ServerAddr> for Address {
+    fn from(addr: ServerAddr) -> Address {
+        match addr {
+            ServerAddr::SocketAddr(sa) => Address::SocketAddress(sa),
+            ServerAddr::DomainName(dn, port) => Address::DomainNameAddress(dn, port),
+        }
+    }
+}
+
+impl From<&ServerAddr> for Address {
+    fn from(addr: &ServerAddr) -> Address {
+        match *addr {
+            ServerAddr::SocketAddr(sa) => Address::SocketAddress(sa),
+            ServerAddr::DomainName(ref dn, port) => Address::DomainNameAddress(dn.clone(), port),
+        }
+    }
+}
+
 /// Configuration for a server
 #[derive(Clone, Debug)]
 pub struct ServerConfig {
@@ -1011,6 +1047,10 @@ pub struct Config {
     pub udp_timeout: Option<Duration>,
     /// Maximum number of UDP Associations, default is unconfigured
     pub udp_max_associations: Option<usize>,
+    /// UDP relay's bind address, it uses `local_addr` by default
+    ///
+    /// Resolving Android's issue: https://github.com/shadowsocks/shadowsocks-android/issues/2571
+    pub udp_bind_addr: Option<ClientConfig>,
     /// `RLIMIT_NOFILE` option for *nix systems
     pub nofile: Option<u64>,
     /// ACL configuration
@@ -1129,6 +1169,7 @@ impl Config {
             config_type,
             udp_timeout: None,
             udp_max_associations: None,
+            udp_bind_addr: None,
             nofile: None,
             acl: None,
             tcp_redir: RedirType::tcp_default(),
diff --git a/src/relay/tcprelay/socks5_local.rs b/src/relay/tcprelay/socks5_local.rs
index cd7ce8628a6d..35a650650263 100644
--- a/src/relay/tcprelay/socks5_local.rs
+++ b/src/relay/tcprelay/socks5_local.rs
@@ -196,7 +196,17 @@ async fn handle_socks5_client(
             Ok(())
         }
         socks5::Command::UdpAssociate => {
-            if udp_conf.enable_udp {
+            if let Some(ref bind_addr) = server.config().udp_bind_addr {
+                debug!("UDP ASSOCIATE {}", addr);
+
+                let rh = TcpResponseHeader::new(socks5::Reply::Succeeded, bind_addr.into());
+                rh.write_to(&mut s).await?;
+
+                // Hold the connection until it ends by its own
+                ignore_until_end(&mut s).await?;
+
+                Ok(())
+            } else if udp_conf.enable_udp {
                 debug!("UDP ASSOCIATE {}", addr);
                 let rh = TcpResponseHeader::new(socks5::Reply::Succeeded, From::from(udp_conf.client_addr));
                 rh.write_to(&mut s).await?;
diff --git a/src/relay/udprelay/socks5_local.rs b/src/relay/udprelay/socks5_local.rs
index ebd8b120e663..2e8d744bebf0 100644
--- a/src/relay/udprelay/socks5_local.rs
+++ b/src/relay/udprelay/socks5_local.rs
@@ -76,8 +76,13 @@ fn assemble_packet(addr: Address, pkt: &[u8]) -> Bytes {
 
 /// Starts a UDP local server
 pub async fn run(context: SharedContext) -> io::Result<()> {
-    let local_addr = context.config().local_addr.as_ref().expect("local config");
-    let bind_addr = local_addr.bind_addr(&context).await?;
+    let bind_addr = match context.config().udp_bind_addr {
+        Some(ref bind_addr) => bind_addr.bind_addr(&context).await?,
+        None => {
+            let local_addr = context.config().local_addr.as_ref().expect("local config");
+            local_addr.bind_addr(&context).await?
+        }
+    };
 
     let l = create_udp_socket(&bind_addr).await?;
     let local_addr = l.local_addr().expect("determine port bound to");