killing UDP association relay tasks properly when server is closed

- ref #557
- ref #271

Author: zonyitoo

crates/shadowsocks-service/src/local/net/udp/association.rs

@@ -54,6 +54,8 @@ where
     context: Arc<ServiceContext>,
     assoc_map: SharedAssociationMap<W>,
     cleanup_abortable: AbortHandle,
+    keepalive_abortable: AbortHandle,
+    keepalive_tx: mpsc::Sender<SocketAddr>,
     balancer: PingBalancer,
 }
 
@@ -63,6 +65,7 @@ where
 {
     fn drop(&mut self) {
         self.cleanup_abortable.abort();
+        self.keepalive_abortable.abort();
     }
 }
 
@@ -98,11 +101,26 @@ where
             cleanup_abortable
         };
 
+        let (keepalive_tx, mut keepalive_rx) = mpsc::channel(256);
+
+        let keepalive_abortable = {
+            let assoc_map = assoc_map.clone();
+            let (keepalive_task, keepalive_abortable) = future::abortable(async move {
+                while let Some(peer_addr) = keepalive_rx.recv().await {
+                    assoc_map.lock().await.get(&peer_addr);
+                }
+            });
+            tokio::spawn(keepalive_task);
+            keepalive_abortable
+        };
+
         UdpAssociationManager {
             respond_writer,
             context,
             assoc_map,
             cleanup_abortable,
+            keepalive_abortable,
+            keepalive_tx,
             balancer,
         }
     }
@@ -120,7 +138,7 @@ where
         let assoc = UdpAssociation::new(
             self.context.clone(),
             peer_addr,
-            self.assoc_map.clone(),
+            self.keepalive_tx.clone(),
             self.balancer.clone(),
             self.respond_writer.clone(),
         );
@@ -160,11 +178,11 @@ where
     fn new(
         context: Arc<ServiceContext>,
         peer_addr: SocketAddr,
-        assoc_map: SharedAssociationMap<W>,
+        keepalive_tx: mpsc::Sender<SocketAddr>,
         balancer: PingBalancer,
         respond_writer: W,
     ) -> UdpAssociation<W> {
-        let (assoc, sender) = UdpAssociationContext::new(context, peer_addr, assoc_map, balancer, respond_writer);
+        let (assoc, sender) = UdpAssociationContext::new(context, peer_addr, keepalive_tx, balancer, respond_writer);
         UdpAssociation { assoc, sender }
     }
 
@@ -219,9 +237,7 @@ enum UdpAssociationProxyState {
 
 impl Drop for UdpAssociationProxyState {
     fn drop(&mut self) {
-        if let UdpAssociationProxyState::Connected { ref abortable, .. } = *self {
-            abortable.abort();
-        }
+        self.abort_inner();
     }
 }
 
@@ -231,16 +247,25 @@ impl UdpAssociationProxyState {
     }
 
     fn reset(&mut self) {
+        self.abort_inner();
         *self = UdpAssociationProxyState::Empty;
     }
 
     fn set_connected(&mut self, socket: Arc<MonProxySocket>, abortable: AbortHandle) {
+        self.abort_inner();
         *self = UdpAssociationProxyState::Connected { socket, abortable };
     }
 
     fn abort(&mut self) {
+        self.abort_inner();
         *self = UdpAssociationProxyState::Aborted;
     }
+
+    fn abort_inner(&mut self) {
+        if let UdpAssociationProxyState::Connected { ref abortable, .. } = *self {
+            abortable.abort();
+        }
+    }
 }
 
 struct UdpAssociationContext<W>
@@ -252,7 +277,7 @@ where
     bypassed_ipv4_socket: SpinMutex<UdpAssociationBypassState>,
     bypassed_ipv6_socket: SpinMutex<UdpAssociationBypassState>,
     proxied_socket: SpinMutex<UdpAssociationProxyState>,
-    assoc_map: SharedAssociationMap<W>,
+    keepalive_tx: mpsc::Sender<SocketAddr>,
     balancer: PingBalancer,
     respond_writer: W,
 }
@@ -273,7 +298,7 @@ where
     fn new(
         context: Arc<ServiceContext>,
         peer_addr: SocketAddr,
-        assoc_map: SharedAssociationMap<W>,
+        keepalive_tx: mpsc::Sender<SocketAddr>,
         balancer: PingBalancer,
         respond_writer: W,
     ) -> (Arc<UdpAssociationContext<W>>, mpsc::Sender<(Address, Bytes)>) {
@@ -288,7 +313,7 @@ where
             bypassed_ipv4_socket: SpinMutex::new(UdpAssociationBypassState::empty()),
             bypassed_ipv6_socket: SpinMutex::new(UdpAssociationBypassState::empty()),
             proxied_socket: SpinMutex::new(UdpAssociationProxyState::empty()),
-            assoc_map,
+            keepalive_tx,
             balancer,
             respond_writer,
         });
@@ -557,7 +582,10 @@ where
                         n
                     );
                     // Keep association alive in map
-                    let _ = self.assoc_map.lock().await.get(&self.peer_addr);
+                    let _ = self
+                        .keepalive_tx
+                        .send_timeout(self.peer_addr, Duration::from_secs(1))
+                        .await;
                     (n, addr)
                 }
                 Err(err) => {
@@ -608,7 +636,10 @@ where
                         n
                     );
                     // Keep association alive in map
-                    let _ = self.assoc_map.lock().await.get(&self.peer_addr);
+                    let _ = self
+                        .keepalive_tx
+                        .send_timeout(self.peer_addr, Duration::from_secs(1))
+                        .await;
                     (n, addr)
                 }
                 Err(err) => {

crates/shadowsocks-service/src/local/tunnel/udprelay.rs

@@ -35,11 +35,14 @@ pub struct UdpTunnel {
     context: Arc<ServiceContext>,
     assoc_map: SharedAssociationMap,
     cleanup_abortable: AbortHandle,
+    keepalive_abortable: AbortHandle,
+    keepalive_tx: mpsc::Sender<SocketAddr>,
 }
 
 impl Drop for UdpTunnel {
     fn drop(&mut self) {
         self.cleanup_abortable.abort();
+        self.keepalive_abortable.abort();
     }
 }
 
@@ -65,10 +68,25 @@ impl UdpTunnel {
             cleanup_abortable
         };
 
+        let (keepalive_tx, mut keepalive_rx) = mpsc::channel(256);
+
+        let keepalive_abortable = {
+            let assoc_map = assoc_map.clone();
+            let (keepalive_task, keepalive_abortable) = future::abortable(async move {
+                while let Some(peer_addr) = keepalive_rx.recv().await {
+                    assoc_map.lock().await.get(&peer_addr);
+                }
+            });
+            tokio::spawn(keepalive_task);
+            keepalive_abortable
+        };
+
         UdpTunnel {
             context,
             assoc_map,
             cleanup_abortable,
+            keepalive_abortable,
+            keepalive_tx,
         }
     }
 
@@ -139,7 +157,7 @@ impl UdpTunnel {
             listener.clone(),
             peer_addr,
             forward_addr.clone(),
-            self.assoc_map.clone(),
+            self.keepalive_tx.clone(),
             balancer.clone(),
         );
 
@@ -169,11 +187,11 @@ impl UdpAssociation {
         inbound: Arc<UdpSocket>,
         peer_addr: SocketAddr,
         forward_addr: Address,
-        assoc_map: SharedAssociationMap,
+        keepalive_tx: mpsc::Sender<SocketAddr>,
         balancer: PingBalancer,
     ) -> UdpAssociation {
         let (assoc, sender) =
-            UdpAssociationContext::new(context, inbound, peer_addr, forward_addr, assoc_map, balancer);
+            UdpAssociationContext::new(context, inbound, peer_addr, forward_addr, keepalive_tx, balancer);
         UdpAssociation { sender, assoc }
     }
 
@@ -197,9 +215,7 @@ enum UdpAssociationState {
 
 impl Drop for UdpAssociationState {
     fn drop(&mut self) {
-        if let UdpAssociationState::Connected { ref abortable, .. } = *self {
-            abortable.abort();
-        }
+        self.abort_inner();
     }
 }
 
@@ -209,16 +225,25 @@ impl UdpAssociationState {
     }
 
     fn reset(&mut self) {
+        self.abort_inner();
         *self = UdpAssociationState::Empty;
     }
 
     fn set_connected(&mut self, socket: Arc<MonProxySocket>, abortable: AbortHandle) {
+        self.abort_inner();
         *self = UdpAssociationState::Connected { socket, abortable };
     }
 
     fn abort(&mut self) {
+        self.abort_inner();
         *self = UdpAssociationState::Aborted;
     }
+
+    fn abort_inner(&mut self) {
+        if let UdpAssociationState::Connected { ref abortable, .. } = *self {
+            abortable.abort();
+        }
+    }
 }
 
 struct UdpAssociationContext {
@@ -227,7 +252,7 @@ struct UdpAssociationContext {
     peer_addr: SocketAddr,
     forward_addr: Address,
     proxied_socket: SpinMutex<UdpAssociationState>,
-    assoc_map: SharedAssociationMap,
+    keepalive_tx: mpsc::Sender<SocketAddr>,
     balancer: PingBalancer,
 }
 
@@ -243,7 +268,7 @@ impl UdpAssociationContext {
         inbound: Arc<UdpSocket>,
         peer_addr: SocketAddr,
         forward_addr: Address,
-        assoc_map: SharedAssociationMap,
+        keepalive_tx: mpsc::Sender<SocketAddr>,
         balancer: PingBalancer,
     ) -> (Arc<UdpAssociationContext>, mpsc::Sender<Bytes>) {
         // Pending packets 1024 should be good enough for a server.
@@ -257,7 +282,7 @@ impl UdpAssociationContext {
             peer_addr,
             forward_addr,
             proxied_socket: SpinMutex::new(UdpAssociationState::empty()),
-            assoc_map,
+            keepalive_tx,
             balancer,
         });
 
@@ -374,7 +399,10 @@ impl UdpAssociationContext {
                 Ok((n, addr)) => {
                     trace!("udp relay {} <- {} received {} bytes", self.peer_addr, addr, n);
                     // Keep association alive in map
-                    let _ = self.assoc_map.lock().await.get(&self.peer_addr);
+                    let _ = self
+                        .keepalive_tx
+                        .send_timeout(self.peer_addr, Duration::from_secs(1))
+                        .await;
                     n
                 }
                 Err(err) => {