port gateway keepalive fix from stable

Cargo.toml

@@ -25,7 +25,7 @@ prost-types = "0.14"
 serde = { version = "1.0", features = ["derive"] }
 syslog = "7.0"
 thiserror = "2.0"
-tokio = { version = "1", features = ["fs", "macros", "rt-multi-thread", "signal", "sync", "time"] }
+tokio = { version = "1", features = ["fs", "macros", "process", "rt-multi-thread", "signal", "sync", "time"] }
 tokio-stream = { version = "0.1", features = [] }
 toml = { version = "1.0", default-features = false, features = [
     "parse",

src/bin/defguard-gateway.rs

@@ -7,6 +7,10 @@ use std::{
     process,
     sync::{Arc, Mutex},
 };
+#[cfg(not(unix))]
+use tokio::signal;
+#[cfg(unix)]
+use tokio::signal::unix::{SignalKind, signal as unix_signal};
 
 use defguard_gateway::{
     CORE_CLIENT_CERT_NAME, GRPC_CA_CERT_NAME, GRPC_CERT_NAME, GRPC_KEY_NAME, VERSION,
@@ -17,13 +21,34 @@ use defguard_gateway::{
     init_syslog,
     logging::init_tracing,
     server::run_http_server,
-    setup::run_setup,
 };
 use defguard_version::Version;
 #[cfg(not(any(target_os = "macos", target_os = "netbsd")))]
 use defguard_wireguard_rs::Kernel;
 use defguard_wireguard_rs::{Userspace, WGApi};
-use tokio::{sync::mpsc, task::JoinSet};
+use tokio::{
+    sync::{mpsc, watch},
+    task::JoinSet,
+};
+
+#[cfg(unix)]
+fn install_shutdown_signal()
+-> Result<impl std::future::Future<Output = std::io::Result<()>>, std::io::Error> {
+    let mut terminate = unix_signal(SignalKind::terminate())?;
+    let mut interrupt = unix_signal(SignalKind::interrupt())?;
+    Ok(async move {
+        tokio::select! {
+            _ = terminate.recv() => Ok(()),
+            _ = interrupt.recv() => Ok(()),
+        }
+    })
+}
+
+#[cfg(not(unix))]
+fn install_shutdown_signal()
+-> Result<impl std::future::Future<Output = std::io::Result<()>>, std::io::Error> {
+    Ok(signal::ctrl_c())
+}
 
 fn load_tls_config(cert_dir: &Path) -> Result<Option<TlsConfig>, GatewayError> {
     let grpc_cert = read_to_string(cert_dir.join(GRPC_CERT_NAME)).ok();
@@ -56,8 +81,10 @@ async fn main() -> Result<(), GatewayError> {
     // selection rustls panics at runtime when it cannot determine which to use.
     defguard_gateway::init_crypto_provider();
 
-    // parse config
+    // parse config and install signal handlers before starting setup or server tasks.
     let config = get_config()?;
+    let shutdown_signal = install_shutdown_signal()?;
+    tokio::pin!(shutdown_signal);
 
     // setup pidfile
     let pid = process::id();
@@ -135,7 +162,7 @@ async fn main() -> Result<(), GatewayError> {
         }
     };
 
-    // Keep track of spawned tasks.
+    // Keep track of background tasks.
     let mut tasks = JoinSet::new();
 
     // Optionally, launch HTTP server to report gateway's health.
@@ -151,44 +178,55 @@ async fn main() -> Result<(), GatewayError> {
     let gateway = Arc::new(Mutex::new(gateway));
     tasks.spawn(run_stats(Arc::clone(&gateway), config.stats_period()));
 
-    // Clone for later.
-    let post_down_clone = config.post_down.clone();
-
-    tasks.spawn(async move {
-        let tls_config = match maybe_tls_config {
-            None => {
-                log::info!(
-                    "gRPC TLS certificates not found in {}. They will be generated during setup.",
-                    config.cert_dir.display()
-                );
-                match run_setup(&config, Arc::clone(&logs_rx)).await {
-                    Ok(tls_config) => tls_config,
-                    Err(err) => {
-                        log::error!("Failed to run setup: {err}");
-                        return;
-                    }
-                }
+    let (shutdown_tx, shutdown_rx) = watch::channel(false);
+    let mut gateway_task = tokio::spawn(run_gateway_loop(
+        config.clone(),
+        gateway,
+        Arc::clone(&logs_rx),
+        maybe_tls_config,
+        shutdown_rx,
+    ));
+    let mut shutdown_error = None;
+
+    tokio::select! {
+        biased;
+        result = &mut shutdown_signal => {
+            if let Err(err) = result {
+                log::error!("Failed to listen for shutdown signal: {err}");
+                shutdown_error = Some(err);
             }
-            Some(tls_config) => {
-                log::info!(
-                    "Using existing gRPC TLS certificates from {}",
-                    config.cert_dir.display()
-                );
-                tls_config
+            let _ = shutdown_tx.send(true);
+            if let Err(err) = (&mut gateway_task).await {
+                log::error!("Gateway task failed: {err}");
             }
-        };
-
-        // Launch gRPC server (with purge-triggered setup loop).
-        run_gateway_loop(config, gateway, Arc::clone(&logs_rx), tls_config).await;
-    });
+        }
+        result = &mut gateway_task => {
+            if let Err(err) = result {
+                log::error!("Gateway task failed: {err}");
+            }
+        }
+        result = tasks.join_next() => {
+            if let Some(Err(err)) = result {
+                log::error!("Background task failed: {err}");
+            }
+            let _ = shutdown_tx.send(true);
+            if let Err(err) = (&mut gateway_task).await {
+                log::error!("Gateway task failed: {err}");
+            }
+        }
+    }
 
-    // Wait for the first task to finish.
-    let _result = tasks.join_next().await;
+    tasks.abort_all();
+    while tasks.join_next().await.is_some() {}
 
-    if let Some(post_down) = &post_down_clone {
+    if let Some(post_down) = &config.post_down {
         log::info!("Executing specified POST_DOWN command: {post_down}");
         execute_command(post_down)?;
     }
 
+    if let Some(err) = shutdown_error {
+        return Err(err.into());
+    }
+
     Ok(())
 }

src/config.rs

@@ -1,6 +1,7 @@
 use std::{
     fs,
     net::{IpAddr, Ipv4Addr, SocketAddr},
+    num::NonZeroU64,
     path::PathBuf,
     time::Duration,
 };
@@ -23,10 +24,21 @@ fn default_stats_period() -> u64 {
     30
 }
 
+const DEFAULT_GRPC_KEEPALIVE_INTERVAL_SECS: NonZeroU64 = NonZeroU64::new(10).unwrap();
+const DEFAULT_GRPC_KEEPALIVE_TIMEOUT_SECS: NonZeroU64 = NonZeroU64::new(20).unwrap();
+
 fn default_core_disconnect_grace_period() -> u64 {
     30
 }
 
+fn default_grpc_keepalive_interval() -> NonZeroU64 {
+    DEFAULT_GRPC_KEEPALIVE_INTERVAL_SECS
+}
+
+fn default_grpc_keepalive_timeout() -> NonZeroU64 {
+    DEFAULT_GRPC_KEEPALIVE_TIMEOUT_SECS
+}
+
 fn default_ifname() -> String {
     String::from("wg0")
 }
@@ -84,6 +96,24 @@ pub struct Config {
     #[serde(default = "default_core_disconnect_grace_period")]
     pub core_disconnect_grace_period: u64,
 
+    /// Seconds between HTTP/2 keepalive pings to Defguard Core.
+    #[arg(
+        long,
+        env = "DEFGUARD_GRPC_KEEPALIVE_INTERVAL",
+        default_value_t = DEFAULT_GRPC_KEEPALIVE_INTERVAL_SECS
+    )]
+    #[serde(default = "default_grpc_keepalive_interval")]
+    pub grpc_keepalive_interval: NonZeroU64,
+
+    /// Seconds to wait for an HTTP/2 keepalive response before closing the connection.
+    #[arg(
+        long,
+        env = "DEFGUARD_GRPC_KEEPALIVE_TIMEOUT",
+        default_value_t = DEFAULT_GRPC_KEEPALIVE_TIMEOUT_SECS
+    )]
+    #[serde(default = "default_grpc_keepalive_timeout")]
+    pub grpc_keepalive_timeout: NonZeroU64,
+
     /// Network interface name (e.g. wg0)
     #[arg(long, short = 'i', env = "DEFGUARD_IFNAME", default_value = "wg0")]
     #[serde(default = "default_ifname")]
@@ -183,6 +213,16 @@ impl Config {
         Duration::from_secs(self.stats_period)
     }
 
+    #[must_use]
+    pub fn grpc_keepalive_interval(&self) -> Duration {
+        Duration::from_secs(self.grpc_keepalive_interval.get())
+    }
+
+    #[must_use]
+    pub fn grpc_keepalive_timeout(&self) -> Duration {
+        Duration::from_secs(self.grpc_keepalive_timeout.get())
+    }
+
     #[must_use]
     pub fn adoption_timeout(&self) -> Duration {
         Duration::from_secs(self.adoption_timeout * 60)
@@ -203,6 +243,8 @@ impl Default for Config {
             userspace: false,
             stats_period: 15,
             core_disconnect_grace_period: 30,
+            grpc_keepalive_interval: default_grpc_keepalive_interval(),
+            grpc_keepalive_timeout: default_grpc_keepalive_timeout(),
             ifname: "wg0".into(),
             pidfile: None,
             use_syslog: false,
@@ -269,3 +311,52 @@ fn test_validate_ifname() {
     assert!(validate_ifname("wg-16-chars-bad_").is_err());
     assert!(validate_ifname("wgMFA-looong-naame").is_err());
 }
+
+#[test]
+fn grpc_keepalive_settings_accept_cli_and_toml_values() {
+    use clap::CommandFactory;
+
+    let defaults = Config::default();
+    assert_eq!(defaults.grpc_keepalive_interval.get(), 10);
+    assert_eq!(defaults.grpc_keepalive_timeout.get(), 20);
+
+    let command = Config::command();
+    let interval = command
+        .get_arguments()
+        .find(|arg| arg.get_id().as_str() == "grpc_keepalive_interval")
+        .unwrap();
+    assert_eq!(
+        interval.get_env().unwrap().to_str(),
+        Some("DEFGUARD_GRPC_KEEPALIVE_INTERVAL")
+    );
+    let timeout = command
+        .get_arguments()
+        .find(|arg| arg.get_id().as_str() == "grpc_keepalive_timeout")
+        .unwrap();
+    assert_eq!(
+        timeout.get_env().unwrap().to_str(),
+        Some("DEFGUARD_GRPC_KEEPALIVE_TIMEOUT")
+    );
+
+    let cli = Config::try_parse_from([
+        "defguard-gateway",
+        "--grpc-keepalive-interval",
+        "15",
+        "--grpc-keepalive-timeout",
+        "25",
+    ])
+    .unwrap();
+    assert_eq!(cli.grpc_keepalive_interval.get(), 15);
+    assert_eq!(cli.grpc_keepalive_timeout.get(), 25);
+    assert!(
+        Config::try_parse_from(["defguard-gateway", "--grpc-keepalive-interval", "0"]).is_err()
+    );
+    assert!(Config::try_parse_from(["defguard-gateway", "--grpc-keepalive-timeout", "0"]).is_err());
+
+    let file: Config =
+        toml::from_str("grpc_keepalive_interval = 12\ngrpc_keepalive_timeout = 24").unwrap();
+    assert_eq!(file.grpc_keepalive_interval.get(), 12);
+    assert_eq!(file.grpc_keepalive_timeout.get(), 24);
+    assert!(toml::from_str::<Config>("grpc_keepalive_interval = 0").is_err());
+    assert!(toml::from_str::<Config>("grpc_keepalive_timeout = 0").is_err());
+}