feat(local): unified TCP & UDP listener creation

shadowsocks · Oct 21, 2023 · f059d30 · f059d30
1 parent bfc4741
commit f059d30
Show file tree

Hide file tree

Showing 10 changed files with 151 additions and 213 deletions.
diff --git a/crates/shadowsocks-service/src/local/dns/server.rs b/crates/shadowsocks-service/src/local/dns/server.rs
@@ -32,15 +32,15 @@ use tokio::{
 
 use shadowsocks::{
     config::Mode,
-    lookup_then,
-    net::{TcpListener, UdpSocket as ShadowUdpSocket},
+    net::TcpListener,
     relay::{udprelay::MAXIMUM_UDP_PAYLOAD_SIZE, Address},
     ServerAddr,
 };
 
 use crate::{
     acl::AccessControl,
     local::{context::ServiceContext, loadbalancing::PingBalancer},
+    net::listener::{create_standard_tcp_listener, create_standard_udp_listener},
 };
 
 use super::{client_cache::DnsClientCache, config::NameServerAddr};
@@ -213,34 +213,18 @@ impl DnsTcpServerBuilder {
                     use tokio::net::TcpListener as TokioTcpListener;
                     use crate::net::launch_activate_socket::get_launch_activate_tcp_listener;
 
-                    let std_listener = get_launch_activate_tcp_listener(&launchd_socket_name)?;
-                    let tokio_listener = TokioTcpListener::from_std(std_listener)?;
-                    TcpListener::from_listener(tokio_listener, self.context.accept_opts())?
-                } else {
-                    match self.bind_addr {
-                        ServerAddr::SocketAddr(ref saddr) => {
-                            TcpListener::bind_with_opts(saddr, self.context.accept_opts()).await?
-                        }
-                        ServerAddr::DomainName(ref dname, port) => {
-                            lookup_then!(self.context.context_ref(), dname, port, |addr| {
-                                TcpListener::bind_with_opts(&addr, self.context.accept_opts()).await
-                            })?
-                            .1
+                    match get_launch_activate_tcp_listener(&launchd_socket_name)? {
+                        Some(std_listener) => {
+                            let tokio_listener = TokioTcpListener::from_std(std_listener)?;
+                            TcpListener::from_listener(tokio_listener, self.context.accept_opts())?
                         }
+                        None => create_standard_tcp_listener(&self.context, &self.bind_addr).await?
                     }
+                } else {
+                    create_standard_tcp_listener(&self.context, &self.bind_addr).await?
                 };
             } else {
-                let listener = match self.bind_addr {
-                    ServerAddr::SocketAddr(ref saddr) => {
-                        TcpListener::bind_with_opts(saddr, self.context.accept_opts()).await?
-                    }
-                    ServerAddr::DomainName(ref dname, port) => {
-                        lookup_then!(self.context.context_ref(), dname, port, |addr| {
-                            TcpListener::bind_with_opts(&addr, self.context.accept_opts()).await
-                        })?
-                        .1
-                    }
-                };
+                let listener = create_standard_tcp_listener(&self.context, &self.bind_addr).await?;
             }
         }
 
@@ -406,35 +390,15 @@ impl DnsUdpServerBuilder {
                     use tokio::net::UdpSocket as TokioUdpSocket;
                     use crate::net::launch_activate_socket::get_launch_activate_udp_socket;
 
-                    let std_socket = get_launch_activate_udp_socket(&launchd_socket_name)?;
-                    TokioUdpSocket::from_std(std_socket)?
+                    match get_launch_activate_udp_socket(&launchd_socket_name)? {
+                        Some(std_socket) => TokioUdpSocket::from_std(std_socket)?,
+                        None => create_standard_udp_listener(&self.context, &self.bind_addr).await?.into()
+                    }
                 } else {
-                    let shadow_socket = match self.bind_addr {
-                        ServerAddr::SocketAddr(ref saddr) => {
-                            ShadowUdpSocket::listen_with_opts(saddr, self.context.accept_opts()).await?
-                        }
-                        ServerAddr::DomainName(ref dname, port) => {
-                            lookup_then!(self.context.context_ref(), dname, port, |addr| {
-                                ShadowUdpSocket::listen_with_opts(&addr, self.context.accept_opts()).await
-                            })?
-                            .1
-                        }
-                    };
-                    shadow_socket.into()
+                    create_standard_udp_listener(&self.context, &self.bind_addr).await?.into()
                 };
             } else {
-                let shadow_socket = match self.bind_addr {
-                    ServerAddr::SocketAddr(ref saddr) => {
-                        ShadowUdpSocket::listen_with_opts(saddr, self.context.accept_opts()).await?
-                    }
-                    ServerAddr::DomainName(ref dname, port) => {
-                        lookup_then!(self.context.context_ref(), dname, port, |addr| {
-                            ShadowUdpSocket::listen_with_opts(&addr, self.context.accept_opts()).await
-                        })?
-                        .1
-                    }
-                };
-                let socket = shadow_socket.into();
+                let socket = create_standard_udp_listener(&self.context, &self.bind_addr).await?.into();
             }
         }
 

diff --git a/crates/shadowsocks-service/src/local/http/server.rs b/crates/shadowsocks-service/src/local/http/server.rs
@@ -16,13 +16,16 @@ use hyper::{
     Server,
 };
 use log::{error, info};
-use shadowsocks::{config::ServerAddr, lookup_then, net::TcpListener};
-
-use crate::local::{
-    context::ServiceContext,
-    http::connector::Connector,
-    loadbalancing::PingBalancer,
-    LOCAL_DEFAULT_KEEPALIVE_TIMEOUT,
+use shadowsocks::{config::ServerAddr, net::TcpListener};
+
+use crate::{
+    local::{
+        context::ServiceContext,
+        http::connector::Connector,
+        loadbalancing::PingBalancer,
+        LOCAL_DEFAULT_KEEPALIVE_TIMEOUT,
+    },
+    net::listener::create_standard_tcp_listener,
 };
 
 use super::{client_cache::ProxyClientCache, dispatcher::HttpDispatcher};
@@ -71,26 +74,18 @@ impl HttpBuilder {
                     use tokio::net::TcpListener as TokioTcpListener;
                     use crate::net::launch_activate_socket::get_launch_activate_tcp_listener;
 
-                    let std_listener = get_launch_activate_tcp_listener(&launchd_socket_name)?;
-                    let tokio_listener = TokioTcpListener::from_std(std_listener)?;
-                    TcpListener::from_listener(tokio_listener, self.context.accept_opts())?
+                    match get_launch_activate_tcp_listener(&launchd_socket_name)? {
+                        Some(std_listener) => {
+                            let tokio_listener = TokioTcpListener::from_std(std_listener)?;
+                            TcpListener::from_listener(tokio_listener, self.context.accept_opts())?
+                        }
+                        None => create_standard_tcp_listener(&self.context, &self.client_config).await?
+                    }
                 } else {
-                    match self.client_config {
-                        ServerAddr::SocketAddr(sa) => TcpListener::bind_with_opts(&sa, self.context.accept_opts().clone()).await,
-                        ServerAddr::DomainName(ref dname, port) => lookup_then!(self.context.context_ref(), dname, port, |addr| {
-                            TcpListener::bind_with_opts(&addr, self.context.accept_opts().clone()).await
-                        })
-                        .map(|(_, b)| b),
-                    }?
+                    create_standard_tcp_listener(&self.context, &self.client_config).await?
                 };
             } else {
-                let listener = match self.client_config {
-                    ServerAddr::SocketAddr(sa) => TcpListener::bind_with_opts(&sa, self.context.accept_opts().clone()).await,
-                    ServerAddr::DomainName(ref dname, port) => lookup_then!(self.context.context_ref(), dname, port, |addr| {
-                        TcpListener::bind_with_opts(&addr, self.context.accept_opts().clone()).await
-                    })
-                    .map(|(_, b)| b),
-                }?;
+                let listener = create_standard_tcp_listener(&self.context, &self.client_config).await?;
             }
         }
 

diff --git a/crates/shadowsocks-service/src/local/socks/server/server.rs b/crates/shadowsocks-service/src/local/socks/server/server.rs
@@ -1,17 +1,18 @@
 use std::{io, net::SocketAddr, sync::Arc, time::Duration};
 
 use log::{error, info};
-use shadowsocks::{config::Mode, lookup_then, net::TcpListener as ShadowTcpListener, ServerAddr};
+use shadowsocks::{config::Mode, net::TcpListener as ShadowTcpListener, ServerAddr};
 use tokio::{net::TcpStream, time};
 
-use crate::local::{context::ServiceContext, loadbalancing::PingBalancer};
+use crate::{
+    local::{context::ServiceContext, loadbalancing::PingBalancer, socks::config::Socks5AuthConfig},
+    net::listener::create_standard_tcp_listener,
+};
 
 #[cfg(feature = "local-socks4")]
 use super::socks4::Socks4TcpHandler;
 use super::socks5::{Socks5TcpHandler, Socks5UdpServer};
 
-use crate::local::socks::config::Socks5AuthConfig;
-
 pub struct SocksTcpServerBuilder {
     context: Arc<ServiceContext>,
     client_config: ServerAddr,
@@ -57,34 +58,18 @@ impl SocksTcpServerBuilder {
                     use tokio::net::TcpListener as TokioTcpListener;
                     use crate::net::launch_activate_socket::get_launch_activate_tcp_listener;
 
-                    let std_listener = get_launch_activate_tcp_listener(&launchd_socket_name)?;
-                    let tokio_listener = TokioTcpListener::from_std(std_listener)?;
-                    ShadowTcpListener::from_listener(tokio_listener, self.context.accept_opts())?
-                } else {
-                    match self.client_config {
-                        ServerAddr::SocketAddr(ref saddr) => {
-                            ShadowTcpListener::bind_with_opts(saddr, self.context.accept_opts()).await?
-                        }
-                        ServerAddr::DomainName(ref dname, port) => {
-                            lookup_then!(self.context.context_ref(), dname, port, |addr| {
-                                ShadowTcpListener::bind_with_opts(&addr, self.context.accept_opts()).await
-                            })?
-                            .1
+                    match get_launch_activate_tcp_listener(&launchd_socket_name)? {
+                        Some(std_listener) => {
+                            let tokio_listener = TokioTcpListener::from_std(std_listener)?;
+                            ShadowTcpListener::from_listener(tokio_listener, self.context.accept_opts())?
                         }
+                        None => create_standard_tcp_listener(&self.context, &self.client_config).await?
                     }
+                } else {
+                    create_standard_tcp_listener(&self.context, &self.client_config).await?
                 };
             } else {
-                let listener = match self.client_config {
-                    ServerAddr::SocketAddr(ref saddr) => {
-                        ShadowTcpListener::bind_with_opts(saddr, self.context.accept_opts()).await?
-                    }
-                    ServerAddr::DomainName(ref dname, port) => {
-                        lookup_then!(self.context.context_ref(), dname, port, |addr| {
-                            ShadowTcpListener::bind_with_opts(&addr, self.context.accept_opts()).await
-                        })?
-                        .1
-                    }
-                };
+                let listener = create_standard_tcp_listener(&self.context, &self.client_config).await?;
             }
         }
 

diff --git a/crates/shadowsocks-service/src/local/socks/server/socks5/udprelay.rs b/crates/shadowsocks-service/src/local/socks/server/socks5/udprelay.rs
@@ -12,8 +12,6 @@ use byte_string::ByteStr;
 use bytes::{BufMut, BytesMut};
 use log::{debug, error, info, trace};
 use shadowsocks::{
-    lookup_then,
-    net::UdpSocket as ShadowUdpSocket,
     relay::{
         socks5::{Address, UdpAssociateHeader},
         udprelay::MAXIMUM_UDP_PAYLOAD_SIZE,
@@ -28,7 +26,7 @@ use crate::{
         loadbalancing::PingBalancer,
         net::{UdpAssociationManager, UdpInboundWrite},
     },
-    net::utils::to_ipv4_mapped,
+    net::{listener::create_standard_udp_listener, utils::to_ipv4_mapped},
 };
 
 pub struct Socks5UdpServerBuilder {
@@ -73,35 +71,15 @@ impl Socks5UdpServerBuilder {
                     use tokio::net::UdpSocket as TokioUdpSocket;
                     use crate::net::launch_activate_socket::get_launch_activate_udp_socket;
 
-                    let std_socket = get_launch_activate_udp_socket(&launchd_socket_name)?;
-                    TokioUdpSocket::from_std(std_socket)?
+                    match get_launch_activate_udp_socket(&launchd_socket_name)? {
+                        Some(std_socket) => TokioUdpSocket::from_std(std_socket)?,
+                        None => create_standard_udp_listener(&self.context, &self.client_config).await?.into()
+                    }
                 } else {
-                    let shadow_socket = match self.client_config {
-                        ServerAddr::SocketAddr(ref saddr) => {
-                            ShadowUdpSocket::listen_with_opts(saddr, self.context.accept_opts()).await?
-                        }
-                        ServerAddr::DomainName(ref dname, port) => {
-                            lookup_then!(self.context.context_ref(), dname, port, |addr| {
-                                ShadowUdpSocket::listen_with_opts(&addr, self.context.accept_opts()).await
-                            })?
-                            .1
-                        }
-                    };
-                    shadow_socket.into()
+                    create_standard_udp_listener(&self.context, &self.client_config).await?.into()
                 };
             } else {
-                let shadow_socket = match self.client_config {
-                    ServerAddr::SocketAddr(ref saddr) => {
-                        ShadowUdpSocket::listen_with_opts(saddr, self.context.accept_opts()).await?
-                    }
-                    ServerAddr::DomainName(ref dname, port) => {
-                        lookup_then!(self.context.context_ref(), dname, port, |addr| {
-                            ShadowUdpSocket::listen_with_opts(&addr, self.context.accept_opts()).await
-                        })?
-                        .1
-                    }
-                };
-                let socket = shadow_socket.into();
+                let socket = create_standard_udp_listener(&self.context, &self.client_config).await?.into();
             }
         }
 

diff --git a/crates/shadowsocks-service/src/local/tunnel/tcprelay.rs b/crates/shadowsocks-service/src/local/tunnel/tcprelay.rs
@@ -3,14 +3,17 @@
 use std::{io, net::SocketAddr, sync::Arc, time::Duration};
 
 use log::{error, info, trace};
-use shadowsocks::{lookup_then, net::TcpListener as ShadowTcpListener, relay::socks5::Address, ServerAddr};
+use shadowsocks::{net::TcpListener as ShadowTcpListener, relay::socks5::Address, ServerAddr};
 use tokio::{net::TcpStream, time};
 
-use crate::local::{
-    context::ServiceContext,
-    loadbalancing::PingBalancer,
-    net::AutoProxyClientStream,
-    utils::{establish_tcp_tunnel, establish_tcp_tunnel_bypassed},
+use crate::{
+    local::{
+        context::ServiceContext,
+        loadbalancing::PingBalancer,
+        net::AutoProxyClientStream,
+        utils::{establish_tcp_tunnel, establish_tcp_tunnel_bypassed},
+    },
+    net::listener::create_standard_tcp_listener,
 };
 
 pub struct TunnelTcpServerBuilder {
@@ -52,34 +55,18 @@ impl TunnelTcpServerBuilder {
                     use tokio::net::TcpListener as TokioTcpListener;
                     use crate::net::launch_activate_socket::get_launch_activate_tcp_listener;
 
-                    let std_listener = get_launch_activate_tcp_listener(&launchd_socket_name)?;
-                    let tokio_listener = TokioTcpListener::from_std(std_listener)?;
-                    ShadowTcpListener::from_listener(tokio_listener, self.context.accept_opts())?
-                } else {
-                    match self.client_config {
-                        ServerAddr::SocketAddr(ref saddr) => {
-                            ShadowTcpListener::bind_with_opts(saddr, self.context.accept_opts()).await?
-                        }
-                        ServerAddr::DomainName(ref dname, port) => {
-                            lookup_then!(self.context.context_ref(), dname, port, |addr| {
-                                ShadowTcpListener::bind_with_opts(&addr, self.context.accept_opts()).await
-                            })?
-                            .1
+                    match get_launch_activate_tcp_listener(&launchd_socket_name)? {
+                        Some(std_listener) => {
+                            let tokio_listener = TokioTcpListener::from_std(std_listener)?;
+                            ShadowTcpListener::from_listener(tokio_listener, self.context.accept_opts())?
                         }
+                        None => create_standard_tcp_listener(&self.context, &self.client_config).await?
                     }
+                } else {
+                    create_standard_tcp_listener(&self.context, &self.client_config).await?
                 };
             } else {
-                let listener = match self.client_config {
-                    ServerAddr::SocketAddr(ref saddr) => {
-                        ShadowTcpListener::bind_with_opts(saddr, self.context.accept_opts()).await?
-                    }
-                    ServerAddr::DomainName(ref dname, port) => {
-                        lookup_then!(self.context.context_ref(), dname, port, |addr| {
-                            ShadowTcpListener::bind_with_opts(&addr, self.context.accept_opts()).await
-                        })?
-                        .1
-                    }
-                };
+                let listener = create_standard_tcp_listener(&self.context, &self.client_config).await?;
             }
         }