diff --git a/easytier-core/src/gateway/proxy/tcp_proxy_service.rs b/easytier-core/src/gateway/proxy/tcp_proxy_service.rs index 7a8a1ac6..964c7b3e 100644 --- a/easytier-core/src/gateway/proxy/tcp_proxy_service.rs +++ b/easytier-core/src/gateway/proxy/tcp_proxy_service.rs @@ -467,7 +467,7 @@ impl; - pub struct SmolTcpStack { ingress_tx: mpsc::Sender, output_rx: Mutex>>>, - net: Arc>>, - listener_tx: mpsc::UnboundedSender, - listener_rx: Mutex>, + listener: Mutex, + _net: Net, tasks: Arc>>, } @@ -68,14 +62,16 @@ impl SmolTcpStack { ), ); net.set_any_ip(true); + let listener = net + .tcp_bind("0.0.0.0:8899".parse().unwrap()) + .await + .map_err(|error| anyhow::anyhow!("bind smoltcp listener failed: {error}"))?; - let (listener_tx, listener_rx) = mpsc::unbounded_channel(); Ok(Arc::new(Self { ingress_tx, output_rx: Mutex::new(Some(stack_stream)), - net: Arc::new(Mutex::new(Some(net))), - listener_tx, - listener_rx: Mutex::new(listener_rx), + listener: Mutex::new(listener), + _net: net, tasks, })) } @@ -99,43 +95,14 @@ impl SmolTcpStack { .ok_or_else(|| anyhow::anyhow!("smoltcp output receiver already taken")) } - pub async fn add_listener(&self) { - let tx = self.listener_tx.clone(); - let locked_net = self.net.lock().await; - let mut tcp = locked_net - .as_ref() - .expect("smoltcp net initialized") - .tcp_bind("0.0.0.0:8899".parse().unwrap()) - .await - .unwrap(); - self.tasks.lock().unwrap().spawn(async move { - let ret = timeout(Duration::from_secs(10), tcp.accept()).await; - if let Ok(accept_ret) = ret { - let _ = - tx.send(accept_ret.map_err(|err| { - anyhow::anyhow!("smol tcp listener accept failed: {:?}", err) - })); - } else { - tracing::error!( - target: "easytier_core::gateway::stack", - "smol tcp listener accept timeout" - ); - } - }); - tracing::info!( - target: "easytier_core::gateway::stack", - "smol tcp listener added" - ); - } - pub async fn accept(&self) -> anyhow::Result<(SocketAddr, Box)> { let (stream, src) = self - .listener_rx + .listener .lock() .await - .recv() + .accept() .await - .ok_or_else(|| anyhow::anyhow!("smoltcp listener closed"))??; + .map_err(|error| anyhow::anyhow!("smoltcp listener accept failed: {error}"))?; tracing::info!( target: "easytier_core::gateway::stack", ?src, @@ -156,3 +123,70 @@ pub fn output_dst_ip(data: &[u8]) -> anyhow::Result { .map_err(|err| anyhow::anyhow!("smoltcp output is not an IPv4 packet: {:?}", err))?; Ok(IpAddr::V4(ipv4.dst_addr())) } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use smoltcp::wire::{Ipv4Address, TcpControl, TcpSeqNumber}; + + use super::SmolTcpStack; + use crate::{ + gateway::smoltcp::tokio_smoltcp::test_utils::{TcpPackets, recv_tcp}, + packet::ZCPacket, + }; + + const LOCAL_ADDR: Ipv4Address = Ipv4Address::new(192, 88, 99, 254); + const LOCAL_PORT: u16 = 8899; + const PACKETS: TcpPackets = TcpPackets::new(LOCAL_ADDR, LOCAL_PORT); + + #[tokio::test] + async fn accepts_concurrent_connections_with_one_logical_listener() { + let stack = SmolTcpStack::new(LOCAL_ADDR).await.unwrap(); + let mut output = stack.take_output_rx().await.unwrap(); + let client_addr = Ipv4Address::new(192, 88, 99, 1); + + for (port, sequence) in [(40000, 1000), (40001, 2000)] { + stack + .send_ingress(ZCPacket::new_with_payload(&PACKETS.syn( + client_addr, + port, + sequence, + ))) + .await + .unwrap(); + } + + let mut syn_acks = Vec::new(); + for _ in 0..2 { + let syn_ack = recv_tcp(&mut output).await; + assert_eq!(syn_ack.control, TcpControl::Syn); + syn_acks.push((syn_ack.dst_port, syn_ack.sequence)); + } + for (port, sequence) in syn_acks { + let client_sequence = if port == 40000 { 1001 } else { 2001 }; + stack + .send_ingress(ZCPacket::new_with_payload(&PACKETS.ack( + client_addr, + port, + TcpSeqNumber(client_sequence), + sequence + 1, + ))) + .await + .unwrap(); + } + + let mut peers = Vec::new(); + let mut streams = Vec::new(); + for _ in 0..2 { + let (peer, stream) = tokio::time::timeout(Duration::from_secs(1), stack.accept()) + .await + .unwrap() + .unwrap(); + peers.push(peer.port()); + streams.push(stream); + } + peers.sort_unstable(); + assert_eq!(peers, vec![40000, 40001]); + } +} diff --git a/easytier-core/src/gateway/smoltcp/tokio_smoltcp/mod.rs b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/mod.rs index 642f6e2a..24850674 100644 --- a/easytier-core/src/gateway/smoltcp/tokio_smoltcp/mod.rs +++ b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/mod.rs @@ -30,6 +30,8 @@ pub mod device; mod reactor; mod socket; mod socket_allocator; +#[cfg(test)] +pub(super) mod test_utils; /// A config for a `Net`. /// diff --git a/easytier-core/src/gateway/smoltcp/tokio_smoltcp/reactor.rs b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/reactor.rs index 73c44e44..46dd7f48 100644 --- a/easytier-core/src/gateway/smoltcp/tokio_smoltcp/reactor.rs +++ b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/reactor.rs @@ -1,15 +1,29 @@ use super::{ device::{BufferDevice, Packet}, - socket_allocator::{BufferSize, SocketAlloctor}, + socket::TCP_LISTENER_MAX_PENDING, + socket_allocator::{ + BufferSize, SocketAlloctor, SocketHandle as OwnedSocketHandle, TCP_LISTENER_TIMEOUT, + TCP_TIMEOUT, + }, }; use futures::{FutureExt, SinkExt, StreamExt, stream::iter}; use parking_lot::{MappedMutexGuard, Mutex, MutexGuard}; use smoltcp::{ - iface::{Context, Interface, SocketHandle}, - socket::{AnySocket, Socket}, + iface::{Context, Interface, SocketHandle as SmolSocketHandle, SocketSet}, + socket::{AnySocket, Socket, tcp}, time::{Duration, Instant}, + wire::{ + IpAddress, IpEndpoint, IpProtocol, Ipv4Packet, Ipv6ExtHeader, Ipv6ExtHeaderRepr, + Ipv6Packet, TcpPacket, + }, +}; +use std::{ + collections::VecDeque, + future::Future, + io, + sync::Arc, + task::{Context as TaskContext, Poll, Waker}, }; -use std::{collections::VecDeque, future::Future, io, sync::Arc}; use tokio::{pin, select, sync::Notify}; use crate::foundation::time::sleep; @@ -17,10 +31,30 @@ use crate::foundation::time::sleep; pub(crate) type BufferInterface = Arc>; const MAX_BURST_SIZE: usize = 100; +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) struct TcpListenerId(u64); + +struct TcpListenerEntry { + id: TcpListenerId, + local_endpoint: IpEndpoint, + // Raw handles stay owned by the listener until accept transfers one to TcpStream. + pending: Vec, + accept_waker: Option, +} + +#[derive(Default)] +struct TcpListenerRegistry { + next_id: u64, + entries: Vec, +} + +type SharedTcpListenerRegistry = Arc>; + pub(crate) struct Reactor { notify: Arc, iface: BufferInterface, socket_allocator: SocketAlloctor, + tcp_listeners: SharedTcpListenerRegistry, } async fn receive( @@ -33,11 +67,148 @@ async fn receive( Ok(()) } +fn tcp_syn_endpoints(packet: &[u8]) -> Option<(IpEndpoint, IpEndpoint)> { + let (src_addr, dst_addr, payload) = match packet.first()? >> 4 { + 4 => { + let packet = Ipv4Packet::new_checked(packet).ok()?; + if packet.next_header() != IpProtocol::Tcp { + return None; + } + ( + IpAddress::Ipv4(packet.src_addr()), + IpAddress::Ipv4(packet.dst_addr()), + packet.payload(), + ) + } + 6 => { + let packet = Ipv6Packet::new_checked(packet).ok()?; + let mut next_header = packet.next_header(); + let mut payload = packet.payload(); + if next_header == IpProtocol::HopByHop { + let header = Ipv6ExtHeader::new_checked(payload).ok()?; + let repr = Ipv6ExtHeaderRepr::parse(&header).ok()?; + next_header = repr.next_header; + payload = &payload[repr.header_len() + repr.data.len()..]; + } + if next_header != IpProtocol::Tcp { + return None; + } + ( + IpAddress::Ipv6(packet.src_addr()), + IpAddress::Ipv6(packet.dst_addr()), + payload, + ) + } + _ => return None, + }; + let packet = TcpPacket::new_checked(payload).ok()?; + if !packet.syn() || packet.ack() || packet.rst() || packet.fin() { + return None; + } + Some(( + IpEndpoint::new(src_addr, packet.src_port()), + IpEndpoint::new(dst_addr, packet.dst_port()), + )) +} + +fn tcp_tuple_exists( + sockets: &SocketSet<'_>, + remote_endpoint: IpEndpoint, + local_endpoint: IpEndpoint, +) -> bool { + sockets.iter().any(|(_, socket)| { + let Socket::Tcp(socket) = socket else { + return false; + }; + socket.state() != smoltcp::socket::tcp::State::Closed + && socket.remote_endpoint() == Some(remote_endpoint) + && socket.local_endpoint() == Some(local_endpoint) + }) +} + +fn tcp_listener_connection_ready(state: tcp::State) -> bool { + matches!(state, tcp::State::Established | tcp::State::CloseWait) +} + +fn prepare_tcp_listener_socket( + remote_endpoint: IpEndpoint, + local_endpoint: IpEndpoint, + listeners: &mut TcpListenerRegistry, + sockets: &mut SocketSet<'static>, + socket_allocator: &SocketAlloctor, +) -> bool { + let Some(listener) = listeners + .entries + .iter_mut() + .find(|listener| listener.local_endpoint == local_endpoint) + else { + return false; + }; + if listener.pending.len() >= TCP_LISTENER_MAX_PENDING { + return false; + } + // Keep the global lookup off the rejected-SYN path. Once pending is full, + // retransmissions can still reach their exact socket without a new Listen socket. + if tcp_tuple_exists(sockets, remote_endpoint, local_endpoint) { + return false; + } + // The socket exists only for this packet. update_tcp_listeners removes it unless + // smoltcp binds it to the SYN's four-tuple. + match socket_allocator.add_tcp_listener_socket(sockets, local_endpoint) { + Ok(handle) => { + listener.pending.push(handle); + true + } + Err(error) => { + tracing::error!( + target: "easytier_core::gateway::smoltcp", + ?error, + ?local_endpoint, + "failed to allocate TCP listener socket" + ); + false + } + } +} + +fn update_tcp_listeners( + listeners: &mut TcpListenerRegistry, + sockets: &mut SocketSet<'static>, +) -> Vec { + let mut wakers = Vec::new(); + for listener in &mut listeners.entries { + let mut index = 0; + let mut ready = false; + while index < listener.pending.len() { + let handle = listener.pending[index]; + let state = sockets.get::(handle).state(); + if matches!(state, tcp::State::Closed | tcp::State::Listen) { + listener.pending.swap_remove(index); + sockets.remove(handle); + continue; + } + if tcp_listener_connection_ready(state) { + let socket = sockets.get_mut::(handle); + if socket.timeout() == Some(TCP_LISTENER_TIMEOUT) { + socket.set_timeout(Some(TCP_TIMEOUT)); + } + ready = true; + } + index += 1; + } + if ready && let Some(waker) = listener.accept_waker.take() { + wakers.push(waker); + } + } + wakers +} + async fn run( mut async_iface: impl super::device::AsyncDevice, iface: BufferInterface, mut device: BufferDevice, socket_allocator: SocketAlloctor, + tcp_listeners: SharedTcpListenerRegistry, notify: Arc, stopper: Arc, ) -> io::Result<()> { @@ -84,24 +255,55 @@ async fn run( } } + let now = Instant::now(); let mut iface = iface.lock(); + let mut listeners = tcp_listeners.lock(); + let mut sockets = socket_allocator.sockets().lock(); + let recv_count = device.avaliable_recv_queue().min(recv_buf.len()); + let mut listener_wakers = Vec::new(); + for packet in recv_buf.drain(..recv_count) { + let listener_socket_added = + if let Some((remote_endpoint, local_endpoint)) = tcp_syn_endpoints(&packet) { + // Advance timers, reclaim stale pending sockets, and preserve packet + // order before exposing a generic Listen socket. + iface.poll(now, &mut device, &mut sockets); + listener_wakers.extend(update_tcp_listeners(&mut listeners, &mut sockets)); + prepare_tcp_listener_socket( + remote_endpoint, + local_endpoint, + &mut listeners, + &mut sockets, + &socket_allocator, + ) + } else { + false + }; + device.push_recv_queue(std::iter::once(packet)); + if listener_socket_added { + // Bind the generic Listen socket to this SYN before another packet + // can be dispatched to it. + iface.poll_ingress_single(now, &mut device, &mut sockets); + listener_wakers.extend(update_tcp_listeners(&mut listeners, &mut sockets)); + } + } - device.push_recv_queue(recv_buf.drain(..device.avaliable_recv_queue().min(recv_buf.len()))); - - iface.poll( - Instant::now(), - &mut device, - &mut socket_allocator.sockets().lock(), - ); + iface.poll(now, &mut device, &mut sockets); + listener_wakers.extend(update_tcp_listeners(&mut listeners, &mut sockets)); // wake up all closed sockets (smoltcp seems have a bug that it doesn't wake up closed sockets) - for (_, socket) in socket_allocator.sockets().lock().iter_mut() { + for (_, socket) in sockets.iter_mut() { if let Socket::Tcp(tcp) = socket && tcp.state() == smoltcp::socket::tcp::State::Closed { tcp.abort(); } } + drop(sockets); + drop(listeners); + drop(iface); + for waker in listener_wakers { + waker.wake(); + } } Ok(()) @@ -118,11 +320,13 @@ impl Reactor { let iface = Arc::new(Mutex::new(iface)); let notify = Arc::new(Notify::new()); let socket_allocator = SocketAlloctor::new(buffer_size); + let tcp_listeners = Arc::new(Mutex::new(TcpListenerRegistry::default())); let fut = run( async_device, iface.clone(), device, socket_allocator.clone(), + tcp_listeners.clone(), notify.clone(), stopper, ); @@ -132,13 +336,14 @@ impl Reactor { notify, iface, socket_allocator, + tcp_listeners, }, fut, ) } pub fn get_socket>( &self, - handle: SocketHandle, + handle: SmolSocketHandle, ) -> MappedMutexGuard<'_, T> { MutexGuard::map( self.socket_allocator.sockets().lock(), @@ -151,6 +356,107 @@ impl Reactor { pub fn socket_allocator(&self) -> &SocketAlloctor { &self.socket_allocator } + pub(super) fn register_tcp_listener( + &self, + local_endpoint: IpEndpoint, + ) -> io::Result { + let mut listeners = self.tcp_listeners.lock(); + if listeners + .entries + .iter() + .any(|listener| listener.local_endpoint == local_endpoint) + { + return Err(io::ErrorKind::AddrInUse.into()); + } + let id = TcpListenerId(listeners.next_id); + listeners.next_id += 1; + listeners.entries.push(TcpListenerEntry { + id, + local_endpoint, + pending: Vec::new(), + accept_waker: None, + }); + Ok(id) + } + pub(super) fn poll_tcp_accept( + &self, + id: TcpListenerId, + cx: &TaskContext<'_>, + ) -> Poll> { + let mut listeners = self.tcp_listeners.lock(); + let Some(listener) = listeners + .entries + .iter_mut() + .find(|listener| listener.id == id) + else { + return Poll::Ready(Err(io::ErrorKind::NotConnected.into())); + }; + let mut sockets = self.socket_allocator.sockets().lock(); + let mut index = 0; + while index < listener.pending.len() { + let handle = listener.pending[index]; + let (state, endpoints) = { + let socket = sockets.get::(handle); + ( + socket.state(), + (socket.remote_endpoint(), socket.local_endpoint()), + ) + }; + if matches!(state, tcp::State::Closed | tcp::State::Listen) { + listener.pending.swap_remove(index); + sockets.remove(handle); + continue; + } + if tcp_listener_connection_ready(state) { + let handle = listener.pending.swap_remove(index); + let (Some(remote_endpoint), Some(local_endpoint)) = endpoints else { + sockets.remove(handle); + return Poll::Ready(Err(io::ErrorKind::NotConnected.into())); + }; + drop(sockets); + drop(listeners); + return Poll::Ready(Ok(( + self.socket_allocator.own_socket(handle), + remote_endpoint, + local_endpoint, + ))); + } + index += 1; + } + if listener + .accept_waker + .as_ref() + .is_none_or(|waker| !waker.will_wake(cx.waker())) + { + listener.accept_waker = Some(cx.waker().clone()); + } + Poll::Pending + } + pub(super) fn unregister_tcp_listener(&self, id: TcpListenerId) { + let mut listeners = self.tcp_listeners.lock(); + let Some(index) = listeners + .entries + .iter() + .position(|listener| listener.id == id) + else { + return; + }; + let listener = listeners.entries.swap_remove(index); + let mut sockets = self.socket_allocator.sockets().lock(); + for handle in listener.pending { + sockets.remove(handle); + } + } + #[cfg(test)] + pub(super) fn tcp_listener_handles(&self, id: TcpListenerId) -> Vec { + self.tcp_listeners + .lock() + .entries + .iter() + .find(|listener| listener.id == id) + .map(|listener| listener.pending.clone()) + .unwrap_or_default() + } pub fn notify(&self) { // The externally driven WASI runtime can park immediately after a // socket enqueues work. Keep one permit when the reactor has not diff --git a/easytier-core/src/gateway/smoltcp/tokio_smoltcp/socket.rs b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/socket.rs index 663ee292..0bab7442 100644 --- a/easytier-core/src/gateway/smoltcp/tokio_smoltcp/socket.rs +++ b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/socket.rs @@ -1,9 +1,11 @@ -use super::{reactor::Reactor, socket_allocator::SocketHandle}; +use super::{ + reactor::{Reactor, TcpListenerId}, + socket_allocator::SocketHandle, +}; use futures::future::{self, poll_fn}; pub use smoltcp::socket::tcp; use smoltcp::socket::udp; use smoltcp::wire::{IpAddress, IpEndpoint}; -use std::mem::replace; use std::net::IpAddr; use std::{ io, @@ -14,11 +16,13 @@ use std::{ }; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; +pub(super) const TCP_LISTENER_MAX_PENDING: usize = 16; + /// A TCP socket server, listening for connections. /// /// You can accept a new connection by using the accept method. pub struct TcpListener { - handle: SocketHandle, + id: TcpListenerId, reactor: Arc, local_addr: SocketAddr, } @@ -32,28 +36,28 @@ impl TcpListener { reactor: Arc, local_endpoint: IpEndpoint, ) -> io::Result { - let handle = reactor.socket_allocator().new_tcp_socket(); - { - let mut socket = reactor.get_socket::(*handle); - socket.listen(local_endpoint).map_err(map_err)?; - } - + let id = reactor.register_tcp_listener(local_endpoint)?; let local_addr = ep2sa(&local_endpoint); Ok(TcpListener { - handle, + id, reactor, local_addr, }) } - pub fn poll_accept(&mut self, cx: &Context<'_>) -> Poll> { - let mut socket = self.reactor.get_socket::(*self.handle); - if socket.state() == tcp::State::Established { - drop(socket); - return Poll::Ready(Ok(TcpStream::accept(self)?)); + pub fn poll_accept(&mut self, cx: &Context<'_>) -> Poll> { + match self.reactor.poll_tcp_accept(self.id, cx) { + Poll::Ready(Ok((handle, remote_endpoint, local_endpoint))) => Poll::Ready(Ok(( + TcpStream { + handle, + reactor: self.reactor.clone(), + local_addr: ep2sa(&local_endpoint), + }, + ep2sa(&remote_endpoint), + ))), + Poll::Ready(Err(error)) => Poll::Ready(Err(error)), + Poll::Pending => Poll::Pending, } - socket.register_send_waker(cx.waker()); - Poll::Pending } pub async fn accept(&mut self) -> io::Result<(TcpStream, SocketAddr)> { poll_fn(|cx| self.poll_accept(cx)).await @@ -63,6 +67,12 @@ impl TcpListener { } } +impl Drop for TcpListener { + fn drop(&mut self) { + self.reactor.unregister_tcp_listener(self.id); + } +} + fn ep2sa(ep: &IpEndpoint) -> SocketAddr { match ep.addr { IpAddress::Ipv4(v4) => SocketAddr::new(IpAddr::V4(v4), ep.port), @@ -121,37 +131,6 @@ impl TcpStream { Ok(tcp) } - fn accept(listener: &mut TcpListener) -> io::Result<(TcpStream, SocketAddr)> { - let reactor = listener.reactor.clone(); - let new_handle = reactor.socket_allocator().new_tcp_socket(); - { - let mut new_socket = reactor.get_socket::(*new_handle); - new_socket - .listen(sa2ep(listener.local_addr)) - .map_err(map_err)?; - } - let (peer_addr, local_addr) = { - let socket = reactor.get_socket::(*listener.handle); - match (socket.remote_endpoint(), socket.local_endpoint()) { - (Some(remote_endpoint), Some(local_endpoint)) => ( - // should be Some, because the state is Established - ep2sa(&remote_endpoint), - ep2sa(&local_endpoint), - ), - _ => return Err(io::ErrorKind::NotConnected.into()), - } - }; - - Ok(( - TcpStream { - handle: replace(&mut listener.handle, new_handle), - reactor, - local_addr, - }, - peer_addr, - )) - } - pub fn local_addr(&self) -> io::Result { Ok(self.local_addr) } @@ -346,3 +325,432 @@ impl UdpSocket { Ok(self.local_addr) } } + +#[cfg(test)] +mod tests { + use std::time::Duration as StdDuration; + + use smoltcp::{ + iface::Config, + phy::{DeviceCapabilities, Medium}, + socket::tcp, + time::Duration, + wire::{HardwareAddress, Ipv4Address, TcpControl, TcpSeqNumber}, + }; + + use super::super::{ + Net, NetConfig, channel_device, + test_utils::{TcpPackets, recv_tcp, recv_tcp_for_port}, + }; + use super::{TCP_LISTENER_MAX_PENDING, TcpListener}; + + const LISTEN_ADDR: Ipv4Address = Ipv4Address::new(10, 126, 126, 1); + const LISTEN_PORT: u16 = 34569; + const PACKETS: TcpPackets = TcpPackets::new(LISTEN_ADDR, LISTEN_PORT); + + type TestNet = ( + Net, + tokio::sync::mpsc::Sender>>, + tokio::sync::mpsc::Receiver>, + ); + + fn test_net() -> TestNet { + let mut capabilities = DeviceCapabilities::default(); + capabilities.max_transmission_unit = 1280; + capabilities.medium = Medium::Ip; + let (device, ingress, egress) = channel_device::ChannelDevice::new(capabilities); + let net = Net::new( + device, + NetConfig::new( + Config::new(HardwareAddress::Ip), + "10.126.126.1/24".parse().unwrap(), + Vec::new(), + None, + ), + ); + (net, ingress, egress) + } + + async fn begin_handshake( + ingress: &tokio::sync::mpsc::Sender>>, + egress: &mut tokio::sync::mpsc::Receiver>, + client_addr: Ipv4Address, + client_port: u16, + client_sequence: i32, + ) -> TcpSeqNumber { + ingress + .send(Ok(PACKETS.syn(client_addr, client_port, client_sequence))) + .await + .unwrap(); + let syn_ack = recv_tcp_for_port(egress, client_port).await; + assert_eq!(syn_ack.control, TcpControl::Syn); + syn_ack.sequence + } + + async fn finish_handshake( + ingress: &tokio::sync::mpsc::Sender>>, + client_addr: Ipv4Address, + client_port: u16, + client_sequence: i32, + server_sequence: TcpSeqNumber, + ) { + ingress + .send(Ok(PACKETS.ack( + client_addr, + client_port, + TcpSeqNumber(client_sequence + 1), + server_sequence + 1, + ))) + .await + .unwrap(); + } + + fn listener_states(listener: &TcpListener) -> Vec { + listener + .reactor + .tcp_listener_handles(listener.id) + .into_iter() + .map(|handle| listener.reactor.get_socket::(handle).state()) + .collect() + } + + #[tokio::test] + async fn duplicate_syn_does_not_consume_another_backlog_slot() { + let (net, ingress, mut egress) = test_net(); + let listener = net + .tcp_bind("10.126.126.1:34569".parse().unwrap()) + .await + .unwrap(); + let client_addr = Ipv4Address::new(10, 126, 126, 2); + + { + let sockets = listener.reactor.socket_allocator().sockets().lock(); + assert_eq!( + sockets + .iter() + .filter(|(_, socket)| matches!(socket, smoltcp::socket::Socket::Tcp(_))) + .count(), + 0 + ); + } + begin_handshake(&ingress, &mut egress, client_addr, 40000, 1000).await; + + ingress + .send(Ok(PACKETS.syn(client_addr, 40000, 1000))) + .await + .unwrap(); + ingress + .send(Ok(PACKETS.syn(client_addr, 40001, 2000))) + .await + .unwrap(); + recv_tcp_for_port(&mut egress, 40001).await; + + let states = listener_states(&listener); + assert_eq!( + states + .iter() + .filter(|state| **state == tcp::State::SynReceived) + .count(), + 2 + ); + } + + #[tokio::test] + async fn listener_accepts_new_client_while_first_handshake_is_stalled() { + let (net, ingress, mut egress) = test_net(); + let mut listener = net + .tcp_bind("10.126.126.1:34569".parse().unwrap()) + .await + .unwrap(); + + let first_addr = Ipv4Address::new(10, 126, 126, 2); + begin_handshake(&ingress, &mut egress, first_addr, 40000, 1000).await; + let states = listener_states(&listener); + assert_eq!( + states + .iter() + .filter(|state| **state == tcp::State::SynReceived) + .count(), + 1 + ); + + let second_addr = Ipv4Address::new(10, 126, 126, 3); + let second_sequence = + begin_handshake(&ingress, &mut egress, second_addr, 40001, 2000).await; + for handle in listener.reactor.tcp_listener_handles(listener.id) { + let socket = listener.reactor.get_socket::(handle); + assert_eq!(socket.timeout(), Some(Duration::from_secs(5))); + } + + finish_handshake(&ingress, second_addr, 40001, 2000, second_sequence).await; + tokio::time::timeout(StdDuration::from_secs(1), async { + loop { + let established_timeout = listener + .reactor + .tcp_listener_handles(listener.id) + .into_iter() + .find_map(|handle| { + let socket = listener.reactor.get_socket::(handle); + (socket.state() == tcp::State::Established).then(|| socket.timeout()) + }); + if let Some(timeout) = established_timeout { + assert_eq!(timeout, Some(Duration::from_secs(60))); + break; + } + tokio::time::sleep(StdDuration::from_millis(1)).await; + } + }) + .await + .unwrap(); + let (stream, peer_addr) = + tokio::time::timeout(StdDuration::from_secs(1), listener.accept()) + .await + .unwrap() + .unwrap(); + assert_eq!(peer_addr, "10.126.126.3:40001".parse().unwrap()); + + let stream_socket = stream.reactor.get_socket::(*stream.handle); + assert_eq!(stream_socket.timeout(), Some(Duration::from_secs(60))); + drop(stream_socket); + assert_eq!(listener_states(&listener), vec![tcp::State::SynReceived]); + + { + let handle = listener.reactor.tcp_listener_handles(listener.id)[0]; + let mut socket = listener.reactor.get_socket::(handle); + socket.set_timeout(Some(Duration::from_millis(10))); + } + listener.reactor.notify(); + tokio::time::timeout(StdDuration::from_secs(1), async { + loop { + if listener + .reactor + .tcp_listener_handles(listener.id) + .is_empty() + { + break; + } + tokio::time::sleep(StdDuration::from_millis(1)).await; + } + }) + .await + .unwrap(); + + begin_handshake(&ingress, &mut egress, first_addr, 40002, 3000).await; + assert_eq!(listener_states(&listener), vec![tcp::State::SynReceived]); + } + + #[tokio::test] + async fn listener_limits_pending_connections() { + let (net, ingress, mut egress) = test_net(); + let mut listener = net + .tcp_bind("10.126.126.1:34569".parse().unwrap()) + .await + .unwrap(); + let client_addr = Ipv4Address::new(10, 126, 126, 2); + let mut syn_ack_sequences = Vec::with_capacity(TCP_LISTENER_MAX_PENDING); + + for index in 0..TCP_LISTENER_MAX_PENDING { + ingress + .send(Ok(PACKETS.syn( + client_addr, + 40000 + index as u16, + 1000 + index as i32, + ))) + .await + .unwrap(); + } + for _ in 0..TCP_LISTENER_MAX_PENDING { + let syn_ack = recv_tcp(&mut egress).await; + assert_eq!(syn_ack.control, TcpControl::Syn); + syn_ack_sequences.push((syn_ack.dst_port, syn_ack.sequence)); + } + syn_ack_sequences.sort_unstable_by_key(|(port, _)| *port); + assert!( + listener_states(&listener) + .iter() + .all(|state| *state == tcp::State::SynReceived) + ); + + ingress + .send(Ok(PACKETS.syn(client_addr, 50000, 3000))) + .await + .unwrap(); + assert_eq!(recv_tcp(&mut egress).await.control, TcpControl::Rst); + + for (port, syn_ack_sequence) in syn_ack_sequences { + let index = port - 40000; + finish_handshake( + &ingress, + client_addr, + port, + 1000 + index as i32, + syn_ack_sequence, + ) + .await; + } + + tokio::time::timeout(StdDuration::from_secs(1), async { + loop { + if listener_states(&listener) + .iter() + .all(|state| *state == tcp::State::Established) + { + break; + } + tokio::time::sleep(StdDuration::from_millis(1)).await; + } + }) + .await + .unwrap(); + + ingress + .send(Ok(PACKETS.syn(client_addr, 50001, 4000))) + .await + .unwrap(); + assert_eq!(recv_tcp(&mut egress).await.control, TcpControl::Rst); + + let mut streams = Vec::with_capacity(TCP_LISTENER_MAX_PENDING + 1); + let mut peer_ports = Vec::with_capacity(TCP_LISTENER_MAX_PENDING + 1); + let (stream, peer_addr) = listener.accept().await.unwrap(); + streams.push(stream); + peer_ports.push(peer_addr.port()); + + let replacement_sequence = + begin_handshake(&ingress, &mut egress, client_addr, 50002, 5000).await; + finish_handshake(&ingress, client_addr, 50002, 5000, replacement_sequence).await; + + for _ in 0..TCP_LISTENER_MAX_PENDING { + let (stream, peer_addr) = + tokio::time::timeout(StdDuration::from_secs(1), listener.accept()) + .await + .unwrap() + .unwrap(); + streams.push(stream); + peer_ports.push(peer_addr.port()); + } + peer_ports.sort_unstable(); + assert_eq!( + peer_ports, + (40000..40000 + TCP_LISTENER_MAX_PENDING as u16) + .chain(std::iter::once(50002)) + .collect::>() + ); + assert!(listener_states(&listener).is_empty()); + } + + #[tokio::test] + async fn closed_pending_connections_are_reclaimed_before_admission() { + let (net, ingress, mut egress) = test_net(); + let listener = net + .tcp_bind("10.126.126.1:34569".parse().unwrap()) + .await + .unwrap(); + let client_addr = Ipv4Address::new(10, 126, 126, 2); + + for index in 0..TCP_LISTENER_MAX_PENDING { + begin_handshake( + &ingress, + &mut egress, + client_addr, + 40000 + index as u16, + 1000 + index as i32, + ) + .await; + } + for handle in listener.reactor.tcp_listener_handles(listener.id) { + listener.reactor.get_socket::(handle).abort(); + } + + begin_handshake(&ingress, &mut egress, client_addr, 50000, 3000).await; + assert_eq!(listener_states(&listener), vec![tcp::State::SynReceived]); + } + + #[tokio::test] + async fn dropping_listener_removes_pending_connections() { + let (net, ingress, mut egress) = test_net(); + let listener = net + .tcp_bind("10.126.126.1:34569".parse().unwrap()) + .await + .unwrap(); + + begin_handshake( + &ingress, + &mut egress, + Ipv4Address::new(10, 126, 126, 2), + 40000, + 1000, + ) + .await; + assert_eq!(listener_states(&listener).len(), 1); + + drop(listener); + let listener = net + .tcp_bind("10.126.126.1:34569".parse().unwrap()) + .await + .unwrap(); + assert!(listener_states(&listener).is_empty()); + } + + #[tokio::test] + async fn old_syn_for_accepted_stream_does_not_create_pending_connection() { + let (net, ingress, mut egress) = test_net(); + let mut listener = net + .tcp_bind("10.126.126.1:34569".parse().unwrap()) + .await + .unwrap(); + let client_addr = Ipv4Address::new(10, 126, 126, 2); + + let server_sequence = + begin_handshake(&ingress, &mut egress, client_addr, 40000, 1000).await; + finish_handshake(&ingress, client_addr, 40000, 1000, server_sequence).await; + let (_stream, _) = tokio::time::timeout(StdDuration::from_secs(1), listener.accept()) + .await + .unwrap() + .unwrap(); + + ingress + .send(Ok(PACKETS.syn(client_addr, 40000, 1000))) + .await + .unwrap(); + ingress + .send(Ok(PACKETS.syn(client_addr, 40001, 2000))) + .await + .unwrap(); + recv_tcp_for_port(&mut egress, 40001).await; + + assert_eq!(listener_states(&listener), vec![tcp::State::SynReceived]); + } + + #[tokio::test] + async fn accepts_connection_closed_by_peer_before_accept() { + let (net, ingress, mut egress) = test_net(); + let mut listener = net + .tcp_bind("10.126.126.1:34569".parse().unwrap()) + .await + .unwrap(); + let client_addr = Ipv4Address::new(10, 126, 126, 2); + + let server_sequence = + begin_handshake(&ingress, &mut egress, client_addr, 40000, 1000).await; + finish_handshake(&ingress, client_addr, 40000, 1000, server_sequence).await; + ingress + .send(Ok(PACKETS.fin( + client_addr, + 40000, + TcpSeqNumber(1001), + server_sequence + 1, + ))) + .await + .unwrap(); + + let (stream, peer_addr) = + tokio::time::timeout(StdDuration::from_secs(1), listener.accept()) + .await + .unwrap() + .unwrap(); + assert_eq!(peer_addr, "10.126.126.2:40000".parse().unwrap()); + let socket = stream.reactor.get_socket::(*stream.handle); + assert_eq!(socket.state(), tcp::State::CloseWait); + assert_eq!(socket.timeout(), Some(Duration::from_secs(60))); + } +} diff --git a/easytier-core/src/gateway/smoltcp/tokio_smoltcp/socket_allocator.rs b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/socket_allocator.rs index 0f01678e..42362249 100644 --- a/easytier-core/src/gateway/smoltcp/tokio_smoltcp/socket_allocator.rs +++ b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/socket_allocator.rs @@ -3,12 +3,16 @@ use smoltcp::{ iface::{SocketHandle as InnerSocketHandle, SocketSet}, socket::{tcp, udp}, time::Duration, + wire::IpEndpoint, }; use std::{ ops::{Deref, DerefMut}, sync::Arc, }; +pub(super) const TCP_TIMEOUT: Duration = Duration::from_secs(60); +pub(super) const TCP_LISTENER_TIMEOUT: Duration = Duration::from_secs(5); + /// `BufferSize` is used to configure the size of the socket buffer. #[derive(Debug, Clone, Copy)] pub struct BufferSize { @@ -59,13 +63,26 @@ impl SocketAlloctor { let handle = set.add(self.alloc_tcp_socket()); SocketHandle::new(handle, self.sockets.clone()) } + pub(super) fn add_tcp_listener_socket( + &self, + sockets: &mut SocketSet<'static>, + local_endpoint: IpEndpoint, + ) -> Result { + let mut socket = self.alloc_tcp_socket(); + socket.set_timeout(Some(TCP_LISTENER_TIMEOUT)); + socket.listen(local_endpoint)?; + Ok(sockets.add(socket)) + } + pub(super) fn own_socket(&self, handle: InnerSocketHandle) -> SocketHandle { + SocketHandle::new(handle, self.sockets.clone()) + } fn alloc_tcp_socket(&self) -> tcp::Socket<'static> { let rx_buffer = tcp::SocketBuffer::new(vec![0; self.buffer_size.tcp_rx_size]); let tx_buffer = tcp::SocketBuffer::new(vec![0; self.buffer_size.tcp_tx_size]); let mut tcp = tcp::Socket::new(rx_buffer, tx_buffer); tcp.set_nagle_enabled(false); tcp.set_keep_alive(Some(Duration::from_secs(10))); - tcp.set_timeout(Some(Duration::from_secs(60))); + tcp.set_timeout(Some(TCP_TIMEOUT)); tcp } diff --git a/easytier-core/src/gateway/smoltcp/tokio_smoltcp/test_utils.rs b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/test_utils.rs new file mode 100644 index 00000000..8ed40d94 --- /dev/null +++ b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/test_utils.rs @@ -0,0 +1,149 @@ +use std::time::Duration; + +use smoltcp::{ + phy::ChecksumCapabilities, + wire::{ + IpAddress, IpProtocol, Ipv4Address, Ipv4Packet, Ipv4Repr, TcpControl, TcpPacket, TcpRepr, + TcpSeqNumber, + }, +}; +use tokio::sync::mpsc; + +pub(crate) struct TcpPackets { + local_addr: Ipv4Address, + local_port: u16, +} + +impl TcpPackets { + pub(crate) const fn new(local_addr: Ipv4Address, local_port: u16) -> Self { + Self { + local_addr, + local_port, + } + } + + fn packet(&self, src_addr: Ipv4Address, repr: TcpRepr<'_>) -> Vec { + let ipv4_repr = Ipv4Repr { + src_addr, + dst_addr: self.local_addr, + next_header: IpProtocol::Tcp, + payload_len: repr.buffer_len(), + hop_limit: 64, + }; + let mut packet = vec![0; ipv4_repr.buffer_len() + repr.buffer_len()]; + let mut ipv4_packet = Ipv4Packet::new_unchecked(&mut packet); + ipv4_repr.emit(&mut ipv4_packet, &ChecksumCapabilities::default()); + repr.emit( + &mut TcpPacket::new_unchecked(ipv4_packet.payload_mut()), + &IpAddress::Ipv4(src_addr), + &IpAddress::Ipv4(self.local_addr), + &ChecksumCapabilities::default(), + ); + packet + } + + pub(crate) fn syn(&self, src_addr: Ipv4Address, src_port: u16, sequence: i32) -> Vec { + self.packet( + src_addr, + TcpRepr { + src_port, + dst_port: self.local_port, + control: TcpControl::Syn, + seq_number: TcpSeqNumber(sequence), + ack_number: None, + window_len: u16::MAX, + window_scale: None, + max_seg_size: Some(1200), + sack_permitted: false, + sack_ranges: [None; 3], + timestamp: None, + payload: &[], + }, + ) + } + + pub(crate) fn ack( + &self, + src_addr: Ipv4Address, + src_port: u16, + sequence: TcpSeqNumber, + ack_number: TcpSeqNumber, + ) -> Vec { + self.control(src_addr, src_port, TcpControl::None, sequence, ack_number) + } + + pub(crate) fn fin( + &self, + src_addr: Ipv4Address, + src_port: u16, + sequence: TcpSeqNumber, + ack_number: TcpSeqNumber, + ) -> Vec { + self.control(src_addr, src_port, TcpControl::Fin, sequence, ack_number) + } + + fn control( + &self, + src_addr: Ipv4Address, + src_port: u16, + control: TcpControl, + sequence: TcpSeqNumber, + ack_number: TcpSeqNumber, + ) -> Vec { + self.packet( + src_addr, + TcpRepr { + src_port, + dst_port: self.local_port, + control, + seq_number: sequence, + ack_number: Some(ack_number), + window_len: u16::MAX, + window_scale: None, + max_seg_size: None, + sack_permitted: false, + sack_ranges: [None; 3], + timestamp: None, + payload: &[], + }, + ) + } +} + +pub(crate) struct TcpPacketSummary { + pub(crate) control: TcpControl, + pub(crate) sequence: TcpSeqNumber, + pub(crate) dst_port: u16, +} + +pub(crate) async fn recv_tcp(receiver: &mut mpsc::Receiver>) -> TcpPacketSummary { + let packet = tokio::time::timeout(Duration::from_secs(1), receiver.recv()) + .await + .unwrap() + .unwrap(); + let ipv4_packet = Ipv4Packet::new_checked(&packet).unwrap(); + let repr = TcpRepr::parse( + &TcpPacket::new_checked(ipv4_packet.payload()).unwrap(), + &IpAddress::Ipv4(ipv4_packet.src_addr()), + &IpAddress::Ipv4(ipv4_packet.dst_addr()), + &ChecksumCapabilities::default(), + ) + .unwrap(); + TcpPacketSummary { + control: repr.control, + sequence: repr.seq_number, + dst_port: repr.dst_port, + } +} + +pub(crate) async fn recv_tcp_for_port( + receiver: &mut mpsc::Receiver>, + dst_port: u16, +) -> TcpPacketSummary { + loop { + let packet = recv_tcp(receiver).await; + if packet.dst_port == dst_port { + return packet; + } + } +}