From 128fa0a05e58c487e503880b0046bf3bbd9361c8 Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Fri, 9 Jan 2026 13:27:41 +0100 Subject: [PATCH 1/6] http server waits for private cookies secret before launching --- src/grpc.rs | 11 ++++++++++- src/http.rs | 38 ++++++++++++++++++++++++-------------- 2 files changed, 34 insertions(+), 15 deletions(-) diff --git a/src/grpc.rs b/src/grpc.rs index e1fef7ec..96c5a653 100644 --- a/src/grpc.rs +++ b/src/grpc.rs @@ -21,11 +21,13 @@ use crate::{ // connected clients type ClientMap = HashMap>>; +static COOKIE_KEY_HEADER: &str = "dg-cookie-key-bin"; pub(crate) struct ProxyServer { current_id: Arc, clients: Arc>, results: Arc>>>, + http_channel: mpsc::UnboundedSender>, pub(crate) connected: Arc, pub(crate) core_version: Arc>>, } @@ -33,8 +35,9 @@ pub(crate) struct ProxyServer { impl ProxyServer { #[must_use] /// Create new `ProxyServer`. - pub(crate) fn new() -> Self { + pub(crate) fn new(http_channel: mpsc::UnboundedSender>) -> Self { Self { + http_channel, current_id: Arc::new(AtomicU64::new(1)), clients: Arc::new(Mutex::new(HashMap::new())), results: Arc::new(Mutex::new(HashMap::new())), @@ -85,6 +88,7 @@ impl Clone for ProxyServer { results: Arc::clone(&self.results), connected: Arc::clone(&self.connected), core_version: Arc::clone(&self.core_version), + http_channel: self.http_channel.clone(), } } } @@ -99,6 +103,11 @@ impl proxy_server::Proxy for ProxyServer { &self, request: Request>, ) -> Result, Status> { + let cookie_key = request.metadata().get_bin(COOKIE_KEY_HEADER).unwrap(); + let key = (cookie_key.to_bytes().unwrap()) + .into_iter() + .collect::>(); + let _ = self.http_channel.send(key); let Some(address) = request.remote_addr() else { error!("Failed to determine client address for request: {request:?}"); return Err(Status::internal("Failed to determine client address")); diff --git a/src/http.rs b/src/http.rs index 40466842..fdc6f4ed 100644 --- a/src/http.rs +++ b/src/http.rs @@ -22,7 +22,7 @@ use defguard_version::{ DefguardComponent, Version, }; use serde::Serialize; -use tokio::{net::TcpListener, sync::oneshot, task::JoinSet}; +use tokio::{net::TcpListener, sync::{mpsc, oneshot}, task::JoinSet}; use tonic::transport::{Identity, Server, ServerTlsConfig}; use tower::ServiceBuilder; use tower_governor::{ @@ -166,24 +166,20 @@ async fn powered_by_header(mut response: Response) -> Response { response } -pub async fn run_server(config: Config) -> anyhow::Result<()> { +pub async fn run_server( + config: Config, +) -> anyhow::Result<()> { info!("Starting Defguard Proxy server"); debug!("Using config: {config:?}"); let mut tasks = JoinSet::new(); - // connect to upstream gRPC server - let grpc_server = ProxyServer::new(); + // Prepare the channel for gRPC -> http server communication. + // The channel sends private cookies key once core connects to gRPC. + let (tx, mut rx) = mpsc::unbounded_channel::>(); - // build application - debug!("Setting up API server"); - let shared_state = AppState { - grpc_server: grpc_server.clone(), - remote_mfa_sessions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), - // Generate secret key for encrypting cookies. - key: Key::generate(), - url: config.url.clone(), - }; + // connect to upstream gRPC server + let grpc_server = ProxyServer::new(tx); // Read gRPC TLS certificate and key. debug!("Configuring certificates for gRPC"); @@ -199,6 +195,7 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> { // Start gRPC server. debug!("Spawning gRPC server"); + let grpc_server_clone = grpc_server.clone(); tasks.spawn(async move { let addr = SocketAddr::new( config @@ -224,7 +221,7 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> { ), )) .layer(DefguardVersionLayer::new(own_version)) - .service(proxy_server::ProxyServer::new(grpc_server)); + .service(proxy_server::ProxyServer::new(grpc_server_clone)); builder .add_service(versioned_service) .serve(addr) @@ -232,6 +229,19 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> { .context("Error running gRPC server") }); + // Wait for core to connect to gRPC and send the key. + let private_cookies_key = rx.recv().await.unwrap(); + + // build application + debug!("Setting up API server"); + let shared_state = AppState { + grpc_server: grpc_server, + remote_mfa_sessions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + // Private cookies encryption key. + key: Key::from(&private_cookies_key), + url: config.url.clone(), + }; + // Setup tower_governor rate-limiter debug!( "Configuring rate limiter, per_second: {}, burst: {}", From 348bbe1c3a92c28482c5cbcff17476af953611f3 Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Mon, 12 Jan 2026 12:16:26 +0100 Subject: [PATCH 2/6] cargo fmt --- src/grpc.rs | 2 +- src/http.rs | 22 ++++++++++++---------- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/src/grpc.rs b/src/grpc.rs index 96c5a653..036fc07b 100644 --- a/src/grpc.rs +++ b/src/grpc.rs @@ -106,7 +106,7 @@ impl proxy_server::Proxy for ProxyServer { let cookie_key = request.metadata().get_bin(COOKIE_KEY_HEADER).unwrap(); let key = (cookie_key.to_bytes().unwrap()) .into_iter() - .collect::>(); + .collect::>(); let _ = self.http_channel.send(key); let Some(address) = request.remote_addr() else { error!("Failed to determine client address for request: {request:?}"); diff --git a/src/http.rs b/src/http.rs index fdc6f4ed..e345e3f7 100644 --- a/src/http.rs +++ b/src/http.rs @@ -22,7 +22,11 @@ use defguard_version::{ DefguardComponent, Version, }; use serde::Serialize; -use tokio::{net::TcpListener, sync::{mpsc, oneshot}, task::JoinSet}; +use tokio::{ + net::TcpListener, + sync::{mpsc, oneshot}, + task::JoinSet, +}; use tonic::transport::{Identity, Server, ServerTlsConfig}; use tower::ServiceBuilder; use tower_governor::{ @@ -166,17 +170,15 @@ async fn powered_by_header(mut response: Response) -> Response { response } -pub async fn run_server( - config: Config, -) -> anyhow::Result<()> { +pub async fn run_server(config: Config) -> anyhow::Result<()> { info!("Starting Defguard Proxy server"); debug!("Using config: {config:?}"); let mut tasks = JoinSet::new(); - // Prepare the channel for gRPC -> http server communication. - // The channel sends private cookies key once core connects to gRPC. - let (tx, mut rx) = mpsc::unbounded_channel::>(); + // Prepare the channel for gRPC -> http server communication. + // The channel sends private cookies key once core connects to gRPC. + let (tx, mut rx) = mpsc::unbounded_channel::>(); // connect to upstream gRPC server let grpc_server = ProxyServer::new(tx); @@ -195,7 +197,7 @@ pub async fn run_server( // Start gRPC server. debug!("Spawning gRPC server"); - let grpc_server_clone = grpc_server.clone(); + let grpc_server_clone = grpc_server.clone(); tasks.spawn(async move { let addr = SocketAddr::new( config @@ -229,8 +231,8 @@ pub async fn run_server( .context("Error running gRPC server") }); - // Wait for core to connect to gRPC and send the key. - let private_cookies_key = rx.recv().await.unwrap(); + // Wait for core to connect to gRPC and send the key. + let private_cookies_key = rx.recv().await.unwrap(); // build application debug!("Setting up API server"); From 35c162981d526978ea237681d1fa79b2f5770a72 Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Tue, 13 Jan 2026 10:58:06 +0100 Subject: [PATCH 3/6] private cookies key error handling --- src/grpc.rs | 34 +++++++++++++++++++++++++++------- src/http.rs | 9 ++++----- 2 files changed, 31 insertions(+), 12 deletions(-) diff --git a/src/grpc.rs b/src/grpc.rs index 036fc07b..763d4ca2 100644 --- a/src/grpc.rs +++ b/src/grpc.rs @@ -8,6 +8,7 @@ use std::{ }, }; +use axum_extra::extract::cookie::Key; use defguard_version::{get_tracing_variables, ComponentInfo, DefguardComponent, Version}; use tokio::sync::{mpsc, oneshot}; use tokio_stream::wrappers::UnboundedReceiverStream; @@ -27,7 +28,7 @@ pub(crate) struct ProxyServer { current_id: Arc, clients: Arc>, results: Arc>>>, - http_channel: mpsc::UnboundedSender>, + http_channel: mpsc::UnboundedSender, pub(crate) connected: Arc, pub(crate) core_version: Arc>>, } @@ -35,7 +36,7 @@ pub(crate) struct ProxyServer { impl ProxyServer { #[must_use] /// Create new `ProxyServer`. - pub(crate) fn new(http_channel: mpsc::UnboundedSender>) -> Self { + pub(crate) fn new(http_channel: mpsc::UnboundedSender) -> Self { Self { http_channel, current_id: Arc::new(AtomicU64::new(1)), @@ -103,11 +104,6 @@ impl proxy_server::Proxy for ProxyServer { &self, request: Request>, ) -> Result, Status> { - let cookie_key = request.metadata().get_bin(COOKIE_KEY_HEADER).unwrap(); - let key = (cookie_key.to_bytes().unwrap()) - .into_iter() - .collect::>(); - let _ = self.http_channel.send(key); let Some(address) = request.remote_addr() else { error!("Failed to determine client address for request: {request:?}"); return Err(Status::internal("Failed to determine client address")); @@ -123,6 +119,30 @@ impl proxy_server::Proxy for ProxyServer { info!("Defguard Core gRPC client connected from: {address}"); + // Retrieve private cookies key from the header. + let cookie_key = request.metadata().get_bin(COOKIE_KEY_HEADER); + let key = match cookie_key { + Some(key) => Key::from(&key.to_bytes().map_err(|err| { + error!("Failed to decode private cookie key: {err:?}"); + Status::internal("Failed to decode private cookie key") + })?), + // If the header is missing, fall back to generating a local key. + // This preserves compatibility with older Core versions that did not + // provide a shared cookie key. In this mode, cookie-based sessions will + // not be shared across proxy instances and HA won't work. + None => { + warn!( + "Private cookie key not provided by Core; falling back to a locally generated key. \ + This typically indicates an older Core version and disables cookie sharing across proxies." + ); + Key::generate() + } + }; + self.http_channel.send(key).map_err(|err| { + error!("Failed to send private cookies key to HTTP server: {err:?}"); + Status::internal("Failed to send private cookies key to HTTP server") + })?; + let (tx, rx) = mpsc::unbounded_channel(); self.clients.lock().unwrap().insert(address, tx); self.connected.store(true, Ordering::Relaxed); diff --git a/src/http.rs b/src/http.rs index e345e3f7..55e75546 100644 --- a/src/http.rs +++ b/src/http.rs @@ -178,7 +178,7 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> { // Prepare the channel for gRPC -> http server communication. // The channel sends private cookies key once core connects to gRPC. - let (tx, mut rx) = mpsc::unbounded_channel::>(); + let (tx, mut rx) = mpsc::unbounded_channel::(); // connect to upstream gRPC server let grpc_server = ProxyServer::new(tx); @@ -231,16 +231,15 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> { .context("Error running gRPC server") }); - // Wait for core to connect to gRPC and send the key. - let private_cookies_key = rx.recv().await.unwrap(); + // Wait for core to connect to gRPC and send private cookies encryption key. + let key = rx.recv().await.unwrap(); // build application debug!("Setting up API server"); let shared_state = AppState { + key, grpc_server: grpc_server, remote_mfa_sessions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), - // Private cookies encryption key. - key: Key::from(&private_cookies_key), url: config.url.clone(), }; From 933a4c8978a919f80712d67f6726eae7733f82dd Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Tue, 13 Jan 2026 11:47:45 +0100 Subject: [PATCH 4/6] fmt --- src/grpc.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/grpc.rs b/src/grpc.rs index e4b759ab..5c20ae27 100644 --- a/src/grpc.rs +++ b/src/grpc.rs @@ -246,16 +246,16 @@ impl proxy_server::Proxy for ProxyServer { // not be shared across proxy instances and HA won't work. None => { warn!( - "Private cookie key not provided by Core; falling back to a locally generated key. \ - This typically indicates an older Core version and disables cookie sharing across proxies." - ); + "Private cookie key not provided by Core; falling back to a locally generated key. \ + This typically indicates an older Core version and disables cookie sharing across proxies." + ); Key::generate() } }; self.http_channel.send(key).map_err(|err| { - error!("Failed to send private cookies key to HTTP server: {err:?}"); - Status::internal("Failed to send private cookies key to HTTP server") - })?; + error!("Failed to send private cookies key to HTTP server: {err:?}"); + Status::internal("Failed to send private cookies key to HTTP server") + })?; let (tx, rx) = mpsc::unbounded_channel(); self.clients From 7749750c5222ffd34c8522c3755ded1dc81733a7 Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Tue, 13 Jan 2026 12:46:32 +0100 Subject: [PATCH 5/6] handle premature http chanel closure --- src/http.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/http.rs b/src/http.rs index 1cf3d2c6..56ee352e 100644 --- a/src/http.rs +++ b/src/http.rs @@ -257,7 +257,9 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> { }); // Wait for core to connect to gRPC and send private cookies encryption key. - let key = rx.recv().await.unwrap(); + let Some(key) = rx.recv().await else { + return Err(anyhow::Error::msg("http channel closed")); + }; // build application debug!("Setting up API server"); From 219869275321850b297a3c60262cf067c46c029e Mon Sep 17 00:00:00 2001 From: Jacek Chmielewski Date: Wed, 14 Jan 2026 11:20:46 +0100 Subject: [PATCH 6/6] clippy fix --- src/http.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/http.rs b/src/http.rs index 56ee352e..8d165df4 100644 --- a/src/http.rs +++ b/src/http.rs @@ -265,7 +265,7 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> { debug!("Setting up API server"); let shared_state = AppState { key, - grpc_server: grpc_server, + grpc_server, remote_mfa_sessions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), url: config.url.clone(), };