From 1292d8a3fd851790bcb9274725f823c2a4e395d9 Mon Sep 17 00:00:00 2001
From: Aleksander <170264518+t-aleksander@users.noreply.github.com>
Date: Tue, 27 Jan 2026 10:45:15 +0100
Subject: [PATCH] fix healthcheck endpoint

---
 src/config.rs |  2 +-
 src/http.rs   | 26 ++++++++++++++++++++++----
 src/main.rs   | 40 +++++++++++++++++-----------------------
 3 files changed, 40 insertions(+), 28 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 2cc5f88..8449a2b 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -9,7 +9,7 @@ fn default_url() -> Url {
     Url::parse("http://localhost:8080").unwrap()
 }
 
-#[derive(Parser, Debug, Deserialize)]
+#[derive(Parser, Debug, Deserialize, Clone)]
 #[command(version)]
 pub struct EnvConfig {
     // port the API server will listen on
diff --git a/src/http.rs b/src/http.rs
index 73ecd61..e3753d4 100644
--- a/src/http.rs
+++ b/src/http.rs
@@ -232,7 +232,11 @@ async fn ensure_configured(
     next.run(request).await
 }
 
-pub async fn run_server(env_config: EnvConfig, config: Configuration) -> anyhow::Result<()> {
+pub async fn run_server(
+    env_config: EnvConfig,
+    config: Option<Configuration>,
+    logs_rx: Option<LogsReceiver>,
+) -> anyhow::Result<()> {
     info!("Starting Defguard Proxy server");
     debug!("Using config: {env_config:?}");
 
@@ -243,12 +247,26 @@ pub async fn run_server(env_config: EnvConfig, config: Configuration) -> anyhow:
     let grpc_server = ProxyServer::new(Arc::clone(&cookie_key));
 
     let server_clone = grpc_server.clone();
-    grpc_server.configure(config);
+    let env_config_clone = env_config.clone();
 
     // Start gRPC server.
-    // TODO: Wait with spawning the HTTP server until gRPC server is ready.
-    debug!("Spawning gRPC server");
+    debug!("Spawning gRPC server task");
     tasks.spawn(async move {
+        let proxy_configuration = if let Some(conf) = config {
+            debug!("Using existing gRPC certificates, skipping setup process");
+            conf
+        } else if let Some(logs_rx) = logs_rx {
+            info!("gRPC certificates not found, running setup process");
+            let conf = run_setup(&env_config_clone, logs_rx).await?;
+            info!("Setup process completed successfully");
+            conf
+        } else {
+            anyhow::bail!(
+                "gRPC certificates not found and logs receiver not available for setup process"
+            );
+        };
+
+        server_clone.configure(proxy_configuration);
         loop {
             info!("Starting gRPC server...");
             let server_to_run = server_clone.clone();
diff --git a/src/main.rs b/src/main.rs
index fe511ac..f6e0c85 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,7 +3,7 @@ use std::{fs::read_to_string, sync::Arc};
 use defguard_proxy::{
     config::get_env_config,
     grpc::Configuration,
-    http::{run_server, run_setup, GRPC_CERT_NAME, GRPC_KEY_NAME},
+    http::{run_server, GRPC_CERT_NAME, GRPC_KEY_NAME},
     logging::init_tracing,
     VERSION,
 };
@@ -24,7 +24,16 @@ async fn main() -> anyhow::Result<()> {
         read_to_string(cert_dir.join(GRPC_KEY_NAME)).ok(),
     );
 
-    let needs_setup = grpc_cert.is_none() || grpc_key.is_none();
+    let proxy_configuration = if let (Some(grpc_cert), Some(grpc_key)) = (grpc_cert, grpc_key) {
+        Some(Configuration {
+            grpc_cert_pem: grpc_cert,
+            grpc_key_pem: grpc_key,
+        })
+    } else {
+        None
+    };
+
+    let needs_setup = proxy_configuration.is_none();
 
     // TODO: The channel size may need to be adjusted or some other approach should be used
     // to avoid dropping log messages.
@@ -39,28 +48,13 @@ async fn main() -> anyhow::Result<()> {
     // read config from env
     tracing::info!("Starting ... version v{}", VERSION);
 
-    let proxy_configuration = if needs_setup {
-        if let Some(logs_rx) = logs_rx {
-            tracing::info!("gRPC certificates not found, running setup process");
-            let proxy_configuration = run_setup(&env_config, Arc::new(Mutex::new(logs_rx))).await?;
-            tracing::info!("Setup process completed successfully");
-            proxy_configuration
-        } else {
-            anyhow::bail!(
-                "gRPC certificates not found and logs receiver not available for setup process"
-            );
-        }
-    } else if let (Some(grpc_cert), Some(grpc_key)) = (grpc_cert, grpc_key) {
-        Configuration {
-            grpc_cert_pem: grpc_cert,
-            grpc_key_pem: grpc_key,
-        }
-    } else {
-        anyhow::bail!("Failed to load gRPC certificates");
-    };
-
     // run API web server
-    run_server(env_config, proxy_configuration).await?;
+    run_server(
+        env_config,
+        proxy_configuration,
+        logs_rx.map(|r| Arc::new(Mutex::new(r))),
+    )
+    .await?;
 
     Ok(())
 }