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"); }