From 3101c97f08e715564144e4762e9077d099e424c8 Mon Sep 17 00:00:00 2001 From: Yousef Hussein Date: Fri, 24 Jul 2026 18:00:10 +0300 Subject: [PATCH] feat(config): Revoke Vault token on shutdown Complete the Vault token lease lifecycle. The runtime already leases and renews its token, but never revoked it: on shutdown the token lingered valid until its TTL expired. Add VaultRuntime::revoke (auth/token/revoke-self) and drive it from graceful shutdown. ConfigManager now owns a CancellationToken and the watcher JoinHandle; ConfigManager::shutdown cancels the watch loop and awaits it, and the loop revokes the token before exiting. Keystone's terminate hook calls it, so revocation completes before the process exits. Revoking invalidates the token (and any leases created with it) immediately instead of leaving valid credentials outstanding. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Yousef Hussein --- Cargo.lock | 1 + crates/config/Cargo.toml | 1 + crates/config/src/lib.rs | 81 +++++++++++++++++++++++++++++++++++-- crates/config/src/vault.rs | 22 ++++++++++ crates/core/src/keystone.rs | 3 ++ 5 files changed, 104 insertions(+), 4 deletions(-) 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(()) } }