diff --git a/Cargo.lock b/Cargo.lock index 008edfcc4..4cc9dd217 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4877,6 +4877,7 @@ dependencies = [ "tempfile", "thiserror 2.0.18", "tokio", + "tokio-util", "tracing", "url", "url-macro", diff --git a/crates/config/Cargo.toml b/crates/config/Cargo.toml index 9e29af9c4..5a7ff31f0 100644 --- a/crates/config/Cargo.toml +++ b/crates/config/Cargo.toml @@ -24,6 +24,7 @@ secrecy = { workspace = true, features = ["serde"] } serde = { workspace = true, features = ["derive"] } serde_json.workspace = true tokio = { workspace = true, features = ["fs", "macros", "rt", "sync", "time"] } +tokio-util.workspace = true tracing.workspace = true url = { workspace = true, features = ["serde"] } url-macro.workspace = true diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 7c29b8108..a14aa10d4 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -57,7 +57,9 @@ use config::{File, FileFormat}; use eyre::{Report, WrapErr}; use notify::{RecommendedWatcher, RecursiveMode, Watcher}; use serde::Deserialize; -use tokio::sync::RwLock; +use tokio::sync::{Mutex, RwLock}; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; use tracing::error; use validator::Validate; @@ -530,6 +532,12 @@ pub struct ConfigManager { pub config: Arc>, /// Notify listeners that something changed. pub notify_tx: tokio::sync::broadcast::Sender<()>, + /// Signals the background watcher to stop and run its teardown (e.g. + /// revoking the Vault token) on graceful shutdown. + shutdown: CancellationToken, + /// Handle to the spawned watcher task, awaited by [`Self::shutdown`] so + /// teardown completes before the process exits. + watcher_handle: Mutex>>, } impl ConfigManager { @@ -539,9 +547,24 @@ impl ConfigManager { Arc::new(Self { config: Arc::new(RwLock::new(config)), notify_tx, + shutdown: CancellationToken::new(), + watcher_handle: Mutex::new(None), }) } + /// Gracefully stop the background watcher. + /// + /// Cancels the watch loop and awaits its completion. When the + /// configuration is Vault-backed, the loop revokes the Vault token as + /// part of its teardown before this returns. Safe to call on an + /// unwatched manager (no-op) and idempotent across repeated calls. + pub async fn shutdown(&self) { + self.shutdown.cancel(); + if let Some(handle) = self.watcher_handle.lock().await.take() { + let _ = handle.await; + } + } + /// Initializes the config, starts the background watcher, /// and returns the manager for the live state. pub async fn watched(config_path: impl Into) -> Result, Report> { @@ -551,16 +574,20 @@ impl ConfigManager { // Initial Load let initial = Config::load_all_with_vault_state(&config_path).await?; + let shutdown = CancellationToken::new(); let manager = Arc::new(Self { config: Arc::new(RwLock::new(initial.config)), notify_tx, + shutdown: shutdown.clone(), + watcher_handle: Mutex::new(None), }); // Spawn Background Watcher let manager_clone = Arc::clone(&manager); - tokio::spawn(async move { - Self::watch_loop(manager_clone, config_path, initial.vault).await; + let handle = tokio::spawn(async move { + Self::watch_loop(manager_clone, config_path, initial.vault, shutdown).await; }); + *manager.watcher_handle.lock().await = Some(handle); Ok(manager) } @@ -572,6 +599,7 @@ impl ConfigManager { manager: Arc, config_path: PathBuf, mut vault_runtime: Option, + shutdown: CancellationToken, ) { let (sync_tx, mut sync_rx) = tokio::sync::mpsc::channel(1); @@ -616,6 +644,9 @@ impl ConfigManager { } }; tokio::select! { + () = shutdown.cancelled() => { + break; + } event = sync_rx.recv() => { if event.is_none() { break; @@ -667,6 +698,16 @@ impl ConfigManager { } } } + + // The loop exited (graceful shutdown or the watcher channel closing). + // For a Vault-backed configuration, revoke the token so it is + // invalidated immediately instead of lingering valid until its TTL + // expires. + if let Some(runtime) = &vault_runtime + && runtime.revoke().await.is_err() + { + error!("Vault token revocation on shutdown failed"); + } } async fn apply_loaded( @@ -702,7 +743,7 @@ mod tests { use tokio::time::{Duration, sleep, timeout}; use super::*; - use crate::vault::tests::{mock_lookup, mock_metadata, mock_renew, mock_secret}; + use crate::vault::tests::{mock_lookup, mock_metadata, mock_renew, mock_revoke, mock_secret}; // `Config::new` is async, but these tests drive it from the synchronous // `temp_env::with_var` closure API, so run it to completion on a local @@ -972,6 +1013,38 @@ mod tests { ); } + #[tokio::test] + #[parallel] + async fn test_vault_token_revoked_on_shutdown() { + let server = MockServer::start(); + let _lookup = mock_lookup(&server, false, 60); + let _metadata = mock_metadata(&server, 1); + let _secret = mock_secret(&server, 1, json!({"password": "version-one"})); + let revoke = mock_revoke(&server); + let mut config_file = NamedTempFile::with_suffix(".conf").unwrap(); + write_vault_config(&mut config_file, &server, 60); + + let manager = ConfigManager::watched(config_file.path()).await.unwrap(); + manager.shutdown().await; + + assert_eq!(revoke.calls(), 1); + } + + #[tokio::test] + #[parallel] + async fn test_shutdown_without_vault_is_noop() { + let mut config_file = NamedTempFile::with_suffix(".conf").unwrap(); + writeln!( + config_file, + "[auth]\nmethods = []\n[database]\nconnection = sqlite://\n" + ) + .unwrap(); + + let manager = ConfigManager::watched(config_file.path()).await.unwrap(); + // Must return promptly and not panic for a non-Vault configuration. + manager.shutdown().await; + } + #[tokio::test] #[parallel] async fn test_vault_version_reload_and_last_known_good_retention() { diff --git a/crates/config/src/vault.rs b/crates/config/src/vault.rs index 9019ab61f..f409dad51 100644 --- a/crates/config/src/vault.rs +++ b/crates/config/src/vault.rs @@ -84,6 +84,8 @@ pub(crate) enum VaultConfigError { MetadataRead, #[error("Vault token renewal failed")] TokenRenewal, + #[error("Vault token revocation failed")] + TokenRevocation, #[error("configuration is invalid after resolving Vault references")] ResolvedConfigurationInvalid, } @@ -111,6 +113,19 @@ impl VaultRuntime { .unwrap_or(self.next_poll) } + /// Revoke the Vault token via `auth/token/revoke-self`. + /// + /// Called during graceful shutdown to invalidate the token (and any + /// leases created with it) immediately, rather than leaving it valid + /// until its TTL expires. Best-effort: the process is stopping, so a + /// failure is surfaced to the caller for logging but cannot be retried. + pub(crate) async fn revoke(&self) -> Result<(), VaultConfigError> { + self.client + .revoke() + .await + .map_err(|_| VaultConfigError::TokenRevocation) + } + pub(crate) async fn renew_if_due(&mut self) -> Result<(), VaultConfigError> { let Some(renewal) = &self.renewal else { return Ok(()); @@ -470,6 +485,13 @@ pub(crate) mod tests { }) } + pub(crate) fn mock_revoke(server: &MockServer) -> Mock<'_> { + server.mock(|when, then| { + when.method(POST).path("/v1/auth/token/revoke-self"); + then.status(204); + }) + } + fn test_vault(server: &MockServer, refresh_interval_seconds: u64) -> VaultSection { VaultSection { address: Url::parse(&server.base_url()).unwrap(), diff --git a/crates/core/src/keystone.rs b/crates/core/src/keystone.rs index fe5adffe7..48ea97b43 100644 --- a/crates/core/src/keystone.rs +++ b/crates/core/src/keystone.rs @@ -227,6 +227,9 @@ impl Service { /// - `Err(KeystoneError)` if an error occurred during termination. pub async fn terminate(&self) -> Result<(), KeystoneError> { info!("Terminating Keystone"); + // Stop the config watcher and, for a Vault-backed configuration, + // revoke the Vault token before the process exits. + self.config_manager.shutdown().await; Ok(()) } }