port keepalive fix from stable

Cargo.toml

@@ -24,7 +24,7 @@ axum-server = { version = "0.8", features = ["tls-rustls"] }
 # match axum-extra -> cookies
 cookie = "0.18.1"
 time = { version = "0.3", default-features = false }
-tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "time"] }
+tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal", "sync", "time"] }
 tokio-stream = "0.1"
 tower-http = { version = "0.7", features = ["fs", "trace", "timeout"] }
 # logging/tracing
@@ -57,7 +57,7 @@ mime_guess = "2.0"
 base64 = "0.23"
 tower = "0.5"
 futures-util = "0.3"
-ammonia = "4.1"
+ammonia = "4.2.4"
 chrono = "0.4"
 rustls = { version = "0.23", default-features = false, features = [
     "aws-lc-rs",
@@ -66,7 +66,7 @@ instant-acme = { version = "0.8", features = ["hyper-rustls", "aws-lc-rs"] }
 reqwest = { version = "0.12", default-features = false, features = ["rustls-tls-native-roots", "json"] }
 
 [dev-dependencies]
-tokio = { version = "1", features = ["net"] }
+tokio = { version = "1", features = ["io-util", "net"] }
 
 [build-dependencies]
 tonic-prost-build = "0.14"

example-config.toml

@@ -6,6 +6,10 @@ http_port = 8080
 # port the gRPC server will listen on
 grpc_port = 50051
 
+# HTTP/2 keepalive ping interval and timeout for the gRPC server, in seconds
+grpc_keepalive_interval_seconds = 10
+grpc_keepalive_timeout_seconds = 20
+
 log_level = "info"
 rate_limit_per_second = 0
 rate_limit_burst = 0

src/config.rs

@@ -1,4 +1,4 @@
-use std::{fs::read_to_string, net::IpAddr, path::PathBuf, time::Duration};
+use std::{fs::read_to_string, net::IpAddr, num::NonZeroU64, path::PathBuf, time::Duration};
 
 use clap::Parser;
 use log::LevelFilter;
@@ -28,6 +28,17 @@ fn default_adoption_timeout() -> u64 {
     10
 }
 
+pub(crate) const DEFAULT_GRPC_KEEPALIVE_INTERVAL_SECONDS: NonZeroU64 = NonZeroU64::new(10).unwrap();
+pub(crate) const DEFAULT_GRPC_KEEPALIVE_TIMEOUT_SECONDS: NonZeroU64 = NonZeroU64::new(20).unwrap();
+
+fn default_grpc_keepalive_interval_seconds() -> NonZeroU64 {
+    DEFAULT_GRPC_KEEPALIVE_INTERVAL_SECONDS
+}
+
+fn default_grpc_keepalive_timeout_seconds() -> NonZeroU64 {
+    DEFAULT_GRPC_KEEPALIVE_TIMEOUT_SECONDS
+}
+
 fn default_rate_limit_per_second() -> u64 {
     0
 }
@@ -54,6 +65,24 @@ pub struct EnvConfig {
     #[serde(default = "default_grpc_port")]
     pub grpc_port: u16,
 
+    /// HTTP/2 keepalive ping interval for the gRPC server, in seconds.
+    #[arg(
+        long,
+        env = "DEFGUARD_PROXY_GRPC_KEEPALIVE_INTERVAL_SECONDS",
+        default_value_t = DEFAULT_GRPC_KEEPALIVE_INTERVAL_SECONDS
+    )]
+    #[serde(default = "default_grpc_keepalive_interval_seconds")]
+    pub grpc_keepalive_interval_seconds: NonZeroU64,
+
+    /// Time to wait for an HTTP/2 keepalive ping acknowledgement, in seconds.
+    #[arg(
+        long,
+        env = "DEFGUARD_PROXY_GRPC_KEEPALIVE_TIMEOUT_SECONDS",
+        default_value_t = DEFAULT_GRPC_KEEPALIVE_TIMEOUT_SECONDS
+    )]
+    #[serde(default = "default_grpc_keepalive_timeout_seconds")]
+    pub grpc_keepalive_timeout_seconds: NonZeroU64,
+
     #[arg(long, env = "DEFGUARD_PROXY_LOG_LEVEL", default_value_t = LevelFilter::Info)]
     #[serde(default = "default_log_level")]
     pub log_level: LevelFilter,
@@ -127,16 +156,109 @@ pub enum ConfigError {
 }
 
 pub fn get_env_config() -> Result<EnvConfig, ConfigError> {
-    // parse CLI arguments to get config file path
     let cli_config = EnvConfig::parse();
 
-    // load config from file if one was specified
-    if let Some(config_path) = cli_config.config_path {
+    let config = if let Some(config_path) = cli_config.config_path {
         info!("Reading configuration from file: {config_path:?}");
         let config_toml = read_to_string(config_path)?;
-        let file_config: EnvConfig = toml::from_str(&config_toml)?;
-        Ok(file_config)
+        toml::from_str(&config_toml)?
     } else {
-        Ok(cli_config)
+        cli_config
+    };
+
+    Ok(config)
+}
+
+#[cfg(test)]
+mod tests {
+    use clap::Parser;
+
+    use super::{
+        DEFAULT_GRPC_KEEPALIVE_INTERVAL_SECONDS, DEFAULT_GRPC_KEEPALIVE_TIMEOUT_SECONDS, EnvConfig,
+    };
+
+    #[test]
+    fn keepalive_defaults_and_toml_values_parse() {
+        let defaults: EnvConfig = toml::from_str("").expect("empty TOML should use defaults");
+        assert_eq!(
+            defaults.grpc_keepalive_interval_seconds.get(),
+            DEFAULT_GRPC_KEEPALIVE_INTERVAL_SECONDS.get()
+        );
+        assert_eq!(
+            defaults.grpc_keepalive_timeout_seconds.get(),
+            DEFAULT_GRPC_KEEPALIVE_TIMEOUT_SECONDS.get()
+        );
+
+        let configured: EnvConfig = toml::from_str(
+            "grpc_keepalive_interval_seconds = 15\ngrpc_keepalive_timeout_seconds = 30",
+        )
+        .expect("keepalive TOML values should parse");
+        assert_eq!(configured.grpc_keepalive_interval_seconds.get(), 15);
+        assert_eq!(configured.grpc_keepalive_timeout_seconds.get(), 30);
+    }
+
+    #[test]
+    fn keepalive_cli_options_parse_seconds() {
+        let config = EnvConfig::try_parse_from([
+            "edge",
+            "--grpc-keepalive-interval-seconds",
+            "15",
+            "--grpc-keepalive-timeout-seconds",
+            "30",
+        ])
+        .expect("keepalive CLI options should parse");
+        assert_eq!(config.grpc_keepalive_interval_seconds.get(), 15);
+        assert_eq!(config.grpc_keepalive_timeout_seconds.get(), 30);
+    }
+
+    #[test]
+    fn keepalive_environment_variables_are_loaded_in_child() {
+        let output = std::process::Command::new(
+            std::env::current_exe().expect("test executable path should exist"),
+        )
+        .args([
+            "--exact",
+            "config::tests::keepalive_environment_variables_child",
+            "--nocapture",
+        ])
+        .env("DEFGUARD_TEST_KEEPALIVE_ENV_CHILD", "1")
+        .env("DEFGUARD_PROXY_GRPC_KEEPALIVE_INTERVAL_SECONDS", "13")
+        .env("DEFGUARD_PROXY_GRPC_KEEPALIVE_TIMEOUT_SECONDS", "27")
+        .output()
+        .expect("environment test child should start");
+        assert!(
+            output.status.success(),
+            "child test failed:\n{}\n{}",
+            String::from_utf8_lossy(&output.stdout),
+            String::from_utf8_lossy(&output.stderr)
+        );
+    }
+
+    #[test]
+    fn keepalive_environment_variables_child() {
+        if std::env::var_os("DEFGUARD_TEST_KEEPALIVE_ENV_CHILD").is_none() {
+            return;
+        }
+
+        let config = EnvConfig::try_parse_from(["edge"]).expect("environment values should parse");
+        assert_eq!(config.grpc_keepalive_interval_seconds.get(), 13);
+        assert_eq!(config.grpc_keepalive_timeout_seconds.get(), 27);
+    }
+
+    #[test]
+    fn keepalive_parsers_reject_zero() {
+        for toml in [
+            "grpc_keepalive_interval_seconds = 0",
+            "grpc_keepalive_timeout_seconds = 0",
+        ] {
+            assert!(toml::from_str::<EnvConfig>(toml).is_err());
+        }
+
+        for args in [
+            ["edge", "--grpc-keepalive-interval-seconds", "0"],
+            ["edge", "--grpc-keepalive-timeout-seconds", "0"],
+        ] {
+            assert!(EnvConfig::try_parse_from(args).is_err());
+        }
     }
 }

src/grpc.rs

@@ -7,6 +7,7 @@ use std::{
         Arc, Mutex, RwLock,
         atomic::{AtomicBool, AtomicU64, Ordering},
     },
+    time::Duration,
 };
 
 use axum_extra::extract::cookie::Key;
@@ -65,6 +66,20 @@ fn resolve_response(results: &RwLock<ResultMap>, id: u64, payload: core_response
     true
 }
 
+fn remove_disconnected_client(
+    clients: &RwLock<ClientMap>,
+    connected: &AtomicBool,
+    address: SocketAddr,
+) {
+    let mut clients = clients
+        .write()
+        .expect("Failed to acquire lock on clients hashmap when removing disconnected client");
+    clients.remove(&address);
+    if clients.is_empty() {
+        connected.store(false, Ordering::Relaxed);
+    }
+}
+
 #[derive(Clone)]
 pub(crate) struct PublicSettingsState {
     pub(crate) display_password_reset: Arc<AtomicBool>,
@@ -205,7 +220,13 @@ impl ProxyServer {
         lock.clone()
     }
 
-    pub(crate) async fn run<F>(self, addr: SocketAddr, shutdown: F) -> Result<(), anyhow::Error>
+    pub(crate) async fn run<F>(
+        self,
+        addr: SocketAddr,
+        keepalive_interval: Duration,
+        keepalive_timeout: Duration,
+        shutdown: F,
+    ) -> Result<(), anyhow::Error>
     where
         F: Future<Output = ()> + Send + 'static,
     {
@@ -224,7 +245,10 @@ impl ProxyServer {
             &tls_config.grpc_key_pem,
             &tls_config.grpc_ca_cert_pem,
         )?;
-        let mut builder = Server::builder().tls_config(tls_config)?;
+        let mut builder = Server::builder()
+            .http2_keepalive_interval(Some(keepalive_interval))
+            .http2_keepalive_timeout(Some(keepalive_timeout))
+            .tls_config(tls_config)?;
 
         let own_version = Version::parse(VERSION)?;
         let versioned_service = ServiceBuilder::new()
@@ -454,14 +478,7 @@ impl proxy_server::Proxy for ProxyServer {
                     }
                 }
                 info!("Defguard core client disconnected: {address}");
-                connected.store(false, Ordering::Relaxed);
-                clients
-                    .write()
-                    .expect(
-                        "Failed to acquire lock on clients hashmap when removing \
-                    disconnected client",
-                    )
-                    .remove(&address);
+                remove_disconnected_client(&clients, &connected, address);
             }
             .instrument(tracing::Span::current()),
         );
@@ -709,6 +726,29 @@ mod tests {
         })
     }
 
+    #[test]
+    fn test_connected_stays_true_until_last_client_disconnects() {
+        let clients = RwLock::new(ClientMap::new());
+        let connected = AtomicBool::new(true);
+        let first_address = SocketAddr::from(([127, 0, 0, 1], 1));
+        let second_address = SocketAddr::from(([127, 0, 0, 1], 2));
+        let (first_tx, _first_rx) = mpsc::unbounded_channel();
+        let (second_tx, _second_rx) = mpsc::unbounded_channel();
+        {
+            let mut clients = clients.write().expect("clients lock poisoned");
+            clients.insert(first_address, first_tx);
+            clients.insert(second_address, second_tx);
+        }
+
+        remove_disconnected_client(&clients, &connected, first_address);
+        assert!(connected.load(Ordering::Relaxed));
+        assert_eq!(clients.read().expect("clients lock poisoned").len(), 1);
+
+        remove_disconnected_client(&clients, &connected, second_address);
+        assert!(!connected.load(Ordering::Relaxed));
+        assert!(clients.read().expect("clients lock poisoned").is_empty());
+    }
+
     /// `send` must stamp the same id into the outgoing `CoreRequest` and into the `results` key,
     /// or the response loop looks up an entry that was never registered.
     #[test]