Multiproxy private cookies

src/grpc.rs

@@ -21,20 +21,23 @@ use crate::{
 
 // connected clients
 type ClientMap = HashMap<SocketAddr, mpsc::UnboundedSender<Result<CoreRequest, Status>>>;
+static COOKIE_KEY_HEADER: &str = "dg-cookie-key-bin";
 
 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<Vec<u8>>,
     pub(crate) connected: Arc<AtomicBool>,
     pub(crate) core_version: Arc<Mutex<Option<Version>>>,
 }
 
 impl ProxyServer {
     #[must_use]
     /// Create new `ProxyServer`.
-    pub(crate) fn new() -> Self {
+    pub(crate) fn new(http_channel: mpsc::UnboundedSender<Vec<u8>>) -> 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<Streaming<CoreResponse>>,
     ) -> Result<Response<Self::BidiStream>, Status> {
+        let cookie_key = request.metadata().get_bin(COOKIE_KEY_HEADER).unwrap();
+        let key = (cookie_key.to_bytes().unwrap())
+            .into_iter()
+            .collect::<Vec<_>>();
+        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"));

src/http.rs

@@ -22,7 +22,11 @@ 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::{
@@ -172,18 +176,12 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> {
 
     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::<Vec<u8>>();
 
-    // 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 +197,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,14 +223,27 @@ 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)
             .await
             .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: {}",