Auto adoption time limit (#310)

Author: jakub-tldr

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,
         }
     }
 }

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<Mutex<Option<String>>>,
     setup_tx: Arc<tokio::sync::Mutex<Option<oneshot::Sender<TlsConfig>>>>,
     setup_rx: Arc<tokio::sync::Mutex<oneshot::Receiver<TlsConfig>>>,
+    adoption_expired: Arc<AtomicBool>,
 }
 
 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<Response<Self::StartStream>, 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"));