feat(kms): implement AppRole login with background renewal and fail-closed expiry

Implement the AppRoleLogin token source (vaultrs approle login +
renew-self) and wire lease-bound credentials through the provider:

- Each successful login/renewal installs a new client generation in the
  ArcSwap; in-flight requests finish on the generation they captured.
- A background renewal task refreshes at half the lease TTL: renewable
  tokens are renewed in place, everything else (or a failed renewal)
  falls back to a fresh login. Auth exchanges run under the typed retry
  policy (OpClass::Auth) and failed cycles retry on a fixed cadence, so
  the provider recovers once Vault does.
- Fail-closed: current() refuses to hand out a token inside the
  configured safety window of its expiry (default: one attempt timeout),
  returning CredentialsUnavailable instead of sending a request whose
  token may lapse mid-flight.
- Refreshes are single-flight: concurrent triggers for the same
  generation coalesce into one login.
- The renewal task's owner handle lives on the KMS service version:
  stop() shuts it down explicitly and reconfigure recycles it via
  cancel-on-drop when the old version is discarded.
- The secret_id file is re-read on every login attempt; missing or empty
  files fail the attempt without contacting Vault. Crate-owned copies of
  tokens and secret_ids are zeroized on drop, and Debug output of every
  credential-carrying type stays redacted (leak regression tests).

The renewal machinery is covered by paused-clock tests driving a
scripted token source: renew-at-half-TTL timing, login fallback,
fail-closed window entry and recovery, prompt task recycling, and
coalesced concurrent refreshes.
This commit is contained in:
overtrue
2026-07-31 00:44:08 +08:00
parent d618c78766
commit 124b32742c
4 changed files with 931 additions and 132 deletions
+45 -30
View File
@@ -14,7 +14,10 @@
//! Vault-based KMS backend implementation using vaultrs
use crate::backends::vault_credentials::{VaultClientHandle, VaultConnectionSettings, VaultCredentialProvider, token_source_for};
use crate::backends::vault_credentials::{
CredentialTaskHandle, VaultClientHandle, VaultConnectionSettings, VaultCredentialPolicy, VaultCredentialProvider,
token_source_for,
};
use crate::backends::{BackendInfo, KmsBackend, KmsClient};
use crate::config::{KmsConfig, VaultConfig};
use crate::encryption::{AesDekCrypto, DataKeyEnvelope, DekCrypto, generate_key_material};
@@ -32,7 +35,7 @@ use vaultrs::{api::kv2::requests::SetSecretRequestOptions, error::ClientError, k
/// Vault KMS client implementation
pub struct VaultKmsClient {
credentials: VaultCredentialProvider,
credentials: Arc<VaultCredentialProvider>,
config: VaultConfig,
/// Mount path for the KV engine (typically "kv" or "secret")
kv_mount: String,
@@ -161,15 +164,18 @@ fn decode_stored_key_material(key_id: &str, encrypted_material: &str) -> Result<
impl VaultKmsClient {
/// Create a new Vault KMS client
///
/// `attempt_timeout` caps every HTTP request issued through this client.
pub async fn new(config: VaultConfig, attempt_timeout: Duration) -> Result<Self> {
let source = token_source_for(&config.auth_method)?;
/// `kms_config` supplies the per-attempt timeout that caps every HTTP
/// request issued through this client, plus the retry and fail-closed
/// budgets for credential refresh.
pub async fn new(config: VaultConfig, kms_config: &KmsConfig) -> Result<Self> {
let settings = VaultConnectionSettings {
address: config.address.clone(),
namespace: config.namespace.clone(),
attempt_timeout,
attempt_timeout: kms_config.effective_timeout(),
};
let credentials = VaultCredentialProvider::new(settings, source).await?;
let source = token_source_for(&config.auth_method, &settings)?;
let policy = VaultCredentialPolicy::from_kms_config(kms_config, &config.auth_method);
let credentials = Arc::new(VaultCredentialProvider::new(settings, source, policy).await?);
info!(address = %config.address, "Vault KMS backend connected");
@@ -185,8 +191,9 @@ impl VaultKmsClient {
/// Snapshot the authenticated Vault client for a single request.
///
/// Every Vault call takes its own snapshot so a credential rotation
/// applies to subsequent calls without interrupting in-flight ones.
fn vault(&self) -> Arc<VaultClientHandle> {
/// applies to subsequent calls without interrupting in-flight ones. Fails
/// closed when the credentials could not be refreshed in time.
fn vault(&self) -> Result<Arc<VaultClientHandle>> {
self.credentials.current()
}
@@ -231,7 +238,7 @@ impl VaultKmsClient {
let path = self.key_version_path(key_id, version);
let record: VaultKeyVersionRecord =
kv2::read(&self.vault().client, &self.kv_mount, &path)
kv2::read(&self.vault()?.client, &self.kv_mount, &path)
.await
.map_err(|e| match e {
ClientError::ResponseWrapError => KmsError::key_version_not_found(key_id, version),
@@ -271,7 +278,7 @@ impl VaultKmsClient {
async fn get_key_data_versioned(&self, key_id: &str) -> Result<(u32, VaultKeyData)> {
let path = self.key_path(key_id);
let metadata = kv2::read_metadata(&self.vault().client, &self.kv_mount, &path)
let metadata = kv2::read_metadata(&self.vault()?.client, &self.kv_mount, &path)
.await
.map_err(|e| match e {
ClientError::ResponseWrapError => KmsError::key_not_found(key_id),
@@ -283,7 +290,7 @@ impl VaultKmsClient {
// Read the exact secret version from the metadata to keep the (cas, data)
// pair consistent even if another writer lands in between.
let key_data: VaultKeyData = kv2::read_version(&self.vault().client, &self.kv_mount, &path, metadata.current_version)
let key_data: VaultKeyData = kv2::read_version(&self.vault()?.client, &self.kv_mount, &path, metadata.current_version)
.await
.map_err(|e| match e {
ClientError::ResponseWrapError => KmsError::key_not_found(key_id),
@@ -303,7 +310,7 @@ impl VaultKmsClient {
let path = self.key_path(key_id);
let written =
kv2::set_with_options(&self.vault().client, &self.kv_mount, &path, key_data, SetSecretRequestOptions { cas })
kv2::set_with_options(&self.vault()?.client, &self.kv_mount, &path, key_data, SetSecretRequestOptions { cas })
.await
.map_err(|e| {
if is_cas_conflict(&e) {
@@ -327,7 +334,8 @@ impl VaultKmsClient {
async fn try_create_key_version_record(&self, key_id: &str, record: &VaultKeyVersionRecord) -> Result<bool> {
let path = self.key_version_path(key_id, record.version);
match kv2::set_with_options(&self.vault().client, &self.kv_mount, &path, record, SetSecretRequestOptions { cas: 0 }).await
match kv2::set_with_options(&self.vault()?.client, &self.kv_mount, &path, record, SetSecretRequestOptions { cas: 0 })
.await
{
Ok(_) => Ok(true),
Err(e) if is_cas_conflict(&e) => Ok(false),
@@ -339,7 +347,7 @@ impl VaultKmsClient {
async fn store_key_data(&self, key_id: &str, key_data: &VaultKeyData) -> Result<()> {
let path = self.key_path(key_id);
kv2::set(&self.vault().client, &self.kv_mount, &path, key_data)
kv2::set(&self.vault()?.client, &self.kv_mount, &path, key_data)
.await
.map_err(|e| KmsError::backend_error(format!("Failed to store key in Vault: {e}")))?;
@@ -389,7 +397,7 @@ impl VaultKmsClient {
async fn get_key_data(&self, key_id: &str) -> Result<VaultKeyData> {
let path = self.key_path(key_id);
let secret: VaultKeyData = kv2::read(&self.vault().client, &self.kv_mount, &path)
let secret: VaultKeyData = kv2::read(&self.vault()?.client, &self.kv_mount, &path)
.await
.map_err(|e| match e {
vaultrs::error::ClientError::ResponseWrapError => KmsError::key_not_found(key_id),
@@ -404,7 +412,7 @@ impl VaultKmsClient {
/// List all keys stored in Vault
async fn list_vault_keys(&self) -> Result<Vec<String>> {
// List keys under the prefix
match kv2::list(&self.vault().client, &self.kv_mount, &self.key_path_prefix).await {
match kv2::list(&self.vault()?.client, &self.kv_mount, &self.key_path_prefix).await {
Ok(keys) => {
let keys = filter_key_directory_entries(keys);
debug!("Found {} keys in Vault", keys.len());
@@ -431,11 +439,11 @@ impl VaultKmsClient {
// record still exists and the deletion can be retried. The reverse order
// would leave orphaned master key material in Vault after the key vanished.
let versions_dir = self.key_versions_dir(key_id);
match kv2::list(&self.vault().client, &self.kv_mount, &versions_dir).await {
match kv2::list(&self.vault()?.client, &self.kv_mount, &versions_dir).await {
Ok(versions) => {
for version in versions {
let version_path = format!("{versions_dir}/{version}");
kv2::delete_metadata(&self.vault().client, &self.kv_mount, &version_path)
kv2::delete_metadata(&self.vault()?.client, &self.kv_mount, &version_path)
.await
.map_err(|e| KmsError::backend_error(format!("Failed to delete key version record from Vault: {e}")))?;
}
@@ -447,7 +455,7 @@ impl VaultKmsClient {
// For this specific key path, we can safely delete the metadata
// since each key has its own unique path under the prefix
kv2::delete_metadata(&self.vault().client, &self.kv_mount, &path)
kv2::delete_metadata(&self.vault()?.client, &self.kv_mount, &path)
.await
.map_err(|e| match e {
vaultrs::error::ClientError::APIError { code: 404, .. } => KmsError::key_not_found(key_id),
@@ -865,10 +873,17 @@ impl VaultKmsBackend {
}
};
let client = VaultKmsClient::new(vault_config, config.effective_timeout()).await?;
let client = VaultKmsClient::new(vault_config, &config).await?;
Ok(Self { client })
}
/// Spawn the background credential renewal task for this backend, if its
/// auth method issues lease-bound tokens. The caller owns the returned
/// handle; dropping it cancels the task.
pub(crate) fn spawn_credential_renewal(&self) -> Option<CredentialTaskHandle> {
self.client.credentials.spawn_renewal_task()
}
/// Update key metadata in Vault storage
async fn update_key_metadata_in_storage(&self, key_id: &str, metadata: &KeyMetadata) -> Result<()> {
// Get the current key data from Vault
@@ -1157,7 +1172,7 @@ mod tests {
tls: None,
};
let client = VaultKmsClient::new(config, Duration::from_secs(30))
let client = VaultKmsClient::new(config, &KmsConfig::default())
.await
.expect("Failed to create Vault client");
@@ -1210,7 +1225,7 @@ mod tests {
#[tokio::test]
async fn test_key_version_paths_stay_under_the_key() {
let client = VaultKmsClient::new(integration_vault_config(), Duration::from_secs(30))
let client = VaultKmsClient::new(integration_vault_config(), &KmsConfig::default())
.await
.expect("client");
@@ -1295,7 +1310,7 @@ mod tests {
#[tokio::test]
async fn test_vault_kv2_backend_info_reports_at_rest_protection() {
let client = VaultKmsClient::new(integration_vault_config(), Duration::from_secs(30))
let client = VaultKmsClient::new(integration_vault_config(), &KmsConfig::default())
.await
.expect("client");
@@ -1327,7 +1342,7 @@ mod tests {
#[tokio::test]
#[ignore] // Requires a running Vault instance (dev mode)
async fn test_vault_kv2_decrypt_after_rotate() {
let client = VaultKmsClient::new(integration_vault_config(), Duration::from_secs(30))
let client = VaultKmsClient::new(integration_vault_config(), &KmsConfig::default())
.await
.expect("client");
@@ -1364,7 +1379,7 @@ mod tests {
#[tokio::test]
#[ignore] // Requires a running Vault instance (dev mode)
async fn test_vault_kv2_rotate_does_not_orphan_legacy_envelopes() {
let client = VaultKmsClient::new(integration_vault_config(), Duration::from_secs(30))
let client = VaultKmsClient::new(integration_vault_config(), &KmsConfig::default())
.await
.expect("client");
@@ -1403,7 +1418,7 @@ mod tests {
#[tokio::test]
#[ignore] // Requires a running Vault instance (dev mode)
async fn test_vault_kv2_envelope_version_tampering_fails_closed() {
let client = VaultKmsClient::new(integration_vault_config(), Duration::from_secs(30))
let client = VaultKmsClient::new(integration_vault_config(), &KmsConfig::default())
.await
.expect("client");
@@ -1446,7 +1461,7 @@ mod tests {
use std::sync::Arc;
let client = Arc::new(
VaultKmsClient::new(integration_vault_config(), Duration::from_secs(30))
VaultKmsClient::new(integration_vault_config(), &KmsConfig::default())
.await
.expect("client"),
);
@@ -1505,7 +1520,7 @@ mod tests {
// Regression: get_key_material previously "self-healed" a decrypt/length failure by
// minting a fresh random master key and overwriting the stored value — destroying the
// original key and making every DEK wrapped by it permanently undecryptable.
let client = VaultKmsClient::new(integration_vault_config(), Duration::from_secs(30))
let client = VaultKmsClient::new(integration_vault_config(), &KmsConfig::default())
.await
.expect("client");
@@ -1543,7 +1558,7 @@ mod tests {
// bootstrap case and silently generated + persisted a fresh master key on the
// read path. Empty material must instead fail closed as MaterialMissing and
// leave the stored record untouched.
let client = VaultKmsClient::new(integration_vault_config(), Duration::from_secs(30))
let client = VaultKmsClient::new(integration_vault_config(), &KmsConfig::default())
.await
.expect("client");
File diff suppressed because it is too large Load Diff
+41 -27
View File
@@ -14,7 +14,10 @@
//! Vault Transit-based KMS backend.
use crate::backends::vault_credentials::{VaultClientHandle, VaultConnectionSettings, VaultCredentialProvider, token_source_for};
use crate::backends::vault_credentials::{
CredentialTaskHandle, VaultClientHandle, VaultConnectionSettings, VaultCredentialPolicy, VaultCredentialProvider,
token_source_for,
};
use crate::backends::{BackendInfo, KmsBackend, KmsClient};
use crate::config::{KmsConfig, VaultTransitConfig};
use crate::encryption::{DataKeyEnvelope, generate_key_material};
@@ -129,7 +132,7 @@ impl From<TransitKeyMetadataPersisted> for TransitKeyMetadata {
}
pub struct VaultTransitKmsClient {
credentials: VaultCredentialProvider,
credentials: Arc<VaultCredentialProvider>,
config: VaultTransitConfig,
/// KV v2 mount path for persisting transit key metadata
metadata_kv_mount: String,
@@ -141,15 +144,18 @@ pub struct VaultTransitKmsClient {
impl VaultTransitKmsClient {
/// Create a new Vault Transit KMS client
///
/// `attempt_timeout` caps every HTTP request issued through this client.
pub async fn new(config: VaultTransitConfig, attempt_timeout: Duration) -> Result<Self> {
let source = token_source_for(&config.auth_method)?;
/// `kms_config` supplies the per-attempt timeout that caps every HTTP
/// request issued through this client, plus the retry and fail-closed
/// budgets for credential refresh.
pub async fn new(config: VaultTransitConfig, kms_config: &KmsConfig) -> Result<Self> {
let settings = VaultConnectionSettings {
address: config.address.clone(),
namespace: config.namespace.clone(),
attempt_timeout,
attempt_timeout: kms_config.effective_timeout(),
};
let credentials = VaultCredentialProvider::new(settings, source).await?;
let source = token_source_for(&config.auth_method, &settings)?;
let policy = VaultCredentialPolicy::from_kms_config(kms_config, &config.auth_method);
let credentials = Arc::new(VaultCredentialProvider::new(settings, source, policy).await?);
Ok(Self {
credentials,
@@ -163,8 +169,9 @@ impl VaultTransitKmsClient {
/// Snapshot the authenticated Vault client for a single request.
///
/// Every Vault call takes its own snapshot so a credential rotation
/// applies to subsequent calls without interrupting in-flight ones.
fn vault(&self) -> Arc<VaultClientHandle> {
/// applies to subsequent calls without interrupting in-flight ones. Fails
/// closed when the credentials could not be refreshed in time.
fn vault(&self) -> Result<Arc<VaultClientHandle>> {
self.credentials.current()
}
@@ -192,7 +199,7 @@ impl VaultTransitKmsClient {
}
async fn read_transit_key(&self, key_id: &str) -> Result<vaultrs::api::transit::responses::ReadKeyResponse> {
key::read(&self.vault().client, &self.config.mount_path, key_id)
key::read(&self.vault()?.client, &self.config.mount_path, key_id)
.await
.or_else(|e| Self::map_vault_error(key_id, e, "read"))
}
@@ -200,7 +207,7 @@ impl VaultTransitKmsClient {
async fn create_transit_key(&self, key_id: &str) -> Result<()> {
let mut builder = CreateKeyRequestBuilder::default();
builder.key_type(KeyType::Aes256Gcm96);
key::create(&self.vault().client, &self.config.mount_path, key_id, Some(&mut builder))
key::create(&self.vault()?.client, &self.config.mount_path, key_id, Some(&mut builder))
.await
.map_err(|e| KmsError::backend_error(format!("Failed to create Vault Transit key {key_id}: {e}")))
}
@@ -217,7 +224,7 @@ impl VaultTransitKmsClient {
builder.associated_data(aad);
}
let response = data::encrypt(&self.vault().client, &self.config.mount_path, key_id, &plaintext_b64, Some(&mut builder))
let response = data::encrypt(&self.vault()?.client, &self.config.mount_path, key_id, &plaintext_b64, Some(&mut builder))
.await
.map_err(|e| KmsError::backend_error(format!("Failed to encrypt data with Vault Transit key {key_id}: {e}")))?;
@@ -235,7 +242,7 @@ impl VaultTransitKmsClient {
builder.associated_data(aad);
}
let response = data::decrypt(&self.vault().client, &self.config.mount_path, key_id, ciphertext, Some(&mut builder))
let response = data::decrypt(&self.vault()?.client, &self.config.mount_path, key_id, ciphertext, Some(&mut builder))
.await
.map_err(|e| KmsError::backend_error(format!("Failed to decrypt data with Vault Transit key {key_id}: {e}")))?;
@@ -250,7 +257,7 @@ impl VaultTransitKmsClient {
async fn read_metadata_from_kv(&self, key_id: &str) -> Result<Option<TransitKeyMetadata>> {
let path = self.metadata_key_path(key_id);
match kv2::read::<TransitKeyMetadataPersisted>(&self.vault().client, &self.metadata_kv_mount, &path).await {
match kv2::read::<TransitKeyMetadataPersisted>(&self.vault()?.client, &self.metadata_kv_mount, &path).await {
Ok(persisted) => Ok(Some(persisted.into())),
Err(vaultrs::error::ClientError::ResponseWrapError)
| Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(None),
@@ -261,7 +268,7 @@ impl VaultTransitKmsClient {
async fn write_metadata_to_kv(&self, key_id: &str, metadata: &TransitKeyMetadata) -> Result<()> {
let path = self.metadata_key_path(key_id);
let persisted: TransitKeyMetadataPersisted = metadata.clone().into();
kv2::set(&self.vault().client, &self.metadata_kv_mount, &path, &persisted)
kv2::set(&self.vault()?.client, &self.metadata_kv_mount, &path, &persisted)
.await
.map(|_| ())
.map_err(|e| KmsError::backend_error(format!("Failed to write transit key metadata to Vault KV: {e}")))
@@ -269,7 +276,7 @@ impl VaultTransitKmsClient {
async fn delete_metadata_from_kv(&self, key_id: &str) -> Result<()> {
let path = self.metadata_key_path(key_id);
match kv2::delete_metadata(&self.vault().client, &self.metadata_kv_mount, &path).await {
match kv2::delete_metadata(&self.vault()?.client, &self.metadata_kv_mount, &path).await {
Ok(_) => Ok(()),
Err(vaultrs::error::ClientError::ResponseWrapError)
| Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(()),
@@ -483,7 +490,7 @@ impl KmsClient for VaultTransitKmsClient {
}
async fn list_keys(&self, request: &ListKeysRequest, _context: Option<&OperationContext>) -> Result<ListKeysResponse> {
let all_keys = key::list(&self.vault().client, &self.config.mount_path)
let all_keys = key::list(&self.vault()?.client, &self.config.mount_path)
.await
.map_err(|e| KmsError::backend_error(format!("Failed to list Vault Transit keys: {e}")))?
.keys;
@@ -553,7 +560,7 @@ impl KmsClient for VaultTransitKmsClient {
}
async fn rotate_key(&self, key_id: &str, _context: Option<&OperationContext>) -> Result<MasterKeyInfo> {
key::rotate(&self.vault().client, &self.config.mount_path, key_id)
key::rotate(&self.vault()?.client, &self.config.mount_path, key_id)
.await
.map_err(|e| KmsError::backend_error(format!("Failed to rotate Vault Transit key {key_id}: {e}")))?;
@@ -576,7 +583,7 @@ impl KmsClient for VaultTransitKmsClient {
}
async fn health_check(&self) -> Result<()> {
key::list(&self.vault().client, &self.config.mount_path)
key::list(&self.vault()?.client, &self.config.mount_path)
.await
.map(|_| ())
.map_err(|e| KmsError::backend_error(format!("Vault Transit health check failed: {e}")))
@@ -612,9 +619,16 @@ impl VaultTransitKmsBackend {
}
};
let client = VaultTransitKmsClient::new(vault_config, config.effective_timeout()).await?;
let client = VaultTransitKmsClient::new(vault_config, &config).await?;
Ok(Self { client })
}
/// Spawn the background credential renewal task for this backend, if its
/// auth method issues lease-bound tokens. The caller owns the returned
/// handle; dropping it cancels the task.
pub(crate) fn spawn_credential_renewal(&self) -> Option<CredentialTaskHandle> {
self.client.credentials.spawn_renewal_task()
}
}
#[async_trait]
@@ -698,7 +712,7 @@ impl KmsBackend for VaultTransitKmsBackend {
let mut update_builder = UpdateKeyConfigurationRequestBuilder::default();
update_builder.deletion_allowed(true);
key::update(
&self.client.vault().client,
&self.client.vault()?.client,
&self.client.config.mount_path,
&key_id,
Some(&mut update_builder),
@@ -708,7 +722,7 @@ impl KmsBackend for VaultTransitKmsBackend {
KmsError::backend_error(format!("Failed to allow deletion of Vault Transit key {key_id}: {e}"))
})?;
}
key::delete(&self.client.vault().client, &self.client.config.mount_path, &key_id)
key::delete(&self.client.vault()?.client, &self.client.config.mount_path, &key_id)
.await
.map_err(|e| KmsError::backend_error(format!("Failed to delete Vault Transit key {key_id}: {e}")))?;
self.client.delete_key_metadata(&key_id).await?;
@@ -798,7 +812,7 @@ mod tests {
let config = test_vault_transit_config();
// --- First "process": create a key and disable it ---
let client1 = VaultTransitKmsClient::new(config.clone(), Duration::from_secs(30))
let client1 = VaultTransitKmsClient::new(config.clone(), &KmsConfig::default())
.await
.expect("Failed to create VaultTransit client");
@@ -821,7 +835,7 @@ mod tests {
assert_eq!(info_after_disable.status, KeyStatus::Disabled, "key must be Disabled after disable_key");
// --- Simulate restart: create a brand new client with empty cache ---
let client2 = VaultTransitKmsClient::new(config, Duration::from_secs(30))
let client2 = VaultTransitKmsClient::new(config, &KmsConfig::default())
.await
.expect("Failed to create second VaultTransit client (restart simulation)");
@@ -851,7 +865,7 @@ mod tests {
async fn test_transit_pending_deletion_survives_restart_simulation() {
let config = test_vault_transit_config();
let client1 = VaultTransitKmsClient::new(config.clone(), Duration::from_secs(30))
let client1 = VaultTransitKmsClient::new(config.clone(), &KmsConfig::default())
.await
.expect("Failed to create VaultTransit client");
@@ -875,7 +889,7 @@ mod tests {
"key must be PendingDeletion after schedule_key_deletion"
);
let client2 = VaultTransitKmsClient::new(config, Duration::from_secs(30))
let client2 = VaultTransitKmsClient::new(config, &KmsConfig::default())
.await
.expect("Failed to create second VaultTransit client (restart simulation)");
@@ -917,7 +931,7 @@ mod tests {
#[tokio::test]
#[ignore] // Requires a running Vault instance with transit engine enabled
async fn test_transit_old_ciphertext_decrypts_after_rotate() {
let client = VaultTransitKmsClient::new(test_vault_transit_config(), Duration::from_secs(30))
let client = VaultTransitKmsClient::new(test_vault_transit_config(), &KmsConfig::default())
.await
.expect("Failed to create VaultTransit client");
+19 -1
View File
@@ -14,6 +14,7 @@
//! KMS service manager for dynamic configuration and runtime management
use crate::backends::vault_credentials::CredentialTaskHandle;
use crate::backends::{KmsBackend, local::LocalKmsBackend};
use crate::config::{BackendConfig, KmsConfig};
use crate::error::{KmsError, Result};
@@ -110,6 +111,10 @@ struct ServiceVersion {
service: Arc<ObjectEncryptionService>,
/// The KMS manager instance
manager: Arc<KmsManager>,
/// Owner of the backend's credential renewal task, if the backend needs
/// one. Stop shuts it down explicitly; reconfigure recycles it through
/// the handle's cancel-on-drop behavior when the old version is discarded.
credential_task: Option<Arc<CredentialTaskHandle>>,
}
#[derive(Clone)]
@@ -331,6 +336,13 @@ impl KmsServiceManager {
current_service: None,
}));
// Shut down the stopped version's credential renewal task before
// reporting stopped, so stop deterministically recycles the background
// task even while in-flight operations still hold the old service Arc.
if let Some(task) = state.current_service.as_ref().and_then(|sv| sv.credential_task.clone()) {
task.shutdown().await;
}
debug!(
event = EVENT_KMS_SERVICE_STATE,
component = LOG_COMPONENT_KMS,
@@ -488,7 +500,10 @@ impl KmsServiceManager {
info!("Creating KMS service version {} with backend: {:?}", version, config.backend);
// Create backend
// Create backend. Vault backends may also spawn a background
// credential renewal task whose owner handle lives on the service
// version, so replacing the version recycles the task.
let mut credential_task = None;
let backend = match &config.backend_config {
BackendConfig::Local(_) => {
info!("Creating Local KMS backend for version {}", version);
@@ -498,11 +513,13 @@ impl KmsServiceManager {
BackendConfig::VaultKv2(_) => {
info!("Creating Vault KV2 KMS backend for version {}", version);
let backend = crate::backends::vault::VaultKmsBackend::new(config.clone()).await?;
credential_task = backend.spawn_credential_renewal().map(Arc::new);
Arc::new(backend) as Arc<dyn KmsBackend>
}
BackendConfig::VaultTransit(_) => {
info!("Creating Vault Transit KMS backend for version {}", version);
let backend = crate::backends::vault_transit::VaultTransitKmsBackend::new(config.clone()).await?;
credential_task = backend.spawn_credential_renewal().map(Arc::new);
Arc::new(backend) as Arc<dyn KmsBackend>
}
BackendConfig::Static(_) => {
@@ -522,6 +539,7 @@ impl KmsServiceManager {
version,
service: encryption_service,
manager: kms_manager,
credential_task,
})
}