Cookie key proto (#235)

* send cookie key via protos
* working key retrieval
* handle InitialInfo in message loop
* new protos

Author: j-chmielewski

proto

@@ -32,7 +32,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(crate) struct Configuration {
@@ -232,28 +231,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.cookie_key.write().unwrap() = Some(key);
-
         let (tx, rx) = mpsc::unbounded_channel();
         self.clients
             .lock()
@@ -266,22 +243,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

@@ -226,6 +226,7 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> {
     tasks.spawn(async move {
         let cert_dir = Path::new(&config.cert_dir);
         if !cert_dir.exists() {
+            debug!("Creating certs directory");
             tokio::fs::create_dir_all(cert_dir).await?;
         }