Merge branch 'dev' of https://github.com/DefGuard/proxy into proxy-pairing-1

Author: t-aleksander

proto

@@ -4,7 +4,7 @@ use std::{
     net::SocketAddr,
     sync::{
         atomic::{AtomicBool, AtomicU64, Ordering},
-        Arc, Mutex,
+        Arc, Mutex, RwLock,
     },
 };
 
@@ -31,7 +31,6 @@ use crate::{
 
 // connected clients
 type ClientMap = HashMap<SocketAddr, mpsc::UnboundedSender<Result<CoreRequest, Status>>>;
-static COOKIE_KEY_HEADER: &str = "dg-cookie-key-bin";
 
 #[derive(Debug, Clone, Default)]
 pub struct Configuration {
@@ -43,18 +42,18 @@ pub(crate) struct ProxyServer {
     current_id: Arc<AtomicU64>,
     clients: Arc<Mutex<ClientMap>>,
     results: Arc<Mutex<HashMap<u64, oneshot::Sender<core_response::Payload>>>>,
-    http_channel: mpsc::UnboundedSender<Key>,
     pub(crate) connected: Arc<AtomicBool>,
     pub(crate) core_version: Arc<Mutex<Option<Version>>>,
     config: Arc<Mutex<Option<Configuration>>>,
+    cookie_key: Arc<RwLock<Option<Key>>>,
 }
 
 impl ProxyServer {
     #[must_use]
     /// Create new `ProxyServer`.
-    pub(crate) fn new(http_channel: mpsc::UnboundedSender<Key>) -> Self {
+    pub(crate) fn new(cookie_key: Arc<RwLock<Option<Key>>>) -> Self {
         Self {
-            http_channel,
+            cookie_key,
             current_id: Arc::new(AtomicU64::new(1)),
             clients: Arc::new(Mutex::new(HashMap::new())),
             results: Arc::new(Mutex::new(HashMap::new())),
@@ -175,7 +174,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(),
+            cookie_key: Arc::clone(&self.cookie_key),
             config: Arc::clone(&self.config),
         }
     }
@@ -213,31 +212,6 @@ impl proxy_server::Proxy for ProxyServer {
         let _guard = span.enter();
 
         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()
@@ -250,22 +224,32 @@ impl proxy_server::Proxy for ProxyServer {
         let clients = Arc::clone(&self.clients);
         let results = Arc::clone(&self.results);
         let connected = Arc::clone(&self.connected);
-        let mut stream = request.into_inner();
+        let cookie_key = Arc::clone(&self.cookie_key);
         tokio::spawn(
             async move {
+                let mut stream = request.into_inner();
                 loop {
                     match stream.message().await {
                         Ok(Some(response)) => {
                             debug!("Received message from Defguard Core ID={}", response.id);
                             connected.store(true, Ordering::Relaxed);
                             if let Some(payload) = response.payload {
-                                let maybe_rx = results.lock().expect("Failed to acquire lock on results hashmap when processing response").remove(&response.id);
-                                if let Some(rx) = maybe_rx {
-                                    if let Err(err) = rx.send(payload) {
-                                        error!("Failed to send message to rx {:?}", err.type_id());
+                                match payload {
+                                    core_response::Payload::InitialInfo(payload) => {
+                                        info!("Received private cookies key");
+                                        let key = Key::from(&payload.private_cookies_key);
+                                        *cookie_key.write().unwrap() = Some(key);
+                                    },
+                                    _ => {
+                                        let maybe_rx = results.lock().expect("Failed to acquire lock on results hashmap when processing response").remove(&response.id);
+                                        if let Some(rx) = maybe_rx {
+                                            if let Err(err) = rx.send(payload) {
+                                                error!("Failed to send message to rx {:?}", err.type_id());
+                                            }
+                                        } else {
+                                            error!("Missing receiver for response #{}", response.id);
+                                        }
                                     }
-                                } else {
-                                    error!("Missing receiver for response #{}", response.id);
                                 }
                             }
                         }

src/grpc.rs

@@ -2,7 +2,7 @@ use std::{
     collections::HashMap,
     net::{IpAddr, Ipv4Addr, SocketAddr},
     path::Path,
-    sync::{atomic::Ordering, Arc},
+    sync::{atomic::Ordering, Arc, LazyLock, RwLock},
     time::Duration,
 };
 
@@ -12,6 +12,7 @@ use axum::{
     extract::{ConnectInfo, FromRef, State},
     http::{header::HeaderValue, Request, Response, StatusCode},
     middleware::{self, Next},
+    response::IntoResponse,
     routing::{get, post},
     serve, Json, Router,
 };
@@ -21,7 +22,7 @@ use defguard_version::{server::DefguardVersionLayer, Version};
 use serde::Serialize;
 use tokio::{
     net::TcpListener,
-    sync::{mpsc, oneshot},
+    sync::{mpsc, oneshot, Mutex},
     task::JoinSet,
 };
 use tower_governor::{
@@ -57,7 +58,7 @@ pub(crate) struct AppState {
     pub(crate) grpc_server: ProxyServer,
     pub(crate) remote_mfa_sessions:
         Arc<tokio::sync::Mutex<HashMap<String, oneshot::Sender<String>>>>,
-    key: Key,
+    cookie_key: Arc<RwLock<Option<Key>>>,
     url: Url,
 }
 
@@ -79,7 +80,10 @@ impl AppState {
 
 impl FromRef<AppState> for Key {
     fn from_ref(state: &AppState) -> Self {
-        state.key.clone()
+        let maybe_key = state.cookie_key.read().unwrap().clone();
+        // We return the dummy key only to satisfy the `FromRef` trait, but it is never
+        // used in practice because of the `ensure_configured` middleware.
+        maybe_key.unwrap_or_else(|| Key::from(&[0; 64]))
     }
 }
 
@@ -206,18 +210,45 @@ pub async fn run_setup(
     Ok(configuration)
 }
 
+/// Middleware that gates all HTTP endpoints except health checks until the proxy
+/// is fully configured.
+///
+/// The proxy cannot safely handle requests that rely on encrypted cookies
+/// (e.g. OpenID / MFA flows) until it receives the cookie encryption key from
+/// the core. This key is provided asynchronously after the core connects.
+///
+/// Until the key is available, only health check endpoints are served and all
+/// other requests return HTTP 503 (Service Unavailable). Once the key is set,
+/// the middleware becomes a no-op and all routes are enabled.
+async fn ensure_configured(
+    State(state): State<AppState>,
+    request: Request<Body>,
+    next: Next,
+) -> Response<Body> {
+    // Allow healthchecks even before core connects and gives us the cookie key.
+    let path = request.uri().path();
+    if matches!(path, "/api/v1/health" | "/api/v1/health-grpc") {
+        return next.run(request).await;
+    }
+
+    // Block all other requests until cookie key is configured.
+    if state.cookie_key.read().unwrap().is_none() {
+        return StatusCode::SERVICE_UNAVAILABLE.into_response();
+    }
+
+    next.run(request).await
+}
+
 pub async fn run_server(env_config: EnvConfig, config: Configuration) -> anyhow::Result<()> {
     info!("Starting Defguard Proxy server");
     debug!("Using config: {env_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::<Key>();
+    let cookie_key = Default::default();
 
     // connect to upstream gRPC server
-    let grpc_server = ProxyServer::new(tx);
+    let grpc_server = ProxyServer::new(Arc::clone(&cookie_key));
+
     let server_clone = grpc_server.clone();
     grpc_server.configure(config);
 
@@ -241,15 +272,10 @@ pub async fn run_server(env_config: EnvConfig, config: Configuration) -> anyhow:
         }
     });
 
-    // Wait for core to connect to gRPC and send private cookies encryption key.
-    let Some(key) = rx.recv().await else {
-        return Err(anyhow::Error::msg("http channel closed"));
-    };
-
     // build application
     debug!("Setting up API server");
     let shared_state = AppState {
-        key,
+        cookie_key,
         grpc_server,
         remote_mfa_sessions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
         url: env_config.url.clone(),
@@ -309,6 +335,10 @@ pub async fn run_server(env_config: EnvConfig, config: Configuration) -> anyhow:
                 .route("/info", get(app_info)),
         )
         .fallback_service(get(handle_404))
+        .layer(middleware::from_fn_with_state(
+            shared_state.clone(),
+            ensure_configured,
+        ))
         .layer(middleware::map_response(powered_by_header))
         .layer(middleware::from_fn_with_state(
             shared_state.clone(),