From 17fcdf127f955f08398839ddbfa0b4a3b90f1929 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maciej=20W=C3=B3jcik?= Date: Fri, 23 Oct 2026 09:59:19 +0200 Subject: [PATCH 1/2] port keepalive fix from stable --- Cargo.toml | 2 +- flake.lock | 18 ++--- src/bin/defguard-gateway.rs | 108 ++++++++++++++++++--------- src/config.rs | 91 +++++++++++++++++++++++ src/gateway.rs | 108 ++++++++++++++++++--------- src/gateway_server.rs | 9 ++- src/lib.rs | 58 ++++++++++----- src/tests/keepalive.rs | 144 ++++++++++++++++++++++++++++++++++++ src/tests/mod.rs | 1 + src/tests/mtls.rs | 30 ++++---- 10 files changed, 453 insertions(+), 116 deletions(-) create mode 100644 src/tests/keepalive.rs diff --git a/Cargo.toml b/Cargo.toml index df7919b1..3871fbde 100644 --- a/Cargo.toml +++ b/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", diff --git a/flake.lock b/flake.lock index 075b766a..e111e7af 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": 1791260994, - "narHash": "sha256-Miqqk/ammqnTUxaoCyvtwoPeLbI1ksFyLxi/CazZpWY=", + "lastModified": 1791956154, + "narHash": "sha256-jBYGPi4udQ9uo4KQp9m2J0546nJVQitG15gvZSaVVAs=", "owner": "NixOS", "repo": "nixpkgs", - "rev": "151fa4e8ddfdd8dd25d945ad94ed54a13de9f6e4", + "rev": "55bfb4d7437bd4e09ab7cba58a43bd68eb310044", "type": "github" }, "original": { @@ -48,11 +48,11 @@ ] }, "locked": { - "lastModified": 1791361643, - "narHash": "sha256-HMdfSb1ONHs1RbSkQtCXm+tWEsTikisWWtXDnv0aQ+I=", + "lastModified": 1792219666, + "narHash": "sha256-RMO08tKD+hFIL5ryWi85VY4aHE5oIpdV1h4kVFctOkw=", "owner": "oxalica", "repo": "rust-overlay", - "rev": "8afee9fa8caa877a4feb65adc33e8d72f3747dc9", + "rev": "a110337465cfb76e3806b81b9db3094c3e6e770a", "type": "github" }, "original": { diff --git a/src/bin/defguard-gateway.rs b/src/bin/defguard-gateway.rs index b1c95942..59887f6d 100644 --- a/src/bin/defguard-gateway.rs +++ b/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>, 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>, std::io::Error> { + Ok(signal::ctrl_c()) +} fn load_tls_config(cert_dir: &Path) -> Result, 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(()) } diff --git a/src/config.rs b/src/config.rs index f76d0827..0f621699 100644 --- a/src/config.rs +++ b/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::("grpc_keepalive_interval = 0").is_err()); + assert!(toml::from_str::("grpc_keepalive_timeout = 0").is_err()); +} diff --git a/src/gateway.rs b/src/gateway.rs index 29e15682..8f36a9dc 100644 --- a/src/gateway.rs +++ b/src/gateway.rs @@ -13,8 +13,8 @@ use std::{ use defguard_wireguard_rs::{WireguardInterfaceApi, net::IpAddrMask}; use tokio::{ - signal, - sync::{mpsc, oneshot}, + sync::{mpsc, oneshot, watch}, + task::JoinSet, time::interval, }; use tonic::Status; @@ -35,44 +35,76 @@ use crate::{ setup::run_setup, }; -/// Keeps the gRPC server running, but allows it to be stopped and -/// restarted in setup mode when a purge request arrives. +fn cleanup_on_quit(config: &Config, gateway: &Arc>) { + if config.clean_on_quit { + gateway.lock().unwrap().purge(); + } +} + +/// Runs initial setup and keeps the gRPC server available until shutdown. pub async fn run_gateway_loop( config: Config, gateway: Arc>, logs_rx: Arc>>, - mut tls_config: TlsConfig, + tls_config: Option, + mut shutdown: watch::Receiver, ) { + let mut tls_config = match tls_config { + Some(tls_config) => { + log::info!( + "Using existing gRPC TLS certificates from {}", + config.cert_dir.display() + ); + tls_config + } + None => { + log::info!( + "gRPC TLS certificates not found in {}. They will be generated during setup.", + config.cert_dir.display() + ); + tokio::select! { + biased; + _ = shutdown.changed() => { + cleanup_on_quit(&config, &gateway); + return; + } + result = run_setup(&config, Arc::clone(&logs_rx)) => match result { + Ok(tls_config) => tls_config, + Err(err) => { + error!("Failed to run setup: {err}"); + return; + } + } + } + } + }; + loop { - // Channel used by the gRPC service to request entering setup mode. - // The purge RPC sends on this channel. let (reset_tx, mut reset_rx) = oneshot::channel(); - // Build and start a gRPC server instance wired with the reset signal. let mut gateway_server = GatewayServer::new(Arc::clone(&gateway), config.cert_dir.clone(), reset_tx); gateway_server.set_tls_config(tls_config.clone()); - let mut server_handle = tokio::spawn(gateway_server.start(config.clone())); + let mut server_tasks = JoinSet::new(); + server_tasks.spawn(gateway_server.start(config.clone())); tokio::select! { biased; - // Handle Ctrl-C - _ = signal::ctrl_c() => { - if config.clean_on_quit { - gateway.lock().unwrap().purge(); - } - break; + _ = shutdown.changed() => { + server_tasks.abort_all(); + while server_tasks.join_next().await.is_some() {} + cleanup_on_quit(&config, &gateway); + return; } // The purge RPC requested setup mode. - signal = &mut reset_rx => { - // Handle channel closed event. - if signal.is_err() { - match server_handle.await { - Ok(Ok(())) => (), - Ok(Err(err)) => { + reset = &mut reset_rx => { + if reset.is_err() { + match server_tasks.join_next().await { + Some(Ok(Ok(()))) | None => (), + Some(Ok(Err(err))) => { error!("gRPC server task failed: {err}"); return; } - Err(err) => { + Some(Err(err)) => { error!("gRPC server task failed: {err}"); return; } @@ -80,29 +112,33 @@ pub async fn run_gateway_loop( break; } - // Entering setup mode - stop current gRPC server. - server_handle.abort(); - let _ = server_handle.await; + server_tasks.abort_all(); + while server_tasks.join_next().await.is_some() {} - // Run setup server to obtain new TLS certs, then loop to restart gRPC. log::info!("Restarting setup server after purge request"); - tls_config = match run_setup(&config, Arc::clone(&logs_rx)).await { - Ok(config) => config, - Err(err) => { - error!("Failed to run setup {err}"); + tls_config = tokio::select! { + biased; + _ = shutdown.changed() => { + cleanup_on_quit(&config, &gateway); return; } - } + result = run_setup(&config, Arc::clone(&logs_rx)) => match result { + Ok(tls_config) => tls_config, + Err(err) => { + error!("Failed to run setup {err}"); + return; + } + } + }; } - // Server exited on its own (error or normal shutdown). - result = &mut server_handle => { + result = server_tasks.join_next() => { match result { - Ok(Ok(())) => (), - Ok(Err(err)) => { + Some(Ok(Ok(()))) | None => (), + Some(Ok(Err(err))) => { error!("gRPC server task failed: {err}"); return; } - Err(err) => { + Some(Err(err)) => { error!("gRPC server task failed: {err}"); return; } diff --git a/src/gateway_server.rs b/src/gateway_server.rs index e5795afa..eee1690a 100644 --- a/src/gateway_server.rs +++ b/src/gateway_server.rs @@ -26,7 +26,7 @@ use crate::{ CORE_CLIENT_CERT_NAME, GRPC_CA_CERT_NAME, GRPC_CERT_NAME, GRPC_KEY_NAME, VERSION, config::Config, error::GatewayError, - execute_command, + execute_command_cancellable, gateway::{Gateway, TlsConfig}, proto::gateway::{CoreRequest, CoreResponse, core_request, core_response, gateway_server}, version::is_core_version_supported, @@ -65,7 +65,7 @@ impl GatewayServer { if let Some(post_up) = &config.post_up { debug!("Executing specified POST_UP command: {post_up}"); - execute_command(post_up)?; + execute_command_cancellable(post_up).await?; } let tls_config = self @@ -89,7 +89,10 @@ impl GatewayServer { let tls_config = server_tls_config(&tls.grpc_cert_pem, &tls.grpc_key_pem, &tls.grpc_ca_cert_pem) .map_err(|e| GatewayError::SetupError(e.to_string()))?; - let mut builder = Server::builder().tls_config(tls_config)?; + let mut builder = Server::builder() + .http2_keepalive_interval(Some(config.grpc_keepalive_interval())) + .http2_keepalive_timeout(Some(config.grpc_keepalive_timeout())) + .tls_config(tls_config)?; // Extract Core client cert serial for pinning. let expected_serial = CertificateInfo::from_der(&tls.core_client_cert_der) diff --git a/src/lib.rs b/src/lib.rs index 1cc0eb6f..a92efb55 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -95,6 +95,29 @@ pub fn init_syslog(config: &Config, pid: u32) -> Result<(), GatewayError> { Ok(()) } +fn command_execution_error(command: &str, err: std::io::Error) -> GatewayError { + error!("Failed to execute command {command}. Error: {err}"); + GatewayError::CommandExecutionFailed { + command: command.to_string(), + error: err.to_string(), + } +} + +fn log_command_output(command: &str, output: &std::process::Output) { + if output.status.success() { + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + + info!("Command {command} executed successfully. Stdout: {stdout}"); + if !stderr.is_empty() { + error!("Stderr:\n{stderr}"); + } + } else { + let stderr = String::from_utf8_lossy(&output.stderr); + error!("Error executing command {command}. Stderr:\n{stderr}"); + } +} + /// Execute command passed as argument. pub fn execute_command(command: &str) -> Result<(), GatewayError> { let mut command_parts = command.split_whitespace(); @@ -103,26 +126,25 @@ pub fn execute_command(command: &str) -> Result<(), GatewayError> { let output = Command::new(command) .args(command_parts) .output() - .map_err(|err| { - error!("Failed to execute command {command}. Error: {err}"); - GatewayError::CommandExecutionFailed { - command: command.to_string(), - error: err.to_string(), - } - })?; + .map_err(|err| command_execution_error(command, err))?; + log_command_output(command, &output); + } + Ok(()) +} - if output.status.success() { - let stdout = String::from_utf8_lossy(&output.stdout); - let stderr = String::from_utf8_lossy(&output.stderr); +/// Kill the direct child process if this future is cancelled. +// ponytail: descendants survive; use process-group termination if hooks need it. +pub(crate) async fn execute_command_cancellable(command: &str) -> Result<(), GatewayError> { + let mut command_parts = command.split_whitespace(); - info!("Command {command} executed successfully. Stdout: {stdout}"); - if !stderr.is_empty() { - error!("Stderr:\n{stderr}"); - } - } else { - let stderr = String::from_utf8_lossy(&output.stderr); - error!("Error executing command {command}. Stderr:\n{stderr}"); - } + if let Some(command) = command_parts.next() { + let output = tokio::process::Command::new(command) + .args(command_parts) + .kill_on_drop(true) + .output() + .await + .map_err(|err| command_execution_error(command, err))?; + log_command_output(command, &output); } Ok(()) } diff --git a/src/tests/keepalive.rs b/src/tests/keepalive.rs new file mode 100644 index 00000000..1f7498be --- /dev/null +++ b/src/tests/keepalive.rs @@ -0,0 +1,144 @@ +use std::{ + net::{IpAddr, Ipv4Addr, SocketAddr}, + time::Duration, +}; + +use defguard_version::{SYSTEM_INFO_HEADER, VERSION_HEADER}; +use tokio::{ + io::{copy, copy_bidirectional, sink}, + net::{TcpListener, TcpStream}, + sync::{mpsc, oneshot}, + time::{sleep, timeout}, +}; +use tokio_stream::wrappers::UnboundedReceiverStream; +use tonic::Request; + +use super::mtls::{TestCerts, connect, init_crypto, spawn_test_gateway}; +use crate::proto::gateway::CoreResponse; + +const RPC_TIMEOUT: Duration = Duration::from_secs(10); + +async fn start_blackhole_relay( + upstream: SocketAddr, +) -> (u16, oneshot::Sender<()>, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0)) + .await + .unwrap(); + let relay_port = listener.local_addr().unwrap().port(); + let (blackhole_tx, blackhole_rx) = oneshot::channel(); + + let relay_task = tokio::spawn(async move { + let (mut client, _) = listener.accept().await.unwrap(); + let mut gateway = TcpStream::connect(upstream).await.unwrap(); + + tokio::select! { + _ = copy_bidirectional(&mut client, &mut gateway) => {} + _ = blackhole_rx => { + let mut client_sink = sink(); + let mut gateway_sink = sink(); + let _ = tokio::join!( + copy(&mut client, &mut client_sink), + copy(&mut gateway, &mut gateway_sink), + ); + } + } + }); + + (relay_port, blackhole_tx, relay_task) +} + +fn core_request( + stream: UnboundedReceiverStream, +) -> Request> { + let mut request = Request::new(stream); + request + .metadata_mut() + .insert(VERSION_HEADER, "2.1.1".parse().unwrap()); + request + .metadata_mut() + .insert(SYSTEM_INFO_HEADER, "Linux;test;x86_64".parse().unwrap()); + request +} + +#[tokio::test] +#[ignore = "waits for the 30-second HTTP/2 keepalive timeout"] +async fn test_silently_dropped_core_connection_releases_slot() { + init_crypto(); + let certs = TestCerts::generate(); + let (gateway_port, shutdown_tx, gateway) = spawn_test_gateway(&certs).await; + let upstream = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), gateway_port); + let (relay_port, blackhole_tx, relay_task) = start_blackhole_relay(upstream).await; + + let identity = Some(( + certs.core_client_cert_pem.as_str(), + certs.core_client_key_pem.as_str(), + )); + let mut first_client = timeout( + RPC_TIMEOUT, + connect(relay_port, &certs.ca_cert_pem, identity), + ) + .await + .expect("first Core connection timed out") + .expect("first Core should connect through the relay"); + let (first_request_tx, first_request_rx) = mpsc::unbounded_channel(); + let first_response = timeout( + RPC_TIMEOUT, + first_client.bidi(core_request(UnboundedReceiverStream::new(first_request_rx))), + ) + .await + .expect("first Core bidi call timed out") + .expect("first Core should be admitted"); + let mut first_response = first_response.into_inner(); + let first_message = timeout(RPC_TIMEOUT, first_response.message()) + .await + .expect("first Core response timed out") + .unwrap(); + assert!(first_message.is_some()); + assert!(gateway.lock().unwrap().client_tx.is_some()); + + blackhole_tx.send(()).unwrap(); + let released = timeout(Duration::from_secs(35), async { + loop { + if gateway.lock().unwrap().client_tx.is_none() { + break; + } + sleep(Duration::from_millis(100)).await; + } + }) + .await; + assert!( + released.is_ok(), + "gateway did not release the slot in 35 seconds" + ); + + let mut second_client = timeout( + RPC_TIMEOUT, + connect(gateway_port, &certs.ca_cert_pem, identity), + ) + .await + .expect("second Core connection timed out") + .expect("second Core should connect directly"); + let (second_request_tx, second_request_rx) = mpsc::unbounded_channel(); + let second_response = timeout( + RPC_TIMEOUT, + second_client.bidi(core_request(UnboundedReceiverStream::new( + second_request_rx, + ))), + ) + .await + .expect("second Core bidi call timed out") + .expect("second Core should be admitted after keepalive closes the stale stream"); + let mut second_response = second_response.into_inner(); + let second_message = timeout(RPC_TIMEOUT, second_response.message()) + .await + .expect("second Core response timed out") + .unwrap(); + assert!(second_message.is_some()); + + drop(first_request_tx); + drop(first_response); + drop(second_request_tx); + let _ = shutdown_tx.send(()); + relay_task.abort(); + let _ = relay_task.await; +} diff --git a/src/tests/mod.rs b/src/tests/mod.rs index 113ec4f5..5da552ba 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -1,3 +1,4 @@ +mod keepalive; mod mock_wgapi; mod mtls; mod recovery; diff --git a/src/tests/mtls.rs b/src/tests/mtls.rs index aa45a3dc..eac6b89a 100644 --- a/src/tests/mtls.rs +++ b/src/tests/mtls.rs @@ -29,13 +29,13 @@ use crate::{ proto::gateway::gateway_client::GatewayClient, }; -struct TestCerts { - ca_cert_pem: String, +pub(super) struct TestCerts { + pub(super) ca_cert_pem: String, gateway_cert_pem: String, gateway_key_pem: String, core_client_cert_der: Vec, - core_client_cert_pem: String, - core_client_key_pem: String, + pub(super) core_client_cert_pem: String, + pub(super) core_client_key_pem: String, wrong_serial_cert_pem: String, wrong_serial_key_pem: String, rogue_client_cert_pem: String, @@ -43,7 +43,7 @@ struct TestCerts { } impl TestCerts { - fn generate() -> Self { + pub(super) fn generate() -> Self { let ca = CertificateAuthority::new("Test CA", "test@test.local", 365).unwrap(); let ca_cert_pem = ca.cert_pem().unwrap(); @@ -103,16 +103,18 @@ fn build_gateway(config: &Config) -> Gateway { /// /// Must be called before any TLS code runs. Safe to call from multiple tests - /// subsequent calls after the first are silently ignored. -fn init_crypto() { +pub(super) fn init_crypto() { let _ = ring::default_provider().install_default(); } /// Spawn a configured `GatewayServer` on a free port. -/// Returns `(bound_port, shutdown_tx)`. +/// Returns the bound port, shutdown sender, and gateway state. /// /// 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_gateway(certs: &TestCerts) -> (u16, oneshot::Sender<()>) { +pub(super) async fn spawn_test_gateway( + certs: &TestCerts, +) -> (u16, oneshot::Sender<()>, Arc>) { let port = TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0)) .unwrap() .local_addr() @@ -152,11 +154,11 @@ async fn spawn_test_gateway(certs: &TestCerts) -> (u16, oneshot::Sender<()>) { sleep(Duration::from_millis(20)).await; } - (port, shutdown_tx) + (port, shutdown_tx, gateway) } /// Build a tonic `GatewayClient` with optional mTLS client identity. -async fn connect( +pub(super) async fn connect( port: u16, ca_cert_pem: &str, client_identity: Option<(&str, &str)>, @@ -217,7 +219,7 @@ async fn test_start_errors_without_tls_config() { async fn test_valid_mtls_client_accepted() { init_crypto(); let certs = TestCerts::generate(); - let (port, shutdown_tx) = spawn_test_gateway(&certs).await; + let (port, shutdown_tx, _) = spawn_test_gateway(&certs).await; let mut client = connect( port, @@ -246,7 +248,7 @@ async fn test_valid_mtls_client_accepted() { async fn test_no_client_cert_rejected() { init_crypto(); let certs = TestCerts::generate(); - let (port, shutdown_tx) = spawn_test_gateway(&certs).await; + let (port, shutdown_tx, _) = spawn_test_gateway(&certs).await; // spawn_test_gateway guarantees the server is ready, so a connect() failure // here is a genuine TLS rejection rather than a startup race. @@ -275,7 +277,7 @@ async fn test_no_client_cert_rejected() { async fn test_wrong_serial_rejected() { init_crypto(); let certs = TestCerts::generate(); - let (port, shutdown_tx) = spawn_test_gateway(&certs).await; + let (port, shutdown_tx, _) = spawn_test_gateway(&certs).await; // This cert is valid (signed by the CA the server trusts) but has a different serial. let mut client = connect( @@ -308,7 +310,7 @@ async fn test_wrong_serial_rejected() { async fn test_rogue_ca_client_rejected() { init_crypto(); let certs = TestCerts::generate(); - let (port, shutdown_tx) = spawn_test_gateway(&certs).await; + let (port, shutdown_tx, _) = spawn_test_gateway(&certs).await; // spawn_test_gateway guarantees the server is ready, so a connect() failure // here is a genuine TLS rejection rather than a startup race. From 77ef710e2b5fcb2fe48c4ef2f7de3e1ee5e9957c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maciej=20W=C3=B3jcik?= Date: Fri, 23 Oct 2026 10:29:29 +0200 Subject: [PATCH 2/2] update deps --- Cargo.lock | 4 ++-- flake.lock | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fae32578..e197bd30 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -813,9 +813,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", ] diff --git a/flake.lock b/flake.lock index e111e7af..3e1f1eff 100644 --- a/flake.lock +++ b/flake.lock @@ -20,11 +20,11 @@ }, "nixpkgs": { "locked": { - "lastModified": 1791956154, - "narHash": "sha256-jBYGPi4udQ9uo4KQp9m2J0546nJVQitG15gvZSaVVAs=", + "lastModified": 1792502884, + "narHash": "sha256-IoNisKcGbUjdd45qCqywGMCYmZ5EnkxezvLII0FWwNA=", "owner": "NixOS", "repo": "nixpkgs", - "rev": "55bfb4d7437bd4e09ab7cba58a43bd68eb310044", + "rev": "02c087f3fe673c9e317f206a627cb6b910336dbe", "type": "github" }, "original": { @@ -48,11 +48,11 @@ ] }, "locked": { - "lastModified": 1792219666, - "narHash": "sha256-RMO08tKD+hFIL5ryWi85VY4aHE5oIpdV1h4kVFctOkw=", + "lastModified": 1792739799, + "narHash": "sha256-M2rr/yPbTHvgAzCyZ+zHJHflFM0XU7t42Sw2skA3TQM=", "owner": "oxalica", "repo": "rust-overlay", - "rev": "a110337465cfb76e3806b81b9db3094c3e6e770a", + "rev": "b0829a2cc9f7b596583331d78a942c1e5170d8a0", "type": "github" }, "original": {