mirror/Easytier
1
0
Fork 0
mirror of https://github.com/EasyTier/EasyTier.git synced 2026-08-05 20:19:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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.
This commit is contained in:
KKRainbow
2026-06-29 10:24:57 +08:00
committed by GitHub
parent 46f1b57367
commit 15e5d89f70
4 changed files with 845 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/ohos.yml
+1 -1
View File
@@ -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
+701 -58
View File
@@ -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
.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<dyn AsyncTcpConnector<S = SocksTcpStream> + 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<ZCPacket> {
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<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
+25 -20
View File
@@ -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.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<AtomicUsize>,
entry: Socks5Entry,
) -> Option<Self> {
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);
}
}
easytier/src/tests/three_node.rs
+118
View File
@@ -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]