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