diff --git a/src/config.rs b/src/config.rs index 0cb4867d..e099930a 100644 --- a/src/config.rs +++ b/src/config.rs @@ -15,6 +15,10 @@ fn default_log_level() -> String { String::from("info") } +fn default_adoption_timeout() -> u64 { + 10 +} + fn default_syslog_socket() -> PathBuf { PathBuf::from("/var/run/log") } @@ -118,6 +122,17 @@ pub struct Config { default_value = "/etc/defguard/certs" )] pub cert_dir: PathBuf, + + /// Time limit in minutes for the auto-adoption process. + /// After this time, gateway 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 Config { @@ -126,6 +141,11 @@ impl Config { Duration::from_secs(self.stats_period) } + #[must_use] + pub fn adoption_timeout(&self) -> Duration { + Duration::from_secs(self.adoption_timeout * 60) + } + /// Return [`SocketAddr`] for gRPC to listen on. #[must_use] pub(crate) fn grpc_socket(&self) -> SocketAddr { @@ -158,6 +178,7 @@ impl Default for Config { disable_firewall_management: false, http_bind_address: None, cert_dir: PathBuf::from("/etc/defguard/certs"), + adoption_timeout: 5, } } } diff --git a/src/setup.rs b/src/setup.rs index c3936f54..440b7a7c 100644 --- a/src/setup.rs +++ b/src/setup.rs @@ -1,4 +1,7 @@ -use std::sync::{Arc, Mutex}; +use std::sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, +}; use defguard_version::{Version, server::DefguardVersionLayer}; use tokio::{ @@ -67,6 +70,7 @@ pub struct GatewaySetupServer { current_session_token: Arc>>, setup_tx: Arc>>>, setup_rx: Arc>>, + adoption_expired: Arc, } impl Clone for GatewaySetupServer { @@ -77,6 +81,7 @@ impl Clone for GatewaySetupServer { current_session_token: Arc::clone(&self.current_session_token), setup_tx: Arc::clone(&self.setup_tx), setup_rx: Arc::clone(&self.setup_rx), + adoption_expired: Arc::clone(&self.adoption_expired), } } } @@ -91,6 +96,7 @@ impl GatewaySetupServer { current_session_token: Arc::new(Mutex::new(None)), setup_tx: Arc::new(tokio::sync::Mutex::new(Some(setup_tx))), setup_rx: Arc::new(tokio::sync::Mutex::new(setup_rx)), + adoption_expired: Arc::new(AtomicBool::new(false)), } } @@ -100,7 +106,25 @@ impl GatewaySetupServer { let setup_rx = Arc::clone(&self.setup_rx); let addr = config.grpc_socket(); - info!("Starting Gateway setup server on {addr} and awaiting configuration from Core"); + let adoption_timeout = config.adoption_timeout(); + info!( + "Starting Gateway setup server on {addr} and awaiting configuration from Core for {} min", + adoption_timeout.as_secs() / 60 + ); + + 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!( + "Gateway adoption expired and is now blocked. Restart the Gateway to enable auto-adoption." + ); + } + _ = cancel_rx => {} + } + }); server_builder .add_service( @@ -122,6 +146,11 @@ impl GatewaySetupServer { }) .await?; + // Skip blocking Gateway adoption if adoption was already done + if server_config.is_some() { + let _ = cancel_tx.send(()); + } + server_config.ok_or_else(|| { GatewayError::SetupError("Failed to receive setup configuration from Core".into()) }) @@ -175,6 +204,11 @@ impl gateway_setup_server::GatewaySetup for GatewaySetupServer { #[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 = "Gateway adoption expired and is now blocked. Restart the Gateway 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"));