From 15e5d89f70c3f64f047c48d60e67c12c74c5bebc Mon Sep 17 00:00:00 2001 From: KKRainbow <443152178@qq.com> Date: Mon, 29 Jun 2026 10:24:57 +0800 Subject: [PATCH] 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 | 2 +- easytier/src/gateway/socks5.rs | 759 +++++++++++++++++++++-- easytier/src/gateway/socks5/dataplane.rs | 45 +- easytier/src/tests/three_node.rs | 118 ++++ 4 files changed, 845 insertions(+), 79 deletions(-) diff --git a/.github/workflows/ohos.yml b/.github/workflows/ohos.yml index 6c4d2951..5c780cc2 100644 --- a/.github/workflows/ohos.yml +++ b/.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 }} - diff --git a/easytier/src/gateway/socks5.rs b/easytier/src/gateway/socks5.rs index 3d92d3e2..d1ef4508 100644 --- a/easytier/src/gateway/socks5.rs +++ b/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>; +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, entries: Socks5EntrySet, @@ -190,9 +271,20 @@ impl AsyncTcpConnector for SmolTcpConnector { entry_type: TCP_ENTRY, }; *self.current_entry.lock().unwrap() = Some(entry.clone()); - self.entries - .insert(entry, Socks5EntryData::Tcp(tmp_listener)); - self.entry_count.fetch_add(1, Ordering::Relaxed); + let (replaced, old_entry_count, new_entry_count) = insert_entry_and_increment_count( + &self.entries, + &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); - self.entry_count.fetch_sub(1, Ordering::Relaxed); + let (removed, old_entry_count, new_entry_count) = + 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() - || peer_mgr_arc.get_msg_dst_peer(&addr.ip()).await.0.is_empty() + let has_smoltcp_net = self.smoltcp_net.is_some(); + 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 + Send> = match (&self.kcp_endpoint, dst_allow_kcp) { - (Some(kcp_endpoint), true) => Box::new(Socks5KcpConnector { - kcp_endpoint: kcp_endpoint.clone(), - peer_mgr: self.peer_mgr.clone(), - src_addr: self.src_addr, - }), - (_, _) => 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), - }), + (Some(kcp_endpoint), true) => { + tracing::trace!( + ?addr, + src_addr = ?self.src_addr, + dst_peer_count = dst_peers.as_ref().map(Vec::len), + "socks5 auto connector selected kcp" + ); + Box::new(Socks5KcpConnector { + kcp_endpoint: kcp_endpoint.clone(), + 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 { - net: self.smoltcp_net.clone().unwrap(), - entries: self.entries.clone(), - entry_count: self.entry_count.clone(), - current_entry: std::sync::Mutex::new(None), - }); + let connector = { + tracing::trace!( + ?addr, + src_addr = ?self.src_addr, + 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 { - if self.entry_count.load(Ordering::Relaxed) == 0 - && !self.socks5_enabled.load(Ordering::Relaxed) - { + let entry_count = self.entry_count.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()), - src: SocketAddr::new( - ipv4.get_destination().into(), - udp_packet.get_destination(), - ), - entry_type: UDP_ENTRY, - } + ( + Socks5Entry { + dst: SocketAddr::new(ipv4.get_source().into(), udp_packet.get_source()), + src: SocketAddr::new( + ipv4.get_destination().into(), + udp_packet.get_destination(), + ), + 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() - && !socks5_enabled.load(Ordering::Relaxed) - && !data_plane_active - { - let _ = net.lock().await.take(); + let active_port_forwards = cancel_tokens.len(); + let is_socks5_enabled = socks5_enabled.load(Ordering::Relaxed); + if active_port_forwards == 0 && !is_socks5_enabled && !data_plane_active { + let had_net = { + 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 + 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 { + 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 { + 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); + } +} diff --git a/easytier/src/gateway/socks5/dataplane.rs b/easytier/src/gateway/socks5/dataplane.rs index ca895a40..4549ef70 100644 --- a/easytier/src/gateway/socks5/dataplane.rs +++ b/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, entry: Socks5Entry, ) -> Self { - if entries - .insert(entry.clone(), Socks5EntryData::DataPlaneRoute) - .is_none() - { - entry_count.fetch_add(1, Ordering::Relaxed); - } + insert_entry_and_increment_count( + &entries, + &entry_count, + entry.clone(), + Socks5EntryData::DataPlaneRoute, + ); Self { entries, entry_count, @@ -78,12 +78,13 @@ impl OwnedRouteEntry { entry_count: Arc, entry: Socks5Entry, ) -> Option { - match entries.entry(entry.clone()) { - Entry::Occupied(_) => return None, - Entry::Vacant(vacant) => { - vacant.insert(Socks5EntryData::DataPlaneRoute); - entry_count.fetch_add(1, Ordering::Relaxed); - } + if !try_insert_entry_and_increment_count( + &entries, + &entry_count, + entry.clone(), + 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) { - entry.insert(Socks5EntryData::Udp(( + try_insert_entry_and_increment_count( + &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); } } diff --git a/easytier/src/tests/three_node.rs b/easytier/src/tests/three_node.rs index 5fa89640..32f262ae 100644 --- a/easytier/src/tests/three_node.rs +++ b/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]