feat(shadowsocks): DatagramTransport & DatagramTransportExt

ProxySocket and MonProxySocket accepts socket as a generic type that implements DatagramTransport.
shadowsocks · Sep 20, 2024 · 60d0576 · 60d0576
1 parent 4e29581
commit 60d0576
Show file tree

Hide file tree

Showing 8 changed files with 246 additions and 145 deletions.
diff --git a/crates/shadowsocks-service/src/local/dns/upstream.rs b/crates/shadowsocks-service/src/local/dns/upstream.rs
@@ -61,7 +61,7 @@ pub enum DnsClient {
         stream: ProxyClientStream<MonProxyStream<ShadowTcpStream>>,
     },
     UdpRemote {
-        socket: MonProxySocket,
+        socket: MonProxySocket<ShadowUdpSocket>,
         ns: Address,
         control: UdpSocketControlData,
         server_windows: LruCache<u64, PacketWindowFilter>,

diff --git a/crates/shadowsocks-service/src/local/net/udp/association.rs b/crates/shadowsocks-service/src/local/net/udp/association.rs
@@ -217,7 +217,7 @@ where
     peer_addr: SocketAddr,
     bypassed_ipv4_socket: Option<ShadowUdpSocket>,
     bypassed_ipv6_socket: Option<ShadowUdpSocket>,
-    proxied_socket: Option<MonProxySocket>,
+    proxied_socket: Option<MonProxySocket<ShadowUdpSocket>>,
     keepalive_tx: mpsc::Sender<SocketAddr>,
     keepalive_flag: bool,
     balancer: PingBalancer,
@@ -409,7 +409,7 @@ where
 
         #[inline]
         async fn receive_from_proxied_opt(
-            socket: &Option<MonProxySocket>,
+            socket: &Option<MonProxySocket<ShadowUdpSocket>>,
             buf: &mut Vec<u8>,
         ) -> io::Result<(usize, Address, Option<UdpSocketControlData>)> {
             match *socket {

diff --git a/crates/shadowsocks-service/src/net/mon_socket.rs b/crates/shadowsocks-service/src/net/mon_socket.rs
@@ -3,24 +3,44 @@
 use std::{io, net::SocketAddr, sync::Arc};
 
 use shadowsocks::{
-    relay::{socks5::Address, udprelay::options::UdpSocketControlData},
+    relay::{
+        socks5::Address,
+        udprelay::{options::UdpSocketControlData, DatagramTransport},
+    },
     ProxySocket,
 };
 
 use super::flow::FlowStat;
 
 /// Monitored `ProxySocket`
-pub struct MonProxySocket {
-    socket: ProxySocket,
+pub struct MonProxySocket<S> {
+    socket: ProxySocket<S>,
     flow_stat: Arc<FlowStat>,
 }
 
-impl MonProxySocket {
+impl<S> MonProxySocket<S> {
     /// Create a new socket with flow monitor
-    pub fn from_socket(socket: ProxySocket, flow_stat: Arc<FlowStat>) -> MonProxySocket {
+    pub fn from_socket(socket: ProxySocket<S>, flow_stat: Arc<FlowStat>) -> MonProxySocket<S> {
         MonProxySocket { socket, flow_stat }
     }
 
+    /// Get the underlying `ProxySocket<S>` immutable reference
+    #[inline]
+    pub fn get_ref(&self) -> &ProxySocket<S> {
+        &self.socket
+    }
+
+    /// Get the flow statistic data
+    #[inline]
+    pub fn flow_stat(&self) -> &FlowStat {
+        &self.flow_stat
+    }
+}
+
+impl<S> MonProxySocket<S>
+where
+    S: DatagramTransport,
+{
     /// Send a UDP packet to addr through proxy
     #[inline]
     pub async fn send(&self, addr: &Address, payload: &[u8]) -> io::Result<()> {
@@ -125,14 +145,4 @@ impl MonProxySocket {
 
         Ok((n, peer_addr, addr, control))
     }
-
-    #[inline]
-    pub fn get_ref(&self) -> &ProxySocket {
-        &self.socket
-    }
-
-    #[inline]
-    pub fn flow_stat(&self) -> &FlowStat {
-        &self.flow_stat
-    }
 }
diff --git a/crates/shadowsocks-service/src/server/udprelay.rs b/crates/shadowsocks-service/src/server/udprelay.rs
@@ -17,7 +17,10 @@ use shadowsocks::{
     config::ServerUser,
     crypto::CipherCategory,
     lookup_then,
-    net::{get_ip_stack_capabilities, AcceptOpts, AddrFamily, UdpSocket as OutboundUdpSocket},
+    net::{
+        get_ip_stack_capabilities, AcceptOpts, AddrFamily, UdpSocket as OutboundUdpSocket,
+        UdpSocket as InboundUdpSocket,
+    },
     relay::{
         socks5::Address,
         udprelay::{options::UdpSocketControlData, ProxySocket, MAXIMUM_UDP_PAYLOAD_SIZE},
@@ -87,7 +90,7 @@ pub struct UdpServer {
     keepalive_tx: mpsc::Sender<NatKey>,
     keepalive_rx: mpsc::Receiver<NatKey>,
     time_to_live: Duration,
-    listener: Arc<MonProxySocket>,
+    listener: Arc<MonProxySocket<InboundUdpSocket>>,
     svr_cfg: ServerConfig,
 }
 
@@ -276,7 +279,7 @@ impl UdpServer {
 
     async fn recv_one_packet(
         context: &ServiceContext,
-        l: &MonProxySocket,
+        l: &MonProxySocket<InboundUdpSocket>,
         buffer: &mut [u8],
     ) -> Option<(usize, SocketAddr, Address, Option<UdpSocketControlData>)> {
         let (n, peer_addr, target_addr, control) = match l.recv_from_with_ctrl(buffer).await {
@@ -316,7 +319,7 @@ impl UdpServer {
 
     async fn send_packet(
         &mut self,
-        listener: &Arc<MonProxySocket>,
+        listener: &Arc<MonProxySocket<InboundUdpSocket>>,
         peer_addr: SocketAddr,
         target_addr: Address,
         control: Option<UdpSocketControlData>,
@@ -394,7 +397,7 @@ impl Drop for UdpAssociation {
 impl UdpAssociation {
     fn new_association(
         context: Arc<ServiceContext>,
-        inbound: Arc<MonProxySocket>,
+        inbound: Arc<MonProxySocket<InboundUdpSocket>>,
         peer_addr: SocketAddr,
         keepalive_tx: mpsc::Sender<NatKey>,
     ) -> UdpAssociation {
@@ -405,7 +408,7 @@ impl UdpAssociation {
     #[cfg(feature = "aead-cipher-2022")]
     fn new_session(
         context: Arc<ServiceContext>,
-        inbound: Arc<MonProxySocket>,
+        inbound: Arc<MonProxySocket<InboundUdpSocket>>,
         peer_addr: SocketAddr,
         keepalive_tx: mpsc::Sender<NatKey>,
         client_session_id: u64,
@@ -447,7 +450,7 @@ struct UdpAssociationContext {
     outbound_ipv6_socket: Option<OutboundUdpSocket>,
     keepalive_tx: mpsc::Sender<NatKey>,
     keepalive_flag: bool,
-    inbound: Arc<MonProxySocket>,
+    inbound: Arc<MonProxySocket<InboundUdpSocket>>,
     // AEAD 2022
     client_session: Option<ClientSessionContext>,
     server_session_id: u64,
@@ -472,7 +475,7 @@ fn generate_server_session_id() -> u64 {
 impl UdpAssociationContext {
     fn create(
         context: Arc<ServiceContext>,
-        inbound: Arc<MonProxySocket>,
+        inbound: Arc<MonProxySocket<InboundUdpSocket>>,
         peer_addr: SocketAddr,
         keepalive_tx: mpsc::Sender<NatKey>,
         client_session_id: Option<u64>,