astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-04-17 10:34:11 +03:00
Code Issues Packages Projects Releases Wiki Activity

ME Probe parallelized

Browse Source
This commit is contained in:
Alexey
2026-02-26 17:56:22 +03:00
parent 92a3529733
commit 7782336264
8 changed files with 411 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

170
src/transport/middle_proxy/pool_nat.rs
View File

@@ -1,7 +1,10 @@
use std::collections::HashMap;
use std::net::{IpAddr, Ipv4Addr};
use std::time::Duration;
use tracing::{info, warn};
use tokio::task::JoinSet;
use tokio::time::timeout;
use tracing::{debug, info, warn};
use crate::error::{ProxyError, Result};
use crate::network::probe::is_bogon;
@@ -10,6 +13,8 @@ use crate::network::stun::{stun_probe_dual, IpFamily, StunProbeResult};
use super::MePool;
use std::time::Instant;
const STUN_BATCH_TIMEOUT: Duration = Duration::from_secs(5);
#[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(|| "stun.l.google.com:19302".to_string());
@@ -22,6 +27,99 @@ pub async fn detect_public_ip() -> Option<IpAddr> {
}
impl MePool {
fn configured_stun_servers(&self) -> Vec<String> {
if !self.nat_stun_servers.is_empty() {
return self.nat_stun_servers.clone();
}
if let Some(s) = &self.nat_stun {
return vec![s.clone()];
}
vec!["stun.l.google.com:19302".to_string()]
}
async fn probe_stun_batch_for_family(
&self,
servers: &[String],
family: IpFamily,
attempt: u8,
) -> (Vec<String>, Option<std::net::SocketAddr>) {
let mut join_set = JoinSet::new();
let mut next_idx = 0usize;
let mut live_servers = Vec::new();
let mut best_by_ip: HashMap<IpAddr, (usize, std::net::SocketAddr)> = HashMap::new();
let concurrency = self.nat_probe_concurrency.max(1);
while next_idx < servers.len() || !join_set.is_empty() {
while next_idx < servers.len() && join_set.len() < concurrency {
let stun_addr = servers[next_idx].clone();
next_idx += 1;
join_set.spawn(async move {
let res = timeout(STUN_BATCH_TIMEOUT, stun_probe_dual(&stun_addr)).await;
(stun_addr, res)
});
}
let Some(task) = join_set.join_next().await else {
break;
};
match task {
Ok((stun_addr, Ok(Ok(res)))) => {
let picked: Option<StunProbeResult> = match family {
IpFamily::V4 => res.v4,
IpFamily::V6 => res.v6,
};
if let Some(result) = picked {
live_servers.push(stun_addr.clone());
let entry = best_by_ip
.entry(result.reflected_addr.ip())
.or_insert((0, result.reflected_addr));
entry.0 += 1;
debug!(
local = %result.local_addr,
reflected = %result.reflected_addr,
family = ?family,
stun = %stun_addr,
"NAT probe: reflected address"
);
}
}
Ok((stun_addr, Ok(Err(e)))) => {
debug!(
error = %e,
stun = %stun_addr,
attempt = attempt + 1,
"NAT probe failed, trying next server"
);
}
Ok((stun_addr, Err(_))) => {
debug!(
stun = %stun_addr,
attempt = attempt + 1,
"NAT probe timeout, trying next server"
);
}
Err(e) => {
debug!(
error = %e,
attempt = attempt + 1,
"NAT probe task join failed"
);
}
}
}
live_servers.sort_unstable();
live_servers.dedup();
let best_reflected = best_by_ip
.into_values()
.max_by_key(|(count, _)| *count)
.map(|(_, addr)| addr);
(live_servers, best_reflected)
}
pub(super) fn translate_ip_for_nat(&self, ip: IpAddr) -> IpAddr {
let nat_ip = self
.nat_ip_cfg
@@ -128,39 +226,51 @@ impl MePool {
}
let attempt = self.nat_probe_attempts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let servers = if !self.nat_stun_servers.is_empty() {
self.nat_stun_servers.clone()
} else if let Some(s) = &self.nat_stun {
vec![s.clone()]
let configured_servers = self.configured_stun_servers();
let live_snapshot = self.nat_stun_live_servers.read().await.clone();
let primary_servers = if live_snapshot.is_empty() {
configured_servers.clone()
} else {
vec!["stun.l.google.com:19302".to_string()]
live_snapshot
};
for stun_addr in servers {
match stun_probe_dual(&stun_addr).await {
Ok(res) => {
let picked: Option<StunProbeResult> = match family {
IpFamily::V4 => res.v4,
IpFamily::V6 => res.v6,
};
if let Some(result) = picked {
info!(local = %result.local_addr, reflected = %result.reflected_addr, family = ?family, stun = %stun_addr, "NAT probe: reflected address");
self.nat_probe_attempts.store(0, std::sync::atomic::Ordering::Relaxed);
if let Ok(mut cache) = self.nat_reflection_cache.try_lock() {
let slot = match family {
IpFamily::V4 => &mut cache.v4,
IpFamily::V6 => &mut cache.v6,
};
*slot = Some((Instant::now(), result.reflected_addr));
}
return Some(result.reflected_addr);
}
}
Err(e) => {
warn!(error = %e, stun = %stun_addr, attempt = attempt + 1, "NAT probe failed, trying next server");
}
}
let (mut live_servers, mut selected_reflected) = self
.probe_stun_batch_for_family(&primary_servers, family, attempt)
.await;
if selected_reflected.is_none() && !configured_servers.is_empty() && primary_servers != configured_servers {
let (rediscovered_live, rediscovered_reflected) = self
.probe_stun_batch_for_family(&configured_servers, family, attempt)
.await;
live_servers = rediscovered_live;
selected_reflected = rediscovered_reflected;
}
let live_server_count = live_servers.len();
if !live_servers.is_empty() {
*self.nat_stun_live_servers.write().await = live_servers;
} else {
self.nat_stun_live_servers.write().await.clear();
}
if let Some(reflected_addr) = selected_reflected {
self.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 let Ok(mut cache) = self.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);
}
let backoff = Duration::from_secs(60 * 2u64.pow((attempt as u32).min(6)));
*self.stun_backoff_until.write().await = Some(Instant::now() + backoff);
None