Conntrack: Firewall reconciliation process-owned&transactional

src/maestro/control_plane.rs

@@ -133,6 +133,21 @@ impl ProcessControlPlane {
         Ok(())
     }
 
+    /// Registers a cooperatively cancelled task whose cleanup future must finish.
+    pub(crate) fn spawn_cooperative<S, F>(&self, spawn: S) -> Result<(), S>
+    where
+        S: FnOnce(CancellationToken) -> F,
+        F: Future<Output = ()> + Send + 'static,
+    {
+        let Some(registration) = self.inner.admission.try_register() else {
+            return Err(spawn);
+        };
+        let cancellation = self.inner.cancellation.clone();
+        self.inner.tasks.spawn(spawn(cancellation));
+        drop(registration);
+        Ok(())
+    }
+
     /// Closes task admission, cancels all owned work, and joins it within the deadline.
     pub(crate) async fn shutdown(&self, timeout: Duration) -> bool {
         let deadline = tokio::time::Instant::now() + timeout;
@@ -254,4 +269,22 @@ mod tests {
         assert!(shutdown.await.unwrap());
         assert!(completed.load(Ordering::Acquire));
     }
+
+    #[tokio::test]
+    async fn cooperative_task_observes_cancellation_and_finishes_cleanup() {
+        let scope = ProcessControlPlane::new();
+        let completed = Arc::new(AtomicBool::new(false));
+        let completed_task = Arc::clone(&completed);
+        assert!(
+            scope
+                .spawn_cooperative(move |cancellation| async move {
+                    cancellation.cancelled().await;
+                    completed_task.store(true, Ordering::Release);
+                })
+                .is_ok()
+        );
+
+        assert!(scope.shutdown(Duration::from_secs(1)).await);
+        assert!(completed.load(Ordering::Acquire));
+    }
 }

src/maestro/orchestrator.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 
 use arc_swap::ArcSwap;
 use tokio::sync::{RwLock, Semaphore, watch};
-use tracing::{error, info};
+use tracing::{error, info, warn};
 
 use crate::api;
 use crate::ip_tracker::UserIpTracker;
@@ -40,7 +40,7 @@ pub(super) async fn run_telemt_core(
         process_started_at,
         process_started_at_epoch_secs,
         startup_tracker,
-        config,
+        mut config,
         config_path,
         has_rust_log,
         effective_log_level,
@@ -48,6 +48,13 @@ pub(super) async fn run_telemt_core(
         logging_guard: _logging_guard,
     } = bootstrap::bootstrap(privilege_drop_requested).await?;
 
+    if privilege_drop_requested && config.server.conntrack_control.inline_conntrack_control {
+        warn!(
+            "Inline conntrack control is disabled when process privileges are dropped"
+        );
+        config.server.conntrack_control.inline_conntrack_control = false;
+    }
+
     let quota_store = Arc::new(QuotaStore::default());
     let connection_authority = Arc::new(UserConnectionAuthority::default());
     let stats = Arc::new(Stats::with_process_authorities(
@@ -367,9 +374,26 @@ pub(super) async fn run_telemt_core(
         .await
         .map_err(std::io::Error::other)?;
 
+    #[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
+        {
+            warn!(
+                "Initial conntrack firewall reconciliation failed; background retries remain active"
+            );
+        }
+        Some(authority)
+    };
+    #[cfg(not(target_os = "linux"))]
+    let conntrack_firewall = None::<crate::conntrack_control::FirewallAuthority>;
+
     drop_after_bind();
 
-    runtime_tasks::spawn_metrics_if_configured(
+    if let Err(error) = runtime_tasks::spawn_metrics_if_configured(
         &runtime.config,
         &startup_tracker,
         active_runtime.clone(),
@@ -377,7 +401,15 @@ pub(super) async fn run_telemt_core(
         tls_full_cert_budget.clone(),
         process_control_plane.clone(),
     )
-    .await?;
+    .await
+    {
+        if let Some(conntrack_firewall) = &conntrack_firewall
+            && !conntrack_firewall.shutdown_and_clear().await
+        {
+            warn!("Conntrack firewall cleanup failed after metrics startup error");
+        }
+        return Err(error.into());
+    }
 
     runtime_watch_tx.send_replace(Some(active_runtime.load_full().watch_state()));
     active_runtime_tx.send_replace(Some(active_runtime.clone()));
@@ -401,6 +433,7 @@ pub(super) async fn run_telemt_core(
         runtime_watch_tx,
         listener_manager,
         web_trace,
+        conntrack_firewall.clone(),
     );
 
     shutdown::spawn_signal_handlers(
@@ -413,6 +446,7 @@ pub(super) async fn run_telemt_core(
         active_runtime,
         quota_state,
         reload_supervisor,
+        conntrack_firewall,
         process_control_plane,
     )
     .await;

src/maestro/reload_supervisor.rs

@@ -8,6 +8,7 @@ use tokio::sync::watch;
 use tokio_util::sync::CancellationToken;
 use tracing::{info, warn};
 
+use crate::conntrack_control::FirewallAuthority;
 use crate::stats::QuotaStore;
 use crate::tls_front::cache::TlsFullCertBudget;
 use crate::web::trace::WebTraceStore;
@@ -33,6 +34,7 @@ pub(crate) struct ReloadSupervisor {
     runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
     listener_manager: Arc<Mutex<ListenerManager>>,
     web_trace: Arc<WebTraceStore>,
+    conntrack_firewall: Option<FirewallAuthority>,
 }
 
 /// Process-owned handle that quiesces reloads before shutdown snapshots the runtime.
@@ -106,6 +108,7 @@ impl ReloadSupervisor {
         runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
         listener_manager: ListenerManager,
         web_trace: Arc<WebTraceStore>,
+        conntrack_firewall: Option<FirewallAuthority>,
     ) -> ReloadSupervisorHandle {
         let listener_manager = Arc::new(Mutex::new(listener_manager));
         let supervisor = Self {
@@ -120,6 +123,7 @@ impl ReloadSupervisor {
             runtime_watch_tx,
             listener_manager: listener_manager.clone(),
             web_trace,
+            conntrack_firewall,
         };
         let control = supervisor.control.clone();
         let shutdown = CancellationToken::new();
@@ -339,6 +343,14 @@ impl ReloadSupervisor {
             old_runtime.stop_accepting_sessions();
             listener_manager.activate_runtime_generation(new_runtime.clone())
         };
+        let conntrack_firewall_published = match &self.conntrack_firewall {
+            Some(conntrack_firewall) => conntrack_firewall.publish(
+                new_runtime.id,
+                new_runtime.config(),
+                new_runtime.stats.clone(),
+            ),
+            None => true,
+        };
         self.web_trace
             .apply_policy(new_runtime.id, &new_runtime.config().web.debug);
         config_watcher_activation.send_replace(true);
@@ -353,6 +365,12 @@ impl ReloadSupervisor {
             .apply_reload(&new_runtime.config().general.log_level);
         self.runtime_watch_tx
             .send_replace(Some(new_runtime.watch_state()));
+        if !conntrack_firewall_published {
+            let warning = "conntrack firewall reconciler is unavailable after runtime activation"
+                .to_string();
+            warn!(reload_id = command.reload_id, warning = %warning);
+            self.control.add_warning(command.reload_id, warning).await;
+        }
 
         info!(
             reload_id = command.reload_id,

src/maestro/reload_supervisor_tests.rs

@@ -61,6 +61,7 @@ async fn fixture(request: ReloadRequest) -> ReloadFixture {
         runtime_watch_tx,
         listener_manager,
         web_trace,
+        conntrack_firewall: None,
     });
     let command = ReloadCommand {
         reload_id: accepted.reload_id,
@@ -275,6 +276,7 @@ async fn quiesce_joins_idle_supervisor_and_rejects_later_submissions() {
             runtime.config().web.debug.clone(),
             &runtime.config().web.limits,
         ),
+        None,
     );
 
     tokio::time::timeout(Duration::from_secs(1), handle.quiesce())

src/maestro/runtime_build.rs

@@ -408,6 +408,12 @@ pub(crate) fn resolve_reload_config(
         fields.push("server.max_connections".to_string());
         effective.server.max_connections = old.server.max_connections;
     }
+    if serde_json::to_value(&old.server.conntrack_control).ok()
+        != serde_json::to_value(&desired.server.conntrack_control).ok()
+    {
+        fields.push("server.conntrack_control".to_string());
+        effective.server.conntrack_control = old.server.conntrack_control.clone();
+    }
     if old.general.direct_relay_buffer_budget_max_bytes
         != desired.general.direct_relay_buffer_budget_max_bytes
     {

src/maestro/runtime_build_tests.rs

@@ -127,6 +127,42 @@ fn process_wide_connection_and_direct_buffer_envelopes_are_restart_only() {
     assert!(!resolved.runtime_changed);
 }
 
+#[test]
+fn conntrack_control_policy_is_restart_only_as_one_process_owned_unit() {
+    let old = ProxyConfig::default();
+    let mut desired = old.clone();
+    desired.server.conntrack_control.inline_conntrack_control =
+        !old.server.conntrack_control.inline_conntrack_control;
+    desired.server.conntrack_control.mode = crate::config::ConntrackMode::Notrack;
+    desired.server.conntrack_control.backend = crate::config::ConntrackBackend::Iptables;
+    desired.server.conntrack_control.profile =
+        crate::config::ConntrackPressureProfile::Aggressive;
+    desired.server.conntrack_control.hybrid_listener_ips =
+        vec!["192.0.2.10".parse().unwrap()];
+    desired
+        .server
+        .conntrack_control
+        .pressure_high_watermark_pct = 90;
+    desired.server.conntrack_control.pressure_low_watermark_pct = 40;
+    desired.server.conntrack_control.delete_budget_per_sec = old
+        .server
+        .conntrack_control
+        .delete_budget_per_sec
+        .saturating_add(1);
+
+    let resolved = resolve_reload_config(&old, &desired).unwrap();
+
+    assert_eq!(
+        resolved.deferred_process_fields,
+        vec!["server.conntrack_control".to_string()]
+    );
+    assert_eq!(
+        serde_json::to_value(&resolved.effective.server.conntrack_control).unwrap(),
+        serde_json::to_value(&old.server.conntrack_control).unwrap()
+    );
+    assert!(!resolved.runtime_changed);
+}
+
 #[test]
 fn mixed_reload_retains_process_state_and_applies_runtime_state() {
     let old = ProxyConfig::default();

src/maestro/shutdown.rs

@@ -23,6 +23,7 @@ use super::control_plane::ProcessControlPlane;
 use super::generation::RuntimeGeneration;
 use super::helpers::{format_uptime, unit_label};
 use super::reload_supervisor::ReloadSupervisorHandle;
+use crate::conntrack_control::FirewallAuthority;
 use crate::quota_state::QuotaStateOwner;
 use crate::stats::Stats;
 use crate::synlimit_control;
@@ -54,6 +55,7 @@ pub(crate) async fn wait_for_shutdown(
     active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
     quota_state: Arc<QuotaStateOwner>,
     reload_supervisor: ReloadSupervisorHandle,
+    conntrack_firewall: Option<FirewallAuthority>,
     process_control_plane: ProcessControlPlane,
 ) {
     let signal = wait_for_shutdown_signal().await;
@@ -63,6 +65,7 @@ pub(crate) async fn wait_for_shutdown(
         active_runtime,
         quota_state,
         reload_supervisor,
+        conntrack_firewall,
         process_control_plane,
     )
     .await;
@@ -95,6 +98,7 @@ async fn perform_shutdown(
     active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
     quota_state: Arc<QuotaStateOwner>,
     reload_supervisor: ReloadSupervisorHandle,
+    conntrack_firewall: Option<FirewallAuthority>,
     process_control_plane: ProcessControlPlane,
 ) {
     let shutdown_started_at = Instant::now();
@@ -126,6 +130,12 @@ async fn perform_shutdown(
         warn!("ME shutdown: pool lifecycle deadline expired");
     }
 
+    if let Some(conntrack_firewall) = conntrack_firewall
+        && !conntrack_firewall.shutdown_and_clear().await
+    {
+        warn!("Conntrack firewall cleanup did not complete successfully");
+    }
+
     if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await {
         warn!(error = %error, "Failed to clear SYN limiter rules during shutdown");
     }