Fix SOCKS5 port forwarding for modified peer data packets (#2391)

Fixes SOCKS5/port-forward handling for peer data packets 
whose source endpoint was rewritten by the KCP or QUIC proxy path.

Keep SOCKS5 entry accounting consistent by centralizing insert/remove 
operations, decrementing only for actual removals, avoiding underflow, 
and resetting counts when entries are retained or cleared after IPv4 changes.

.github/workflows/ohos.yml

@@ -35,6 +35,7 @@ jobs:
         with:
           gui: false
           pnpm: false
+          token: ${{ secrets.GITHUB_TOKEN }}
 
       - uses: actions-rust-lang/setup-rust-toolchain@v1
         with:
@@ -243,4 +244,3 @@ jobs:
              ohpm publish easytier-release.har
           fi
           curl --header "Content-Type: application/json" --request POST --data "{}" ${{ secrets.CODEARTS_WEBHOOKS }}
-

easytier/src/gateway/socks5.rs

@@ -32,7 +32,7 @@ use crate::{
     tunnel::packet_def::{PacketType, ZCPacket},
 };
 use anyhow::Context;
-use dashmap::DashMap;
+use dashmap::{DashMap, mapref::entry::Entry};
 use pnet::packet::{
     Packet, ip::IpNextHeaderProtocols, ipv4::Ipv4Packet, tcp::TcpPacket, udp::UdpPacket,
 };
@@ -164,6 +164,87 @@ struct Socks5Entry {
 
 type Socks5EntrySet = Arc<DashMap<Socks5Entry, Socks5EntryData>>;
 
+fn increment_entry_count(entry_count: &AtomicUsize) -> (usize, usize) {
+    let old_entry_count = entry_count
+        .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| {
+            count.checked_add(1)
+        })
+        .unwrap_or_else(|count| count);
+    (old_entry_count, old_entry_count.saturating_add(1))
+}
+
+fn decrement_entry_count(entry_count: &AtomicUsize) -> (usize, usize) {
+    decrement_entry_count_by(entry_count, 1)
+}
+
+fn decrement_entry_count_by(entry_count: &AtomicUsize, delta: usize) -> (usize, usize) {
+    if delta == 0 {
+        let current = entry_count.load(Ordering::Relaxed);
+        return (current, current);
+    }
+
+    let old_entry_count = entry_count
+        .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| {
+            Some(count.saturating_sub(delta))
+        })
+        .unwrap_or_else(|count| count);
+    (old_entry_count, old_entry_count.saturating_sub(delta))
+}
+
+fn insert_entry_and_increment_count(
+    entries: &Socks5EntrySet,
+    entry_count: &AtomicUsize,
+    entry: Socks5Entry,
+    data: Socks5EntryData,
+) -> (bool, usize, usize) {
+    match entries.entry(entry) {
+        Entry::Occupied(mut occupied) => {
+            occupied.insert(data);
+            let current = entry_count.load(Ordering::Relaxed);
+            (true, current, current)
+        }
+        Entry::Vacant(vacant) => {
+            // Keep the count update inside the VacantEntry shard lock so bulk clear
+            // cannot observe the inserted entry before its count is reserved.
+            let (old_entry_count, new_entry_count) = increment_entry_count(entry_count);
+            vacant.insert(data);
+            (false, old_entry_count, new_entry_count)
+        }
+    }
+}
+
+fn try_insert_entry_and_increment_count(
+    entries: &Socks5EntrySet,
+    entry_count: &AtomicUsize,
+    entry: Socks5Entry,
+    data: Socks5EntryData,
+) -> bool {
+    match entries.entry(entry) {
+        Entry::Occupied(_) => false,
+        Entry::Vacant(vacant) => {
+            // See insert_entry_and_increment_count for why the count is reserved first.
+            increment_entry_count(entry_count);
+            vacant.insert(data);
+            true
+        }
+    }
+}
+
+fn remove_entry_and_decrement_count(
+    entries: &Socks5EntrySet,
+    entry_count: &AtomicUsize,
+    entry: &Socks5Entry,
+) -> (bool, usize, usize) {
+    let removed = entries.remove(entry).is_some();
+    let (old_entry_count, new_entry_count) = if removed {
+        decrement_entry_count(entry_count)
+    } else {
+        let current = entry_count.load(Ordering::Relaxed);
+        (current, current)
+    };
+    (removed, old_entry_count, new_entry_count)
+}
+
 struct SmolTcpConnector {
     net: Arc<Net>,
     entries: Socks5EntrySet,
@@ -190,9 +271,20 @@ impl AsyncTcpConnector for SmolTcpConnector {
             entry_type: TCP_ENTRY,
         };
         *self.current_entry.lock().unwrap() = Some(entry.clone());
-        self.entries
+        let (replaced, old_entry_count, new_entry_count) = insert_entry_and_increment_count(
-            .insert(entry, Socks5EntryData::Tcp(tmp_listener));
+            &self.entries,
-        self.entry_count.fetch_add(1, Ordering::Relaxed);
+            &self.entry_count,
+            entry.clone(),
+            Socks5EntryData::Tcp(tmp_listener),
+        );
+        tracing::trace!(
+            ?entry,
+            replaced,
+            old_entry_count,
+            new_entry_count,
+            entries_len = self.entries.len(),
+            "socks5 inserted smoltcp tcp connector entry"
+        );
 
         if addr.ip() == local_addr {
             let modified_addr =
@@ -220,8 +312,16 @@ impl Drop for SmolTcpConnector {
     fn drop(&mut self) {
         if let Some(entry) = self.current_entry.lock().unwrap().take() {
             tracing::debug!("drop smoltcp connector entry {:?}", entry);
-            self.entries.remove(&entry);
+            let (removed, old_entry_count, new_entry_count) =
-            self.entry_count.fetch_sub(1, Ordering::Relaxed);
+                remove_entry_and_decrement_count(&self.entries, &self.entry_count, &entry);
+            tracing::trace!(
+                ?entry,
+                removed,
+                old_entry_count,
+                new_entry_count,
+                entries_len = self.entries.len(),
+                "socks5 removed smoltcp tcp connector entry"
+            );
         }
     }
 }
@@ -294,11 +394,26 @@ impl AsyncTcpConnector for Socks5AutoConnector {
             addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), addr.port());
         }
 
-        if self.smoltcp_net.is_none()
+        let has_smoltcp_net = self.smoltcp_net.is_some();
-            || peer_mgr_arc.get_msg_dst_peer(&addr.ip()).await.0.is_empty()
+        let dst_peers = if has_smoltcp_net && !addr.ip().is_loopback() {
+            Some(peer_mgr_arc.get_msg_dst_peer(&addr.ip()).await.0)
+        } else {
+            None
+        };
+
+        if !has_smoltcp_net
+            || dst_peers.as_ref().is_some_and(Vec::is_empty)
             || addr.ip().is_loopback()
         {
             // cannot find dst in virtual network, so try connect to dst directly
+            tracing::trace!(
+                ?addr,
+                src_addr = ?self.src_addr,
+                has_smoltcp_net,
+                dst_peer_count = dst_peers.as_ref().map(Vec::len),
+                is_loopback = addr.ip().is_loopback(),
+                "socks5 auto connector falling back to kernel tcp connect"
+            );
             return Ok(SocksTcpStream::Tcp(
                 tcp_connect_with_timeout(addr, timeout_s).await?,
             ));
@@ -310,25 +425,51 @@ impl AsyncTcpConnector for Socks5AutoConnector {
         #[cfg(feature = "kcp")]
         let connector: Box<dyn AsyncTcpConnector<S = SocksTcpStream> + Send> =
             match (&self.kcp_endpoint, dst_allow_kcp) {
-                (Some(kcp_endpoint), true) => Box::new(Socks5KcpConnector {
+                (Some(kcp_endpoint), true) => {
-                    kcp_endpoint: kcp_endpoint.clone(),
+                    tracing::trace!(
-                    peer_mgr: self.peer_mgr.clone(),
+                        ?addr,
-                    src_addr: self.src_addr,
+                        src_addr = ?self.src_addr,
-                }),
+                        dst_peer_count = dst_peers.as_ref().map(Vec::len),
-                (_, _) => Box::new(SmolTcpConnector {
+                        "socks5 auto connector selected kcp"
-                    net: self.smoltcp_net.clone().unwrap(),
+                    );
-                    entries: self.entries.clone(),
+                    Box::new(Socks5KcpConnector {
-                    entry_count: self.entry_count.clone(),
+                        kcp_endpoint: kcp_endpoint.clone(),
-                    current_entry: std::sync::Mutex::new(None),
+                        peer_mgr: self.peer_mgr.clone(),
-                }),
+                        src_addr: self.src_addr,
+                    })
+                }
+                (_, _) => {
+                    tracing::trace!(
+                        ?addr,
+                        src_addr = ?self.src_addr,
+                        dst_peer_count = dst_peers.as_ref().map(Vec::len),
+                        dst_allow_kcp,
+                        has_kcp_endpoint = self.kcp_endpoint.is_some(),
+                        "socks5 auto connector selected smoltcp"
+                    );
+                    Box::new(SmolTcpConnector {
+                        net: self.smoltcp_net.clone().unwrap(),
+                        entries: self.entries.clone(),
+                        entry_count: self.entry_count.clone(),
+                        current_entry: std::sync::Mutex::new(None),
+                    })
+                }
             };
         #[cfg(not(feature = "kcp"))]
-        let connector = Box::new(SmolTcpConnector {
+        let connector = {
-            net: self.smoltcp_net.clone().unwrap(),
+            tracing::trace!(
-            entries: self.entries.clone(),
+                ?addr,
-            entry_count: self.entry_count.clone(),
+                src_addr = ?self.src_addr,
-            current_entry: std::sync::Mutex::new(None),
+                dst_peer_count = dst_peers.as_ref().map(Vec::len),
-        });
+                "socks5 auto connector selected smoltcp"
+            );
+            Box::new(SmolTcpConnector {
+                net: self.smoltcp_net.clone().unwrap(),
+                entries: self.entries.clone(),
+                entry_count: self.entry_count.clone(),
+                current_entry: std::sync::Mutex::new(None),
+            })
+        };
 
         let ret = connector.tcp_connect(addr, timeout_s).await;
         self.inner_connector.lock().replace(Box::new(connector));
@@ -503,25 +644,89 @@ pub struct Socks5Server {
 #[async_trait::async_trait]
 impl PeerPacketFilter for Socks5Server {
     async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> Option<ZCPacket> {
-        if self.entry_count.load(Ordering::Relaxed) == 0
+        let entry_count = self.entry_count.load(Ordering::Relaxed);
-            && !self.socks5_enabled.load(Ordering::Relaxed)
+        let socks5_enabled = self.socks5_enabled.load(Ordering::Relaxed);
-        {
+        if entry_count == 0 && !socks5_enabled && self.entries.is_empty() {
+            if tracing::enabled!(tracing::Level::TRACE)
+                && let Some(hdr) = packet.peer_manager_header()
+                && matches!(
+                    hdr.packet_type,
+                    x if x == PacketType::Data as u8
+                        || x == PacketType::DataWithKcpSrcModified as u8
+                        || x == PacketType::DataWithQuicSrcModified as u8
+                )
+            {
+                if let Some(ipv4) = Ipv4Packet::new(packet.payload()) {
+                    let (tcp_src_port, tcp_dst_port, tcp_flags) =
+                        if ipv4.get_next_level_protocol() == IpNextHeaderProtocols::Tcp {
+                            TcpPacket::new(ipv4.payload())
+                                .map(|tcp| {
+                                    (
+                                        Some(tcp.get_source()),
+                                        Some(tcp.get_destination()),
+                                        Some(tcp.get_flags()),
+                                    )
+                                })
+                                .unwrap_or((None, None, None))
+                        } else {
+                            (None, None, None)
+                        };
+                    tracing::trace!(
+                        packet_type = hdr.packet_type,
+                        from_peer_id = hdr.from_peer_id.get(),
+                        to_peer_id = hdr.to_peer_id.get(),
+                        ipv4_src = %ipv4.get_source(),
+                        ipv4_dst = %ipv4.get_destination(),
+                        next_protocol = ?ipv4.get_next_level_protocol(),
+                        ?tcp_src_port,
+                        ?tcp_dst_port,
+                        ?tcp_flags,
+                        entry_count,
+                        socks5_enabled,
+                        "socks5 fast gate passed packet from peer"
+                    );
+                } else {
+                    tracing::trace!(
+                        packet_type = hdr.packet_type,
+                        from_peer_id = hdr.from_peer_id.get(),
+                        to_peer_id = hdr.to_peer_id.get(),
+                        entry_count,
+                        socks5_enabled,
+                        "socks5 fast gate passed non-ipv4 packet from peer"
+                    );
+                }
+            }
+            return Some(packet);
+        }
+        let hdr = packet.peer_manager_header().unwrap();
+        let is_modified_src_packet = matches!(
+            hdr.packet_type,
+            x if x == PacketType::DataWithKcpSrcModified as u8
+                || x == PacketType::DataWithQuicSrcModified as u8
+        );
+        if hdr.packet_type != PacketType::Data as u8 && !is_modified_src_packet {
+            return Some(packet);
+        }
+        if is_modified_src_packet && hdr.from_peer_id != hdr.to_peer_id {
+            tracing::trace!(
+                packet_type = hdr.packet_type,
+                from_peer_id = hdr.from_peer_id.get(),
+                to_peer_id = hdr.to_peer_id.get(),
+                "socks5 passed non-loopback modified-source packet from peer"
+            );
             return Some(packet);
         }
 
-        let hdr = packet.peer_manager_header().unwrap();
-        if hdr.packet_type != PacketType::Data as u8 {
-            return Some(packet);
-        };
-
         let payload_bytes = packet.payload();
 
-        let ipv4 = Ipv4Packet::new(payload_bytes).unwrap();
+        let Some(ipv4) = Ipv4Packet::new(payload_bytes) else {
+            return Some(packet);
+        };
         if ipv4.get_version() != 4 {
             return Some(packet);
         }
 
-        let entry_key = match ipv4.get_next_level_protocol() {
+        let (entry_key, tcp_flags) = match ipv4.get_next_level_protocol() {
             IpNextHeaderProtocols::Tcp => {
                 let Some(tcp_packet) = TcpPacket::new(ipv4.payload()) else {
                     return Some(packet);
@@ -546,11 +751,11 @@ impl PeerPacketFilter for Socks5Server {
                         entry_type: TCP_LISTEN_ENTRY,
                     }
                 };
-                entry
+                (entry, Some(tcp_packet.get_flags()))
             }
 
             IpNextHeaderProtocols::Udp => {
-                if IpReassembler::is_packet_fragmented(&ipv4) && !self.entries.is_empty() {
+                if IpReassembler::is_packet_fragmented(&ipv4) {
                     let ipv4_src: IpAddr = ipv4.get_source().into();
                     // only send to smoltcp if the ipv4 src is in the entries
                     let is_in_entries = self.entries.iter().any(|x| x.key().dst.ip() == ipv4_src);
@@ -562,7 +767,19 @@ impl PeerPacketFilter for Socks5Server {
                     if is_in_entries {
                         // if the packet is fragmented, no matther what the payload is, need send it to both smoltcp and kernel tun. because
                         // we cannot determine the udp port of the packet.
-                        let _ = self.packet_sender.try_send(packet.clone()).ok();
+                        match self.packet_sender.try_send(packet.clone()) {
+                            Ok(()) => tracing::trace!(
+                                ?ipv4_src,
+                                entry_count = self.entry_count.load(Ordering::Relaxed),
+                                "socks5 delivered fragmented packet from peer to smoltcp"
+                            ),
+                            Err(err) => tracing::trace!(
+                                ?ipv4_src,
+                                ?err,
+                                entry_count = self.entry_count.load(Ordering::Relaxed),
+                                "socks5 failed to deliver fragmented packet from peer to smoltcp"
+                            ),
+                        }
                     }
                     return Some(packet);
                 }
@@ -570,14 +787,17 @@ impl PeerPacketFilter for Socks5Server {
                 let Some(udp_packet) = UdpPacket::new(ipv4.payload()) else {
                     return Some(packet);
                 };
-                Socks5Entry {
+                (
-                    dst: SocketAddr::new(ipv4.get_source().into(), udp_packet.get_source()),
+                    Socks5Entry {
-                    src: SocketAddr::new(
+                        dst: SocketAddr::new(ipv4.get_source().into(), udp_packet.get_source()),
-                        ipv4.get_destination().into(),
+                        src: SocketAddr::new(
-                        udp_packet.get_destination(),
+                            ipv4.get_destination().into(),
-                    ),
+                            udp_packet.get_destination(),
-                    entry_type: UDP_ENTRY,
+                        ),
-                }
+                        entry_type: UDP_ENTRY,
+                    },
+                    None,
+                )
             }
             _ => {
                 return Some(packet);
@@ -585,12 +805,41 @@ impl PeerPacketFilter for Socks5Server {
         };
 
         if !self.entries.contains_key(&entry_key) {
+            tracing::trace!(
+                ?entry_key,
+                ?tcp_flags,
+                ipv4_src = %ipv4.get_source(),
+                ipv4_dst = %ipv4.get_destination(),
+                entry_count = self.entry_count.load(Ordering::Relaxed),
+                socks5_enabled = self.socks5_enabled.load(Ordering::Relaxed),
+                "socks5 no entry for packet from peer"
+            );
             return Some(packet);
         }
 
-        tracing::trace!(?entry_key, ?ipv4, "socks5 found entry for packet from peer");
+        tracing::trace!(
+            ?entry_key,
+            ?tcp_flags,
+            ?ipv4,
+            entry_count = self.entry_count.load(Ordering::Relaxed),
+            "socks5 found entry for packet from peer"
+        );
 
-        let _ = self.packet_sender.try_send(packet).ok();
+        match self.packet_sender.try_send(packet) {
+            Ok(()) => tracing::trace!(
+                ?entry_key,
+                ?tcp_flags,
+                entry_count = self.entry_count.load(Ordering::Relaxed),
+                "socks5 delivered packet from peer to smoltcp"
+            ),
+            Err(err) => tracing::trace!(
+                ?entry_key,
+                ?tcp_flags,
+                ?err,
+                entry_count = self.entry_count.load(Ordering::Relaxed),
+                "socks5 failed to deliver packet from peer to smoltcp"
+            ),
+        }
 
         None
     }
@@ -655,11 +904,22 @@ impl Socks5Server {
                 #[cfg(not(feature = "ffi-dataplane"))]
                 let data_plane_active = false;
 
-                if cancel_tokens.is_empty()
+                let active_port_forwards = cancel_tokens.len();
-                    && !socks5_enabled.load(Ordering::Relaxed)
+                let is_socks5_enabled = socks5_enabled.load(Ordering::Relaxed);
-                    && !data_plane_active
+                if active_port_forwards == 0 && !is_socks5_enabled && !data_plane_active {
-                {
+                    let had_net = {
-                    let _ = net.lock().await.take();
+                        let mut net_guard = net.lock().await;
+                        net_guard.take().is_some()
+                    };
+                    tracing::trace!(
+                        had_net,
+                        active_port_forwards,
+                        is_socks5_enabled,
+                        data_plane_active,
+                        entry_count = entry_count.load(Ordering::Relaxed),
+                        entries_len = entries.len(),
+                        "socks5 net update waiting for consumers"
+                    );
                     #[cfg(feature = "ffi-dataplane")]
                     let _ = data_plane_net_ready.send_replace(false);
                     port_forward_list_change_notifier.notified().await;
@@ -670,13 +930,34 @@ impl Socks5Server {
 
                 let cur_ipv4 = global_ctx.get_ipv4();
                 if prev_ipv4 != cur_ipv4 {
+                    let old_ipv4 = prev_ipv4;
                     prev_ipv4 = cur_ipv4;
 
+                    tracing::trace!(
+                        ?old_ipv4,
+                        ?cur_ipv4,
+                        old_entry_count = entry_count.load(Ordering::Relaxed),
+                        old_entries_len = entries.len(),
+                        udp_client_count = udp_client_map.len(),
+                        "socks5 net update resetting entries for ipv4 change"
+                    );
+                    let mut removed_entries = 0;
                     entries.retain(|_, _| {
-                        entry_count.fetch_sub(1, Ordering::Relaxed);
+                        removed_entries += 1;
                         false
                     });
+                    let (_, new_entry_count) =
+                        decrement_entry_count_by(&entry_count, removed_entries);
                     udp_client_map.clear();
+                    tracing::trace!(
+                        ?old_ipv4,
+                        ?cur_ipv4,
+                        removed_entries,
+                        new_entry_count,
+                        new_entries_len = entries.len(),
+                        udp_client_count = udp_client_map.len(),
+                        "socks5 net update reset entries complete"
+                    );
 
                     if let Some(cur_ipv4) = cur_ipv4 {
                         net.lock().await.replace(Socks5ServerNet::new(
@@ -686,12 +967,23 @@ impl Socks5Server {
                             packet_recv.clone(),
                             entries.clone(),
                         ));
+                        tracing::trace!(
+                            ?cur_ipv4,
+                            entry_count = entry_count.load(Ordering::Relaxed),
+                            entries_len = entries.len(),
+                            "socks5 net update installed smoltcp net"
+                        );
                         // Wake any data-plane callers waiting in
                         // `wait_data_plane_net` for the smoltcp net to appear.
                         #[cfg(feature = "ffi-dataplane")]
                         let _ = data_plane_net_ready.send_replace(true);
                     } else {
                         let _ = net.lock().await.take();
+                        tracing::trace!(
+                            entry_count = entry_count.load(Ordering::Relaxed),
+                            entries_len = entries.len(),
+                            "socks5 net update removed smoltcp net"
+                        );
                         #[cfg(feature = "ffi-dataplane")]
                         let _ = data_plane_net_ready.send_replace(false);
                     }
@@ -774,6 +1066,13 @@ impl Socks5Server {
         peer_manager
             .add_packet_process_pipeline(Box::new(self.clone()))
             .await;
+        tracing::trace!(
+            cfg_count = cfgs.len(),
+            cancel_token_count = self.cancel_tokens.len(),
+            entry_count = self.entry_count.load(Ordering::Relaxed),
+            entries_len = self.entries.len(),
+            "socks5 peer packet pipeline registered"
+        );
 
         self.run_net_update_task().await;
 
@@ -806,6 +1105,7 @@ impl Socks5Server {
         connector: Box<dyn AsyncTcpConnector<S = SocksTcpStream> + Send>,
         dst_addr: SocketAddr,
     ) {
+        tracing::trace!(?dst_addr, "port forward: connecting to destination");
         let outgoing_socket = match connector.tcp_connect(dst_addr, 10).await {
             Ok(socket) => socket,
             Err(e) => {
@@ -813,6 +1113,7 @@ impl Socks5Server {
                 return;
             }
         };
+        tracing::trace!(?dst_addr, "port forward: connected to destination");
 
         let mut outgoing_socket = outgoing_socket;
         match tokio::io::copy_bidirectional(&mut incoming_socket, &mut outgoing_socket).await {
@@ -899,12 +1200,30 @@ impl Socks5Server {
                     dst_addr
                 );
 
+                let (smoltcp_net, net_ipv4) = {
+                    let net_guard = net.lock().await;
+                    (
+                        net_guard.as_ref().map(|net| net.smoltcp_net.clone()),
+                        net_guard.as_ref().map(|net| net.ipv4_addr),
+                    )
+                };
+                tracing::trace!(
+                    ?bind_addr,
+                    ?dst_addr,
+                    client_addr = ?addr,
+                    has_smoltcp_net = smoltcp_net.is_some(),
+                    ?net_ipv4,
+                    entry_count = entry_count.load(Ordering::Relaxed),
+                    entries_len = entries.len(),
+                    "port forward: preparing connector"
+                );
+
                 let connector = Socks5AutoConnector {
                     #[cfg(feature = "kcp")]
                     kcp_endpoint: kcp_endpoint.clone(),
                     peer_mgr: peer_mgr.clone(),
                     entries: entries.clone(),
-                    smoltcp_net: net.lock().await.as_ref().map(|net| net.smoltcp_net.clone()),
+                    smoltcp_net,
                     src_addr: addr,
                     entry_count: entry_count.clone(),
                     inner_connector: parking_lot::Mutex::new(None),
@@ -1040,11 +1359,12 @@ impl Socks5Server {
                             )
                         };
                         let socks_udp = Arc::new(sokcs_udp);
-                        entries.insert(
+                        insert_entry_and_increment_count(
+                            &entries,
+                            &entry_count,
                             client_info.entry_key.clone(),
                             Socks5EntryData::Udp((socks_udp.clone(), udp_client_key.clone())),
                         );
-                        entry_count.fetch_add(1, Ordering::Relaxed);
 
                         let socks = socket.clone();
                         let client_addr = addr;
@@ -1107,16 +1427,18 @@ impl Socks5Server {
                     now.duration_since(client_info.last_active.load()).as_secs() < 600
                 });
                 udp_forward_task.retain(|k, _| udp_client_map.contains_key(k));
+                let mut removed_entries = 0;
                 entries.retain(|_, data| match data {
                     Socks5EntryData::Udp((_, udp_client_key)) => {
                         let keep = udp_client_map.contains_key(udp_client_key);
                         if !keep {
-                            entry_count.fetch_sub(1, Ordering::Relaxed);
+                            removed_entries += 1;
                         }
                         keep
                     }
                     _ => true,
                 });
+                decrement_entry_count_by(&entry_count, removed_entries);
 
                 udp_client_map.shrink_to_fit();
                 udp_forward_task.shrink_to_fit();
@@ -1128,3 +1450,324 @@ impl Socks5Server {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+
+    use pnet::packet::{
+        MutablePacket,
+        ip::IpNextHeaderProtocols,
+        ipv4::{self, MutableIpv4Packet},
+        tcp::{self, MutableTcpPacket, TcpFlags},
+    };
+
+    use super::*;
+    use crate::peers::tests::create_mock_peer_manager;
+
+    fn build_tcp_packet(src: SocketAddr, dst: SocketAddr) -> Vec<u8> {
+        let mut buf = vec![0u8; 40];
+        let src_ip = match src.ip() {
+            IpAddr::V4(ip) => ip,
+            IpAddr::V6(_) => panic!("test only supports ipv4"),
+        };
+        let dst_ip = match dst.ip() {
+            IpAddr::V4(ip) => ip,
+            IpAddr::V6(_) => panic!("test only supports ipv4"),
+        };
+
+        {
+            let mut ip_packet = MutableIpv4Packet::new(&mut buf).unwrap();
+            ip_packet.set_version(4);
+            ip_packet.set_header_length(5);
+            ip_packet.set_total_length(40);
+            ip_packet.set_ttl(64);
+            ip_packet.set_next_level_protocol(IpNextHeaderProtocols::Tcp);
+            ip_packet.set_source(src_ip);
+            ip_packet.set_destination(dst_ip);
+
+            let mut tcp_packet = MutableTcpPacket::new(ip_packet.payload_mut()).unwrap();
+            tcp_packet.set_source(src.port());
+            tcp_packet.set_destination(dst.port());
+            tcp_packet.set_data_offset(5);
+            tcp_packet.set_flags(TcpFlags::SYN | TcpFlags::ACK);
+            tcp_packet.set_window(65535);
+            tcp_packet.set_checksum(tcp::ipv4_checksum(
+                &tcp_packet.to_immutable(),
+                &src_ip,
+                &dst_ip,
+            ));
+
+            ip_packet.set_checksum(ipv4::checksum(&ip_packet.to_immutable()));
+        }
+
+        buf
+    }
+
+    fn build_udp_followup_fragment(src: Ipv4Addr, dst: Ipv4Addr) -> Vec<u8> {
+        let mut buf = vec![0u8; 28];
+        {
+            let mut ip_packet = MutableIpv4Packet::new(&mut buf).unwrap();
+            ip_packet.set_version(4);
+            ip_packet.set_header_length(5);
+            ip_packet.set_total_length(28);
+            ip_packet.set_ttl(64);
+            ip_packet.set_next_level_protocol(IpNextHeaderProtocols::Udp);
+            ip_packet.set_fragment_offset(1);
+            ip_packet.set_source(src);
+            ip_packet.set_destination(dst);
+            ip_packet
+                .payload_mut()
+                .copy_from_slice(&[0xde, 0xad, 0xbe, 0xef, 0xca, 0xfe, 0xba, 0xbe]);
+
+            ip_packet.set_checksum(ipv4::checksum(&ip_packet.to_immutable()));
+        }
+
+        buf
+    }
+
+    #[tokio::test]
+    async fn socks5_consumes_modified_data_when_entry_matches() {
+        let peer_manager = create_mock_peer_manager().await;
+        let server = Socks5Server::new(peer_manager.get_global_ctx(), peer_manager, None);
+
+        let local = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 144, 144, 1)), 40000);
+        let remote = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 144, 144, 3)), 22);
+        let entry = Socks5Entry {
+            src: local,
+            dst: remote,
+            entry_type: TCP_ENTRY,
+        };
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        insert_entry_and_increment_count(
+            &server.entries,
+            &server.entry_count,
+            entry,
+            Socks5EntryData::Tcp(listener),
+        );
+
+        for packet_type in [
+            PacketType::DataWithKcpSrcModified,
+            PacketType::DataWithQuicSrcModified,
+        ] {
+            let mut packet = ZCPacket::new_with_payload(&build_tcp_packet(remote, local));
+            packet.fill_peer_manager_hdr(1, 1, packet_type as u8);
+
+            let result = server.try_process_packet_from_peer(packet).await;
+            assert!(result.is_none());
+
+            let mut receiver = server.packet_recv.lock().await;
+            let received = receiver.try_recv().unwrap();
+            assert_eq!(
+                received.peer_manager_header().unwrap().packet_type,
+                packet_type as u8
+            );
+        }
+    }
+
+    #[tokio::test]
+    async fn socks5_passes_through_unmatched_or_malformed_modified_data() {
+        let peer_manager = create_mock_peer_manager().await;
+        let server = Socks5Server::new(peer_manager.get_global_ctx(), peer_manager, None);
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        insert_entry_and_increment_count(
+            &server.entries,
+            &server.entry_count,
+            Socks5Entry {
+                src: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 144, 144, 1)), 40000),
+                dst: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 144, 144, 3)), 22),
+                entry_type: TCP_ENTRY,
+            },
+            Socks5EntryData::Tcp(listener),
+        );
+
+        let unmatched_local = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 144, 144, 1)), 40001);
+        let remote = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 144, 144, 3)), 22);
+        let mut unmatched_packet =
+            ZCPacket::new_with_payload(&build_tcp_packet(remote, unmatched_local));
+        unmatched_packet.fill_peer_manager_hdr(1, 2, PacketType::DataWithKcpSrcModified as u8);
+        let result = server.try_process_packet_from_peer(unmatched_packet).await;
+        assert!(result.is_some());
+
+        let mut malformed_packet = ZCPacket::new_with_payload(&[0u8; 8]);
+        malformed_packet.fill_peer_manager_hdr(1, 2, PacketType::DataWithQuicSrcModified as u8);
+        let result = server.try_process_packet_from_peer(malformed_packet).await;
+        assert!(result.is_some());
+
+        let mut receiver = server.packet_recv.lock().await;
+        assert!(receiver.try_recv().is_err());
+    }
+
+    #[tokio::test]
+    async fn socks5_passes_through_non_loopback_modified_data_even_when_entry_matches() {
+        let peer_manager = create_mock_peer_manager().await;
+        let server = Socks5Server::new(peer_manager.get_global_ctx(), peer_manager, None);
+
+        let local = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 144, 144, 1)), 40000);
+        let remote = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 144, 144, 3)), 22);
+        let entry = Socks5Entry {
+            src: local,
+            dst: remote,
+            entry_type: TCP_ENTRY,
+        };
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        insert_entry_and_increment_count(
+            &server.entries,
+            &server.entry_count,
+            entry,
+            Socks5EntryData::Tcp(listener),
+        );
+
+        let mut packet = ZCPacket::new_with_payload(&build_tcp_packet(remote, local));
+        packet.fill_peer_manager_hdr(1, 2, PacketType::DataWithKcpSrcModified as u8);
+
+        let result = server.try_process_packet_from_peer(packet).await;
+        assert!(result.is_some());
+
+        let mut receiver = server.packet_recv.lock().await;
+        assert!(receiver.try_recv().is_err());
+    }
+
+    #[tokio::test]
+    async fn socks5_mirrors_fragmented_udp_even_when_entry_count_is_stale_zero() {
+        let peer_manager = create_mock_peer_manager().await;
+        let server = Socks5Server::new(peer_manager.get_global_ctx(), peer_manager, None);
+
+        let local = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 144, 144, 1)), 40000);
+        let remote = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 144, 144, 3)), 53);
+        let udp_socket = Arc::new(tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap());
+        server.entries.insert(
+            Socks5Entry {
+                src: local,
+                dst: remote,
+                entry_type: UDP_ENTRY,
+            },
+            Socks5EntryData::Udp((
+                Arc::new(SocksUdpSocket::UdpSocket(udp_socket)),
+                UdpClientKey {
+                    client_addr: local,
+                    dst_addr: remote,
+                },
+            )),
+        );
+        assert_eq!(server.entry_count.load(Ordering::Relaxed), 0);
+
+        let mut packet = ZCPacket::new_with_payload(&build_udp_followup_fragment(
+            match remote.ip() {
+                IpAddr::V4(ip) => ip,
+                IpAddr::V6(_) => unreachable!(),
+            },
+            match local.ip() {
+                IpAddr::V4(ip) => ip,
+                IpAddr::V6(_) => unreachable!(),
+            },
+        ));
+        packet.fill_peer_manager_hdr(1, 2, PacketType::Data as u8);
+
+        let result = server.try_process_packet_from_peer(packet).await;
+        assert!(result.is_some());
+
+        let mut receiver = server.packet_recv.lock().await;
+        let received = receiver.try_recv().unwrap();
+        assert_eq!(
+            received.peer_manager_header().unwrap().packet_type,
+            PacketType::Data as u8
+        );
+    }
+
+    #[test]
+    fn decrement_entry_count_does_not_underflow() {
+        let entry_count = AtomicUsize::new(0);
+
+        let (old_entry_count, new_entry_count) = decrement_entry_count(&entry_count);
+
+        assert_eq!(old_entry_count, 0);
+        assert_eq!(new_entry_count, 0);
+        assert_eq!(entry_count.load(Ordering::Relaxed), 0);
+    }
+
+    #[tokio::test]
+    async fn removing_missing_entry_does_not_decrement_entry_count() {
+        let entries = Arc::new(DashMap::new());
+        let entry_count = AtomicUsize::new(1);
+        let entry = Socks5Entry {
+            src: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 42, 0, 2)), 40000),
+            dst: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 42, 0, 1)), 22),
+            entry_type: TCP_ENTRY,
+        };
+
+        let (removed, old_entry_count, new_entry_count) =
+            remove_entry_and_decrement_count(&entries, &entry_count, &entry);
+
+        assert!(!removed);
+        assert_eq!(old_entry_count, 1);
+        assert_eq!(new_entry_count, 1);
+        assert_eq!(entry_count.load(Ordering::Relaxed), 1);
+    }
+
+    #[tokio::test]
+    async fn removing_present_entry_decrements_entry_count_once() {
+        let entries = Arc::new(DashMap::new());
+        let entry_count = AtomicUsize::new(0);
+        let entry = Socks5Entry {
+            src: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 42, 0, 2)), 40000),
+            dst: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 42, 0, 1)), 22),
+            entry_type: TCP_ENTRY,
+        };
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        insert_entry_and_increment_count(
+            &entries,
+            &entry_count,
+            entry.clone(),
+            Socks5EntryData::Tcp(listener),
+        );
+
+        let (removed, old_entry_count, new_entry_count) =
+            remove_entry_and_decrement_count(&entries, &entry_count, &entry);
+        let (removed_again, old_entry_count_again, new_entry_count_again) =
+            remove_entry_and_decrement_count(&entries, &entry_count, &entry);
+
+        assert!(removed);
+        assert_eq!(old_entry_count, 1);
+        assert_eq!(new_entry_count, 0);
+        assert!(!removed_again);
+        assert_eq!(old_entry_count_again, 0);
+        assert_eq!(new_entry_count_again, 0);
+        assert_eq!(entry_count.load(Ordering::Relaxed), 0);
+    }
+
+    #[tokio::test]
+    async fn replacing_present_entry_does_not_increment_entry_count() {
+        let entries = Arc::new(DashMap::new());
+        let entry_count = AtomicUsize::new(0);
+        let entry = Socks5Entry {
+            src: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 42, 0, 2)), 40000),
+            dst: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 42, 0, 1)), 22),
+            entry_type: TCP_ENTRY,
+        };
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let replacement = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+
+        let (replaced, old_entry_count, new_entry_count) = insert_entry_and_increment_count(
+            &entries,
+            &entry_count,
+            entry.clone(),
+            Socks5EntryData::Tcp(listener),
+        );
+        let (replaced_again, old_entry_count_again, new_entry_count_again) =
+            insert_entry_and_increment_count(
+                &entries,
+                &entry_count,
+                entry,
+                Socks5EntryData::Tcp(replacement),
+            );
+
+        assert!(!replaced);
+        assert_eq!(old_entry_count, 0);
+        assert_eq!(new_entry_count, 1);
+        assert!(replaced_again);
+        assert_eq!(old_entry_count_again, 1);
+        assert_eq!(new_entry_count_again, 1);
+        assert_eq!(entry_count.load(Ordering::Relaxed), 1);
+    }
+}

easytier/src/gateway/socks5/dataplane.rs

@@ -26,7 +26,6 @@ use std::{
 };
 
 use anyhow::Context as _;
-use dashmap::mapref::entry::Entry;
 use hotpath::instant::Instant;
 use tokio::io::{AsyncRead, AsyncWrite};
 
@@ -35,6 +34,7 @@ use crate::{common::error::Error, gateway::fast_socks5::server::AsyncTcpConnecto
 use super::{
     Socks5AutoConnector, Socks5Entry, Socks5EntryData, Socks5EntrySet, Socks5Server,
     SocksTcpStream, SocksUdpSocket, TCP_ENTRY, TCP_LISTEN_ENTRY, UDP_ENTRY, UdpClientKey,
+    decrement_entry_count, insert_entry_and_increment_count, try_insert_entry_and_increment_count,
 };
 use crate::gateway::tokio_smoltcp::{Net, TcpListener};
 
@@ -59,12 +59,12 @@ impl OwnedRouteEntry {
         entry_count: Arc<AtomicUsize>,
         entry: Socks5Entry,
     ) -> Self {
-        if entries
+        insert_entry_and_increment_count(
-            .insert(entry.clone(), Socks5EntryData::DataPlaneRoute)
+            &entries,
-            .is_none()
+            &entry_count,
-        {
+            entry.clone(),
-            entry_count.fetch_add(1, Ordering::Relaxed);
+            Socks5EntryData::DataPlaneRoute,
-        }
+        );
         Self {
             entries,
             entry_count,
@@ -78,12 +78,13 @@ impl OwnedRouteEntry {
         entry_count: Arc<AtomicUsize>,
         entry: Socks5Entry,
     ) -> Option<Self> {
-        match entries.entry(entry.clone()) {
+        if !try_insert_entry_and_increment_count(
-            Entry::Occupied(_) => return None,
+            &entries,
-            Entry::Vacant(vacant) => {
+            &entry_count,
-                vacant.insert(Socks5EntryData::DataPlaneRoute);
+            entry.clone(),
-                entry_count.fetch_add(1, Ordering::Relaxed);
+            Socks5EntryData::DataPlaneRoute,
-            }
+        ) {
+            return None;
         }
         Some(Self {
             entries,
@@ -96,7 +97,7 @@ impl OwnedRouteEntry {
 impl Drop for OwnedRouteEntry {
     fn drop(&mut self) {
         if self.entries.remove(&self.entry).is_some() {
-            self.entry_count.fetch_sub(1, Ordering::Relaxed);
+            decrement_entry_count(&self.entry_count);
         }
     }
 }
@@ -224,16 +225,18 @@ impl DataPlaneUdpSocket {
             dst: addr,
             entry_type: UDP_ENTRY,
         };
-        if let Entry::Vacant(entry) = self.entries.entry(key) {
+        try_insert_entry_and_increment_count(
-            entry.insert(Socks5EntryData::Udp((
+            &self.entries,
+            &self.entry_count,
+            key,
+            Socks5EntryData::Udp((
                 self.socket.clone(),
                 UdpClientKey {
                     client_addr: self.local_addr,
                     dst_addr: addr,
                 },
-            )));
+            )),
-            self.entry_count.fetch_add(1, Ordering::Relaxed);
+        );
-        }
         self.socket.send_to(buf, addr).await
     }
 
@@ -244,13 +247,15 @@ impl DataPlaneUdpSocket {
 
 impl Drop for DataPlaneUdpSocket {
     fn drop(&mut self) {
+        let mut removed_entries = 0;
         self.entries.retain(|_, data| match data {
             Socks5EntryData::Udp((socket, _)) if Arc::ptr_eq(socket, &self.socket) => {
-                self.entry_count.fetch_sub(1, Ordering::Relaxed);
+                removed_entries += 1;
                 false
             }
             _ => true,
         });
+        super::decrement_entry_count_by(&self.entry_count, removed_entries);
     }
 }
 

easytier/src/tests/three_node.rs

@@ -2121,6 +2121,124 @@ pub async fn port_forward_test(
     drop_insts(_insts).await;
 }
 
+#[rstest::rstest]
+#[case(false, false)]
+#[case(true, false)]
+#[case(true, true)]
+#[serial_test::serial]
+#[tokio::test]
+pub async fn port_forward_with_inbound_default_drop_acl_test(
+    #[case] dhcp: bool,
+    #[case] enable_quic_proxy: bool,
+) {
+    use crate::proto::acl::*;
+
+    let acl = Acl {
+        acl_v1: Some(AclV1 {
+            chains: vec![Chain {
+                name: "drop_unsolicited_inbound".to_string(),
+                chain_type: ChainType::Inbound as i32,
+                enabled: true,
+                default_action: Action::Drop as i32,
+                ..Default::default()
+            }],
+            ..Default::default()
+        }),
+    };
+
+    let insts = init_three_node_ex(
+        "udp",
+        |cfg| {
+            if cfg.get_inst_name() == "inst1" {
+                if dhcp {
+                    cfg.set_ipv4(None);
+                    cfg.set_dhcp(true);
+                }
+                cfg.set_acl(Some(acl.clone()));
+                cfg.set_port_forwards(vec![
+                    PortForwardConfig {
+                        bind_addr: "0.0.0.0:23456".parse().unwrap(),
+                        dst_addr: "10.144.144.3:23456".parse().unwrap(),
+                        proto: "tcp".to_string(),
+                    },
+                    PortForwardConfig {
+                        bind_addr: "0.0.0.0:23457".parse().unwrap(),
+                        dst_addr: "10.1.2.4:23457".parse().unwrap(),
+                        proto: "tcp".to_string(),
+                    },
+                ]);
+
+                let mut flags = cfg.get_flags();
+                flags.no_tun = true;
+                flags.enable_kcp_proxy = false;
+                flags.enable_quic_proxy = enable_quic_proxy;
+                cfg.set_flags(flags);
+            } else if cfg.get_inst_name() == "inst3" {
+                cfg.add_proxy_cidr("10.1.2.0/24".parse().unwrap(), None)
+                    .unwrap();
+
+                let mut flags = cfg.get_flags();
+                flags.disable_kcp_input = true;
+                flags.disable_quic_input = !enable_quic_proxy;
+                cfg.set_flags(flags);
+            } else if cfg.get_inst_name() == "inst2" {
+                let mut flags = cfg.get_flags();
+                flags.disable_relay_kcp = true;
+                cfg.set_flags(flags);
+            }
+
+            cfg
+        },
+        false,
+    )
+    .await;
+
+    if dhcp {
+        wait_for_condition(
+            || async { insts[0].get_global_ctx().get_ipv4().is_some() },
+            Duration::from_secs(5),
+        )
+        .await;
+    }
+
+    for (bind_port, server_ns) in [(23456, "net_c"), (23457, "net_d")] {
+        let tcp_listener =
+            TcpTunnelListener::new(format!("tcp://0.0.0.0:{bind_port}").parse().unwrap());
+        let tcp_connector =
+            TcpTunnelConnector::new(format!("tcp://127.0.0.1:{bind_port}").parse().unwrap());
+
+        let mut buf = vec![0; 64];
+        rand::thread_rng().fill(&mut buf[..]);
+
+        let result = _tunnel_pingpong_netns_with_timeout(
+            tcp_listener,
+            tcp_connector,
+            NetNS::new(Some(server_ns.into())),
+            NetNS::new(Some("net_a".into())),
+            buf,
+            Duration::from_secs(1),
+        )
+        .await;
+
+        let stats = insts[0].get_global_ctx().get_acl_filter().get_stats();
+        println!(
+            "port forward source bind_port={} dhcp={} enable_quic_proxy={} ACL stats: {}",
+            bind_port, dhcp, enable_quic_proxy, stats
+        );
+
+        assert!(
+            result.is_ok(),
+            "port-forward TCP should complete through outbound ACL state, bind_port={}, dhcp={}, enable_quic_proxy={}; stats: {}",
+            bind_port,
+            dhcp,
+            enable_quic_proxy,
+            stats,
+        );
+    }
+
+    drop_insts(insts).await;
+}
+
 #[rstest::rstest]
 #[serial_test::serial]
 #[tokio::test]