diff --git a/src/transport/middle_proxy/handshake.rs b/src/transport/middle_proxy/handshake.rs index 181bae1..b43d240 100644 --- a/src/transport/middle_proxy/handshake.rs +++ b/src/transport/middle_proxy/handshake.rs @@ -8,6 +8,7 @@ use std::net::{IpAddr, SocketAddr}; use std::os::fd::{AsRawFd, RawFd}; #[cfg(target_os = "linux")] use std::os::raw::c_int; +use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::{Duration, Instant}; @@ -252,7 +253,7 @@ impl MePool { /// Perform full ME RPC handshake on an established TCP stream. /// Returns cipher keys/ivs and split halves; does not register writer. pub(crate) async fn handshake_only( - &self, + self: &Arc, stream: TcpStream, addr: SocketAddr, upstream_egress: Option, diff --git a/src/transport/middle_proxy/pool.rs b/src/transport/middle_proxy/pool.rs index 6d58df2..95ba2c2 100644 --- a/src/transport/middle_proxy/pool.rs +++ b/src/transport/middle_proxy/pool.rs @@ -407,6 +407,7 @@ pub(super) struct BindingPolicyCore { } pub(super) struct NatRuntimeCore { + pub(super) discovery: super::pool_nat::NatDiscovery, pub(super) nat_ip_cfg: Option, pub(super) nat_ip_detected: Arc>>, pub(super) nat_probe: bool, diff --git a/src/transport/middle_proxy/pool/construction.rs b/src/transport/middle_proxy/pool/construction.rs index 5eede32..83bd429 100644 --- a/src/transport/middle_proxy/pool/construction.rs +++ b/src/transport/middle_proxy/pool/construction.rs @@ -279,6 +279,7 @@ impl MePool { me_bind_stale_ttl_secs: AtomicU64::new(me_bind_stale_ttl_secs), }), nat_runtime: Arc::new(NatRuntimeCore { + discovery: super::super::pool_nat::NatDiscovery::default(), nat_ip_cfg: nat_ip, nat_ip_detected: Arc::new(RwLock::new(None)), nat_probe, diff --git a/src/transport/middle_proxy/pool_nat.rs b/src/transport/middle_proxy/pool_nat.rs index 9c1a15c..23ff765 100644 --- a/src/transport/middle_proxy/pool_nat.rs +++ b/src/transport/middle_proxy/pool_nat.rs @@ -1,7 +1,9 @@ use std::collections::HashMap; use std::net::IpAddr; +use std::sync::Arc; use std::time::Duration; +use tokio::sync::oneshot; use tokio::task::JoinSet; use tokio::time::timeout; use tracing::{debug, info}; @@ -18,7 +20,18 @@ use std::time::Instant; const STUN_BATCH_TIMEOUT: Duration = Duration::from_secs(5); const STUN_BATCH_TCP_FALLBACK_TIMEOUT: Duration = Duration::from_secs(12); +const STUN_CACHE_TTL: Duration = Duration::from_secs(600); +const HTTP_FAILURE_TTL: Duration = Duration::from_secs(60); +// Discovery ownership and source-bound cache admission for this pool generation. +mod state; +pub(super) use state::NatDiscovery; + +// Local network regressions exercise cancellation before discovery publication. +#[cfg(test)] +mod tests; + +/// Probes the configured or default STUN endpoint with both address families. #[allow(dead_code)] pub async fn stun_probe(stun_addr: Option) -> Result { let stun_addr = stun_addr.unwrap_or_else(|| { @@ -35,6 +48,7 @@ pub async fn stun_probe(stun_addr: Option) -> Result Option { let urls = crate::config::defaults::default_http_ip_detect_urls(); @@ -151,6 +165,7 @@ impl MePool { (live_servers, best_reflected) } + /// Applies the configured or HTTP-detected NAT IP to non-public addresses. pub(super) fn translate_ip_for_nat(&self, ip: IpAddr) -> IpAddr { let nat_ip = self.nat_runtime.nat_ip_cfg.or_else(|| { self.nat_runtime @@ -177,6 +192,7 @@ impl MePool { } } + /// Selects the NAT source IP without replacing the actual TCP source port. pub(super) fn translate_our_addr_with_reflection( &self, addr: std::net::SocketAddr, @@ -203,7 +219,8 @@ impl MePool { std::net::SocketAddr::new(ip, addr.port()) } - pub(super) async fn maybe_detect_nat_ip(&self, local_ip: IpAddr) -> Option { + /// Waits for pool-owned HTTP discovery without owning its cancellation. + pub(super) async fn maybe_detect_nat_ip(self: &Arc, local_ip: IpAddr) -> Option { if self.nat_runtime.nat_ip_cfg.is_some() { return self.nat_runtime.nat_ip_cfg; } @@ -220,103 +237,117 @@ impl MePool { return Some(ip); } - let Some(ip) = detect_public_ipv4_http(&self.nat_runtime.http_ip_detect_urls).await else { - return None; - }; - { - let mut guard = self.nat_runtime.nat_ip_detected.write().await; - *guard = Some(IpAddr::V4(ip)); + let mut flight = self.nat_runtime.discovery.http.clone().lock_owned().await; + if let Some(ip) = *self.nat_runtime.nat_ip_detected.read().await { + return Some(ip); } - info!(public_ip = %ip, "Auto-detected public IP for NAT translation"); - Some(IpAddr::V4(ip)) + if flight.is_some_and(|until| Instant::now() < until) { + return None; + } + + let pool = self.clone(); + let (send, receive) = oneshot::channel(); + // Transfer singleflight ownership before the producer can be cancelled or polled. + self.lifecycle + .spawn_producer(async move { + let detected = detect_public_ipv4_http(&pool.nat_runtime.http_ip_detect_urls) + .await + .map(IpAddr::V4); + if let Some(ip) = detected { + *pool.nat_runtime.nat_ip_detected.write().await = Some(ip); + *flight = None; + info!(public_ip = %ip, "Auto-detected public IP for NAT translation"); + } else { + *flight = Some(Instant::now() + HTTP_FAILURE_TTL); + } + let _ = send.send(detected); + }) + .ok()?; + receive.await.ok().flatten() } + /// Waits for a pool-owned STUN refresh in the exact source-address context. pub(super) async fn maybe_reflect_public_addr( - &self, + self: &Arc, family: IpFamily, bind_ip: Option, ) -> Option { - const STUN_CACHE_TTL: Duration = Duration::from_secs(600); - let use_shared_cache = bind_ip.is_none(); - if !use_shared_cache { - match (family, bind_ip) { - (IpFamily::V4, Some(IpAddr::V4(_))) - | (IpFamily::V6, Some(IpAddr::V6(_))) - | (_, None) => {} - _ => return None, - } + if let Some(bind_ip) = bind_ip { + return self.reflect_bound_addr(family, bind_ip).await; } - // Backoff window - if use_shared_cache - && let Some(until) = *self.nat_runtime.stun_backoff_until.read().await - && Instant::now() < until - { - if let Ok(cache) = self.nat_runtime.nat_reflection_cache.try_lock() { - let slot = match family { - IpFamily::V4 => cache.v4, - IpFamily::V6 => cache.v6, - }; - return slot.map(|(_, addr)| addr); - } - return None; + if let Some(cached) = self.cached_shared_reflection(family).await { + return cached; + } + let singleflight = match family { + IpFamily::V4 => self.nat_runtime.nat_reflection_singleflight_v4.clone(), + IpFamily::V6 => self.nat_runtime.nat_reflection_singleflight_v6.clone(), + }; + let flight = singleflight.lock_owned().await; + if let Some(cached) = self.cached_shared_reflection(family).await { + return cached; } - if use_shared_cache && let Ok(mut cache) = self.nat_runtime.nat_reflection_cache.try_lock() - { - let slot = match family { - IpFamily::V4 => &mut cache.v4, - IpFamily::V6 => &mut cache.v6, - }; - if let Some((ts, addr)) = slot - && ts.elapsed() < STUN_CACHE_TTL - { - return Some(*addr); - } - } - - let _singleflight_guard = if use_shared_cache { - Some(match family { - IpFamily::V4 => self.nat_runtime.nat_reflection_singleflight_v4.lock().await, - IpFamily::V6 => self.nat_runtime.nat_reflection_singleflight_v6.lock().await, + let pool = self.clone(); + let (send, receive) = oneshot::channel(); + self.lifecycle + .spawn_producer(async move { + let _flight = flight; + let attempt = pool + .nat_runtime + .nat_probe_attempts + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let reflected = pool.refresh_stun(family, None, attempt).await; + if let Some(addr) = reflected { + // Cache readers must never cause a completed refresh to be discarded. + let mut cache = pool.nat_runtime.nat_reflection_cache.lock().await; + let slot = match family { + IpFamily::V4 => &mut cache.v4, + IpFamily::V6 => &mut cache.v6, + }; + *slot = Some((Instant::now(), addr)); + pool.nat_runtime + .nat_probe_attempts + .store(0, std::sync::atomic::Ordering::Relaxed); + } else { + let backoff = Duration::from_secs(60 * 2u64.pow((attempt as u32).min(6))); + *pool.nat_runtime.stun_backoff_until.write().await = + Some(Instant::now() + backoff); + } + let _ = send.send(reflected); }) - } else { - None + .ok()?; + receive.await.ok().flatten() + } + + async fn cached_shared_reflection( + &self, + family: IpFamily, + ) -> Option> { + let backoff = self + .nat_runtime + .stun_backoff_until + .read() + .await + .is_some_and(|until| Instant::now() < until); + let cache = self.nat_runtime.nat_reflection_cache.lock().await; + let slot = match family { + IpFamily::V4 => cache.v4, + IpFamily::V6 => cache.v6, }; - - if use_shared_cache - && let Some(until) = *self.nat_runtime.stun_backoff_until.read().await - && Instant::now() < until - { - if let Ok(cache) = self.nat_runtime.nat_reflection_cache.try_lock() { - let slot = match family { - IpFamily::V4 => cache.v4, - IpFamily::V6 => cache.v6, - }; - return slot.map(|(_, addr)| addr); - } - return None; + // Preserve the existing stale-result fallback during shared failure backoff. + if backoff { + return Some(slot.map(|(_, addr)| addr)); } + slot.filter(|(created, _)| created.elapsed() < STUN_CACHE_TTL) + .map(|(_, addr)| Some(addr)) + } - if use_shared_cache && let Ok(mut cache) = self.nat_runtime.nat_reflection_cache.try_lock() - { - let slot = match family { - IpFamily::V4 => &mut cache.v4, - IpFamily::V6 => &mut cache.v6, - }; - if let Some((ts, addr)) = slot - && ts.elapsed() < STUN_CACHE_TTL - { - return Some(*addr); - } - } - - let attempt = if use_shared_cache { - self.nat_runtime - .nat_probe_attempts - .fetch_add(1, std::sync::atomic::Ordering::Relaxed) - } else { - 0 - }; + async fn refresh_stun( + &self, + family: IpFamily, + bind_ip: Option, + attempt: u8, + ) -> Option { let configured_servers = self.configured_stun_servers(); let live_snapshot = self.nat_runtime.nat_stun_live_servers.read().await.clone(); let primary_servers = if live_snapshot.is_empty() { @@ -348,33 +379,14 @@ impl MePool { } if let Some(reflected_addr) = selected_reflected { - if use_shared_cache { - self.nat_runtime - .nat_probe_attempts - .store(0, std::sync::atomic::Ordering::Relaxed); - } info!( family = ?family, live_servers = live_server_count, "STUN-Quorum reached, IP: {}", reflected_addr.ip() ); - if use_shared_cache - && let Ok(mut cache) = self.nat_runtime.nat_reflection_cache.try_lock() - { - let slot = match family { - IpFamily::V4 => &mut cache.v4, - IpFamily::V6 => &mut cache.v6, - }; - *slot = Some((Instant::now(), reflected_addr)); - } return Some(reflected_addr); } - - if use_shared_cache { - let backoff = Duration::from_secs(60 * 2u64.pow((attempt as u32).min(6))); - *self.nat_runtime.stun_backoff_until.write().await = Some(Instant::now() + backoff); - } None } } diff --git a/src/transport/middle_proxy/pool_nat/state.rs b/src/transport/middle_proxy/pool_nat/state.rs new file mode 100644 index 0000000..1c08954 --- /dev/null +++ b/src/transport/middle_proxy/pool_nat/state.rs @@ -0,0 +1,119 @@ +use std::collections::{BTreeSet, HashMap}; +use std::net::{IpAddr, SocketAddr}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use tokio::sync::{Mutex, oneshot}; +use tokio_util::sync::CancellationToken; + +use super::{IpFamily, MePool, STUN_CACHE_TTL}; + +// Context retirement and TTL tests use the same gates as production discovery. +#[cfg(test)] +mod tests; + +/// Pool-local discovery gates; bound entries are limited to configured egress addresses. +#[derive(Default)] +pub(in crate::transport::middle_proxy) struct NatDiscovery { + pub(super) http: Arc>>, + bound: Mutex>>, +} + +#[derive(Default)] +struct BoundDiscovery { + state: Arc>, + cancel: CancellationToken, +} + +#[derive(Default)] +struct BoundDiscoveryState { + reflection: Option<(Instant, SocketAddr)>, + retry_after: Option, + attempt: u8, +} + +impl NatDiscovery { + async fn bound_context( + &self, + family: IpFamily, + bind_ip: IpAddr, + allowed: &BTreeSet, + ) -> Option> { + let mut contexts = self.bound.lock().await; + // Retired producers cannot republish into the map or a replacement context. + contexts.retain(|(entry_family, ip), context| { + let keep = *entry_family != family || allowed.contains(ip); + if !keep { + context.cancel.cancel(); + } + keep + }); + if !allowed.contains(&bind_ip) { + return None; + } + Some(contexts.entry((family, bind_ip)).or_default().clone()) + } +} + +impl MePool { + /// Refreshes one configured source without borrowing another source's cached reflection. + pub(super) async fn reflect_bound_addr( + self: &Arc, + family: IpFamily, + bind_ip: IpAddr, + ) -> Option { + if bind_ip.is_ipv6() != matches!(family, IpFamily::V6) { + return None; + } + let allowed = self + .upstream + .as_ref()? + .direct_bind_ips_for_family(bind_ip.is_ipv6()) + .await; + let context = self + .nat_runtime + .discovery + .bound_context(family, bind_ip, &allowed) + .await?; + let mut flight = tokio::select! { + biased; + _ = context.cancel.cancelled() => return None, + guard = context.state.clone().lock_owned() => guard, + }; + if let Some((created, addr)) = flight.reflection + && created.elapsed() < STUN_CACHE_TTL + { + return Some(addr); + } + if flight + .retry_after + .is_some_and(|until| Instant::now() < until) + { + return None; + } + + let pool = self.clone(); + let (send, receive) = oneshot::channel(); + self.lifecycle.spawn_producer(async move { + tokio::select! { + biased; + _ = context.cancel.cancelled() => {} + _ = async { + let reflected = pool.refresh_stun(family, Some(bind_ip), flight.attempt).await; + if let Some(addr) = reflected { + flight.reflection = Some((Instant::now(), addr)); + flight.retry_after = None; + flight.attempt = 0; + } else { + // Completed failures must also survive a short reconnect deadline. + let backoff = Duration::from_secs(60 * 2u64.pow(flight.attempt as u32)); + flight.retry_after = Some(Instant::now() + backoff); + flight.attempt = flight.attempt.saturating_add(1).min(6); + } + let _ = send.send(reflected); + } => {} + } + }).ok()?; + receive.await.ok().flatten() + } +} diff --git a/src/transport/middle_proxy/pool_nat/state/tests.rs b/src/transport/middle_proxy/pool_nat/state/tests.rs new file mode 100644 index 0000000..6909f19 --- /dev/null +++ b/src/transport/middle_proxy/pool_nat/state/tests.rs @@ -0,0 +1,146 @@ +use super::*; +use crate::transport::middle_proxy::pool_nat::tests::lifecycle::set_bind_addresses; +use crate::transport::middle_proxy::pool_nat::tests::{ + PUBLIC_IP, TEST_DEADLINE, nat_pool, stun_server, +}; +use std::net::Ipv4Addr; +use std::sync::atomic::Ordering; +use tokio::time::timeout; + +#[tokio::test] +async fn bound_positive_ttl_refreshes_without_publishing_to_default_cache() { + let stun = stun_server("127.0.0.1:0".parse().unwrap(), PUBLIC_IP.into()).await; + stun.release.add_permits(16); + let mut pool = nat_pool(vec![stun.addr.to_string()], vec![]).await; + set_bind_addresses(&mut pool, &["127.0.0.1"]); + let ip = IpAddr::V4(Ipv4Addr::LOCALHOST); + let allowed = BTreeSet::from([ip]); + let context = pool + .nat_runtime + .discovery + .bound_context(IpFamily::V4, ip, &allowed) + .await + .unwrap(); + assert!( + pool.maybe_reflect_public_addr(IpFamily::V4, Some(ip)) + .await + .is_some() + ); + assert!( + pool.maybe_reflect_public_addr(IpFamily::V4, Some(ip)) + .await + .is_some() + ); + assert_eq!(stun.requests.load(Ordering::Relaxed), 1); + context.state.lock().await.reflection.as_mut().unwrap().0 = Instant::now() - STUN_CACHE_TTL; + assert!( + pool.maybe_reflect_public_addr(IpFamily::V4, Some(ip)) + .await + .is_some() + ); + assert_eq!(stun.requests.load(Ordering::Relaxed), 2); + assert_eq!(pool.nat_runtime.nat_reflection_cache.lock().await.v4, None); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} + +#[tokio::test] +async fn context_churn_is_bounded_and_retirement_does_not_mix_families() { + let discovery = NatDiscovery::default(); + let v6: IpAddr = "::1".parse().unwrap(); + let v6_context = discovery + .bound_context(IpFamily::V6, v6, &BTreeSet::from([v6])) + .await + .unwrap(); + let mut previous: Option> = None; + for suffix in 1..=1000u32 { + let ip = IpAddr::V4(Ipv4Addr::from(0x7f000000 + suffix)); + let context = discovery + .bound_context(IpFamily::V4, ip, &BTreeSet::from([ip])) + .await + .unwrap(); + if let Some(old) = previous.replace(context.clone()) { + assert!(old.cancel.is_cancelled()); + } + assert!(!context.cancel.is_cancelled()); + assert!(!v6_context.cancel.is_cancelled()); + assert_eq!(discovery.bound.lock().await.len(), 2); + } + let denied = IpAddr::V4(Ipv4Addr::LOCALHOST); + assert!( + discovery + .bound_context(IpFamily::V4, denied, &BTreeSet::new()) + .await + .is_none() + ); + assert!(previous.unwrap().cancel.is_cancelled()); + assert_eq!(discovery.bound.lock().await.len(), 1); +} + +#[tokio::test] +async fn retiring_bound_context_cancels_its_producer_and_prevents_late_publication() { + let stun = stun_server("127.0.0.1:0".parse().unwrap(), PUBLIC_IP.into()).await; + let mut pool = nat_pool(vec![stun.addr.to_string()], vec![]).await; + set_bind_addresses(&mut pool, &["127.0.0.1"]); + let ip = IpAddr::V4(Ipv4Addr::LOCALHOST); + let context = pool + .nat_runtime + .discovery + .bound_context(IpFamily::V4, ip, &BTreeSet::from([ip])) + .await + .unwrap(); + let waiter_pool = pool.clone(); + let waiter = tokio::spawn(async move { + waiter_pool + .maybe_reflect_public_addr(IpFamily::V4, Some(ip)) + .await + }); + stun.received().await; + assert!( + pool.nat_runtime + .discovery + .bound_context(IpFamily::V4, ip, &BTreeSet::new()) + .await + .is_none() + ); + assert_eq!(timeout(TEST_DEADLINE, waiter).await.unwrap().unwrap(), None); + stun.release.add_permits(16); + assert_eq!(context.state.lock().await.reflection, None); + assert!(pool.nat_runtime.discovery.bound.lock().await.is_empty()); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} + +#[tokio::test] +async fn bound_failure_retry_deadline_expires_and_success_resets_backoff() { + let stun = stun_server("127.0.0.1:0".parse().unwrap(), PUBLIC_IP.into()).await; + stun.release.add_permits(16); + let mut pool = nat_pool(vec![stun.addr.to_string()], vec![]).await; + set_bind_addresses(&mut pool, &["127.0.0.1"]); + let ip = IpAddr::V4(Ipv4Addr::LOCALHOST); + let context = pool + .nat_runtime + .discovery + .bound_context(IpFamily::V4, ip, &BTreeSet::from([ip])) + .await + .unwrap(); + { + let mut state = context.state.lock().await; + state.retry_after = Some(Instant::now() + Duration::from_secs(60)); + state.attempt = 3; + } + assert!( + pool.maybe_reflect_public_addr(IpFamily::V4, Some(ip)) + .await + .is_none() + ); + assert_eq!(stun.requests.load(Ordering::Relaxed), 0); + context.state.lock().await.retry_after = Some(Instant::now() - Duration::from_secs(1)); + assert!( + pool.maybe_reflect_public_addr(IpFamily::V4, Some(ip)) + .await + .is_some() + ); + assert_eq!(context.state.lock().await.attempt, 0); + assert_eq!(context.state.lock().await.retry_after, None); + assert_eq!(stun.requests.load(Ordering::Relaxed), 1); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} diff --git a/src/transport/middle_proxy/pool_nat/tests.rs b/src/transport/middle_proxy/pool_nat/tests.rs new file mode 100644 index 0000000..0a1c68c --- /dev/null +++ b/src/transport/middle_proxy/pool_nat/tests.rs @@ -0,0 +1,294 @@ +use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; + +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, UdpSocket}; +use tokio::sync::Semaphore; +use tokio::task::JoinHandle; + +use super::*; +use crate::network::probe::NetworkDecision; +use crate::transport::middle_proxy::pool_writer_security_tests::make_pool_with_decision; + +/// Public IPv4 used as source material by local discovery and encrypted ME fixtures. +pub(super) const PUBLIC_IP: Ipv4Addr = Ipv4Addr::new(91, 108, 56, 1); +/// Bounds asynchronous fixture assertions without relying on external services. +pub(super) const TEST_DEADLINE: Duration = Duration::from_secs(2); + +// Concurrency and end-to-end regressions share gated local discovery servers. +pub(super) mod lifecycle; +mod recovery; + +/// Gated local discovery service with observable request counts. +pub(super) struct DiscoveryServer { + /// Listening address selected by the kernel. + pub(super) addr: SocketAddr, + /// Number of requests actually received from clients. + pub(super) requests: Arc, + arrived: Arc, + /// One permit releases one discovery response. + pub(super) release: Arc, + task: JoinHandle<()>, +} + +impl Drop for DiscoveryServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +impl DiscoveryServer { + /// Waits until a real discovery request reaches the server. + pub(super) async fn received(&self) { + tokio::time::timeout(TEST_DEADLINE, self.arrived.acquire()) + .await + .expect("discovery request must reach the local server") + .unwrap() + .forget(); + } +} + +/// Serves valid STUN responses with a deliberately different reflected port. +pub(super) async fn stun_server(bind: SocketAddr, reflected: IpAddr) -> DiscoveryServer { + stun_server_response(bind, Some(reflected), Duration::ZERO).await +} + +async fn stun_server_response( + bind: SocketAddr, + reflected: Option, + delay: Duration, +) -> DiscoveryServer { + let socket = UdpSocket::bind(bind).await.unwrap(); + let addr = socket.local_addr().unwrap(); + let requests = Arc::new(AtomicUsize::new(0)); + let arrived = Arc::new(Semaphore::new(0)); + let release = Arc::new(Semaphore::new(0)); + let count = requests.clone(); + let signal = arrived.clone(); + let gate = release.clone(); + let task = tokio::spawn(async move { + let mut request = [0u8; 512]; + loop { + let (len, peer) = socket.recv_from(&mut request).await.unwrap(); + assert!(len >= 20); + count.fetch_add(1, Ordering::Relaxed); + signal.add_permits(1); + gate.acquire().await.unwrap().forget(); + tokio::time::sleep(delay).await; + let Some(reflected) = reflected else { + socket.send_to(&[0u8], peer).await.unwrap(); + continue; + }; + let (family, ip) = match reflected { + IpAddr::V4(ip) => (1u8, ip.octets().to_vec()), + IpAddr::V6(ip) => (2u8, ip.octets().to_vec()), + }; + let attribute_len = 4 + ip.len(); + let mut response = vec![0u8; 24 + attribute_len]; + response[..2].copy_from_slice(&0x0101u16.to_be_bytes()); + response[2..4].copy_from_slice(&((4 + attribute_len) as u16).to_be_bytes()); + response[4..20].copy_from_slice(&request[4..20]); + response[20..22].copy_from_slice(&1u16.to_be_bytes()); + response[22..24].copy_from_slice(&(attribute_len as u16).to_be_bytes()); + response[25] = family; + response[26..28].copy_from_slice(&45678u16.to_be_bytes()); + response[28..].copy_from_slice(&ip); + socket.send_to(&response, peer).await.unwrap(); + } + }); + DiscoveryServer { + addr, + requests, + arrived, + release, + task, + } +} + +async fn http_server(body: &'static str) -> DiscoveryServer { + http_server_with_delay(body, Duration::ZERO).await +} + +async fn http_server_with_delay(body: &'static str, delay: Duration) -> DiscoveryServer { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let requests = Arc::new(AtomicUsize::new(0)); + let arrived = Arc::new(Semaphore::new(0)); + let release = Arc::new(Semaphore::new(0)); + let count = requests.clone(); + let signal = arrived.clone(); + let gate = release.clone(); + let task = tokio::spawn(async move { + loop { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = [0u8; 4096]; + let _ = stream.read(&mut request).await.unwrap(); + count.fetch_add(1, Ordering::Relaxed); + signal.add_permits(1); + gate.acquire().await.unwrap().forget(); + tokio::time::sleep(delay).await; + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ); + let _ = stream.write_all(response.as_bytes()).await; + } + }); + DiscoveryServer { + addr, + requests, + arrived, + release, + task, + } +} + +/// Builds an isolated dual-stack pool using only caller-supplied discovery services. +pub(super) async fn nat_pool(servers: Vec, urls: Vec) -> Arc { + let mut pool = make_pool_with_decision(NetworkDecision { + ipv4_me: true, + ipv6_me: true, + ..NetworkDecision::default() + }) + .await; + let pool_mut = Arc::get_mut(&mut pool).unwrap(); + let nat = Arc::get_mut(&mut pool_mut.nat_runtime).unwrap(); + nat.nat_probe = true; + nat.nat_stun_servers = servers; + nat.http_ip_detect_urls = urls; + nat.nat_probe_concurrency = 2; + pool +} + +async fn wait_for_reflection(pool: &MePool, family: IpFamily) -> SocketAddr { + tokio::time::timeout(TEST_DEADLINE, async { + loop { + let cached = { + let cache = pool.nat_runtime.nat_reflection_cache.lock().await; + match family { + IpFamily::V4 => cache.v4, + IpFamily::V6 => cache.v6, + } + }; + if let Some((created, addr)) = cached + && created.elapsed() < Duration::from_secs(600) + { + return addr; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("pool-owned STUN refresh must publish after its caller is cancelled") +} + +#[tokio::test] +async fn expired_stun_refresh_survives_reconnect_timeout() { + let fast = stun_server("127.0.0.1:0".parse().unwrap(), PUBLIC_IP.into()).await; + let slow = stun_server("127.0.0.1:0".parse().unwrap(), PUBLIC_IP.into()).await; + fast.release.add_permits(16); + let pool = nat_pool(vec![fast.addr.to_string(), slow.addr.to_string()], vec![]).await; + pool.nat_runtime.nat_reflection_cache.lock().await.v4 = Some(( + Instant::now() - Duration::from_secs(601), + SocketAddr::new(PUBLIC_IP.into(), 1), + )); + let caller_pool = pool.clone(); + let caller = tokio::spawn(async move { + tokio::time::timeout( + Duration::from_millis(100), + caller_pool.maybe_reflect_public_addr(IpFamily::V4, None), + ) + .await + }); + fast.received().await; + slow.received().await; + assert!(caller.await.unwrap().is_err()); + slow.release.add_permits(16); + assert_eq!( + wait_for_reflection(&pool, IpFamily::V4).await.ip(), + PUBLIC_IP + ); + assert_eq!(fast.requests.load(Ordering::Relaxed), 1); + assert_eq!(slow.requests.load(Ordering::Relaxed), 1); + assert_eq!( + pool.nat_runtime.nat_probe_attempts.load(Ordering::Relaxed), + 0 + ); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} + +#[tokio::test] +async fn http_discovery_survives_reconnect_timeout() { + let http = http_server("91.108.56.1").await; + let pool = nat_pool(vec![], vec![format!("http://{}/", http.addr)]).await; + let caller_pool = pool.clone(); + let caller = tokio::spawn(async move { + tokio::time::timeout( + Duration::from_millis(100), + caller_pool.maybe_detect_nat_ip(Ipv4Addr::LOCALHOST.into()), + ) + .await + }); + http.received().await; + assert!(caller.await.unwrap().is_err()); + http.release.add_permits(16); + tokio::time::timeout(TEST_DEADLINE, async { + while pool.nat_runtime.nat_ip_detected.read().await.is_none() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("pool-owned HTTP discovery must publish after caller cancellation"); + assert_eq!( + *pool.nat_runtime.nat_ip_detected.read().await, + Some(PUBLIC_IP.into()) + ); + assert_eq!(http.requests.load(Ordering::Relaxed), 1); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} + +#[tokio::test] +async fn failed_http_discovery_is_reused_during_negative_cache_window() { + let http = http_server("not an IP address").await; + http.release.add_permits(16); + let pool = nat_pool(vec![], vec![format!("http://{}/", http.addr)]).await; + for _ in 0..2 { + assert!( + pool.maybe_detect_nat_ip(Ipv4Addr::LOCALHOST.into()) + .await + .is_none() + ); + } + assert_eq!( + http.requests.load(Ordering::Relaxed), + 1, + "a failed HTTP lookup must not restart ahead of every STUN/handshake attempt" + ); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} + +#[tokio::test] +async fn successful_stun_publication_waits_for_cache_reader() { + let server = stun_server("127.0.0.1:0".parse().unwrap(), PUBLIC_IP.into()).await; + let pool = nat_pool(vec![server.addr.to_string()], vec![]).await; + let caller_pool = pool.clone(); + let caller = tokio::spawn(async move { + caller_pool + .maybe_reflect_public_addr(IpFamily::V4, None) + .await + }); + server.received().await; + let cache_reader = pool.nat_runtime.nat_reflection_cache.lock().await; + server.release.add_permits(16); + tokio::time::sleep(Duration::from_millis(50)).await; + drop(cache_reader); + assert_eq!(caller.await.unwrap().unwrap().ip(), PUBLIC_IP); + assert_eq!( + wait_for_reflection(&pool, IpFamily::V4).await.ip(), + PUBLIC_IP + ); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} diff --git a/src/transport/middle_proxy/pool_nat/tests/lifecycle.rs b/src/transport/middle_proxy/pool_nat/tests/lifecycle.rs new file mode 100644 index 0000000..1365af3 --- /dev/null +++ b/src/transport/middle_proxy/pool_nat/tests/lifecycle.rs @@ -0,0 +1,342 @@ +use super::*; +use crate::config::{UpstreamConfig, UpstreamType}; +use crate::transport::UpstreamManager; + +/// Gives an isolated pool the same finite bind-address authority as production egress. +pub(in crate::transport::middle_proxy::pool_nat) fn set_bind_addresses( + pool: &mut Arc, + addresses: &[&str], +) { + let manager = UpstreamManager::new( + vec![UpstreamConfig { + upstream_type: UpstreamType::Direct { + interface: None, + bind_addresses: Some(addresses.iter().map(|value| value.to_string()).collect()), + bindtodevice: None, + }, + weight: 1, + enabled: true, + scopes: String::new(), + selected_scope: String::new(), + ipv4: None, + ipv6: None, + prefer: None, + }], + 1, + 0, + 1000, + 1, + 1, + true, + pool.stats.clone(), + ); + Arc::get_mut(pool).unwrap().upstream = Some(Arc::new(manager)); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancelled_waiter_storm_shares_one_http_and_stun_producer() { + let http = http_server("91.108.56.1").await; + let stun = stun_server("127.0.0.1:0".parse().unwrap(), PUBLIC_IP.into()).await; + let pool = nat_pool( + vec![stun.addr.to_string()], + vec![format!("http://{}/", http.addr)], + ) + .await; + let mut waiters = JoinSet::new(); + for _ in 0..32 { + let pool = pool.clone(); + waiters.spawn(async move { + tokio::join!( + pool.maybe_detect_nat_ip(Ipv4Addr::LOCALHOST.into()), + pool.maybe_reflect_public_addr(IpFamily::V4, None), + ); + }); + } + http.received().await; + stun.received().await; + waiters.abort_all(); + while waiters.join_next().await.is_some() {} + http.release.add_permits(64); + stun.release.add_permits(64); + assert_eq!( + pool.maybe_detect_nat_ip(Ipv4Addr::LOCALHOST.into()).await, + Some(PUBLIC_IP.into()) + ); + assert_eq!( + pool.maybe_reflect_public_addr(IpFamily::V4, None) + .await + .unwrap() + .ip(), + PUBLIC_IP + ); + assert_eq!(http.requests.load(Ordering::Relaxed), 1); + assert_eq!(stun.requests.load(Ordering::Relaxed), 1); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} + +#[tokio::test] +async fn shutdown_cancels_discovery_and_releases_all_waiters() { + let http = http_server("91.108.56.1").await; + let stun = stun_server("127.0.0.1:0".parse().unwrap(), PUBLIC_IP.into()).await; + let pool = nat_pool( + vec![stun.addr.to_string()], + vec![format!("http://{}/", http.addr)], + ) + .await; + let waiter_pool = pool.clone(); + let waiter = tokio::spawn(async move { + tokio::join!( + waiter_pool.maybe_detect_nat_ip(Ipv4Addr::LOCALHOST.into()), + waiter_pool.maybe_reflect_public_addr(IpFamily::V4, None), + ) + }); + http.received().await; + stun.received().await; + assert!(pool.shutdown_until(TEST_DEADLINE).await); + assert_eq!( + timeout(TEST_DEADLINE, waiter).await.unwrap().unwrap(), + (None, None) + ); + http.release.add_permits(64); + stun.release.add_permits(64); + assert_eq!( + pool.maybe_detect_nat_ip(Ipv4Addr::LOCALHOST.into()).await, + None + ); + assert_eq!( + pool.maybe_reflect_public_addr(IpFamily::V4, None).await, + None + ); + assert_eq!(*pool.nat_runtime.nat_ip_detected.read().await, None); + assert_eq!(pool.nat_runtime.nat_reflection_cache.lock().await.v4, None); + assert_eq!(http.requests.load(Ordering::Relaxed), 1); + assert_eq!(stun.requests.load(Ordering::Relaxed), 1); + assert!( + pool.nat_runtime + .nat_reflection_singleflight_v4 + .try_lock() + .is_ok() + ); + assert!(pool.nat_runtime.discovery.http.try_lock().is_ok()); +} + +#[tokio::test] +async fn bound_refresh_survives_cancellation_and_reuses_only_its_source() { + let stun = stun_server("127.0.0.1:0".parse().unwrap(), PUBLIC_IP.into()).await; + let mut pool = nat_pool(vec![stun.addr.to_string()], vec![]).await; + set_bind_addresses(&mut pool, &["127.0.0.1", "127.0.0.2"]); + let bind = Some(Ipv4Addr::LOCALHOST.into()); + let waiter_pool = pool.clone(); + let waiter = tokio::spawn(async move { + waiter_pool + .maybe_reflect_public_addr(IpFamily::V4, bind) + .await + }); + stun.received().await; + waiter.abort(); + assert!(waiter.await.unwrap_err().is_cancelled()); + stun.release.add_permits(64); + assert_eq!( + pool.maybe_reflect_public_addr(IpFamily::V4, bind) + .await + .unwrap() + .ip(), + PUBLIC_IP + ); + assert_eq!(stun.requests.load(Ordering::Relaxed), 1); + assert_eq!(pool.nat_runtime.nat_reflection_cache.lock().await.v4, None); + let other = Some("127.0.0.2".parse().unwrap()); + assert_eq!( + pool.maybe_reflect_public_addr(IpFamily::V4, other) + .await + .unwrap() + .ip(), + PUBLIC_IP + ); + assert_eq!(stun.requests.load(Ordering::Relaxed), 2); + assert!( + pool.maybe_reflect_public_addr(IpFamily::V4, Some("127.0.0.3".parse().unwrap())) + .await + .is_none() + ); + assert!( + pool.maybe_reflect_public_addr(IpFamily::V6, bind) + .await + .is_none() + ); + assert_eq!(stun.requests.load(Ordering::Relaxed), 2); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} + +#[tokio::test] +async fn failed_bound_refresh_is_observable_after_caller_cancellation() { + let stun = stun_server_response("127.0.0.1:0".parse().unwrap(), None, Duration::ZERO).await; + let mut pool = nat_pool(vec![stun.addr.to_string()], vec![]).await; + set_bind_addresses(&mut pool, &["127.0.0.1"]); + let bind = Some(Ipv4Addr::LOCALHOST.into()); + let waiter_pool = pool.clone(); + let waiter = tokio::spawn(async move { + waiter_pool + .maybe_reflect_public_addr(IpFamily::V4, bind) + .await + }); + stun.received().await; + waiter.abort(); + let _ = waiter.await; + // One response completes the original probe; a restarted probe would block forever here. + stun.release.add_permits(1); + assert_eq!( + timeout( + TEST_DEADLINE, + pool.maybe_reflect_public_addr(IpFamily::V4, bind) + ) + .await + .unwrap(), + None + ); + assert_eq!( + timeout( + TEST_DEADLINE, + pool.maybe_reflect_public_addr(IpFamily::V4, bind) + ) + .await + .unwrap(), + None + ); + assert_eq!(stun.requests.load(Ordering::Relaxed), 1); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} + +#[tokio::test] +async fn http_negative_cache_expires_without_waiting_a_minute() { + let http = http_server("invalid").await; + http.release.add_permits(64); + let pool = nat_pool(vec![], vec![format!("http://{}/", http.addr)]).await; + assert_eq!( + pool.maybe_detect_nat_ip(Ipv4Addr::LOCALHOST.into()).await, + None + ); + *pool.nat_runtime.discovery.http.lock().await = Some(Instant::now() - Duration::from_secs(1)); + assert_eq!( + pool.maybe_detect_nat_ip(Ipv4Addr::LOCALHOST.into()).await, + None + ); + assert_eq!(http.requests.load(Ordering::Relaxed), 2); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} + +#[tokio::test] +async fn ipv6_refresh_is_independent_of_blocked_ipv4_refresh() { + let v4 = stun_server("127.0.0.1:0".parse().unwrap(), PUBLIC_IP.into()).await; + let public_v6: IpAddr = "2001:4860:4860::8888".parse().unwrap(); + let v6 = stun_server("[::1]:0".parse().unwrap(), public_v6).await; + v6.release.add_permits(16); + let pool = nat_pool(vec![v4.addr.to_string(), v6.addr.to_string()], vec![]).await; + let waiter_pool = pool.clone(); + let waiter = tokio::spawn(async move { + waiter_pool + .maybe_reflect_public_addr(IpFamily::V4, None) + .await + }); + v4.received().await; + assert_eq!( + timeout( + TEST_DEADLINE, + pool.maybe_reflect_public_addr(IpFamily::V6, None) + ) + .await + .unwrap() + .unwrap() + .ip(), + public_v6 + ); + assert!(!waiter.is_finished()); + v4.release.add_permits(16); + assert_eq!(waiter.await.unwrap().unwrap().ip(), PUBLIC_IP); + assert_eq!( + pool.nat_runtime + .nat_reflection_cache + .lock() + .await + .v4 + .unwrap() + .1 + .ip(), + PUBLIC_IP + ); + assert_eq!( + pool.nat_runtime + .nat_reflection_cache + .lock() + .await + .v6 + .unwrap() + .1 + .ip(), + public_v6 + ); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} + +#[tokio::test] +async fn failed_shared_refresh_publishes_backoff_after_caller_cancellation() { + let stun = stun_server_response("127.0.0.1:0".parse().unwrap(), None, Duration::ZERO).await; + let pool = nat_pool(vec![stun.addr.to_string()], vec![]).await; + let waiter_pool = pool.clone(); + let waiter = tokio::spawn(async move { + waiter_pool + .maybe_reflect_public_addr(IpFamily::V4, None) + .await + }); + stun.received().await; + waiter.abort(); + let _ = waiter.await; + stun.release.add_permits(1); + assert_eq!( + timeout( + TEST_DEADLINE, + pool.maybe_reflect_public_addr(IpFamily::V4, None) + ) + .await + .unwrap(), + None + ); + assert!(pool.nat_runtime.stun_backoff_until.read().await.unwrap() > Instant::now()); + assert_eq!( + pool.maybe_reflect_public_addr(IpFamily::V4, None).await, + None + ); + assert_eq!(stun.requests.load(Ordering::Relaxed), 1); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} + +#[tokio::test] +async fn shutdown_before_producer_first_poll_releases_owned_gates() { + let http = http_server("91.108.56.1").await; + let stun = stun_server("127.0.0.1:0".parse().unwrap(), PUBLIC_IP.into()).await; + let pool = nat_pool( + vec![stun.addr.to_string()], + vec![format!("http://{}/", http.addr)], + ) + .await; + let mut http_waiter = + tokio_test::task::spawn(pool.maybe_detect_nat_ip(Ipv4Addr::LOCALHOST.into())); + let mut stun_waiter = + tokio_test::task::spawn(pool.maybe_reflect_public_addr(IpFamily::V4, None)); + assert!(http_waiter.poll().is_pending()); + assert!(stun_waiter.poll().is_pending()); + // On this current-thread runtime neither producer can run before shutdown admission closes. + pool.begin_shutdown(); + assert!(pool.shutdown_until(TEST_DEADLINE).await); + assert_eq!(http_waiter.poll(), std::task::Poll::Ready(None)); + assert_eq!(stun_waiter.poll(), std::task::Poll::Ready(None)); + assert_eq!(http.requests.load(Ordering::Relaxed), 0); + assert_eq!(stun.requests.load(Ordering::Relaxed), 0); + assert!(pool.nat_runtime.discovery.http.try_lock().is_ok()); + assert!( + pool.nat_runtime + .nat_reflection_singleflight_v4 + .try_lock() + .is_ok() + ); +} diff --git a/src/transport/middle_proxy/pool_nat/tests/recovery.rs b/src/transport/middle_proxy/pool_nat/tests/recovery.rs new file mode 100644 index 0000000..911fb32 --- /dev/null +++ b/src/transport/middle_proxy/pool_nat/tests/recovery.rs @@ -0,0 +1,266 @@ +use super::*; +use crate::crypto::{SecureRandom, derive_middleproxy_keys}; +use crate::transport::middle_proxy::codec::{ + RpcChecksumMode, build_handshake_payload, build_nonce_payload, build_rpc_frame, + cbc_decrypt_inplace, cbc_encrypt_padded, parse_handshake_flags, parse_nonce_payload, + read_rpc_frame_plaintext, rpc_crc, +}; +use crate::transport::middle_proxy::me_health_monitor; +use crate::transport::middle_proxy::pool::MeFamilyRuntimeState; +use tokio::net::TcpStream; + +struct MiddlePeer { + addr: SocketAddr, + handshakes: Arc, + cancelled: Arc, + task: JoinHandle<()>, +} + +impl Drop for MiddlePeer { + fn drop(&mut self) { + self.task.abort(); + } +} + +impl MiddlePeer { + async fn new(server_ip: Ipv4Addr) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let handshakes = Arc::new(AtomicUsize::new(0)); + let cancelled = Arc::new(AtomicUsize::new(0)); + let succeeded = handshakes.clone(); + let closed = cancelled.clone(); + let task = tokio::spawn(async move { + let mut pending = JoinSet::new(); + let mut established = Vec::new(); + loop { + tokio::select! { + accepted = listener.accept() => { + let (stream, _) = accepted.unwrap(); + pending.spawn(accept_handshake(stream, server_ip)); + } + completed = pending.join_next(), if !pending.is_empty() => { + match completed.unwrap().expect("fake ME handshake task must not panic") { + Some(stream) => { + established.push(stream); + succeeded.fetch_add(1, Ordering::Release); + } + None => { closed.fetch_add(1, Ordering::Release); } + } + } + } + } + }); + Self { + addr, + handshakes, + cancelled, + task, + } + } +} + +async fn accept_handshake(mut stream: TcpStream, server_ip: Ipv4Addr) -> Option { + let client_port = stream.peer_addr().unwrap().port(); + let server_port = stream.local_addr().unwrap().port(); + let (seq, payload) = match read_rpc_frame_plaintext(&mut stream).await { + Ok(frame) => frame, + // Reconnect may time out during NAT discovery, before sending any nonce. + Err(ProxyError::Io(error)) if error.kind() == std::io::ErrorKind::UnexpectedEof => { + return None; + } + Err(error) => panic!("invalid client nonce: {error}"), + }; + assert_eq!(seq, -2); + let (selector, _, timestamp, client_nonce) = parse_nonce_payload(&payload).unwrap(); + let server_nonce = [0x42u8; 16]; + let nonce = build_nonce_payload(selector, timestamp, &server_nonce); + stream + .write_all(&build_rpc_frame(-2, &nonce, RpcChecksumMode::Crc32)) + .await + .unwrap(); + let mut client_ip = PUBLIC_IP.octets(); + client_ip.reverse(); + let mut server_ip = server_ip.octets(); + server_ip.reverse(); + let (read_key, read_iv) = derive_middleproxy_keys( + &server_nonce, + &client_nonce, + ×tamp.to_le_bytes(), + Some(&server_ip), + &client_port.to_le_bytes(), + b"CLIENT", + Some(&client_ip), + &server_port.to_le_bytes(), + &[1u8; 32], + None, + None, + ); + let (write_key, write_iv) = derive_middleproxy_keys( + &server_nonce, + &client_nonce, + ×tamp.to_le_bytes(), + Some(&server_ip), + &client_port.to_le_bytes(), + b"SERVER", + Some(&client_ip), + &server_port.to_le_bytes(), + &[1u8; 32], + None, + None, + ); + let mut encrypted = [0u8; 48]; + stream.read_exact(&mut encrypted).await.unwrap(); + cbc_decrypt_inplace(&read_key, &read_iv, &mut encrypted).unwrap(); + assert_eq!(u32::from_le_bytes(encrypted[..4].try_into().unwrap()), 44); + assert_eq!(i32::from_le_bytes(encrypted[4..8].try_into().unwrap()), -1); + assert_eq!( + u32::from_le_bytes(encrypted[40..44].try_into().unwrap()), + rpc_crc(RpcChecksumMode::Crc32, &encrypted[..40]) + ); + parse_handshake_flags(&encrypted[8..40]).unwrap(); + assert_eq!(&encrypted[16..20], &client_ip); + assert_eq!( + u16::from_le_bytes(encrypted[20..22].try_into().unwrap()), + client_port + ); + assert_eq!( + u16::from_le_bytes(encrypted[32..34].try_into().unwrap()), + server_port + ); + let reply = build_handshake_payload(server_ip, server_port, client_ip, client_port, 0); + let (encrypted, _) = cbc_encrypt_padded( + &write_key, + &write_iv, + &build_rpc_frame(-1, &reply, RpcChecksumMode::Crc32), + ) + .unwrap(); + stream.write_all(&encrypted).await.unwrap(); + Some(stream) +} + +async fn suppressed_pool_recovers_after_delayed_discovery(http_slow: bool) { + // Every lookup exceeds me_one_timeout_ms, not just the first request of the outage. + let delay = Duration::from_millis(1500); + let stun = stun_server_response( + "127.0.0.1:0".parse().unwrap(), + Some(PUBLIC_IP.into()), + delay, + ) + .await; + let http = http_server_with_delay("91.108.56.1", delay).await; + stun.release.add_permits(64); + http.release.add_permits(64); + let peer = MiddlePeer::new(if http_slow { + PUBLIC_IP + } else { + Ipv4Addr::LOCALHOST + }) + .await; + let urls = if http_slow { + vec![format!("http://{}/", http.addr)] + } else { + vec![] + }; + let pool = nat_pool(vec![stun.addr.to_string()], urls).await; + pool.health_runtime + .me_health_interval_ms_unhealthy + .store(10, Ordering::Relaxed); + pool.health_runtime + .me_health_interval_ms_healthy + .store(10, Ordering::Relaxed); + pool.single_endpoint_runtime + .me_single_endpoint_outage_backoff_min_ms + .store(100, Ordering::Relaxed); + pool.single_endpoint_runtime + .me_single_endpoint_outage_backoff_max_ms + .store(100, Ordering::Relaxed); + pool.single_endpoint_runtime + .me_single_endpoint_shadow_writers + .store(0, Ordering::Relaxed); + pool.update_proxy_maps( + HashMap::from([(2, vec![(peer.addr.ip(), peer.addr.port())])]), + None, + ) + .await; + let generation = pool.current_generation(); + let now = MePool::now_epoch_secs(); + pool.set_family_runtime_state( + IpFamily::V4, + MeFamilyRuntimeState::Suppressed, + now, + now + 3600, + 5, + 0, + ); + pool.nat_runtime.nat_reflection_cache.lock().await.v4 = Some(( + if http_slow { + Instant::now() + } else { + Instant::now() - Duration::from_secs(601) + }, + SocketAddr::new(PUBLIC_IP.into(), 45678), + )); + let monitor = tokio::spawn(me_health_monitor( + pool.clone(), + Arc::new(SecureRandom::new()), + 0, + )); + if http_slow { + http.received().await; + } else { + stun.received().await; + } + timeout(Duration::from_secs(5), async { + while peer.cancelled.load(Ordering::Acquire) == 0 { + assert!(!peer.task.is_finished()); + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the actual 1200ms reconnect must cancel at least one pre-KDF TCP connection"); + assert_eq!(peer.handshakes.load(Ordering::Acquire), 0); + let expected = 2 * pool.required_writers_for_dc(1); + timeout(Duration::from_secs(10), async { + while !pool.admission_ready_full_floor().await + || peer.handshakes.load(Ordering::Acquire) < expected + { + assert!( + !peer.task.is_finished(), + "fake ME must validate every encrypted handshake" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("suppressed pool must autonomously regain its full floor without restart"); + monitor.abort(); + let _ = monitor.await; + assert_eq!(pool.current_generation(), generation); + assert_eq!(pool.reinit.status.load().pending_hardswap_generation, 0); + assert_eq!( + pool.writer_connect_active_reserved.load(Ordering::Acquire), + 0 + ); + assert_eq!(pool.writers.read().await.len(), expected); + assert!(peer.handshakes.load(Ordering::Acquire) >= expected); + assert_eq!( + http.requests.load(Ordering::Relaxed), + usize::from(http_slow) + ); + assert_eq!( + stun.requests.load(Ordering::Relaxed), + usize::from(!http_slow) + ); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} + +#[tokio::test] +async fn suppressed_pool_recovers_full_floor_after_expired_stun_refresh() { + suppressed_pool_recovers_after_delayed_discovery(false).await; +} + +#[tokio::test] +async fn suppressed_pool_recovers_full_floor_after_slow_http_with_fresh_stun() { + suppressed_pool_recovers_after_delayed_discovery(true).await; +} diff --git a/src/transport/upstream.rs b/src/transport/upstream.rs index 49141b8..667ad1f 100644 --- a/src/transport/upstream.rs +++ b/src/transport/upstream.rs @@ -339,6 +339,8 @@ pub struct UpstreamManager { // Upstream manager configuration, DNS resolution, and API snapshots. mod manager_config; +// Current direct egress addresses eligible for source-bound NAT discovery. +mod nat_contexts; // Latency- and scope-aware upstream selection. mod selection; // Direct and proxied connection establishment. diff --git a/src/transport/upstream/nat_contexts.rs b/src/transport/upstream/nat_contexts.rs new file mode 100644 index 0000000..0727142 --- /dev/null +++ b/src/transport/upstream/nat_contexts.rs @@ -0,0 +1,122 @@ +use super::*; + +impl UpstreamManager { + /// Resolves this generation's direct bind candidates without advancing round-robin state. + pub(crate) async fn direct_bind_ips_for_family(&self, ipv6: bool) -> BTreeSet { + let guard = self.upstreams.read().await; + let target = SocketAddr::new( + if ipv6 { + std::net::Ipv6Addr::UNSPECIFIED.into() + } else { + std::net::Ipv4Addr::UNSPECIFIED.into() + }, + 0, + ); + let mut allowed = BTreeSet::new(); + for state in guard.iter() { + let config = &state.config; + let UpstreamType::Direct { + interface, + bind_addresses, + .. + } = &config.upstream_type + else { + continue; + }; + if (if ipv6 { config.ipv6 } else { config.ipv4 }) == Some(false) { + continue; + } + if let Some(addresses) = bind_addresses.as_ref().filter(|v| !v.is_empty()) { + for address in addresses { + if !address + .parse::() + .is_ok_and(|ip| ip.is_ipv6() == ipv6) + { + continue; + } + if let Some(ip) = Self::resolve_bind_address( + interface, + &Some(vec![address.clone()]), + target, + None, + true, + ) { + allowed.insert(ip); + } + } + } else if let Some(ip) = + Self::resolve_bind_address(interface, bind_addresses, target, None, true) + { + allowed.insert(ip); + } + } + allowed + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config(source: &str) -> UpstreamConfig { + toml::from_str(source).unwrap() + } + + #[tokio::test] + async fn candidates_follow_generation_policy_without_advancing_round_robin() { + let manager = UpstreamManager::new( + vec![ + config("type = 'direct'\nbind_addresses = ['127.0.0.1', '127.0.0.2', '::1']"), + config("type = 'direct'\ninterface = '127.0.0.3'"), + config("type = 'direct'\ninterface = '127.0.0.4'\nenabled = false"), + config("type = 'direct'\ninterface = '127.0.0.5'\nipv4 = false"), + config("type = 'socks5'\naddress = '127.0.0.1:1080'\ninterface = '127.0.0.6'"), + ], + 1, + 0, + 1000, + 1, + 1, + true, + Arc::new(Stats::new()), + ); + manager.upstreams.write().await[0].healthy = false; + let expected = ["127.0.0.1", "127.0.0.2", "127.0.0.3"] + .into_iter() + .map(|value| value.parse().unwrap()) + .collect(); + assert_eq!(manager.direct_bind_ips_for_family(false).await, expected); + assert_eq!( + manager.direct_bind_ips_for_family(true).await, + BTreeSet::from(["::1".parse().unwrap()]) + ); + for state in manager.upstreams.read().await.iter() { + assert_eq!(state.bind_rr.load(Ordering::Relaxed), 0); + } + } + + #[cfg(target_os = "linux")] + #[tokio::test] + async fn named_interface_candidates_match_direct_socket_binding() { + let manager = UpstreamManager::new( + vec![config( + "type = 'direct'\ninterface = 'lo'\nbind_addresses = ['127.0.0.1', '::1']", + )], + 1, + 0, + 1000, + 1, + 1, + true, + Arc::new(Stats::new()), + ); + assert_eq!( + manager.direct_bind_ips_for_family(false).await, + BTreeSet::from(["127.0.0.1".parse().unwrap()]) + ); + assert_eq!( + manager.direct_bind_ips_for_family(true).await, + BTreeSet::from(["::1".parse().unwrap()]) + ); + } +}