diff --git a/easytier-core/src/listener/mod.rs b/easytier-core/src/listener/mod.rs index 1d1007bf..d9f972c8 100644 --- a/easytier-core/src/listener/mod.rs +++ b/easytier-core/src/listener/mod.rs @@ -145,7 +145,6 @@ impl ListenerFactory { #[derive(Debug, Clone)] pub struct ListenerManagerOptions { - pub max_listen_retries: usize, pub listen_retry_delay: Duration, pub accept_retry_delay: Duration, } @@ -153,9 +152,8 @@ pub struct ListenerManagerOptions { impl Default for ListenerManagerOptions { fn default() -> Self { Self { - max_listen_retries: 5, - listen_retry_delay: Duration::from_secs(1), - accept_retry_delay: Duration::from_secs(1), + listen_retry_delay: Duration::from_millis(100), + accept_retry_delay: Duration::from_millis(100), } } } @@ -361,7 +359,7 @@ async fn run_listener( Accepted: Send + 'static, H: AcceptedSocketHandler + ?Sized + 'static, { - let mut listen_error_count = 0; + let mut listen_error_count: usize = 0; loop { let registered_listener = match initial_listener.take() { Some(listener) => listener, @@ -373,18 +371,15 @@ async fn run_listener( RegisteredListener::new(listener, events.clone(), registry.clone()) } Err(error) => { - listen_error_count += 1; - let will_retry = listen_error_count <= options.max_listen_retries; + listen_error_count = listen_error_count.saturating_add(1); events.emit(CoreEvent::ListenerAddFailed { url: listener.local_url(), error: format!("{error:?}"), retry_count: listen_error_count, - will_retry, + will_retry: true, }); tracing::error!(?error, ?listener, "listener listen error"); - if !will_retry { - return; - } + drop(listener); crate::foundation::time::sleep(options.listen_retry_delay).await; continue; } @@ -392,7 +387,7 @@ async fn run_listener( } }; let mut listener = registered_listener.listener; - let _registration = registered_listener.registration; + let registration = registered_listener.registration; loop { let listener_url = listener.local_url(); @@ -404,7 +399,6 @@ async fn run_listener( error: format!("{error:?}"), }); tracing::error!(?error, ?listener, "listener accept error"); - crate::foundation::time::sleep(options.accept_retry_delay).await; break; } }; @@ -423,6 +417,9 @@ async fn run_listener( } }); } + drop(listener); + drop(registration); + crate::foundation::time::sleep(options.accept_retry_delay).await; } } @@ -651,7 +648,6 @@ mod tests { ListenerManagerOptions { accept_retry_delay: Duration::from_millis(1), listen_retry_delay: Duration::from_millis(1), - max_listen_retries: 0, }, ); @@ -704,7 +700,6 @@ mod tests { ListenerManagerOptions { accept_retry_delay: Duration::from_millis(1), listen_retry_delay: Duration::from_millis(1), - max_listen_retries: 2, }, ); @@ -751,7 +746,6 @@ mod tests { ListenerManagerOptions { accept_retry_delay: Duration::from_millis(1), listen_retry_delay: Duration::from_millis(1), - max_listen_retries: 0, }, ); @@ -815,7 +809,6 @@ mod tests { ListenerManagerOptions { accept_retry_delay: Duration::from_millis(1), listen_retry_delay: Duration::from_millis(1), - max_listen_retries: 0, }, ); @@ -936,7 +929,6 @@ mod tests { ListenerManagerOptions { accept_retry_delay: Duration::from_millis(1), listen_retry_delay: Duration::from_millis(1), - max_listen_retries: 0, }, ); let accepts = Arc::new(Mutex::new(VecDeque::from([Ok::<_, anyhow::Error>( @@ -999,12 +991,13 @@ mod tests { } #[tokio::test] - async fn listen_retry_exhaustion_keeps_in_flight_handler_tasks_until_manager_drop() { + async fn listen_recovery_keeps_in_flight_handler_tasks_until_manager_drop() { let (drop_tx, mut drop_rx) = tokio::sync::oneshot::channel(); - let listen_results = Arc::new(Mutex::new(VecDeque::from([ - Ok(()), - Err(anyhow::anyhow!("bind failed")), - ]))); + let listen_results = Arc::new(Mutex::new( + std::iter::once(Ok(())) + .chain((0..8).map(|_| Err(anyhow::anyhow!("bind failed")))) + .collect::>(), + )); let accepts = Arc::new(Mutex::new(VecDeque::from([ Ok::<_, anyhow::Error>(DropSignal(Some(drop_tx))), Err(anyhow::anyhow!("accept failed")), @@ -1016,7 +1009,6 @@ mod tests { ListenerManagerOptions { accept_retry_delay: Duration::from_millis(1), listen_retry_delay: Duration::from_millis(1), - max_listen_retries: 0, }, ); let listener_results = listen_results.clone(); @@ -1024,7 +1016,7 @@ mod tests { manager.add_listener( move || { Box::new(DropSignalListener { - url: "mock://retry-exhausted".parse().unwrap(), + url: "mock://retrying".parse().unwrap(), listen_results: listener_results.clone(), accepts: listener_accepts.clone(), }) @@ -1035,15 +1027,22 @@ mod tests { manager.run().await.unwrap(); crate::foundation::time::timeout(Duration::from_secs(1), async { loop { - if events.events.lock().unwrap().iter().any(|event| { - matches!( + let recovered = { + let events = events.events.lock().unwrap(); + assert!(events.iter().all(|event| !matches!( event, CoreEvent::ListenerAddFailed { will_retry: false, .. } - ) - }) { + ))); + events + .iter() + .filter(|event| matches!(event, CoreEvent::ListenerAdded { .. })) + .count() + == 2 + }; + if recovered { break; } crate::foundation::time::sleep(Duration::from_millis(1)).await; diff --git a/easytier-core/src/socket/mod.rs b/easytier-core/src/socket/mod.rs index d39164df..80a6237c 100644 --- a/easytier-core/src/socket/mod.rs +++ b/easytier-core/src/socket/mod.rs @@ -6,6 +6,7 @@ //! egress, environment facts, and the WASI mechanism backend) live in //! [`crate::host`]. +pub mod retry; pub mod ring; pub mod tcp; pub mod udp; diff --git a/easytier-core/src/socket/retry.rs b/easytier-core/src/socket/retry.rs new file mode 100644 index 00000000..273a5ecc --- /dev/null +++ b/easytier-core/src/socket/retry.rs @@ -0,0 +1,85 @@ +//! Retry runtime listener failures until the owning task is cancelled. + +use super::SocketListener; + +pub async fn wait() { + crate::foundation::time::sleep(std::time::Duration::from_millis(100)).await; +} + +/// For owners whose listener can be restarted in place after an accept error. +/// The first listen remains fallible so startup errors reach the caller. +pub async fn accept_with_rebind(listener: &mut L) -> L::Accepted { + loop { + match listener.accept().await { + Ok(accepted) => return accepted, + Err(error) => { + tracing::warn!(target: "CORE", url = %listener.local_url(), error = %format_args!("{error:#}"), "listener failed; rebinding"); + } + } + loop { + wait().await; + match listener.listen().await { + Ok(()) => break, + Err(error) => { + tracing::warn!(target: "CORE", url = %listener.local_url(), error = %format_args!("{error:#}"), "listener bind failed; retrying"); + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::{collections::VecDeque, io}; + + #[derive(Debug, Default)] + struct Listener { + accepts: VecDeque>, + binds: VecDeque>, + } + #[async_trait::async_trait] + impl SocketListener for Listener { + type Accepted = usize; + async fn listen(&mut self) -> anyhow::Result<()> { + Ok(self.binds.pop_front().unwrap_or(Ok(()))?) + } + async fn accept(&mut self) -> anyhow::Result { + Ok(self.accepts.pop_front().expect("unexpected retry")?) + } + fn local_url(&self) -> url::Url { + "tcp://127.0.0.1:22020".parse().unwrap() + } + } + + #[tokio::test(start_paused = true)] + async fn unknown_accept_and_repeated_bind_failures_do_not_stop_recovery() { + let mut listener = Listener { + accepts: VecDeque::from([Err(io::Error::other("unknown failure")), Ok(7)]), + binds: (0..10) + .map(|_| Err(io::ErrorKind::PermissionDenied.into())) + .collect(), + }; + let start = tokio::time::Instant::now(); + assert_eq!(accept_with_rebind(&mut listener).await, 7); + assert_eq!(start.elapsed(), std::time::Duration::from_millis(1100)); + } + + #[tokio::test(start_paused = true)] + async fn rebind_wait_is_cancelled_with_its_owner() { + let mut listener = Listener { + accepts: VecDeque::from([Err(io::ErrorKind::BrokenPipe.into()), Ok(7)]), + binds: VecDeque::from([Err(io::ErrorKind::AddrInUse.into()), Ok(())]), + }; + assert!( + tokio::time::timeout( + std::time::Duration::from_millis(150), + accept_with_rebind(&mut listener) + ) + .await + .is_err() + ); + assert_eq!(listener.binds.len(), 1); + assert_eq!(listener.accepts.len(), 1); + } +} diff --git a/easytier-core/src/socket/tcp.rs b/easytier-core/src/socket/tcp.rs index 1256ec3e..476e17b4 100644 --- a/easytier-core/src/socket/tcp.rs +++ b/easytier-core/src/socket/tcp.rs @@ -363,15 +363,14 @@ where type Accepted = AcceptedTcpSocket; async fn listen(&mut self) -> anyhow::Result<()> { - if self.listener.is_some() { - return Ok(()); - } - + // Calling listen again replaces the old binding, including after an error. + self.listener = None; let listener = self.factory.bind_tcp(self.options.clone()).await?; let local_addr = listener.local_addr()?; self.url .set_port(Some(local_addr.port())) .map_err(|_| anyhow::anyhow!("failed to update tcp listener port for {}", self.url))?; + self.options.bind.local_addr = Some(local_addr); self.listener = Some(listener); Ok(()) } @@ -750,4 +749,25 @@ mod tests { assert_eq!(accepted.peer_addr().unwrap(), peer_addr); } + #[tokio::test] + async fn tcp_socket_listener_releases_failed_binding_and_keeps_assigned_port() { + let bound_addr: SocketAddr = "127.0.0.1:12010".parse().unwrap(); + let listener = Arc::new(MockTcpListener::new( + bound_addr, + vec![Err(io::ErrorKind::BrokenPipe.into())], + )); + let factory = Arc::new(MockTcpListenerFactory::new(listener)); + let mut socket_listener = TcpSocketListener::new_with_options( + "tcp://127.0.0.1:0".parse().unwrap(), + TcpListenOptions::direct_connect("127.0.0.1:0".parse().unwrap()), + factory.clone(), + ); + socket_listener.listen().await.unwrap(); + assert!(socket_listener.accept().await.is_err()); + assert!(socket_listener.listener.is_some()); + socket_listener.listen().await.unwrap(); + let binds = factory.binds.lock().unwrap(); + assert_eq!(binds.len(), 2); + assert_eq!(binds[1].bind.local_addr, Some(bound_addr)); + } } diff --git a/easytier-core/src/socket/udp/layer.rs b/easytier-core/src/socket/udp/layer.rs index f9dc4be2..cd272b31 100644 --- a/easytier-core/src/socket/udp/layer.rs +++ b/easytier-core/src/socket/udp/layer.rs @@ -295,7 +295,15 @@ where let accept = self.classified_accept(protocol)?; accept.accept_enabled.store(true, Ordering::Relaxed); let mut accepted_rx = accept.accepted_rx.lock().await; - accepted_rx.recv().await.ok_or_else(|| { + let mut shutdown = self.session_shutdown_tx.subscribe(); + // The layer owns the classified senders, so recv() alone does not + // finish when the socket receive task fails. + let session = tokio::select! { + biased; + _ = shutdown.wait_for(|closed| *closed) => None, + session = accepted_rx.recv() => session, + }; + session.ok_or_else(|| { io::Error::new( io::ErrorKind::UnexpectedEof, format!("{protocol:?} udp listener closed"), @@ -1175,7 +1183,7 @@ where self.factory.clone(), )); let mut session = layer.open_classified_session(request.protocol, request.remote_addr)?; - session._cleanup.layer_guard = Some(Box::new(layer)); + session.keep_layer_alive(layer); Ok(session) } } diff --git a/easytier-core/src/socket/udp/listener.rs b/easytier-core/src/socket/udp/listener.rs index 27c4e634..f8646978 100644 --- a/easytier-core/src/socket/udp/listener.rs +++ b/easytier-core/src/socket/udp/listener.rs @@ -116,15 +116,16 @@ where type Accepted = UdpSession; async fn listen(&mut self) -> anyhow::Result<()> { - if self.layer.is_some() { - return Ok(()); - } + self.layer = None; + self.socket = None; + *self.layer_ref.lock().unwrap() = None; let socket = self.factory.bind_udp(self.request.bind.clone()).await?; let local_addr = socket.local_addr()?; self.url .set_port(Some(local_addr.port())) .map_err(|_| anyhow::anyhow!("failed to update udp listener port for {}", self.url))?; + self.request.bind.local_addr = Some(local_addr); let layer = Arc::new(UdpSessionLayer::new_with_stun_responder( socket.clone(), diff --git a/easytier-core/src/socket/udp/session.rs b/easytier-core/src/socket/udp/session.rs index 4b170292..4f0f02e0 100644 --- a/easytier-core/src/socket/udp/session.rs +++ b/easytier-core/src/socket/udp/session.rs @@ -541,7 +541,7 @@ pub(crate) struct UdpSessionCleanup { session_close: Option, shutdown: Option>, tasks: Vec>, - pub(super) layer_guard: Option>, + layer_guard: Arc>>>, } impl std::fmt::Debug for UdpSessionCleanup { @@ -550,7 +550,10 @@ impl std::fmt::Debug for UdpSessionCleanup { .field("has_session_close", &self.session_close.is_some()) .field("has_shutdown", &self.shutdown.is_some()) .field("tasks", &self.tasks.len()) - .field("has_layer_guard", &self.layer_guard.is_some()) + .field( + "has_layer_guard", + &self.layer_guard.lock().unwrap().is_some(), + ) .finish() } } @@ -594,10 +597,12 @@ impl UdpSession { rings.session_send_rx, close.clone(), )); + let layer_guard = Arc::new(StdMutex::new(None)); let shutdown_task = tokio::spawn(close_udp_session_on_shutdown( shutdown, close.clone(), send_task.abort_handle(), + layer_guard.clone(), )); Self { @@ -612,7 +617,7 @@ impl UdpSession { session_close: Some(close), shutdown: None, tasks: vec![send_task, shutdown_task], - layer_guard: None, + layer_guard, }, } } @@ -621,7 +626,12 @@ impl UdpSession { where T: Send + Sync + 'static, { - self._cleanup.layer_guard = Some(Box::new(layer_guard)); + let mut guard = self._cleanup.layer_guard.lock().unwrap(); + // Accept may finish after shutdown. Synchronize with cleanup so a + // closed session cannot retain a failed listener's socket. + if !*self.closed.borrow() { + *guard = Some(Box::new(layer_guard)); + } } pub(crate) fn into_tunnel_parts(self) -> UdpSessionTunnelParts { @@ -899,12 +909,26 @@ async fn close_udp_session_on_shutdown( mut shutdown: watch::Receiver, close: UdpSessionClose, send_task: tokio::task::AbortHandle, + layer_guard: Arc>>>, ) { - if !*shutdown.borrow() { - let _ = shutdown.changed().await; + let mut closed = close.close.subscribe(); + if !*closed.borrow() { + if *shutdown.borrow() { + close.close(); + } else { + tokio::select! { + biased; + _ = closed.changed() => {}, + _ = shutdown.changed() => { + if !*closed.borrow() { + close.close(); + } + }, + } + } } - close.close(); send_task.abort(); + layer_guard.lock().unwrap().take(); } pub(super) fn dispatch_payload_to_session( @@ -934,6 +958,31 @@ pub(super) fn dispatch_payload_to_session( mod tests { use super::*; + #[tokio::test] + async fn closed_session_cleanup_preserves_replacement_with_same_key() { + let sessions = Arc::new(DashMap::new()); + let key = UdpSessionKey::new("127.0.0.1:12001".parse().unwrap(), 1); + let old = create_udp_session_rings(); + sessions.insert(key, udp_session_registry_entry(&old)); + let close = UdpSessionClose::easy_tier(key, old.close_tx.clone(), sessions.clone()); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + let send_task = tokio::spawn(std::future::pending::<()>()); + let mut cleanup = Box::pin(close_udp_session_on_shutdown( + shutdown_rx, + close.clone(), + send_task.abort_handle(), + Arc::new(StdMutex::new(None)), + )); + assert!(futures::poll!(&mut cleanup).is_pending()); + close.close(); + let replacement = create_udp_session_rings(); + sessions.insert(key, udp_session_registry_entry(&replacement)); + shutdown_tx.send(true).unwrap(); + cleanup.await; + assert!(sessions.contains_key(&key)); + assert!(!*replacement.close_rx.borrow()); + } + impl UdpSessionClose { fn signal_only(close: watch::Sender) -> Self { Self { diff --git a/easytier-core/src/socket/udp/tests.rs b/easytier-core/src/socket/udp/tests.rs index db3070e5..d3ba3f4f 100644 --- a/easytier-core/src/socket/udp/tests.rs +++ b/easytier-core/src/socket/udp/tests.rs @@ -1459,12 +1459,16 @@ async fn udp_session_recv_loop_error_closes_registered_sessions() { let socket = Arc::new(MockVirtualUdpSocket::new(local_addr, Vec::new())); let sessions = Arc::new(DashMap::new()); let (session_shutdown_tx, session_shutdown_rx) = watch::channel(false); - let session = create_test_easy_tier_mux_session_with_shutdown( + let mut session = create_test_easy_tier_mux_session_with_shutdown( socket.clone(), UdpSessionKey::new(peer_addr, conn_id), sessions.clone(), session_shutdown_rx, ); + let guard = Arc::new(()); + let weak_guard = Arc::downgrade(&guard); + session.keep_layer_alive(guard); + assert!(weak_guard.upgrade().is_some()); let pending_connects = Arc::new(DashMap::new()); let classified_sessions = Arc::new(DashMap::new()); let classified_accepts = create_classified_udp_session_accepts(); @@ -1507,6 +1511,18 @@ async fn udp_session_recv_loop_error_closes_registered_sessions() { .await .unwrap(); assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof); + + tokio::time::timeout(Duration::from_secs(1), async { + while weak_guard.upgrade().is_some() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let guard = Arc::new(()); + let weak_guard = Arc::downgrade(&guard); + session.keep_layer_alive(guard); + assert!(weak_guard.upgrade().is_none()); } #[tokio::test] @@ -2088,12 +2104,15 @@ impl VirtualUdpSocketFactory for MockVirtualUdpSocketFactory { async fn bind_udp(&self, options: UdpBindOptions) -> anyhow::Result> { self.bind_options.lock().unwrap().push(options.clone()); - let local_addr = options.local_addr.unwrap_or_else(|| { - SocketAddr::from(( - [127, 0, 0, 1], - self.next_port.fetch_add(1, Ordering::Relaxed), - )) - }); + let local_addr = options + .local_addr + .filter(|addr| addr.port() != 0) + .unwrap_or_else(|| { + SocketAddr::from(( + [127, 0, 0, 1], + self.next_port.fetch_add(1, Ordering::Relaxed), + )) + }); let incoming = self .incoming .lock() @@ -2148,6 +2167,59 @@ async fn udp_session_socket_listener_builds_port_bound_bind_options() { )); } +#[tokio::test] +async fn classified_accept_observes_receive_failure_before_and_during_wait() { + for protocol in [UdpSessionProtocol::WireGuard, UdpSessionProtocol::Quic] { + for failed_before_accept in [false, true] { + let socket = Arc::new(MockVirtualUdpSocket::new( + "127.0.0.1:12000".parse().unwrap(), + Vec::new(), + )); + let layer = UdpSessionLayer::new(socket); + if failed_before_accept { + // The mux sender closes when the receive task has exited. + assert!( + tokio::time::timeout(Duration::from_secs(1), layer.accept()) + .await + .unwrap() + .is_err() + ); + } + let mut accept = Box::pin(layer.accept_classified_session(protocol)); + if !failed_before_accept { + assert!(futures::poll!(&mut accept).is_pending()); + } + let error = tokio::time::timeout(Duration::from_secs(1), accept) + .await + .unwrap() + .unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof); + } + } +} + +#[tokio::test] +async fn udp_listener_releases_failed_binding_and_keeps_assigned_port() { + let factory = Arc::new(MockVirtualUdpSocketFactory::new(13000)); + let mut listener = UdpSessionSocketListener::new( + "udp://127.0.0.1:0".parse().unwrap(), + "127.0.0.1:0".parse().unwrap(), + factory.clone(), + ); + listener.listen().await.unwrap(); + let bound_addr = listener.bound_socket().unwrap().local_addr().unwrap(); + let old_socket = Arc::downgrade(&listener.bound_socket().unwrap()); + factory.sockets.lock().unwrap().clear(); + assert!(listener.accept().await.is_err()); + listener.listen().await.unwrap(); + assert!(old_socket.upgrade().is_none()); + assert_eq!( + listener.bound_socket().unwrap().local_addr().unwrap(), + bound_addr + ); + assert_eq!(factory.bind_options().len(), 2); +} + #[tokio::test] async fn udp_session_socket_listener_accepts_easy_tier_mux_session() { let local_addr = SocketAddr::from(([127, 0, 0, 1], 11010)); diff --git a/easytier-web/src/client_manager/listener_tests.rs b/easytier-web/src/client_manager/listener_tests.rs index d298f800..5ec436e9 100644 --- a/easytier-web/src/client_manager/listener_tests.rs +++ b/easytier-web/src/client_manager/listener_tests.rs @@ -7,6 +7,39 @@ use easytier_core::connectivity::protocol::raw::TunnelDialer; use std::sync::atomic::AtomicUsize; use tokio::{io::AsyncReadExt, net::TcpStream, sync::Notify, time::timeout}; +#[derive(Debug)] +struct FailingListener { + inner: RuntimeRpcListener, + attempts: Arc>>, + errors_remaining: usize, + error_kind: std::io::ErrorKind, +} + +#[async_trait::async_trait] +impl SocketListener for FailingListener { + type Accepted = Box; + + async fn listen(&mut self) -> anyhow::Result<()> { + self.inner.listen().await + } + + async fn accept(&mut self) -> anyhow::Result { + self.attempts + .lock() + .unwrap() + .push(tokio::time::Instant::now()); + if self.errors_remaining > 0 { + self.errors_remaining -= 1; + return Err(std::io::Error::from(self.error_kind).into()); + } + self.inner.accept().await + } + + fn local_url(&self) -> url::Url { + self.inner.local_url() + } +} + #[derive(Debug)] struct TestListener { inner: RuntimeRpcListener, @@ -61,6 +94,66 @@ async fn manager() -> ClientManager { ) } +#[tokio::test] +async fn accept_errors_rebind_listener_with_fixed_delay() { + for error_kind in [std::io::ErrorKind::OutOfMemory, std::io::ErrorKind::Other] { + let mut manager = manager().await; + let attempts = Arc::new(std::sync::Mutex::new(Vec::new())); + let url = manager + .add_listener(FailingListener { + inner: runtime_rpc_listener("127.0.0.1:0".parse().unwrap()), + attempts: attempts.clone(), + errors_remaining: 2, + error_kind, + }) + .await + .unwrap(); + wait_until(|| attempts.lock().unwrap().len() >= 3).await; + assert!(manager.is_running()); + + let client = runtime_rpc_dialer(url).connect().await.unwrap(); + let _client = timeout( + Duration::from_secs(1), + web_security::upgrade_client_tunnel(client), + ) + .await + .expect("listener did not recover after accept errors") + .unwrap(); + wait_until(|| manager.client_sessions.len() == 1).await; + { + let attempts = attempts.lock().unwrap(); + assert!(attempts.len() >= 3); + for pair in attempts[..3].windows(2) { + assert!(pair[1] - pair[0] >= Duration::from_millis(100)); + } + } + assert!(manager.is_running()); + drop(manager); + wait_until(|| Arc::strong_count(&attempts) == 1).await; + } +} + +#[tokio::test] +async fn manager_shutdown_cancels_accept_error_retries() { + let mut manager = manager().await; + let attempts = Arc::new(std::sync::Mutex::new(Vec::new())); + let url = manager + .add_listener(FailingListener { + inner: runtime_rpc_listener("127.0.0.1:0".parse().unwrap()), + attempts: attempts.clone(), + errors_remaining: usize::MAX, + error_kind: std::io::ErrorKind::Other, + }) + .await + .unwrap(); + wait_until(|| !attempts.lock().unwrap().is_empty()).await; + assert!(manager.is_running()); + drop(manager); + wait_until(|| Arc::strong_count(&attempts) == 1).await; + let addr = url.socket_addrs(|| None).unwrap()[0]; + assert!(TcpStream::connect(addr).await.is_err()); +} + #[tokio::test] async fn slow_handshakes_do_not_block_other_clients() { let mut manager = manager().await; diff --git a/easytier-web/src/client_manager/mod.rs b/easytier-web/src/client_manager/mod.rs index 71ee2456..6c7c4c3e 100644 --- a/easytier-web/src/client_manager/mod.rs +++ b/easytier-web/src/client_manager/mod.rs @@ -199,7 +199,6 @@ impl ClientManager { self.listeners_cnt.fetch_add(1, Ordering::Relaxed); let sessions = self.client_sessions.clone(); let storage = self.storage.weak_ref(); - let listeners_cnt = self.listeners_cnt.clone(); let next_session_epoch = self.next_session_epoch.clone(); let geoip_db = self.geoip_db.clone(); let heartbeat_policy = self.heartbeat_policy; @@ -207,15 +206,15 @@ impl ClientManager { let webhook_config = self.webhook_config.clone(); self.tasks.spawn(async move { let mut handshakes = JoinSet::new(); - 'accept: loop { + loop { // Some listeners include a WebSocket upgrade in accept(). Keep // that future alive while processing completed handshakes. - let accepting = listener.accept(); + let accepting = easytier_core::socket::retry::accept_with_rebind(&mut listener); tokio::pin!(accepting); loop { let result = tokio::select! { accepted = &mut accepting, if handshakes.len() < MAX_PENDING_HANDSHAKES => { - let Ok(tunnel) = accepted else { break 'accept }; + let tunnel = accepted; handshakes.spawn(web_security::accept_or_upgrade_server_tunnel(tunnel)); break; } @@ -258,7 +257,6 @@ impl ClientManager { session.mark_route_ready(); } } - listeners_cnt.fetch_sub(1, Ordering::Relaxed); }); Ok(local_url) diff --git a/easytier/src/socket/udp_src/windows.rs b/easytier/src/socket/udp_src/windows.rs index 813d429f..d489d35c 100644 --- a/easytier/src/socket/udp_src/windows.rs +++ b/easytier/src/socket/udp_src/windows.rs @@ -46,6 +46,14 @@ pub(crate) async fn recv_from_with_dst_ip( loop { match recv_from_with_dst_ip_once(socket, buf) { Err(err) if err.kind() == io::ErrorKind::Interrupted => continue, + // WSARecvMsg consumed this oversized datagram. Discard it + // just as the Unix adapter discards MSG_TRUNC packets. + Err(err) + if err.raw_os_error() + == Some(windows::Win32::Networking::WinSock::WSAEMSGSIZE.0) => + { + continue; + } ret => break ret, } } diff --git a/easytier/src/tunnel/wireguard.rs b/easytier/src/tunnel/wireguard.rs index d8d80eee..6821eb6d 100644 --- a/easytier/src/tunnel/wireguard.rs +++ b/easytier/src/tunnel/wireguard.rs @@ -609,12 +609,8 @@ impl WgTunnelListener { }; tracing::info!("New peer: {}", addr); - let mut wg = WgPeer::new( - Box::new(session_listener.clone()), - session, - config.clone(), - addr, - ); + // Accepted sessions already keep their UDP layer alive until close. + let mut wg = WgPeer::new(Box::new(()), session, config.clone(), addr); let (stream, sink) = wg.start_and_get_tunnel().split(); wg.spawn_session_recv_task(None); let tunnel = Box::new(TunnelWrapper::new( diff --git a/easytier/src/vpn_portal/wireguard.rs b/easytier/src/vpn_portal/wireguard.rs index 5fa52c1f..8e0525be 100644 --- a/easytier/src/vpn_portal/wireguard.rs +++ b/easytier/src/vpn_portal/wireguard.rs @@ -9,7 +9,7 @@ mod engine; use std::{ collections::{BTreeMap, BTreeSet}, - fmt, + fmt, io, net::{Ipv6Addr, SocketAddr, SocketAddrV6}, sync::{Arc, Mutex as StdMutex, RwLock, Weak}, }; @@ -32,16 +32,51 @@ use sha2::Sha256; use tokio::{sync::mpsc, task::JoinSet}; use crate::{ - common::global_ctx::ArcGlobalCtx, - host_runtime::{NativeHostRuntime, native_host_runtime}, + common::global_ctx::ArcGlobalCtx, host_runtime::native_host_runtime, socket::udp::RuntimeUdpSocket, }; use self::engine::{DerivedClient, PortalEngine}; +// Endpoints retain this family-specific slot, not a socket generation. Holding +// the read guard during send lets the reader release the old binding on failure. +struct PortalSocket { + socket: tokio::sync::RwLock>>, + options: UdpBindOptions, +} + +impl PortalSocket { + async fn send_to(&self, data: &[u8], remote: SocketAddr) -> io::Result { + let socket = self.socket.read().await; + match socket.as_ref() { + Some(socket) => socket.send_to(data, remote).await, + None => Err(io::Error::new( + io::ErrorKind::WouldBlock, + "portal socket is rebinding", + )), + } + } + + async fn rebind(&self) { + self.socket.write().await.take(); + loop { + easytier_core::socket::retry::wait().await; + match native_host_runtime().bind_udp(self.options.clone()).await { + Ok(socket) => { + *self.socket.write().await = Some(socket); + return; + } + Err(error) => { + tracing::warn!(target: "CORE", %error, address = ?self.options.local_addr, "WireGuard portal UDP bind failed; retrying") + } + } + } + } +} + struct WireGuardPortalListener { url: url::Url, - sockets: Vec>, + sockets: Vec>, receiver: mpsc::UnboundedReceiver, engine: Arc, tasks: JoinSet>, @@ -66,15 +101,24 @@ impl SocketListener for WireGuardPortalListener { return Ok(()); } self.listened = true; - for socket in &self.sockets { - let socket = socket.clone(); + for socket in self.sockets.drain(..) { let engine = self.engine.clone(); self.tasks.spawn(async move { loop { - let datagram = socket - .recv_session_datagram() - .await - .context("WireGuard portal UDP receive failed")?; + let datagram = { + // Do not hold a read lock while waiting for traffic, or + // retain this generation while trying to rebind its port. + let current = socket.socket.read().await.as_ref().unwrap().clone(); + current.recv_session_datagram().await + }; + let datagram = match datagram { + Ok(datagram) => datagram, + Err(error) => { + tracing::warn!(target: "CORE", %error, address = ?socket.options.local_addr, "WireGuard portal UDP socket failed; rebinding"); + socket.rebind().await; + continue; + } + }; engine .handle_datagram(socket.clone(), datagram.remote_addr, &datagram.payload) .await; @@ -191,20 +235,21 @@ impl WireGuardPortalHost { async fn bind_socket( &self, - runtime: &NativeHostRuntime, address: SocketAddr, only_v6: bool, - ) -> anyhow::Result> { + ) -> anyhow::Result> { let context = SocketContext::default() .with_socket_mark(self.global_ctx.get_flags().socket_mark) .with_netns(self.global_ctx.net_ns.name().map(NetNamespace::new)); - runtime - .bind_udp( - UdpBindOptions::port_bound_listener(address) - .with_context(context) - .with_only_v6(only_v6), - ) - .await + let mut options = UdpBindOptions::port_bound_listener(address) + .with_context(context) + .with_only_v6(only_v6); + let socket = native_host_runtime().bind_udp(options.clone()).await?; + options.local_addr = Some(socket.local_addr()?); + Ok(Arc::new(PortalSocket { + socket: tokio::sync::RwLock::new(Some(socket)), + options, + })) } async fn apply_client_updates(&self, clients: &[PortalClientConfig]) -> anyhow::Result<()> { @@ -285,16 +330,15 @@ impl PortalHost for WireGuardPortalHost { .setup .as_ref() .map_err(|error| anyhow::anyhow!(error.clone()))?; - let runtime = native_host_runtime(); let primary = self - .bind_socket(&runtime, self.config.wireguard_listen, false) + .bind_socket(self.config.wireguard_listen, false) .await .context("failed to bind WireGuard VPN portal")?; - let local = primary.local_addr()?; - let mut sockets = vec![primary.clone()]; + let local = primary.options.local_addr.unwrap(); + let mut sockets = vec![primary]; if let Some(v6_address) = secondary_ipv6_bind_address(self.config.wireguard_listen, local.port()) - && let Ok(v6) = self.bind_socket(&runtime, v6_address, true).await + && let Ok(v6) = self.bind_socket(v6_address, true).await { sockets.push(v6); } @@ -391,6 +435,58 @@ pub(crate) fn test_wireguard_keys( mod tests { use super::*; + #[tokio::test] + async fn portal_socket_rebinds_while_endpoints_keep_their_reference() { + let client = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let mut options = UdpBindOptions::port_bound_listener("127.0.0.1:0".parse().unwrap()) + .with_need_protect(false); + let initial = native_host_runtime() + .bind_udp(options.clone()) + .await + .unwrap(); + let local = initial.local_addr().unwrap(); + options.local_addr = Some(local); + let old_socket = Arc::downgrade(&initial); + let socket = Arc::new(PortalSocket { + socket: tokio::sync::RwLock::new(Some(initial)), + options, + }); + let endpoint = socket.clone(); + let mut rebind = Box::pin(socket.rebind()); + assert!(futures::poll!(&mut rebind).is_pending()); + assert!(old_socket.upgrade().is_none()); + assert_eq!( + endpoint + .send_to(b"during-rebind", client.local_addr().unwrap()) + .await + .unwrap_err() + .kind(), + io::ErrorKind::WouldBlock + ); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(10), rebind) + .await + .is_err() + ); + assert!(socket.socket.read().await.is_none()); + socket.rebind().await; + + endpoint + .send_to(b"after-rebind", client.local_addr().unwrap()) + .await + .unwrap(); + let mut buffer = [0; 32]; + let (len, source) = tokio::time::timeout( + std::time::Duration::from_secs(1), + client.recv_from(&mut buffer), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(&buffer[..len], b"after-rebind"); + assert_eq!(source, local); + } + #[test] fn named_wireguard_keys_are_stable_and_client_scoped() { let master = [7; 32]; diff --git a/easytier/src/vpn_portal/wireguard/engine.rs b/easytier/src/vpn_portal/wireguard/engine.rs index a10a71cb..069bfa8e 100644 --- a/easytier/src/vpn_portal/wireguard/engine.rs +++ b/easytier/src/vpn_portal/wireguard/engine.rs @@ -18,10 +18,7 @@ use boringtun::{ }, x25519::{PublicKey, StaticSecret}, }; -use easytier_core::{ - gateway::vpn_portal::{PortalClientConfig, PortalSession}, - socket::udp::VirtualUdpSocket, -}; +use easytier_core::gateway::vpn_portal::{PortalClientConfig, PortalSession}; use rand::rngs::OsRng; use tokio::{ sync::{Mutex, mpsc, watch}, @@ -29,7 +26,7 @@ use tokio::{ }; use tokio_util::sync::CancellationToken; -use crate::socket::udp::RuntimeUdpSocket; +use super::PortalSocket; const MIN_WIREGUARD_PACKET_CAPACITY: usize = 148; // We pre-verify through this shared limiter and BoringTun verifies again inside // each Tunn. Doubling the threshold preserves the intended 100 datagrams/s @@ -95,12 +92,12 @@ impl EngineSlots { #[derive(Clone)] struct Endpoint { - socket: Arc, + socket: Arc, remote: SocketAddr, } impl ClientSession { - fn update_endpoint(&mut self, socket: Arc, remote: SocketAddr) { + fn update_endpoint(&mut self, socket: Arc, remote: SocketAddr) { let changed = self .endpoint .as_ref() @@ -223,7 +220,7 @@ impl PortalEngine { pub(super) async fn handle_datagram( self: &Arc, - socket: Arc, + socket: Arc, remote: SocketAddr, datagram: &[u8], ) { @@ -375,7 +372,7 @@ impl PortalEngine { fn new_session( self: &Arc, slot: &Arc, - socket: Arc, + socket: Arc, remote: SocketAddr, ) -> ClientSession { let generation = slot.next_generation.fetch_add(1, Ordering::Relaxed); @@ -385,7 +382,7 @@ impl PortalEngine { let engine = Arc::downgrade(self); let slot_for_task = Arc::downgrade(slot); let mut tasks = JoinSet::new(); - tasks.spawn(async move { + let forward_packets = async move { while let Some(payload) = to_client.recv().await { let Some(engine) = engine.upgrade() else { return; @@ -400,6 +397,13 @@ impl PortalEngine { if let (Some(engine), Some(slot)) = (engine.upgrade(), slot_for_task.upgrade()) { engine.expire_if_current(slot, generation).await; } + }; + let cancel = self.cancel.clone(); + tasks.spawn(async move { + tokio::select! { + _ = cancel.cancelled() => {}, + _ = forward_packets => {}, + } }); ClientSession { generation,