ME NAT Discovery Cancellation correctness

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

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<Mutex<Option<Instant>>>,
+    bound: Mutex<HashMap<(IpFamily, IpAddr), Arc<BoundDiscovery>>>,
+}
+
+#[derive(Default)]
+struct BoundDiscovery {
+    state: Arc<Mutex<BoundDiscoveryState>>,
+    cancel: CancellationToken,
+}
+
+#[derive(Default)]
+struct BoundDiscoveryState {
+    reflection: Option<(Instant, SocketAddr)>,
+    retry_after: Option<Instant>,
+    attempt: u8,
+}
+
+impl NatDiscovery {
+    async fn bound_context(
+        &self,
+        family: IpFamily,
+        bind_ip: IpAddr,
+        allowed: &BTreeSet<IpAddr>,
+    ) -> Option<Arc<BoundDiscovery>> {
+        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<Self>,
+        family: IpFamily,
+        bind_ip: IpAddr,
+    ) -> Option<SocketAddr> {
+        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()
+    }
+}

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<Arc<BoundDiscovery>> = 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);
+}

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<AtomicUsize>,
+    arrived: Arc<Semaphore>,
+    /// One permit releases one discovery response.
+    pub(super) release: Arc<Semaphore>,
+    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<IpAddr>,
+    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<String>, urls: Vec<String>) -> Arc<MePool> {
+    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);
+}

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<MePool>,
+    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()
+    );
+}

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<AtomicUsize>,
+    cancelled: Arc<AtomicUsize>,
+    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<TcpStream> {
+    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,
+        &timestamp.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,
+        &timestamp.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;
+}