From 318a0b0686db9ee4b01e9c3c3cbaefe01315fcae Mon Sep 17 00:00:00 2001 From: KKRainbow <5665404+KKRainbow@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:23:10 +0800 Subject: [PATCH] fix(gateway): make port-forward reload preserve live listeners (#2710) Reuse listeners by protocol and bind address when their destination changes, keeping established TCP streams connected to the old target. Prepare all new sockets before persistence or configuration publication so bind and storage failures leave committed forwarding usable. Commit prepared resources with shared TOML and runtime publication. Serialize UDP target changes with flow insertion, and track response tasks through eviction so removal waits for every socket reference. Keep cleanup and its operation lock alive if the caller stops waiting. Add TCP and UDP regressions for target changes, failed updates, retries, established streams, immediate rebinds, and cancelled cleanup waiters. --- easytier-core/src/gateway/mod.rs | 2 +- easytier-core/src/gateway/port_forward.rs | 393 +++++++++++++----- .../src/gateway/port_forward/reload_tests.rs | 290 +++++++++++++ easytier-core/src/host/testkit.rs | 65 ++- easytier-core/src/instance/mod.rs | 55 ++- easytier-core/src/instance/tests.rs | 262 ++++++++++++ .../src/management/full/config_patch.rs | 22 +- easytier-core/src/peers/peer_manager.rs | 15 +- 8 files changed, 982 insertions(+), 122 deletions(-) create mode 100644 easytier-core/src/gateway/port_forward/reload_tests.rs diff --git a/easytier-core/src/gateway/mod.rs b/easytier-core/src/gateway/mod.rs index d85ba0eb..748e4046 100644 --- a/easytier-core/src/gateway/mod.rs +++ b/easytier-core/src/gateway/mod.rs @@ -27,6 +27,6 @@ pub use dataplane::{ DataPlaneTcpListener, DataPlaneTcpStream, DataPlaneUdpSocket, }; #[cfg(feature = "proxy-smoltcp-stack")] -pub(crate) use port_forward::PortForwardAdapter; +pub(crate) use port_forward::{PortForwardAdapter, PreparedPortForwards}; #[cfg(feature = "proxy-smoltcp-stack")] pub(crate) use socks5::Socks5GatewayAdapter; diff --git a/easytier-core/src/gateway/port_forward.rs b/easytier-core/src/gateway/port_forward.rs index a3a0a88f..dd02e12a 100644 --- a/easytier-core/src/gateway/port_forward.rs +++ b/easytier-core/src/gateway/port_forward.rs @@ -1,6 +1,7 @@ //! Host port-forward adapter backed by the data-plane runtime. use std::{ + collections::HashMap, net::SocketAddr, sync::{ Arc, @@ -15,11 +16,11 @@ use quanta::Instant; use tokio::{ select, sync::{Mutex, OwnedMutexGuard, OwnedSemaphorePermit, Semaphore}, - task::JoinSet, + task::{JoinHandle, JoinSet}, }; use tokio_util::{ - sync::{CancellationToken, DropGuard}, - task::AbortOnDropHandle, + sync::CancellationToken, + task::{AbortOnDropHandle, TaskTracker}, }; use crate::{ @@ -94,11 +95,119 @@ struct UdpClientReservation { admission: OwnedMutexGuard<()>, } +struct PortForwardListener { + config: Arc>, + cancel: CancellationToken, + task: AbortOnDropHandle<()>, + udp_tasks: TaskTracker, +} + +enum PortForwardSocket { + Tcp(Arc), + Udp(Arc<::Socket>), +} + +/// Bound additions stay private until persistence and validation have succeeded. +/// Dropping a prepared update leaves every existing listener and target intact. +pub(crate) struct PreparedPortForwards +where + H: VirtualTcpSocketFactory + VirtualTcpListenerFactory + VirtualUdpSocketFactory, +{ + adapter: Arc>, + configs: HashMap<(String, SocketAddr), PortForwardConfig>, + sockets: Vec<(PortForwardConfig, PortForwardSocket)>, + consumer_lease: Option, + // Release prepared sockets before allowing another update to bind them. + operation: OwnedMutexGuard<()>, +} + +impl PreparedPortForwards +where + H: VirtualTcpSocketFactory + VirtualTcpListenerFactory + VirtualUdpSocketFactory, +{ + /// All fallible work is complete. Publish targets and start the bound sockets + /// without yielding, then retain the operation lock until removals release + /// their sockets, even if the caller stops waiting for this cleanup. + pub(crate) fn commit(self) -> JoinHandle<()> { + let Self { + adapter, + operation, + configs, + sockets, + consumer_lease, + } = self; + let _runtime = adapter.runtime_handle.as_ref().map(|handle| handle.enter()); + let mut listeners = adapter.listeners.lock().unwrap(); + let mut removed = Vec::new(); + let mut response_tasks = Vec::new(); + let removed_keys = listeners + .keys() + .filter(|key| !configs.contains_key(*key)) + .cloned() + .collect::>(); + for key in removed_keys { + let listener = listeners.remove(&key).unwrap(); + listener.cancel.cancel(); + listener.task.abort(); + { + let current = listener.config.lock().unwrap(); + tracing::info!(?current, "port-forward removed by runtime config reload"); + adapter.remove_udp_clients(¤t, &mut response_tasks); + } + removed.push(listener); + } + for (key, listener) in listeners.iter() { + let mut current = listener.config.lock().unwrap(); + if let Some(next) = configs.get(key) + && *current != *next + { + adapter.remove_udp_clients(¤t, &mut response_tasks); + *current = next.clone(); + adapter + .events + .emit(CoreEvent::GatewayPortForwardAdded(next.clone())); + } + } + for (cfg, socket) in sockets { + let key = (cfg.proto.clone(), cfg.bind_addr); + let listener = match socket { + PortForwardSocket::Tcp(socket) => adapter.add_tcp_port_forward(&cfg, socket), + PortForwardSocket::Udp(socket) => adapter.add_udp_port_forward(&cfg, socket), + }; + listeners.insert(key, listener); + adapter.events.emit(CoreEvent::GatewayPortForwardAdded(cfg)); + } + drop(listeners); + let mut lease = adapter.consumer_lease.lock().unwrap(); + if let Some(consumer_lease) = consumer_lease { + lease.replace(consumer_lease); + } + let released_lease = if configs.is_empty() { + lease.take() + } else { + None + }; + drop(lease); + tokio::spawn(async move { + let _operation = operation; + for listener in removed { + let _ = listener.task.await; + listener.udp_tasks.close(); + listener.udp_tasks.wait().await; + } + for task in response_tasks { + let _ = task.await; + } + drop(released_lease); + }) + } +} + pub(crate) struct PortForwardAdapter where H: VirtualTcpSocketFactory + VirtualTcpListenerFactory + VirtualUdpSocketFactory, { - operation: Mutex<()>, + operation: Arc>, started: AtomicBool, runtime_handle: Option, runtime_config: CoreRuntimeConfigStore, @@ -107,12 +216,12 @@ where socket_context: SocketContext, events: Arc, tasks: Arc>>, - cancel_tokens: Arc>, + listeners: std::sync::Mutex>, udp_clients: Arc>>>, udp_response_tasks: Arc>>, udp_client_admission: Arc>, udp_client_slots: Arc, - consumer_lease: Mutex>, + consumer_lease: std::sync::Mutex>, } impl PortForwardAdapter @@ -128,7 +237,7 @@ where runtime_handle: Option, ) -> Arc { Arc::new(Self { - operation: Mutex::new(()), + operation: Arc::new(Mutex::new(())), started: AtomicBool::new(false), runtime_handle, runtime_config, @@ -137,17 +246,17 @@ where socket_context, events, tasks: Arc::new(std::sync::Mutex::new(JoinSet::new())), - cancel_tokens: Arc::new(DashMap::new()), + listeners: std::sync::Mutex::new(HashMap::new()), udp_clients: Arc::new(DashMap::new()), udp_response_tasks: Arc::new(DashMap::new()), udp_client_admission: Arc::new(Mutex::new(())), udp_client_slots: Arc::new(Semaphore::new(MAX_ACTIVE_UDP_CLIENTS)), - consumer_lease: Mutex::new(None), + consumer_lease: std::sync::Mutex::new(None), }) } - pub(crate) async fn start(&self) -> anyhow::Result<()> { - let _operation = self.operation.lock().await; + pub(crate) async fn start(self: &Arc) -> anyhow::Result<()> { + let operation = self.operation.clone().lock_owned().await; if self.started.load(Ordering::Acquire) { return Ok(()); } @@ -163,110 +272,157 @@ where .gateway .port_forwards .clone(); - if let Err(error) = self.apply_port_forwards(&cfgs).await { - self.stop_inner().await; - return Err(error); - } + let prepared = match self.prepare_inner(&cfgs, operation).await { + Ok(prepared) => prepared, + Err(error) => { + self.stop_inner().await; + return Err(error); + } + }; + let cleanup = prepared.commit(); self.started.store(true, Ordering::Release); + let _ = cleanup.await; Ok(()) } - pub(crate) async fn reload(self: &Arc, cfgs: &[PortForwardConfig]) -> anyhow::Result<()> { + pub(crate) async fn prepare( + self: &Arc, + cfgs: &[PortForwardConfig], + ) -> anyhow::Result>> { if let Some(runtime) = &self.runtime_handle { let adapter = self.clone(); let cfgs = cfgs.to_vec(); - // Binding and spawning must both happen on the instance executor, - // including reloads reached through management config patches. + // Binding must happen on the instance executor, including reloads + // reached through management config patches. return AbortOnDropHandle::new( - runtime.spawn(async move { adapter.reload_on_runtime(&cfgs).await }), + runtime.spawn(async move { adapter.prepare_on_runtime(&cfgs).await }), ) .await - .map_err(|error| anyhow::anyhow!("port-forward reload task failed: {error}"))?; + .map_err(|error| anyhow::anyhow!("port-forward prepare task failed: {error}"))?; } - self.reload_on_runtime(cfgs).await + self.prepare_on_runtime(cfgs).await } - async fn reload_on_runtime(&self, cfgs: &[PortForwardConfig]) -> anyhow::Result<()> { - let _operation = self.operation.lock().await; + async fn prepare_on_runtime( + self: &Arc, + cfgs: &[PortForwardConfig], + ) -> anyhow::Result>> { + let operation = self.operation.clone().lock_owned().await; if !self.started.load(Ordering::Acquire) { - return Ok(()); + return Ok(None); } - self.apply_port_forwards(cfgs).await + self.prepare_inner(cfgs, operation).await.map(Some) } - async fn apply_port_forwards(&self, cfgs: &[PortForwardConfig]) -> anyhow::Result<()> { + async fn prepare_inner( + self: &Arc, + cfgs: &[PortForwardConfig], + operation: OwnedMutexGuard<()>, + ) -> anyhow::Result> { + let mut configs = HashMap::new(); for cfg in cfgs { - if !matches!(cfg.proto.to_lowercase().as_str(), "tcp" | "udp") { + let mut cfg = cfg.clone(); + cfg.proto.make_ascii_lowercase(); + if !matches!(cfg.proto.as_str(), "tcp" | "udp") { anyhow::bail!( "unsupported protocol: {}, only support udp / tcp", cfg.proto ); } - } - - if !cfgs.is_empty() { - let mut consumer_lease = self.consumer_lease.lock().await; - if consumer_lease.is_none() { - consumer_lease.replace(self.data_plane.acquire_consumer_lease()?); + let key = (cfg.proto.clone(), cfg.bind_addr); + if let Some(previous) = configs.insert(key, cfg.clone()) { + anyhow::ensure!( + previous == cfg, + "conflicting port-forward targets for {} {}", + cfg.proto, + cfg.bind_addr + ); } } - - self.cancel_tokens.retain(|current, _| { - let keep = cfgs.iter().any(|next| { - if next.dst_addr.ip().is_unspecified() { - current.bind_addr == next.bind_addr && current.proto == next.proto - } else { - current == next - } - }); - if !keep { - tracing::info!(?current, "port-forward removed by runtime config reload"); - } - keep - }); - self.udp_clients - .retain(|key, _| self.cancel_tokens.contains_key(&key.forward)); - self.udp_response_tasks - .retain(|key, _| self.udp_clients.contains_key(key)); + let consumer_lease = if !configs.is_empty() && self.consumer_lease.lock().unwrap().is_none() + { + Some(self.data_plane.acquire_consumer_lease()?) + } else { + None + }; + let mut sockets = Vec::new(); for cfg in cfgs { - if !self.cancel_tokens.contains_key(cfg) { - self.add_port_forward(cfg.clone()).await?; + let mut cfg = cfg.clone(); + cfg.proto.make_ascii_lowercase(); + let key = (cfg.proto.clone(), cfg.bind_addr); + if self.listeners.lock().unwrap().contains_key(&key) + || sockets + .iter() + .any(|(bound, _): &(PortForwardConfig, PortForwardSocket)| { + bound.proto == cfg.proto && bound.bind_addr == cfg.bind_addr + }) + { + continue; } - } - if cfgs.is_empty() { - self.consumer_lease.lock().await.take(); - } - Ok(()) - } - - async fn add_port_forward(&self, cfg: PortForwardConfig) -> anyhow::Result<()> { - match cfg.proto.to_lowercase().as_str() { - "tcp" => self.add_tcp_port_forward(&cfg).await?, - "udp" => self.add_udp_port_forward(&cfg).await?, - _ => { - anyhow::bail!( - "unsupported protocol: {}, only support udp / tcp", - cfg.proto + let socket = if cfg.proto == "tcp" { + PortForwardSocket::Tcp( + self.host + .bind_tcp(self.tcp_listen_options(cfg.bind_addr)) + .await?, ) - } + } else { + PortForwardSocket::Udp( + self.host + .bind_udp( + UdpBindOptions::port_forward(cfg.bind_addr) + .with_context(self.socket_context.clone()), + ) + .await?, + ) + }; + sockets.push((cfg, socket)); } - self.events.emit(CoreEvent::GatewayPortForwardAdded(cfg)); - Ok(()) + Ok(PreparedPortForwards { + adapter: self.clone(), + operation, + configs, + sockets, + consumer_lease, + }) } - async fn add_tcp_port_forward(&self, cfg: &PortForwardConfig) -> anyhow::Result<()> { - let (bind_addr, dst_addr) = (cfg.bind_addr, cfg.dst_addr); + fn tcp_listen_options(&self, bind_addr: SocketAddr) -> TcpListenOptions { let options = TcpListenOptions::port_forward(bind_addr); let bind = options .bind .clone() .with_context(self.socket_context.clone()); - let options = options.with_bind(bind); - let listener = self.host.bind_tcp(options.clone()).await?; + options.with_bind(bind) + } + + fn remove_udp_clients(&self, cfg: &PortForwardConfig, tasks: &mut Vec>) { + self.udp_clients.retain(|key, _| key.forward != *cfg); + let keys = self + .udp_response_tasks + .iter() + .filter(|entry| entry.key().forward == *cfg) + .map(|entry| entry.key().clone()) + .collect::>(); + for key in keys { + if let Some((_, task)) = self.udp_response_tasks.remove(&key) { + task.abort(); + tasks.push(task); + } + } + } + + fn add_tcp_port_forward( + &self, + cfg: &PortForwardConfig, + listener: Arc, + ) -> PortForwardListener { + let (bind_addr, dst_addr) = (cfg.bind_addr, cfg.dst_addr); + let options = self.tcp_listen_options(bind_addr); tracing::info!(?bind_addr, ?dst_addr, "TCP port-forward listener bound"); let cancel = CancellationToken::new(); - self.cancel_tokens - .insert(cfg.clone(), cancel.clone().drop_guard()); + let config = Arc::new(std::sync::Mutex::new(cfg.clone())); + let task_cancel = cancel.clone(); + let task_config = config.clone(); let data_plane = self.data_plane.clone(); let connections = Arc::new(std::sync::Mutex::new(JoinSet::new())); @@ -275,7 +431,8 @@ where "TCP port-forward connections", )); let host = self.host.clone(); - self.tasks.lock().unwrap().spawn(async move { + let task = AbortOnDropHandle::new(tokio::spawn(async move { + let cancel = task_cancel; let mut listener = Some(listener); loop { let Some((incoming, source_addr)) = @@ -283,6 +440,7 @@ where else { break; }; + let dst_addr = task_config.lock().unwrap().dst_addr; tracing::info!( ?bind_addr, ?source_addr, @@ -317,22 +475,27 @@ where ?dst_addr, "TCP port-forward listener task stopped" ); - }); - Ok(()) + })); + PortForwardListener { + config, + cancel, + task, + udp_tasks: TaskTracker::new(), + } } - async fn add_udp_port_forward(&self, cfg: &PortForwardConfig) -> anyhow::Result<()> { - let (bind_addr, dst_addr) = (cfg.bind_addr, cfg.dst_addr); - let forward = cfg.clone(); - let socket = self - .host - .bind_udp( - UdpBindOptions::port_forward(bind_addr).with_context(self.socket_context.clone()), - ) - .await?; + fn add_udp_port_forward( + &self, + cfg: &PortForwardConfig, + socket: Arc<::Socket>, + ) -> PortForwardListener { + let bind_addr = cfg.bind_addr; let cancel = CancellationToken::new(); - self.cancel_tokens - .insert(cfg.clone(), cancel.clone().drop_guard()); + let config = Arc::new(std::sync::Mutex::new(cfg.clone())); + let task_cancel = cancel.clone(); + let task_config = config.clone(); + let udp_tasks = TaskTracker::new(); + let response_tracker = udp_tasks.clone(); let data_plane = self.data_plane.clone(); let host = self.host.clone(); @@ -341,7 +504,8 @@ where let response_tasks = self.udp_response_tasks.clone(); let client_admission = self.udp_client_admission.clone(); let client_slots = self.udp_client_slots.clone(); - self.tasks.lock().unwrap().spawn(async move { + let task = AbortOnDropHandle::new(tokio::spawn(async move { + let cancel = task_cancel; let adapter = UdpFlowFactory { data_plane, host, @@ -360,6 +524,8 @@ where } }, }; + let forward = task_config.lock().unwrap().clone(); + let dst_addr = forward.dst_addr; let key = UdpClientKey { client_addr, forward: forward.clone(), @@ -394,6 +560,13 @@ where continue; }; let UdpClientReservation { slot, admission } = reservation; + // Commit uses this same lock when replacing the target + // and removing its flows. Do not publish an old flow + // after either operation, including after an async open. + let current = task_config.lock().unwrap(); + if *current != forward || cancel.is_cancelled() { + continue; + } let slot = Arc::new(slot); let client = Arc::new(UdpClientInfo { flow: flow.clone(), @@ -408,7 +581,7 @@ where let response_slot = slot; response_tasks.insert( key.clone(), - AbortOnDropHandle::new(tokio::spawn(async move { + AbortOnDropHandle::new(response_tracker.spawn(async move { let _slot = response_slot; let mut buf = vec![0u8; MAX_UDP_PAYLOAD_SIZE]; loop { @@ -435,6 +608,7 @@ where } })), ); + drop(current); drop(admission); client } @@ -444,9 +618,14 @@ where tracing::error!(?error, ?dst_addr, "send UDP data-plane packet failed"); } } - }); + })); - Ok(()) + PortForwardListener { + config, + cancel, + task, + udp_tasks, + } } fn start_udp_reaper(&self) { @@ -468,14 +647,29 @@ where async fn stop_inner(&self) { self.started.store(false, Ordering::Release); + let listeners = std::mem::take(&mut *self.listeners.lock().unwrap()); tracing::info!( - forward_count = self.cancel_tokens.len(), + forward_count = listeners.len(), "port-forward adapter stopping" ); - self.cancel_tokens.clear(); + for listener in listeners.values() { + listener.cancel.cancel(); + listener.task.abort(); + } + // First stop publishers, then remove response handles. The trackers + // include responses whose handles the reaper/eviction already dropped. + let mut trackers = Vec::new(); + for (_, listener) in listeners { + let _ = listener.task.await; + trackers.push(listener.udp_tasks); + } self.udp_response_tasks.clear(); self.udp_clients.clear(); - self.consumer_lease.lock().await.take(); + for tracker in trackers { + tracker.close(); + tracker.wait().await; + } + self.consumer_lease.lock().unwrap().take(); let mut tasks = { let mut tasks = self.tasks.lock().unwrap(); std::mem::replace(&mut *tasks, JoinSet::new()) @@ -639,6 +833,9 @@ where } } +#[cfg(test)] +mod reload_tests; + #[cfg(test)] mod tests { use std::{io, sync::atomic::AtomicUsize}; diff --git a/easytier-core/src/gateway/port_forward/reload_tests.rs b/easytier-core/src/gateway/port_forward/reload_tests.rs new file mode 100644 index 00000000..01a87d83 --- /dev/null +++ b/easytier-core/src/gateway/port_forward/reload_tests.rs @@ -0,0 +1,290 @@ +use std::{io, sync::Weak}; + +use tokio::sync::{Notify, mpsc, oneshot}; + +use super::*; +use crate::{ + config::{IpPrefix, peers::PeerRuntimeSnapshot, runtime::CoreRuntimeConfig}, + host::{ + packet::{HostPacketReceiver, host_packet_channel}, + testkit::TestHost, + }, + peers::peer_manager::{PeerManagerCore, PortablePeerManagerConfig}, + socket::tcp::TcpConnectOptions, +}; + +type Datagram = (Vec, SocketAddr); + +struct ReloadUdpSocket { + address: SocketAddr, + incoming: Option>>, + sent: mpsc::UnboundedSender, +} + +#[async_trait::async_trait] +impl VirtualUdpSocket for ReloadUdpSocket { + fn local_addr(&self) -> io::Result { + Ok(self.address) + } + + async fn send_to(&self, data: &[u8], address: SocketAddr) -> io::Result { + self.sent.send((data.to_vec(), address)).unwrap(); + Ok(data.len()) + } + + async fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { + let Some(incoming) = &self.incoming else { + return std::future::pending().await; + }; + let (data, address) = incoming.lock().await.recv().await.unwrap(); + buf[..data.len()].copy_from_slice(&data); + Ok((data.len(), address)) + } +} + +struct ReloadUdpHost { + incoming: std::sync::Mutex>>, + sent: mpsc::UnboundedSender, + sockets: std::sync::Mutex>>, + open_started: Notify, + open_release: std::sync::Mutex>>, +} + +#[async_trait::async_trait] +impl VirtualTcpSocketFactory for ReloadUdpHost { + type Socket = ::Socket; + + async fn connect_tcp(&self, options: TcpConnectOptions) -> anyhow::Result { + TestHost::default().connect_tcp(options).await + } +} + +#[async_trait::async_trait] +impl VirtualTcpListenerFactory for ReloadUdpHost { + type Listener = ::Listener; + + async fn bind_tcp(&self, options: TcpListenOptions) -> anyhow::Result> { + TestHost::default().bind_tcp(options).await + } +} + +#[async_trait::async_trait] +impl VirtualUdpSocketFactory for ReloadUdpHost { + type Socket = ReloadUdpSocket; + + async fn bind_udp(&self, options: UdpBindOptions) -> anyhow::Result> { + let address = options.local_addr.unwrap(); + let incoming = if address.port() != 0 { + self.incoming.lock().unwrap().take().map(Mutex::new) + } else { + let release = self.open_release.lock().unwrap().take(); + if let Some(release) = release { + self.open_started.notify_one(); + release.await.unwrap(); + } + None + }; + let socket = Arc::new(ReloadUdpSocket { + address, + incoming, + sent: self.sent.clone(), + }); + self.sockets.lock().unwrap().push(Arc::downgrade(&socket)); + Ok(socket) + } +} + +struct UdpReloadTest { + adapter: Arc>, + host: Arc, + incoming: mpsc::UnboundedSender, + sent: mpsc::UnboundedReceiver, + _peer_manager: Arc, + _packets: HostPacketReceiver, +} + +impl UdpReloadTest { + async fn new() -> Self { + let mut runtime = PeerRuntimeSnapshot::default().runtime; + runtime.core.routes.ipv4 = Some(IpPrefix::new("10.144.0.2".parse().unwrap(), 24).unwrap()); + let peer_config = PortablePeerManagerConfig::new(runtime); + let runtime_config = CoreRuntimeConfigStore::new( + CoreRuntimeConfig::default(), + Arc::new(peer_config.snapshot.clone()), + ); + let (packet_sender, packets) = host_packet_channel(); + let peer_manager = + Arc::new(PeerManagerCore::new_portable_for_test(peer_config, packet_sender).unwrap()); + let (incoming, receiver) = mpsc::unbounded_channel(); + let (sent, sent_receiver) = mpsc::unbounded_channel(); + let host = Arc::new(ReloadUdpHost { + incoming: std::sync::Mutex::new(Some(receiver)), + sent, + sockets: std::sync::Mutex::new(Vec::new()), + open_started: Notify::new(), + open_release: std::sync::Mutex::new(None), + }); + let data_plane = DataPlaneRuntime::new( + runtime_config.clone(), + peer_manager.clone(), + None, + host.clone(), + SocketContext::default(), + ); + data_plane.start_runtime().await.unwrap(); + let adapter = PortForwardAdapter::new( + runtime_config, + data_plane, + host.clone(), + SocketContext::default(), + Arc::new(()), + None, + ); + adapter.start().await.unwrap(); + Self { + adapter, + host, + incoming, + sent: sent_receiver, + _peer_manager: peer_manager, + _packets: packets, + } + } + + async fn apply(&self, port: u16) -> PortForwardConfig { + let config = PortForwardConfig { + proto: "udp".to_owned(), + bind_addr: "127.0.0.1:19491".parse().unwrap(), + dst_addr: SocketAddr::from(([10, 144, 0, 2], port)), + }; + self.adapter + .prepare(std::slice::from_ref(&config)) + .await + .unwrap() + .unwrap() + .commit() + .await + .unwrap(); + config + } + + fn send(&self, payload: &[u8]) { + self.incoming + .send((payload.to_vec(), "127.0.0.1:40000".parse().unwrap())) + .unwrap(); + } + + async fn receive(&mut self) -> Datagram { + tokio::time::timeout(Duration::from_secs(2), self.sent.recv()) + .await + .unwrap() + .unwrap() + } + + async fn stop(&self) { + self.adapter.stop().await; + self.adapter.data_plane.stop_runtime().await; + } +} + +#[tokio::test] +async fn udp_target_update_reuses_listener_and_releases_old_flow() { + let mut test = UdpReloadTest::new().await; + let first = test.apply(19381).await; + test.send(b"first"); + assert_eq!(test.receive().await, (b"first".to_vec(), first.dst_addr)); + let second = test.apply(19382).await; + assert!(test.adapter.udp_clients.is_empty()); + { + let sockets = test.host.sockets.lock().unwrap(); + assert_eq!( + sockets.len(), + 2, + "target changes must not bind a new listener" + ); + assert!(sockets[0].upgrade().is_some()); + assert!(sockets[1].upgrade().is_none(), "old flow must be released"); + } + test.send(b"second"); + assert_eq!(test.receive().await, (b"second".to_vec(), second.dst_addr)); + assert_eq!(test.adapter.udp_clients.len(), 1); + test.stop().await; +} + +#[tokio::test] +async fn udp_target_update_during_open_discards_old_flow() { + let mut test = UdpReloadTest::new().await; + test.apply(19381).await; + let (release, waiting) = oneshot::channel(); + *test.host.open_release.lock().unwrap() = Some(waiting); + test.send(b"old"); + tokio::time::timeout(Duration::from_secs(2), test.host.open_started.notified()) + .await + .unwrap(); + let next = test.apply(19382).await; + release.send(()).unwrap(); + test.send(b"new"); + assert_eq!(test.receive().await, (b"new".to_vec(), next.dst_addr)); + assert!( + test.adapter + .udp_clients + .iter() + .all(|entry| entry.key().forward == next) + ); + assert!(test.host.sockets.lock().unwrap()[1].upgrade().is_none()); + test.stop().await; +} + +#[cfg(not(any(target_arch = "wasm32", target_os = "wasi")))] +#[tokio::test] +async fn udp_removal_waits_for_response_tasks_already_removed_from_map() { + let test = UdpReloadTest::new().await; + let forward = test.apply(19381).await; + let socket = test.host.sockets.lock().unwrap()[0].upgrade().unwrap(); + let tracker = test + .adapter + .listeners + .lock() + .unwrap() + .get(&(forward.proto.clone(), forward.bind_addr)) + .unwrap() + .udp_tasks + .clone(); + let (started, ready) = oneshot::channel(); + let (release, waiting) = std::sync::mpsc::channel(); + // A running blocking task makes the delay between abort and socket release + // deterministic, as can happen when a response task is currently executing. + let task = tracker.spawn_blocking(move || { + started.send(()).unwrap(); + waiting.recv().unwrap(); + drop(socket); + }); + let key = UdpClientKey { + client_addr: "127.0.0.1:40000".parse().unwrap(), + forward, + }; + test.adapter + .udp_response_tasks + .insert(key, AbortOnDropHandle::new(task)); + ready.await.unwrap(); + test.adapter.udp_response_tasks.clear(); + let cleanup = test.adapter.prepare(&[]).await.unwrap().unwrap().commit(); + drop(cleanup); + let mut next = std::pin::pin!(test.adapter.prepare(&[])); + let result = tokio::time::timeout(Duration::from_millis(20), &mut next).await; + let pending = result.is_err(); + release.send(()).unwrap(); + let prepared = match result { + Ok(prepared) => prepared, + Err(_) => tokio::time::timeout(Duration::from_secs(2), next) + .await + .unwrap(), + }; + prepared.unwrap().unwrap().commit().await.unwrap(); + assert!( + pending, + "dropping the cleanup waiter released the operation lock before the socket" + ); + assert!(test.host.sockets.lock().unwrap()[0].upgrade().is_none()); + test.stop().await; +} diff --git a/easytier-core/src/host/testkit.rs b/easytier-core/src/host/testkit.rs index 6ab47745..1853dd1c 100644 --- a/easytier-core/src/host/testkit.rs +++ b/easytier-core/src/host/testkit.rs @@ -6,11 +6,12 @@ //! accept, UDP sockets never receive, and DNS answers only literal IP hosts. use std::{ + collections::HashMap, io, net::{IpAddr, SocketAddr}, pin::Pin, sync::{ - Arc, + Arc, Mutex, atomic::{AtomicUsize, Ordering}, }, task::{Context, Poll}, @@ -70,11 +71,16 @@ impl VirtualTcpSocket for TestTcpSocket { pub struct TestTcpListener { address: SocketAddr, active_listeners: Arc, + listeners: Option>>, + incoming: Option>>, } impl Drop for TestTcpListener { fn drop(&mut self) { self.active_listeners.fetch_sub(1, Ordering::Relaxed); + if let Some(listeners) = &self.listeners { + listeners.lock().unwrap().remove(&self.address); + } } } @@ -87,6 +93,12 @@ impl VirtualTcpListener for TestTcpListener { } async fn accept(&self) -> io::Result<(Self::Socket, SocketAddr)> { + if let Some(incoming) = &self.incoming { + let socket = incoming.lock().await.recv().await.ok_or_else(|| { + io::Error::new(io::ErrorKind::ConnectionAborted, "test listener closed") + })?; + return Ok((socket, "127.0.0.1:20001".parse().unwrap())); + } std::future::pending().await } } @@ -108,6 +120,8 @@ impl VirtualUdpSocket for TestUdpSocket { } } +pub type TestTcpListeners = HashMap>; + #[derive(Default)] pub struct TestHost { pub tcp_binds: AtomicUsize, @@ -115,21 +129,42 @@ pub struct TestHost { pub udp_binds: AtomicUsize, pub proxy_nat_connections: Option>, + pub port_forward_connections: + Option>, + pub tcp_listeners: Option>>, + pub reject_tcp_bind: Mutex>, pub reject_socks5_listener: bool, } +impl TestHost { + #[cfg(all(feature = "management", feature = "proxy-smoltcp-stack"))] + pub fn connect_to_listener(&self, address: SocketAddr) -> tokio::io::DuplexStream { + let (socket, peer) = tokio::io::duplex(1024); + self.tcp_listeners + .as_ref() + .unwrap() + .lock() + .unwrap() + .get(&address) + .expect("test listener must be bound") + .send(TestTcpSocket(socket)) + .unwrap(); + peer + } +} + #[async_trait::async_trait] impl VirtualTcpSocketFactory for TestHost { type Socket = TestTcpSocket; async fn connect_tcp(&self, options: TcpConnectOptions) -> anyhow::Result { - if options.purpose != TcpSocketPurpose::ProxyNat { - anyhow::bail!("test host does not connect non-proxy TCP sockets"); + let connections = match options.purpose { + TcpSocketPurpose::ProxyNat => &self.proxy_nat_connections, + TcpSocketPurpose::PortForward => &self.port_forward_connections, + _ => anyhow::bail!("test host does not connect non-proxy TCP sockets"), } - let connections = self - .proxy_nat_connections - .as_ref() - .ok_or_else(|| anyhow::anyhow!("test host proxy NAT is disabled"))?; + .as_ref() + .ok_or_else(|| anyhow::anyhow!("test host TCP connections are disabled"))?; let (socket, peer) = tokio::io::duplex(1024); connections .send((options.remote_addr, peer)) @@ -159,10 +194,26 @@ impl VirtualTcpListenerFactory for TestHost { } else { address }; + if *self.reject_tcp_bind.lock().unwrap() == Some(address) { + anyhow::bail!("test host rejected TCP bind to {address}"); + } + let incoming = if let Some(listeners) = &self.tcp_listeners { + let mut listeners = listeners.lock().unwrap(); + if listeners.contains_key(&address) { + return Err(io::Error::new(io::ErrorKind::AddrInUse, address.to_string()).into()); + } + let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); + listeners.insert(address, sender); + Some(tokio::sync::Mutex::new(receiver)) + } else { + None + }; self.active_tcp_listeners.fetch_add(1, Ordering::Relaxed); Ok(Arc::new(TestTcpListener { address, active_listeners: self.active_tcp_listeners.clone(), + listeners: self.tcp_listeners.clone(), + incoming, })) } } diff --git a/easytier-core/src/instance/mod.rs b/easytier-core/src/instance/mod.rs index d9df9d7d..43303d7c 100644 --- a/easytier-core/src/instance/mod.rs +++ b/easytier-core/src/instance/mod.rs @@ -936,26 +936,57 @@ where anyhow::bail!("runtime config cannot update while instance is stopping or stopped"); } self.validate_runtime_config_capabilities(&config)?; + #[cfg(feature = "proxy-smoltcp-stack")] + let prepared = self.prepare_port_forwards(&config).await?; + self.commit_runtime_config_under_operation( + config, + #[cfg(feature = "proxy-smoltcp-stack")] + prepared, + || (), + ) + .await + } + + #[cfg(feature = "proxy-smoltcp-stack")] + pub(crate) async fn prepare_port_forwards( + &self, + config: &CoreInstanceRuntimeConfig, + ) -> anyhow::Result>> { + self.port_forward_adapter + .prepare(&config.services.gateway.port_forwards) + .await + } + + pub(crate) async fn commit_runtime_config_under_operation( + &self, + config: CoreInstanceRuntimeConfig, + #[cfg(feature = "proxy-smoltcp-stack")] prepared: Option< + crate::gateway::PreparedPortForwards, + >, + commit_config: impl FnOnce(), + ) -> anyhow::Result<()> { #[cfg(feature = "test-utils")] let reload_acl = self.runtime_config.snapshot().services.acl != config.services.acl; - let published = self.peer_manager.update_runtime_config(config).await?; + let (published, cleanup) = self + .peer_manager + .update_runtime_config_with(config, || { + #[cfg(feature = "proxy-smoltcp-stack")] + let cleanup = prepared.map(|prepared| prepared.commit()); + #[cfg(not(feature = "proxy-smoltcp-stack"))] + let cleanup: Option> = None; + commit_config(); + cleanup + }) + .await?; #[cfg(feature = "test-utils")] if reload_acl { self.acl_reload_count.fetch_add(1, Ordering::Relaxed); } self.proxy_cidr_table .update_snapshot(proxy_cidr_snapshot(&published)); - #[cfg(feature = "proxy-smoltcp-stack")] - self.port_forward_adapter - .reload( - &self - .runtime_config - .snapshot() - .services - .gateway - .port_forwards, - ) - .await?; + if let Some(cleanup) = cleanup { + let _ = cleanup.await; + } Ok(()) } diff --git a/easytier-core/src/instance/tests.rs b/easytier-core/src/instance/tests.rs index b5cf282c..cde3f038 100644 --- a/easytier-core/src/instance/tests.rs +++ b/easytier-core/src/instance/tests.rs @@ -2981,6 +2981,268 @@ virtual_ip = "10.82.0.2/24" assert_eq!(host.active_tcp_listeners.load(Ordering::Relaxed), 0); } + #[cfg(all(feature = "management", feature = "proxy-smoltcp-stack"))] + mod port_forward_reload { + use super::*; + use easytier_proto::api::config::{ + ConfigPatchAction, InstanceConfigPatch, PortForwardPatch, + }; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::sync::mpsc::UnboundedReceiver; + + fn forward(bind_port: u16, destination_port: u16) -> PortForwardConfig { + PortForwardConfig { + bind_addr: SocketAddr::from(([127, 0, 0, 1], bind_port)), + dst_addr: SocketAddr::from(([10, 144, 0, 2], destination_port)), + proto: "tcp".to_owned(), + } + } + + async fn start() -> ( + Arc>, + Arc, + UnboundedReceiver<(SocketAddr, tokio::io::DuplexStream)>, + ) { + let (connections, receiver) = tokio::sync::mpsc::unbounded_channel(); + let host = Arc::new(TestHost { + port_forward_connections: Some(connections), + tcp_listeners: Some(Arc::new(std::sync::Mutex::new(Default::default()))), + ..Default::default() + }); + let config = + TomlConfig::new_from_str("instance_name = \"port-forward-reload\"").unwrap(); + config.set_port_forwards(vec![forward(19491, 19381)]); + let (packet_sink, _packet_receiver) = tokio::sync::mpsc::channel(16); + let instance = CoreInstance::from_toml( + config, + adapters_with_host(host.clone(), None, Arc::new(packet_sink)), + ) + .unwrap(); + instance.start().await.unwrap(); + (instance, host, receiver) + } + + async fn assert_destination( + host: &TestHost, + receiver: &mut UnboundedReceiver<(SocketAddr, tokio::io::DuplexStream)>, + config: &PortForwardConfig, + ) -> (tokio::io::DuplexStream, tokio::io::DuplexStream) { + let mut incoming = host.connect_to_listener(config.bind_addr); + let outgoing = tokio::time::timeout(Duration::from_secs(2), async { + let (destination, mut outgoing) = receiver.recv().await.unwrap(); + assert_eq!(destination, config.dst_addr); + incoming.write_all(b"request").await.unwrap(); + let mut data = [0; 7]; + outgoing.read_exact(&mut data).await.unwrap(); + assert_eq!(&data, b"request"); + outgoing + }) + .await + .expect("port forward must accept and forward to its committed destination"); + (incoming, outgoing) + } + + fn patch( + removed: &[PortForwardConfig], + added: &[PortForwardConfig], + ) -> InstanceConfigPatch { + InstanceConfigPatch { + port_forwards: removed + .iter() + .map(|config| (ConfigPatchAction::Remove, config)) + .chain(added.iter().map(|config| (ConfigPatchAction::Add, config))) + .map(|(action, config)| PortForwardPatch { + action: action as i32, + cfg: Some(config.clone().into()), + }) + .collect(), + ..Default::default() + } + } + + fn persistence(fail: bool) -> RecordingConfigPatchPersistence { + RecordingConfigPatchPersistence { + writes: std::sync::Mutex::new(Vec::new()), + fail: AtomicBool::new(fail), + } + } + + #[tokio::test] + async fn destination_changes_reuse_listener() { + let (instance, host, mut connections) = start().await; + let (mut established_client, mut established_destination) = + assert_destination(&host, &mut connections, &forward(19491, 19381)).await; + let initial_binds = host.tcp_binds.load(Ordering::Relaxed); + + for port in [19382, 19383, 19381] { + let config = forward(19491, port); + let mut runtime = (*instance.runtime_config.snapshot()).clone(); + runtime.services.gateway.port_forwards = vec![config.clone()]; + instance.update_runtime_config(runtime).await.unwrap(); + + assert_eq!(host.tcp_binds.load(Ordering::Relaxed), initial_binds); + assert_eq!(host.active_tcp_listeners.load(Ordering::Relaxed), 1); + assert_destination(&host, &mut connections, &config).await; + tokio::time::timeout(Duration::from_secs(2), async { + let payload = port.to_be_bytes(); + established_client.write_all(&payload).await.unwrap(); + let mut received = [0; 2]; + established_destination + .read_exact(&mut received) + .await + .unwrap(); + assert_eq!(received, payload); + established_destination.write_all(&payload).await.unwrap(); + established_client.read_exact(&mut received).await.unwrap(); + assert_eq!(received, payload); + }) + .await + .expect("changing the destination must preserve established TCP connections"); + } + instance.stop().await; + } + + #[tokio::test] + async fn failed_runtime_bind_keeps_committed_snapshot() { + let (instance, host, mut connections) = start().await; + let original = forward(19491, 19381); + let rejected = forward(19492, 19382); + *host.reject_tcp_bind.lock().unwrap() = Some(rejected.bind_addr); + let mut runtime = (*instance.runtime_config.snapshot()).clone(); + runtime.services.gateway.port_forwards = vec![original.clone(), rejected]; + + assert!( + instance + .update_runtime_config(runtime.clone()) + .await + .is_err() + ); + assert_eq!( + instance + .runtime_config + .snapshot() + .services + .gateway + .port_forwards, + std::slice::from_ref(&original) + ); + assert_eq!(host.active_tcp_listeners.load(Ordering::Relaxed), 1); + assert_destination(&host, &mut connections, &original).await; + + *host.reject_tcp_bind.lock().unwrap() = None; + instance.update_runtime_config(runtime).await.unwrap(); + assert_eq!(host.active_tcp_listeners.load(Ordering::Relaxed), 2); + instance.stop().await; + } + + #[tokio::test] + async fn failed_patch_bind_releases_prepared_listener_and_retries() { + let (instance, host, mut connections) = start().await; + let original = forward(19491, 19381); + let added = [forward(19492, 19382), forward(19493, 19383)]; + let before = instance.toml_config().unwrap().dump(); + let persistence = persistence(false); + let patch = patch(&[], &added); + *host.reject_tcp_bind.lock().unwrap() = Some(added[1].bind_addr); + let initial_binds = host.tcp_binds.load(Ordering::Relaxed); + + let error = + crate::management::apply_config_patch(&instance, patch.clone(), Some(&persistence)) + .await + .unwrap_err(); + + assert!(error.to_string().contains("test host rejected TCP bind")); + assert_eq!(host.tcp_binds.load(Ordering::Relaxed), initial_binds + 2); + assert_eq!(host.active_tcp_listeners.load(Ordering::Relaxed), 1); + assert_eq!(instance.toml_config().unwrap().dump(), before); + assert_eq!( + instance + .runtime_config + .snapshot() + .services + .gateway + .port_forwards, + std::slice::from_ref(&original) + ); + assert!(persistence.writes.lock().unwrap().is_empty()); + assert_destination(&host, &mut connections, &original).await; + + *host.reject_tcp_bind.lock().unwrap() = None; + crate::management::apply_config_patch(&instance, patch, Some(&persistence)) + .await + .unwrap(); + assert_eq!(host.active_tcp_listeners.load(Ordering::Relaxed), 3); + assert_eq!(persistence.writes.lock().unwrap().len(), 1); + for config in &added { + assert_destination(&host, &mut connections, config).await; + } + instance.stop().await; + } + + #[tokio::test] + async fn failed_patch_persistence_keeps_destination_and_releases_prepared_listener() { + let (instance, host, mut connections) = start().await; + let original = forward(19491, 19381); + let replacement = forward(19491, 19382); + let added = forward(19492, 19383); + let before = instance.toml_config().unwrap().dump(); + let persistence = persistence(true); + let patch = patch( + std::slice::from_ref(&original), + &[replacement.clone(), added.clone()], + ); + + let error = + crate::management::apply_config_patch(&instance, patch.clone(), Some(&persistence)) + .await + .unwrap_err(); + + assert!( + error + .to_string() + .contains("injected config persistence failure") + ); + assert_eq!(host.active_tcp_listeners.load(Ordering::Relaxed), 1); + assert_eq!(instance.toml_config().unwrap().dump(), before); + assert_eq!( + instance + .runtime_config + .snapshot() + .services + .gateway + .port_forwards, + std::slice::from_ref(&original) + ); + assert!(persistence.writes.lock().unwrap().is_empty()); + assert_destination(&host, &mut connections, &original).await; + + persistence.fail.store(false, Ordering::Relaxed); + crate::management::apply_config_patch(&instance, patch, Some(&persistence)) + .await + .unwrap(); + assert_eq!(host.active_tcp_listeners.load(Ordering::Relaxed), 2); + assert_eq!(persistence.writes.lock().unwrap().len(), 1); + assert_destination(&host, &mut connections, &replacement).await; + assert_destination(&host, &mut connections, &added).await; + instance.stop().await; + } + + #[tokio::test] + async fn removed_listener_can_be_rebound_immediately() { + let (instance, host, mut connections) = start().await; + let original = (*instance.runtime_config.snapshot()).clone(); + let mut empty = original.clone(); + empty.services.gateway.port_forwards.clear(); + + instance.update_runtime_config(empty).await.unwrap(); + assert_eq!(host.active_tcp_listeners.load(Ordering::Relaxed), 0); + instance.update_runtime_config(original).await.unwrap(); + assert_eq!(host.active_tcp_listeners.load(Ordering::Relaxed), 1); + assert_destination(&host, &mut connections, &forward(19491, 19381)).await; + instance.stop().await; + } + } + #[cfg(feature = "proxy-packet")] #[tokio::test] async fn runtime_core_instance_owns_wrapped_transport_source_nat() { diff --git a/easytier-core/src/management/full/config_patch.rs b/easytier-core/src/management/full/config_patch.rs index d4714d97..226bc458 100644 --- a/easytier-core/src/management/full/config_patch.rs +++ b/easytier-core/src/management/full/config_patch.rs @@ -51,9 +51,25 @@ where // Preserve the existing ordered partial-commit contract: earlier valid // sub-patches remain applied if a later sub-patch fails. let patch_result: anyhow::Result<(bool, bool)> = async { - let result = patch_port_forwards(&candidate, patch.port_forwards); - validate_persist_and_commit_candidate(instance, &config, &candidate, persistence).await?; - result?; + if !patch.port_forwards.is_empty() { + patch_port_forwards(&candidate, patch.port_forwards)?; + // Port forwards have fallible host resources. Reserve them before the + // durable write, then commit resources and readback together. A failed + // bind or persistence operation must leave the previous rule usable. + let normalized = validate_candidate(instance, &candidate)?; + let runtime = runtime_config_from_normalized(&normalized); + #[cfg(feature = "proxy-smoltcp-stack")] + let prepared = instance.prepare_port_forwards(&runtime).await?; + persist_candidate_if_changed(instance, &config, &candidate, persistence).await?; + instance + .commit_runtime_config_under_operation( + runtime, + #[cfg(feature = "proxy-smoltcp-stack")] + prepared, + || config.replace_from_snapshot(&candidate), + ) + .await?; + } let result = patch_acl(&candidate, patch.acl); validate_persist_and_commit_candidate(instance, &config, &candidate, persistence).await?; diff --git a/easytier-core/src/peers/peer_manager.rs b/easytier-core/src/peers/peer_manager.rs index b95414c5..7fc2f248 100644 --- a/easytier-core/src/peers/peer_manager.rs +++ b/easytier-core/src/peers/peer_manager.rs @@ -1336,6 +1336,16 @@ impl PeerManagerCore { &self, config: CoreInstanceRuntimeConfig, ) -> anyhow::Result> { + self.update_runtime_config_with(config, || ()) + .await + .map(|(published, ())| published) + } + + pub(crate) async fn update_runtime_config_with( + &self, + config: CoreInstanceRuntimeConfig, + before_publish: impl FnOnce() -> T, + ) -> anyhow::Result<(Arc, T)> { let _update = self.runtime_config_update.lock().await; let current = self.runtime_config.snapshot(); let refresh_acl_groups = current.peer.peer_group_memberships @@ -1345,6 +1355,9 @@ impl PeerManagerCore { let next_acl = reload_acl.then(|| config.services.acl.clone()); let next_built_acl = next_acl.as_ref().map(AclRuleConfig::build).transpose()?; + // Commit prepared host resources only after all fallible validation, + // with no cancellation point before publishing their configuration. + let committed = before_publish(); self.set_avoid_relay_data_preference(config.peer.avoid_relay_data_preference); let published = self .runtime_config @@ -1357,7 +1370,7 @@ impl PeerManagerCore { if refresh_acl_groups { self.route.refresh_acl_groups().await; } - Ok(published) + Ok((published, committed)) } /// Keeps this manager's ACL rules and group assignments aligned with a