From df74df03fd1e96e4727dfe9f61e1c1dfcc6c14dc Mon Sep 17 00:00:00 2001 From: Kuba <78603704+jakub-tldr@users.noreply.github.com> Date: Thu, 16 Apr 2026 18:37:23 +0200 Subject: [PATCH 1/6] add config --- src/config.rs | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/src/config.rs b/src/config.rs index 7f53d25..23b0813 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,4 +1,4 @@ -use std::{fs::read_to_string, net::IpAddr, path::PathBuf}; +use std::{fs::read_to_string, net::IpAddr, path::PathBuf, time::Duration}; use clap::Parser; use log::LevelFilter; @@ -9,6 +9,10 @@ fn default_url() -> Url { Url::parse("http://localhost:8080").unwrap() } +fn default_adoption_timeout() -> u64 { + 10 +} + #[derive(Parser, Debug, Deserialize, Clone)] #[command(version)] pub struct EnvConfig { @@ -88,6 +92,24 @@ pub struct EnvConfig { /// Use Let's Encrypt staging environment for ACME issuance. #[arg(long, env = "DEFGUARD_PROXY_ACME_STAGING", default_value_t = false)] pub acme_staging: bool, + + /// Time limit in minutes for the auto-adoption process. + /// After this time Edge will reject adoption attempts until restarted. + #[arg( + long, + short = 't', + env = "DEFGUARD_ADOPTION_TIMEOUT", + default_value = "10" + )] + #[serde(default = "default_adoption_timeout")] + pub adoption_timeout: u64, +} + +impl EnvConfig { + #[must_use] + pub fn adoption_timeout(&self) -> Duration { + Duration::from_secs(self.adoption_timeout * 60) + } } #[derive(thiserror::Error, Debug)] From a01fa32118f2326d97d0640463dcf994b5546fa7 Mon Sep 17 00:00:00 2001 From: Kuba <78603704+jakub-tldr@users.noreply.github.com> Date: Thu, 16 Apr 2026 18:37:50 +0200 Subject: [PATCH 2/6] block adopting edge after specified time --- src/http.rs | 15 +++++++++------ src/setup.rs | 39 +++++++++++++++++++++++++++++++++++---- 2 files changed, 44 insertions(+), 10 deletions(-) diff --git a/src/http.rs b/src/http.rs index 9085d54..10d34e4 100644 --- a/src/http.rs +++ b/src/http.rs @@ -205,12 +205,15 @@ pub async fn run_setup( cert_dir.display() ); let configuration = setup_server - .await_initial_setup(SocketAddr::new( - env_config - .grpc_bind_address - .unwrap_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED)), - env_config.grpc_port, - )) + .await_initial_setup( + SocketAddr::new( + env_config + .grpc_bind_address + .unwrap_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED)), + env_config.grpc_port, + ), + env_config, + ) .await?; info!("Generated new gRPC TLS certificates and signed by Defguard Core"); diff --git a/src/setup.rs b/src/setup.rs index 684639b..3d92dde 100644 --- a/src/setup.rs +++ b/src/setup.rs @@ -1,18 +1,22 @@ use std::{ net::SocketAddr, - sync::{Arc, LazyLock, Mutex}, + sync::{ + Arc, LazyLock, Mutex, + atomic::{AtomicBool, Ordering}, + }, }; use defguard_version::{ DefguardComponent, Version, server::{DefguardVersionLayer, grpc::DefguardVersionInterceptor}, }; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, oneshot}; use tokio_stream::wrappers::UnboundedReceiverStream; use tonic::{Request, Response, Status, transport::Server}; use crate::{ CommsChannel, LogsReceiver, MIN_CORE_VERSION, VERSION, + config::EnvConfig, error::ApiError, grpc::Configuration, proto::{CertificateInfo, DerPayload, LogEntry, proxy_setup_server}, @@ -32,6 +36,7 @@ pub(crate) struct ProxySetupServer { key_pair: Arc>>, logs_rx: LogsReceiver, current_session_token: Arc>>, + adoption_expired: Arc, } impl Clone for ProxySetupServer { @@ -40,6 +45,7 @@ impl Clone for ProxySetupServer { key_pair: Arc::clone(&self.key_pair), logs_rx: Arc::clone(&self.logs_rx), current_session_token: Arc::clone(&self.current_session_token), + adoption_expired: Arc::clone(&self.adoption_expired), } } } @@ -50,6 +56,7 @@ impl ProxySetupServer { key_pair: Arc::new(Mutex::new(None)), logs_rx, current_session_token: Arc::new(Mutex::new(None)), + adoption_expired: Arc::new(AtomicBool::new(false)), } } @@ -62,11 +69,29 @@ impl ProxySetupServer { pub(crate) async fn await_initial_setup( &self, addr: SocketAddr, + config: &EnvConfig, ) -> Result { - info!("gRPC waiting for setup connection from Core on {addr}"); + let adoption_timeout = config.adoption_timeout(); + info!( + "gRPC waiting for setup connection from Core on {addr} for {} seconds", + adoption_timeout.as_secs() + ); + let adoption_expired = Arc::clone(&self.adoption_expired); + let (cancel_tx, cancel_rx) = oneshot::channel::<()>(); + tokio::spawn(async move { + tokio::select! { + _ = tokio::time::sleep(adoption_timeout) => { + adoption_expired.store(true, Ordering::Relaxed); + error!( + "Edge adoption expired and is now blocked. Restart the Edge to enable auto-adoption." + ); + } + _ = cancel_rx => {} + } + }); let own_version = Version::parse(VERSION)?; - debug!("Proxy version: {}", VERSION); + debug!("Edge version: {}", VERSION); let config_slot: Arc>> = Arc::new(tokio::sync::Mutex::new(None)); @@ -107,6 +132,7 @@ impl ProxySetupServer { ApiError::Unexpected("No configuration received after setup".into()) })?; + let _ = cancel_tx.send(()); Ok(configuration) } @@ -158,6 +184,11 @@ impl proxy_setup_server::ProxySetup for ProxySetupServer { #[instrument(skip(self, request))] async fn start(&self, request: Request<()>) -> Result, Status> { debug!("Core initiated setup process, preparing to stream logs"); + if self.adoption_expired.load(Ordering::Relaxed) { + let error_message = "Edge adoption expired and is now blocked. Restart the Edge to enable auto-adoption."; + error!("{error_message}"); + return Err(Status::failed_precondition(error_message)); + } if self.is_setup_in_progress() { error!("Setup already in progress, rejecting new setup request"); return Err(Status::resource_exhausted("Setup already in progress")); From 88ff7445dae1e21dc34de2a46d618a67349f2ab0 Mon Sep 17 00:00:00 2001 From: Kuba <78603704+jakub-tldr@users.noreply.github.com> Date: Thu, 16 Apr 2026 18:45:22 +0200 Subject: [PATCH 3/6] edge --> proxy --- src/setup.rs | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/setup.rs b/src/setup.rs index 3d92dde..e0c4398 100644 --- a/src/setup.rs +++ b/src/setup.rs @@ -140,7 +140,7 @@ impl ProxySetupServer { let in_progress = self .current_session_token .lock() - .expect("Failed to acquire lock on current session token during proxy setup") + .expect("Failed to acquire lock on current session token during Edge setup") .is_some(); debug!("Setup in progress check: {}", in_progress); in_progress @@ -150,7 +150,7 @@ impl ProxySetupServer { debug!("Terminating setup session"); self.current_session_token .lock() - .expect("Failed to acquire lock on current session token during proxy setup") + .expect("Failed to acquire lock on current session token during Edge setup") .take(); debug!("Setup session terminated"); } @@ -159,7 +159,7 @@ impl ProxySetupServer { debug!("Establishing new setup session with Core"); self.current_session_token .lock() - .expect("Failed to acquire lock on current session token during proxy setup") + .expect("Failed to acquire lock on current session token during Edge setup") .replace(token); debug!("Setup session established"); } @@ -169,7 +169,7 @@ impl ProxySetupServer { let is_valid = (*self .current_session_token .lock() - .expect("Failed to acquire lock on current session token during proxy setup")) + .expect("Failed to acquire lock on current session token during Edge setup")) .as_ref() .is_some_and(|t| t == token); debug!("Authorization validation result: {}", is_valid); @@ -205,7 +205,7 @@ impl proxy_setup_server::ProxySetup for ProxySetupServer { debug!("Setup session authenticated successfully"); self.initialize_setup_session(token.to_string()); - debug!("Preparing to forward Proxy logs to Core in real-time"); + debug!("Preparing to forward Edge logs to Core in real-time"); let logs_rx = self.logs_rx.clone(); let (tx, rx) = mpsc::unbounded_channel(); @@ -239,7 +239,7 @@ impl proxy_setup_server::ProxySetup for ProxySetupServer { self_clone.clear_setup_session(); }); - debug!("Log stream established, Core will now receive real-time Proxy logs"); + debug!("Log stream established, Core will now receive real-time Edge logs"); Ok(Response::new(UnboundedReceiverStream::new(rx))) } @@ -305,7 +305,7 @@ impl proxy_setup_server::ProxySetup for ProxySetupServer { self.key_pair .lock() - .expect("Failed to acquire lock on key pair during proxy setup when trying to store generated key pair") + .expect("Failed to acquire lock on key pair during Edge setup when trying to store generated key pair") .replace(key_pair); debug!("Encoding Certificate Signing Request for transmission"); @@ -362,17 +362,17 @@ impl proxy_setup_server::ProxySetup for ProxySetupServer { let key_pair = self .key_pair .lock() - .expect("Failed to acquire lock on key pair during proxy setup when trying to receive certificate") + .expect("Failed to acquire lock on key pair during Edge setup when trying to receive certificate") .take(); if let Some(kp) = key_pair { kp } else { error!( - "Key pair not found during Proxy setup. Key pair generation step might have failed." + "Key pair not found during Edge setup. Key pair generation step might have failed." ); self.clear_setup_session(); return Err(Status::internal( - "Key pair not found during Proxy setup. Key pair generation step might have failed.", + "Key pair not found during Edge setup. Key pair generation step might have failed.", )); } }; @@ -384,7 +384,7 @@ impl proxy_setup_server::ProxySetup for ProxySetupServer { debug!("Passing configuration to gRPC server for finalization"); match SETUP_CHANNEL.0.lock().await.send(Some(configuration)).await { - Ok(()) => info!("Proxy configuration passed to gRPC server successfully"), + Ok(()) => info!("Edge configuration passed to gRPC server successfully"), Err(err) => { error!("Failed to send configuration to gRPC server: {err}"); self.clear_setup_session(); From f8c975c4f1b0c8a2792679b8d2508ed240a30f54 Mon Sep 17 00:00:00 2001 From: Kuba <78603704+jakub-tldr@users.noreply.github.com> Date: Fri, 17 Apr 2026 10:57:45 +0200 Subject: [PATCH 4/6] apply suggestions --- src/setup.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/setup.rs b/src/setup.rs index e0c4398..f24d156 100644 --- a/src/setup.rs +++ b/src/setup.rs @@ -66,6 +66,11 @@ impl ProxySetupServer { /// `GetCsr`, `SendCert`. The server shuts down as soon as `SendCert` deposits a /// `Configuration` into `SETUP_CHANNEL`, after which this function returns the received /// gRPC configuration (locally generated key pair and remotely signed certificate). + /// + /// A timeout is started in the background using `config.adoption_timeout()`. If the timeout + /// elapses before setup completes, the `adoption_expired` flag is set and subsequent `Start` + /// requests are rejected with `failed_precondition` until the Edge is restarted. + /// On successful adoption the timeout is cancelled. pub(crate) async fn await_initial_setup( &self, addr: SocketAddr, @@ -73,8 +78,8 @@ impl ProxySetupServer { ) -> Result { let adoption_timeout = config.adoption_timeout(); info!( - "gRPC waiting for setup connection from Core on {addr} for {} seconds", - adoption_timeout.as_secs() + "gRPC waiting for setup connection from Core on {addr} for {} minutes", + adoption_timeout.as_secs() / 60 ); let adoption_expired = Arc::clone(&self.adoption_expired); @@ -132,7 +137,9 @@ impl ProxySetupServer { ApiError::Unexpected("No configuration received after setup".into()) })?; + // Skip blocking Edge adoption if adoption was already done let _ = cancel_tx.send(()); + Ok(configuration) } From 56ddb1963f67d980ce020facced21ab828874c57 Mon Sep 17 00:00:00 2001 From: Kuba <78603704+jakub-tldr@users.noreply.github.com> Date: Fri, 17 Apr 2026 11:10:42 +0200 Subject: [PATCH 5/6] minutes -> min --- src/setup.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/setup.rs b/src/setup.rs index f24d156..d231316 100644 --- a/src/setup.rs +++ b/src/setup.rs @@ -68,7 +68,7 @@ impl ProxySetupServer { /// gRPC configuration (locally generated key pair and remotely signed certificate). /// /// A timeout is started in the background using `config.adoption_timeout()`. If the timeout - /// elapses before setup completes, the `adoption_expired` flag is set and subsequent `Start` + /// elapses before setup completes, the `adoption_expired` flag is set and incoming `Start` /// requests are rejected with `failed_precondition` until the Edge is restarted. /// On successful adoption the timeout is cancelled. pub(crate) async fn await_initial_setup( @@ -78,7 +78,7 @@ impl ProxySetupServer { ) -> Result { let adoption_timeout = config.adoption_timeout(); info!( - "gRPC waiting for setup connection from Core on {addr} for {} minutes", + "gRPC waiting for setup connection from Core on {addr} for {} min", adoption_timeout.as_secs() / 60 ); From f4992e5672d988529f9eff40512483d0208c95d7 Mon Sep 17 00:00:00 2001 From: Kuba <78603704+jakub-tldr@users.noreply.github.com> Date: Fri, 17 Apr 2026 11:21:28 +0200 Subject: [PATCH 6/6] auto-adoption -> adoption --- src/setup.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/setup.rs b/src/setup.rs index d231316..344ee34 100644 --- a/src/setup.rs +++ b/src/setup.rs @@ -89,7 +89,7 @@ impl ProxySetupServer { _ = tokio::time::sleep(adoption_timeout) => { adoption_expired.store(true, Ordering::Relaxed); error!( - "Edge adoption expired and is now blocked. Restart the Edge to enable auto-adoption." + "Edge adoption expired and is now blocked. Restart the Edge to enable adoption." ); } _ = cancel_rx => {} @@ -192,7 +192,8 @@ impl proxy_setup_server::ProxySetup for ProxySetupServer { async fn start(&self, request: Request<()>) -> Result, Status> { debug!("Core initiated setup process, preparing to stream logs"); if self.adoption_expired.load(Ordering::Relaxed) { - let error_message = "Edge adoption expired and is now blocked. Restart the Edge to enable auto-adoption."; + let error_message = + "Edge adoption expired and is now blocked. Restart the Edge to enable adoption."; error!("{error_message}"); return Err(Status::failed_precondition(error_message)); }