Antireplay Improvements + DC Ping
- Fix: LruCache::get type ambiguity in stats/mod.rs - Changed `self.cache.get(&key.into())` to `self.cache.get(key)` (key is already &[u8], resolved via Box<[u8]>: Borrow<[u8]>) - Changed `self.cache.peek(&key)` / `.pop(&key)` to `.peek(key.as_ref())` / `.pop(key.as_ref())` (explicit &[u8] instead of &Box<[u8]>) - Startup DC ping with RTT display and improved health-check (all DCs, RTT tracking, EMA latency, 30s interval): - Implemented `LatencyEma` – exponential moving average (α=0.3) for RTT - `connect()` – measures RTT of each real connection and updates EMA - `ping_all_dcs()` – pings all 5 DCs via each upstream, returns `Vec<StartupPingResult>` with RTT or error - `run_health_checks(prefer_ipv6)` – accepts IPv6 preference parameter, rotates DC between cycles (DC1→DC2→...→DC5→DC1...), interval reduced to 30s from 60s, failed checks now mark upstream as unhealthy after 3 consecutive fails - `DcPingResult` / `StartupPingResult` – public structures for display - DC Ping at startup: calls `upstream_manager.ping_all_dcs()` before accept loop, outputs table via `println!` (always visible) - Health checks with `prefer_ipv6`: `run_health_checks(prefer_ipv6)` receives the parameter - Exported `StartupPingResult` and `DcPingResult` - Summary: Startup DC ping with RTT, rotational health-check with EMA latency tracking, 30-second interval, correct unhealthy marking after 3 fails. Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
125
src/main.rs
125
src/main.rs
@@ -26,11 +26,6 @@ use crate::transport::{create_listener, ListenOptions, UpstreamManager};
|
||||
use crate::util::ip::detect_ip;
|
||||
use crate::stream::BufferPool;
|
||||
|
||||
/// Parse command-line arguments.
|
||||
///
|
||||
/// Usage: telemt [config_path] [--silent] [--log-level <level>]
|
||||
///
|
||||
/// Returns (config_path, silent_flag, log_level_override)
|
||||
fn parse_cli() -> (String, bool, Option<String>) {
|
||||
let mut config_path = "config.toml".to_string();
|
||||
let mut silent = false;
|
||||
@@ -40,33 +35,23 @@ fn parse_cli() -> (String, bool, Option<String>) {
|
||||
let mut i = 0;
|
||||
while i < args.len() {
|
||||
match args[i].as_str() {
|
||||
"--silent" | "-s" => {
|
||||
silent = true;
|
||||
}
|
||||
"--silent" | "-s" => { silent = true; }
|
||||
"--log-level" => {
|
||||
i += 1;
|
||||
if i < args.len() {
|
||||
log_level = Some(args[i].clone());
|
||||
}
|
||||
if i < args.len() { log_level = Some(args[i].clone()); }
|
||||
}
|
||||
s if s.starts_with("--log-level=") => {
|
||||
log_level = Some(s.trim_start_matches("--log-level=").to_string());
|
||||
}
|
||||
"--help" | "-h" => {
|
||||
eprintln!("Usage: telemt [config.toml] [OPTIONS]");
|
||||
eprintln!();
|
||||
eprintln!("Options:");
|
||||
eprintln!(" --silent, -s Suppress info logs (only warn/error)");
|
||||
eprintln!(" --log-level <LEVEL> Set log level: debug|verbose|normal|silent");
|
||||
eprintln!(" --silent, -s Suppress info logs");
|
||||
eprintln!(" --log-level <LEVEL> debug|verbose|normal|silent");
|
||||
eprintln!(" --help, -h Show this help");
|
||||
std::process::exit(0);
|
||||
}
|
||||
s if !s.starts_with('-') => {
|
||||
config_path = s.to_string();
|
||||
}
|
||||
other => {
|
||||
eprintln!("Unknown option: {}", other);
|
||||
}
|
||||
s if !s.starts_with('-') => { config_path = s.to_string(); }
|
||||
other => { eprintln!("Unknown option: {}", other); }
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
@@ -76,20 +61,17 @@ fn parse_cli() -> (String, bool, Option<String>) {
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
// 1. Parse CLI arguments
|
||||
let (config_path, cli_silent, cli_log_level) = parse_cli();
|
||||
|
||||
// 2. Load config (tracing not yet initialized — errors go to stderr)
|
||||
let config = match ProxyConfig::load(&config_path) {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
if std::path::Path::new(&config_path).exists() {
|
||||
eprintln!("[telemt] Error: Failed to load config '{}': {}", config_path, e);
|
||||
eprintln!("[telemt] Error: {}", e);
|
||||
std::process::exit(1);
|
||||
} else {
|
||||
let default = ProxyConfig::default();
|
||||
let toml_str = toml::to_string_pretty(&default).unwrap();
|
||||
std::fs::write(&config_path, toml_str).unwrap();
|
||||
std::fs::write(&config_path, toml::to_string_pretty(&default).unwrap()).unwrap();
|
||||
eprintln!("[telemt] Created default config at {}", config_path);
|
||||
default
|
||||
}
|
||||
@@ -97,80 +79,90 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
};
|
||||
|
||||
if let Err(e) = config.validate() {
|
||||
eprintln!("[telemt] Error: Invalid configuration: {}", e);
|
||||
eprintln!("[telemt] Invalid config: {}", e);
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
// 3. Determine effective log level
|
||||
// Priority: RUST_LOG env > CLI flags > config file > default (normal)
|
||||
let effective_log_level = if cli_silent {
|
||||
LogLevel::Silent
|
||||
} else if let Some(ref level_str) = cli_log_level {
|
||||
LogLevel::from_str_loose(level_str)
|
||||
} else if let Some(ref s) = cli_log_level {
|
||||
LogLevel::from_str_loose(s)
|
||||
} else {
|
||||
config.general.log_level.clone()
|
||||
};
|
||||
|
||||
// 4. Initialize tracing
|
||||
let filter = if std::env::var("RUST_LOG").is_ok() {
|
||||
// RUST_LOG takes absolute priority
|
||||
EnvFilter::from_default_env()
|
||||
} else {
|
||||
EnvFilter::new(effective_log_level.to_filter_str())
|
||||
};
|
||||
|
||||
fmt()
|
||||
.with_env_filter(filter)
|
||||
.init();
|
||||
fmt().with_env_filter(filter).init();
|
||||
|
||||
// 5. Log startup info (operational — respects log level)
|
||||
info!("Telemt MTProxy v{}", env!("CARGO_PKG_VERSION"));
|
||||
info!("Log level: {}", effective_log_level);
|
||||
info!(
|
||||
"Modes: classic={} secure={} tls={}",
|
||||
info!("Modes: classic={} secure={} tls={}",
|
||||
config.general.modes.classic,
|
||||
config.general.modes.secure,
|
||||
config.general.modes.tls
|
||||
);
|
||||
config.general.modes.tls);
|
||||
info!("TLS domain: {}", config.censorship.tls_domain);
|
||||
info!(
|
||||
"Mask: {} -> {}:{}",
|
||||
info!("Mask: {} -> {}:{}",
|
||||
config.censorship.mask,
|
||||
config.censorship.mask_host.as_deref().unwrap_or(&config.censorship.tls_domain),
|
||||
config.censorship.mask_port
|
||||
);
|
||||
config.censorship.mask_port);
|
||||
|
||||
if config.censorship.tls_domain == "www.google.com" {
|
||||
warn!("Using default tls_domain (www.google.com). Consider setting a custom domain.");
|
||||
}
|
||||
|
||||
let prefer_ipv6 = config.general.prefer_ipv6;
|
||||
let config = Arc::new(config);
|
||||
let stats = Arc::new(Stats::new());
|
||||
let rng = Arc::new(SecureRandom::new());
|
||||
|
||||
// Initialize ReplayChecker
|
||||
let replay_checker = Arc::new(ReplayChecker::new(
|
||||
config.access.replay_check_len,
|
||||
Duration::from_secs(config.access.replay_window_secs),
|
||||
));
|
||||
|
||||
// Initialize Upstream Manager
|
||||
let upstream_manager = Arc::new(UpstreamManager::new(config.upstreams.clone()));
|
||||
|
||||
// Initialize Buffer Pool (16KB buffers, max 4096 cached ≈ 64MB)
|
||||
let buffer_pool = Arc::new(BufferPool::with_config(16 * 1024, 4096));
|
||||
|
||||
// Start health checks
|
||||
// === Startup DC Ping ===
|
||||
println!("=== Telegram DC Connectivity ===");
|
||||
let ping_results = upstream_manager.ping_all_dcs(prefer_ipv6).await;
|
||||
for upstream_result in &ping_results {
|
||||
println!(" via {}", upstream_result.upstream_name);
|
||||
for dc in &upstream_result.results {
|
||||
match (&dc.rtt_ms, &dc.error) {
|
||||
(Some(rtt), _) => {
|
||||
println!(" DC{} ({:>21}): {:.0}ms", dc.dc_idx, dc.dc_addr, rtt);
|
||||
}
|
||||
(None, Some(err)) => {
|
||||
println!(" DC{} ({:>21}): FAIL ({})", dc.dc_idx, dc.dc_addr, err);
|
||||
}
|
||||
(None, None) => {
|
||||
println!(" DC{} ({:>21}): FAIL", dc.dc_idx, dc.dc_addr);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
println!("================================");
|
||||
|
||||
// Start background tasks
|
||||
let um_clone = upstream_manager.clone();
|
||||
tokio::spawn(async move {
|
||||
um_clone.run_health_checks().await;
|
||||
um_clone.run_health_checks(prefer_ipv6).await;
|
||||
});
|
||||
|
||||
let rc_clone = replay_checker.clone();
|
||||
tokio::spawn(async move {
|
||||
rc_clone.run_periodic_cleanup().await;
|
||||
});
|
||||
|
||||
// Detect public IP (once at startup)
|
||||
let detected_ip = detect_ip().await;
|
||||
debug!("Detected IPs: v4={:?} v6={:?}", detected_ip.ipv4, detected_ip.ipv6);
|
||||
|
||||
// 6. Start listeners
|
||||
let mut listeners = Vec::new();
|
||||
|
||||
for listener_conf in &config.server.listeners {
|
||||
@@ -185,7 +177,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let listener = TcpListener::from_std(socket.into())?;
|
||||
info!("Listening on {}", addr);
|
||||
|
||||
// Determine public IP for tg:// links
|
||||
let public_ip = if let Some(ip) = listener_conf.announce_ip {
|
||||
ip
|
||||
} else if listener_conf.ip.is_unspecified() {
|
||||
@@ -198,30 +189,26 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
listener_conf.ip
|
||||
};
|
||||
|
||||
// 7. Print proxy links (always visible — uses println!, not tracing)
|
||||
if !config.show_link.is_empty() {
|
||||
println!("--- Proxy Links ({}) ---", public_ip);
|
||||
for user_name in &config.show_link {
|
||||
if let Some(secret) = config.access.users.get(user_name) {
|
||||
println!("[{}]", user_name);
|
||||
|
||||
if config.general.modes.classic {
|
||||
println!(" Classic: tg://proxy?server={}&port={}&secret={}",
|
||||
println!(" Classic: tg://proxy?server={}&port={}&secret={}",
|
||||
public_ip, config.server.port, secret);
|
||||
}
|
||||
|
||||
if config.general.modes.secure {
|
||||
println!(" DD: tg://proxy?server={}&port={}&secret=dd{}",
|
||||
println!(" DD: tg://proxy?server={}&port={}&secret=dd{}",
|
||||
public_ip, config.server.port, secret);
|
||||
}
|
||||
|
||||
if config.general.modes.tls {
|
||||
let domain_hex = hex::encode(&config.censorship.tls_domain);
|
||||
println!(" EE-TLS: tg://proxy?server={}&port={}&secret=ee{}{}",
|
||||
println!(" EE-TLS: tg://proxy?server={}&port={}&secret=ee{}{}",
|
||||
public_ip, config.server.port, secret, domain_hex);
|
||||
}
|
||||
} else {
|
||||
warn!("User '{}' in show_link not found in users", user_name);
|
||||
warn!("User '{}' in show_link not found", user_name);
|
||||
}
|
||||
}
|
||||
println!("------------------------");
|
||||
@@ -236,11 +223,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
}
|
||||
|
||||
if listeners.is_empty() {
|
||||
error!("No listeners could be started. Exiting.");
|
||||
error!("No listeners. Exiting.");
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
// 8. Accept loop
|
||||
for listener in listeners {
|
||||
let config = config.clone();
|
||||
let stats = stats.clone();
|
||||
@@ -262,14 +248,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = ClientHandler::new(
|
||||
stream,
|
||||
peer_addr,
|
||||
config,
|
||||
stats,
|
||||
upstream_manager,
|
||||
replay_checker,
|
||||
buffer_pool,
|
||||
rng
|
||||
stream, peer_addr, config, stats,
|
||||
upstream_manager, replay_checker, buffer_pool, rng
|
||||
).run().await {
|
||||
debug!(peer = %peer_addr, error = %e, "Connection error");
|
||||
}
|
||||
@@ -284,7 +264,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
});
|
||||
}
|
||||
|
||||
// 9. Wait for shutdown signal
|
||||
match signal::ctrl_c().await {
|
||||
Ok(()) => info!("Shutting down..."),
|
||||
Err(e) => error!("Signal error: {}", e),
|
||||
|
||||
Reference in New Issue
Block a user