astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-10-21 08:14:24 +03:00
Code Issues Packages Projects Releases Wiki Activity

ME NAT Discovery Cancellation correctness

Browse Source 
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-10-03 12:36:29 +03:00
parent d65efcc479
commit 84d5e98c11
11 changed files with 1410 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/transport/middle_proxy/handshake.rs
+2 -1
View File
@@ -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<Self>,
stream: TcpStream,
addr: SocketAddr,
upstream_egress: Option<UpstreamEgressInfo>,
src/transport/middle_proxy/pool.rs
+1
View File
@@ -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<IpAddr>,
pub(super) nat_ip_detected: Arc<RwLock<Option<IpAddr>>>,
pub(super) nat_probe: bool,
src/transport/middle_proxy/pool/construction.rs
+1
View File
@@ -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,
src/transport/middle_proxy/pool_nat.rs
+115 -103
View File
@@ -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<String>) -> Result<crate::network::stun::DualStunResult> {
let stun_addr = stun_addr.unwrap_or_else(|| {
@@ -35,6 +48,7 @@ pub async fn stun_probe(stun_addr: Option<String>) -> Result<crate::network::stu
stun_probe_dual_with_tcp_fallback(&stun_addr, false).await
}
/// Detects a public IPv4 address using the default HTTP endpoints.
#[allow(dead_code)]
pub async fn detect_public_ip() -> Option<IpAddr> {
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<IpAddr> {
/// Waits for pool-owned HTTP discovery without owning its cancellation.
pub(super) async fn maybe_detect_nat_ip(self: &Arc<Self>, local_ip: IpAddr) -> Option<IpAddr> {
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<Self>,
family: IpFamily,
bind_ip: Option<IpAddr>,
) -> Option<std::net::SocketAddr> {
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<Option<std::net::SocketAddr>> {
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<IpAddr>,
attempt: u8,
) -> Option<std::net::SocketAddr> {
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
}
}
src/transport/middle_proxy/pool_nat/state.rs
+119
View File
@@ -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
+146
View File
@@ -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
+294
View File
@@ -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
+342
View File
@@ -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
+266
View File
@@ -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;
}
src/transport/upstream.rs
+2
View File
@@ -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.
src/transport/upstream/nat_contexts.rs
+122
View File
@@ -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<IpAddr> {
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::<IpAddr>()
.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()])
);
}
}