From 7592951b808c4839cb9d77cc4bc642f124915951 Mon Sep 17 00:00:00 2001 From: lightfirst Date: Sun, 11 Oct 2026 01:35:33 +0800 Subject: [PATCH] fix(dataplane): route UDP by bound local endpoints (#2652) Register one UDP flow at bind so sockets can receive from new peers before their first send. Allocate and reserve ports within each instance independently of TUN and Host sockets. Release the endpoint on close. Search the ephemeral range once and preserve concurrent cursor progress. Route local UDP fragments by the first fragment's bound endpoint. Buffer earlier tails in a bounded shallow cache and forward original packets to one recipient. Preserve cached ownership across binding changes, expire idle entries, and clear state on address changes and runtime shutdown. This avoids orphan fragments exhausting smoltcp's reassembly slots while keeping original packet sizes for TUN and proxy consumers. Randomize each new stack's seed to avoid replaying IP IDs on rebuild. Document accepted residual collision-related UDP loss or downstream delivery instead of adding fragment completion tracking. Allow peer filters to return a batch, resuming each packet at the next filter without repeating earlier filters. The single-packet path does not allocate a batch queue. Keep ZCPacket and IpReassembler unchanged. Cover endpoint ownership, ephemeral allocation, out-of-order fragments, cache limits and expiry, binding changes, batch continuation, and large UDP datagrams after more than 16 complete unrelated fragmented datagrams. Exercise two recreated sending stacks with controlled seeds to verify that new fragments do not inherit an old downstream route. Co-authored-by: KKRainbow <5665404+KKRainbow@users.noreply.github.com> --- docs/data-plane-runtime-plan.md | 25 +- easytier-core/src/gateway/dataplane/flow.rs | 23 +- easytier-core/src/gateway/dataplane/mod.rs | 124 +++-- easytier-core/src/gateway/dataplane/packet.rs | 119 +++-- easytier-core/src/gateway/dataplane/stack.rs | 10 +- easytier-core/src/gateway/dataplane/tests.rs | 500 +++++++++++++----- easytier-core/src/gateway/dataplane/udp.rs | 31 +- .../src/gateway/dataplane/udp_fragments.rs | 381 +++++++++++++ .../src/gateway/proxy/icmp_proxy_service.rs | 8 +- .../src/gateway/proxy/tcp_proxy_service.rs | 8 +- .../src/gateway/proxy/udp_proxy_service.rs | 8 +- .../src/gateway/proxy/wrapped_transport.rs | 22 +- .../src/gateway/smoltcp/tokio_smoltcp/mod.rs | 35 +- easytier-core/src/peers/mod.rs | 21 +- easytier-core/src/peers/peer_manager.rs | 108 +++- 15 files changed, 1103 insertions(+), 320 deletions(-) create mode 100644 easytier-core/src/gateway/dataplane/udp_fragments.rs diff --git a/docs/data-plane-runtime-plan.md b/docs/data-plane-runtime-plan.md index 947c09b8..7cd1a45a 100644 --- a/docs/data-plane-runtime-plan.md +++ b/docs/data-plane-runtime-plan.md @@ -290,7 +290,11 @@ Each resource owns explicit leases: lease through the backend connection wrapper; - a listener owns its wildcard listener lease; - an accepted stream owns an independent exact flow lease; -- a UDP socket owns the destination flow leases it created; +- a UDP socket owns one bound local-endpoint lease, independent of its remote + destinations; an unconnected socket can receive before its first send; +- UDP binds allocate and reserve ports within the instance, independently of + TUN and Host sockets. Bound endpoints are owned by the data plane; unmatched + traffic continues through the peer-packet pipeline; - a direct or local stream reports its actual address and does not fabricate a smoltcp address. @@ -305,12 +309,25 @@ flow table using pointer equality. Packet classification behavior must remain compatible: -- malformed or unmatched packets pass to the next pipeline; +- malformed or unmatched nonfragmented packets pass to the next pipeline; - exact TCP flows win over wildcard listeners; - only supported EasyTier data packet kinds are consumed; - modified-source packets that are not local loopback traffic pass through; -- fragmented UDP response routing remains supported; -- a packet is consumed only when the owning flow exists. +- local UDP fragments follow the route selected by the first fragment's + destination port. Earlier tails wait for that decision, then continue as + original packets, without mirroring or changing packet sizes; +- a shallow cache retains up to 1,024 routes and 1 MiB of pending packet + bytes, with a one-second idle timeout. Expiry and LRU eviction discard + pending fragments; they never fall back to a different route. Each new + data-plane smoltcp stack uses a random seed to avoid replaying its IP ID + sequence. Residual 16-bit ID collisions can still reuse a cached route, + causing UDP loss or unintended downstream delivery. This tradeoff is + accepted to avoid tracking fragment ranges and completion. Cache hits + refresh the idle timeout, so repeated collisions can prolong the effect; +- a closed UDP binding cannot transfer its cached fragments to a new binding. + Address changes and runtime shutdown clear the cache; +- released batches resume at the next filter in the current pipeline snapshot; +- nonfragmented packets are consumed only when the owning flow exists. `DataPlaneRuntime`, not the SOCKS5 Adapter, registers and owns the peer-packet pipeline. diff --git a/easytier-core/src/gateway/dataplane/flow.rs b/easytier-core/src/gateway/dataplane/flow.rs index 8ed78ff3..d386e630 100644 --- a/easytier-core/src/gateway/dataplane/flow.rs +++ b/easytier-core/src/gateway/dataplane/flow.rs @@ -1,7 +1,7 @@ //! Shared data-plane flow registration and ownership. use std::{ - net::{IpAddr, SocketAddr}, + net::SocketAddr, sync::{ Arc, atomic::{AtomicUsize, Ordering}, @@ -26,6 +26,16 @@ pub(crate) struct FlowKey { pub kind: FlowKind, } +impl FlowKey { + pub(super) fn udp_bind(local_addr: SocketAddr) -> Self { + Self { + src: local_addr, + dst: SocketAddr::from(([0, 0, 0, 0], 0)), + kind: FlowKind::Udp, + } + } +} + #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(crate) struct FlowCountChange { pub previous: usize, @@ -136,14 +146,12 @@ impl FlowTable { self.entries.is_empty() } - pub fn contains_key(&self, entry: &FlowKey) -> bool { - self.entries.contains_key(entry) + pub(super) fn registration(&self, entry: &FlowKey) -> Option { + self.entries.get(entry).map(|flow| flow.registration) } - pub fn contains_destination_ip(&self, destination: IpAddr) -> bool { - self.entries - .iter() - .any(|entry| entry.key().dst.ip() == destination) + pub fn contains_key(&self, entry: &FlowKey) -> bool { + self.entries.contains_key(entry) } #[cfg(test)] @@ -376,7 +384,6 @@ mod tests { assert!(!table.try_insert(first.clone(), 2)); assert!(table.try_insert(second.clone(), 3)); assert_eq!(table.count(), 2); - assert!(table.contains_destination_ip(first.dst.ip())); let retained = table.retain(|entry, _| entry == &second); assert_eq!(retained.removed, 1); diff --git a/easytier-core/src/gateway/dataplane/mod.rs b/easytier-core/src/gateway/dataplane/mod.rs index 81597af1..b8b2a8d7 100644 --- a/easytier-core/src/gateway/dataplane/mod.rs +++ b/easytier-core/src/gateway/dataplane/mod.rs @@ -42,7 +42,7 @@ use crate::{ }, packet::{PacketType, ZCPacket}, peers::{ - PeerPacketFilter, + PeerPacketFilter, PeerPacketFilterResult, peer_manager::{PeerManagerCore, PipelineRegistrationGuard}, }, socket::{ @@ -52,7 +52,7 @@ use crate::{ VirtualTcpListener, VirtualTcpListenerFactory, VirtualTcpSocket, VirtualTcpSocketFactory, }, - udp::{UdpBindOptions, VirtualUdpSocket, VirtualUdpSocketFactory}, + udp::VirtualUdpSocketFactory, }, }; @@ -69,6 +69,9 @@ mod tcp; #[cfg(test)] mod tests; mod udp; +mod udp_fragments; + +use udp_fragments::UdpFragments; use self::{ deadline::{DataPlaneDeadline, DataPlaneIoDeadline}, @@ -169,7 +172,6 @@ pub(super) enum FlowData { Udp, } -const UDP_ENTRY: FlowKind = FlowKind::Udp; const TCP_ENTRY: FlowKind = FlowKind::Tcp; const TCP_LISTEN_ENTRY: FlowKind = FlowKind::TcpListen; @@ -195,6 +197,7 @@ where net: Arc>>, pub(super) entries: FlowSet, + udp_fragments: Arc>, data_plane_consumers: Arc, // Tracks whether the smoltcp `net` is ready for data-plane callers. @@ -207,8 +210,8 @@ impl PeerPacketFilter for DataPlaneRuntime where H: VirtualTcpSocketFactory + VirtualTcpListenerFactory + VirtualUdpSocketFactory, { - async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> Option { - if self.entries.is_idle() { + async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> PeerPacketFilterResult { + if self.entries.is_idle() && !packet::is_udp_fragment(&packet) { if tracing::enabled!(tracing::Level::TRACE) && let Some(hdr) = packet.peer_manager_header() && matches!( @@ -257,11 +260,11 @@ where ); } } - return Some(packet); + return PeerPacketFilterResult::Pass(packet); } let route = self.entries.route_peer_packet(&packet, true); let (entry_key, tcp_flags) = match route { - PeerPacketRoute::Pass => return Some(packet), + PeerPacketRoute::Pass => return PeerPacketFilterResult::Pass(packet), PeerPacketRoute::Unmatched { entry, tcp_flags } => { tracing::trace!( entry_key = ?entry, @@ -271,34 +274,20 @@ where entry_count = self.entries.count(), "data plane has no flow for packet from peer" ); - return Some(packet); + return PeerPacketFilterResult::Pass(packet); } PeerPacketRoute::Deliver { entry, tcp_flags } => (entry, tcp_flags), - PeerPacketRoute::FragmentedUdp { source, mirror } => { - let source: IpAddr = source.into(); - tracing::trace!( - is_in_entries = mirror, - "ipv4 src = {:?}, check need send both smoltcp and kernel tun", - source - ); - if mirror { - // 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. - match self.packet_sender.try_send(packet.clone()) { - Ok(()) => tracing::trace!( - ?source, - entry_count = self.entries.count(), - "data plane delivered fragmented packet from peer to smoltcp" - ), - Err(err) => tracing::trace!( - ?source, - ?err, - entry_count = self.entries.count(), - "data plane failed to deliver fragmented packet from peer to smoltcp" - ), - } + PeerPacketRoute::FragmentedUdp { destination } => { + let local_ip = Self::runtime_ipv4(&self.runtime_config).map(|ip| ip.address()); + if Some(destination) != local_ip { + return PeerPacketFilterResult::Pass(packet); } - return Some(packet); + return self.udp_fragments.lock().unwrap().process( + packet, + local_ip, + &self.entries, + &self.packet_sender, + ); } }; @@ -327,7 +316,7 @@ where ), } - None + PeerPacketFilterResult::Consumed } } @@ -359,6 +348,7 @@ where net: Arc::new(Mutex::new(None)), entries: Arc::new(FlowTable::default()), + udp_fragments: Arc::new(std::sync::Mutex::new(UdpFragments::default())), data_plane_consumers: Arc::new(DataPlaneConsumers::new()), data_plane_net_ready: tokio::sync::watch::channel(false).0, @@ -392,12 +382,21 @@ where let peer_manager = self.peer_manager.clone(); let packet_recv = self.packet_recv.clone(); let entries = self.entries.clone(); + let udp_fragments = self.udp_fragments.clone(); let data_plane_consumers = self.data_plane_consumers.clone(); let data_plane_net_ready = self.data_plane_net_ready.clone(); self.runtime_tasks.lock().unwrap().spawn(async move { let mut prev_ipv4 = None; + let mut fragment_gc = tokio::time::interval(Duration::from_secs(1)); let mut peer_changes = runtime_config.subscribe_peer_runtime_changes(); loop { + // Run housekeeping regardless of which select branch woke us. + let cur_ipv4 = Self::runtime_ipv4(&runtime_config); + { + let mut fragments = udp_fragments.lock().unwrap(); + fragments.set_local_ip(cur_ipv4.map(|ip| ip.address())); + fragments.remove_expired(); + } let data_plane_active = data_plane_consumers.has_consumers(); if !data_plane_active { @@ -429,12 +428,12 @@ where let _ = data_plane_net_ready.send_replace(false); select! { _ = peer_changes.changed() => {} + _ = fragment_gc.tick() => {} _ = data_plane_consumers.changed() => {} } continue; } - let cur_ipv4 = Self::runtime_ipv4(&runtime_config); if prev_ipv4 != cur_ipv4 { let old_ipv4 = prev_ipv4; prev_ipv4 = cur_ipv4; @@ -487,6 +486,7 @@ where select! { _ = peer_changes.changed() => {} + _ = fragment_gc.tick() => {} _ = data_plane_consumers.changed() => {} } } @@ -549,6 +549,7 @@ where let _ = self.data_plane_net_ready.send_replace(false); self.entries.clear(); Self::shutdown_tasks(&self.runtime_tasks).await; + *self.udp_fragments.lock().unwrap() = UdpFragments::default(); } pub(crate) async fn stop_runtime(&self) { @@ -943,19 +944,44 @@ where ) -> DataPlaneResult { let data_plane_ref = self.acquire_data_plane_ref()?; let (ipv4_addr, smoltcp_net, generation) = self.wait_data_plane_net(deadline).await?; - let reservation_options = UdpBindOptions::port_lease(SocketAddr::new( - IpAddr::V4(Ipv4Addr::UNSPECIFIED), - local_port, - )) - .with_context(self.socket_context.clone()); - let reservation = generation - .while_open(deadline.run(self.host.bind_udp(reservation_options))) - .await?; - let reserved_port = reservation - .local_addr() - .map_err(DataPlaneError::from)? - .port(); - let bind_addr = SocketAddr::new(IpAddr::V4(ipv4_addr.address()), reserved_port); + + // The bound endpoint is also the port lease. Reserve it before creating + // the smoltcp socket, which does not check other sockets for duplicates. + let try_register = |port| { + let local_addr = SocketAddr::new(ipv4_addr.address().into(), port); + FlowLease::try_register( + self.entries.clone(), + FlowKey::udp_bind(local_addr), + FlowData::Udp, + ) + .map(|flow| (local_addr, flow)) + }; + let (bind_addr, bind_flow) = if local_port != 0 { + generation.ensure_open()?; + deadline.remaining()?; + try_register(local_port).ok_or_else(|| { + DataPlaneError::new( + DataPlaneErrorKind::AddressInUse, + format!("data-plane UDP port {local_port} is already bound"), + ) + })? + } else { + let mut binding = None; + for port in smoltcp_net.ephemeral_ports() { + generation.ensure_open()?; + deadline.remaining()?; + binding = try_register(port); + if binding.is_some() { + break; + } + } + binding.ok_or_else(|| { + DataPlaneError::new( + DataPlaneErrorKind::ResourceLimit, + "no free data-plane UDP ports", + ) + })? + }; let smol = generation .while_open(deadline.run(smoltcp_net.udp_bind(bind_addr))) .await?; @@ -965,10 +991,8 @@ where Ok(DataPlaneUdpSocket { socket, - flows: self.entries.clone(), - routes: std::sync::Mutex::new(std::collections::HashMap::new()), + _bind_flow: bind_flow, local_addr, - _reservation: reservation, _data_plane_lease: data_plane_ref, generation, }) diff --git a/easytier-core/src/gateway/dataplane/packet.rs b/easytier-core/src/gateway/dataplane/packet.rs index d122b634..623b78b7 100644 --- a/easytier-core/src/gateway/dataplane/packet.rs +++ b/easytier-core/src/gateway/dataplane/packet.rs @@ -4,10 +4,7 @@ use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use smoltcp::wire::{IPV4_HEADER_LEN, IpProtocol, Ipv4Packet, TcpPacket, UdpPacket}; -use crate::{ - gateway::proxy::ip_reassembler::IpReassembler, - packet::{PacketType, ZCPacket}, -}; +use crate::packet::{PacketType, ZCPacket}; use super::flow::{FlowKey, FlowKind, FlowTable}; @@ -22,7 +19,7 @@ enum ClassifiedPeerPacket { entry: FlowKey, }, FragmentedUdp { - source: Ipv4Addr, + destination: Ipv4Addr, }, Unsupported, } @@ -38,8 +35,7 @@ pub(crate) enum PeerPacketRoute { tcp_flags: Option, }, FragmentedUdp { - source: Ipv4Addr, - mirror: bool, + destination: Ipv4Addr, }, } fn classify_peer_ipv4_payload(payload: &[u8]) -> ClassifiedPeerPacket { @@ -72,9 +68,9 @@ fn classify_peer_ipv4_payload(payload: &[u8]) -> ClassifiedPeerPacket { } } IpProtocol::Udp => { - if IpReassembler::is_packet_fragmented(&ipv4) { + if ipv4.frag_offset() != 0 || ipv4.more_frags() { return ClassifiedPeerPacket::FragmentedUdp { - source: ipv4.src_addr(), + destination: ipv4.dst_addr(), }; } let Ok(udp) = UdpPacket::new_checked(ipv4.payload()) else { @@ -92,6 +88,34 @@ fn classify_peer_ipv4_payload(payload: &[u8]) -> ClassifiedPeerPacket { } } +pub(super) fn is_udp_fragment(packet: &ZCPacket) -> bool { + is_data_plane_packet(packet) + && Ipv4Packet::new_checked(packet.payload()).is_ok_and(|ip| { + ip.version() == 4 + && usize::from(ip.header_len()) >= IPV4_HEADER_LEN + && ip.next_header() == IpProtocol::Udp + && (ip.frag_offset() != 0 || ip.more_frags()) + }) +} + +fn is_data_plane_packet(packet: &ZCPacket) -> bool { + let Some(header) = packet.peer_manager_header() else { + return false; + }; + let is_modified_source = matches!( + header.packet_type, + x if x == PacketType::DataWithKcpSrcModified as u8 + || x == PacketType::DataWithQuicSrcModified as u8 + ); + if header.packet_type != PacketType::Data as u8 && !is_modified_source { + return false; + } + if is_modified_source && header.from_peer_id != header.to_peer_id { + return false; + } + true +} + pub(super) fn tcp_flags>(tcp: &TcpPacket) -> u8 { u8::from(tcp.fin()) | (u8::from(tcp.syn()) << 1) @@ -108,18 +132,7 @@ impl FlowTable { packet: &ZCPacket, allow_tcp_listen_fallback: bool, ) -> PeerPacketRoute { - let Some(header) = packet.peer_manager_header() else { - return PeerPacketRoute::Pass; - }; - let is_modified_source = matches!( - header.packet_type, - x if x == PacketType::DataWithKcpSrcModified as u8 - || x == PacketType::DataWithQuicSrcModified as u8 - ); - if header.packet_type != PacketType::Data as u8 && !is_modified_source { - return PeerPacketRoute::Pass; - } - if is_modified_source && header.from_peer_id != header.to_peer_id { + if !is_data_plane_packet(packet) { return PeerPacketRoute::Pass; } @@ -144,12 +157,9 @@ impl FlowTable { }; (entry, Some(flags)) } - ClassifiedPeerPacket::Udp { entry } => (entry, None), - ClassifiedPeerPacket::FragmentedUdp { source } => { - return PeerPacketRoute::FragmentedUdp { - source, - mirror: self.contains_destination_ip(source.into()), - }; + ClassifiedPeerPacket::Udp { entry } => (FlowKey::udp_bind(entry.src), None), + ClassifiedPeerPacket::FragmentedUdp { destination } => { + return PeerPacketRoute::FragmentedUdp { destination }; } ClassifiedPeerPacket::Unsupported => return PeerPacketRoute::Pass, }; @@ -234,7 +244,7 @@ mod tests { assert_eq!( classify_peer_ipv4_payload(&fragmented), ClassifiedPeerPacket::FragmentedUdp { - source: Ipv4Addr::new(10, 1, 1, 2), + destination: Ipv4Addr::new(10, 2, 2, 3), } ); } @@ -314,34 +324,51 @@ mod tests { ); } #[test] - fn flow_table_routes_fragmented_udp_by_source_ip() { + fn flow_table_routes_udp_by_bound_port() { let mut packet = ipv4_packet(IpProtocol::Udp, 8); - Ipv4Packet::new_unchecked(&mut packet).set_frag_offset(8); + let mut ipv4 = Ipv4Packet::new_unchecked(&mut packet); + let mut udp = UdpPacket::new_unchecked(ipv4.payload_mut()); + udp.set_src_port(1234); + udp.set_dst_port(4321); + udp.set_len(8); let table = FlowTable::default(); - + let entry = FlowKey::udp_bind("10.2.2.3:4321".parse().unwrap()); assert_eq!( table.route_peer_ipv4_payload(&packet, false), - PeerPacketRoute::FragmentedUdp { - source: Ipv4Addr::new(10, 1, 1, 2), - mirror: false, + PeerPacketRoute::Unmatched { + entry: entry.clone(), + tcp_flags: None } ); - - table.insert( - FlowKey { - src: "10.2.2.3:4321".parse().unwrap(), - dst: "10.1.1.2:1234".parse().unwrap(), - kind: FlowKind::Udp, - }, - (), - ); + table.insert(entry.clone(), ()); assert_eq!( table.route_peer_ipv4_payload(&packet, false), - PeerPacketRoute::FragmentedUdp { - source: Ipv4Addr::new(10, 1, 1, 2), - mirror: true, + PeerPacketRoute::Deliver { + entry, + tcp_flags: None } ); + for more_frags in [false, true] { + let mut ipv4 = Ipv4Packet::new_unchecked(&mut packet); + ipv4.set_frag_offset(8); + ipv4.set_more_frags(more_frags); + assert_eq!( + table.route_peer_ipv4_payload(&packet, false), + PeerPacketRoute::FragmentedUdp { + destination: Ipv4Addr::new(10, 2, 2, 3), + } + ); + } + for len in 0..smoltcp::wire::UDP_HEADER_LEN { + let mut short = ipv4_packet(IpProtocol::Udp, len); + Ipv4Packet::new_unchecked(&mut short).set_more_frags(true); + assert_eq!( + table.route_peer_ipv4_payload(&short, false), + PeerPacketRoute::FragmentedUdp { + destination: Ipv4Addr::new(10, 2, 2, 3) + } + ); + } } #[test] fn flow_table_routes_loopback_modified_source_packets() { diff --git a/easytier-core/src/gateway/dataplane/stack.rs b/easytier-core/src/gateway/dataplane/stack.rs index 19b664cb..61e5b2bc 100644 --- a/easytier-core/src/gateway/dataplane/stack.rs +++ b/easytier-core/src/gateway/dataplane/stack.rs @@ -77,7 +77,9 @@ impl SmoltcpPlane { tracing::debug!("smoltcp-to-peer bridge stopped"); }); - let interface_config = smoltcp::iface::Config::new(smoltcp::wire::HardwareAddress::Ip); + // A fixed seed replays IP IDs after the last consumer closes and the + // stack is recreated, potentially reusing a remote shallow-cache route. + let interface_config = Self::interface_config(rand::random()); let net = Net::new( device, NetConfig::new( @@ -117,6 +119,12 @@ impl SmoltcpPlane { } } + pub(super) fn interface_config(random_seed: u64) -> smoltcp::iface::Config { + let mut config = smoltcp::iface::Config::new(smoltcp::wire::HardwareAddress::Ip); + config.random_seed = random_seed; + config + } + pub(super) fn lease(&self) -> DataPlaneIoGuard { self.generation.clone() } diff --git a/easytier-core/src/gateway/dataplane/tests.rs b/easytier-core/src/gateway/dataplane/tests.rs index 40b01715..250bd00e 100644 --- a/easytier-core/src/gateway/dataplane/tests.rs +++ b/easytier-core/src/gateway/dataplane/tests.rs @@ -1,6 +1,6 @@ use std::net::{IpAddr, Ipv4Addr, SocketAddr}; -use smoltcp::wire::{IpAddress, IpProtocol, Ipv4Packet, TcpPacket}; +use smoltcp::wire::{IpAddress, IpProtocol, Ipv4Packet, TcpPacket, UdpPacket}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use super::*; @@ -36,6 +36,7 @@ fn test_gateway() -> Arc> { packet_recv: Arc::new(Mutex::new(packet_recv)), net: Arc::new(Mutex::new(None)), entries: Arc::new(FlowTable::default()), + udp_fragments: Arc::new(std::sync::Mutex::new(UdpFragments::default())), data_plane_consumers: Arc::new(DataPlaneConsumers::new()), data_plane_net_ready: tokio::sync::watch::channel(false).0, pipeline_guard: Mutex::new(None), @@ -196,26 +197,45 @@ fn build_tcp_packet(src: SocketAddr, dst: SocketAddr) -> Vec { buf } -fn build_udp_followup_fragment(src: Ipv4Addr, dst: Ipv4Addr) -> Vec { - let mut buf = vec![0u8; 28]; - { - let mut ip_packet = Ipv4Packet::new_unchecked(&mut buf); - ip_packet.set_version(4); - ip_packet.set_header_len(20); - ip_packet.set_total_len(28); - ip_packet.set_hop_limit(64); - ip_packet.set_next_header(IpProtocol::Udp); - ip_packet.set_frag_offset(8); - ip_packet.set_src_addr(src); - ip_packet.set_dst_addr(dst); - ip_packet - .payload_mut() - .copy_from_slice(&[0xde, 0xad, 0xbe, 0xef, 0xca, 0xfe, 0xba, 0xbe]); +pub(super) const UDP_LOCAL_IP: Ipv4Addr = Ipv4Addr::new(10, 144, 144, 1); +const UDP_LOCAL: SocketAddr = SocketAddr::new(IpAddr::V4(UDP_LOCAL_IP), 40000); +const UDP_REMOTE: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 144, 144, 3)), 53); - ip_packet.fill_checksum(); +pub(super) fn build_udp_fragments() -> [ZCPacket; 2] { + let src_ip = Ipv4Addr::new(10, 144, 144, 3); + let dst_ip = UDP_LOCAL_IP; + let mut udp = [0; 24]; + { + let mut packet = UdpPacket::new_unchecked(&mut udp[..]); + packet.set_src_port(UDP_REMOTE.port()); + packet.set_dst_port(UDP_LOCAL.port()); + packet.set_len(24); + packet.payload_mut().copy_from_slice(b"fragmented-first"); + packet.fill_checksum(&IpAddress::Ipv4(src_ip), &IpAddress::Ipv4(dst_ip)); } - buf + [(&udp[..16], 0, true), (&udp[16..], 16, false)].map(|(payload, offset, more_fragments)| { + let mut buf = vec![0; 20 + payload.len()]; + let mut ip = Ipv4Packet::new_unchecked(&mut buf); + ip.set_version(4); + ip.set_header_len(20); + ip.set_total_len((20 + payload.len()) as u16); + ip.set_ident(42); + ip.set_hop_limit(64); + ip.set_next_header(IpProtocol::Udp); + ip.set_src_addr(src_ip); + ip.set_dst_addr(dst_ip); + ip.set_frag_offset(offset); + ip.set_more_frags(more_fragments); + ip.payload_mut().copy_from_slice(payload); + ip.fill_checksum(); + let mut packet = ZCPacket::new_with_payload(&buf); + packet.fill_peer_manager_hdr(1, 2, PacketType::Data as u8); + let header = packet.mut_peer_manager_header().unwrap(); + header.set_latency_first(true).set_no_proxy(true); + header.forward_counter = 3; + packet + }) } #[tokio::test] @@ -440,18 +460,6 @@ async fn data_plane_sessions_report_udp_truncation() { .unwrap() .unwrap(); - let warmup = session_b - .submit_udp_send(socket_b, addr_a, b"warmup".to_vec()) - .unwrap(); - wait_for_session_completion(&session_b).await; - session_b - .take_result_with(warmup, |outcome| match outcome { - Ok(DataPlaneOperationResult::UdpSent { len }) => Some(*len), - _ => None, - }) - .unwrap() - .unwrap(); - let receive = session_b.submit_udp_receive(socket_b, 2).unwrap(); let send = session_a .submit_udp_send(socket_a, addr_b, b"ping".to_vec()) @@ -541,37 +549,39 @@ async fn listener_and_accepted_stream_own_independent_flow_lifetimes() { stop_data_plane_pair(&a, &b).await; } +async fn send_udp(sender: &DataPlaneUdpSocket, receiver: &DataPlaneUdpSocket, data: &[u8]) { + sender.send_to(data, receiver.local_addr()).await.unwrap(); + let mut buf = vec![0; data.len()]; + let (len, from) = tokio::time::timeout(Duration::from_secs(10), receiver.recv_from(&mut buf)) + .await + .expect("UDP receive timed out") + .unwrap(); + assert_eq!(&buf[..len], data); + assert_eq!(from, sender.local_addr()); +} + #[tokio::test] async fn data_plane_udp_pingpong() { let (a, b) = setup_data_plane_pair().await; let timeout = Duration::from_secs(10); - let socket_a = a.gateway.data_plane_udp_bind(0, timeout).await.unwrap(); - let socket_b = b.gateway.data_plane_udp_bind(0, timeout).await.unwrap(); - let addr_a = SocketAddr::new(a.ip.address().into(), socket_a.local_addr().port()); - let addr_b = SocketAddr::new(b.ip.address().into(), socket_b.local_addr().port()); - - socket_b.send_to(b"warmup", addr_a).await.unwrap(); - socket_a.send_to(b"ping", addr_b).await.unwrap(); - let mut buf = [0u8; 16]; - let (len, from) = tokio::time::timeout(timeout, socket_b.recv_from(&mut buf)) - .await - .expect("receive ping timed out") - .unwrap(); - assert_eq!(&buf[..len], b"ping"); - assert_eq!(from, addr_a); - - socket_b.send_to(b"pong", addr_a).await.unwrap(); - loop { - let (len, from) = tokio::time::timeout(timeout, socket_a.recv_from(&mut buf)) - .await - .expect("receive pong timed out") - .unwrap(); - if &buf[..len] == b"pong" { - assert_eq!(from, addr_b); - break; - } + let first = a.gateway.data_plane_udp_bind(0, timeout).await.unwrap(); + let second = a.gateway.data_plane_udp_bind(0, timeout).await.unwrap(); + let receiver = b.gateway.data_plane_udp_bind(0, timeout).await.unwrap(); + assert_ne!(first.local_addr(), second.local_addr()); + for sender in [&first, &second] { + assert_ne!(sender.local_addr().port(), 0); + assert_eq!(sender.local_addr().ip(), IpAddr::V4(a.ip.address())); + send_udp(sender, &receiver, b"ping").await; + send_udp(&receiver, sender, b"pong").await; } - + // Sending to different destinations must not create per-remote flows. + assert_eq!(a.gateway.entries.count(), 2); + assert_eq!(b.gateway.entries.count(), 1); + assert_eq!(a.gateway.host.udp_binds.load(Ordering::Relaxed), 0); + assert_eq!(b.gateway.host.udp_binds.load(Ordering::Relaxed), 0); + drop((first, second, receiver)); + assert_eq!(a.gateway.entries.count(), 0); + assert_eq!(b.gateway.entries.count(), 0); stop_data_plane_pair(&a, &b).await; } @@ -581,60 +591,125 @@ async fn data_plane_udp_carries_maximum_ipv4_payload() { let timeout = Duration::from_secs(10); let socket_a = a.gateway.data_plane_udp_bind(0, timeout).await.unwrap(); let socket_b = b.gateway.data_plane_udp_bind(0, timeout).await.unwrap(); - let addr_a = SocketAddr::new(a.ip.address().into(), socket_a.local_addr().port()); - let addr_b = SocketAddr::new(b.ip.address().into(), socket_b.local_addr().port()); - - socket_b.send_to(b"warmup", addr_a).await.unwrap(); - let payload = vec![0x5a; 65_507]; - socket_a.send_to(&payload, addr_b).await.unwrap(); - let mut received = vec![0; payload.len()]; - let (len, from) = tokio::time::timeout(timeout, socket_b.recv_from(&mut received)) - .await - .expect("receive maximum UDP payload timed out") - .unwrap(); - assert_eq!(from, addr_a); - assert_eq!(len, payload.len()); - assert_eq!(received, payload); - + send_udp(&socket_a, &socket_b, &vec![0x5a; 65_507]).await; stop_data_plane_pair(&a, &b).await; } #[tokio::test] -async fn udp_socket_drop_releases_every_destination_flow() { +async fn udp_fixed_ports_are_owned_by_each_instance_until_socket_drop() { let (a, b) = setup_data_plane_pair().await; let timeout = Duration::from_secs(10); - let socket = a.gateway.data_plane_udp_bind(0, timeout).await.unwrap(); - let first = SocketAddr::new(b.ip.address().into(), 31001); - let second = SocketAddr::new(b.ip.address().into(), 31002); - - socket.send_to(b"one", first).await.unwrap(); - socket.send_to(b"two", second).await.unwrap(); - assert_eq!(a.gateway.entries.count(), 2); - - drop(socket); + let first = a.gateway.data_plane_udp_bind(48000, timeout).await.unwrap(); + let second = b.gateway.data_plane_udp_bind(48000, timeout).await.unwrap(); + let error = a + .gateway + .data_plane_udp_bind(48000, timeout) + .await + .err() + .unwrap(); + assert_eq!(error.kind(), DataPlaneErrorKind::AddressInUse); + assert_eq!(a.gateway.entries.count(), 1); + send_udp(&first, &second, b"same-port").await; + let addr = first.local_addr(); + drop(first); assert_eq!(a.gateway.entries.count(), 0); - + let rebound = a.gateway.data_plane_udp_bind(48000, timeout).await.unwrap(); + assert_eq!(rebound.local_addr(), addr); + send_udp(&second, &rebound, b"rebound").await; stop_data_plane_pair(&a, &b).await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn concurrent_udp_binds_have_exactly_one_owner() { + let host = Arc::new(TestHost::default()); + let endpoint = data_plane_endpoint(host, "10.126.132.1/24".parse().unwrap()); + endpoint.peer_manager.run().await.unwrap(); + endpoint.gateway.start_runtime().await.unwrap(); + let barrier = Arc::new(tokio::sync::Barrier::new(8)); + let mut tasks = JoinSet::new(); + for _ in 0..8 { + let gateway = endpoint.gateway.clone(); + let barrier = barrier.clone(); + tasks.spawn(async move { + barrier.wait().await; + gateway + .data_plane_udp_bind(48000, Duration::from_secs(10)) + .await + }); + } + let mut sockets = Vec::new(); + while let Some(result) = tasks.join_next().await { + match result.unwrap() { + Ok(socket) => sockets.push(socket), + Err(error) => assert_eq!(error.kind(), DataPlaneErrorKind::AddressInUse), + } + } + assert_eq!(sockets.len(), 1); + assert_eq!(endpoint.gateway.entries.count(), 1); + assert_eq!(endpoint.gateway.host.udp_binds.load(Ordering::Relaxed), 0); + drop(sockets); + assert_eq!(endpoint.gateway.entries.count(), 0); + endpoint.gateway.stop_runtime().await; + endpoint.peer_manager.clear_resources().await; +} + #[tokio::test] -async fn udp_socket_owns_a_host_port_reservation() { +async fn udp_ephemeral_allocation_advances_past_collisions() { + let endpoint = data_plane_endpoint( + Arc::new(TestHost::default()), + "10.126.132.1/24".parse().unwrap(), + ); + endpoint.peer_manager.run().await.unwrap(); + endpoint.gateway.start_runtime().await.unwrap(); + let gateway = &endpoint.gateway; + let timeout = Duration::from_secs(10); + let fixed = gateway.data_plane_udp_bind(11000, timeout).await.unwrap(); + for port in 10001..11000 { + let addr = SocketAddr::new(endpoint.ip.address().into(), port); + gateway + .entries + .insert(FlowKey::udp_bind(addr), FlowData::Udp); + } + let socket = gateway.data_plane_udp_bind(0, timeout).await.unwrap(); + assert_eq!(socket.local_addr().port(), 11001); + let net = gateway.net.lock().await.as_ref().unwrap().net.clone(); + assert_eq!(net.get_port(), 11002); + let next = gateway.data_plane_udp_bind(0, timeout).await.unwrap(); + assert_eq!(next.local_addr().port(), 11003); + + let mut slow = net.ephemeral_ports(); + let last = net.ephemeral_ports().nth(10).unwrap(); + slow.nth(2); + assert_eq!(net.get_port(), last + 1); + drop((fixed, socket, next)); + gateway.stop_runtime().await; + endpoint.peer_manager.clear_resources().await; +} + +#[tokio::test] +async fn udp_ephemeral_bind_wraps_and_reports_exhaustion() { let host = Arc::new(TestHost::default()); let endpoint = data_plane_endpoint(host.clone(), "10.126.132.1/24".parse().unwrap()); endpoint.peer_manager.run().await.unwrap(); endpoint.gateway.start_runtime().await.unwrap(); + let gateway = &endpoint.gateway; + let timeout = Duration::from_secs(10); + let fixed = gateway.data_plane_udp_bind(60001, timeout).await.unwrap(); + // Leave only the port reached after wrapping past the end of the range. + for port in 10001..60001 { + let addr = SocketAddr::new(endpoint.ip.address().into(), port); + gateway + .entries + .insert(FlowKey::udp_bind(addr), FlowData::Udp); + } + let socket = gateway.data_plane_udp_bind(0, timeout).await.unwrap(); + assert_eq!(socket.local_addr().port(), 10000); + let error = gateway.data_plane_udp_bind(0, timeout).await.err().unwrap(); + assert_eq!(error.kind(), DataPlaneErrorKind::ResourceLimit); + assert_eq!(host.udp_binds.load(Ordering::Relaxed), 0); - let socket = endpoint - .gateway - .data_plane_udp_bind(0, Duration::from_secs(1)) - .await - .unwrap(); - - assert_eq!(socket.local_addr().port(), 20002); - assert_eq!(host.udp_binds.load(Ordering::Relaxed), 1); - - drop(socket); - endpoint.gateway.stop_runtime().await; + drop((fixed, socket)); + gateway.stop_runtime().await; endpoint.peer_manager.clear_resources().await; } @@ -833,7 +908,7 @@ async fn data_plane_consumes_modified_data_when_entry_matches() { packet.fill_peer_manager_hdr(1, 1, packet_type as u8); let result = gateway.try_process_packet_from_peer(packet).await; - assert!(result.is_none()); + assert!(matches!(result, PeerPacketFilterResult::Consumed)); let mut receiver = gateway.packet_recv.lock().await; let received = receiver.try_recv().unwrap(); @@ -864,12 +939,12 @@ async fn data_plane_passes_through_unmatched_or_malformed_modified_data() { ZCPacket::new_with_payload(&build_tcp_packet(remote, unmatched_local)); unmatched_packet.fill_peer_manager_hdr(1, 2, PacketType::DataWithKcpSrcModified as u8); let result = gateway.try_process_packet_from_peer(unmatched_packet).await; - assert!(result.is_some()); + assert!(matches!(result, PeerPacketFilterResult::Pass(_))); let mut malformed_packet = ZCPacket::new_with_payload(&[0u8; 8]); malformed_packet.fill_peer_manager_hdr(1, 2, PacketType::DataWithQuicSrcModified as u8); let result = gateway.try_process_packet_from_peer(malformed_packet).await; - assert!(result.is_some()); + assert!(matches!(result, PeerPacketFilterResult::Pass(_))); let mut receiver = gateway.packet_recv.lock().await; assert!(receiver.try_recv().is_err()); @@ -897,47 +972,204 @@ async fn data_plane_passes_through_non_loopback_modified_data_when_entry_matches packet.fill_peer_manager_hdr(1, 2, PacketType::DataWithKcpSrcModified as u8); let result = gateway.try_process_packet_from_peer(packet).await; - assert!(result.is_some()); + assert!(matches!(result, PeerPacketFilterResult::Pass(_))); let mut receiver = gateway.packet_recv.lock().await; assert!(receiver.try_recv().is_err()); } #[tokio::test] -async fn data_plane_mirrors_fragmented_udp_when_entry_matches() { - let gateway = test_gateway(); - - 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); - gateway.entries.insert( - FlowKey { - src: local, - dst: remote, - kind: UDP_ENTRY, - }, - FlowData::Udp, - ); - assert_eq!(gateway.entries.count(), 1); - - 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 = gateway.try_process_packet_from_peer(packet).await; - assert!(result.is_some()); - - let mut receiver = gateway.packet_recv.lock().await; - let received = receiver.try_recv().unwrap(); - assert_eq!( - received.peer_manager_header().unwrap().packet_type, - PacketType::Data as u8 - ); +async fn data_plane_routes_udp_fragments_together() { + for port in [None, Some(UDP_LOCAL.port() + 1), Some(UDP_LOCAL.port())] { + for reverse in [false, true] { + let gateway = test_gateway(); + gateway.runtime_config.update_peer_with(|peer| { + peer.runtime.core.routes.ipv4 = + Some(IpPrefix::new(UDP_LOCAL_IP.into(), 24).unwrap()); + }); + if let Some(port) = port { + gateway.entries.insert( + FlowKey::udp_bind(SocketAddr::new(UDP_LOCAL_IP.into(), port)), + FlowData::Udp, + ); + } + let mut downstream = Vec::new(); + let mut delivered = Vec::new(); + let mut expected = Vec::new(); + // More than smoltcp's 16 slots, with every tail arriving before any first. + for index in if reverse { [1, 0] } else { [0, 1] } { + for id in 0..32 { + let mut packet = build_udp_fragments().into_iter().nth(index).unwrap(); + let mut ip = Ipv4Packet::new_unchecked(packet.mut_payload()); + ip.set_ident(id); + ip.fill_checksum(); + expected.push(packet.clone().into_bytes()); + match gateway.try_process_packet_from_peer(packet).await { + PeerPacketFilterResult::Consumed => {} + PeerPacketFilterResult::Pass(packet) => { + downstream.push(packet.into_bytes()) + } + PeerPacketFilterResult::PassBatch(packets) => { + downstream.extend(packets.into_iter().map(ZCPacket::into_bytes)) + } + } + let mut receiver = gateway.packet_recv.lock().await; + while let Ok(packet) = receiver.try_recv() { + delivered.push(packet.into_bytes()); + } + } + } + expected.sort(); + downstream.sort(); + delivered.sort(); + if port == Some(UDP_LOCAL.port()) { + assert_eq!(delivered, expected); + assert!(downstream.is_empty()); + } else { + assert_eq!(downstream, expected); + assert!(delivered.is_empty()); + } + } + } +} + +#[tokio::test] +async fn data_plane_only_buffers_eligible_local_udp_fragments() { + let gateway = test_gateway(); + gateway.runtime_config.update_peer_with(|peer| { + peer.runtime.core.routes.ipv4 = Some(IpPrefix::new(UDP_LOCAL_IP.into(), 24).unwrap()); + }); + gateway + .entries + .insert(FlowKey::udp_bind(UDP_LOCAL), FlowData::Udp); + for (destination, packet_type, to_peer, buffered) in [ + (UDP_LOCAL_IP, PacketType::Data, 2, true), + (Ipv4Addr::new(10, 144, 144, 9), PacketType::Data, 2, false), + (UDP_LOCAL_IP, PacketType::DataWithKcpSrcModified, 1, true), + (UDP_LOCAL_IP, PacketType::DataWithKcpSrcModified, 2, false), + (UDP_LOCAL_IP, PacketType::DataWithQuicSrcModified, 1, true), + (UDP_LOCAL_IP, PacketType::DataWithQuicSrcModified, 2, false), + ] { + let [_, mut tail] = build_udp_fragments(); + let mut ip = Ipv4Packet::new_unchecked(tail.mut_payload()); + ip.set_dst_addr(destination); + ip.fill_checksum(); + tail.fill_peer_manager_hdr(1, to_peer, packet_type as u8); + let expected = tail.clone().into_bytes(); + let result = gateway.try_process_packet_from_peer(tail).await; + if buffered { + assert!(matches!(result, PeerPacketFilterResult::Consumed)); + } else { + let PeerPacketFilterResult::Pass(packet) = result else { + panic!("unrelated fragment buffered") + }; + assert_eq!(packet.into_bytes(), expected); + } + assert!(gateway.packet_recv.lock().await.try_recv().is_err()); + } +} + +#[tokio::test] +async fn data_plane_udp_receives_fragments_after_unbound_datagrams() { + let endpoint = data_plane_endpoint( + Arc::new(TestHost::default()), + "10.144.144.1/24".parse().unwrap(), + ); + endpoint.peer_manager.run().await.unwrap(); + endpoint.gateway.start_runtime().await.unwrap(); + let timeout = Duration::from_secs(10); + let socket = endpoint + .gateway + .data_plane_udp_bind(UDP_LOCAL.port(), timeout) + .await + .unwrap(); + // Unrelated complete datagrams must never leave orphan tails in smoltcp. + for id in 0..32 { + for mut packet in build_udp_fragments() { + let mut ip = Ipv4Packet::new_unchecked(packet.mut_payload()); + ip.set_ident(id); + if ip.frag_offset() == 0 { + let mut udp = UdpPacket::new_unchecked(ip.payload_mut()); + udp.set_dst_port(UDP_LOCAL.port() + 1); + udp.set_checksum(0); + } + ip.fill_checksum(); + assert!(matches!( + endpoint.gateway.try_process_packet_from_peer(packet).await, + PeerPacketFilterResult::Pass(_) + )); + } + } + let [first, tail] = build_udp_fragments(); + for packet in [tail, first] { + assert!(matches!( + endpoint.gateway.try_process_packet_from_peer(packet).await, + PeerPacketFilterResult::Consumed + )); + } + let mut data = [0; 32]; + let (len, from) = tokio::time::timeout(timeout, socket.recv_from(&mut data)) + .await + .unwrap() + .unwrap(); + assert_eq!(&data[..len], b"fragmented-first"); + assert_eq!(from, UDP_REMOTE); + // A datagram larger than either MTU remains a set of original small packets. + for port in [UDP_LOCAL.port(), UDP_LOCAL.port() + 1] { + let mut udp = vec![0; 8 + 4096]; + let mut udp_packet = UdpPacket::new_unchecked(&mut udp); + udp_packet.set_src_port(UDP_REMOTE.port()); + udp_packet.set_dst_port(port); + udp_packet.set_len(4104); + udp_packet.payload_mut().fill(0xab); + udp_packet.fill_checksum(&UDP_REMOTE.ip().into(), &UDP_LOCAL.ip().into()); + let [template, _] = build_udp_fragments(); + let mut fragments = Vec::new(); + for (index, payload) in udp.chunks(1200).enumerate() { + let mut bytes = template.payload()[..20].to_vec(); + bytes.extend_from_slice(payload); + let mut ip = Ipv4Packet::new_unchecked(&mut bytes); + ip.set_ident(port); + ip.set_total_len((20 + payload.len()) as u16); + ip.set_frag_offset((index * 1200) as u16); + ip.set_more_frags((index + 1) * 1200 < udp.len()); + ip.fill_checksum(); + let mut packet = ZCPacket::new_with_payload(&bytes); + packet.fill_peer_manager_hdr(1, 2, PacketType::Data as u8); + fragments.push(packet); + } + let mut expected = fragments + .iter() + .map(|p| p.clone().into_bytes()) + .collect::>(); + let mut passed = Vec::new(); + for packet in fragments.into_iter().rev() { + match endpoint.gateway.try_process_packet_from_peer(packet).await { + PeerPacketFilterResult::Consumed => {} + PeerPacketFilterResult::Pass(_) => panic!("tails should wait for the first"), + PeerPacketFilterResult::PassBatch(packets) => { + assert!(packets.iter().all(|packet| packet.payload_len() <= 1280)); + passed.extend(packets.into_iter().map(ZCPacket::into_bytes)); + } + } + } + if port == UDP_LOCAL.port() { + assert!(passed.is_empty()); + let mut data = vec![0; 4096]; + let (len, from) = tokio::time::timeout(timeout, socket.recv_from(&mut data)) + .await + .unwrap() + .unwrap(); + assert_eq!(len, 4096); + assert_eq!(data, vec![0xab; 4096]); + assert_eq!(from, UDP_REMOTE); + } else { + expected.sort(); + passed.sort(); + assert_eq!(passed, expected); + } + } + drop(socket); + endpoint.gateway.stop_runtime().await; + endpoint.peer_manager.clear_resources().await; } diff --git a/easytier-core/src/gateway/dataplane/udp.rs b/easytier-core/src/gateway/dataplane/udp.rs index 03e5e782..0c658bda 100644 --- a/easytier-core/src/gateway/dataplane/udp.rs +++ b/easytier-core/src/gateway/dataplane/udp.rs @@ -1,23 +1,13 @@ //! UDP socket resources exposed by the data plane. -use std::{ - any::Any, - collections::{HashMap, hash_map::Entry}, - net::SocketAddr, - sync::{Arc, Mutex}, -}; +use std::{net::SocketAddr, sync::Arc}; -use super::{ - DataPlaneIoGuard, DataPlaneLease, DataPlaneUdpIo, FlowData, FlowKey, FlowLease, FlowSet, - UDP_ENTRY, -}; +use super::{DataPlaneIoGuard, DataPlaneLease, DataPlaneUdpIo, FlowData, FlowLease}; pub struct DataPlaneUdpSocket { pub(super) socket: Arc, - pub(super) flows: FlowSet, - pub(super) routes: Mutex>>, + pub(super) _bind_flow: FlowLease, pub(super) local_addr: SocketAddr, - pub(super) _reservation: Arc, pub(super) _data_plane_lease: DataPlaneLease, pub(super) generation: DataPlaneIoGuard, } @@ -31,21 +21,6 @@ impl DataPlaneUdpSocket { self.generation .ensure_open() .map_err(|error| error.into_io_error())?; - let key = FlowKey { - src: self.local_addr, - dst: addr, - kind: UDP_ENTRY, - }; - if let Entry::Vacant(route) = self.routes.lock().unwrap().entry(addr) { - let lease = FlowLease::try_register(self.flows.clone(), key, FlowData::Udp) - .ok_or_else(|| { - std::io::Error::new( - std::io::ErrorKind::AddrInUse, - "data-plane UDP flow already exists", - ) - })?; - route.insert(lease); - } tokio::select! { biased; _ = self.generation.closed() => Err(self.generation.closed_io_error()), diff --git a/easytier-core/src/gateway/dataplane/udp_fragments.rs b/easytier-core/src/gateway/dataplane/udp_fragments.rs new file mode 100644 index 00000000..6fd0e556 --- /dev/null +++ b/easytier-core/src/gateway/dataplane/udp_fragments.rs @@ -0,0 +1,381 @@ +//! Route original UDP fragments together without assembling their payloads. + +use std::{ + net::{Ipv4Addr, SocketAddr}, + time::Duration, +}; + +use ordered_hash_map::OrderedHashMap; +use smoltcp::wire::{Ipv4Packet, UDP_HEADER_LEN, UdpPacket}; +use tokio::{sync::mpsc, time::Instant}; + +use super::{FlowKey, FlowSet, PeerPacketFilterResult}; +use crate::packet::ZCPacket; + +#[derive(Clone, Copy, Hash, PartialEq, Eq)] +struct UdpFragmentKey { + source: Ipv4Addr, + destination: Ipv4Addr, + id: u16, + from_peer: u32, + to_peer: u32, + packet_type: u8, + exit_node: bool, + no_proxy: bool, + not_send_to_tun: bool, +} + +struct UdpFragmentRoute { + entry: FlowKey, + registration: Option, +} + +struct UdpFragment { + last_seen: Instant, + route: Option, + pending: Vec, + bytes: usize, +} + +impl UdpFragment { + fn new(now: Instant) -> Self { + Self { + last_seen: now, + route: None, + pending: Vec::new(), + bytes: 0, + } + } +} + +#[derive(Default)] +pub(super) struct UdpFragments { + local_ip: Option, + packets: OrderedHashMap, + bytes: usize, +} + +impl UdpFragments { + const CAPACITY: usize = 1024; + // Budget retained packet bytes, not allocator overhead. + const MAX_BYTES: usize = 1024 * 1024; + // A shallow cache cannot detect completion. Randomizing each sending stack's + // seed avoids deterministic IP ID reuse after a restart, but 16-bit IDs can + // still collide. We accept occasional UDP loss or delivery to the cached + // downstream route instead of tracking fragment ranges to detect completion. + // Hits refresh this idle timeout, so repeated collisions can prolong it. + // Expired/evicted unknown tails wait for a first fragment. + const TIMEOUT: Duration = Duration::from_secs(1); + + pub(super) fn set_local_ip(&mut self, local_ip: Option) { + if self.local_ip != local_ip { + *self = Self { + local_ip, + ..Self::default() + }; + } + } + + fn evict_oldest(&mut self) { + if let Some(packet) = self.packets.pop_front() { + self.bytes -= packet.bytes; + } + } + + pub(super) fn remove_expired(&mut self) { + let now = Instant::now(); + while self + .packets + .front() + .is_some_and(|packet| now.duration_since(packet.last_seen) >= Self::TIMEOUT) + { + self.evict_oldest(); + } + } + + pub(super) fn process( + &mut self, + packet: ZCPacket, + local_ip: Option, + flows: &FlowSet, + sender: &mpsc::Sender, + ) -> PeerPacketFilterResult { + self.set_local_ip(local_ip); + // The caller has already classified a valid IPv4 UDP fragment. + let ip = Ipv4Packet::new_checked(packet.payload()).unwrap(); + if Some(ip.dst_addr()) != local_ip { + return PeerPacketFilterResult::Pass(packet); + } + let payload = ip.payload(); + if !ip.verify_checksum() + || payload.is_empty() + || (ip.more_frags() && !payload.len().is_multiple_of(8)) + || usize::from(ip.frag_offset()) + payload.len() > usize::from(u16::MAX) - 20 + { + return PeerPacketFilterResult::Consumed; + } + let port = if ip.frag_offset() == 0 { + if payload.len() < UDP_HEADER_LEN { + return PeerPacketFilterResult::Consumed; + } + // Its length covers the whole datagram, not just this fragment. + let udp = UdpPacket::new_unchecked(payload); + if usize::from(udp.len()) < UDP_HEADER_LEN { + return PeerPacketFilterResult::Consumed; + } + Some(udp.dst_port()) + } else { + None + }; + let header = packet.peer_manager_header().unwrap(); + let key = UdpFragmentKey { + source: ip.src_addr(), + destination: ip.dst_addr(), + id: ip.ident(), + from_peer: header.from_peer_id.get(), + to_peer: header.to_peer_id.get(), + packet_type: header.packet_type, + exit_node: header.is_exit_node(), + no_proxy: header.is_no_proxy(), + not_send_to_tun: header.is_not_send_to_tun(), + }; + self.remove_expired(); + let now = Instant::now(); + let mut context = self + .packets + .get_mut(&key) + .map(|context| std::mem::replace(context, UdpFragment::new(now))) + .unwrap_or_else(|| UdpFragment::new(now)); + self.bytes -= context.bytes; + context.last_seen = now; + if context.route.is_none() + && let Some(port) = port + { + let entry = FlowKey::udp_bind(SocketAddr::new(key.destination.into(), port)); + let registration = flows.registration(&entry); + context.route = Some(UdpFragmentRoute { + entry, + registration, + }); + } + let result = if let Some(route) = &context.route { + context.bytes = 0; + // Preserve ownership across close/rebind, including downstream routes. + if let Some(registration) = route.registration { + if flows.registration(&route.entry) == Some(registration) { + for packet in std::iter::once(packet).chain(context.pending.drain(..)) { + if let Err(error) = sender.try_send(packet) { + tracing::trace!(?error, "data plane fragment queue full or closed"); + } + } + } + context.pending = Vec::new(); + PeerPacketFilterResult::Consumed + } else if context.pending.is_empty() { + PeerPacketFilterResult::Pass(packet) + } else { + context.pending.push(packet); + PeerPacketFilterResult::PassBatch(std::mem::take(&mut context.pending)) + } + } else { + context.bytes += packet.buf_len(); + context.pending.push(packet); + PeerPacketFilterResult::Consumed + }; + self.bytes += context.bytes; + // Replacing an entry moves it to the back, keeping expiry and LRU order. + self.packets.insert(key, context); + while self.packets.len() > Self::CAPACITY || self.bytes > Self::MAX_BYTES { + self.evict_oldest(); + } + result + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::super::{ + FlowData, FlowTable, + tests::{UDP_LOCAL_IP, build_udp_fragments}, + }; + use super::*; + + #[tokio::test] + async fn recreated_sending_stack_does_not_replay_fragment_routes() { + use super::super::stack::SmoltcpPlane; + use crate::gateway::smoltcp::{Net, NetConfig, channel_device}; + + let flows = Arc::new(FlowTable::default()); + let local = SocketAddr::new(UDP_LOCAL_IP.into(), 40000); + flows.insert(FlowKey::udp_bind(local), FlowData::Udp); + let (sender, mut receiver) = mpsc::channel(16); + let mut cache = UdpFragments::default(); + // Use controlled seeds: random 16-bit IDs may legitimately collide. + // Each iteration drops the old sender and creates a real new stack. + for (seed, port) in [(1, 40001), (2, 40000)] { + let mut capabilities = smoltcp::phy::DeviceCapabilities::default(); + capabilities.medium = smoltcp::phy::Medium::Ip; + capabilities.max_transmission_unit = 1284; + let (device, _input, mut output) = channel_device::ChannelDevice::new(capabilities); + let net = Net::new( + device, + NetConfig::new( + SmoltcpPlane::interface_config(seed), + "10.144.144.3/24".parse().unwrap(), + Vec::new(), + None, + ), + ); + let socket = net + .udp_bind("10.144.144.3:53".parse().unwrap()) + .await + .unwrap(); + socket + .send_to(&[0xab; 4096], SocketAddr::new(UDP_LOCAL_IP.into(), port)) + .await + .unwrap(); + let mut count = 0; + loop { + let bytes = tokio::time::timeout(Duration::from_secs(1), output.recv()) + .await + .unwrap() + .unwrap(); + let ip = Ipv4Packet::new_checked(bytes.as_slice()).unwrap(); + assert!(ip.more_frags() || ip.frag_offset() != 0); + let last = !ip.more_frags(); + // Keep the old route live regardless of test-machine scheduling. + for context in cache.packets.values_mut() { + context.last_seen = Instant::now(); + } + let mut packet = ZCPacket::new_with_payload(&bytes); + packet.fill_peer_manager_hdr(1, 2, crate::packet::PacketType::Data as u8); + let result = cache.process(packet, Some(UDP_LOCAL_IP), &flows, &sender); + if port == local.port() { + assert!(matches!(result, PeerPacketFilterResult::Consumed)); + assert_eq!(receiver.try_recv().unwrap().payload(), bytes); + } else { + assert!(matches!(result, PeerPacketFilterResult::Pass(_))); + assert!(receiver.try_recv().is_err()); + } + count += 1; + if last { + break; + } + } + assert_eq!(count, 4); + } + assert_eq!(cache.packets.len(), 2); + } + + #[test] + fn routes_survive_binding_changes() { + for bound in [false, true] { + let flows = Arc::new(FlowTable::default()); + let entry = FlowKey::udp_bind(SocketAddr::new(UDP_LOCAL_IP.into(), 40000)); + if bound { + flows.insert(entry.clone(), FlowData::Udp); + } + let (sender, mut receiver) = mpsc::channel(4); + let mut cache = UdpFragments::default(); + let [first, tail] = build_udp_fragments(); + let result = cache.process(first, Some(UDP_LOCAL_IP), &flows, &sender); + assert_eq!(matches!(result, PeerPacketFilterResult::Consumed), bound); + if bound { + receiver.try_recv().unwrap(); + } + flows.clear(); + flows.insert(entry, FlowData::Udp); + let result = cache.process(tail, Some(UDP_LOCAL_IP), &flows, &sender); + assert_eq!(matches!(result, PeerPacketFilterResult::Consumed), bound); + assert!( + receiver.try_recv().is_err(), + "old fragments reached a new binding" + ); + } + } + + #[test] + fn cache_expires_and_resets_on_address_change() { + let flows = Arc::new(FlowTable::default()); + let (sender, _receiver) = mpsc::channel(4); + let mut cache = UdpFragments::default(); + let [first, tail] = build_udp_fragments(); + cache.process(tail.clone(), Some(UDP_LOCAL_IP), &flows, &sender); + assert!(cache.bytes > 0); + for context in cache.packets.values_mut() { + context.last_seen -= UdpFragments::TIMEOUT; + } + cache.remove_expired(); + assert!(cache.packets.is_empty()); + assert_eq!(cache.bytes, 0); + assert!(matches!( + cache.process(first, Some(UDP_LOCAL_IP), &flows, &sender), + PeerPacketFilterResult::Pass(_) + )); + cache.process(tail, Some(UDP_LOCAL_IP), &flows, &sender); + cache.set_local_ip(None); + assert!(cache.packets.is_empty()); + assert_eq!(cache.bytes, 0); + } + + #[test] + fn cache_limits_pending_bytes_and_evicts_least_recent_entry() { + let flows = Arc::new(FlowTable::default()); + let (sender, _receiver) = mpsc::channel(4); + let mut cache = UdpFragments::default(); + for id in 0..=UdpFragments::CAPACITY as u16 { + let [_, mut tail] = build_udp_fragments(); + let mut ip = Ipv4Packet::new_unchecked(tail.mut_payload()); + ip.set_ident(id); + ip.fill_checksum(); + cache.process(tail, Some(UDP_LOCAL_IP), &flows, &sender); + } + assert_eq!(cache.packets.len(), UdpFragments::CAPACITY); + assert!(!cache.packets.keys().any(|key| key.id == 0)); + let [_, tail] = build_udp_fragments(); + cache.process(tail.clone(), Some(UDP_LOCAL_IP), &flows, &sender); + let refreshed = *cache.packets.keys().find(|key| key.id == 42).unwrap(); + cache.evict_oldest(); + assert!(cache.packets.contains_key(&refreshed)); + cache.set_local_ip(None); + let mut payload = tail.payload().to_vec(); + payload.resize(65516, 0); + let mut ip = Ipv4Packet::new_unchecked(&mut payload); + ip.set_total_len(65516); + ip.fill_checksum(); + let mut large = ZCPacket::new_with_payload(&payload); + large.fill_peer_manager_hdr(1, 2, crate::packet::PacketType::Data as u8); + for _ in 0..=UdpFragments::MAX_BYTES / large.buf_len() { + cache.process(large.clone(), Some(UDP_LOCAL_IP), &flows, &sender); + assert!(cache.bytes <= UdpFragments::MAX_BYTES); + } + assert_eq!( + cache.bytes, 0, + "over-budget pending datagram was not evicted" + ); + assert!(cache.packets.is_empty()); + } + + #[test] + fn pending_fragments_do_not_cross_peer_or_delivery_flags() { + let flows = Arc::new(FlowTable::default()); + let (sender, _receiver) = mpsc::channel(4); + for change_peer in [true, false] { + let mut cache = UdpFragments::default(); + let [first, mut tail] = build_udp_fragments(); + if change_peer { + tail.mut_peer_manager_header().unwrap().from_peer_id.set(9); + } else { + tail.mut_peer_manager_header().unwrap().set_no_proxy(false); + } + cache.process(tail, Some(UDP_LOCAL_IP), &flows, &sender); + assert!(matches!( + cache.process(first, Some(UDP_LOCAL_IP), &flows, &sender), + PeerPacketFilterResult::Pass(_) + )); + assert!(cache.bytes > 0); + } + } +} diff --git a/easytier-core/src/gateway/proxy/icmp_proxy_service.rs b/easytier-core/src/gateway/proxy/icmp_proxy_service.rs index 1c167f4c..14f4e48e 100644 --- a/easytier-core/src/gateway/proxy/icmp_proxy_service.rs +++ b/easytier-core/src/gateway/proxy/icmp_proxy_service.rs @@ -15,7 +15,7 @@ use tokio::{ use crate::{ packet::ZCPacket, peers::{ - PeerPacketFilter, + PeerPacketFilter, PeerPacketFilterResult, peer_manager::{PeerManagerCore, PipelineRegistrationGuard}, }, }; @@ -249,11 +249,11 @@ struct IcmpProxyServiceFilter { #[async_trait::async_trait] impl PeerPacketFilter for IcmpProxyServiceFilter { - async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> Option { + async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> PeerPacketFilterResult { let Some(service) = self.service.upgrade() else { - return Some(packet); + return PeerPacketFilterResult::Pass(packet); }; - service.handle_peer_packet(packet).await + service.handle_peer_packet(packet).await.into() } } diff --git a/easytier-core/src/gateway/proxy/tcp_proxy_service.rs b/easytier-core/src/gateway/proxy/tcp_proxy_service.rs index 805bcb11..35da9d50 100644 --- a/easytier-core/src/gateway/proxy/tcp_proxy_service.rs +++ b/easytier-core/src/gateway/proxy/tcp_proxy_service.rs @@ -11,7 +11,7 @@ use crate::{ foundation::time::timeout, packet::ZCPacket, peers::{ - NicPacketFilter, PeerPacketFilter, + NicPacketFilter, PeerPacketFilter, PeerPacketFilterResult, peer_manager::{PeerManagerCore, PipelineRegistrationGuard}, }, socket::{ @@ -544,11 +544,11 @@ struct TcpProxyServiceFilter< impl PeerPacketFilter for TcpProxyServiceFilter { - async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> Option { + async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> PeerPacketFilterResult { let Some(service) = self.service.upgrade() else { - return Some(packet); + return PeerPacketFilterResult::Pass(packet); }; - service.handle_peer_packet(packet).await + service.handle_peer_packet(packet).await.into() } } diff --git a/easytier-core/src/gateway/proxy/udp_proxy_service.rs b/easytier-core/src/gateway/proxy/udp_proxy_service.rs index e06d4e1a..f741915a 100644 --- a/easytier-core/src/gateway/proxy/udp_proxy_service.rs +++ b/easytier-core/src/gateway/proxy/udp_proxy_service.rs @@ -9,8 +9,8 @@ use tokio::sync::mpsc::{self, Receiver, Sender}; use tokio::task::JoinSet; use crate::packet::ZCPacket; -use crate::peers::PeerPacketFilter; use crate::peers::peer_manager::{PeerManagerCore, PipelineRegistrationGuard}; +use crate::peers::{PeerPacketFilter, PeerPacketFilterResult}; use super::cidr_table::ProxyCidrTable; use super::traits::{UdpProxyResponseSink, UdpProxyRuntime}; @@ -208,10 +208,10 @@ struct UdpProxyServiceFilter { #[async_trait::async_trait] impl PeerPacketFilter for UdpProxyServiceFilter { - async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> Option { + async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> PeerPacketFilterResult { let Some(service) = self.service.upgrade() else { - return Some(packet); + return PeerPacketFilterResult::Pass(packet); }; - service.handle_peer_packet(packet).await + service.handle_peer_packet(packet).await.into() } } diff --git a/easytier-core/src/gateway/proxy/wrapped_transport.rs b/easytier-core/src/gateway/proxy/wrapped_transport.rs index 02280ae3..e061ec9a 100644 --- a/easytier-core/src/gateway/proxy/wrapped_transport.rs +++ b/easytier-core/src/gateway/proxy/wrapped_transport.rs @@ -12,7 +12,7 @@ use crate::{ listener::RunningListenerRegistry, packet::{PacketType, ZCPacket, ZCPacketType}, peers::{ - PeerPacketFilter, + PeerPacketFilter, PeerPacketFilterResult, peer_manager::{PeerManagerCore, PipelineRegistrationGuard}, }, process_runtime::ProtectedTcpPortRegistry, @@ -147,15 +147,15 @@ struct WrappedTransportPeerFilter { #[async_trait] impl PeerPacketFilter for WrappedTransportPeerFilter { - async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> Option { + async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> PeerPacketFilterResult { let Some(header) = packet.peer_manager_header() else { - return Some(packet); + return PeerPacketFilterResult::Pass(packet); }; if header.packet_type != self.transport.incoming_packet_type(self.role) as u8 { - return Some(packet); + return PeerPacketFilterResult::Pass(packet); } let Some(engine) = self.engine.upgrade() else { - return Some(packet); + return PeerPacketFilterResult::Pass(packet); }; let from_peer_id = header.from_peer_id.get(); if let Err(error) = engine @@ -169,7 +169,7 @@ impl PeerPacketFilter for WrappedTransportPeerFilter { "failed to inject wrapped transport packet" ); } - None + PeerPacketFilterResult::Consumed } } @@ -797,7 +797,10 @@ mod tests { let mut packet = ZCPacket::new_with_payload(b"payload"); packet.fill_peer_manager_hdr(7, 1, PacketType::KcpDst as u8); - assert!(filter.try_process_packet_from_peer(packet).await.is_none()); + assert!(matches!( + filter.try_process_packet_from_peer(packet).await, + PeerPacketFilterResult::Consumed + )); assert_eq!( *engine.injections.lock().unwrap(), [( @@ -811,7 +814,10 @@ mod tests { drop(engine); let mut packet = ZCPacket::new_with_payload(b"stale"); packet.fill_peer_manager_hdr(7, 1, PacketType::KcpDst as u8); - assert!(filter.try_process_packet_from_peer(packet).await.is_some()); + assert!(matches!( + filter.try_process_packet_from_peer(packet).await, + PeerPacketFilterResult::Pass(_) + )); } #[test] diff --git a/easytier-core/src/gateway/smoltcp/tokio_smoltcp/mod.rs b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/mod.rs index 6184d43e..3b97e18f 100644 --- a/easytier-core/src/gateway/smoltcp/tokio_smoltcp/mod.rs +++ b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/mod.rs @@ -82,6 +82,9 @@ impl std::fmt::Debug for Net { } impl Net { + const EPHEMERAL_PORT_START: u16 = 10000; + const EPHEMERAL_PORT_END: u16 = 60001; + /// Creates a new `Net` instance. It panics if the medium is not supported. pub fn new(device: D, config: NetConfig) -> Net { Self::new2(device, config) @@ -119,18 +122,44 @@ impl Net { Net { reactor: Arc::new(reactor), ip_addr: config.ip_addr, - from_port: AtomicU16::new(10001), + from_port: AtomicU16::new(Self::EPHEMERAL_PORT_START + 1), stopper, _fut: AbortOnDropHandle::new(tokio::spawn(fut)), } } + fn next_port(port: u16) -> u16 { + if port >= Self::EPHEMERAL_PORT_END { + Self::EPHEMERAL_PORT_START + } else { + port + 1 + } + } + pub fn get_port(&self) -> u16 { self.from_port - .try_update(Ordering::SeqCst, Ordering::SeqCst, |x| { - Some(if x > 60000 { 10000 } else { x + 1 }) + .try_update(Ordering::SeqCst, Ordering::SeqCst, |port| { + Some(Self::next_port(port)) }) .unwrap() } + /// Visit each ephemeral port once, independently of concurrent allocations. + pub(crate) fn ephemeral_ports(&self) -> impl Iterator { + let start = self.get_port(); + std::iter::once(start).chain( + (start + 1..=Self::EPHEMERAL_PORT_END) + .chain(Self::EPHEMERAL_PORT_START..start) + .inspect(|&port| { + // Follow the scan, but never overwrite another allocator's progress. + let _ = self.from_port.compare_exchange( + port, + Self::next_port(port), + Ordering::SeqCst, + Ordering::SeqCst, + ); + }), + ) + } + /// Creates a new TcpListener, which will be bound to the specified address. pub async fn tcp_bind(&self, addr: SocketAddr) -> io::Result { let addr = self.set_address(addr); diff --git a/easytier-core/src/peers/mod.rs b/easytier-core/src/peers/mod.rs index abd980b7..b6a3e889 100644 --- a/easytier-core/src/peers/mod.rs +++ b/easytier-core/src/peers/mod.rs @@ -159,11 +159,28 @@ pub async fn recv_packet_from_chan( .map(|envelope| envelope.packet) } +/// Returned packets continue at the next peer-packet filter. +#[derive(Debug)] +pub enum PeerPacketFilterResult { + Consumed, + Pass(ZCPacket), + PassBatch(Vec), +} + +impl From> for PeerPacketFilterResult { + fn from(packet: Option) -> Self { + match packet { + Some(packet) => Self::Pass(packet), + None => Self::Consumed, + } + } +} + #[async_trait::async_trait] #[auto_impl::auto_impl(Arc)] pub trait PeerPacketFilter { - async fn try_process_packet_from_peer(&self, zc_packet: ZCPacket) -> Option { - Some(zc_packet) + async fn try_process_packet_from_peer(&self, zc_packet: ZCPacket) -> PeerPacketFilterResult { + PeerPacketFilterResult::Pass(zc_packet) } } diff --git a/easytier-core/src/peers/peer_manager.rs b/easytier-core/src/peers/peer_manager.rs index 0ab7b081..14496f0c 100644 --- a/easytier-core/src/peers/peer_manager.rs +++ b/easytier-core/src/peers/peer_manager.rs @@ -50,7 +50,7 @@ use crate::{ use super::{ BoxNicPacketFilter, BoxPeerPacketFilter, PacketRecvChanReceiver, PeerConnectionOrigin, - PeerPacketFilter, PeerPacketIngress, + PeerPacketFilter, PeerPacketFilterResult, PeerPacketIngress, acl::AclFilter, conn::{ peer_conn::{PeerConn, PeerConnId}, @@ -453,7 +453,7 @@ struct NicPacketProcessor { #[async_trait::async_trait] impl PeerPacketFilter for NicPacketProcessor { - async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> Option { + async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> PeerPacketFilterResult { let hdr = packet.peer_manager_header().unwrap(); if hdr.packet_type == PacketType::Data as u8 && !hdr.is_not_send_to_tun() { if hdr.is_encrypted() || hdr.is_compressed() { @@ -464,7 +464,7 @@ impl PeerPacketFilter for NicPacketProcessor { compressed = hdr.is_compressed(), "dropping packet before nic because it is not fully decoded" ); - return None; + return PeerPacketFilterResult::Consumed; } tracing::trace!(?packet, "send packet to nic channel"); if let Err(error) = self @@ -476,9 +476,9 @@ impl PeerPacketFilter for NicPacketProcessor { "dropping packet because nic channel cannot accept it" ); } - None + PeerPacketFilterResult::Consumed } else { - Some(packet) + PeerPacketFilterResult::Pass(packet) } } } @@ -489,16 +489,16 @@ struct PeerRpcPacketProcessor { #[async_trait::async_trait] impl PeerPacketFilter for PeerRpcPacketProcessor { - async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> Option { + async fn try_process_packet_from_peer(&self, packet: ZCPacket) -> PeerPacketFilterResult { let hdr = packet.peer_manager_header().unwrap(); if hdr.packet_type == PacketType::TaRpc as u8 || hdr.packet_type == PacketType::RpcReq as u8 || hdr.packet_type == PacketType::RpcResp as u8 { self.peer_rpc_tspt_sender.send(packet).unwrap(); - None + PeerPacketFilterResult::Consumed } else { - Some(packet) + PeerPacketFilterResult::Pass(packet) } } } @@ -3175,25 +3175,33 @@ impl PeerPacketRouter { return; } - let mut processed = false; - let mut zc_packet = Some(ret); - tracing::trace!(?zc_packet, "try_process_packet_from_peer"); + tracing::trace!(?ret, "try_process_packet_from_peer"); let pipelines = self.peer_packet_process_pipeline.load_full(); - for pipeline in pipelines.iter().rev() { - if let Some(filter) = pipeline.filter_if_active() { - zc_packet = filter - .try_process_packet_from_peer(zc_packet.unwrap()) - .await; + process_peer_packet(&pipelines, ret).await; + } + } +} + +async fn process_peer_packet(pipelines: &[PeerPipelineEntry], packet: ZCPacket) { + let mut first = Some((packet, pipelines.len())); + // The ordinary single-packet path does not allocate this queue. + let mut pending = Vec::new(); + 'packets: while let Some((mut packet, next)) = first.take().or_else(|| pending.pop()) { + for (index, pipeline) in pipelines[..next].iter().enumerate().rev() { + let Some(filter) = pipeline.filter_if_active() else { + continue; + }; + match filter.try_process_packet_from_peer(packet).await { + PeerPacketFilterResult::Consumed => continue 'packets, + PeerPacketFilterResult::Pass(next) => packet = next, + PeerPacketFilterResult::PassBatch(packets) => { + // Resume each packet after the filter that produced it. + pending.extend(packets.into_iter().rev().map(|packet| (packet, index))); + continue 'packets; } - if zc_packet.is_none() { - processed = true; - break; - } - } - if !processed { - tracing::error!(?zc_packet, "unhandled packet"); } } + tracing::error!(?packet, "unhandled packet"); } } @@ -3538,6 +3546,58 @@ mod tests { } } + #[tokio::test] + async fn packet_batches_resume_at_next_filter_in_order() { + struct Filter { + stage: u8, + seen: Arc>>, + } + #[async_trait::async_trait] + impl PeerPacketFilter for Filter { + async fn try_process_packet_from_peer( + &self, + packet: ZCPacket, + ) -> PeerPacketFilterResult { + self.seen + .lock() + .unwrap() + .push((self.stage, packet.payload()[0])); + match self.stage { + 0 => PeerPacketFilterResult::Consumed, + 1 | 2 => { + let mut extra = packet.clone(); + extra.mut_payload()[0] += if self.stage == 1 { 10 } else { 1 }; + PeerPacketFilterResult::PassBatch(vec![packet, extra]) + } + _ => PeerPacketFilterResult::Pass(packet), + } + } + } + let seen = Arc::new(std::sync::Mutex::new(Vec::new())); + let pipelines = (0..4) + .map(|stage| { + permanent_peer_pipeline_entry(Box::new(Filter { + stage, + seen: seen.clone(), + })) + }) + .collect::>(); + process_peer_packet(&pipelines, ZCPacket::new_with_payload(&[0])).await; + assert_eq!( + *seen.lock().unwrap(), + vec![ + (3, 0), + (2, 0), + (1, 0), + (0, 0), + (0, 10), + (1, 1), + (0, 1), + (0, 11), + ] + ); + } + struct RuntimeConfigStunInfoSource(CoreRuntimeConfigStore); impl PeerStunInfoSource for RuntimeConfigStunInfoSource { @@ -3607,7 +3667,7 @@ mod tests { .await; assert!(result.is_ok(), "a full host queue blocked the peer router"); - assert!(result.unwrap().is_none()); + assert!(matches!(result.unwrap(), PeerPacketFilterResult::Consumed)); assert_eq!(nic_receiver.try_recv().unwrap().payload(), b"queued"); }