diff --git a/Cargo.lock b/Cargo.lock index f995cb75..ba97ba55 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -979,9 +979,9 @@ dependencies = [ [[package]] name = "ctutils" -version = "0.4.2" +version = "0.4.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +checksum = "38d28fb6aa7dd680b3ffd1a3e64e115c177bdba63c0ae7f9cd48129ff05c4fa0" dependencies = [ "cmov", ] @@ -4182,9 +4182,9 @@ checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" [[package]] name = "rustls-webpki" -version = "0.103.15" +version = "0.103.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +checksum = "4443b1b6d39d7c663139942cc4558262562d2c84123bb6282ef177009ca8b281" dependencies = [ "aws-lc-rs", "ring", @@ -5830,9 +5830,9 @@ dependencies = [ [[package]] name = "webpki-root-certs" -version = "1.0.9" +version = "1.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b96554aa2acc8ccdb7e1c9a58a7a68dd5d13bccc69cd124cb09406db612a1c9b" +checksum = "e49d0e65abc00c142e8f9428da991da9f1a250c39be97459901eb12b9b49f62a" dependencies = [ "rustls-pki-types", ] diff --git a/Cargo.toml b/Cargo.toml index 230d8c4f..ad92a1ad 100644 --- a/Cargo.toml +++ b/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" diff --git a/example-config.toml b/example-config.toml index 066b626f..06c92e79 100644 --- a/example-config.toml +++ b/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 diff --git a/flake.lock b/flake.lock index ba0e2d8a..3e1f1eff 100644 --- a/flake.lock +++ b/flake.lock @@ -5,11 +5,11 @@ "systems": "systems" }, "locked": { - "lastModified": 1731533236, - "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", + "lastModified": 1791914995, + "narHash": "sha256-Zf/xfQEAasnwnmEfHqW/4hjDE/s44yn2hD4RuVIZRjA=", "owner": "numtide", "repo": "flake-utils", - "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", + "rev": "b7a79d3fea1e214f6d668d521cecec0c01e26544", "type": "github" }, "original": { @@ -20,11 +20,11 @@ }, "nixpkgs": { "locked": { - "lastModified": 1790822859, - "narHash": "sha256-69xHQhAeMAD2wDXO7T2pcOZIF9Sga2W+JkmY2a11Ops=", + "lastModified": 1792502884, + "narHash": "sha256-IoNisKcGbUjdd45qCqywGMCYmZ5EnkxezvLII0FWwNA=", "owner": "NixOS", "repo": "nixpkgs", - "rev": "c59305bab2065cfecc4944690d9eedbb56f3a9fa", + "rev": "02c087f3fe673c9e317f206a627cb6b910336dbe", "type": "github" }, "original": { @@ -48,11 +48,11 @@ ] }, "locked": { - "lastModified": 1790928989, - "narHash": "sha256-5r3MN/Rx7N9tLapCBmUq6/zCc3C9J7gNif1HybD1fJQ=", + "lastModified": 1792739799, + "narHash": "sha256-M2rr/yPbTHvgAzCyZ+zHJHflFM0XU7t42Sw2skA3TQM=", "owner": "oxalica", "repo": "rust-overlay", - "rev": "368fee9beaab04ca6fe7af28db63caa9badb22fa", + "rev": "b0829a2cc9f7b596583331d78a942c1e5170d8a0", "type": "github" }, "original": { diff --git a/src/config.rs b/src/config.rs index 0a0d6375..61966026 100644 --- a/src/config.rs +++ b/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 { - // 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::(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()); + } } } diff --git a/src/grpc.rs b/src/grpc.rs index 93ffc097..21b7ba29 100644 --- a/src/grpc.rs +++ b/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, id: u64, payload: core_response true } +fn remove_disconnected_client( + clients: &RwLock, + 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, @@ -205,7 +220,13 @@ impl ProxyServer { lock.clone() } - pub(crate) async fn run(self, addr: SocketAddr, shutdown: F) -> Result<(), anyhow::Error> + pub(crate) async fn run( + self, + addr: SocketAddr, + keepalive_interval: Duration, + keepalive_timeout: Duration, + shutdown: F, + ) -> Result<(), anyhow::Error> where F: Future + 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] diff --git a/src/http.rs b/src/http.rs index b51893a6..d9d77c4e 100644 --- a/src/http.rs +++ b/src/http.rs @@ -1,6 +1,7 @@ #[cfg(unix)] use std::{fs::Permissions, os::unix::fs::PermissionsExt}; use std::{ + future::Future, io::ErrorKind, net::{IpAddr, Ipv4Addr, SocketAddr}, path::Path, @@ -31,7 +32,7 @@ use cookie::CookieBuilder; use defguard_version::{Version, server::DefguardVersionLayer}; use serde::Serialize; use tokio::{ - sync::{broadcast, mpsc, oneshot}, + sync::{broadcast, mpsc, oneshot, watch}, task::JoinSet, }; use tower_governor::{ @@ -62,10 +63,71 @@ pub(crate) const PASSWORD_RESET_COOKIE_PATH: &str = "/api/v1/password-reset"; const DEFGUARD_CORE_CONNECTED_HEADER: &str = "defguard-core-connected"; const DEFGUARD_CORE_VERSION_HEADER: &str = "defguard-core-version"; const RATE_LIMITER_CLEANUP_PERIOD: Duration = Duration::from_secs(60); +// Drain active work within Docker's default 10-second stop window. +const SERVER_SHUTDOWN_GRACE_PERIOD: Duration = Duration::from_secs(5); +const SERVER_SHUTDOWN_FORCE_TIMEOUT: Duration = Duration::from_secs(6); +const SERVER_SHUTDOWN_DEADLINE: Duration = Duration::from_secs(7); const X_FORWARDED_FOR: &str = "x-forwarded-for"; /// Default request body size limit applied globally to every route. const REQUEST_BODY_LIMIT: usize = 256 * 1024; // 256 KB +async fn wait_for_server(server: &mut F, timeout: Duration) -> anyhow::Result +where + F: Future> + Unpin, +{ + match tokio::time::timeout(timeout, server).await { + Ok(result) => { + result?; + Ok(true) + } + Err(_) => Ok(false), + } +} + +enum ServerWaitOutcome { + Finished, + Shutdown, + TimedOut, +} + +async fn receive_https_cert_update(receiver: &mut broadcast::Receiver) -> Option { + loop { + match receiver.recv().await { + Err(broadcast::error::RecvError::Lagged(_)) => { + warn!("Missed HTTP server update; waiting for the latest one"); + } + Ok(update) => return Some(update), + Err(broadcast::error::RecvError::Closed) => return None, + } + } +} + +async fn wait_for_server_or_shutdown( + server: &mut F, + shutdown_rx: &mut watch::Receiver, + restart_timeout: Duration, +) -> anyhow::Result +where + F: Future> + Unpin, +{ + tokio::select! { + result = &mut *server => { + result?; + Ok(ServerWaitOutcome::Finished) + } + _ = shutdown_rx.changed() => { + if !wait_for_server(server, SERVER_SHUTDOWN_FORCE_TIMEOUT).await? { + warn!("Server did not stop before the shutdown deadline"); + } + Ok(ServerWaitOutcome::Shutdown) + } + _ = tokio::time::sleep(restart_timeout) => { + warn!("Server did not stop before the restart deadline"); + Ok(ServerWaitOutcome::TimedOut) + } + } +} + /// Maximum time a single request may take before the server returns 408. const REQUEST_TIMEOUT: Duration = Duration::from_secs(30); // Header name constants not yet present in the `http` crate v1.x standard set. @@ -504,11 +566,13 @@ pub async fn run_server( env_config: EnvConfig, tls_config: Option, logs_rx: Option, + shutdown_signal: impl Future>, ) -> anyhow::Result<()> { info!("Starting Defguard Proxy server"); debug!("Using config: {env_config:?}"); let mut tasks = JoinSet::new(); + let (shutdown_tx, shutdown_rx) = watch::channel(false); let cookie_key = Arc::default(); let (reset_tx, mut reset_rx) = tokio::sync::broadcast::channel(1); let (https_cert_tx, https_cert_rx) = broadcast::channel::<(String, String)>(1); @@ -548,25 +612,39 @@ pub async fn run_server( let server_clone = grpc_server.clone(); let env_config_clone = env_config.clone(); + let mut grpc_shutdown_rx = shutdown_rx.clone(); // Start gRPC server. debug!("Spawning gRPC server task"); tasks.spawn(async move { let mut proxy_tls_config = tls_config; + let keepalive_interval = + Duration::from_secs(env_config.grpc_keepalive_interval_seconds.get()); + let keepalive_timeout = + Duration::from_secs(env_config.grpc_keepalive_timeout_seconds.get()); loop { + if *grpc_shutdown_rx.borrow() { + break; + } + let configuration = if let Some(conf) = proxy_tls_config.clone() { debug!("Using existing gRPC certificates, skipping setup process"); conf } else { info!("gRPC certificates not found, running setup process"); - let conf = run_setup(&env_config_clone, Arc::clone(&logs_rx)).await?; + let setup = run_setup(&env_config_clone, Arc::clone(&logs_rx)); + tokio::pin!(setup); + let configuration = tokio::select! { + result = &mut setup => result?, + _ = grpc_shutdown_rx.changed() => return Ok(()), + }; info!("Setup process completed successfully"); - proxy_tls_config = Some(conf.clone()); - conf + configuration }; - + proxy_tls_config = Some(configuration.clone()); server_clone.configure(configuration); + info!("Starting gRPC server..."); let server_to_run = server_clone.clone(); let addr = SocketAddr::new( @@ -575,21 +653,21 @@ pub async fn run_server( .unwrap_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED)), env_config.grpc_port, ); - let mut shutdown_rx = reset_rx.resubscribe(); - let mut server_task = tokio::spawn(async move { - server_to_run - .run(addr, async move { - let _ = shutdown_rx.recv().await; - }) - .await - }); + let mut reset_server_rx = reset_rx.resubscribe(); + let mut shutdown_server_rx = grpc_shutdown_rx.clone(); + let server = + server_to_run.run(addr, keepalive_interval, keepalive_timeout, async move { + tokio::select! { + _ = reset_server_rx.recv() => {}, + _ = shutdown_server_rx.changed() => {}, + } + }); + tokio::pin!(server); tokio::select! { - result = &mut server_task => { - match result { - Ok(Ok(())) => {} - Ok(Err(err)) => error!("gRPC server error: {err:?}, restarting..."), - Err(err) => error!("gRPC server task error: {err:?}, restarting..."), + result = &mut server => { + if let Err(err) = result { + error!("gRPC server error: {err:?}, restarting..."); } } result = reset_rx.recv() => { @@ -599,10 +677,25 @@ pub async fn run_server( } else { error!("Reset channel closed; gRPC server will keep running"); } - let _ = server_task.await; + if !wait_for_server(&mut server, SERVER_SHUTDOWN_FORCE_TIMEOUT).await? { + warn!("gRPC server did not stop after reset; dropping it"); + } + } + _ = grpc_shutdown_rx.changed() => { + info!("Shutdown signal received, draining gRPC server"); + if !wait_for_server(&mut server, SERVER_SHUTDOWN_FORCE_TIMEOUT).await? { + warn!("gRPC server did not drain before the shutdown deadline"); + } + return Ok(()); } } + + if *grpc_shutdown_rx.borrow() { + return Ok(()); + } } + + Ok(()) }); // build application @@ -625,15 +718,21 @@ pub async fn run_server( let governor_limiter = conf.limiter().clone(); // Start background task to cleanup rate-limiter data - tokio::spawn(async move { + let mut rate_limit_shutdown_rx = shutdown_rx.clone(); + tasks.spawn(async move { loop { - tokio::time::sleep(RATE_LIMITER_CLEANUP_PERIOD).await; - debug!( - "Cleaning-up rate limiter storage, current size: {}", - governor_limiter.len() - ); - governor_limiter.retain_recent(); + tokio::select! { + _ = rate_limit_shutdown_rx.changed() => break, + _ = tokio::time::sleep(RATE_LIMITER_CLEANUP_PERIOD) => { + debug!( + "Cleaning-up rate limiter storage, current size: {}", + governor_limiter.len() + ); + governor_limiter.retain_recent(); + } + } } + Ok(()) }); info!( "Configured rate limiter, per_second: {}, burst: {}", @@ -656,6 +755,7 @@ pub async fn run_server( // Start web server. debug!("Spawning API web server"); + let mut http_shutdown_rx = shutdown_rx.clone(); tasks.spawn(async move { let http_addr = SocketAddr::new( env_config @@ -675,13 +775,17 @@ pub async fn run_server( let mut port80_pause_rx = port80_pause_rx; loop { + if *http_shutdown_rx.borrow() { + break; + } + let handle = axum_server::Handle::new(); let handle_clone = handle.clone(); tls_active.store(current_tls.is_some(), Ordering::Relaxed); let app_service = app.clone().into_make_service_with_connect_info::(); let tls_certs = current_tls.clone(); - let mut server_task = tokio::spawn(async move { + let server = async move { if let Some((cert_pem, key_pem)) = tls_certs { let tls_config = build_tls_config(&cert_pem, &key_pem).await?; info!("API web server is listening on https://{https_addr}"); @@ -698,50 +802,65 @@ pub async fn run_server( .await .context("Error running HTTP server") } - }); + }; + tokio::pin!(server); tokio::select! { - result = &mut server_task => { - match result { - Ok(Ok(())) => {} - Ok(Err(err)) => return Err(err), - Err(join_err) => { - return Err(anyhow::anyhow!("Web server task panicked: {join_err}")); - } - } + result = &mut server => { + result?; break; } - result = https_cert_rx.recv() => { + result = receive_https_cert_update(&mut https_cert_rx) => { match result { - Ok(certs) => { + Some(certs) => { info!("Received new HTTPS certificates, restarting web server with TLS"); - current_tls = Some(certs); handle.graceful_shutdown(Some(Duration::from_secs(30))); - let _ = server_task.await; - } - Err(broadcast::error::RecvError::Lagged(_)) => { - warn!("Missed HTTPS certificate update; will apply next one"); + if matches!( + wait_for_server_or_shutdown( + &mut server, + &mut http_shutdown_rx, + Duration::from_secs(31), + ).await?, + ServerWaitOutcome::Shutdown + ) { + return Ok(()); + } + current_tls = Some(certs); } - Err(broadcast::error::RecvError::Closed) => { + None => { error!("HTTPS certificate channel closed unexpectedly"); - break; + handle.graceful_shutdown(Some(SERVER_SHUTDOWN_GRACE_PERIOD)); + if !wait_for_server(&mut server, SERVER_SHUTDOWN_FORCE_TIMEOUT).await? { + warn!("HTTP server did not drain after certificate channel closed"); + } + return Ok(()); } } } - result = clear_https_rx.recv() => { + result = receive_https_cert_update(&mut clear_https_rx) => { match result { - Ok(()) => { + Some(()) => { info!("Received ClearHttpsCerts, restarting web server without TLS"); - current_tls = None; handle.graceful_shutdown(Some(Duration::from_secs(30))); - let _ = server_task.await; - } - Err(broadcast::error::RecvError::Lagged(_)) => { - warn!("Missed ClearHttpsCerts update; will apply next one"); + if matches!( + wait_for_server_or_shutdown( + &mut server, + &mut http_shutdown_rx, + Duration::from_secs(31), + ).await?, + ServerWaitOutcome::Shutdown + ) { + return Ok(()); + } + current_tls = None; } - Err(broadcast::error::RecvError::Closed) => { + None => { error!("ClearHttpsCerts channel closed unexpectedly"); - break; + handle.graceful_shutdown(Some(SERVER_SHUTDOWN_GRACE_PERIOD)); + if !wait_for_server(&mut server, SERVER_SHUTDOWN_FORCE_TIMEOUT).await? { + warn!("HTTP server did not drain after certificate channel closed"); + } + return Ok(()); } } } @@ -757,13 +876,37 @@ pub async fn run_server( } => { info!("ACME task requested port 80; pausing HTTP server"); handle.graceful_shutdown(Some(Duration::from_secs(10))); - let _ = server_task.await; - // Port 80 is now free - tell the ACME task to proceed. + match wait_for_server_or_shutdown( + &mut server, + &mut http_shutdown_rx, + Duration::from_secs(11), + ).await? { + ServerWaitOutcome::Finished => {} + ServerWaitOutcome::Shutdown => return Ok(()), + ServerWaitOutcome::TimedOut => { + return Err(anyhow::anyhow!("HTTP server did not stop for ACME hand-off")); + } + } + if *http_shutdown_rx.borrow() { + return Ok(()); + } let _ = ready_tx.send(()); - // Wait for the ACME task to finish with port 80. - let _ = done_rx.await; + tokio::select! { + _ = done_rx => {}, + _ = http_shutdown_rx.changed() => return Ok(()), + } + if *http_shutdown_rx.borrow() { + return Ok(()); + } info!("ACME task released port 80; restarting HTTP server"); - // Loop will restart the server automatically. + } + _ = http_shutdown_rx.changed() => { + info!("Shutdown signal received, draining HTTP server"); + handle.graceful_shutdown(Some(SERVER_SHUTDOWN_GRACE_PERIOD)); + if !wait_for_server(&mut server, SERVER_SHUTDOWN_FORCE_TIMEOUT).await? { + warn!("HTTP server did not drain before the shutdown deadline"); + } + return Ok(()); } } } @@ -771,10 +914,41 @@ pub async fn run_server( Ok(()) }); - // TODO: Possibly switch to using select! macro info!("Defguard Proxy HTTP server initialization complete"); - while let Some(Ok(result)) = tasks.join_next().await { - result?; + tokio::pin!(shutdown_signal); + loop { + tokio::select! { + result = &mut shutdown_signal => { + info!("Shutdown signal received, draining server tasks"); + let _ = shutdown_tx.send(true); + let all_tasks_stopped = tokio::time::timeout( + SERVER_SHUTDOWN_DEADLINE, + async { + while let Some(result) = tasks.join_next().await { + match result { + Ok(Ok(())) => {} + Ok(Err(err)) => error!("Server task failed during shutdown: {err:?}"), + Err(err) if err.is_cancelled() => {}, + Err(err) => error!("Server task join failed during shutdown: {err}"), + } + } + }, + ) + .await + .is_ok(); + if !all_tasks_stopped { + warn!("Server tasks exceeded the shutdown deadline; aborting remaining tasks"); + tasks.abort_all(); + while tasks.join_next().await.is_some() {} + } + result?; + break; + } + result = tasks.join_next() => match result { + Some(Ok(result)) => result?, + Some(Err(_)) | None => break, + } + } } Ok(()) @@ -783,6 +957,98 @@ pub async fn run_server( #[cfg(test)] mod tests { use super::*; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpStream, + }; + + #[tokio::test] + async fn test_lagged_update_drains_active_http_request() { + let addr = std::net::TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0)) + .unwrap() + .local_addr() + .unwrap(); + let request_started = Arc::new(tokio::sync::Notify::new()); + let release_request = Arc::new(tokio::sync::Notify::new()); + let app = Router::new().route( + "/", + get({ + let request_started = Arc::clone(&request_started); + let release_request = Arc::clone(&release_request); + move || { + let request_started = Arc::clone(&request_started); + let release_request = Arc::clone(&release_request); + async move { + request_started.notify_one(); + release_request.notified().await; + "request drained" + } + } + }), + ); + let handle = axum_server::Handle::new(); + let server_handle = handle.clone(); + let mut server = Box::pin(async move { + axum_server::bind(addr) + .handle(server_handle) + .serve(app.into_make_service()) + .await + .map_err(anyhow::Error::from) + }); + let (update_tx, mut update_rx) = broadcast::channel(1); + let mut update = Box::pin(receive_https_cert_update(&mut update_rx)); + let mut request = tokio::spawn(async move { + let mut stream = TcpStream::connect(addr).await?; + stream + .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n") + .await?; + let mut response = Vec::new(); + stream.read_to_end(&mut response).await?; + Ok::<_, std::io::Error>(response) + }); + + tokio::select! { + _ = request_started.notified() => {}, + result = &mut server => panic!("HTTP server stopped before the request started: {result:?}"), + _ = &mut update => panic!("update arrived before the test sent one"), + _ = tokio::time::sleep(Duration::from_secs(3)) => panic!("HTTP request did not start"), + } + assert!(!request.is_finished(), "handler should still be waiting"); + + update_tx + .send(("old cert".to_owned(), "old key".to_owned())) + .unwrap(); + update_tx + .send(("latest cert".to_owned(), "latest key".to_owned())) + .unwrap(); + let latest = tokio::select! { + result = &mut server => panic!("HTTP server stopped while the request was active: {result:?}"), + update = &mut update => update.expect("latest certificate update should arrive"), + }; + assert_eq!(latest.0, "latest cert"); + assert!( + !request.is_finished(), + "lagged update must not cancel the handler" + ); + + handle.graceful_shutdown(Some(Duration::from_secs(2))); + release_request.notify_one(); + assert!( + tokio::time::timeout( + Duration::from_secs(3), + wait_for_server(&mut server, Duration::from_secs(2)) + ) + .await + .expect("server should finish after the handler drains") + .expect("server should stop cleanly") + ); + let response = tokio::time::timeout(Duration::from_secs(2), &mut request) + .await + .expect("HTTP response should finish") + .expect("request task should not panic") + .expect("HTTP request should succeed"); + assert!(response.starts_with(b"HTTP/1.1 200")); + } #[test] fn test_session_cookie_sets_security_attributes_in_both_modes() { diff --git a/src/main.rs b/src/main.rs index d674c4a7..efc521e5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -64,6 +64,23 @@ fn load_tls_config(cert_dir: &Path) -> anyhow::Result> { } } +async fn shutdown_signal() -> anyhow::Result<()> { + #[cfg(unix)] + { + use tokio::signal::unix::{SignalKind, signal}; + + let mut terminate = signal(SignalKind::terminate())?; + tokio::select! { + result = tokio::signal::ctrl_c() => result?, + _ = terminate.recv() => {} + } + } + #[cfg(not(unix))] + tokio::signal::ctrl_c().await?; + + Ok(()) +} + #[tokio::main] async fn main() -> anyhow::Result<()> { // Install the aws-lc-rs CryptoProvider as the process-wide default for rustls. @@ -98,6 +115,7 @@ async fn main() -> anyhow::Result<()> { env_config, maybe_tls_config, logs_rx.map(|r| Arc::new(Mutex::new(r))), + shutdown_signal(), ) .await?; diff --git a/src/tests/mtls.rs b/src/tests/mtls.rs index a2c695f2..bc1ebef6 100644 --- a/src/tests/mtls.rs +++ b/src/tests/mtls.rs @@ -1,5 +1,10 @@ use std::{ + env::temp_dir, net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener}, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, time::Duration, }; @@ -9,18 +14,30 @@ use defguard_certs::{ use futures_util::stream; use rustls::crypto::aws_lc_rs; use tokio::{ - net::TcpStream, + io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}, + net::{TcpListener as AsyncTcpListener, TcpStream}, spawn, - sync::oneshot, + sync::{Mutex, mpsc, oneshot}, time::{Instant, sleep}, }; use tonic::{ Code, Request, Status, + service::Interceptor, transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity}, }; use super::support::{cookie_key, test_proxy_server}; -use crate::{grpc::TlsConfig, proto::proxy_client::ProxyClient}; +use crate::{ + config::{ + DEFAULT_GRPC_KEEPALIVE_INTERVAL_SECONDS, DEFAULT_GRPC_KEEPALIVE_TIMEOUT_SECONDS, EnvConfig, + }, + grpc::TlsConfig, + http::run_server, + proto::{ + CoreResponse, EnrollmentStartResponse, InitialInfo, core_request, core_response, + proxy_client::ProxyClient, + }, +}; struct TestCerts { /// PEM-encoded CA certificate (used as the trust root for both server and client validation). @@ -109,12 +126,13 @@ fn init_crypto() { /// Spawn a configured `ProxyServer` on an OS-assigned port. /// -/// Returns `(bound_addr, shutdown_tx)`. Drop / send `shutdown_tx` to stop the server. +/// Returns `(bound_addr, shutdown_tx, connected)`. Send `shutdown_tx` to stop the server. /// /// Waits until the server is accepting TCP connections before returning, so /// callers do not need a fixed sleep to avoid startup races. -async fn spawn_test_proxy(certs: &TestCerts) -> (SocketAddr, oneshot::Sender<()>) { +async fn spawn_test_proxy(certs: &TestCerts) -> (SocketAddr, oneshot::Sender<()>, Arc) { let server = test_proxy_server(cookie_key()); + let connected = Arc::clone(&server.connected); server.configure(make_tls_config(certs)); // Find a free port, drop the listener, pass the addr to run(). @@ -127,9 +145,14 @@ async fn spawn_test_proxy(certs: &TestCerts) -> (SocketAddr, oneshot::Sender<()> let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); spawn(async move { let _ = server - .run(addr, async move { - let _ = shutdown_rx.await; - }) + .run( + addr, + Duration::from_secs(DEFAULT_GRPC_KEEPALIVE_INTERVAL_SECONDS.get()), + Duration::from_secs(DEFAULT_GRPC_KEEPALIVE_TIMEOUT_SECONDS.get()), + async move { + let _ = shutdown_rx.await; + }, + ) .await; }); @@ -148,7 +171,96 @@ async fn spawn_test_proxy(certs: &TestCerts) -> (SocketAddr, oneshot::Sender<()> sleep(Duration::from_millis(20)).await; } - (addr, shutdown_tx) + (addr, shutdown_tx, connected) +} + +async fn wait_for_listener(addr: SocketAddr) { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if TcpStream::connect(addr).await.is_ok() { + return; + } + assert!(Instant::now() < deadline, "server did not start on {addr}"); + sleep(Duration::from_millis(20)).await; + } +} + +async fn send_enrollment_request(addr: SocketAddr) -> Vec { + let mut stream = TcpStream::connect(addr) + .await + .expect("Failed to connect to HTTP server"); + let body = r#"{"token":"shutdown-test"}"#; + let request = format!( + "POST /api/v1/enrollment/start HTTP/1.1\r\nHost: {addr}\r\nUser-Agent: shutdown-test\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + stream + .write_all(request.as_bytes()) + .await + .expect("Failed to write enrollment request"); + + let mut response = Vec::new(); + tokio::time::timeout(Duration::from_secs(10), stream.read_to_end(&mut response)) + .await + .expect("HTTP request did not finish") + .expect("Failed to read HTTP response"); + response +} + +async fn relay_bytes( + mut reader: R, + mut writer: W, + forward_bytes: Arc, +) -> tokio::io::Result<()> +where + R: AsyncRead + Unpin, + W: AsyncWrite + Unpin, +{ + let mut buffer = [0; 8192]; + loop { + let bytes_read = reader.read(&mut buffer).await?; + if bytes_read == 0 { + return Ok(()); + } + if forward_bytes.load(Ordering::Relaxed) { + writer.write_all(&buffer[..bytes_read]).await?; + } + } +} + +async fn spawn_test_relay( + target: SocketAddr, +) -> (SocketAddr, Arc, tokio::task::JoinHandle<()>) { + let listener = AsyncTcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0)) + .await + .expect("Failed to bind test relay"); + let addr = listener.local_addr().expect("Failed to get relay address"); + let forward_bytes = Arc::new(AtomicBool::new(true)); + let relay_forward_bytes = Arc::clone(&forward_bytes); + let relay_task = spawn(async move { + let (client, _) = listener + .accept() + .await + .expect("Failed to accept relay client"); + let upstream = TcpStream::connect(target) + .await + .expect("Failed to connect relay to test server"); + let (client_read, client_write) = client.into_split(); + let (upstream_read, upstream_write) = upstream.into_split(); + let client_to_server = relay_bytes( + client_read, + upstream_write, + Arc::clone(&relay_forward_bytes), + ); + let server_to_client = relay_bytes(upstream_read, client_write, relay_forward_bytes); + let result = tokio::select! { + result = client_to_server => result, + result = server_to_client => result, + }; + assert!(result.is_ok(), "TCP relay failed: {result:?}"); + }); + + (addr, forward_bytes, relay_task) } /// Build a tonic `ProxyClient` using the given CA and optional client identity. @@ -193,7 +305,12 @@ async fn run_errors_without_tls_config() { // configure() is deliberately NOT called. let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0); let result = server - .run(addr, futures_util::future::pending::<()>()) + .run( + addr, + Duration::from_secs(DEFAULT_GRPC_KEEPALIVE_INTERVAL_SECONDS.get()), + Duration::from_secs(DEFAULT_GRPC_KEEPALIVE_TIMEOUT_SECONDS.get()), + futures_util::future::pending::<()>(), + ) .await; assert!(result.is_err(), "expected Err, got Ok"); assert!( @@ -214,7 +331,7 @@ async fn run_errors_without_tls_config() { async fn valid_mtls_client_accepted() { init_crypto(); let certs = TestCerts::generate(); - let (addr, shutdown_tx) = spawn_test_proxy(&certs).await; + let (addr, shutdown_tx, _) = spawn_test_proxy(&certs).await; let mut client = connect( addr, @@ -235,12 +352,215 @@ async fn valid_mtls_client_accepted() { let _ = shutdown_tx.send(()); } +#[tokio::test] +#[ignore = "takes about 6 seconds; run explicitly for shutdown timeout coverage"] +async fn test_shutdown_drains_active_enrollment_request_and_bidi_stream() { + init_crypto(); + let certs = TestCerts::generate(); + let http_addr = TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0)) + .unwrap() + .local_addr() + .unwrap(); + let grpc_addr = loop { + let addr = TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0)) + .unwrap() + .local_addr() + .unwrap(); + if addr != http_addr { + break addr; + } + }; + let mut env_config: EnvConfig = toml::from_str("").expect("default config should parse"); + env_config.http_bind_address = Some(IpAddr::V4(Ipv4Addr::LOCALHOST)); + env_config.grpc_bind_address = Some(IpAddr::V4(Ipv4Addr::LOCALHOST)); + env_config.http_port = http_addr.port(); + env_config.grpc_port = grpc_addr.port(); + env_config.https_port = 0; + env_config.cert_dir = temp_dir(); + + let tls_config = make_tls_config(&certs); + let (_logs_tx, logs_rx) = mpsc::channel(10); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let server = spawn(async move { + run_server( + env_config, + Some(tls_config), + Some(Arc::new(Mutex::new(logs_rx))), + async move { shutdown_rx.await.map_err(anyhow::Error::from) }, + ) + .await + }); + wait_for_listener(http_addr).await; + + let mut core_client = connect( + grpc_addr, + &certs.ca_cert_pem, + Some((&certs.core_client_cert_pem, &certs.core_client_key_pem)), + ) + .await + .expect("Core should connect to the gRPC server"); + let (core_tx, core_rx) = mpsc::unbounded_channel(); + assert!( + core_tx + .send(CoreResponse { + id: 0, + payload: Some(core_response::Payload::InitialInfo(InitialInfo { + private_cookies_key: vec![0; 64], + })), + }) + .is_ok() + ); + let mut request = Request::new(tokio_stream::wrappers::UnboundedReceiverStream::new( + core_rx, + )); + let mut interceptor = + defguard_version::client::ClientVersionInterceptor::new(crate::MIN_CORE_VERSION.clone()); + let version_headers = interceptor + .call(Request::new(())) + .expect("Core version headers should be valid"); + *request.metadata_mut() = version_headers.metadata().clone(); + let response = core_client + .bidi(request) + .await + .expect("Core should open the bidi stream"); + let mut core_stream = response.into_inner(); + + let mut http_response = tokio::spawn(send_enrollment_request(http_addr)); + let Ok(core_request) = tokio::time::timeout(Duration::from_secs(3), async { + loop { + tokio::select! { + response = core_stream.message() => { + break response + .expect("bidi stream should stay open") + .expect("Core request should arrive"); + } + response = &mut http_response => { + let response = response.expect("HTTP request task should not panic"); + let status_line = response.split(|byte| *byte == b'\n').next().unwrap_or_default(); + assert!( + status_line.starts_with(b"HTTP/1.1 503"), + "expected setup middleware to respond until the cookie key arrives, got: {}", + String::from_utf8_lossy(&response) + ); + http_response = tokio::spawn(send_enrollment_request(http_addr)); + } + } + } + }) + .await + else { + http_response.abort(); + panic!("Core did not receive enrollment request after initial info"); + }; + assert!(matches!( + core_request.payload, + Some(core_request::Payload::EnrollmentStart(_)) + )); + + let shutdown_started = Instant::now(); + let _ = shutdown_tx.send(()); + sleep(Duration::from_millis(25)).await; + assert!( + !server.is_finished(), + "server should wait for the active HTTP request and bidi stream" + ); + assert!( + !http_response.is_finished(), + "enrollment request should still be waiting on Core" + ); + + assert!( + core_tx + .send(CoreResponse { + id: core_request.id, + payload: Some(core_response::Payload::EnrollmentStart( + EnrollmentStartResponse { + deadline_timestamp: time::OffsetDateTime::now_utc().unix_timestamp() + 60, + ..Default::default() + }, + )), + }) + .is_ok() + ); + let response = tokio::time::timeout(Duration::from_secs(3), &mut http_response) + .await + .expect("in-flight HTTP request should finish during shutdown") + .expect("HTTP request task should not panic"); + assert!( + response.starts_with(b"HTTP/1.1 200"), + "expected enrollment response, got: {}", + String::from_utf8_lossy(&response) + ); + + sleep(Duration::from_millis(25)).await; + assert!( + !server.is_finished(), + "the open bidi stream should keep gRPC in its graceful-drain period" + ); + tokio::time::timeout(Duration::from_secs(8), server) + .await + .expect("server should enforce a bounded shutdown for the open bidi stream") + .expect("server task should not panic") + .expect("run_server should shut down cleanly"); + assert!( + shutdown_started.elapsed() >= Duration::from_secs(5), + "server should give the active bidi stream its grace period" + ); + + drop(core_tx); + drop(core_stream); + drop(core_client); +} + +#[tokio::test] +#[ignore = "takes about 30 seconds; run explicitly for keepalive coverage"] +async fn test_silent_bidi_connection_drop_is_detected_by_keepalive() { + init_crypto(); + let certs = TestCerts::generate(); + let (server_addr, shutdown_tx, connected) = spawn_test_proxy(&certs).await; + let (relay_addr, forward_bytes, relay_task) = spawn_test_relay(server_addr).await; + let mut client = connect( + relay_addr, + &certs.ca_cert_pem, + Some((&certs.core_client_cert_pem, &certs.core_client_key_pem)), + ) + .await + .expect("TLS handshake should succeed with valid client cert"); + + let mut interceptor = + defguard_version::client::ClientVersionInterceptor::new(crate::MIN_CORE_VERSION.clone()); + let version_headers = interceptor + .call(Request::new(())) + .expect("Core version headers should be valid"); + let mut request = Request::new(stream::pending::()); + *request.metadata_mut() = version_headers.metadata().clone(); + let response = client + .bidi(request) + .await + .expect("valid Core client should open a bidi stream"); + assert!(connected.load(Ordering::Relaxed)); + + forward_bytes.store(false, Ordering::Relaxed); + tokio::time::timeout(Duration::from_secs(35), async { + while connected.load(Ordering::Relaxed) { + sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Edge did not detect the silently dropped Core within 35 seconds"); + + let _ = shutdown_tx.send(()); + drop(response); + drop(client); + relay_task.await.expect("test relay task failed"); +} + /// A client that presents no certificate must be rejected at the TLS layer. #[tokio::test] async fn no_client_cert_rejected() { init_crypto(); let certs = TestCerts::generate(); - let (addr, shutdown_tx) = spawn_test_proxy(&certs).await; + let (addr, shutdown_tx, _) = spawn_test_proxy(&certs).await; // spawn_test_proxy guarantees the server is ready, so a connect() failure // here is a genuine TLS rejection rather than a startup race. @@ -269,7 +589,7 @@ async fn no_client_cert_rejected() { async fn wrong_serial_rejected() { init_crypto(); let certs = TestCerts::generate(); - let (addr, shutdown_tx) = spawn_test_proxy(&certs).await; + let (addr, shutdown_tx, _) = spawn_test_proxy(&certs).await; // This cert is valid (signed by the CA the server trusts) but has a different serial. let mut client = connect( @@ -302,7 +622,7 @@ async fn wrong_serial_rejected() { async fn rogue_ca_client_rejected() { init_crypto(); let certs = TestCerts::generate(); - let (addr, shutdown_tx) = spawn_test_proxy(&certs).await; + let (addr, shutdown_tx, _) = spawn_test_proxy(&certs).await; // spawn_test_proxy guarantees the server is ready, so a connect() failure // here is a genuine TLS rejection rather than a startup race.