From d65efcc479f6a02d688c1392d2b5220aba6aea07 Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Fri, 2 Oct 2026 23:51:14 +0300 Subject: [PATCH] Fix autonomous ME recovery from Direct fallback Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com> --- src/config/defaults.rs | 3 +- .../tests/load_basic_tests/conntrack_tests.rs | 66 +++ src/conntrack_control/firewall/actor.rs | 77 ++- .../firewall/tests/startup_admission.rs | 267 ++++++++++ src/maestro/admission.rs | 24 +- src/maestro/admission/tests.rs | 210 ++++++++ src/maestro/orchestrator.rs | 16 +- src/service/mod.rs | 39 +- src/transport/middle_proxy/health/recovery.rs | 26 +- .../middle_proxy/health/recovery/tests.rs | 148 ++++++ src/transport/middle_proxy/mod.rs | 5 + .../middle_proxy/pool/writer_admission.rs | 8 +- src/transport/middle_proxy/pool_writer.rs | 4 + .../middle_proxy/pool_writer/publication.rs | 12 +- .../pool_writer/recovery_tests.rs | 466 ++++++++++++++++++ .../tests/admission_test_support.rs | 98 ++++ 16 files changed, 1423 insertions(+), 46 deletions(-) create mode 100644 src/conntrack_control/firewall/tests/startup_admission.rs create mode 100644 src/maestro/admission/tests.rs create mode 100644 src/transport/middle_proxy/health/recovery/tests.rs create mode 100644 src/transport/middle_proxy/pool_writer/recovery_tests.rs create mode 100644 src/transport/middle_proxy/tests/admission_test_support.rs diff --git a/src/config/defaults.rs b/src/config/defaults.rs index b2ad135..635bed2 100644 --- a/src/config/defaults.rs +++ b/src/config/defaults.rs @@ -59,7 +59,7 @@ const DEFAULT_ME_POOL_DRAIN_SOFT_EVICT_BUDGET_PER_CORE: u16 = 16; const DEFAULT_ME_POOL_DRAIN_SOFT_EVICT_COOLDOWN_MS: u64 = 1000; const DEFAULT_USER_MAX_UNIQUE_IPS_WINDOW_SECS: u64 = 30; const DEFAULT_ACCEPT_PERMIT_TIMEOUT_MS: u64 = 250; -const DEFAULT_CONNTRACK_CONTROL_ENABLED: bool = true; +const DEFAULT_CONNTRACK_CONTROL_ENABLED: bool = false; const DEFAULT_CONNTRACK_PRESSURE_HIGH_WATERMARK_PCT: u8 = 85; const DEFAULT_CONNTRACK_PRESSURE_LOW_WATERMARK_PCT: u8 = 70; const DEFAULT_CONNTRACK_DELETE_BUDGET_PER_SEC: u64 = 4096; @@ -244,6 +244,7 @@ pub(crate) fn default_accept_permit_timeout_ms() -> u64 { DEFAULT_ACCEPT_PERMIT_TIMEOUT_MS } +/// Keeps privileged conntrack control disabled unless explicitly requested. pub(crate) fn default_conntrack_control_enabled() -> bool { DEFAULT_CONNTRACK_CONTROL_ENABLED } diff --git a/src/config/tests/load_basic_tests/conntrack_tests.rs b/src/config/tests/load_basic_tests/conntrack_tests.rs index e8fb3e3..714e644 100644 --- a/src/config/tests/load_basic_tests/conntrack_tests.rs +++ b/src/config/tests/load_basic_tests/conntrack_tests.rs @@ -1,5 +1,71 @@ use super::*; +#[test] +fn conntrack_control_is_opt_in_by_default() { + assert!(!default_conntrack_control_enabled()); + assert!( + !ProxyConfig::default() + .server + .conntrack_control + .inline_conntrack_control + ); + let cfg = load_config_from_temp_toml( + r#" + [censorship] + tls_domain = "example.com" + + [access.users] + user = "00000000000000000000000000000000" + "#, + ); + assert!(!cfg.server.conntrack_control.inline_conntrack_control); + assert!( + !cfg.server + .conntrack_control + .inline_conntrack_control_explicit + ); + assert_eq!( + serde_json::to_value(&cfg.server.conntrack_control).unwrap()["inline_conntrack_control"], + false, + ); +} + +#[test] +fn conntrack_control_preserves_explicit_enablement() { + for enabled in [false, true] { + let cfg = load_config_from_temp_toml(&format!( + r#" + [server.conntrack_control] + inline_conntrack_control = {enabled} + backend = "nftables" + + [censorship] + tls_domain = "example.com" + + [access.users] + user = "00000000000000000000000000000000" + "#, + )); + assert_eq!( + cfg.server.conntrack_control.inline_conntrack_control, + enabled + ); + assert!( + cfg.server + .conntrack_control + .inline_conntrack_control_explicit + ); + assert_eq!( + cfg.server.conntrack_control.backend, + ConntrackBackend::Nftables + ); + assert_eq!( + serde_json::to_value(&cfg.server.conntrack_control).unwrap()["inline_conntrack_control"], + enabled, + ); + } +} + #[test] fn conntrack_pressure_high_watermark_out_of_range_is_rejected() { let toml = r#" diff --git a/src/conntrack_control/firewall/actor.rs b/src/conntrack_control/firewall/actor.rs index 13d6d3f..da68a8f 100644 --- a/src/conntrack_control/firewall/actor.rs +++ b/src/conntrack_control/firewall/actor.rs @@ -26,15 +26,21 @@ const RETRY_DELAYS: [Duration; 6] = [ Duration::from_secs(30), ]; +/// Reports whether a generation's desired firewall policy was confirmed. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(super) enum ReconcileOutcome { + /// The applied policy matches the accepted generation's desired policy. Applied, + /// The attempt failed without confirming the desired policy. Failed, } +/// Associates a reconciliation result with its accepted runtime generation. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(super) struct ReconcileStatus { + /// Runtime generation whose desired policy was attempted. pub(super) generation: u64, + /// Result of the latest completed attempt for this generation. pub(super) outcome: ReconcileOutcome, } @@ -51,25 +57,14 @@ pub(crate) struct FirewallAuthority { } impl FirewallAuthority { - /// Starts the single process-owned firewall reconciler. - pub(crate) fn spawn(control_plane: &ProcessControlPlane) -> Result { - let (desired_tx, desired_rx) = watch::channel(None); - let (status_tx, status_rx) = watch::channel(None); - let terminal = CancellationToken::new(); - let closed = Arc::new(AtomicBool::new(false)); - let completed_flag = Arc::new(AtomicBool::new(false)); - let cleanup_succeeded = Arc::new(AtomicBool::new(false)); - let completed = Arc::new(Notify::new()); - let actor = FirewallReconciler::new( - SystemCommandRunner, - desired_rx, - status_tx, - terminal.clone(), - Arc::clone(&closed), - Arc::clone(&completed_flag), - Arc::clone(&cleanup_succeeded), - Arc::clone(&completed), - ); + /// Starts reconciliation only when enabled and CAP_NET_ADMIN is available. + pub(crate) fn spawn( + control_plane: &ProcessControlPlane, + config: &ProxyConfig, + ) -> Result, String> { + let Some((authority, actor)) = Self::prepare(config, SystemCommandRunner) else { + return Ok(None); + }; control_plane .spawn_cooperative(move |process_cancellation| async move { actor.run(process_cancellation).await; @@ -78,7 +73,37 @@ impl FirewallAuthority { "process control-plane admission closed before conntrack firewall startup" .to_string() })?; - Ok(Self { + Ok(Some(authority)) + } + + // Admission precedes allocation so an opted-out process never owns firewall cleanup. + // Keeping construction separate permits fake runners without changing async Send bounds. + fn prepare(config: &ProxyConfig, runner: R) -> Option<(Self, FirewallReconciler)> + where + R: FirewallCommandRunner + 'static, + { + if !config.server.conntrack_control.inline_conntrack_control || !runner.has_cap_net_admin() + { + return None; + } + let (desired_tx, desired_rx) = watch::channel(None); + let (status_tx, status_rx) = watch::channel(None); + let terminal = CancellationToken::new(); + let closed = Arc::new(AtomicBool::new(false)); + let completed_flag = Arc::new(AtomicBool::new(false)); + let cleanup_succeeded = Arc::new(AtomicBool::new(false)); + let completed = Arc::new(Notify::new()); + let actor = FirewallReconciler::new( + runner, + desired_rx, + status_tx, + terminal.clone(), + Arc::clone(&closed), + Arc::clone(&completed_flag), + Arc::clone(&cleanup_succeeded), + Arc::clone(&completed), + ); + let authority = Self { desired_tx, status_rx, terminal, @@ -86,7 +111,8 @@ impl FirewallAuthority { completed_flag, cleanup_succeeded, completed, - }) + }; + Some((authority, actor)) } /// Publishes policy only after its runtime generation becomes active. @@ -169,6 +195,7 @@ impl Drop for CompletionGuard { } } +/// Serializes confirmed-state transitions and retains cleanup ownership until shutdown. pub(super) struct FirewallReconciler { runner: R, desired_rx: watch::Receiver>, @@ -186,6 +213,7 @@ impl FirewallReconciler where R: FirewallCommandRunner + 'static, { + /// Constructs an admitted actor with unknown state pending startup recovery. pub(super) fn new( runner: R, desired_rx: watch::Receiver>, @@ -214,6 +242,7 @@ where } } + /// Reconciles accepted generations and completes bounded cleanup on cancellation. pub(super) async fn run(mut self, process_cancellation: CancellationToken) { let mut current = None; let mut retry_index = 0usize; @@ -337,6 +366,7 @@ where let _completion = &self.completion; } + /// Accepts fenced publications without replacing telemetry with stale desired state. pub(super) fn take_latest_desired(&mut self) -> Option { let next = self.desired_rx.borrow_and_update().clone()?; if next.generation < self.last_generation { @@ -369,3 +399,8 @@ where Some(next) } } + +// Exercises startup admission and cleanup ownership with fake helper processes. +#[cfg(test)] +#[path = "tests/startup_admission.rs"] +mod startup_admission_tests; diff --git a/src/conntrack_control/firewall/tests/startup_admission.rs b/src/conntrack_control/firewall/tests/startup_admission.rs new file mode 100644 index 0000000..b68ac0a --- /dev/null +++ b/src/conntrack_control/firewall/tests/startup_admission.rs @@ -0,0 +1,267 @@ +use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; + +use crate::config::{ConntrackBackend, ConntrackMode}; + +use super::super::command::{CommandErrorKind, CommandSpec, classify_command_error}; +use super::super::transaction::reconcile_once; +use super::*; + +const NETLINK_ERROR: &str = + "src/mnl.c:68: Unable to initialize Netlink socket: Address family not supported by protocol"; +const PERMISSION_ERROR: &str = "iptables v1.8.11 (nf_tables): Could not fetch rule set generation id: Permission denied (you must be root)"; + +#[derive(Clone)] +struct StartupRunner { + has_cap_net_admin: bool, + cap_probes: Arc, + calls: Arc>>, + failure: Option<&'static str>, +} + +impl StartupRunner { + fn new(has_cap_net_admin: bool, failure: Option<&'static str>) -> Self { + Self { + has_cap_net_admin, + cap_probes: Arc::new(AtomicUsize::new(0)), + calls: Arc::new(Mutex::new(Vec::new())), + failure, + } + } + + fn calls(&self) -> Vec { + self.calls.lock().unwrap().clone() + } +} + +impl FirewallCommandRunner for StartupRunner { + fn available(&self, _binary: &str) -> bool { + true + } + + fn has_cap_net_admin(&self) -> bool { + self.cap_probes.fetch_add(1, Ordering::Relaxed); + self.has_cap_net_admin + } + + async fn run(&self, spec: CommandSpec) -> Result<(), CommandError> { + self.calls.lock().unwrap().push(spec.clone()); + if let Some(message) = self.failure { + return Err(CommandError { + kind: classify_command_error(spec.binary, &spec.args, message), + message: message.to_string(), + }); + } + if (spec.binary == "nft" && spec.args.first().map(String::as_str) == Some("delete")) + || (matches!(spec.binary, "iptables" | "ip6tables") + && matches!( + spec.args.get(2).map(String::as_str), + Some("-C" | "-D" | "-F" | "-X") + )) + { + return Err(CommandError { + kind: CommandErrorKind::NotFound, + message: "injected absent owned object".to_string(), + }); + } + Ok(()) + } +} + +fn start_with_runner( + config: &ProxyConfig, + runner: StartupRunner, + scope: &ProcessControlPlane, +) -> Option { + let (authority, actor) = FirewallAuthority::prepare(config, runner)?; + assert!( + scope + .spawn_cooperative(move |cancellation| actor.run(cancellation)) + .is_ok() + ); + Some(authority) +} + +#[tokio::test(start_paused = true)] +async fn disabled_startup_never_acquires_cleanup_ownership() { + for explicit in [false, true] { + for has_cap in [true, false] { + for backend in [ + ConntrackBackend::Auto, + ConntrackBackend::Nftables, + ConntrackBackend::Iptables, + ] { + for mode in [ + ConntrackMode::Tracked, + ConntrackMode::Notrack, + ConntrackMode::Hybrid, + ] { + for message in [NETLINK_ERROR, PERMISSION_ERROR] { + let mut config = ProxyConfig::default(); + if explicit { + config.server.conntrack_control.inline_conntrack_control = false; + config + .server + .conntrack_control + .inline_conntrack_control_explicit = true; + } + config.server.conntrack_control.backend = backend; + config.server.conntrack_control.mode = mode; + let runner = StartupRunner::new(has_cap, Some(message)); + let scope = ProcessControlPlane::new(); + let authority = start_with_runner(&config, runner.clone(), &scope); + assert!( + authority.is_none(), + "{explicit} {has_cap} {backend:?} {mode:?}" + ); + assert_eq!(runner.cap_probes.load(Ordering::Relaxed), 0); + tokio::time::advance(Duration::from_secs(90)).await; + assert!(scope.shutdown(Duration::from_secs(1)).await); + assert!(runner.calls().is_empty()); + } + } + } + } + } +} + +#[tokio::test(start_paused = true)] +async fn enabled_startup_without_capability_creates_no_firewall_work() { + let mut config = ProxyConfig::default(); + config.server.conntrack_control.inline_conntrack_control = true; + for mode in [ConntrackMode::Tracked, ConntrackMode::Notrack] { + config.server.conntrack_control.mode = mode; + let runner = StartupRunner::new(false, Some(PERMISSION_ERROR)); + let scope = ProcessControlPlane::new(); + assert!(start_with_runner(&config, runner.clone(), &scope).is_none()); + assert_eq!(runner.cap_probes.load(Ordering::Relaxed), 1); + tokio::time::advance(Duration::from_secs(90)).await; + assert!(scope.shutdown(Duration::from_secs(1)).await); + assert!(runner.calls().is_empty()); + } +} + +#[tokio::test(start_paused = true)] +async fn admitted_startup_preserves_recovery_idle_and_shutdown_cleanup() { + for mode in [ConntrackMode::Tracked, ConntrackMode::Notrack] { + let mut config = ProxyConfig::default(); + config.server.conntrack_control.inline_conntrack_control = true; + config.server.conntrack_control.mode = mode; + config.server.conntrack_control.backend = ConntrackBackend::Nftables; + let runner = StartupRunner::new(true, None); + let scope = ProcessControlPlane::new(); + let authority = start_with_runner(&config, runner.clone(), &scope).unwrap(); + let stats = Arc::new(Stats::new()); + assert!( + authority + .publish_initial(1, Arc::new(config), stats.clone()) + .await + ); + assert!(stats.get_conntrack_rule_apply_ok()); + for binary in ["nft", "iptables", "ip6tables"] { + assert!(runner.calls().iter().any(|call| call.binary == binary)); + } + let calls = runner.calls().len(); + tokio::time::advance(Duration::from_secs(90)).await; + tokio::task::yield_now().await; + assert_eq!(runner.calls().len(), calls); + assert!(authority.shutdown_and_clear().await); + assert!(!stats.get_conntrack_rule_apply_ok()); + assert!(runner.calls().len() > calls); + assert!(scope.shutdown(Duration::from_secs(1)).await); + } +} + +#[tokio::test(start_paused = true)] +async fn admitted_process_cancellation_preserves_cleanup() { + let mut config = ProxyConfig::default(); + config.server.conntrack_control.inline_conntrack_control = true; + let runner = StartupRunner::new(true, None); + let scope = ProcessControlPlane::new(); + let authority = start_with_runner(&config, runner.clone(), &scope).unwrap(); + assert!(scope.shutdown(Duration::from_secs(1)).await); + assert!(authority.shutdown_and_clear().await); + assert!(!runner.calls().is_empty()); +} + +#[test] +fn netlink_and_permission_errors_are_not_classified_as_absence() { + for spec in [ + CommandSpec::new("nft", ["delete", "table", "inet", "telemt_conntrack"]), + CommandSpec::new("nft", ["-f", "-"]), + CommandSpec::new( + "iptables", + ["-t", "raw", "-D", "PREROUTING", "-j", "TELEMT_NOTRACK"], + ), + CommandSpec::new("iptables", ["-t", "raw", "-F", "TELEMT_NT_A"]), + CommandSpec::new("ip6tables", ["-t", "raw", "-X", "TELEMT_NT_B"]), + CommandSpec::new("iptables-restore", ["--noflush"]), + CommandSpec::new("ip6tables-restore", ["--noflush"]), + ] { + for message in [ + NETLINK_ERROR, + PERMISSION_ERROR, + "iptables: Failed to initialize nft: Address family not supported by protocol", + ] { + assert_eq!( + classify_command_error(spec.binary, &spec.args, message), + CommandErrorKind::Failed, + "{spec:?} {message}", + ); + } + } +} + +#[tokio::test(start_paused = true)] +async fn admitted_genuine_failures_remain_unknown_and_retry() { + for message in [NETLINK_ERROR, PERMISSION_ERROR] { + let mut config = ProxyConfig::default(); + config.server.conntrack_control.inline_conntrack_control = true; + config.server.conntrack_control.mode = ConntrackMode::Notrack; + config.server.conntrack_control.backend = ConntrackBackend::Nftables; + let runner = StartupRunner::new(true, Some(message)); + let mut applied = AppliedState::Unknown; + let requested = DesiredState { + generation: 1, + policy: DesiredPolicy::from_config(&config), + stats: Arc::new(Stats::new()), + }; + let error = reconcile_once(&runner, &runner, &mut applied, &requested) + .await + .unwrap_err(); + assert!(error.message.contains(message)); + assert_eq!(applied, AppliedState::Unknown); + assert!(runner.calls().iter().all(|call| call.stdin.is_none())); + + let scope = ProcessControlPlane::new(); + let authority = start_with_runner(&config, runner.clone(), &scope).unwrap(); + let stats = Arc::new(Stats::new()); + assert!( + !authority + .publish_initial(1, Arc::new(config), stats.clone()) + .await + ); + assert!(!stats.get_conntrack_rule_apply_ok()); + assert_eq!(stats.get_conntrack_rule_reconcile_error_total(), 1); + let mut status = authority.status_rx.clone(); + assert_eq!( + status.borrow_and_update().as_ref().unwrap().outcome, + ReconcileOutcome::Failed, + ); + let calls = runner.calls().len(); + tokio::time::advance(Duration::from_millis(999)).await; + tokio::task::yield_now().await; + assert_eq!(runner.calls().len(), calls); + tokio::time::advance(Duration::from_millis(1)).await; + status.changed().await.unwrap(); + assert_eq!( + status.borrow().as_ref().unwrap().outcome, + ReconcileOutcome::Failed, + ); + assert_eq!(stats.get_conntrack_rule_reconcile_error_total(), 2); + assert!(!authority.shutdown_and_clear().await); + assert!(authority.completed_flag.load(Ordering::Acquire)); + assert!(scope.shutdown(Duration::from_secs(1)).await); + assert!(runner.calls().iter().all(|call| call.stdin.is_none())); + } +} diff --git a/src/maestro/admission.rs b/src/maestro/admission.rs index b6aadc3..4d2a05e 100644 --- a/src/maestro/admission.rs +++ b/src/maestro/admission.rs @@ -13,6 +13,11 @@ use super::generation::RuntimeTaskScope; const STARTUP_FALLBACK_AFTER: Duration = Duration::from_secs(80); const RUNTIME_FALLBACK_AFTER: Duration = Duration::from_secs(6); +// Admission regressions cover notification loss without network startup. +#[cfg(test)] +mod tests; + +/// Keeps generation admission and routing synchronized with periodic ME readiness. pub(crate) async fn configure_admission_gate( config: &Arc, me_pool: Option>, @@ -76,11 +81,18 @@ pub(crate) async fn configure_admission_gate( } else { Some(Instant::now()) }; + let mut config_watch_open = true; + let mut me_ready_watch_open = true; loop { tokio::select! { - changed = config_rx_gate.changed() => { + changed = config_rx_gate.changed(), if config_watch_open => { if changed.is_err() { - break; + config_watch_open = false; + warn!( + watch_channel = "config", + "Admission config watch closed; continuing readiness polling with last configuration" + ); + continue; } let cfg = config_rx_gate.borrow_and_update().clone(); admission_poll_ms = cfg.general.me_admission_poll_ms.max(1); @@ -88,9 +100,13 @@ pub(crate) async fn configure_admission_gate( fast_fallback_enabled = cfg.general.me2dc_fallback && cfg.general.me2dc_fast; continue; } - changed = me_ready_rx_gate.changed() => { + changed = me_ready_rx_gate.changed(), if me_ready_watch_open => { if changed.is_err() { - break; + me_ready_watch_open = false; + warn!( + watch_channel = "me_ready", + "Admission ME readiness watch closed; continuing periodic polling" + ); } } _ = tokio::time::sleep(Duration::from_millis(admission_poll_ms)) => {} diff --git a/src/maestro/admission/tests.rs b/src/maestro/admission/tests.rs new file mode 100644 index 0000000..572b268 --- /dev/null +++ b/src/maestro/admission/tests.rs @@ -0,0 +1,210 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +use tracing::{Event, Subscriber}; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::{Context, SubscriberExt}; + +use super::*; +use crate::transport::middle_proxy::admission_test_support::IdlePoolFixture; + +const POLL_MS: u64 = 20; + +struct ClosureWarnings(Arc); + +impl Layer for ClosureWarnings { + fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) { + struct ChannelField(bool); + impl tracing::field::Visit for ChannelField { + fn record_debug( + &mut self, + field: &tracing::field::Field, + _value: &dyn std::fmt::Debug, + ) { + self.0 |= field.name() == "watch_channel"; + } + } + if event.metadata().level() == &tracing::Level::WARN { + let mut visitor = ChannelField(false); + event.record(&mut visitor); + if visitor.0 { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + } +} + +struct GateFixture { + pool: IdlePoolFixture, + route: Arc, + admission: watch::Receiver, + config_tx: Option>>, + ready_tx: Option>, + config: Arc, + scope: RuntimeTaskScope, +} + +impl GateFixture { + async fn new() -> Self { + let pool = IdlePoolFixture::new(false).await; + let mut cfg = ProxyConfig::default(); + cfg.general.use_middle_proxy = true; + cfg.general.me2dc_fallback = true; + cfg.general.me2dc_fast = true; + cfg.general.me_admission_poll_ms = POLL_MS; + let config = Arc::new(cfg); + let (config_tx, config_rx) = watch::channel(config.clone()); + let (ready_tx, ready_rx) = watch::channel(0); + let (admission_tx, mut admission) = watch::channel(false); + let route = Arc::new(RouteRuntimeController::new(RelayRouteMode::Direct)); + let scope = RuntimeTaskScope::new(); + configure_admission_gate( + &config, + Some(pool.pool()), + Arc::new(RwLock::new(Some(pool.pool()))), + route.clone(), + &admission_tx, + config_rx, + ready_rx, + scope.clone(), + ) + .await; + drop(admission_tx); + admission.borrow_and_update(); + tokio::task::yield_now().await; + Self { + pool, + route, + admission, + config_tx: Some(config_tx), + ready_tx: Some(ready_tx), + config, + scope, + } + } + + async fn poll(&self) { + tokio::task::yield_now().await; + tokio::time::advance(Duration::from_millis(POLL_MS)).await; + tokio::task::yield_now().await; + } + + async fn expect_route(&self, mode: RelayRouteMode) { + let mut rx = self.route.subscribe(); + tokio::time::timeout(Duration::from_secs(1), async { + loop { + if rx.borrow_and_update().mode == mode { + return; + } + rx.changed().await.unwrap(); + } + }) + .await + .expect("the generation must keep sampling readiness without client traffic"); + } +} + +async fn recovery_after_watch_closure(close_config: bool, close_ready: bool) { + let warnings = Arc::new(AtomicUsize::new(0)); + let subscriber = tracing_subscriber::registry().with(ClosureWarnings(warnings.clone())); + let _default = tracing::subscriber::set_default(subscriber); + let mut fixture = GateFixture::new().await; + assert!(*fixture.admission.borrow()); + if close_config { + fixture.config_tx.take(); + } + if close_ready { + fixture.ready_tx.take(); + } + for _ in 0..4 { + tokio::task::yield_now().await; + } + let expected_warnings = usize::from(close_config) + usize::from(close_ready); + fixture.pool.set_ready(true); + for _ in 0..4 { + tokio::task::yield_now().await; + } + assert_eq!(fixture.route.snapshot().mode, RelayRouteMode::Direct); + fixture.poll().await; + fixture.expect_route(RelayRouteMode::Middle).await; + fixture.pool.set_ready(false); + fixture.poll().await; + fixture.expect_route(RelayRouteMode::Direct).await; + fixture.pool.set_ready(true); + fixture.poll().await; + fixture.expect_route(RelayRouteMode::Middle).await; + for _ in 0..5 { + fixture.poll().await; + } + assert_eq!(warnings.load(Ordering::Relaxed), expected_warnings); + fixture.scope.stop().await; + assert!(fixture.admission.has_changed().is_err()); +} + +#[tokio::test(start_paused = true)] +async fn idle_gate_recovers_after_readiness_watch_closes() { + recovery_after_watch_closure(false, true).await; +} + +#[tokio::test(start_paused = true)] +async fn idle_gate_recovers_after_config_watch_closes() { + recovery_after_watch_closure(true, false).await; +} + +#[tokio::test(start_paused = true)] +async fn idle_gate_recovers_after_both_watches_close_without_spinning() { + recovery_after_watch_closure(true, true).await; +} + +#[tokio::test(start_paused = true)] +async fn closed_config_watch_preserves_last_effective_fallback_policy() { + let mut fixture = GateFixture::new().await; + let mut cfg = fixture.config.as_ref().clone(); + cfg.general.me2dc_fallback = false; + fixture + .config_tx + .as_ref() + .unwrap() + .send_replace(Arc::new(cfg)); + fixture.poll().await; + fixture.expect_route(RelayRouteMode::Middle).await; + assert!(!*fixture.admission.borrow_and_update()); + fixture.config_tx.take(); + fixture.ready_tx.take(); + fixture.poll().await; + fixture.pool.set_ready(true); + fixture.poll().await; + tokio::time::timeout(Duration::from_secs(1), fixture.admission.changed()) + .await + .expect("restored readiness must open admission after notification loss") + .unwrap(); + assert!(*fixture.admission.borrow_and_update()); + fixture.pool.set_ready(false); + fixture.poll().await; + tokio::time::timeout(Duration::from_secs(1), fixture.admission.changed()) + .await + .unwrap() + .unwrap(); + assert!(!*fixture.admission.borrow_and_update()); + assert_eq!(fixture.route.snapshot().mode, RelayRouteMode::Middle); + fixture.scope.stop().await; +} + +#[tokio::test(start_paused = true)] +async fn generation_cancellation_stops_gate_after_notification_loss() { + let mut fixture = GateFixture::new().await; + fixture.config_tx.take(); + fixture.ready_tx.take(); + fixture.poll().await; + fixture.pool.set_ready(true); + fixture.poll().await; + fixture.expect_route(RelayRouteMode::Middle).await; + fixture.scope.stop().await; + assert!(fixture.admission.has_changed().is_err()); + let state = fixture.route.snapshot(); + fixture.pool.set_ready(false); + for _ in 0..3 { + fixture.poll().await; + } + assert_eq!(fixture.route.snapshot(), state); + assert!(*fixture.admission.borrow()); +} diff --git a/src/maestro/orchestrator.rs b/src/maestro/orchestrator.rs index c8d78fb..9aa61ba 100644 --- a/src/maestro/orchestrator.rs +++ b/src/maestro/orchestrator.rs @@ -375,17 +375,21 @@ pub(super) async fn run_telemt_core( #[cfg(target_os = "linux")] let conntrack_firewall = { - let authority = crate::conntrack_control::FirewallAuthority::spawn(&process_control_plane) - .map_err(std::io::Error::other)?; - if !authority - .publish_initial(1, runtime.config.clone(), stats.clone()) - .await + let authority = crate::conntrack_control::FirewallAuthority::spawn( + &process_control_plane, + &runtime.config, + ) + .map_err(std::io::Error::other)?; + if let Some(authority) = &authority + && !authority + .publish_initial(1, runtime.config.clone(), stats.clone()) + .await { warn!( "Initial conntrack firewall reconciliation failed; background retries remain active" ); } - Some(authority) + authority }; #[cfg(not(target_os = "linux"))] let conntrack_firewall = None::; diff --git a/src/service/mod.rs b/src/service/mod.rs index bb8e38b..a31758d 100644 --- a/src/service/mod.rs +++ b/src/service/mod.rs @@ -104,11 +104,12 @@ pub fn generate_service_file(init_system: InitSystem, opts: &ServiceOptions) -> InitSystem::Systemd => generate_systemd_unit(opts), InitSystem::OpenRC => generate_openrc_script(opts), InitSystem::FreeBSDRc => generate_freebsd_rc_script(opts), - InitSystem::Unknown => generate_systemd_unit(opts), // Default to systemd format + // Preserve systemd generation when init detection is unavailable. + InitSystem::Unknown => generate_systemd_unit(opts), } } -/// Generates an enhanced systemd unit file. +/// Generates a hardened systemd unit with Netlink access for optional firewall helpers. fn generate_systemd_unit(opts: &ServiceOptions) -> String { let user_line = opts.user.map(|u| format!("User={}", u)).unwrap_or_default(); let group_line = opts @@ -151,14 +152,14 @@ PrivateDevices=true ProtectKernelTunables=true ProtectKernelModules=true ProtectControlGroups=true -RestrictAddressFamilies=AF_INET AF_INET6 AF_UNIX +RestrictAddressFamilies=AF_INET AF_INET6 AF_UNIX AF_NETLINK RestrictNamespaces=true RestrictRealtime=true RestrictSUIDSGID=true MemoryDenyWriteExecute=true LockPersonality=true -# Allow binding to privileged ports and writing to specific paths +# Allow privileged port binding, optional firewall helpers, and runtime state writes AmbientCapabilities=CAP_NET_BIND_SERVICE CAP_NET_ADMIN CapabilityBoundingSet=CAP_NET_BIND_SERVICE CAP_NET_ADMIN ReadWritePaths=/etc/telemt /var/run /var/lib/telemt @@ -233,9 +234,11 @@ fn generate_freebsd_rc_script(opts: &ServiceOptions) -> String { # Add the following lines to /etc/rc.conf to enable telemt: # # telemt_enable="YES" -# telemt_config="/etc/telemt/config.toml" # optional -# telemt_user="telemt" # optional -# telemt_group="telemt" # optional +# The configuration path can be overridden in rc.conf. +# telemt_config="/etc/telemt/config.toml" +# The service user and group can be overridden in rc.conf. +# telemt_user="telemt" +# telemt_group="telemt" # . /etc/rc.subr @@ -355,6 +358,28 @@ mod tests { assert!(unit.contains("PIDFile=")); } + #[test] + fn systemd_unit_allows_netlink_without_removing_hardening() { + let unit = generate_systemd_unit(&ServiceOptions::default()); + let families = unit + .lines() + .find_map(|line| line.strip_prefix("RestrictAddressFamilies=")) + .unwrap(); + assert_eq!( + families.split_whitespace().collect::>(), + ["AF_INET", "AF_INET6", "AF_UNIX", "AF_NETLINK"], + ); + for directive in [ + "NoNewPrivileges=true", + "ProtectSystem=strict", + "RestrictNamespaces=true", + "AmbientCapabilities=CAP_NET_BIND_SERVICE CAP_NET_ADMIN", + "CapabilityBoundingSet=CAP_NET_BIND_SERVICE CAP_NET_ADMIN", + ] { + assert!(unit.lines().any(|line| line == directive), "{directive}"); + } + } + #[test] fn test_openrc_script_generation() { let opts = ServiceOptions::default(); diff --git a/src/transport/middle_proxy/health/recovery.rs b/src/transport/middle_proxy/health/recovery.rs index 709ca0b..b9455ce 100644 --- a/src/transport/middle_proxy/health/recovery.rs +++ b/src/transport/middle_proxy/health/recovery.rs @@ -1,5 +1,10 @@ use super::*; +// Outage regressions use loopback sockets and explicit family-state snapshots. +#[cfg(test)] +mod tests; + +/// Reports whether replacing an endpoint could disrupt a bound client. pub(super) fn has_bound_clients_on_endpoint( writer_ids: &[u64], bound_clients_by_writer: &HashMap, @@ -9,6 +14,7 @@ pub(super) fn has_bound_clients_on_endpoint( .any(|writer_id| bound_clients_by_writer.get(writer_id).copied().unwrap_or(0) > 0) } +/// Restores missing single-endpoint coverage under the bounded outage retry policy. pub(super) async fn recover_single_endpoint_outage( pool: &Arc, rng: &Arc, @@ -53,13 +59,21 @@ pub(super) async fn recover_single_endpoint_outage( pool.stats .increment_me_single_endpoint_outage_reconnect_attempt_total(); + let generation = pool.current_generation(); let bypass_quarantine = pool.single_endpoint_outage_disable_quarantine(); let attempt_ok = if bypass_quarantine { pool.stats .increment_me_single_endpoint_quarantine_bypass_total(); match tokio::time::timeout( pool.reconnect_runtime.me_one_timeout, - pool.connect_one_for_dc(endpoint, key.0, rng.as_ref()), + pool.connect_one_with_generation_contour_for_dc_with_intent( + endpoint, + rng.as_ref(), + generation, + WriterContour::Active, + key.0, + WriterOpenIntent::Coverage, + ), ) .await { @@ -88,7 +102,14 @@ pub(super) async fn recover_single_endpoint_outage( let one_endpoint = [endpoint]; match tokio::time::timeout( pool.reconnect_runtime.me_one_timeout, - pool.connect_endpoints_round_robin(key.0, &one_endpoint, rng.as_ref()), + pool.connect_endpoints_round_robin_with_generation_contour( + key.0, + &one_endpoint, + rng.as_ref(), + generation, + WriterContour::Active, + WriterOpenIntent::Coverage, + ), ) .await { @@ -142,6 +163,7 @@ pub(super) async fn recover_single_endpoint_outage( ); } +/// Rotates an idle single-endpoint shadow without removing the existing coverage first. pub(super) async fn maybe_rotate_single_endpoint_shadow( pool: &Arc, rng: &Arc, diff --git a/src/transport/middle_proxy/health/recovery/tests.rs b/src/transport/middle_proxy/health/recovery/tests.rs new file mode 100644 index 0000000..863240e --- /dev/null +++ b/src/transport/middle_proxy/health/recovery/tests.rs @@ -0,0 +1,148 @@ +use std::sync::atomic::Ordering; + +use super::*; +use crate::transport::middle_proxy::admission_test_support::unregistered_writer; +use crate::transport::middle_proxy::pool_writer_security_tests::make_pool; + +async fn single_endpoint_reaches_tcp_with_a_full_cap(bypass_quarantine: bool) { + let pool = make_pool().await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = listener.local_addr().unwrap(); + let donor = SocketAddr::new(endpoint.ip(), 1); + pool.update_proxy_maps( + HashMap::from([ + (2, vec![(endpoint.ip(), endpoint.port())]), + (3, vec![(donor.ip(), donor.port())]), + ]), + None, + ) + .await; + pool.floor_runtime + .me_adaptive_floor_cpu_cores_override + .store(1, Ordering::Relaxed); + pool.floor_runtime + .me_adaptive_floor_max_active_writers_per_core + .store(1, Ordering::Relaxed); + pool.floor_runtime + .me_adaptive_floor_max_active_writers_global + .store(1, Ordering::Relaxed); + pool.single_endpoint_runtime + .me_single_endpoint_outage_disable_quarantine + .store(bypass_quarantine, Ordering::Relaxed); + let (writer, _receiver) = unregistered_writer( + &pool, + 1, + 3, + donor, + pool.current_generation(), + WriterContour::Active, + ); + pool.writers.write().await.push(writer); + if bypass_quarantine { + pool.endpoint_quarantine + .lock() + .await + .insert(endpoint, Instant::now() + Duration::from_secs(60)); + } + let rng = Arc::new(SecureRandom::new()); + let key = (2, IpFamily::V4); + let mut backoff = HashMap::new(); + let mut next_attempt = HashMap::new(); + let semaphore = Arc::new(Semaphore::new(1)); + let required = pool.required_writers_for_dc(1); + // TCP acceptance proves capacity admission; closing the peer deliberately rejects the handshake. + tokio::time::timeout(Duration::from_secs(1), async { + tokio::join!( + recover_single_endpoint_outage( + &pool, + &rng, + key, + endpoint, + required, + &mut backoff, + &mut next_attempt, + &semaphore, + ), + async { + let (stream, _) = listener.accept().await.unwrap(); + drop(stream); + }, + ); + }) + .await + .expect("missing-DC recovery must reach TCP even when another DC fills the cap"); + assert_eq!( + pool.stats + .get_me_single_endpoint_outage_reconnect_attempt_total(), + 1 + ); + assert_eq!( + pool.stats + .get_me_single_endpoint_outage_reconnect_success_total(), + 0 + ); + assert_eq!( + pool.stats.get_me_single_endpoint_quarantine_bypass_total(), + u64::from(bypass_quarantine) + ); + assert!(next_attempt.contains_key(&key)); + assert!(backoff.contains_key(&key)); + assert_eq!(semaphore.available_permits(), 1); + assert_eq!( + pool.writer_connect_active_reserved.load(Ordering::Acquire), + 0 + ); +} + +#[tokio::test] +async fn single_endpoint_quarantine_bypass_recovers_coverage_at_cap() { + single_endpoint_reaches_tcp_with_a_full_cap(true).await; +} + +#[tokio::test] +async fn single_endpoint_round_robin_recovers_coverage_at_cap() { + single_endpoint_reaches_tcp_with_a_full_cap(false).await; +} + +#[tokio::test] +async fn restored_coverage_preserves_family_suppression_hysteresis() { + let pool = make_pool().await; + let now = MePool::now_epoch_secs(); + let until = now + 3600; + pool.set_family_runtime_state( + IpFamily::V4, + MeFamilyRuntimeState::Suppressed, + now, + until, + 5, + 0, + ); + super::super::update_family_runtime_state(&pool, IpFamily::V4, false); + assert_eq!( + pool.family_runtime_state(IpFamily::V4), + MeFamilyRuntimeState::Suppressed + ); + assert_eq!(pool.family_suppressed_until_epoch_secs(IpFamily::V4), until); + assert_eq!(pool.family_recover_success_streak(IpFamily::V4), 0); + + pool.set_family_runtime_state( + IpFamily::V4, + MeFamilyRuntimeState::Suppressed, + now, + now - 1, + 5, + 0, + ); + super::super::update_family_runtime_state(&pool, IpFamily::V4, false); + assert_eq!( + pool.family_runtime_state(IpFamily::V4), + MeFamilyRuntimeState::Recovering + ); + assert_eq!(pool.family_recover_success_streak(IpFamily::V4), 1); + super::super::update_family_runtime_state(&pool, IpFamily::V4, false); + assert_eq!( + pool.family_runtime_state(IpFamily::V4), + MeFamilyRuntimeState::Healthy + ); + assert_eq!(pool.family_fail_streak(IpFamily::V4), 0); +} diff --git a/src/transport/middle_proxy/mod.rs b/src/transport/middle_proxy/mod.rs index ead2fa2..968aa2d 100644 --- a/src/transport/middle_proxy/mod.rs +++ b/src/transport/middle_proxy/mod.rs @@ -49,6 +49,11 @@ mod send; mod send_adversarial_tests; mod wire; +/// Network-free writer fixtures shared by admission and recovery regressions. +#[cfg(test)] +#[path = "tests/admission_test_support.rs"] +pub(crate) mod admission_test_support; + use bytes::Bytes; use tokio::sync::OwnedSemaphorePermit; diff --git a/src/transport/middle_proxy/pool/writer_admission.rs b/src/transport/middle_proxy/pool/writer_admission.rs index 722eed3..a8afd76 100644 --- a/src/transport/middle_proxy/pool/writer_admission.rs +++ b/src/transport/middle_proxy/pool/writer_admission.rs @@ -61,13 +61,13 @@ impl Drop for WriterOpenReservation<'_> { } impl MePool { - /// Computes the authoritative active-writer floor across enabled families and DCs. + /// Computes reservation coverage across configured families, including suppressed ones. pub(in crate::transport::middle_proxy) async fn active_coverage_required_total(&self) -> usize { - let now_epoch_secs = Self::now_epoch_secs(); let mut required_total = 0usize; let endpoint_snapshot = self.endpoint_snapshot.load_full(); - if self.family_enabled_for_drain_coverage(IpFamily::V4, now_epoch_secs) { + // Suppression must not remove the capacity needed to recover active coverage. + if self.decision.ipv4_me { for addrs in endpoint_snapshot.map_v4.values() { let mut endpoints = HashSet::::new(); for (ip, port) in addrs.iter().copied() { @@ -79,7 +79,7 @@ impl MePool { } } - if self.family_enabled_for_drain_coverage(IpFamily::V6, now_epoch_secs) { + if self.decision.ipv6_me { for addrs in endpoint_snapshot.map_v6.values() { let mut endpoints = HashSet::::new(); for (ip, port) in addrs.iter().copied() { diff --git a/src/transport/middle_proxy/pool_writer.rs b/src/transport/middle_proxy/pool_writer.rs index a02c868..7bf97c0 100644 --- a/src/transport/middle_proxy/pool_writer.rs +++ b/src/transport/middle_proxy/pool_writer.rs @@ -33,6 +33,10 @@ mod publication; mod replacement; mod runtime; +// Recovery regressions exercise the serialized production publication boundary. +#[cfg(test)] +mod recovery_tests; + struct PreparedWriter<'a> { writer: MeWriter, tx: mpsc::Sender, diff --git a/src/transport/middle_proxy/pool_writer/publication.rs b/src/transport/middle_proxy/pool_writer/publication.rs index 5647f95..212ac45 100644 --- a/src/transport/middle_proxy/pool_writer/publication.rs +++ b/src/transport/middle_proxy/pool_writer/publication.rs @@ -84,7 +84,17 @@ impl MePool { crate::network::IpFamily::V6 }; let now_epoch_secs = Self::now_epoch_secs(); - if !self.family_enabled_for_drain_coverage(family, now_epoch_secs) { + let configured = match family { + crate::network::IpFamily::V4 => self.decision.ipv4_me, + crate::network::IpFamily::V6 => self.decision.ipv6_me, + }; + // Recovery coverage must remain publishable during suppression; the serialized + // per-group floor check below still bounds it and generation authority is unchanged. + let recovery_coverage = + contour == WriterContour::Active && intent == WriterOpenIntent::Coverage; + if !configured + || (self.is_family_temporarily_suppressed(family, now_epoch_secs) && !recovery_coverage) + { return Err(ProxyError::Proxy( "ME writer family lost publication authority".into(), )); diff --git a/src/transport/middle_proxy/pool_writer/recovery_tests.rs b/src/transport/middle_proxy/pool_writer/recovery_tests.rs new file mode 100644 index 0000000..ee7859f --- /dev/null +++ b/src/transport/middle_proxy/pool_writer/recovery_tests.rs @@ -0,0 +1,466 @@ +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; + +use super::*; +use crate::network::{IpFamily, probe::NetworkDecision}; +use crate::transport::middle_proxy::admission_test_support::unregistered_writer; +use crate::transport::middle_proxy::pool::MeFamilyRuntimeState; +use crate::transport::middle_proxy::pool_writer_security_tests::{ + make_pool, make_pool_with_decision, +}; + +fn cap_at_one(pool: &MePool) { + pool.floor_runtime + .me_adaptive_floor_cpu_cores_override + .store(1, Ordering::Relaxed); + pool.floor_runtime + .me_adaptive_floor_max_active_writers_per_core + .store(1, Ordering::Relaxed); + pool.floor_runtime + .me_adaptive_floor_max_active_writers_global + .store(1, Ordering::Relaxed); +} + +fn suppress(pool: &MePool, family: IpFamily) { + let now = MePool::now_epoch_secs(); + pool.set_family_runtime_state( + family, + MeFamilyRuntimeState::Suppressed, + now, + now + 3600, + 5, + 0, + ); +} + +fn prepared<'a>( + pool: &'a Arc, + writer: MeWriter, + intent: WriterOpenIntent, + reservation: WriterOpenReservation<'a>, + task_started: Arc, +) -> PreparedWriter<'a> { + let cancel = writer.cancel.clone(); + PreparedWriter { + tx: writer.tx.clone(), + byte_budget: writer.byte_budget.clone(), + writer, + task_registration: pool.lifecycle.try_register().unwrap(), + writer_task: Box::pin(async move { + task_started.store(true, Ordering::Release); + cancel.cancelled().await; + }), + intent, + _open_reservation: reservation, + } +} + +#[tokio::test] +async fn suppressed_active_coverage_publishes_only_the_missing_floor() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.update_proxy_maps(HashMap::from([(2, vec![(addr.ip(), addr.port())])]), None) + .await; + cap_at_one(&pool); + suppress(&pool, IpFamily::V4); + let required = pool.required_writers_for_dc(1); + let mut receivers = Vec::new(); + for id in 1..=required as u64 { + let reservation = pool + .reserve_writer_open(WriterContour::Active, WriterOpenIntent::Coverage, 2, addr) + .await + .expect("missing coverage must have a reservation despite suppression"); + let (writer, receiver) = unregistered_writer( + &pool, + id, + 2, + addr, + pool.current_generation(), + WriterContour::Active, + ); + receivers.push(receiver); + pool.publish_connected_writer(prepared( + &pool, + writer, + WriterOpenIntent::Coverage, + reservation, + Arc::new(AtomicBool::new(false)), + )) + .await + .expect("authenticated coverage must publish without waiting for suppression expiry"); + assert_eq!(pool.conn_count.load(Ordering::Relaxed), id as usize); + assert_eq!( + pool.writer_connect_active_reserved.load(Ordering::Acquire), + 0 + ); + } + assert!( + pool.reserve_writer_open(WriterContour::Active, WriterOpenIntent::Coverage, 2, addr) + .await + .is_none() + ); + let (extra, _receiver) = unregistered_writer( + &pool, + 99, + 2, + addr, + pool.current_generation(), + WriterContour::Active, + ); + assert!( + pool.authorize_writer_publication_capacity( + &extra, + WriterContour::Active, + WriterOpenIntent::Coverage, + &pool.writers.read().await, + ) + .is_err() + ); + assert_eq!( + pool.family_runtime_state(IpFamily::V4), + MeFamilyRuntimeState::Suppressed + ); + assert!(pool.shutdown_until(Duration::from_secs(1)).await); +} + +#[tokio::test] +async fn concurrent_suppressed_coverage_publications_cannot_overfill_a_group() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.update_proxy_maps(HashMap::from([(2, vec![(addr.ip(), addr.port())])]), None) + .await; + cap_at_one(&pool); + suppress(&pool, IpFamily::V4); + let required = pool.required_writers_for_dc(1); + let mut receivers = Vec::new(); + let mut publications = Vec::new(); + for id in 1..=(required + 1) as u64 { + let reservation = pool + .reserve_writer_open(WriterContour::Active, WriterOpenIntent::Coverage, 2, addr) + .await + .unwrap(); + let (writer, receiver) = unregistered_writer( + &pool, + id, + 2, + addr, + pool.current_generation(), + WriterContour::Active, + ); + receivers.push(receiver); + publications.push(pool.publish_connected_writer(prepared( + &pool, + writer, + WriterOpenIntent::Coverage, + reservation, + Arc::new(AtomicBool::new(false)), + ))); + } + let guard = pool.writers.write().await; + let results = futures::future::join_all(publications); + tokio::pin!(results); + tokio::select! { + biased; + results = &mut results => panic!("publication bypassed the writer barrier: {results:?}"), + _ = tokio::task::yield_now() => {} + } + drop(guard); + let results = results.await; + assert_eq!( + results.iter().filter(|result| result.is_ok()).count(), + required + ); + assert_eq!(results.iter().filter(|result| result.is_err()).count(), 1); + assert_eq!(pool.writers.read().await.len(), required); + assert_eq!(pool.conn_count.load(Ordering::Relaxed), required); + assert_eq!( + pool.registry.writer_idle_since_snapshot().await.len(), + required + ); + assert_eq!( + pool.writer_connect_active_reserved.load(Ordering::Acquire), + 0 + ); + assert!(pool.shutdown_until(Duration::from_secs(1)).await); +} + +#[tokio::test] +async fn committed_warm_writer_uses_resolved_active_coverage_authority() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.update_proxy_maps(HashMap::from([(2, vec![(addr.ip(), addr.port())])]), None) + .await; + suppress(&pool, IpFamily::V4); + let reservation = pool + .reserve_writer_open(WriterContour::Warm, WriterOpenIntent::Coverage, 2, addr) + .await + .unwrap(); + let (writer, _receiver) = unregistered_writer( + &pool, + 1, + 2, + addr, + pool.current_generation(), + WriterContour::Warm, + ); + pool.publish_connected_writer(prepared( + &pool, + writer, + WriterOpenIntent::Coverage, + reservation, + Arc::new(AtomicBool::new(false)), + )) + .await + .expect("a committed generation must resolve warm coverage to its active role"); + assert_eq!( + pool.writers.read().await[0].contour.load(Ordering::Acquire), + WriterContour::Active.as_u8() + ); + assert_eq!(pool.writer_connect_warm_reserved.load(Ordering::Acquire), 0); + assert!(pool.shutdown_until(Duration::from_secs(1)).await); +} + +#[tokio::test] +async fn suppressed_families_keep_their_full_coverage_reservation_budget() { + let pool = make_pool_with_decision(NetworkDecision { + ipv4_me: true, + ipv6_me: true, + ..NetworkDecision::default() + }) + .await; + let v4 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + let v6 = SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 443); + pool.update_proxy_maps( + HashMap::from([(2, vec![(v4.ip(), v4.port())])]), + Some(HashMap::from([(2, vec![(v6.ip(), v6.port())])])), + ) + .await; + cap_at_one(&pool); + suppress(&pool, IpFamily::V4); + suppress(&pool, IpFamily::V6); + let required = pool.required_writers_for_dc(1); + assert_eq!(pool.active_coverage_required_total().await, required * 4); + + let mut receivers = Vec::new(); + let mut writers = pool.writers.write().await; + for (index, (dc, addr)) in [(2, v4), (-2, v4), (-2, v6)].into_iter().enumerate() { + for offset in 0..required { + let (writer, receiver) = unregistered_writer( + &pool, + (index * required + offset + 1) as u64, + dc, + addr, + pool.current_generation(), + WriterContour::Active, + ); + writers.push(writer); + receivers.push(receiver); + } + } + drop(writers); + let reservation = pool + .reserve_writer_open(WriterContour::Active, WriterOpenIntent::Coverage, 2, v6) + .await + .expect("other covered groups must not consume the suppressed family's floor budget"); + assert_eq!( + pool.writer_connect_active_reserved.load(Ordering::Acquire), + 1 + ); + drop(reservation); + assert_eq!( + pool.writer_connect_active_reserved.load(Ordering::Acquire), + 0 + ); +} + +#[tokio::test] +async fn coverage_reservations_remain_bounded_and_release_on_drop() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.update_proxy_maps(HashMap::from([(2, vec![(addr.ip(), addr.port())])]), None) + .await; + cap_at_one(&pool); + suppress(&pool, IpFamily::V4); + let limit = pool.required_writers_for_dc(1) * 2 + + pool + .reconnect_runtime + .me_reconnect_max_concurrent_per_dc + .max(1) as usize; + let mut reservations = Vec::new(); + for _ in 0..limit { + reservations.push( + pool.reserve_writer_open(WriterContour::Active, WriterOpenIntent::Coverage, 2, addr) + .await + .expect("the configured coverage budget must remain available"), + ); + } + assert!( + pool.reserve_writer_open(WriterContour::Active, WriterOpenIntent::Coverage, 2, addr,) + .await + .is_none() + ); + assert_eq!( + pool.writer_connect_active_reserved.load(Ordering::Acquire), + limit + ); + drop(reservations); + assert_eq!( + pool.writer_connect_active_reserved.load(Ordering::Acquire), + 0 + ); +} + +#[tokio::test] +async fn suppression_exception_does_not_admit_growth_replacement_or_pending_warm() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.update_proxy_maps(HashMap::from([(2, vec![(addr.ip(), addr.port())])]), None) + .await; + suppress(&pool, IpFamily::V4); + for (contour, intent) in [ + (WriterContour::Active, WriterOpenIntent::Normal), + (WriterContour::Active, WriterOpenIntent::Replacement), + (WriterContour::Warm, WriterOpenIntent::Coverage), + (WriterContour::Warm, WriterOpenIntent::Normal), + ] { + let (writer, _receiver) = + unregistered_writer(&pool, 1, 2, addr, pool.current_generation(), contour); + assert!( + pool.authorize_writer_publication_capacity(&writer, contour, intent, &[]) + .is_err() + ); + } +} + +#[tokio::test] +async fn coverage_exception_never_enables_a_disabled_family() { + let pool = make_pool().await; + let v4 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + let v6 = SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 443); + pool.update_proxy_maps( + HashMap::from([(2, vec![(v4.ip(), v4.port())])]), + Some(HashMap::from([(2, vec![(v6.ip(), v6.port())])])), + ) + .await; + suppress(&pool, IpFamily::V6); + assert_eq!( + pool.active_coverage_required_total().await, + pool.required_writers_for_dc(1) * 2 + ); + let (writer, _receiver) = unregistered_writer( + &pool, + 1, + 2, + v6, + pool.current_generation(), + WriterContour::Active, + ); + assert!( + pool.authorize_writer_publication_capacity( + &writer, + WriterContour::Active, + WriterOpenIntent::Coverage, + &[], + ) + .is_err() + ); +} + +#[tokio::test] +async fn suppressed_coverage_cannot_publish_a_stale_generation() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.update_proxy_maps(HashMap::from([(2, vec![(addr.ip(), addr.port())])]), None) + .await; + suppress(&pool, IpFamily::V4); + let reservation = pool + .reserve_writer_open(WriterContour::Active, WriterOpenIntent::Coverage, 2, addr) + .await + .unwrap(); + let (writer, _receiver) = unregistered_writer( + &pool, + 1, + 2, + addr, + pool.current_generation().saturating_sub(1), + WriterContour::Active, + ); + let task_started = Arc::new(AtomicBool::new(false)); + let error = pool + .publish_connected_writer(prepared( + &pool, + writer, + WriterOpenIntent::Coverage, + reservation, + task_started.clone(), + )) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("generation lost publication authority") + ); + assert!(pool.writers.read().await.is_empty()); + assert!(pool.registry.writer_idle_since_snapshot().await.is_empty()); + assert_eq!(pool.conn_count.load(Ordering::Relaxed), 0); + assert_eq!( + pool.writer_connect_active_reserved.load(Ordering::Acquire), + 0 + ); + assert!(!task_started.load(Ordering::Acquire)); + assert!(pool.shutdown_until(Duration::from_secs(1)).await); +} + +#[tokio::test] +async fn suppressed_coverage_rechecks_endpoints_after_publication_waits() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.update_proxy_maps(HashMap::from([(2, vec![(addr.ip(), addr.port())])]), None) + .await; + suppress(&pool, IpFamily::V4); + let reservation = pool + .reserve_writer_open(WriterContour::Active, WriterOpenIntent::Coverage, 2, addr) + .await + .unwrap(); + let (writer, _receiver) = unregistered_writer( + &pool, + 1, + 2, + addr, + pool.current_generation(), + WriterContour::Active, + ); + let task_started = Arc::new(AtomicBool::new(false)); + let publication = pool.publish_connected_writer(prepared( + &pool, + writer, + WriterOpenIntent::Coverage, + reservation, + task_started.clone(), + )); + tokio::pin!(publication); + let guard = pool.writers.write().await; + tokio::select! { + biased; + result = &mut publication => panic!("publication bypassed the writer barrier: {result:?}"), + _ = tokio::task::yield_now() => {} + } + pool.update_proxy_maps(HashMap::from([(2, vec![(addr.ip(), 444)])]), None) + .await; + drop(guard); + let error = publication.await.unwrap_err(); + assert!( + error + .to_string() + .contains("target changed before publication") + ); + assert!(pool.writers.read().await.is_empty()); + assert!(pool.registry.writer_idle_since_snapshot().await.is_empty()); + assert_eq!(pool.conn_count.load(Ordering::Relaxed), 0); + assert_eq!( + pool.writer_connect_active_reserved.load(Ordering::Acquire), + 0 + ); + assert!(!task_started.load(Ordering::Acquire)); + assert!(pool.shutdown_until(Duration::from_secs(1)).await); +} diff --git a/src/transport/middle_proxy/tests/admission_test_support.rs b/src/transport/middle_proxy/tests/admission_test_support.rs new file mode 100644 index 0000000..1f5b6a6 --- /dev/null +++ b/src/transport/middle_proxy/tests/admission_test_support.rs @@ -0,0 +1,98 @@ +use std::collections::HashMap; +use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering}; +use std::time::Instant; + +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use super::codec::WriterCommand; +use super::pool::{MePool, MeWriter, WriterContour}; +use super::pool_writer_security_tests::make_pool; + +/// Owns idle writer visibility without spawning ME transports or client sessions. +pub(crate) struct IdlePoolFixture { + pool: Arc, + draining: Vec>, + _receivers: Vec>, +} + +impl IdlePoolFixture { + /// Builds both signed DC groups with a controllable readiness quorum. + pub(crate) async fn new(ready: bool) -> Self { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.update_proxy_maps(HashMap::from([(2, vec![(addr.ip(), addr.port())])]), None) + .await; + let mut draining = Vec::new(); + let mut receivers = Vec::new(); + { + let mut writers = pool.writers.write().await; + for (id, dc) in [(1, 2), (2, -2)] { + let (writer, receiver) = unregistered_writer( + &pool, + id, + dc, + addr, + pool.current_generation(), + WriterContour::Active, + ); + writer.draining.store(!ready, Ordering::Release); + draining.push(writer.draining.clone()); + receivers.push(receiver); + writers.push(writer); + } + } + Self { + pool, + draining, + _receivers: receivers, + } + } + + /// Returns the pool sampled by the production admission gate. + pub(crate) fn pool(&self) -> Arc { + self.pool.clone() + } + + /// Changes the fixture quorum without notifying the admission gate. + pub(crate) fn set_ready(&self, ready: bool) { + for draining in &self.draining { + draining.store(!ready, Ordering::Release); + } + } +} + +/// Builds a writer whose command receiver remains owned by the test. +pub(super) fn unregistered_writer( + pool: &Arc, + id: u64, + dc: i32, + addr: SocketAddr, + generation: u64, + contour: WriterContour, +) -> (MeWriter, mpsc::Receiver) { + let (tx, rx) = mpsc::channel(8); + ( + MeWriter { + id, + addr, + source_ip: addr.ip(), + writer_dc: dc, + generation, + contour: Arc::new(AtomicU8::new(contour.as_u8())), + created_at: Instant::now(), + tx, + byte_budget: pool.new_writer_byte_budget(), + cancel: CancellationToken::new(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + }, + rx, + ) +}