diff --git a/src/api/mod.rs b/src/api/mod.rs index 1677a0c..968ea2d 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -20,9 +20,9 @@ use tokio::sync::{Mutex, RwLock, Semaphore, watch}; use tokio::time::timeout; use tracing::{debug, info, warn}; -use crate::config::{ApiGrayAction, ProxyConfig}; +use crate::config::ApiGrayAction; use crate::ip_tracker::UserIpTracker; -use crate::maestro::generation::RuntimeGeneration; +use crate::maestro::generation::{RuntimeGeneration, RuntimeWatchState}; use crate::maestro::reload::{ReloadControl, ReloadRequest, ReloadSubmitError}; use crate::proxy::route_mode::RouteRuntimeController; use crate::proxy::shared_state::ProxySharedState; @@ -240,8 +240,6 @@ pub async fn serve( route_runtime: Arc, proxy_shared: Arc, upstream_manager: Arc, - config_rx: watch::Receiver>, - admission_rx: watch::Receiver, config_path: PathBuf, quota_state_path: PathBuf, detected_ips_rx: watch::Receiver<(Option, Option)>, @@ -249,6 +247,7 @@ pub async fn serve( startup_tracker: Arc, reload_control: ReloadControl, mut active_runtime_rx: watch::Receiver>>>, + mut runtime_watch_rx: watch::Receiver>, ) { let active_runtime = loop { if let Some(active_runtime) = active_runtime_rx.borrow().clone() { @@ -259,6 +258,17 @@ pub async fn serve( return; } }; + let initial_watch_state = loop { + if let Some(watch_state) = runtime_watch_rx.borrow().clone() { + break watch_state; + } + if runtime_watch_rx.changed().await.is_err() { + warn!("Runtime watch channel closed before API bootstrap"); + return; + } + }; + let config_rx = initial_watch_state.config_rx.clone(); + let admission_rx = initial_watch_state.admission_rx.clone(); let listener = match TcpListener::bind(listen).await { Ok(listener) => listener, Err(error) => { @@ -306,8 +316,7 @@ pub async fn serve( }); spawn_runtime_watchers( - config_rx.clone(), - admission_rx.clone(), + runtime_watch_rx, runtime_state.clone(), shared.runtime_events.clone(), ); diff --git a/src/api/runtime_watch.rs b/src/api/runtime_watch.rs index 0485e55..71a352f 100644 --- a/src/api/runtime_watch.rs +++ b/src/api/runtime_watch.rs @@ -4,63 +4,283 @@ use std::time::{SystemTime, UNIX_EPOCH}; use tokio::sync::watch; -use crate::config::ProxyConfig; +use crate::maestro::generation::RuntimeWatchState; use super::ApiRuntimeState; use super::events::ApiEventStore; pub(super) fn spawn_runtime_watchers( - config_rx: watch::Receiver>, - admission_rx: watch::Receiver, + runtime_watch_rx: watch::Receiver>, runtime_state: Arc, runtime_events: Arc, ) { - let mut config_rx_reload = config_rx; - let runtime_state_reload = runtime_state.clone(); - let runtime_events_reload = runtime_events.clone(); - tokio::spawn(async move { - loop { - if config_rx_reload.changed().await.is_err() { - break; - } - runtime_state_reload - .config_reload_count - .fetch_add(1, Ordering::Relaxed); - runtime_state_reload - .last_config_reload_epoch_secs - .store(now_epoch_secs(), Ordering::Relaxed); - runtime_events_reload.record("config.reload.applied", "config receiver updated"); - } - }); + spawn_config_watcher( + runtime_watch_rx.clone(), + runtime_state.clone(), + runtime_events.clone(), + ); + spawn_admission_watcher(runtime_watch_rx, runtime_state, runtime_events); +} - let mut admission_rx_watch = admission_rx; +fn spawn_config_watcher( + mut runtime_watch_rx: watch::Receiver>, + runtime_state: Arc, + runtime_events: Arc, +) { tokio::spawn(async move { - runtime_state - .admission_open - .store(*admission_rx_watch.borrow(), Ordering::Relaxed); - runtime_events.record( - "admission.state", - format!("accepting_new_connections={}", *admission_rx_watch.borrow()), - ); + let Some(mut current) = runtime_watch_rx.borrow().clone() else { + return; + }; loop { - if admission_rx_watch.changed().await.is_err() { - break; + tokio::select! { + biased; + changed = runtime_watch_rx.changed() => { + if changed.is_err() { + break; + } + let Some(next) = runtime_watch_rx.borrow().clone() else { + continue; + }; + if next.generation_id != current.generation_id { + current = next; + record_config_reload( + &runtime_state, + &runtime_events, + format!("runtime generation {} activated", current.generation_id), + ); + } + } + changed = current.config_rx.changed() => { + if changed.is_err() { + let Some(next) = wait_for_new_generation( + &mut runtime_watch_rx, + current.generation_id, + ).await else { + break; + }; + current = next; + record_config_reload( + &runtime_state, + &runtime_events, + format!("runtime generation {} activated", current.generation_id), + ); + continue; + } + if active_generation_id(&runtime_watch_rx) != Some(current.generation_id) { + continue; + } + record_config_reload( + &runtime_state, + &runtime_events, + format!("generation {} config receiver updated", current.generation_id), + ); + } } - let admission_open = *admission_rx_watch.borrow(); - runtime_state - .admission_open - .store(admission_open, Ordering::Relaxed); - runtime_events.record( - "admission.state", - format!("accepting_new_connections={}", admission_open), - ); } }); } +fn spawn_admission_watcher( + mut runtime_watch_rx: watch::Receiver>, + runtime_state: Arc, + runtime_events: Arc, +) { + tokio::spawn(async move { + let Some(mut current) = runtime_watch_rx.borrow().clone() else { + return; + }; + record_admission_state(&runtime_state, &runtime_events, ¤t); + loop { + tokio::select! { + biased; + changed = runtime_watch_rx.changed() => { + if changed.is_err() { + break; + } + let Some(next) = runtime_watch_rx.borrow().clone() else { + continue; + }; + if next.generation_id != current.generation_id { + current = next; + record_admission_state(&runtime_state, &runtime_events, ¤t); + } + } + changed = current.admission_rx.changed() => { + if changed.is_err() { + let Some(next) = wait_for_new_generation( + &mut runtime_watch_rx, + current.generation_id, + ).await else { + break; + }; + current = next; + record_admission_state(&runtime_state, &runtime_events, ¤t); + continue; + } + if active_generation_id(&runtime_watch_rx) == Some(current.generation_id) { + record_admission_state(&runtime_state, &runtime_events, ¤t); + } + } + } + } + }); +} + +fn active_generation_id( + runtime_watch_rx: &watch::Receiver>, +) -> Option { + runtime_watch_rx + .borrow() + .as_ref() + .map(|state| state.generation_id) +} + +async fn wait_for_new_generation( + runtime_watch_rx: &mut watch::Receiver>, + previous_generation_id: u64, +) -> Option { + loop { + if let Some(state) = runtime_watch_rx.borrow().clone() + && state.generation_id != previous_generation_id + { + return Some(state); + } + if runtime_watch_rx.changed().await.is_err() { + return None; + } + } +} + +fn record_config_reload( + runtime_state: &ApiRuntimeState, + runtime_events: &ApiEventStore, + context: String, +) { + runtime_state + .config_reload_count + .fetch_add(1, Ordering::Relaxed); + runtime_state + .last_config_reload_epoch_secs + .store(now_epoch_secs(), Ordering::Relaxed); + runtime_events.record("config.reload.applied", context); +} + +fn record_admission_state( + runtime_state: &ApiRuntimeState, + runtime_events: &ApiEventStore, + current: &RuntimeWatchState, +) { + let admission_open = *current.admission_rx.borrow(); + runtime_state + .admission_open + .store(admission_open, Ordering::Relaxed); + runtime_events.record( + "admission.state", + format!( + "generation={} accepting_new_connections={}", + current.generation_id, admission_open + ), + ); +} + fn now_epoch_secs() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_secs() } + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::ProxyConfig; + use std::sync::atomic::{AtomicBool, AtomicU64}; + use std::time::Duration; + + fn state( + generation_id: u64, + ) -> ( + RuntimeWatchState, + watch::Sender>, + watch::Sender, + ) { + let (config_tx, config_rx) = watch::channel(Arc::new(ProxyConfig::default())); + let (admission_tx, admission_rx) = watch::channel(true); + ( + RuntimeWatchState { + generation_id, + config_rx, + admission_rx, + }, + config_tx, + admission_tx, + ) + } + + fn runtime_state() -> Arc { + Arc::new(ApiRuntimeState { + process_started_at_epoch_secs: 1, + config_reload_count: AtomicU64::new(0), + last_config_reload_epoch_secs: AtomicU64::new(0), + admission_open: AtomicBool::new(false), + }) + } + + async fn wait_for_count(runtime_state: &ApiRuntimeState, expected: u64) { + tokio::time::timeout(Duration::from_secs(1), async { + loop { + if runtime_state.config_reload_count.load(Ordering::Relaxed) == expected { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } + + #[tokio::test] + async fn watchers_follow_only_the_active_generation() { + let (initial, initial_config_tx, initial_admission_tx) = state(1); + let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(initial)); + let runtime_state = runtime_state(); + let events = Arc::new(ApiEventStore::new(16)); + spawn_runtime_watchers(runtime_watch_rx, runtime_state.clone(), events.clone()); + tokio::task::yield_now().await; + + assert_eq!(runtime_state.config_reload_count.load(Ordering::Relaxed), 0); + initial_config_tx.send_replace(Arc::new(ProxyConfig::default())); + wait_for_count(&runtime_state, 1).await; + + let (next, next_config_tx, next_admission_tx) = state(2); + runtime_watch_tx.send_replace(Some(next)); + wait_for_count(&runtime_state, 2).await; + + initial_config_tx.send_replace(Arc::new(ProxyConfig::default())); + initial_admission_tx.send_replace(false); + tokio::task::yield_now().await; + assert_eq!(runtime_state.config_reload_count.load(Ordering::Relaxed), 2); + assert!(runtime_state.admission_open.load(Ordering::Relaxed)); + + next_config_tx.send_replace(Arc::new(ProxyConfig::default())); + next_admission_tx.send_replace(false); + wait_for_count(&runtime_state, 3).await; + tokio::time::timeout(Duration::from_secs(1), async { + while runtime_state.admission_open.load(Ordering::Relaxed) { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + let snapshot = events.snapshot(16); + assert_eq!( + snapshot + .events + .iter() + .filter(|event| event.event_type == "config.reload.applied") + .count(), + 3 + ); + } +} diff --git a/src/api/runtime_zero.rs b/src/api/runtime_zero.rs index b17615f..e9dce7f 100644 --- a/src/api/runtime_zero.rs +++ b/src/api/runtime_zero.rs @@ -1,4 +1,5 @@ use std::sync::atomic::Ordering; +use std::time::{SystemTime, UNIX_EPOCH}; use serde::Serialize; @@ -162,7 +163,7 @@ pub(super) fn build_system_info_data( build_time_utc, rustc_version, process_started_at_epoch_secs: shared.runtime_state.process_started_at_epoch_secs, - uptime_seconds: shared.stats.uptime_secs(), + uptime_seconds: process_uptime_seconds(shared.runtime_state.process_started_at_epoch_secs), config_path: shared.config_path.display().to_string(), config_hash: revision.to_string(), config_reload_count: shared @@ -173,6 +174,18 @@ pub(super) fn build_system_info_data( } } +fn process_uptime_seconds(process_started_at_epoch_secs: u64) -> f64 { + let now_epoch_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + process_uptime_seconds_at(process_started_at_epoch_secs, now_epoch_secs) +} + +fn process_uptime_seconds_at(process_started_at_epoch_secs: u64, now_epoch_secs: u64) -> f64 { + now_epoch_secs.saturating_sub(process_started_at_epoch_secs) as f64 +} + pub(super) async fn build_runtime_gates_data( shared: &ApiShared, cfg: &ProxyConfig, @@ -339,3 +352,14 @@ fn me_writer_pick_mode_label(mode: MeWriterPickMode) -> &'static str { MeWriterPickMode::P2c => "p2c", } } + +#[cfg(test)] +mod tests { + use super::process_uptime_seconds_at; + + #[test] + fn process_uptime_is_monotonic_and_saturating() { + assert_eq!(process_uptime_seconds_at(100, 135), 35.0); + assert_eq!(process_uptime_seconds_at(135, 100), 0.0); + } +} diff --git a/src/maestro/generation.rs b/src/maestro/generation.rs index f4f52de..365a5bb 100644 --- a/src/maestro/generation.rs +++ b/src/maestro/generation.rs @@ -22,6 +22,14 @@ use crate::transport::middle_proxy::MePool; const SESSION_STOP_TIMEOUT: Duration = Duration::from_secs(5); const BACKGROUND_STOP_TIMEOUT: Duration = Duration::from_secs(5); +/// Process-visible control-plane receivers for one active runtime generation. +#[derive(Clone)] +pub(crate) struct RuntimeWatchState { + pub(crate) generation_id: u64, + pub(crate) config_rx: watch::Receiver>, + pub(crate) admission_rx: watch::Receiver, +} + /// Cancellation and join ownership for one generation's background tasks. #[derive(Clone)] pub(crate) struct RuntimeTaskScope { @@ -134,6 +142,15 @@ impl RuntimeGeneration { self.config_rx.borrow().clone() } + /// Returns receivers used by process-scoped observers of this generation. + pub(crate) fn watch_state(&self) -> RuntimeWatchState { + RuntimeWatchState { + generation_id: self.id, + config_rx: self.config_rx.clone(), + admission_rx: self.admission_rx.clone(), + } + } + pub(crate) async fn current_me_pool(&self) -> Option> { if let Some(pool) = &self.me_pool { return Some(pool.clone()); diff --git a/src/maestro/me_startup.rs b/src/maestro/me_startup.rs index 65b37d1..024230c 100644 --- a/src/maestro/me_startup.rs +++ b/src/maestro/me_startup.rs @@ -1,9 +1,11 @@ #![allow(clippy::too_many_arguments)] +use std::future::Future; use std::sync::Arc; use std::time::Duration; use tokio::sync::{RwLock, watch}; +use tokio_util::task::AbortOnDropHandle; use tracing::{error, info, warn}; use crate::config::ProxyConfig; @@ -17,8 +19,58 @@ use crate::stats::Stats; use crate::transport::UpstreamManager; use crate::transport::middle_proxy::MePool; +use super::generation::RuntimeTaskScope; use super::helpers::load_startup_proxy_config_snapshot; +async fn supervise_me_task(task_name: &'static str, mut task: F) +where + F: FnMut() -> Fut, + Fut: Future + Send + 'static, +{ + loop { + let result = AbortOnDropHandle::new(tokio::spawn(task())).await; + match result { + Ok(()) => warn!(task = task_name, "Middle-End supervisor task exited unexpectedly, restarting"), + Err(error) => { + error!(task = task_name, error = %error, "Middle-End supervisor task panicked, restarting in 1s"); + tokio::time::sleep(Duration::from_secs(1)).await; + } + } + } +} + +fn spawn_me_supervisors( + task_scope: RuntimeTaskScope, + pool: Arc, + rng: Arc, + min_connections: usize, +) { + let health_pool = pool.clone(); + let health_rng = rng; + task_scope.spawn(supervise_me_task("health_monitor", move || { + let pool = health_pool.clone(); + let rng = health_rng.clone(); + async move { + crate::transport::middle_proxy::me_health_monitor(pool, rng, min_connections).await; + } + })); + + let drain_pool = pool.clone(); + task_scope.spawn(supervise_me_task("drain_timeout_enforcer", move || { + let pool = drain_pool.clone(); + async move { + crate::transport::middle_proxy::me_drain_timeout_enforcer(pool).await; + } + })); + + task_scope.spawn(supervise_me_task("zombie_writer_watchdog", move || { + let pool = pool.clone(); + async move { + crate::transport::middle_proxy::me_zombie_writer_watchdog(pool).await; + } + })); +} + pub(crate) async fn initialize_me_pool( use_middle_proxy: bool, config: &ProxyConfig, @@ -30,6 +82,7 @@ pub(crate) async fn initialize_me_pool( stats: Arc, api_me_pool: Arc>>>, me_ready_tx: watch::Sender, + task_scope: RuntimeTaskScope, ) -> Option> { if !use_middle_proxy { return None; @@ -52,15 +105,8 @@ pub(crate) async fn initialize_me_pool( .as_ref() .map(|tag| hex::decode(tag).expect("general.ad_tag must be validated before startup")); - // ============================================================= - // CRITICAL: Download Telegram proxy-secret (NOT user secret!) - // - // C MTProxy uses TWO separate secrets: - // -S flag = 16-byte user secret for client obfuscation - // --aes-pwd = 32-512 byte binary file for ME RPC auth - // - // proxy-secret is from: https://core.telegram.org/getProxySecret - // ============================================================= + // The Telegram proxy-secret authenticates ME RPC and is distinct from client secrets. + // It corresponds to the C MTProxy --aes-pwd input and may be fetched from Telegram. let proxy_secret_path = config.general.proxy_secret_path.as_deref(); let pool_size = config.general.middle_proxy_pool_size.max(1); let proxy_secret = loop { @@ -319,143 +365,70 @@ pub(crate) async fn initialize_me_pool( let rng_bg = rng.clone(); let startup_tracker_bg = startup_tracker.clone(); let me_ready_tx_bg = me_ready_tx.clone(); + let task_scope_bg = task_scope.clone(); let retry_limit = if me_init_retry_attempts == 0 { String::from("unlimited") } else { me_init_retry_attempts.to_string() }; - std::thread::spawn(move || { - let runtime = match tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - { - Ok(runtime) => runtime, - Err(error) => { - error!(error = %error, "Failed to build background runtime for ME initialization"); - return; - } - }; - runtime.block_on(async move { - let mut init_attempt: u32 = 0; - loop { - init_attempt = init_attempt.saturating_add(1); - startup_tracker_bg.set_me_init_attempt(init_attempt).await; - match pool_bg.init(pool_size, &rng_bg).await { - Ok(()) => { - startup_tracker_bg.set_me_last_error(None).await; - startup_tracker_bg - .complete_component( - COMPONENT_ME_POOL_INIT_STAGE1, - Some("ME pool initialized".to_string()), - ) - .await; - startup_tracker_bg - .set_me_status(StartupMeStatus::Ready, "ready") - .await; - me_ready_tx_bg.send_modify(|version| { - *version = version.saturating_add(1); - }); - info!( + task_scope.spawn(async move { + let mut init_attempt: u32 = 0; + loop { + init_attempt = init_attempt.saturating_add(1); + startup_tracker_bg.set_me_init_attempt(init_attempt).await; + match pool_bg.init(pool_size, &rng_bg).await { + Ok(()) => { + startup_tracker_bg.set_me_last_error(None).await; + startup_tracker_bg + .complete_component( + COMPONENT_ME_POOL_INIT_STAGE1, + Some("ME pool initialized".to_string()), + ) + .await; + startup_tracker_bg + .set_me_status(StartupMeStatus::Ready, "ready") + .await; + me_ready_tx_bg.send_modify(|version| { + *version = version.saturating_add(1); + }); + info!( + attempt = init_attempt, + "Middle-End pool initialized successfully" + ); + spawn_me_supervisors( + task_scope_bg, + pool_bg.clone(), + rng_bg.clone(), + pool_size, + ); + break; + } + Err(e) => { + startup_tracker_bg + .set_me_last_error(Some(e.to_string())) + .await; + if init_attempt >= me_init_warn_after_attempts { + warn!( + error = %e, attempt = init_attempt, - "Middle-End pool initialized successfully" + retry_limit = %retry_limit, + retry_in_secs = 2, + "ME pool is not ready yet; retrying background initialization" + ); + } else { + info!( + error = %e, + attempt = init_attempt, + retry_limit = %retry_limit, + retry_in_secs = 2, + "ME pool startup warmup: retrying background initialization" ); - - // ── Supervised background tasks ────────────────── - // Each task runs inside a nested tokio::spawn so - // that a panic is caught via JoinHandle and the - // outer loop restarts the task automatically. - let pool_health = pool_bg.clone(); - let rng_health = rng_bg.clone(); - let min_conns = pool_size; - tokio::spawn(async move { - loop { - let p = pool_health.clone(); - let r = rng_health.clone(); - let res = tokio::spawn(async move { - crate::transport::middle_proxy::me_health_monitor( - p, r, min_conns, - ) - .await; - }) - .await; - match res { - Ok(()) => warn!("me_health_monitor exited unexpectedly, restarting"), - Err(e) => { - error!(error = %e, "me_health_monitor panicked, restarting in 1s"); - tokio::time::sleep(Duration::from_secs(1)).await; - } - } - } - }); - let pool_drain_enforcer = pool_bg.clone(); - tokio::spawn(async move { - loop { - let p = pool_drain_enforcer.clone(); - let res = tokio::spawn(async move { - crate::transport::middle_proxy::me_drain_timeout_enforcer(p).await; - }) - .await; - match res { - Ok(()) => warn!("me_drain_timeout_enforcer exited unexpectedly, restarting"), - Err(e) => { - error!(error = %e, "me_drain_timeout_enforcer panicked, restarting in 1s"); - tokio::time::sleep(Duration::from_secs(1)).await; - } - } - } - }); - let pool_watchdog = pool_bg.clone(); - tokio::spawn(async move { - loop { - let p = pool_watchdog.clone(); - let res = tokio::spawn(async move { - crate::transport::middle_proxy::me_zombie_writer_watchdog(p).await; - }) - .await; - match res { - Ok(()) => warn!("me_zombie_writer_watchdog exited unexpectedly, restarting"), - Err(e) => { - error!(error = %e, "me_zombie_writer_watchdog panicked, restarting in 1s"); - tokio::time::sleep(Duration::from_secs(1)).await; - } - } - } - }); - // CRITICAL: keep the current-thread runtime - // alive. Without this, block_on() returns, - // the Runtime is dropped, and ALL spawned - // background tasks (health monitor, drain - // enforcer, zombie watchdog) are silently - // cancelled — causing the draining-writer - // leak that brought us here. - std::future::pending::<()>().await; - unreachable!(); - } - Err(e) => { - startup_tracker_bg.set_me_last_error(Some(e.to_string())).await; - if init_attempt >= me_init_warn_after_attempts { - warn!( - error = %e, - attempt = init_attempt, - retry_limit = %retry_limit, - retry_in_secs = 2, - "ME pool is not ready yet; retrying background initialization" - ); - } else { - info!( - error = %e, - attempt = init_attempt, - retry_limit = %retry_limit, - retry_in_secs = 2, - "ME pool startup warmup: retrying background initialization" - ); - } - pool_bg.reset_stun_state(); - tokio::time::sleep(Duration::from_secs(2)).await; } + pool_bg.reset_stun_state(); + tokio::time::sleep(Duration::from_secs(2)).await; } } - }); + } }); startup_tracker .set_me_status(StartupMeStatus::Initializing, "background_init") @@ -490,70 +463,12 @@ pub(crate) async fn initialize_me_pool( "Middle-End pool initialized successfully" ); - // ── Supervised background tasks ────────────────── - let pool_clone = pool.clone(); - let rng_clone = rng.clone(); - let min_conns = pool_size; - tokio::spawn(async move { - loop { - let p = pool_clone.clone(); - let r = rng_clone.clone(); - let res = tokio::spawn(async move { - crate::transport::middle_proxy::me_health_monitor( - p, r, min_conns, - ) - .await; - }) - .await; - match res { - Ok(()) => warn!( - "me_health_monitor exited unexpectedly, restarting" - ), - Err(e) => { - error!(error = %e, "me_health_monitor panicked, restarting in 1s"); - tokio::time::sleep(Duration::from_secs(1)).await; - } - } - } - }); - let pool_drain_enforcer = pool.clone(); - tokio::spawn(async move { - loop { - let p = pool_drain_enforcer.clone(); - let res = tokio::spawn(async move { - crate::transport::middle_proxy::me_drain_timeout_enforcer(p).await; - }) - .await; - match res { - Ok(()) => warn!( - "me_drain_timeout_enforcer exited unexpectedly, restarting" - ), - Err(e) => { - error!(error = %e, "me_drain_timeout_enforcer panicked, restarting in 1s"); - tokio::time::sleep(Duration::from_secs(1)).await; - } - } - } - }); - let pool_watchdog = pool.clone(); - tokio::spawn(async move { - loop { - let p = pool_watchdog.clone(); - let res = tokio::spawn(async move { - crate::transport::middle_proxy::me_zombie_writer_watchdog(p).await; - }) - .await; - match res { - Ok(()) => warn!( - "me_zombie_writer_watchdog exited unexpectedly, restarting" - ), - Err(e) => { - error!(error = %e, "me_zombie_writer_watchdog panicked, restarting in 1s"); - tokio::time::sleep(Duration::from_secs(1)).await; - } - } - } - }); + spawn_me_supervisors( + task_scope.clone(), + pool.clone(), + rng.clone(), + pool_size, + ); break Some(pool); } @@ -666,3 +581,38 @@ pub(crate) async fn initialize_me_pool( } } } + +#[cfg(test)] +mod tests { + use super::*; + use tokio::sync::Notify; + + struct DropSignal(Arc); + + impl Drop for DropSignal { + fn drop(&mut self) { + self.0.notify_one(); + } + } + + #[tokio::test] + async fn scoped_supervisor_aborts_its_current_child() { + let scope = RuntimeTaskScope::new(); + let dropped = Arc::new(Notify::new()); + let dropped_for_task = dropped.clone(); + scope.spawn(supervise_me_task("test", move || { + let dropped = dropped_for_task.clone(); + async move { + let _signal = DropSignal(dropped); + std::future::pending::<()>().await; + } + })); + tokio::task::yield_now().await; + + scope.stop().await; + + tokio::time::timeout(Duration::from_secs(1), dropped.notified()) + .await + .unwrap(); + } +} diff --git a/src/maestro/mod.rs b/src/maestro/mod.rs index dd2cadf..6b2af2a 100644 --- a/src/maestro/mod.rs +++ b/src/maestro/mod.rs @@ -503,7 +503,6 @@ async fn run_telemt_core( config.access.cidr_rate_limits.clone(), ); - let (api_config_tx, api_config_rx) = watch::channel(Arc::new(config.clone())); let (detected_ips_tx, detected_ips_rx) = watch::channel((None::, None::)); let initial_direct_first = config.general.use_middle_proxy && config.general.me2dc_fallback; let initial_admission_open = !config.general.use_middle_proxy || initial_direct_first; @@ -511,6 +510,8 @@ async fn run_telemt_core( let (reload_control, reload_commands) = reload::ReloadControl::channel(1); let (active_runtime_tx, active_runtime_rx) = watch::channel(None::>>); + let (runtime_watch_tx, runtime_watch_rx) = + watch::channel(None::); let initial_route_mode = if !config.general.use_middle_proxy || initial_direct_first { RelayRouteMode::Direct } else { @@ -544,14 +545,13 @@ async fn run_telemt_core( let upstream_manager_api = upstream_manager.clone(); let route_runtime_api = route_runtime.clone(); let proxy_shared_api = shared_state.clone(); - let config_rx_api = api_config_rx.clone(); - let admission_rx_api = admission_rx.clone(); let config_path_api = config_path.clone(); let quota_state_path_api = quota_state_path.clone(); let startup_tracker_api = startup_tracker.clone(); let detected_ips_rx_api = detected_ips_rx.clone(); let reload_control_api = reload_control.clone(); let active_runtime_rx_api = active_runtime_rx.clone(); + let runtime_watch_rx_api = runtime_watch_rx.clone(); tokio::spawn(async move { api::serve( listen, @@ -561,8 +561,6 @@ async fn run_telemt_core( route_runtime_api, proxy_shared_api, upstream_manager_api, - config_rx_api, - admission_rx_api, config_path_api, quota_state_path_api, detected_ips_rx_api, @@ -570,6 +568,7 @@ async fn run_telemt_core( startup_tracker_api, reload_control_api, active_runtime_rx_api, + runtime_watch_rx_api, ) .await; }); @@ -610,8 +609,9 @@ async fn run_telemt_core( upstream_manager.clone(), &startup_tracker, runtime_task_scope.clone(), + tls_bootstrap::TlsBootstrapPolicy::BestEffort, ) - .await; + .await?; startup_tracker .start_component( @@ -737,6 +737,7 @@ async fn run_telemt_core( stats.clone(), api_me_pool.clone(), me_ready_tx.clone(), + runtime_task_scope.clone(), ) .await }; @@ -824,7 +825,6 @@ async fn run_telemt_core( rng.clone(), ip_tracker.clone(), beobachten.clone(), - api_config_tx.clone(), me_pool.clone(), shared_state.clone(), me_ready_tx.clone(), @@ -869,6 +869,7 @@ async fn run_telemt_core( stats_bg.clone(), api_me_pool_bg.clone(), me_ready_tx_bg.clone(), + task_scope_bg.clone(), ) .await; if let Some(pool) = pool { @@ -961,6 +962,7 @@ async fn run_telemt_core( quota_store, detected_ips_tx, runtime_log_filter, + runtime_watch_tx.clone(), ); let bound = listeners::bind_listeners( @@ -988,12 +990,12 @@ async fn run_telemt_core( // On Unix, caller supplies privilege drop after bind (may require root for port < 1024). drop_after_bind(); - synlimit_control::reconcile_synlimit_rules(&config).await; - synlimit_control::spawn_synlimit_controller(config_rx.clone()); + let synlimit_controller = synlimit_control::spawn_synlimit_controller(runtime_watch_rx); runtime_tasks::spawn_metrics_if_configured(&config, &startup_tracker, active_runtime.clone()) .await; + runtime_watch_tx.send_replace(Some(active_runtime.load_full().watch_state())); active_runtime_tx.send_replace(Some(active_runtime.clone())); runtime_tasks::mark_runtime_ready(&startup_tracker).await; @@ -1004,7 +1006,13 @@ async fn run_telemt_core( #[cfg(unix)] listeners::spawn_unix_accept_loop(unix_listener, active_runtime.clone()); - shutdown::wait_for_shutdown(process_started_at, active_runtime, quota_state_path).await; + shutdown::wait_for_shutdown( + process_started_at, + active_runtime, + quota_state_path, + synlimit_controller, + ) + .await; Ok(()) } diff --git a/src/maestro/reload.rs b/src/maestro/reload.rs index b2b3bbc..0a7ab58 100644 --- a/src/maestro/reload.rs +++ b/src/maestro/reload.rs @@ -138,7 +138,11 @@ pub(crate) struct ReloadStatus { pub(crate) started_at_epoch_secs: Option, #[serde(skip_serializing_if = "Option::is_none")] pub(crate) finished_at_epoch_secs: Option, - #[serde(default, skip_serializing_if = "Vec::is_empty")] + #[serde( + rename = "deferred_process_fields", + default, + skip_serializing_if = "Vec::is_empty" + )] pub(crate) deferred_fields: Vec, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub(crate) warnings: Vec, @@ -428,6 +432,31 @@ mod tests { assert!(ReloadRequest::from_query(Some("timeout_secs=30")).is_err()); } + #[test] + fn status_uses_documented_deferred_process_fields_key() { + let status = ReloadStatus { + reload_id: 1, + target_generation: 2, + config_revision: "revision".to_string(), + state: ReloadPhase::Succeeded, + mode: ReloadMode::Instant, + failure_policy: ReloadFailurePolicy::KeepNew, + requested_at_epoch_secs: 10, + started_at_epoch_secs: Some(11), + finished_at_epoch_secs: Some(12), + deferred_fields: vec!["server.listeners".to_string()], + warnings: Vec::new(), + error: None, + }; + let value = serde_json::to_value(status).unwrap(); + + assert_eq!( + value["deferred_process_fields"], + serde_json::json!(["server.listeners"]) + ); + assert!(value.get("deferred_fields").is_none()); + } + #[tokio::test] async fn coordinator_rejects_concurrent_reload_and_releases_terminal_slot() { let (control, mut receiver) = ReloadControl::channel(1); diff --git a/src/maestro/reload_supervisor.rs b/src/maestro/reload_supervisor.rs index aa88cce..f71f58a 100644 --- a/src/maestro/reload_supervisor.rs +++ b/src/maestro/reload_supervisor.rs @@ -8,7 +8,7 @@ use tracing::{info, warn}; use crate::stats::QuotaStore; -use super::generation::RuntimeGeneration; +use super::generation::{RuntimeGeneration, RuntimeWatchState}; use super::reload::{ ReloadCommand, ReloadCommandReceiver, ReloadControl, ReloadFailurePolicy, ReloadMode, ReloadPhase, @@ -24,6 +24,48 @@ pub(crate) struct ReloadSupervisor { quota_store: Arc, detected_ips_tx: watch::Sender<(Option, Option)>, runtime_log_filter: RuntimeLogFilter, + runtime_watch_tx: watch::Sender>, +} + +#[derive(Debug, PartialEq, Eq)] +enum RevisionGateAction { + Proceed, + Warn(String), + Rollback(String), +} + +fn revision_gate_action( + accepted_revision: &str, + current_revision: Result, + failure_policy: ReloadFailurePolicy, +) -> RevisionGateAction { + let warning = match current_revision { + Ok(current) if current == accepted_revision => return RevisionGateAction::Proceed, + Ok(current) => format!( + "config revision changed during preparation: accepted={} current={}", + accepted_revision, current + ), + Err(error) => format!("config revision verification failed: {}", error), + }; + match failure_policy { + ReloadFailurePolicy::KeepNew => RevisionGateAction::Warn(warning), + ReloadFailurePolicy::Rollback => RevisionGateAction::Rollback(warning), + } +} + +async fn stop_background_and_middle_end(generation: &RuntimeGeneration) -> bool { + generation.stop_background_tasks().await; + let Some(pool) = generation.current_me_pool().await else { + return false; + }; + tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all()) + .await + .is_err() +} + +async fn cleanup_candidate(generation: &RuntimeGeneration) -> bool { + generation.stop_sessions().await; + stop_background_and_middle_end(generation).await } impl ReloadSupervisor { @@ -36,6 +78,7 @@ impl ReloadSupervisor { quota_store: Arc, detected_ips_tx: watch::Sender<(Option, Option)>, runtime_log_filter: RuntimeLogFilter, + runtime_watch_tx: watch::Sender>, ) { let supervisor = Self { active_runtime, @@ -45,6 +88,7 @@ impl ReloadSupervisor { quota_store, detected_ips_tx, runtime_log_filter, + runtime_watch_tx, }; tokio::spawn(supervisor.run()); } @@ -81,30 +125,23 @@ impl ReloadSupervisor { } }; - if let Ok(current_revision) = - crate::api::config_store::current_revision_for_maestro(&self.config_path).await - && current_revision != command.config_revision - { - let warning = format!( - "config revision changed during preparation: accepted={} current={}", - command.config_revision, current_revision - ); - if command.request.failure_policy == ReloadFailurePolicy::Rollback { + let revision_action = revision_gate_action( + &command.config_revision, + crate::api::config_store::current_revision_for_maestro(&self.config_path).await, + command.request.failure_policy, + ); + match revision_action { + RevisionGateAction::Proceed => {} + RevisionGateAction::Warn(warning) => { + self.control.add_warning(command.reload_id, warning).await; + } + RevisionGateAction::Rollback(warning) => { + let _ = cleanup_candidate(&prepared.generation).await; self.runtime_log_filter .apply_reload(&old_runtime.config().general.log_level); - prepared.generation.stop_sessions().await; - if let Some(pool) = prepared.generation.current_me_pool().await { - let _ = tokio::time::timeout( - Duration::from_secs(2), - pool.shutdown_send_close_conn_all(), - ) - .await; - } - prepared.generation.stop_background_tasks().await; self.control.rolled_back(command.reload_id, warning).await; return; } - self.control.add_warning(command.reload_id, warning).await; } self.control @@ -112,37 +149,26 @@ impl ReloadSupervisor { .await; let new_runtime = prepared.generation; old_runtime.stop_accepting_sessions(); - let replaced = self.active_runtime.swap(new_runtime.clone()); if let Err(error) = crate::network::dns_overrides::install_entries( &new_runtime.config().network.dns_overrides, ) { let message = format!("runtime DNS activation failed: {}", error); if command.request.failure_policy == ReloadFailurePolicy::Rollback { - let candidate = self.active_runtime.swap(replaced.clone()); - replaced.resume_accepting_sessions(); + old_runtime.resume_accepting_sessions(); + let _ = cleanup_candidate(&new_runtime).await; self.runtime_log_filter - .apply_reload(&replaced.config().general.log_level); - let _ = crate::network::dns_overrides::install_entries( - &replaced.config().network.dns_overrides, - ); - candidate.stop_sessions().await; - if let Some(pool) = candidate.current_me_pool().await { - let _ = tokio::time::timeout( - Duration::from_secs(2), - pool.shutdown_send_close_conn_all(), - ) - .await; - } - candidate.stop_background_tasks().await; + .apply_reload(&old_runtime.config().general.log_level); self.control.rolled_back(command.reload_id, message).await; return; } self.control.add_warning(command.reload_id, message).await; } + let replaced = self.active_runtime.swap(new_runtime.clone()); self.detected_ips_tx.send_replace(prepared.detected_ips); self.runtime_log_filter .apply_reload(&new_runtime.config().general.log_level); - crate::synlimit_control::reconcile_synlimit_rules(&new_runtime.config()).await; + self.runtime_watch_tx + .send_replace(Some(new_runtime.watch_state())); info!( reload_id = command.reload_id, @@ -177,11 +203,7 @@ impl ReloadSupervisor { } } - if let Some(pool) = replaced.current_me_pool().await - && tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all()) - .await - .is_err() - { + if stop_background_and_middle_end(&replaced).await { let warning = format!( "generation {} Middle-End close broadcast timed out", replaced.id @@ -189,9 +211,46 @@ impl ReloadSupervisor { warn!(reload_id = command.reload_id, warning = %warning); self.control.add_warning(command.reload_id, warning).await; } - replaced.stop_background_tasks().await; self.control .succeed(command.reload_id, new_runtime.id) .await; } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn revision_gate_proceeds_only_on_verified_match() { + assert_eq!( + revision_gate_action( + "accepted", + Ok("accepted".to_string()), + ReloadFailurePolicy::Rollback, + ), + RevisionGateAction::Proceed + ); + } + + #[test] + fn revision_gate_applies_failure_policy_to_mismatch_and_read_error() { + for result in [ + Ok("changed".to_string()), + Err("read failed".to_string()), + ] { + assert!(matches!( + revision_gate_action( + "accepted", + result.clone(), + ReloadFailurePolicy::KeepNew, + ), + RevisionGateAction::Warn(_) + )); + assert!(matches!( + revision_gate_action("accepted", result, ReloadFailurePolicy::Rollback), + RevisionGateAction::Rollback(_) + )); + } + } +} diff --git a/src/maestro/runtime_build.rs b/src/maestro/runtime_build.rs index 4ec1062..1cf7e5b 100644 --- a/src/maestro/runtime_build.rs +++ b/src/maestro/runtime_build.rs @@ -34,7 +34,7 @@ pub(crate) struct PreparedRuntime { pub(crate) async fn prepare_runtime( generation_id: u64, - mut config: ProxyConfig, + config: ProxyConfig, config_path: &Path, quota_store: Arc, runtime_log_filter: RuntimeLogFilter, @@ -112,8 +112,10 @@ pub(crate) async fn prepare_runtime( upstream_manager.clone(), &startup_tracker, task_scope.clone(), + tls_bootstrap::TlsBootstrapPolicy::RequireReady, ) - .await; + .await + .map_err(|error| error.to_string())?; let beobachten = Arc::new(BeobachtenStore::new()); let rng = Arc::new(SecureRandom::new()); @@ -140,11 +142,20 @@ pub(crate) async fn prepare_runtime( stats.clone(), me_pool_runtime.clone(), me_ready_tx.clone(), + task_scope.clone(), ) .await }; - if config.general.use_middle_proxy && me_pool.is_none() && !direct_first_startup { - config.general.use_middle_proxy = false; + if strict_middle_proxy_unavailable( + config.general.use_middle_proxy, + direct_first_startup, + me_pool.is_some(), + ) { + task_scope.stop().await; + return Err( + "Middle-End pool is required but did not become ready during reload preparation" + .to_string(), + ); } let config = Arc::new(config); @@ -159,7 +170,6 @@ pub(crate) async fn prepare_runtime( config.server.max_connections as usize }; let max_connections = Arc::new(Semaphore::new(max_connections_limit)); - let (api_config_tx, _) = watch::channel(config.clone()); let watches = runtime_tasks::spawn_runtime_tasks( &config, config_path, @@ -175,7 +185,6 @@ pub(crate) async fn prepare_runtime( rng.clone(), ip_tracker.clone(), beobachten.clone(), - api_config_tx, me_pool.clone(), proxy_shared.clone(), me_ready_tx.clone(), @@ -226,6 +235,7 @@ pub(crate) async fn prepare_runtime( stats_bg.clone(), me_pool_runtime_bg.clone(), me_ready_tx_bg.clone(), + task_scope_bg.clone(), ) .await; if let Some(pool) = pool { @@ -291,6 +301,14 @@ pub(crate) async fn prepare_runtime( }) } +fn strict_middle_proxy_unavailable( + use_middle_proxy: bool, + direct_first_startup: bool, + pool_available: bool, +) -> bool { + use_middle_proxy && !direct_first_startup && !pool_available +} + pub(crate) fn deferred_process_fields(old: &ProxyConfig, new: &ProxyConfig) -> Vec { let mut fields = Vec::new(); if old.server.port != new.server.port @@ -354,4 +372,12 @@ mod tests { new.censorship.tls_domain = "reload.example".to_string(); assert!(deferred_process_fields(&old, &new).is_empty()); } + + #[test] + fn strict_middle_proxy_requires_a_prepared_pool() { + assert!(strict_middle_proxy_unavailable(true, false, false)); + assert!(!strict_middle_proxy_unavailable(true, false, true)); + assert!(!strict_middle_proxy_unavailable(true, true, false)); + assert!(!strict_middle_proxy_unavailable(false, false, false)); + } } diff --git a/src/maestro/runtime_tasks.rs b/src/maestro/runtime_tasks.rs index 7bd3976..1b896e6 100644 --- a/src/maestro/runtime_tasks.rs +++ b/src/maestro/runtime_tasks.rs @@ -103,7 +103,6 @@ pub(crate) async fn spawn_runtime_tasks( rng: Arc, ip_tracker: Arc, beobachten: Arc, - api_config_tx: watch::Sender>, me_pool_for_policy: Option>, shared_state: Arc, me_ready_tx: watch::Sender, @@ -166,18 +165,6 @@ pub(crate) async fn spawn_runtime_tasks( Some("config hot-reload watcher started".to_string()), ) .await; - let mut config_rx_api_bridge = config_rx.clone(); - let api_config_tx_bridge = api_config_tx.clone(); - task_scope.spawn(async move { - loop { - if config_rx_api_bridge.changed().await.is_err() { - break; - } - let cfg = config_rx_api_bridge.borrow_and_update().clone(); - api_config_tx_bridge.send_replace(cfg); - } - }); - let stats_policy = stats.clone(); let upstream_policy = upstream_manager.clone(); let mut config_rx_policy = config_rx.clone(); diff --git a/src/maestro/shutdown.rs b/src/maestro/shutdown.rs index fcb739a..922a69e 100644 --- a/src/maestro/shutdown.rs +++ b/src/maestro/shutdown.rs @@ -50,9 +50,17 @@ pub(crate) async fn wait_for_shutdown( process_started_at: Instant, active_runtime: Arc>, quota_state_path: PathBuf, + synlimit_controller: tokio::task::JoinHandle<()>, ) { let signal = wait_for_shutdown_signal().await; - perform_shutdown(signal, process_started_at, active_runtime, quota_state_path).await; + perform_shutdown( + signal, + process_started_at, + active_runtime, + quota_state_path, + synlimit_controller, + ) + .await; } /// Waits for any shutdown signal (SIGINT, SIGTERM, SIGQUIT). @@ -81,6 +89,7 @@ async fn perform_shutdown( process_started_at: Instant, active_runtime: Arc>, quota_state_path: PathBuf, + synlimit_controller: tokio::task::JoinHandle<()>, ) { let runtime = active_runtime.load_full(); let stats = runtime.stats.as_ref(); @@ -96,12 +105,9 @@ async fn perform_shutdown( let uptime_secs = process_started_at.elapsed().as_secs(); info!("Uptime: {}", format_uptime(uptime_secs)); - if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await { - warn!(error = %error, "Failed to clear SYN limiter rules during shutdown"); - } - // Graceful ME pool shutdown runtime.stop_sessions().await; + runtime.stop_background_tasks().await; if let Some(pool) = runtime.current_me_pool().await { match tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all()) .await @@ -117,7 +123,12 @@ async fn perform_shutdown( } } } - runtime.stop_background_tasks().await; + + synlimit_controller.abort(); + let _ = synlimit_controller.await; + if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await { + warn!(error = %error, "Failed to clear SYN limiter rules during shutdown"); + } match crate::quota_state::save_quota_state("a_state_path, stats).await { Ok(()) => { diff --git a/src/maestro/tls_bootstrap.rs b/src/maestro/tls_bootstrap.rs index 7b7075f..58bd442 100644 --- a/src/maestro/tls_bootstrap.rs +++ b/src/maestro/tls_bootstrap.rs @@ -5,6 +5,7 @@ use rand::RngExt; use tracing::warn; use crate::config::ProxyConfig; +use crate::error::{ProxyError, Result}; use crate::startup::{COMPONENT_TLS_FRONT_BOOTSTRAP, StartupTracker}; use crate::tls_front::TlsFrontCache; use crate::tls_front::fetcher::TlsFetchStrategy; @@ -12,6 +13,82 @@ use crate::transport::UpstreamManager; use super::generation::RuntimeTaskScope; +/// Readiness requirement for TLS-front cache initialization. +#[derive(Clone, Copy)] +pub(crate) enum TlsBootstrapPolicy { + BestEffort, + RequireReady, +} + +#[derive(Clone)] +struct TlsFetchContext { + cache: Arc, + domains: Vec, + mask_host: String, + primary_domain: String, + mask_unix_sock: Option, + tls_fetch_scope: Option, + upstream_manager: Arc, + strategy: TlsFetchStrategy, + port: u16, + proxy_protocol: u8, +} + +impl TlsFetchContext { + async fn fetch_all(&self, failure_message: &'static str) { + let mut join = tokio::task::JoinSet::new(); + for domain in self.domains.clone() { + let cache = self.cache.clone(); + let host = tls_fetch_host_for_domain( + &self.mask_host, + &self.primary_domain, + &domain, + ); + let unix_sock = self.mask_unix_sock.clone(); + let scope = self.tls_fetch_scope.clone(); + let upstream = self.upstream_manager.clone(); + let strategy = self.strategy.clone(); + let port = self.port; + let proxy_protocol = self.proxy_protocol; + join.spawn(async move { + match crate::tls_front::fetcher::fetch_real_tls_with_strategy( + &host, + port, + &domain, + &strategy, + Some(upstream), + scope.as_deref(), + proxy_protocol, + unix_sock.as_deref(), + ) + .await + { + Ok(result) => cache.update_from_fetch(&domain, result).await, + Err(error) => warn!(domain = %domain, error = %error, failure_message), + } + }); + } + while let Some(result) = join.join_next().await { + if let Err(error) = result { + warn!(error = %error, "TLS emulation fetch task join failed"); + } + } + } + + async fn fetch_all_with_budget(&self, phase: &'static str) { + if tokio::time::timeout(self.strategy.total_budget, self.fetch_all(phase)) + .await + .is_err() + { + warn!( + phase, + timeout_ms = self.strategy.total_budget.as_millis(), + "TLS emulation fetch budget exhausted" + ); + } + } +} + fn tls_fetch_host_for_domain(mask_host: &str, primary_tls_domain: &str, domain: &str) -> String { if mask_host.eq_ignore_ascii_case(primary_tls_domain) { domain.to_string() @@ -20,13 +97,24 @@ fn tls_fetch_host_for_domain(mask_host: &str, primary_tls_domain: &str, domain: } } +fn readiness_error(default_domains: &[String]) -> Option { + (!default_domains.is_empty()).then(|| { + format!( + "TLS-front profiles are not ready for domains: {}", + default_domains.join(", ") + ) + }) +} + +/// Initializes the TLS-front cache and generation-owned refresh tasks. pub(crate) async fn bootstrap_tls_front( config: &ProxyConfig, tls_domains: &[String], upstream_manager: Arc, startup_tracker: &Arc, task_scope: RuntimeTaskScope, -) -> Option> { + policy: TlsBootstrapPolicy, +) -> Result>> { startup_tracker .start_component( COMPONENT_TLS_FRONT_BOOTSTRAP, @@ -34,26 +122,38 @@ pub(crate) async fn bootstrap_tls_front( ) .await; - let tls_cache: Option> = if config.censorship.tls_emulation { - let cache = Arc::new(TlsFrontCache::new( - tls_domains, - config.censorship.fake_cert_len, - &config.censorship.tls_front_dir, - )); - cache.load_from_disk().await; + if !config.censorship.tls_emulation { + startup_tracker + .skip_component( + COMPONENT_TLS_FRONT_BOOTSTRAP, + Some("censorship.tls_emulation is false".to_string()), + ) + .await; + return Ok(None); + } - let port = config.censorship.mask_port; - let proxy_protocol = config.censorship.mask_proxy_protocol; - let mask_host = config + let cache = Arc::new(TlsFrontCache::new( + tls_domains, + config.censorship.fake_cert_len, + &config.censorship.tls_front_dir, + )); + cache.load_from_disk().await; + + let tls_fetch = config.censorship.tls_fetch.clone(); + let fetch_context = TlsFetchContext { + cache: cache.clone(), + domains: tls_domains.to_vec(), + mask_host: config .censorship .mask_host .clone() - .unwrap_or_else(|| config.censorship.tls_domain.clone()); - let mask_unix_sock = config.censorship.mask_unix_sock.clone(); - let tls_fetch_scope = (!config.censorship.tls_fetch_scope.is_empty()) - .then(|| config.censorship.tls_fetch_scope.clone()); - let tls_fetch = config.censorship.tls_fetch.clone(); - let fetch_strategy = TlsFetchStrategy { + .unwrap_or_else(|| config.censorship.tls_domain.clone()), + primary_domain: config.censorship.tls_domain.clone(), + mask_unix_sock: config.censorship.mask_unix_sock.clone(), + tls_fetch_scope: (!config.censorship.tls_fetch_scope.is_empty()) + .then(|| config.censorship.tls_fetch_scope.clone()), + upstream_manager, + strategy: TlsFetchStrategy { profiles: tls_fetch.profiles, strict_route: tls_fetch.strict_route, attempt_timeout: Duration::from_millis(tls_fetch.attempt_timeout_ms.max(1)), @@ -61,150 +161,71 @@ pub(crate) async fn bootstrap_tls_front( grease_enabled: tls_fetch.grease_enabled, deterministic: tls_fetch.deterministic, profile_cache_ttl: Duration::from_secs(tls_fetch.profile_cache_ttl_secs), - }; - let fetch_timeout = fetch_strategy.total_budget; + }, + port: config.censorship.mask_port, + proxy_protocol: config.censorship.mask_proxy_protocol, + }; - let cache_initial = cache.clone(); - let domains_initial = tls_domains.to_vec(); - let host_initial = mask_host.clone(); - let primary_initial = config.censorship.tls_domain.clone(); - let unix_sock_initial = mask_unix_sock.clone(); - let scope_initial = tls_fetch_scope.clone(); - let upstream_initial = upstream_manager.clone(); - let strategy_initial = fetch_strategy.clone(); - task_scope.spawn(async move { - let mut join = tokio::task::JoinSet::new(); - for domain in domains_initial { - let cache_domain = cache_initial.clone(); - let host_domain = - tls_fetch_host_for_domain(&host_initial, &primary_initial, &domain); - let unix_sock_domain = unix_sock_initial.clone(); - let scope_domain = scope_initial.clone(); - let upstream_domain = upstream_initial.clone(); - let strategy_domain = strategy_initial.clone(); - join.spawn(async move { - match crate::tls_front::fetcher::fetch_real_tls_with_strategy( - &host_domain, - port, - &domain, - &strategy_domain, - Some(upstream_domain), - scope_domain.as_deref(), - proxy_protocol, - unix_sock_domain.as_deref(), - ) + match policy { + TlsBootstrapPolicy::BestEffort => { + let initial_fetch = fetch_context.clone(); + let fake_cert_len = config.censorship.fake_cert_len; + task_scope.spawn(async move { + initial_fetch + .fetch_all_with_budget("TLS emulation initial fetch failed") + .await; + for domain in initial_fetch + .cache + .default_profile_domains(&initial_fetch.domains) .await - { - Ok(res) => cache_domain.update_from_fetch(&domain, res).await, - Err(e) => { - warn!(domain = %domain, error = %e, "TLS emulation initial fetch failed") - } - } - }); - } - while let Some(res) = join.join_next().await { - if let Err(e) = res { - warn!(error = %e, "TLS emulation initial fetch task join failed"); - } - } - }); - - let cache_timeout = cache.clone(); - let domains_timeout = tls_domains.to_vec(); - let fake_cert_len = config.censorship.fake_cert_len; - task_scope.spawn(async move { - tokio::time::sleep(fetch_timeout).await; - for domain in domains_timeout { - let cached = cache_timeout.get(&domain).await; - if cached.domain == "default" { + { warn!( domain = %domain, - timeout_secs = fetch_timeout.as_secs(), + timeout_ms = initial_fetch.strategy.total_budget.as_millis(), fake_cert_len, "TLS-front fetch not ready within timeout; using cache/default fake cert fallback" ); } + }); + } + TlsBootstrapPolicy::RequireReady => { + fetch_context + .fetch_all_with_budget("TLS emulation initial fetch failed") + .await; + let default_domains = cache.default_profile_domains(tls_domains).await; + if let Some(error) = readiness_error(&default_domains) { + startup_tracker + .fail_component(COMPONENT_TLS_FRONT_BOOTSTRAP, Some(error.clone())) + .await; + return Err(ProxyError::Proxy(error)); } - }); - - let cache_refresh = cache.clone(); - let domains_refresh = tls_domains.to_vec(); - let host_refresh = mask_host.clone(); - let primary_refresh = config.censorship.tls_domain.clone(); - let unix_sock_refresh = mask_unix_sock.clone(); - let scope_refresh = tls_fetch_scope.clone(); - let upstream_refresh = upstream_manager.clone(); - let strategy_refresh = fetch_strategy.clone(); - task_scope.spawn(async move { - loop { - let base_secs = rand::rng().random_range(4 * 3600..=6 * 3600); - let jitter_secs = rand::rng().random_range(0..=7200); - tokio::time::sleep(Duration::from_secs(base_secs + jitter_secs)).await; - - let mut join = tokio::task::JoinSet::new(); - for domain in domains_refresh.clone() { - let cache_domain = cache_refresh.clone(); - let host_domain = - tls_fetch_host_for_domain(&host_refresh, &primary_refresh, &domain); - let unix_sock_domain = unix_sock_refresh.clone(); - let scope_domain = scope_refresh.clone(); - let upstream_domain = upstream_refresh.clone(); - let strategy_domain = strategy_refresh.clone(); - join.spawn(async move { - match crate::tls_front::fetcher::fetch_real_tls_with_strategy( - &host_domain, - port, - &domain, - &strategy_domain, - Some(upstream_domain), - scope_domain.as_deref(), - proxy_protocol, - unix_sock_domain.as_deref(), - ) - .await - { - Ok(res) => cache_domain.update_from_fetch(&domain, res).await, - Err(e) => { - warn!(domain = %domain, error = %e, "TLS emulation refresh failed") - } - } - }); - } - - while let Some(res) = join.join_next().await { - if let Err(e) = res { - warn!(error = %e, "TLS emulation refresh task join failed"); - } - } - } - }); - - Some(cache) - } else { - startup_tracker - .skip_component( - COMPONENT_TLS_FRONT_BOOTSTRAP, - Some("censorship.tls_emulation is false".to_string()), - ) - .await; - None - }; - - if tls_cache.is_some() { - startup_tracker - .complete_component( - COMPONENT_TLS_FRONT_BOOTSTRAP, - Some("tls front cache is initialized".to_string()), - ) - .await; + } } - tls_cache + let refresh_context = fetch_context; + task_scope.spawn(async move { + loop { + let base_secs = rand::rng().random_range(4 * 3600..=6 * 3600); + let jitter_secs = rand::rng().random_range(0..=7200); + tokio::time::sleep(Duration::from_secs(base_secs + jitter_secs)).await; + refresh_context + .fetch_all_with_budget("TLS emulation refresh failed") + .await; + } + }); + + startup_tracker + .complete_component( + COMPONENT_TLS_FRONT_BOOTSTRAP, + Some("tls front cache is initialized".to_string()), + ) + .await; + Ok(Some(cache)) } #[cfg(test)] mod tests { - use super::tls_fetch_host_for_domain; + use super::{readiness_error, tls_fetch_host_for_domain}; #[test] fn tls_fetch_host_uses_each_domain_when_mask_host_is_primary_default() { @@ -221,4 +242,13 @@ mod tests { "origin.example" ); } + + #[test] + fn readiness_rejects_only_default_profiles() { + assert!(readiness_error(&[]).is_none()); + assert_eq!( + readiness_error(&["front.example".to_string()]), + Some("TLS-front profiles are not ready for domains: front.example".to_string()) + ); + } } diff --git a/src/synlimit_control/mod.rs b/src/synlimit_control/mod.rs index 448a57a..b87f2b2 100644 --- a/src/synlimit_control/mod.rs +++ b/src/synlimit_control/mod.rs @@ -4,6 +4,7 @@ use tokio::sync::watch; use tracing::warn; use crate::config::{ProxyConfig, SynLimitMode}; +use crate::maestro::generation::RuntimeWatchState; mod command; mod iptables; @@ -15,28 +16,105 @@ use self::model::{SynLimitNamespace, synlimit_namespace, synlimit_targets}; static ACTIVE_SYNLIMIT_NAMESPACE: Mutex> = Mutex::new(None); -pub(crate) fn spawn_synlimit_controller(config_rx: watch::Receiver>) { +/// Spawns the process-scoped SYN limiter reconciler for active generations. +pub(crate) fn spawn_synlimit_controller( + runtime_watch_rx: watch::Receiver>, +) -> tokio::task::JoinHandle<()> { if !cfg!(target_os = "linux") { - if has_synlimit_config(&config_rx.borrow()) { - warn!("SYN limiter is configured but unsupported on this OS; skipping netfilter rules"); - } - return; + return tokio::spawn(watch_active_runtime_configs( + runtime_watch_rx, + |_generation_id, cfg| async move { + if has_synlimit_config(&cfg) { + warn!( + "SYN limiter is configured but unsupported on this OS; skipping netfilter rules" + ); + } + }, + )); } - tokio::spawn(async move { - wait_for_config_channel_close_and_reconcile(config_rx).await; - if let Err(error) = clear_synlimit_rules_all_backends().await { - warn!(error = %error, "Failed to clear SYN limiter rules after config channel close"); - } - }); + tokio::spawn(watch_active_runtime_configs( + runtime_watch_rx, + |_generation_id, cfg| async move { + reconcile_synlimit_rules(&cfg).await; + }, + )) } -async fn wait_for_config_channel_close_and_reconcile( - mut config_rx: watch::Receiver>, -) { - while config_rx.changed().await.is_ok() { - let cfg = config_rx.borrow_and_update().clone(); - reconcile_synlimit_rules(&cfg).await; +async fn watch_active_runtime_configs( + mut runtime_watch_rx: watch::Receiver>, + mut on_config: F, +) +where + F: FnMut(u64, Arc) -> Fut, + Fut: std::future::Future, +{ + let mut current = loop { + if let Some(state) = runtime_watch_rx.borrow().clone() { + break state; + } + if runtime_watch_rx.changed().await.is_err() { + return; + } + }; + let initial_config = current.config_rx.borrow().clone(); + on_config(current.generation_id, initial_config).await; + + loop { + tokio::select! { + biased; + changed = runtime_watch_rx.changed() => { + if changed.is_err() { + break; + } + let Some(next) = runtime_watch_rx.borrow().clone() else { + continue; + }; + if next.generation_id != current.generation_id { + current = next; + let config = current.config_rx.borrow().clone(); + on_config(current.generation_id, config).await; + } + } + changed = current.config_rx.changed() => { + if changed.is_err() { + let Some(next) = wait_for_new_runtime( + &mut runtime_watch_rx, + current.generation_id, + ).await else { + break; + }; + current = next; + let config = current.config_rx.borrow().clone(); + on_config(current.generation_id, config).await; + continue; + } + let active_generation_id = runtime_watch_rx + .borrow() + .as_ref() + .map(|state| state.generation_id); + if active_generation_id == Some(current.generation_id) { + let cfg = current.config_rx.borrow_and_update().clone(); + on_config(current.generation_id, cfg).await; + } + } + } + } +} + +async fn wait_for_new_runtime( + runtime_watch_rx: &mut watch::Receiver>, + previous_generation_id: u64, +) -> Option { + loop { + if let Some(state) = runtime_watch_rx.borrow().clone() + && state.generation_id != previous_generation_id + { + return Some(state); + } + if runtime_watch_rx.changed().await.is_err() { + return None; + } } } @@ -168,3 +246,70 @@ fn has_synlimit_config(cfg: &ProxyConfig) -> bool { .iter() .any(|listener| !matches!(listener.synlimit, SynLimitMode::Off)) } + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + use tokio::sync::mpsc; + + fn runtime_state( + generation_id: u64, + max_connections: u32, + ) -> ( + RuntimeWatchState, + watch::Sender>, + watch::Sender, + ) { + let mut config = ProxyConfig::default(); + config.server.max_connections = max_connections; + let (config_tx, config_rx) = watch::channel(Arc::new(config)); + let (admission_tx, admission_rx) = watch::channel(true); + ( + RuntimeWatchState { + generation_id, + config_rx, + admission_rx, + }, + config_tx, + admission_tx, + ) + } + + #[tokio::test] + async fn config_watcher_ignores_retired_generation_updates() { + let (initial, initial_config_tx, _initial_admission_tx) = runtime_state(1, 10); + let (runtime_tx, runtime_rx) = watch::channel(Some(initial)); + let (observed_tx, mut observed_rx) = mpsc::unbounded_channel(); + let watcher = tokio::spawn(watch_active_runtime_configs( + runtime_rx, + move |generation_id, cfg| { + let observed_tx = observed_tx.clone(); + async move { + let _ = observed_tx.send((generation_id, cfg.server.max_connections)); + } + }, + )); + + assert_eq!(observed_rx.recv().await, Some((1, 10))); + let (next, next_config_tx, _next_admission_tx) = runtime_state(2, 20); + runtime_tx.send_replace(Some(next)); + assert_eq!(observed_rx.recv().await, Some((2, 20))); + + let mut stale = ProxyConfig::default(); + stale.server.max_connections = 30; + initial_config_tx.send_replace(Arc::new(stale)); + assert!( + tokio::time::timeout(Duration::from_millis(50), observed_rx.recv()) + .await + .is_err() + ); + + let mut active = ProxyConfig::default(); + active.server.max_connections = 40; + next_config_tx.send_replace(Arc::new(active)); + assert_eq!(observed_rx.recv().await, Some((2, 40))); + + watcher.abort(); + } +} diff --git a/src/tls_front/cache.rs b/src/tls_front/cache.rs index c63a539..42e9d20 100644 --- a/src/tls_front/cache.rs +++ b/src/tls_front/cache.rs @@ -190,6 +190,22 @@ impl TlsFrontCache { (snapshot, suppressed) } + /// Returns configured domains that still resolve to the synthetic default profile. + pub(crate) async fn default_profile_domains(&self, domains: &[String]) -> Vec { + let guard = self.memory.read().await; + domains + .iter() + .filter(|domain| { + guard + .get(domain.as_str()) + .unwrap_or(&self.default) + .domain + == "default" + }) + .cloned() + .collect() + } + fn full_cert_sent_shard_index(client_ip: IpAddr) -> usize { let mut hasher = DefaultHasher::new(); client_ip.hash(&mut hasher); @@ -546,6 +562,23 @@ mod tests { assert!(!cert_info_matches_domain(&cached)); } + #[tokio::test] + async fn default_profile_domains_reports_only_unprepared_entries() { + let domains = vec!["ready.example".to_string(), "pending.example".to_string()]; + let cache = TlsFrontCache::new(&domains, 1024, "tlsfront-test-cache"); + cache + .set( + "ready.example", + cached_with_cert_info("ready.example", None, Vec::new()), + ) + .await; + + assert_eq!( + cache.default_profile_domains(&domains).await, + vec!["pending.example".to_string()] + ); + } + #[tokio::test] async fn test_take_full_cert_budget_for_ip_uses_ttl() { let cache = TlsFrontCache::new(&["example.com".to_string()], 1024, "tlsfront-test-cache"); diff --git a/src/tls_front/fetcher.rs b/src/tls_front/fetcher.rs index 22f8548..3bcc3c5 100644 --- a/src/tls_front/fetcher.rs +++ b/src/tls_front/fetcher.rs @@ -916,26 +916,22 @@ async fn connect_tcp_with_upstream( strict_route: bool, ) -> Result { if let Some(manager) = upstream { - let resolved = if let Some(addr) = resolve_socket_addr(host, port) { - Some(addr) - } else { - match tokio::net::lookup_host((host, port)).await { - Ok(mut addrs) => addrs.find(|a| a.is_ipv4()), - Err(e) => { - if strict_route { - return Err(anyhow!( - "upstream route DNS resolution failed for {host}:{port}: {e}" - )); - } - warn!( - host = %host, - port = port, - scope = ?scope, - error = %e, - "Upstream DNS resolution failed, using direct connect" - ); - None + let resolved = match manager.resolve_hostname(host, port).await { + Ok(addr) => Some(addr), + Err(e) => { + if strict_route { + return Err(anyhow!( + "upstream route DNS resolution failed for {host}:{port}: {e}" + )); } + warn!( + host = %host, + port = port, + scope = ?scope, + error = %e, + "Upstream DNS resolution failed, using direct connect" + ); + None } }; @@ -955,6 +951,9 @@ async fn connect_tcp_with_upstream( error = %e, "Upstream connect failed, using direct connect" ); + return Ok(UpstreamStream::Tcp( + timeout(connect_timeout, TcpStream::connect(addr)).await??, + )); } } } else if strict_route {