mirror of
https://github.com/EasyTier/EasyTier.git
synced 2026-10-16 09:36:14 -08:00
fix(proxy): isolate TCP flows and recover KCP control loss (#2569)
* fix(proxy): distinguish TCP flows sharing source ports Key TCP NAT entries by source and destination so connections sharing a source port can reach different targets independently. Allocate a unique translated source port for each flow and restore the original addresses and ports on the return path. Keep connection identity during accept and cleanup so a replaced flow cannot lose its new mapping. Give each smoltcp proxy mode its own listener port and retain the original destination for mapped subnets. Include regression coverage for SYN retries, concurrent insertion, stale cleanup, translated port wraparound and exhaustion, plus the recorded cross-platform and six-mode traffic validation. * fix(proxy): preserve half-closed streams and pending KCP accepts Propagate EOF to the opposite writer while keeping reverse traffic alive until both directions finish. This allows a server to reply after reading request EOF and a client to send after reading response EOF. Pin kcp-sys to b37ee660 for retained pending accepts, receive-buffer draining before EOF, early FIN handling, atomic state cleanup, and keepalives that start only after the handshake completes. Cover both shutdown directions, empty requests and buffered transfers in gateway tests and the TCP/KCP/QUIC kernel/smoltcp integration matrix. Include the half-close validation record and its remaining loss limits. * fix(proxy): enable backward-compatible KCP control recovery Pin kcp-sys to 26853356 and update both the root and OHOS workspace lockfiles. The dependency negotiates reliable control per connection while preserving legacy behavior when either endpoint lacks support. New peers recover lost handshake and FIN control packets, acknowledge FIN without closing the reverse writer, and retain bounded close state for duplicate packets. Queue saturation cannot block reliable stateless responses. Compatibility tests retain one historical baseline. Include the protocol design and validation records for deterministic loss, real old endpoints, three platforms, the full integration matrix, and proxy traffic. Preserve mixed-version failures and the documented minor scheduling issue rather than claiming legacy loss recovery. Validate the OHOS workspace with --locked: 12 library tests and the binding test build check pass, along with formatting.
This commit is contained in:
1 parent
286e0f4a0d
commit
5477d4bca2
13 files changed
+1859
-148
No files matched your search
@@ -9,7 +9,7 @@ use std::{
|
||||
|
||||
use cidr::Ipv4Inet;
|
||||
use crossbeam::atomic::AtomicCell;
|
||||
use dashmap::DashMap;
|
||||
use dashmap::{DashMap, mapref::entry::Entry};
|
||||
use smoltcp::wire::{IpAddress, IpProtocol, Ipv4Packet, TcpPacket};
|
||||
|
||||
use crate::packet::{PacketType, ZCPacket};
|
||||
@@ -25,6 +25,16 @@ pub(crate) enum TcpProxyMode {
|
||||
QuicSrc,
|
||||
}
|
||||
|
||||
impl TcpProxyMode {
|
||||
pub(super) const fn smoltcp_listener_port(self) -> u16 {
|
||||
match self {
|
||||
Self::Tcp => 8899,
|
||||
Self::KcpSrc => 8900,
|
||||
Self::QuicSrc => 8901,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum TcpNatEntryState {
|
||||
SynReceived,
|
||||
@@ -44,24 +54,30 @@ pub struct TcpNatEntrySnapshot {
|
||||
pub state: TcpNatEntryState,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
|
||||
struct TcpNatFlowKey {
|
||||
src: SocketAddr,
|
||||
mapped_dst: SocketAddr,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct TcpNatEntry {
|
||||
id: TcpNatEntryId,
|
||||
src: SocketAddr,
|
||||
flow: TcpNatFlowKey,
|
||||
translated_src: SocketAddr,
|
||||
real_dst: SocketAddr,
|
||||
mapped_dst: SocketAddr,
|
||||
start_time: Instant,
|
||||
start_time_unix_secs: u64,
|
||||
state: AtomicCell<TcpNatEntryState>,
|
||||
}
|
||||
|
||||
impl TcpNatEntry {
|
||||
fn new(src: SocketAddr, real_dst: SocketAddr, mapped_dst: SocketAddr) -> Self {
|
||||
fn new(flow: TcpNatFlowKey, translated_src: SocketAddr, real_dst: SocketAddr) -> Self {
|
||||
Self {
|
||||
id: uuid::Uuid::new_v4(),
|
||||
src,
|
||||
flow,
|
||||
translated_src,
|
||||
real_dst,
|
||||
mapped_dst,
|
||||
start_time: Instant::now(),
|
||||
start_time_unix_secs: SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
@@ -76,7 +92,7 @@ impl TcpNatEntry {
|
||||
}
|
||||
|
||||
pub fn src(&self) -> SocketAddr {
|
||||
self.src
|
||||
self.flow.src
|
||||
}
|
||||
|
||||
pub fn real_dst(&self) -> SocketAddr {
|
||||
@@ -84,7 +100,7 @@ impl TcpNatEntry {
|
||||
}
|
||||
|
||||
pub fn mapped_dst(&self) -> SocketAddr {
|
||||
self.mapped_dst
|
||||
self.flow.mapped_dst
|
||||
}
|
||||
|
||||
pub fn state(&self) -> TcpNatEntryState {
|
||||
@@ -97,9 +113,9 @@ impl TcpNatEntry {
|
||||
|
||||
fn snapshot(&self) -> TcpNatEntrySnapshot {
|
||||
TcpNatEntrySnapshot {
|
||||
src: self.src,
|
||||
src: self.src(),
|
||||
dst: self.real_dst,
|
||||
mapped_dst: self.mapped_dst,
|
||||
mapped_dst: self.mapped_dst(),
|
||||
start_time: self.start_time_unix_secs,
|
||||
state: self.state(),
|
||||
}
|
||||
@@ -127,6 +143,7 @@ pub(crate) struct TcpProxyNicContext {
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum TcpProxyPacketAction {
|
||||
Handled { new_syn: bool },
|
||||
Drop,
|
||||
Pass,
|
||||
}
|
||||
|
||||
@@ -134,9 +151,10 @@ pub(crate) enum TcpProxyPacketAction {
|
||||
pub(crate) struct TcpProxyEngine {
|
||||
cidr_table: Arc<ProxyCidrTable>,
|
||||
local_port: AtomicU16,
|
||||
syn_map: DashMap<SocketAddr, Arc<TcpNatEntry>>,
|
||||
next_translated_port: AtomicU16,
|
||||
flow_map: DashMap<TcpNatFlowKey, Arc<TcpNatEntry>>,
|
||||
translated_src_map: DashMap<SocketAddr, Arc<TcpNatEntry>>,
|
||||
conn_map: DashMap<uuid::Uuid, Arc<TcpNatEntry>>,
|
||||
addr_conn_map: DashMap<SocketAddr, Arc<TcpNatEntry>>,
|
||||
}
|
||||
|
||||
impl TcpProxyEngine {
|
||||
@@ -144,9 +162,10 @@ impl TcpProxyEngine {
|
||||
Self {
|
||||
cidr_table,
|
||||
local_port: AtomicU16::new(0),
|
||||
syn_map: DashMap::new(),
|
||||
next_translated_port: AtomicU16::new(1),
|
||||
flow_map: DashMap::new(),
|
||||
translated_src_map: DashMap::new(),
|
||||
conn_map: DashMap::new(),
|
||||
addr_conn_map: DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -158,6 +177,61 @@ impl TcpProxyEngine {
|
||||
self.local_port.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
fn allocate_entry(
|
||||
&self,
|
||||
flow: TcpNatFlowKey,
|
||||
real_dst: SocketAddr,
|
||||
) -> Option<Arc<TcpNatEntry>> {
|
||||
let local_port = self.local_port();
|
||||
for _ in 0..u16::MAX {
|
||||
let translated_port = self
|
||||
.next_translated_port
|
||||
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |port| {
|
||||
Some(if port == u16::MAX { 1 } else { port + 1 })
|
||||
})
|
||||
.expect("translated port counter update cannot fail");
|
||||
if translated_port == local_port {
|
||||
continue;
|
||||
}
|
||||
let translated_src = SocketAddr::new(flow.src.ip(), translated_port);
|
||||
let Entry::Vacant(slot) = self.translated_src_map.entry(translated_src) else {
|
||||
continue;
|
||||
};
|
||||
let entry = Arc::new(TcpNatEntry::new(flow, translated_src, real_dst));
|
||||
slot.insert(entry.clone());
|
||||
return Some(entry);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn entry_for_syn(
|
||||
&self,
|
||||
flow: TcpNatFlowKey,
|
||||
real_dst: SocketAddr,
|
||||
) -> Option<(Arc<TcpNatEntry>, bool)> {
|
||||
match self.flow_map.entry(flow) {
|
||||
Entry::Occupied(mut slot) => {
|
||||
let entry = slot.get();
|
||||
if !matches!(
|
||||
entry.state(),
|
||||
TcpNatEntryState::ClosingSrc
|
||||
| TcpNatEntryState::ClosingDst
|
||||
| TcpNatEntryState::Closed
|
||||
) {
|
||||
return Some((entry.clone(), false));
|
||||
}
|
||||
let entry = self.allocate_entry(flow, real_dst)?;
|
||||
slot.insert(entry.clone());
|
||||
Some((entry, true))
|
||||
}
|
||||
Entry::Vacant(slot) => {
|
||||
let entry = self.allocate_entry(flow, real_dst)?;
|
||||
slot.insert(entry.clone());
|
||||
Some((entry, true))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn check_packet_from_peer_fast(
|
||||
&self,
|
||||
mode: TcpProxyMode,
|
||||
@@ -230,28 +304,36 @@ impl TcpProxyEngine {
|
||||
|
||||
let source_ip = ip_packet.src_addr();
|
||||
let source_port = tcp_packet.src_port();
|
||||
let src = SocketAddr::V4(SocketAddrV4::new(source_ip, source_port));
|
||||
let dest_ip = ip_packet.dst_addr();
|
||||
let dest_port = tcp_packet.dst_port();
|
||||
let flow = TcpNatFlowKey {
|
||||
src: SocketAddr::V4(SocketAddrV4::new(source_ip, source_port)),
|
||||
mapped_dst: SocketAddr::V4(SocketAddrV4::new(dest_ip, dest_port)),
|
||||
};
|
||||
|
||||
let mut new_syn = false;
|
||||
if tcp_packet.syn() && !tcp_packet.ack() {
|
||||
let dest_ip = ip_packet.dst_addr();
|
||||
let dest_port = tcp_packet.dst_port();
|
||||
let mapped_dst = SocketAddr::V4(SocketAddrV4::new(dest_ip, dest_port));
|
||||
let is_syn = tcp_packet.syn() && !tcp_packet.ack();
|
||||
let (entry, new_syn) = if is_syn {
|
||||
let real_dst = SocketAddr::V4(SocketAddrV4::new(real_dst_ip, dest_port));
|
||||
let Some(entry) = self.entry_for_syn(flow, real_dst) else {
|
||||
tracing::error!(?flow, "tcp proxy translated source ports exhausted");
|
||||
return TcpProxyPacketAction::Drop;
|
||||
};
|
||||
entry
|
||||
} else {
|
||||
let Some(entry) = self.flow_map.get(&flow) else {
|
||||
return TcpProxyPacketAction::Pass;
|
||||
};
|
||||
(entry.clone(), false)
|
||||
};
|
||||
|
||||
let old_val = self
|
||||
.syn_map
|
||||
.insert(src, Arc::new(TcpNatEntry::new(src, real_dst, mapped_dst)));
|
||||
if new_syn {
|
||||
tracing::info!(
|
||||
?src,
|
||||
?real_dst,
|
||||
?mapped_dst,
|
||||
old_entry = ?old_val,
|
||||
src = ?entry.src(),
|
||||
translated_src = ?entry.translated_src,
|
||||
real_dst = ?entry.real_dst(),
|
||||
mapped_dst = ?entry.mapped_dst(),
|
||||
"tcp syn received"
|
||||
);
|
||||
new_syn = true;
|
||||
} else if !self.addr_conn_map.contains_key(&src) && !self.syn_map.contains_key(&src) {
|
||||
return TcpProxyPacketAction::Pass;
|
||||
}
|
||||
|
||||
let mut ip_packet = Ipv4Packet::new_checked(payload_bytes).expect("checked ipv4 packet");
|
||||
@@ -263,6 +345,7 @@ impl TcpProxyEngine {
|
||||
{
|
||||
let mut tcp_packet =
|
||||
TcpPacket::new_checked(ip_packet.payload_mut()).expect("checked tcp packet");
|
||||
tcp_packet.set_src_port(entry.translated_src.port());
|
||||
tcp_packet.set_dst_port(ctx.local_port);
|
||||
tcp_packet.fill_checksum(&IpAddress::Ipv4(source), &IpAddress::Ipv4(local_ip));
|
||||
}
|
||||
@@ -312,17 +395,13 @@ impl TcpProxyEngine {
|
||||
}
|
||||
|
||||
tracing::trace!(?dst_addr, "tcp packet try find entry");
|
||||
let entry = if let Some(entry) = self.addr_conn_map.get(&dst_addr) {
|
||||
entry.clone()
|
||||
} else {
|
||||
let Some(syn_entry) = self.syn_map.get(&dst_addr) else {
|
||||
return false;
|
||||
};
|
||||
syn_entry.clone()
|
||||
let Some(entry) = self.translated_src_map.get(&dst_addr) else {
|
||||
return false;
|
||||
};
|
||||
assert_eq!(entry.src, dst_addr);
|
||||
let entry = entry.clone();
|
||||
assert_eq!(entry.translated_src, dst_addr);
|
||||
|
||||
let IpAddr::V4(mapped_dst_ip) = entry.mapped_dst.ip() else {
|
||||
let IpAddr::V4(mapped_dst_ip) = entry.mapped_dst().ip() else {
|
||||
panic!("v4 nat entry src ip is not v4");
|
||||
};
|
||||
|
||||
@@ -346,6 +425,7 @@ impl TcpProxyEngine {
|
||||
let mut tcp_packet =
|
||||
TcpPacket::new_checked(ip_packet.payload_mut()).expect("checked tcp packet");
|
||||
tcp_packet.set_src_port(entry.real_dst.port());
|
||||
tcp_packet.set_dst_port(entry.src().port());
|
||||
tcp_packet.fill_checksum(&IpAddress::Ipv4(mapped_dst_ip), &IpAddress::Ipv4(dst));
|
||||
}
|
||||
ip_packet.fill_checksum();
|
||||
@@ -366,53 +446,90 @@ impl TcpProxyEngine {
|
||||
}
|
||||
}
|
||||
|
||||
let (_, entry) = self.syn_map.remove(&socket_addr)?;
|
||||
if entry.state() != TcpNatEntryState::SynReceived {
|
||||
let entry = self.translated_src_map.get(&socket_addr)?.clone();
|
||||
if entry
|
||||
.state
|
||||
.compare_exchange(
|
||||
TcpNatEntryState::SynReceived,
|
||||
TcpNatEntryState::ConnectingDst,
|
||||
)
|
||||
.is_err()
|
||||
{
|
||||
if entry.state() == TcpNatEntryState::Closed {
|
||||
self.remove_indices(&entry);
|
||||
}
|
||||
return None;
|
||||
}
|
||||
|
||||
entry.set_state(TcpNatEntryState::ConnectingDst);
|
||||
self.addr_conn_map.insert(entry.src, entry.clone());
|
||||
let old_nat_val = self.conn_map.insert(entry.id, entry.clone());
|
||||
assert!(old_nat_val.is_none());
|
||||
Some(entry)
|
||||
}
|
||||
|
||||
fn remove_indices(&self, entry: &TcpNatEntry) {
|
||||
self.flow_map
|
||||
.remove_if(&entry.flow, |_, current| current.id == entry.id);
|
||||
self.translated_src_map
|
||||
.remove_if(&entry.translated_src, |_, current| current.id == entry.id);
|
||||
}
|
||||
|
||||
pub fn remove_entry(&self, entry_id: TcpNatEntryId) {
|
||||
let Some((_, entry)) = self.conn_map.remove(&entry_id) else {
|
||||
return;
|
||||
};
|
||||
self.addr_conn_map
|
||||
.remove_if(&entry.src, |_, current| current.id == entry.id);
|
||||
self.remove_indices(&entry);
|
||||
if self.conn_map.capacity() - self.conn_map.len() > 16 {
|
||||
self.conn_map.shrink_to_fit();
|
||||
}
|
||||
if self.addr_conn_map.capacity() - self.addr_conn_map.len() > 16 {
|
||||
self.addr_conn_map.shrink_to_fit();
|
||||
if self.flow_map.capacity() - self.flow_map.len() > 16 {
|
||||
self.flow_map.shrink_to_fit();
|
||||
}
|
||||
if self.translated_src_map.capacity() - self.translated_src_map.len() > 16 {
|
||||
self.translated_src_map.shrink_to_fit();
|
||||
}
|
||||
}
|
||||
|
||||
pub fn cleanup_expired_syn(&self, timeout: Duration) {
|
||||
self.syn_map.retain(|_, entry| {
|
||||
if entry.start_time.elapsed() > timeout {
|
||||
tracing::warn!(?entry, "syn nat entry expired");
|
||||
entry.set_state(TcpNatEntryState::Closed);
|
||||
false
|
||||
} else {
|
||||
true
|
||||
}
|
||||
});
|
||||
self.syn_map.shrink_to_fit();
|
||||
pub fn clear(&self) {
|
||||
for entry in self.flow_map.iter() {
|
||||
entry.set_state(TcpNatEntryState::Closed);
|
||||
}
|
||||
for entry in self.conn_map.iter() {
|
||||
entry.set_state(TcpNatEntryState::Closed);
|
||||
}
|
||||
self.flow_map.clear();
|
||||
self.translated_src_map.clear();
|
||||
self.conn_map.clear();
|
||||
}
|
||||
|
||||
pub fn is_tcp_proxy_connection(&self, src: SocketAddr) -> bool {
|
||||
self.syn_map.contains_key(&src) || self.addr_conn_map.contains_key(&src)
|
||||
pub fn cleanup_expired_syn(&self, timeout: Duration) {
|
||||
self.flow_map.retain(|_, entry| {
|
||||
let expired = entry.start_time.elapsed() > timeout
|
||||
&& entry
|
||||
.state
|
||||
.compare_exchange(TcpNatEntryState::SynReceived, TcpNatEntryState::Closed)
|
||||
.is_ok();
|
||||
if expired {
|
||||
tracing::warn!(?entry, "syn nat entry expired");
|
||||
self.translated_src_map
|
||||
.remove_if(&entry.translated_src, |_, current| current.id == entry.id);
|
||||
}
|
||||
!expired
|
||||
});
|
||||
self.flow_map.shrink_to_fit();
|
||||
self.translated_src_map.shrink_to_fit();
|
||||
}
|
||||
|
||||
pub fn is_tcp_proxy_flow(&self, src: SocketAddr, mapped_dst: SocketAddr) -> bool {
|
||||
self.flow_map
|
||||
.contains_key(&TcpNatFlowKey { src, mapped_dst })
|
||||
}
|
||||
|
||||
pub fn list_entries(&self) -> Vec<TcpNatEntrySnapshot> {
|
||||
let mut entries = Vec::new();
|
||||
for entry in self.syn_map.iter() {
|
||||
entries.push(entry.value().snapshot());
|
||||
for entry in self.flow_map.iter() {
|
||||
if entry.state() == TcpNatEntryState::SynReceived {
|
||||
entries.push(entry.value().snapshot());
|
||||
}
|
||||
}
|
||||
for entry in self.conn_map.iter() {
|
||||
entries.push(entry.value().snapshot());
|
||||
@@ -503,6 +620,32 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
fn packet_src(packet: &ZCPacket) -> SocketAddrV4 {
|
||||
let ipv4 = Ipv4Packet::new_checked(packet.payload()).unwrap();
|
||||
let tcp = TcpPacket::new_checked(ipv4.payload()).unwrap();
|
||||
SocketAddrV4::new(ipv4.src_addr(), tcp.src_port())
|
||||
}
|
||||
|
||||
fn nic_ctx() -> TcpProxyNicContext {
|
||||
TcpProxyNicContext {
|
||||
local_inet: Some("10.144.144.204/24".parse().unwrap()),
|
||||
local_port: 8899,
|
||||
my_peer_id: 2,
|
||||
smoltcp_enabled: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smoltcp_listener_ports_are_unique_per_proxy_mode() {
|
||||
let tcp = TcpProxyMode::Tcp.smoltcp_listener_port();
|
||||
let kcp = TcpProxyMode::KcpSrc.smoltcp_listener_port();
|
||||
let quic = TcpProxyMode::QuicSrc.smoltcp_listener_port();
|
||||
|
||||
assert_ne!(tcp, kcp);
|
||||
assert_ne!(tcp, quic);
|
||||
assert_ne!(kcp, quic);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn peer_syn_creates_entry_and_rewrites_to_local_stack() {
|
||||
let engine = tcp_engine();
|
||||
@@ -522,7 +665,7 @@ mod tests {
|
||||
"10.144.144.204".parse::<Ipv4Addr>().unwrap()
|
||||
);
|
||||
let tcp = TcpPacket::new_checked(ipv4.payload()).unwrap();
|
||||
assert_eq!(tcp.src_port(), src.port());
|
||||
assert_ne!(tcp.src_port(), src.port());
|
||||
assert_eq!(tcp.dst_port(), 8899);
|
||||
|
||||
let entries = engine.list_entries();
|
||||
@@ -545,25 +688,18 @@ mod tests {
|
||||
engine.try_handle_peer_packet(TcpProxyMode::Tcp, &mut request, peer_ctx()),
|
||||
TcpProxyPacketAction::Handled { new_syn: true }
|
||||
));
|
||||
let translated_src = packet_src(&request);
|
||||
let entry = engine
|
||||
.accept_connection(
|
||||
SocketAddr::V4(src),
|
||||
SocketAddr::V4(translated_src),
|
||||
Some("10.144.144.204/24".parse().unwrap()),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(entry.state(), TcpNatEntryState::ConnectingDst);
|
||||
|
||||
let local = SocketAddrV4::new("10.144.144.204".parse().unwrap(), 8899);
|
||||
let mut response = build_tcp_packet(local, src, false, true);
|
||||
assert!(engine.try_process_packet_from_nic(
|
||||
&mut response,
|
||||
TcpProxyNicContext {
|
||||
local_inet: Some("10.144.144.204/24".parse().unwrap()),
|
||||
local_port: 8899,
|
||||
my_peer_id: 2,
|
||||
smoltcp_enabled: false,
|
||||
},
|
||||
));
|
||||
let mut response = build_tcp_packet(local, translated_src, false, true);
|
||||
assert!(engine.try_process_packet_from_nic(&mut response, nic_ctx()));
|
||||
|
||||
let hdr: &PeerManagerHeader = response.peer_manager_header().unwrap();
|
||||
assert!(hdr.is_no_proxy());
|
||||
@@ -573,6 +709,111 @@ mod tests {
|
||||
let tcp = TcpPacket::new_checked(ipv4.payload()).unwrap();
|
||||
assert_eq!(tcp.src_port(), mapped_dst.port());
|
||||
assert_eq!(tcp.dst_port(), src.port());
|
||||
|
||||
engine.remove_entry(entry.id());
|
||||
assert!(!engine.is_tcp_proxy_flow(SocketAddr::V4(src), SocketAddr::V4(mapped_dst),));
|
||||
assert!(
|
||||
engine
|
||||
.translated_src_map
|
||||
.get(&SocketAddr::V4(translated_src))
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_source_port_to_mapped_and_real_destinations_stay_distinct() {
|
||||
let engine = tcp_engine();
|
||||
let src = SocketAddrV4::new("10.144.144.206".parse().unwrap(), 50000);
|
||||
let mapped_dst = SocketAddrV4::new("10.10.10.42".parse().unwrap(), 80);
|
||||
let real_dst = SocketAddrV4::new("127.0.0.42".parse().unwrap(), 80);
|
||||
|
||||
let mut mapped_request = build_tcp_packet(src, mapped_dst, true, false);
|
||||
assert_eq!(
|
||||
engine.try_handle_peer_packet(TcpProxyMode::Tcp, &mut mapped_request, peer_ctx()),
|
||||
TcpProxyPacketAction::Handled { new_syn: true }
|
||||
);
|
||||
let mapped_translated_src = packet_src(&mapped_request);
|
||||
engine
|
||||
.accept_connection(
|
||||
SocketAddr::V4(mapped_translated_src),
|
||||
Some("10.144.144.204/24".parse().unwrap()),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let mut real_request = build_tcp_packet(src, real_dst, true, false);
|
||||
real_request
|
||||
.mut_peer_manager_header()
|
||||
.unwrap()
|
||||
.set_exit_node(true);
|
||||
let mut real_ctx = peer_ctx();
|
||||
real_ctx.enable_exit_node = true;
|
||||
assert_eq!(
|
||||
engine.try_handle_peer_packet(TcpProxyMode::Tcp, &mut real_request, real_ctx),
|
||||
TcpProxyPacketAction::Handled { new_syn: true }
|
||||
);
|
||||
let real_translated_src = packet_src(&real_request);
|
||||
engine
|
||||
.accept_connection(
|
||||
SocketAddr::V4(real_translated_src),
|
||||
Some("10.144.144.204/24".parse().unwrap()),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_ne!(mapped_translated_src, real_translated_src);
|
||||
|
||||
let local = SocketAddrV4::new("10.144.144.204".parse().unwrap(), 8899);
|
||||
for (translated_src, expected_source) in [
|
||||
(mapped_translated_src, mapped_dst),
|
||||
(real_translated_src, real_dst),
|
||||
] {
|
||||
let mut response = build_tcp_packet(local, translated_src, false, true);
|
||||
assert!(engine.try_process_packet_from_nic(&mut response, nic_ctx()));
|
||||
|
||||
let ipv4 = Ipv4Packet::new_checked(response.payload()).unwrap();
|
||||
assert_eq!(ipv4.src_addr(), *expected_source.ip());
|
||||
assert_eq!(ipv4.dst_addr(), *src.ip());
|
||||
let tcp = TcpPacket::new_checked(ipv4.payload()).unwrap();
|
||||
assert_eq!(tcp.src_port(), expected_source.port());
|
||||
assert_eq!(tcp.dst_port(), src.port());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clear_discards_accepted_entries_before_restart() {
|
||||
let engine = tcp_engine();
|
||||
let src = SocketAddrV4::new("10.144.144.206".parse().unwrap(), 50000);
|
||||
let mapped_dst = SocketAddrV4::new("10.10.10.42".parse().unwrap(), 80);
|
||||
let mut request = build_tcp_packet(src, mapped_dst, true, false);
|
||||
assert_eq!(
|
||||
engine.try_handle_peer_packet(TcpProxyMode::Tcp, &mut request, peer_ctx()),
|
||||
TcpProxyPacketAction::Handled { new_syn: true }
|
||||
);
|
||||
let old_entry = engine
|
||||
.accept_connection(
|
||||
SocketAddr::V4(packet_src(&request)),
|
||||
Some("10.144.144.204/24".parse().unwrap()),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
engine.clear();
|
||||
|
||||
assert_eq!(old_entry.state(), TcpNatEntryState::Closed);
|
||||
assert!(engine.flow_map.is_empty());
|
||||
assert!(engine.translated_src_map.is_empty());
|
||||
assert!(engine.conn_map.is_empty());
|
||||
|
||||
let mut retry = build_tcp_packet(src, mapped_dst, true, false);
|
||||
assert_eq!(
|
||||
engine.try_handle_peer_packet(TcpProxyMode::Tcp, &mut retry, peer_ctx()),
|
||||
TcpProxyPacketAction::Handled { new_syn: true }
|
||||
);
|
||||
let new_entry = engine
|
||||
.accept_connection(
|
||||
SocketAddr::V4(packet_src(&retry)),
|
||||
Some("10.144.144.204/24".parse().unwrap()),
|
||||
)
|
||||
.unwrap();
|
||||
assert_ne!(old_entry.id(), new_entry.id());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -585,19 +826,264 @@ mod tests {
|
||||
engine.try_handle_peer_packet(TcpProxyMode::Tcp, &mut request, peer_ctx()),
|
||||
TcpProxyPacketAction::Handled { new_syn: true }
|
||||
));
|
||||
let entry = engine.syn_map.get(&SocketAddr::V4(src)).unwrap().clone();
|
||||
let translated_src = packet_src(&request);
|
||||
let entry = engine
|
||||
.translated_src_map
|
||||
.get(&SocketAddr::V4(translated_src))
|
||||
.unwrap()
|
||||
.clone();
|
||||
entry.set_state(TcpNatEntryState::Closed);
|
||||
|
||||
assert!(
|
||||
engine
|
||||
.accept_connection(
|
||||
SocketAddr::V4(src),
|
||||
SocketAddr::V4(translated_src),
|
||||
Some("10.144.144.204/24".parse().unwrap()),
|
||||
)
|
||||
.is_none()
|
||||
);
|
||||
assert!(engine.syn_map.get(&SocketAddr::V4(src)).is_none());
|
||||
assert!(engine.addr_conn_map.get(&SocketAddr::V4(src)).is_none());
|
||||
assert!(engine.flow_map.is_empty());
|
||||
assert!(engine.translated_src_map.is_empty());
|
||||
assert!(engine.conn_map.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn retransmitted_syn_keeps_the_pending_and_accepted_mapping() {
|
||||
let engine = tcp_engine();
|
||||
let src = "10.144.144.206:50000".parse().unwrap();
|
||||
let dst = "10.10.10.42:80".parse().unwrap();
|
||||
let mut first = build_tcp_packet(src, dst, true, false);
|
||||
assert_eq!(
|
||||
engine.try_handle_peer_packet(TcpProxyMode::Tcp, &mut first, peer_ctx()),
|
||||
TcpProxyPacketAction::Handled { new_syn: true }
|
||||
);
|
||||
let translated_src = SocketAddr::V4(packet_src(&first));
|
||||
|
||||
for accepted in [false, true] {
|
||||
if accepted {
|
||||
engine.accept_connection(translated_src, None).unwrap();
|
||||
engine.cleanup_expired_syn(Duration::ZERO);
|
||||
}
|
||||
let mut retry = build_tcp_packet(src, dst, true, false);
|
||||
assert_eq!(
|
||||
engine.try_handle_peer_packet(TcpProxyMode::Tcp, &mut retry, peer_ctx()),
|
||||
TcpProxyPacketAction::Handled { new_syn: false }
|
||||
);
|
||||
assert_eq!(SocketAddr::V4(packet_src(&retry)), translated_src);
|
||||
assert_eq!(engine.flow_map.len(), 1);
|
||||
assert_eq!(engine.translated_src_map.len(), 1);
|
||||
assert_eq!(engine.list_entries().len(), 1);
|
||||
}
|
||||
assert!(engine.accept_connection(translated_src, None).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn removing_closing_connection_preserves_replacement_flow() {
|
||||
for state in [
|
||||
TcpNatEntryState::ClosingSrc,
|
||||
TcpNatEntryState::ClosingDst,
|
||||
TcpNatEntryState::Closed,
|
||||
] {
|
||||
let engine = tcp_engine();
|
||||
let flow = TcpNatFlowKey {
|
||||
src: "10.144.144.206:50000".parse().unwrap(),
|
||||
mapped_dst: "10.10.10.42:80".parse().unwrap(),
|
||||
};
|
||||
let real_dst = "127.0.0.42:80".parse().unwrap();
|
||||
let (old, _) = engine.entry_for_syn(flow, real_dst).unwrap();
|
||||
engine.accept_connection(old.translated_src, None).unwrap();
|
||||
old.set_state(state);
|
||||
let (replacement, new_syn) = engine.entry_for_syn(flow, real_dst).unwrap();
|
||||
assert!(new_syn);
|
||||
assert_ne!(old.translated_src, replacement.translated_src);
|
||||
assert_eq!(engine.translated_src_map.len(), 2);
|
||||
|
||||
engine.remove_entry(old.id());
|
||||
assert!(engine.is_tcp_proxy_flow(flow.src, flow.mapped_dst));
|
||||
assert!(!engine.translated_src_map.contains_key(&old.translated_src));
|
||||
assert_eq!(engine.flow_map.get(&flow).unwrap().id(), replacement.id());
|
||||
assert_eq!(
|
||||
engine
|
||||
.accept_connection(replacement.translated_src, None)
|
||||
.unwrap()
|
||||
.id(),
|
||||
replacement.id()
|
||||
);
|
||||
|
||||
let local = "10.144.144.204:8899".parse().unwrap();
|
||||
let SocketAddr::V4(translated_src) = replacement.translated_src else {
|
||||
unreachable!();
|
||||
};
|
||||
let mut response = build_tcp_packet(local, translated_src, false, true);
|
||||
assert!(engine.try_process_packet_from_nic(&mut response, nic_ctx()));
|
||||
let ip = Ipv4Packet::new_checked(response.payload()).unwrap();
|
||||
let tcp = TcpPacket::new_checked(ip.payload()).unwrap();
|
||||
assert_eq!(SocketAddr::V4(packet_src(&response)), flow.mapped_dst);
|
||||
assert_eq!(tcp.dst_port(), flow.src.port());
|
||||
assert!(ip.verify_checksum());
|
||||
assert!(tcp.verify_checksum(
|
||||
&IpAddress::Ipv4(ip.src_addr()),
|
||||
&IpAddress::Ipv4(ip.dst_addr())
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expired_syn_releases_both_indices_without_removing_accepted_flow() {
|
||||
let engine = tcp_engine();
|
||||
let pending_flow = TcpNatFlowKey {
|
||||
src: "10.144.144.206:50000".parse().unwrap(),
|
||||
mapped_dst: "10.10.10.42:80".parse().unwrap(),
|
||||
};
|
||||
let accepted_flow = TcpNatFlowKey {
|
||||
mapped_dst: "10.10.10.43:80".parse().unwrap(),
|
||||
..pending_flow
|
||||
};
|
||||
let real_dst = "127.0.0.42:80".parse().unwrap();
|
||||
let (pending, _) = engine.entry_for_syn(pending_flow, real_dst).unwrap();
|
||||
let (accepted, _) = engine.entry_for_syn(accepted_flow, real_dst).unwrap();
|
||||
engine
|
||||
.accept_connection(accepted.translated_src, None)
|
||||
.unwrap();
|
||||
|
||||
engine.cleanup_expired_syn(Duration::ZERO);
|
||||
|
||||
assert_eq!(pending.state(), TcpNatEntryState::Closed);
|
||||
assert!(
|
||||
engine
|
||||
.accept_connection(pending.translated_src, None)
|
||||
.is_none()
|
||||
);
|
||||
assert!(!engine.is_tcp_proxy_flow(pending_flow.src, pending_flow.mapped_dst));
|
||||
assert!(engine.is_tcp_proxy_flow(accepted_flow.src, accepted_flow.mapped_dst));
|
||||
assert_eq!(engine.translated_src_map.len(), 1);
|
||||
assert_eq!(engine.conn_map.len(), 1);
|
||||
assert_eq!(accepted.state(), TcpNatEntryState::ConnectingDst);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn translated_port_wrap_skips_live_ports_and_listener() {
|
||||
let engine = tcp_engine();
|
||||
engine.set_local_port(8899);
|
||||
let flow = TcpNatFlowKey {
|
||||
src: "10.144.144.206:50000".parse().unwrap(),
|
||||
mapped_dst: "10.10.10.42:80".parse().unwrap(),
|
||||
};
|
||||
let real_dst = "127.0.0.42:80".parse().unwrap();
|
||||
engine
|
||||
.next_translated_port
|
||||
.store(u16::MAX, Ordering::Relaxed);
|
||||
let last = engine.allocate_entry(flow, real_dst).unwrap();
|
||||
let first = engine.allocate_entry(flow, real_dst).unwrap();
|
||||
assert_eq!(last.translated_src.port(), u16::MAX);
|
||||
assert_eq!(first.translated_src.port(), 1);
|
||||
|
||||
engine
|
||||
.next_translated_port
|
||||
.store(u16::MAX, Ordering::Relaxed);
|
||||
let next = engine.allocate_entry(flow, real_dst).unwrap();
|
||||
assert_eq!(next.translated_src.port(), 2);
|
||||
assert_eq!(
|
||||
engine
|
||||
.translated_src_map
|
||||
.get(&last.translated_src)
|
||||
.unwrap()
|
||||
.id(),
|
||||
last.id()
|
||||
);
|
||||
assert_eq!(
|
||||
engine
|
||||
.translated_src_map
|
||||
.get(&first.translated_src)
|
||||
.unwrap()
|
||||
.id(),
|
||||
first.id()
|
||||
);
|
||||
|
||||
engine.next_translated_port.store(8899, Ordering::Relaxed);
|
||||
let after_listener = engine.allocate_entry(flow, real_dst).unwrap();
|
||||
assert_eq!(after_listener.translated_src.port(), 8900);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exhausted_translated_ports_are_isolated_and_reusable_after_cleanup() {
|
||||
let engine = tcp_engine();
|
||||
engine.set_local_port(8899);
|
||||
let src_ip = "10.144.144.206".parse().unwrap();
|
||||
let dst = "10.10.10.42:80".parse().unwrap();
|
||||
let real_dst = "127.0.0.42:80".parse().unwrap();
|
||||
for port in 1..u16::MAX {
|
||||
let flow = TcpNatFlowKey {
|
||||
src: SocketAddr::V4(SocketAddrV4::new(src_ip, port)),
|
||||
mapped_dst: SocketAddr::V4(dst),
|
||||
};
|
||||
assert!(engine.entry_for_syn(flow, real_dst).is_some());
|
||||
}
|
||||
assert_eq!(engine.translated_src_map.len(), usize::from(u16::MAX) - 1);
|
||||
let src = SocketAddrV4::new(src_ip, u16::MAX);
|
||||
let mut exhausted = build_tcp_packet(src, dst, true, false);
|
||||
assert_eq!(
|
||||
engine.try_handle_peer_packet(TcpProxyMode::Tcp, &mut exhausted, peer_ctx()),
|
||||
TcpProxyPacketAction::Drop
|
||||
);
|
||||
assert!(!engine.is_tcp_proxy_flow(SocketAddr::V4(src), SocketAddr::V4(dst)));
|
||||
|
||||
let other_flow = TcpNatFlowKey {
|
||||
src: "10.144.144.207:50000".parse().unwrap(),
|
||||
mapped_dst: SocketAddr::V4(dst),
|
||||
};
|
||||
assert!(engine.entry_for_syn(other_flow, real_dst).is_some());
|
||||
|
||||
let released_addr = SocketAddr::V4(SocketAddrV4::new(src_ip, 1));
|
||||
let released = engine.accept_connection(released_addr, None).unwrap();
|
||||
released.set_state(TcpNatEntryState::Closed);
|
||||
engine.remove_entry(released.id());
|
||||
let mut retry = build_tcp_packet(src, dst, true, false);
|
||||
assert_eq!(
|
||||
engine.try_handle_peer_packet(TcpProxyMode::Tcp, &mut retry, peer_ctx()),
|
||||
TcpProxyPacketAction::Handled { new_syn: true }
|
||||
);
|
||||
assert_eq!(SocketAddr::V4(packet_src(&retry)), released_addr);
|
||||
assert!(engine.accept_connection(released_addr, None).is_some());
|
||||
|
||||
engine.clear();
|
||||
assert!(engine.flow_map.is_empty());
|
||||
assert!(engine.translated_src_map.is_empty());
|
||||
assert!(engine.conn_map.is_empty());
|
||||
assert!(engine.entry_for_syn(other_flow, real_dst).is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn concurrent_syns_share_one_mapping_and_only_one_accept() {
|
||||
let engine = tcp_engine();
|
||||
let barrier = std::sync::Barrier::new(8);
|
||||
let flow = TcpNatFlowKey {
|
||||
src: "10.144.144.206:50000".parse().unwrap(),
|
||||
mapped_dst: "10.10.10.42:80".parse().unwrap(),
|
||||
};
|
||||
let results = std::thread::scope(|scope| {
|
||||
let tasks: Vec<_> = (0..8)
|
||||
.map(|_| {
|
||||
scope.spawn(|| {
|
||||
barrier.wait();
|
||||
let (entry, new_syn) = engine.entry_for_syn(flow, flow.mapped_dst).unwrap();
|
||||
let accepted = engine
|
||||
.accept_connection(entry.translated_src, None)
|
||||
.is_some();
|
||||
(entry.id(), new_syn, accepted)
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
tasks
|
||||
.into_iter()
|
||||
.map(|task| task.join().unwrap())
|
||||
.collect::<Vec<_>>()
|
||||
});
|
||||
assert!(results.iter().all(|result| result.0 == results[0].0));
|
||||
assert_eq!(results.iter().filter(|result| result.1).count(), 1);
|
||||
assert_eq!(results.iter().filter(|result| result.2).count(), 1);
|
||||
assert_eq!(engine.flow_map.len(), 1);
|
||||
assert_eq!(engine.translated_src_map.len(), 1);
|
||||
assert_eq!(engine.conn_map.len(), 1);
|
||||
}
|
||||
}
|
||||
@@ -3,7 +3,8 @@ use std::sync::{Arc, Weak, atomic::Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use atomic_shim::AtomicU64;
|
||||
use tokio::io::{AsyncWriteExt, copy};
|
||||
use parking_lot::Mutex;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::task::JoinSet;
|
||||
|
||||
use crate::{
|
||||
@@ -34,10 +35,10 @@ use crate::gateway::smoltcp::{SmolTcpStack, output_dst_ip};
|
||||
fn spawn_tcp_proxy_task(
|
||||
lifecycle: &AtomicU64,
|
||||
expected_generation: u64,
|
||||
tasks: &std::sync::Mutex<JoinSet<()>>,
|
||||
tasks: &Mutex<JoinSet<()>>,
|
||||
task: impl Future<Output = ()> + Send + 'static,
|
||||
) -> bool {
|
||||
let mut tasks = tasks.lock().unwrap();
|
||||
let mut tasks = tasks.lock();
|
||||
if lifecycle.load(Ordering::Acquire) != expected_generation {
|
||||
return false;
|
||||
}
|
||||
@@ -57,12 +58,12 @@ pub struct TcpProxyService<
|
||||
connector: Arc<C>,
|
||||
engine: Arc<TcpProxyEngine>,
|
||||
mode: TcpProxyMode,
|
||||
peer_pipeline_guard: std::sync::Mutex<Option<PipelineRegistrationGuard>>,
|
||||
nic_pipeline_guard: std::sync::Mutex<Option<PipelineRegistrationGuard>>,
|
||||
kernel_listener: std::sync::Mutex<Option<Arc<F::Listener>>>,
|
||||
peer_pipeline_guard: Mutex<Option<PipelineRegistrationGuard>>,
|
||||
nic_pipeline_guard: Mutex<Option<PipelineRegistrationGuard>>,
|
||||
kernel_listener: Mutex<Option<Arc<F::Listener>>>,
|
||||
#[cfg(feature = "proxy-smoltcp-stack")]
|
||||
smoltcp_stack: std::sync::Mutex<Option<Arc<SmolTcpStack>>>,
|
||||
tasks: std::sync::Mutex<JoinSet<()>>,
|
||||
smoltcp_stack: Mutex<Option<Arc<SmolTcpStack>>>,
|
||||
tasks: Mutex<JoinSet<()>>,
|
||||
lifecycle: AtomicU64,
|
||||
}
|
||||
|
||||
@@ -86,12 +87,12 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
connector,
|
||||
engine: Arc::new(TcpProxyEngine::new(cidr_table)),
|
||||
mode,
|
||||
peer_pipeline_guard: std::sync::Mutex::new(None),
|
||||
nic_pipeline_guard: std::sync::Mutex::new(None),
|
||||
kernel_listener: std::sync::Mutex::new(None),
|
||||
peer_pipeline_guard: Mutex::new(None),
|
||||
nic_pipeline_guard: Mutex::new(None),
|
||||
kernel_listener: Mutex::new(None),
|
||||
#[cfg(feature = "proxy-smoltcp-stack")]
|
||||
smoltcp_stack: std::sync::Mutex::new(None),
|
||||
tasks: std::sync::Mutex::new(JoinSet::new()),
|
||||
smoltcp_stack: Mutex::new(None),
|
||||
tasks: Mutex::new(JoinSet::new()),
|
||||
lifecycle: AtomicU64::new(0),
|
||||
})
|
||||
}
|
||||
@@ -120,6 +121,8 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
}
|
||||
};
|
||||
|
||||
self.engine.clear();
|
||||
|
||||
let snapshot = self.runtime.proxy_runtime_snapshot();
|
||||
let start_result = if snapshot.smoltcp_enabled {
|
||||
#[cfg(feature = "proxy-smoltcp-stack")]
|
||||
@@ -157,13 +160,15 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
|
||||
pub fn stop(&self) {
|
||||
self.stop_resources();
|
||||
self.tasks.lock().unwrap().abort_all();
|
||||
self.tasks.lock().abort_all();
|
||||
self.engine.clear();
|
||||
}
|
||||
|
||||
pub(crate) async fn stop_and_wait(&self) {
|
||||
self.stop_resources();
|
||||
let mut tasks = std::mem::take(&mut *self.tasks.lock().unwrap());
|
||||
let mut tasks = std::mem::take(&mut *self.tasks.lock());
|
||||
tasks.shutdown().await;
|
||||
self.engine.clear();
|
||||
}
|
||||
|
||||
fn stop_resources(&self) {
|
||||
@@ -183,17 +188,17 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
break;
|
||||
}
|
||||
}
|
||||
if let Some(guard) = self.peer_pipeline_guard.lock().unwrap().take() {
|
||||
if let Some(guard) = self.peer_pipeline_guard.lock().take() {
|
||||
guard.close();
|
||||
}
|
||||
if let Some(guard) = self.nic_pipeline_guard.lock().unwrap().take() {
|
||||
if let Some(guard) = self.nic_pipeline_guard.lock().take() {
|
||||
guard.close();
|
||||
}
|
||||
if let Some(listener) = self.kernel_listener.lock().unwrap().take() {
|
||||
if let Some(listener) = self.kernel_listener.lock().take() {
|
||||
drop(listener);
|
||||
}
|
||||
#[cfg(feature = "proxy-smoltcp-stack")]
|
||||
if let Some(stack) = self.smoltcp_stack.lock().unwrap().take() {
|
||||
if let Some(stack) = self.smoltcp_stack.lock().take() {
|
||||
drop(stack);
|
||||
}
|
||||
}
|
||||
@@ -204,7 +209,7 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
}
|
||||
|
||||
pub async fn register_peer_pipeline(self: &Arc<Self>) {
|
||||
if self.peer_pipeline_guard.lock().unwrap().is_some() {
|
||||
if self.peer_pipeline_guard.lock().is_some() {
|
||||
return;
|
||||
}
|
||||
let peer_guard = self
|
||||
@@ -213,11 +218,11 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
service: Arc::downgrade(self),
|
||||
}))
|
||||
.await;
|
||||
self.peer_pipeline_guard.lock().unwrap().replace(peer_guard);
|
||||
self.peer_pipeline_guard.lock().replace(peer_guard);
|
||||
}
|
||||
|
||||
pub async fn register_nic_pipeline(self: &Arc<Self>) {
|
||||
if self.nic_pipeline_guard.lock().unwrap().is_some() {
|
||||
if self.nic_pipeline_guard.lock().is_some() {
|
||||
return;
|
||||
}
|
||||
let nic_guard = self
|
||||
@@ -226,7 +231,7 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
service: Arc::downgrade(self),
|
||||
}))
|
||||
.await;
|
||||
self.nic_pipeline_guard.lock().unwrap().replace(nic_guard);
|
||||
self.nic_pipeline_guard.lock().replace(nic_guard);
|
||||
}
|
||||
|
||||
fn spawn_syn_cleanup(self: &Arc<Self>, generation: u64) {
|
||||
@@ -244,7 +249,7 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
}
|
||||
|
||||
fn drain_completed_tasks(&self) {
|
||||
let mut tasks = self.tasks.lock().unwrap();
|
||||
let mut tasks = self.tasks.lock();
|
||||
while let Some(result) = tasks.try_join_next() {
|
||||
if let Err(err) = result {
|
||||
tracing::warn!(?err, "tcp proxy task finished with error");
|
||||
@@ -273,10 +278,7 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
)
|
||||
.await?;
|
||||
self.engine.set_local_port(listener.local_addr()?.port());
|
||||
self.kernel_listener
|
||||
.lock()
|
||||
.unwrap()
|
||||
.replace(listener.clone());
|
||||
self.kernel_listener.lock().replace(listener.clone());
|
||||
|
||||
let service = Arc::downgrade(self);
|
||||
let _ = spawn_tcp_proxy_task(&self.lifecycle, generation, &self.tasks, async move {
|
||||
@@ -309,7 +311,7 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
.local_inet
|
||||
.map(|inet| inet.address())
|
||||
.unwrap_or(std::net::Ipv4Addr::new(192, 88, 99, 254));
|
||||
let stack = SmolTcpStack::new(local_ip).await?;
|
||||
let stack = SmolTcpStack::new(local_ip, self.mode.smoltcp_listener_port()).await?;
|
||||
self.engine.set_local_port(stack.local_port());
|
||||
|
||||
let mut output_rx = stack.take_output_rx().await?;
|
||||
@@ -350,7 +352,7 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
}
|
||||
});
|
||||
|
||||
self.smoltcp_stack.lock().unwrap().replace(stack);
|
||||
self.smoltcp_stack.lock().replace(stack);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -429,7 +431,7 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
entry.set_state(TcpNatEntryState::Connected);
|
||||
}
|
||||
|
||||
let ret = copy_bidirectional_no_shutdown(src_stream.as_mut(), dst_stream.as_mut()).await;
|
||||
let ret = copy_bidirectional(src_stream.as_mut(), dst_stream.as_mut()).await;
|
||||
tracing::info!(nat_entry = ?entry, ret = ?ret, "nat tcp connection closed");
|
||||
|
||||
entry.set_state(TcpNatEntryState::ClosingSrc);
|
||||
@@ -462,9 +464,11 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
smoltcp_enabled: snapshot.smoltcp_enabled,
|
||||
},
|
||||
);
|
||||
let TcpProxyPacketAction::Handled { new_syn: _new_syn } = action else {
|
||||
return Some(packet);
|
||||
};
|
||||
match action {
|
||||
TcpProxyPacketAction::Handled { new_syn: _ } => {}
|
||||
TcpProxyPacketAction::Drop => return None,
|
||||
TcpProxyPacketAction::Pass => return Some(packet),
|
||||
}
|
||||
|
||||
if snapshot.smoltcp_enabled {
|
||||
#[cfg(feature = "proxy-smoltcp-stack")]
|
||||
@@ -486,7 +490,7 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
|
||||
#[cfg(feature = "proxy-smoltcp-stack")]
|
||||
async fn handle_smoltcp_packet(&self, packet: ZCPacket) {
|
||||
let stack = self.smoltcp_stack.lock().unwrap().clone();
|
||||
let stack = self.smoltcp_stack.lock().clone();
|
||||
let Some(stack) = stack else {
|
||||
tracing::error!("smoltcp stack is not started");
|
||||
return;
|
||||
@@ -510,24 +514,13 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
}
|
||||
}
|
||||
|
||||
async fn copy_bidirectional_no_shutdown(
|
||||
async fn copy_bidirectional(
|
||||
src: &mut dyn TcpProxyStream,
|
||||
dst: &mut dyn TcpProxyStream,
|
||||
) -> Result<(), ProxyRuntimeError> {
|
||||
let (mut src_reader, mut src_writer) = tokio::io::split(src);
|
||||
let (mut dst_reader, mut dst_writer) = tokio::io::split(dst);
|
||||
let src_to_dst = copy(&mut src_reader, &mut dst_writer);
|
||||
let dst_to_src = copy(&mut dst_reader, &mut src_writer);
|
||||
tokio::pin!(src_to_dst);
|
||||
tokio::pin!(dst_to_src);
|
||||
tokio::select! {
|
||||
result = &mut src_to_dst => {
|
||||
result?;
|
||||
}
|
||||
result = &mut dst_to_src => {
|
||||
result?;
|
||||
}
|
||||
}
|
||||
// Forward EOF to the opposite writer while continuing to relay its reply.
|
||||
// Returning on the first EOF would discard responses to half-closed requests.
|
||||
tokio::io::copy_bidirectional(src, dst).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -575,6 +568,57 @@ impl<R: TcpProxyRuntime + 'static, F: VirtualTcpListenerFactory, C: TcpProxyDest
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use tokio::io::AsyncReadExt;
|
||||
|
||||
async fn response_after_half_close(source_closes_first: bool, request_size: usize) {
|
||||
// Small buffers force the relay to make progress in both directions;
|
||||
// the response is deliberately larger than either relay buffer.
|
||||
let (mut client, mut src) = tokio::io::duplex(64);
|
||||
let (mut dst, mut server) = tokio::io::duplex(64);
|
||||
let request = vec![0x35; request_size];
|
||||
let response = vec![0xa7; 64 * 1024];
|
||||
let relay = async {
|
||||
if source_closes_first {
|
||||
copy_bidirectional(&mut src, &mut dst).await
|
||||
} else {
|
||||
copy_bidirectional(&mut dst, &mut src).await
|
||||
}
|
||||
.unwrap();
|
||||
};
|
||||
let requester = async {
|
||||
client.write_all(&request).await.unwrap();
|
||||
client.shutdown().await.unwrap();
|
||||
let mut received = Vec::new();
|
||||
client.read_to_end(&mut received).await.unwrap();
|
||||
assert_eq!(received, response);
|
||||
};
|
||||
let responder = async {
|
||||
let mut received = Vec::new();
|
||||
server.read_to_end(&mut received).await.unwrap();
|
||||
assert_eq!(received, request);
|
||||
server.write_all(&response).await.unwrap();
|
||||
server.shutdown().await.unwrap();
|
||||
};
|
||||
timeout(Duration::from_secs(5), async {
|
||||
tokio::join!(relay, requester, responder);
|
||||
})
|
||||
.await
|
||||
.expect("half-closed relay did not finish");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn source_half_close_preserves_response() {
|
||||
for request_size in [0, 32 * 1024] {
|
||||
response_after_half_close(true, request_size).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn destination_half_close_preserves_response() {
|
||||
for request_size in [0, 32 * 1024] {
|
||||
response_after_half_close(false, request_size).await;
|
||||
}
|
||||
}
|
||||
|
||||
struct DropSignal(Arc<AtomicBool>);
|
||||
|
||||
@@ -595,7 +639,7 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn stop_fence_linearizes_task_registration() {
|
||||
let lifecycle = AtomicU64::new(1);
|
||||
let tasks = std::sync::Mutex::new(JoinSet::new());
|
||||
let tasks = Mutex::new(JoinSet::new());
|
||||
let accepted_dropped = Arc::new(AtomicBool::new(false));
|
||||
|
||||
assert!(spawn_tcp_proxy_task(
|
||||
@@ -605,7 +649,7 @@ mod tests {
|
||||
pending_task(accepted_dropped.clone()),
|
||||
));
|
||||
lifecycle.store(2, Ordering::Release);
|
||||
let mut stopping = std::mem::take(&mut *tasks.lock().unwrap());
|
||||
let mut stopping = std::mem::take(&mut *tasks.lock());
|
||||
stopping.shutdown().await;
|
||||
assert!(accepted_dropped.load(Ordering::Acquire));
|
||||
|
||||
@@ -627,7 +671,7 @@ mod tests {
|
||||
&tasks,
|
||||
pending_task(current_dropped.clone()),
|
||||
));
|
||||
let mut current = std::mem::take(&mut *tasks.lock().unwrap());
|
||||
let mut current = std::mem::take(&mut *tasks.lock());
|
||||
current.shutdown().await;
|
||||
assert!(current_dropped.load(Ordering::Acquire));
|
||||
}
|
||||
|
||||
@@ -128,11 +128,11 @@ pub struct WrappedTcpProxyNicContext {
|
||||
pub async fn try_process_wrapped_tcp_packet_from_nic<ConnectionLookup, AllowCheck, AllowCheckFut>(
|
||||
zc_packet: &mut ZCPacket,
|
||||
ctx: WrappedTcpProxyNicContext,
|
||||
is_tcp_proxy_connection: ConnectionLookup,
|
||||
is_tcp_proxy_flow: ConnectionLookup,
|
||||
check_dst_allowed: AllowCheck,
|
||||
) -> bool
|
||||
where
|
||||
ConnectionLookup: Fn(SocketAddr) -> bool,
|
||||
ConnectionLookup: Fn(SocketAddr, SocketAddr) -> bool,
|
||||
AllowCheck: FnOnce(Ipv4Addr) -> AllowCheckFut,
|
||||
AllowCheckFut: Future<Output = bool>,
|
||||
{
|
||||
@@ -156,6 +156,7 @@ where
|
||||
let src_ip = ip_packet.src_addr();
|
||||
let dst_ip = ip_packet.dst_addr();
|
||||
let src_port = tcp_packet.src_port();
|
||||
let dst_port = tcp_packet.dst_port();
|
||||
let is_syn = tcp_packet.syn() && !tcp_packet.ack();
|
||||
|
||||
if is_syn {
|
||||
@@ -171,7 +172,10 @@ where
|
||||
);
|
||||
return false;
|
||||
}
|
||||
} else if !is_tcp_proxy_connection(SocketAddr::V4(SocketAddrV4::new(src_ip, src_port))) {
|
||||
} else if !is_tcp_proxy_flow(
|
||||
SocketAddr::V4(SocketAddrV4::new(src_ip, src_port)),
|
||||
SocketAddr::V4(SocketAddrV4::new(dst_ip, dst_port)),
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -366,7 +370,7 @@ mod tests {
|
||||
try_process_wrapped_tcp_packet_from_nic(
|
||||
&mut packet,
|
||||
context(WrappedTcpProxyTransport::Kcp),
|
||||
|_| false,
|
||||
|_, _| false,
|
||||
|_| async { true },
|
||||
)
|
||||
.await
|
||||
@@ -387,7 +391,7 @@ mod tests {
|
||||
!try_process_wrapped_tcp_packet_from_nic(
|
||||
&mut packet,
|
||||
context(WrappedTcpProxyTransport::Kcp),
|
||||
|_| false,
|
||||
|_, _| false,
|
||||
|_| async { false },
|
||||
)
|
||||
.await
|
||||
@@ -407,7 +411,9 @@ mod tests {
|
||||
try_process_wrapped_tcp_packet_from_nic(
|
||||
&mut packet,
|
||||
context(WrappedTcpProxyTransport::Quic),
|
||||
|addr| addr == SocketAddr::V4(src),
|
||||
|src_addr, dst_addr| {
|
||||
src_addr == SocketAddr::V4(src) && dst_addr == SocketAddr::V4(dst)
|
||||
},
|
||||
|_| async { false },
|
||||
)
|
||||
.await
|
||||
@@ -428,7 +434,7 @@ mod tests {
|
||||
!try_process_wrapped_tcp_packet_from_nic(
|
||||
&mut packet,
|
||||
context(WrappedTcpProxyTransport::Quic),
|
||||
|_| false,
|
||||
|_, _| false,
|
||||
|_| async { true },
|
||||
)
|
||||
.await
|
||||
@@ -445,7 +451,7 @@ mod tests {
|
||||
!try_process_wrapped_tcp_packet_from_nic(
|
||||
&mut packet,
|
||||
context(WrappedTcpProxyTransport::Kcp),
|
||||
|_| false,
|
||||
|_, _| false,
|
||||
|_| async { true },
|
||||
)
|
||||
.await
|
||||
|
||||
@@ -260,7 +260,7 @@ where
|
||||
local_ipv4: snapshot.virtual_ipv4,
|
||||
smoltcp_enabled: snapshot.smoltcp_enabled,
|
||||
},
|
||||
move |src| connection_engine.is_tcp_proxy_connection(src),
|
||||
move |src, mapped_dst| connection_engine.is_tcp_proxy_flow(src, mapped_dst),
|
||||
move |dst_ip| async move {
|
||||
match transport {
|
||||
WrappedTransportKind::Kcp => {
|
||||
|
||||
@@ -11,6 +11,7 @@ use crate::gateway::proxy::traits::TcpProxyStream;
|
||||
|
||||
pub struct SmolTcpStack {
|
||||
ingress_tx: mpsc::Sender<ZCPacket>,
|
||||
local_port: u16,
|
||||
output_rx: Mutex<Option<mpsc::Receiver<Vec<u8>>>>,
|
||||
listener: Mutex<TcpListener>,
|
||||
_net: Net,
|
||||
@@ -18,7 +19,7 @@ pub struct SmolTcpStack {
|
||||
}
|
||||
|
||||
impl SmolTcpStack {
|
||||
pub async fn new(local_ip: Ipv4Addr) -> anyhow::Result<Arc<Self>> {
|
||||
pub async fn new(local_ip: Ipv4Addr, local_port: u16) -> anyhow::Result<Arc<Self>> {
|
||||
let tasks = Arc::new(std::sync::Mutex::new(JoinSet::new()));
|
||||
let mut cap = smoltcp::phy::DeviceCapabilities::default();
|
||||
cap.max_transmission_unit = 1280;
|
||||
@@ -63,12 +64,13 @@ impl SmolTcpStack {
|
||||
);
|
||||
net.set_any_ip(true);
|
||||
let listener = net
|
||||
.tcp_bind("0.0.0.0:8899".parse().unwrap())
|
||||
.tcp_bind(SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), local_port))
|
||||
.await
|
||||
.map_err(|error| anyhow::anyhow!("bind smoltcp listener failed: {error}"))?;
|
||||
|
||||
Ok(Arc::new(Self {
|
||||
ingress_tx,
|
||||
local_port,
|
||||
output_rx: Mutex::new(Some(stack_stream)),
|
||||
listener: Mutex::new(listener),
|
||||
_net: net,
|
||||
@@ -77,7 +79,7 @@ impl SmolTcpStack {
|
||||
}
|
||||
|
||||
pub fn local_port(&self) -> u16 {
|
||||
8899
|
||||
self.local_port
|
||||
}
|
||||
|
||||
pub async fn send_ingress(&self, packet: ZCPacket) -> anyhow::Result<()> {
|
||||
@@ -137,12 +139,12 @@ mod tests {
|
||||
};
|
||||
|
||||
const LOCAL_ADDR: Ipv4Address = Ipv4Address::new(192, 88, 99, 254);
|
||||
const LOCAL_PORT: u16 = 8899;
|
||||
const LOCAL_PORT: u16 = 8900;
|
||||
const PACKETS: TcpPackets = TcpPackets::new(LOCAL_ADDR, LOCAL_PORT);
|
||||
|
||||
#[tokio::test]
|
||||
async fn accepts_concurrent_connections_with_one_logical_listener() {
|
||||
let stack = SmolTcpStack::new(LOCAL_ADDR).await.unwrap();
|
||||
let stack = SmolTcpStack::new(LOCAL_ADDR, LOCAL_PORT).await.unwrap();
|
||||
let mut output = stack.take_output_rx().await.unwrap();
|
||||
let client_addr = Ipv4Address::new(192, 88, 99, 1);
|
||||
|
||||
|
||||
Reference in new issue
Block a user