From cad8246ffb9f05ca515cbc7be94b76ce3d899ba7 Mon Sep 17 00:00:00 2001 From: Zhengchao An Date: Fri, 31 Jul 2026 08:51:02 +0800 Subject: [PATCH] feat(kms): route Vault operations through the retry policy engine (#5495) * test(kms): add a scripted loopback Vault for policy wiring tests A minimal HTTP/1.1 responder that serves canned Vault responses in order and records the method/path sequence, so wiring tests can assert exactly how many requests a code path performed (retries, read-confirm) without a live Vault server. * feat(kms): route Vault operations through the retry policy engine Wire every outbound vaultrs call in the KV2 and Transit backends through policy::execute, completing the wiring half of the operation policy work (the engine landed separately): - Reads (KV2 read/read_metadata/read_version/list, transit read/list/ encrypt/decrypt, health checks) run as ReadIdempotent: bounded retries with exponential backoff and jitter on 429, recoverable 5xx, and connection-level failures; 400/401/403/404 stay fatal. - Writes (KV2 set/CAS set/delete_metadata, transit create/update/rotate/ delete, metadata writes) run as MutatingNonIdempotent: exactly one attempt under the per-attempt timeout, never replayed. CAS conflicts in the rotation protocol pass through unchanged as the concurrency signal they are. - Each attempt takes a fresh credential snapshot, so a retry after a credential rotation uses the new token. - Read-confirm recovery for lost create responses: when a create finds an existing key that is exactly what it would have produced (same algorithm, enabled, usable material, and for request-level creates the same usage/description/tags), it reports the stored key as the create result instead of KeyAlreadyExists. Any divergence keeps failing. - Deletes treat already-deleted records as completed deletes (KV2 version records; transit metadata already did), so re-running an interrupted deletion converges. - A failed existence pre-check inside create now fails the create instead of falling through to a blind overwrite (fail closed). - The policy module sheds its allow(dead_code) now that it is wired. Wiring tests run against a scripted loopback Vault and assert request counts and endpoints for the retry, single-attempt, CAS-conflict, and read-confirm paths. Refs rustfs/backlog#1569 (part of rustfs/backlog#1562) --- crates/kms/Cargo.toml | 3 +- crates/kms/src/backends/mod.rs | 2 + crates/kms/src/backends/scripted_vault.rs | 159 +++++++ crates/kms/src/backends/vault.rs | 485 ++++++++++++++++---- crates/kms/src/backends/vault_transit.rs | 522 ++++++++++++++++++---- crates/kms/src/lib.rs | 3 - crates/kms/src/policy.rs | 30 +- 7 files changed, 1029 insertions(+), 175 deletions(-) create mode 100644 crates/kms/src/backends/scripted_vault.rs diff --git a/crates/kms/Cargo.toml b/crates/kms/Cargo.toml index e492f0e83..3c46a63a3 100644 --- a/crates/kms/Cargo.toml +++ b/crates/kms/Cargo.toml @@ -75,7 +75,8 @@ anyhow = { workspace = true } insta = { workspace = true, features = ["yaml", "json"] } tempfile = { workspace = true } temp-env = { workspace = true } -tokio = { workspace = true, features = ["test-util"] } +# "net" backs the scripted loopback Vault used by the policy wiring tests. +tokio = { workspace = true, features = ["net", "test-util"] } [features] default = [] diff --git a/crates/kms/src/backends/mod.rs b/crates/kms/src/backends/mod.rs index ae2d531b9..eecc236e8 100644 --- a/crates/kms/src/backends/mod.rs +++ b/crates/kms/src/backends/mod.rs @@ -24,6 +24,8 @@ use std::collections::HashMap; #[cfg(test)] mod contract_tests; pub mod local; +#[cfg(test)] +pub(crate) mod scripted_vault; pub mod static_kms; pub mod vault; pub(crate) mod vault_credentials; diff --git a/crates/kms/src/backends/scripted_vault.rs b/crates/kms/src/backends/scripted_vault.rs new file mode 100644 index 000000000..19453ea46 --- /dev/null +++ b/crates/kms/src/backends/scripted_vault.rs @@ -0,0 +1,159 @@ +// Copyright 2024 RustFS Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Minimal scripted HTTP responder standing in for a Vault server. +//! +//! Wiring tests need to observe how many Vault requests a code path performs +//! (retries, read-confirm recovery) without a live Vault. The responder serves +//! one canned response per incoming request in order, closes the connection +//! after each response, and records the `METHOD /path` sequence for +//! assertions. It intentionally implements just enough HTTP/1.1 for the +//! `vaultrs` reqwest client: no keep-alive, no chunked bodies. + +use std::sync::{Arc, Mutex}; + +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; + +/// One canned HTTP response. +pub(crate) struct ScriptedResponse { + status: u16, + body: String, +} + +impl ScriptedResponse { + /// A 200 response carrying `data` inside the standard Vault envelope. + pub(crate) fn ok(data: serde_json::Value) -> Self { + Self { + status: 200, + body: serde_json::json!({ + "request_id": "scripted", + "lease_id": "", + "lease_duration": 0, + "renewable": false, + "data": data, + }) + .to_string(), + } + } + + /// An error response in Vault's `{"errors": [...]}` format. + pub(crate) fn error(status: u16, message: &str) -> Self { + Self { + status, + body: serde_json::json!({ "errors": [message] }).to_string(), + } + } +} + +/// A scripted stand-in Vault listening on a loopback port. +pub(crate) struct ScriptedVault { + /// Base address (`http://127.0.0.1:port`) to point a Vault client at. + pub(crate) address: String, + requests: Arc>>, +} + +impl ScriptedVault { + /// Bind a loopback listener and serve `responses` one per request. + /// + /// Requests beyond the script get a 599 error so a test that under-scripts + /// fails loudly instead of hanging. + pub(crate) async fn serve(responses: Vec) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind scripted vault listener"); + let address = format!("http://{}", listener.local_addr().expect("scripted vault local addr")); + let requests = Arc::new(Mutex::new(Vec::new())); + let recorded = Arc::clone(&requests); + + tokio::spawn(async move { + let mut responses = responses.into_iter(); + loop { + let Ok((mut stream, _)) = listener.accept().await else { + return; + }; + let Some(request_line) = read_request(&mut stream).await else { + continue; + }; + recorded + .lock() + .expect("scripted vault request log poisoned") + .push(request_line); + let response = responses + .next() + .unwrap_or_else(|| ScriptedResponse::error(599, "scripted vault: script exhausted")); + let payload = format!( + "HTTP/1.1 {} Scripted\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response.status, + response.body.len(), + response.body + ); + let _ = stream.write_all(payload.as_bytes()).await; + let _ = stream.shutdown().await; + } + }); + + Self { address, requests } + } + + /// The `METHOD /path` lines of every request served so far, in order. + pub(crate) fn requests(&self) -> Vec { + self.requests.lock().expect("scripted vault request log poisoned").clone() + } +} + +/// Read one HTTP/1.1 request (head plus content-length body) and return its +/// `METHOD /path` line. Draining the body before responding keeps the client +/// from seeing a connection reset while it is still writing. +async fn read_request(stream: &mut TcpStream) -> Option { + let mut buffer = Vec::new(); + let mut chunk = [0u8; 4096]; + let head_end = loop { + if let Some(position) = buffer.windows(4).position(|window| window == b"\r\n\r\n") { + break position + 4; + } + let read = stream.read(&mut chunk).await.ok()?; + if read == 0 { + return None; + } + buffer.extend_from_slice(&chunk[..read]); + }; + + let head = String::from_utf8_lossy(&buffer[..head_end]).into_owned(); + let mut lines = head.lines(); + let request_line = lines.next()?; + let mut parts = request_line.split_whitespace(); + let method = parts.next()?; + let path = parts.next()?; + // rustify appends a lone "?" when an endpoint has no query parameters; + // strip it so assertions can use the plain path. + let path = path.strip_suffix('?').unwrap_or(path); + + let content_length: usize = lines + .filter_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse().ok())? + }) + .next() + .unwrap_or(0); + let mut remaining = content_length.saturating_sub(buffer.len() - head_end); + while remaining > 0 { + let read = stream.read(&mut chunk).await.ok()?; + if read == 0 { + break; + } + remaining = remaining.saturating_sub(read); + } + + Some(format!("{method} {path}")) +} diff --git a/crates/kms/src/backends/vault.rs b/crates/kms/src/backends/vault.rs index 610c70348..ebc417833 100644 --- a/crates/kms/src/backends/vault.rs +++ b/crates/kms/src/backends/vault.rs @@ -25,14 +25,17 @@ use crate::backends::{ use crate::config::{KmsConfig, VaultConfig}; use crate::encryption::{AesDekCrypto, DataKeyEnvelope, DekCrypto, generate_key_material}; use crate::error::{KmsError, Result}; +use crate::policy::{self, AttemptError, OpClass, RetryPolicy}; use crate::types::*; use async_trait::async_trait; use base64::{Engine as _, engine::general_purpose}; use jiff::Zoned; use serde::{Deserialize, Serialize}; use std::collections::HashMap; +use std::future::Future; use std::sync::Arc; use std::time::Duration; +use tokio_util::sync::CancellationToken; use tracing::{debug, info, warn}; use vaultrs::{api::kv2::requests::SetSecretRequestOptions, error::ClientError, kv2}; @@ -46,6 +49,13 @@ pub struct VaultKmsClient { key_path_prefix: String, /// DEK encryption implementation dek_crypto: AesDekCrypto, + /// Budgets wrapping every outbound Vault call (see `crate::policy`). + retry: RetryPolicy, + /// Cancellation point for the operation executor: aborts in-flight + /// attempts and backoff sleeps. Owned by the client and currently never + /// triggered — shutdown drops the whole client — but kept as the single + /// hook a future lifecycle owner can cancel through. + cancel: CancellationToken, } /// Key data stored in Vault @@ -192,6 +202,8 @@ impl VaultKmsClient { key_path_prefix: config.key_path_prefix.clone(), config, dek_crypto: AesDekCrypto::new(), + retry: RetryPolicy::from_config(kms_config), + cancel: CancellationToken::new(), }) } @@ -204,6 +216,19 @@ impl VaultKmsClient { self.credentials.current() } + /// Run one Vault call under the operation policy. + /// + /// The closure performs a single classified attempt and takes a fresh + /// credential snapshot per attempt, so a retry after a credential rotation + /// uses the new token. + async fn run(&self, operation: &'static str, class: OpClass, attempt: F) -> Result + where + F: FnMut() -> Fut, + Fut: Future>, + { + policy::execute(operation, class, &self.retry, &self.cancel, attempt).await + } + /// Get the full path for a key in Vault fn key_path(&self, key_id: &str) -> String { format!("{}/{}", self.key_path_prefix, key_id) @@ -244,14 +269,20 @@ impl VaultKmsClient { async fn get_key_version_record(&self, key_id: &str, version: u32) -> Result { let path = self.key_version_path(key_id, version); - let record: VaultKeyVersionRecord = - 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), - ClientError::APIError { code: 404, .. } => KmsError::key_version_not_found(key_id, version), - _ => KmsError::backend_error(format!("Failed to read key version record from Vault: {e}")), - })?; + let path = path.as_str(); + let record: VaultKeyVersionRecord = self + .run("vault_kv2_read_key_version", OpClass::ReadIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + kv2::read(&vault.client, &self.kv_mount, path).await.map_err(|e| { + AttemptError::from_vaultrs(e, |e| match e { + ClientError::ResponseWrapError | ClientError::APIError { code: 404, .. } => { + KmsError::key_version_not_found(key_id, version) + } + e => KmsError::backend_error(format!("Failed to read key version record from Vault: {e}")), + }) + }) + }) + .await?; if record.version != version { return Err(KmsError::material_corrupt( @@ -284,26 +315,42 @@ impl VaultKmsClient { /// later write can be check-and-set against exactly this snapshot. async fn get_key_data_versioned(&self, key_id: &str) -> Result<(u32, VaultKeyData)> { let path = self.key_path(key_id); + let path = path.as_str(); - 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), - ClientError::APIError { code: 404, .. } => KmsError::key_not_found(key_id), - _ => KmsError::backend_error(format!("Failed to read key metadata from Vault: {e}")), - })?; + let metadata = self + .run("vault_kv2_read_key_metadata", OpClass::ReadIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + kv2::read_metadata(&vault.client, &self.kv_mount, path).await.map_err(|e| { + AttemptError::from_vaultrs(e, |e| match e { + ClientError::ResponseWrapError | ClientError::APIError { code: 404, .. } => { + KmsError::key_not_found(key_id) + } + e => KmsError::backend_error(format!("Failed to read key metadata from Vault: {e}")), + }) + }) + }) + .await?; let cas = u32::try_from(metadata.current_version) .map_err(|_| KmsError::backend_error(format!("KV2 secret version for key {key_id} exceeds u32")))?; // 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) - .await - .map_err(|e| match e { - ClientError::ResponseWrapError => KmsError::key_not_found(key_id), - ClientError::APIError { code: 404, .. } => KmsError::key_not_found(key_id), - _ => KmsError::backend_error(format!("Failed to read key from Vault: {e}")), - })?; + let secret_version = metadata.current_version; + let key_data: VaultKeyData = self + .run("vault_kv2_read_key_at_version", OpClass::ReadIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + kv2::read_version(&vault.client, &self.kv_mount, path, secret_version) + .await + .map_err(|e| { + AttemptError::from_vaultrs(e, |e| match e { + ClientError::ResponseWrapError | ClientError::APIError { code: 404, .. } => { + KmsError::key_not_found(key_id) + } + e => KmsError::backend_error(format!("Failed to read key from Vault: {e}")), + }) + }) + }) + .await?; Ok((cas, key_data)) } @@ -315,19 +362,29 @@ impl VaultKmsClient { /// further check-and-set writes. async fn cas_store_key_data(&self, key_id: &str, key_data: &VaultKeyData, cas: u32) -> Result { let path = self.key_path(key_id); + let path = path.as_str(); - let written = - kv2::set_with_options(&self.vault()?.client, &self.kv_mount, &path, key_data, SetSecretRequestOptions { cas }) - .await - .map_err(|e| { - if is_cas_conflict(&e) { - KmsError::invalid_operation(format!( - "Concurrent modification of key {key_id} detected, retry the rotation" - )) - } else { - KmsError::backend_error(format!("Failed to store key in Vault: {e}")) - } - })?; + // Single attempt: replaying a lost-response write would double-apply + // the mutation, and a CAS conflict is a normal concurrency signal that + // must reach the caller untouched. + let written = self + .run("vault_kv2_cas_write_key", OpClass::MutatingNonIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + kv2::set_with_options(&vault.client, &self.kv_mount, path, key_data, SetSecretRequestOptions { cas }) + .await + .map_err(|e| { + AttemptError::from_vaultrs(e, |e| { + if is_cas_conflict(&e) { + KmsError::invalid_operation(format!( + "Concurrent modification of key {key_id} detected, retry the rotation" + )) + } else { + KmsError::backend_error(format!("Failed to store key in Vault: {e}")) + } + }) + }) + }) + .await?; u32::try_from(written.version) .map_err(|_| KmsError::backend_error(format!("KV2 secret version for key {key_id} exceeds u32"))) @@ -340,23 +397,42 @@ impl VaultKmsClient { /// existing record is acceptable. The record is never overwritten. async fn try_create_key_version_record(&self, key_id: &str, record: &VaultKeyVersionRecord) -> Result { let path = self.key_version_path(key_id, record.version); + let path = path.as_str(); - 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), - Err(e) => Err(KmsError::backend_error(format!("Failed to store key version record in Vault: {e}"))), - } + // Single attempt: the create-only CAS makes a duplicate replay fail + // with a conflict, which the caller resolves by reading the record + // back, so retrying here would only mask that recovery path. + self.run("vault_kv2_create_key_version", OpClass::MutatingNonIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + match kv2::set_with_options(&vault.client, &self.kv_mount, path, record, SetSecretRequestOptions { cas: 0 }).await { + Ok(_) => Ok(true), + Err(e) if is_cas_conflict(&e) => Ok(false), + Err(e) => Err(AttemptError::from_vaultrs(e, |e| { + KmsError::backend_error(format!("Failed to store key version record in Vault: {e}")) + })), + } + }) + .await } /// Store key data in Vault async fn store_key_data(&self, key_id: &str, key_data: &VaultKeyData) -> Result<()> { let path = self.key_path(key_id); + let path = path.as_str(); - 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}")))?; + // Single attempt: this is a whole-record overwrite without a CAS + // precondition, so a replay after a lost response could clobber a + // concurrent writer. + self.run("vault_kv2_write_key", OpClass::MutatingNonIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + kv2::set(&vault.client, &self.kv_mount, path, key_data) + .await + .map(|_| ()) + .map_err(|e| { + AttemptError::from_vaultrs(e, |e| KmsError::backend_error(format!("Failed to store key in Vault: {e}"))) + }) + }) + .await?; debug!("Stored key {} in Vault at path {}", key_id, path); Ok(()) @@ -404,14 +480,21 @@ impl VaultKmsClient { /// Retrieve key data from Vault async fn get_key_data(&self, key_id: &str) -> Result { let path = self.key_path(key_id); + let path = path.as_str(); - 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), - vaultrs::error::ClientError::APIError { code: 404, .. } => KmsError::key_not_found(key_id), - _ => KmsError::backend_error(format!("Failed to read key from Vault: {e}")), - })?; + let secret: VaultKeyData = self + .run("vault_kv2_read_key", OpClass::ReadIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + kv2::read(&vault.client, &self.kv_mount, path).await.map_err(|e| { + AttemptError::from_vaultrs(e, |e| match e { + ClientError::ResponseWrapError | ClientError::APIError { code: 404, .. } => { + KmsError::key_not_found(key_id) + } + e => KmsError::backend_error(format!("Failed to read key from Vault: {e}")), + }) + }) + }) + .await?; debug!("Retrieved key {} from Vault, tags: {:?}", key_id, secret.tags); Ok(secret) @@ -419,56 +502,87 @@ impl VaultKmsClient { /// List all keys stored in Vault async fn list_vault_keys(&self) -> Result> { - // List keys under the prefix - match kv2::list(&self.vault()?.client, &self.kv_mount, &self.key_path_prefix).await { - Ok(keys) => { + // List keys under the prefix; `None` means the prefix does not exist + // yet (no keys were ever created). + let keys = self + .run("vault_kv2_list_keys", OpClass::ReadIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + match kv2::list(&vault.client, &self.kv_mount, &self.key_path_prefix).await { + Ok(keys) => Ok(Some(keys)), + Err(ClientError::ResponseWrapError) | Err(ClientError::APIError { code: 404, .. }) => Ok(None), + Err(e) => Err(AttemptError::from_vaultrs(e, |e| { + KmsError::backend_error(format!("Failed to list keys in Vault: {e}")) + })), + } + }) + .await?; + + match keys { + Some(keys) => { let keys = filter_key_directory_entries(keys); debug!("Found {} keys in Vault", keys.len()); Ok(keys) } - Err(vaultrs::error::ClientError::ResponseWrapError) => { - // No keys exist yet - Ok(Vec::new()) - } - Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => { - // Path doesn't exist - no keys exist yet + None => { debug!("Key path doesn't exist in Vault (404), returning empty list"); Ok(Vec::new()) } - Err(e) => Err(KmsError::backend_error(format!("Failed to list keys in Vault: {e}"))), } } /// Physically delete a key from Vault storage async fn delete_key(&self, key_id: &str) -> Result<()> { let path = self.key_path(key_id); + let path = path.as_str(); // Purge immutable version records first: if any purge fails, the top-level // 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 { - Ok(versions) => { - for version in versions { - let version_path = format!("{versions_dir}/{version}"); - 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}")))?; + let versions_dir = versions_dir.as_str(); + // `None` means no version records exist (the key was never rotated). + let versions = self + .run("vault_kv2_list_key_versions", OpClass::ReadIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + match kv2::list(&vault.client, &self.kv_mount, versions_dir).await { + Ok(versions) => Ok(Some(versions)), + Err(ClientError::ResponseWrapError) | Err(ClientError::APIError { code: 404, .. }) => Ok(None), + Err(e) => Err(AttemptError::from_vaultrs(e, |e| { + KmsError::backend_error(format!("Failed to list key version records in Vault: {e}")) + })), } - } - // No version records exist (the key was never rotated). - Err(ClientError::ResponseWrapError) | Err(ClientError::APIError { code: 404, .. }) => {} - Err(e) => return Err(KmsError::backend_error(format!("Failed to list key version records in Vault: {e}"))), + }) + .await?; + for version in versions.unwrap_or_default() { + let version_path = format!("{versions_dir}/{version}"); + let version_path = version_path.as_str(); + self.run("vault_kv2_delete_key_version", OpClass::MutatingNonIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + match kv2::delete_metadata(&vault.client, &self.kv_mount, version_path).await { + // A version record that is already gone is a completed + // delete (e.g. this deletion is being re-run after a lost + // response), not a failure. + Ok(_) | Err(ClientError::ResponseWrapError) | Err(ClientError::APIError { code: 404, .. }) => Ok(()), + Err(e) => Err(AttemptError::from_vaultrs(e, |e| { + KmsError::backend_error(format!("Failed to delete key version record from Vault: {e}")) + })), + } + }) + .await?; } // 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) - .await - .map_err(|e| match e { - vaultrs::error::ClientError::APIError { code: 404, .. } => KmsError::key_not_found(key_id), - _ => KmsError::backend_error(format!("Failed to delete key metadata from Vault: {e}")), - })?; + self.run("vault_kv2_delete_key", OpClass::MutatingNonIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + kv2::delete_metadata(&vault.client, &self.kv_mount, path).await.map_err(|e| { + AttemptError::from_vaultrs(e, |e| match e { + ClientError::APIError { code: 404, .. } => KmsError::key_not_found(key_id), + e => KmsError::backend_error(format!("Failed to delete key metadata from Vault: {e}")), + }) + }) + }) + .await?; debug!("Permanently deleted key {} metadata from Vault at path {}", key_id, path); Ok(()) @@ -586,9 +700,43 @@ impl KmsClient for VaultKmsClient { async fn create_key(&self, key_id: &str, algorithm: &str, _context: Option<&OperationContext>) -> Result { debug!("Creating master key: {} with algorithm: {}", key_id, algorithm); - // Check if key already exists - if self.get_key_data(key_id).await.is_ok() { - return Err(KmsError::key_already_exists(key_id)); + // Existence pre-check with read-confirm recovery: a create whose + // response was lost gets retried by callers, and used to be + // misreported as KeyAlreadyExists. If the stored key is exactly what + // this create would have produced (same algorithm, active, usable + // material), report the stored key as the create result. Anything + // else keeps failing: create never adopts a key it would not have + // produced. A failed pre-check read must fail the create rather than + // fall through to a blind overwrite of a possibly existing key. + match self.get_key_data(key_id).await { + Ok(existing) => { + return if existing.algorithm == algorithm + && existing.status == KeyStatus::Active + && decode_stored_key_material(key_id, &existing.encrypted_key_material).is_ok() + { + info!( + key_id, + "Vault KMS create found an identical active key; treating it as a recovered create" + ); + Ok(MasterKeyInfo { + key_id: key_id.to_string(), + version: existing.version, + algorithm: existing.algorithm, + usage: existing.usage, + status: existing.status, + description: existing.description, + metadata: existing.metadata, + created_at: existing.created_at, + rotated_at: None, + created_by: None, + deletion_date: existing.deletion_date, + }) + } else { + Err(KmsError::key_already_exists(key_id)) + }; + } + Err(KmsError::KeyNotFound { .. }) => {} + Err(error) => return Err(error), } // Generate key material @@ -1203,8 +1351,183 @@ impl KmsBackend for VaultKmsBackend { #[cfg(test)] mod tests { use super::*; + use crate::backends::scripted_vault::{ScriptedResponse, ScriptedVault}; use crate::config::{VaultAuthMethod, VaultConfig}; + /// Vault + KMS config pair pointing at a scripted loopback Vault. + fn scripted_configs(address: &str) -> (VaultConfig, KmsConfig) { + let vault_config = VaultConfig { + address: address.to_string(), + auth_method: VaultAuthMethod::Token { + token: "scripted-token".to_string(), + }, + kv_mount: "secret".to_string(), + key_path_prefix: "rustfs/kms/keys".to_string(), + mount_path: "transit".to_string(), + namespace: None, + tls: None, + }; + let kms_config = KmsConfig { + timeout: Duration::from_secs(5), + retry_attempts: 3, + ..KmsConfig::default() + }; + (vault_config, kms_config) + } + + async fn scripted_client(responses: Vec) -> (ScriptedVault, VaultKmsClient) { + let vault = ScriptedVault::serve(responses).await; + let (vault_config, kms_config) = scripted_configs(&vault.address); + let client = VaultKmsClient::new(vault_config, &kms_config) + .await + .expect("scripted Vault client"); + (vault, client) + } + + fn healthy_key_data() -> VaultKeyData { + VaultKeyData { + algorithm: "AES_256".to_string(), + usage: KeyUsage::EncryptDecrypt, + created_at: Zoned::now(), + status: KeyStatus::Active, + version: 1, + description: None, + metadata: HashMap::new(), + tags: HashMap::new(), + deletion_date: None, + encrypted_key_material: general_purpose::STANDARD.encode([0x42u8; 32]), + baseline_version: None, + } + } + + /// KV2 read payload (the `data` field of the Vault envelope) for a key record. + fn kv2_read_data(key_data: &VaultKeyData) -> serde_json::Value { + serde_json::json!({ + "data": serde_json::to_value(key_data).expect("serialize key data"), + "metadata": { + "created_time": "2026-01-01T00:00:00Z", + "deletion_time": "", + "custom_metadata": null, + "destroyed": false, + "version": 1, + }, + }) + } + + #[tokio::test] + async fn wired_read_retries_transient_status_then_succeeds() { + let (vault, client) = scripted_client(vec![ + ScriptedResponse::error(503, "temporarily unavailable"), + ScriptedResponse::ok(kv2_read_data(&healthy_key_data())), + ]) + .await; + + let key_data = client + .get_key_data("wired-key") + .await + .expect("read must retry past a transient 503"); + assert_eq!(key_data.algorithm, "AES_256"); + + let requests = vault.requests(); + assert_eq!(requests.len(), 2, "one failed attempt plus one retry: {requests:?}"); + assert!( + requests + .iter() + .all(|line| line == "GET /v1/secret/data/rustfs/kms/keys/wired-key"), + "both attempts must hit the same read endpoint: {requests:?}" + ); + } + + #[tokio::test] + async fn wired_read_does_not_retry_permission_errors() { + let (vault, client) = scripted_client(vec![ScriptedResponse::error(403, "permission denied")]).await; + + client + .get_key_data("wired-key") + .await + .expect_err("a 403 must fail the read outright"); + + let requests = vault.requests(); + assert_eq!(requests.len(), 1, "fatal statuses must not be retried: {requests:?}"); + } + + #[tokio::test] + async fn wired_write_is_never_retried_on_transient_status() { + let (vault, client) = scripted_client(vec![ScriptedResponse::error(503, "sealed")]).await; + + let error = client + .store_key_data("wired-key", &healthy_key_data()) + .await + .expect_err("the scripted 503 must fail the write"); + assert!(matches!(error, KmsError::BackendError { .. }), "got {error:?}"); + + let requests = vault.requests(); + assert_eq!( + requests, + vec!["POST /v1/secret/data/rustfs/kms/keys/wired-key".to_string()], + "a non-idempotent write must run exactly once even on a retryable status" + ); + } + + #[tokio::test] + async fn wired_cas_conflict_is_surfaced_without_retry() { + let (vault, client) = scripted_client(vec![ScriptedResponse::error( + 400, + "check-and-set parameter did not match the current version", + )]) + .await; + + let error = client + .cas_store_key_data("wired-key", &healthy_key_data(), 7) + .await + .expect_err("the scripted CAS conflict must fail the write"); + assert!( + matches!(error, KmsError::InvalidOperation { .. }), + "a CAS conflict is a concurrency signal, not a backend failure: {error:?}" + ); + + let requests = vault.requests(); + assert_eq!(requests.len(), 1, "a CAS conflict must never be retried: {requests:?}"); + } + + #[tokio::test] + async fn wired_create_key_read_confirms_identical_existing_key() { + // The stored key is exactly what create_key("wired-key", "AES_256") + // would have produced, so a retried create whose first response was + // lost recovers by reading it back instead of failing. + let (vault, client) = scripted_client(vec![ScriptedResponse::ok(kv2_read_data(&healthy_key_data()))]).await; + + let recovered = client + .create_key("wired-key", "AES_256", None) + .await + .expect("an identical active key must read-confirm as a recovered create"); + assert_eq!(recovered.version, 1); + assert_eq!(recovered.algorithm, "AES_256"); + + let requests = vault.requests(); + assert_eq!( + requests, + vec!["GET /v1/secret/data/rustfs/kms/keys/wired-key".to_string()], + "a recovered create must not write anything" + ); + } + + #[tokio::test] + async fn wired_create_key_still_fails_on_mismatched_existing_key() { + let mut disabled = healthy_key_data(); + disabled.status = KeyStatus::Disabled; + let (vault, client) = scripted_client(vec![ScriptedResponse::ok(kv2_read_data(&disabled))]).await; + + let error = client + .create_key("wired-key", "AES_256", None) + .await + .expect_err("a non-active existing key must keep failing the create"); + assert!(matches!(error, KmsError::KeyAlreadyExists { .. }), "got {error:?}"); + + let requests = vault.requests(); + assert_eq!(requests.len(), 1, "the mismatch must be decided from the single read: {requests:?}"); + } + /// Poison matrix for the read-side material gate. Every corruption class must fail /// closed with its typed error; reintroducing any "self-heal" (regenerate on empty or /// undecodable material) turns one of these expected errors into an Ok and fails the diff --git a/crates/kms/src/backends/vault_transit.rs b/crates/kms/src/backends/vault_transit.rs index b6d16b521..79d5219b7 100644 --- a/crates/kms/src/backends/vault_transit.rs +++ b/crates/kms/src/backends/vault_transit.rs @@ -24,15 +24,19 @@ use crate::backends::{ use crate::config::{KmsConfig, VaultTransitConfig}; use crate::encryption::{DataKeyEnvelope, generate_key_material}; use crate::error::{KmsError, Result}; +use crate::policy::{self, AttemptError, OpClass, RetryPolicy}; use crate::types::*; use async_trait::async_trait; use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; use jiff::Zoned; use serde::{Deserialize, Serialize}; use std::collections::{BTreeMap, HashMap}; +use std::future::Future; use std::sync::Arc; use std::time::Duration; use tokio::sync::RwLock; +use tokio_util::sync::CancellationToken; +use tracing::info; use vaultrs::{ api::transit::{ KeyType, @@ -147,6 +151,13 @@ pub struct VaultTransitKmsClient { /// Path prefix under metadata_kv_mount for storing transit key metadata records metadata_key_prefix: String, metadata_cache: RwLock>, + /// Budgets wrapping every outbound Vault call (see `crate::policy`). + retry: RetryPolicy, + /// Cancellation point for the operation executor: aborts in-flight + /// attempts and backoff sleeps. Owned by the client and currently never + /// triggered — shutdown drops the whole client — but kept as the single + /// hook a future lifecycle owner can cancel through. + cancel: CancellationToken, } impl VaultTransitKmsClient { @@ -171,6 +182,8 @@ impl VaultTransitKmsClient { metadata_key_prefix: config.metadata_key_prefix.clone(), config, metadata_cache: RwLock::new(HashMap::new()), + retry: RetryPolicy::from_config(kms_config), + cancel: CancellationToken::new(), }) } @@ -183,6 +196,19 @@ impl VaultTransitKmsClient { self.credentials.current() } + /// Run one Vault call under the operation policy. + /// + /// The closure performs a single classified attempt and takes a fresh + /// credential snapshot per attempt, so a retry after a credential rotation + /// uses the new token. + async fn run(&self, operation: &'static str, class: OpClass, attempt: F) -> Result + where + F: FnMut() -> Fut, + Fut: Future>, + { + policy::execute(operation, class, &self.retry, &self.cancel, attempt).await + } + fn canonicalize_context(encryption_context: &HashMap) -> Result> { if encryption_context.is_empty() { return Ok(None); @@ -196,28 +222,40 @@ impl VaultTransitKmsClient { Ok(Some(BASE64.encode(serialized))) } - fn map_vault_error(key_id: &str, error: vaultrs::error::ClientError, operation: &str) -> Result { + fn map_vault_error(key_id: &str, error: vaultrs::error::ClientError, operation: &str) -> KmsError { match error { - vaultrs::error::ClientError::ResponseWrapError => Err(KmsError::key_not_found(key_id)), - vaultrs::error::ClientError::APIError { code: 404, .. } => Err(KmsError::key_not_found(key_id)), - other => Err(KmsError::backend_error(format!( - "Vault Transit {operation} failed for key {key_id}: {other}" - ))), + vaultrs::error::ClientError::ResponseWrapError => KmsError::key_not_found(key_id), + vaultrs::error::ClientError::APIError { code: 404, .. } => KmsError::key_not_found(key_id), + other => KmsError::backend_error(format!("Vault Transit {operation} failed for key {key_id}: {other}")), } } async fn read_transit_key(&self, key_id: &str) -> Result { - key::read(&self.vault()?.client, &self.config.mount_path, key_id) - .await - .or_else(|e| Self::map_vault_error(key_id, e, "read")) + self.run("vault_transit_read_key", OpClass::ReadIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + key::read(&vault.client, &self.config.mount_path, key_id) + .await + .map_err(|e| AttemptError::from_vaultrs(e, |e| Self::map_vault_error(key_id, e, "read"))) + }) + .await } 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)) - .await - .map_err(|e| KmsError::backend_error(format!("Failed to create Vault Transit key {key_id}: {e}"))) + // Single attempt: create carries external side effects and the caller + // owns the read-confirm recovery for lost responses. + self.run("vault_transit_create_key", OpClass::MutatingNonIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + let mut builder = CreateKeyRequestBuilder::default(); + builder.key_type(KeyType::Aes256Gcm96); + key::create(&vault.client, &self.config.mount_path, key_id, Some(&mut builder)) + .await + .map_err(|e| { + AttemptError::from_vaultrs(e, |e| { + KmsError::backend_error(format!("Failed to create Vault Transit key {key_id}: {e}")) + }) + }) + }) + .await } async fn transit_encrypt( @@ -227,14 +265,26 @@ impl VaultTransitKmsClient { encryption_context: &HashMap, ) -> Result { let plaintext_b64 = BASE64.encode(plaintext); - let mut builder = EncryptDataRequestBuilder::default(); - if let Some(aad) = Self::canonicalize_context(encryption_context)? { - builder.associated_data(aad); - } + let plaintext_b64 = plaintext_b64.as_str(); + let aad = Self::canonicalize_context(encryption_context)?; + let aad = aad.as_deref(); - 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}")))?; + let response = self + .run("vault_transit_encrypt", OpClass::ReadIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + let mut builder = EncryptDataRequestBuilder::default(); + if let Some(aad) = aad { + builder.associated_data(aad); + } + data::encrypt(&vault.client, &self.config.mount_path, key_id, plaintext_b64, Some(&mut builder)) + .await + .map_err(|e| { + AttemptError::from_vaultrs(e, |e| { + KmsError::backend_error(format!("Failed to encrypt data with Vault Transit key {key_id}: {e}")) + }) + }) + }) + .await?; Ok(response.ciphertext) } @@ -245,14 +295,25 @@ impl VaultTransitKmsClient { ciphertext: &str, encryption_context: &HashMap, ) -> Result> { - let mut builder = DecryptDataRequestBuilder::default(); - if let Some(aad) = Self::canonicalize_context(encryption_context)? { - builder.associated_data(aad); - } + let aad = Self::canonicalize_context(encryption_context)?; + let aad = aad.as_deref(); - 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}")))?; + let response = self + .run("vault_transit_decrypt", OpClass::ReadIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + let mut builder = DecryptDataRequestBuilder::default(); + if let Some(aad) = aad { + builder.associated_data(aad); + } + data::decrypt(&vault.client, &self.config.mount_path, key_id, ciphertext, Some(&mut builder)) + .await + .map_err(|e| { + AttemptError::from_vaultrs(e, |e| { + KmsError::backend_error(format!("Failed to decrypt data with Vault Transit key {key_id}: {e}")) + }) + }) + }) + .await?; BASE64 .decode(response.plaintext) @@ -265,33 +326,93 @@ impl VaultTransitKmsClient { async fn read_metadata_from_kv(&self, key_id: &str) -> Result> { let path = self.metadata_key_path(key_id); - match kv2::read::(&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), - Err(e) => Err(KmsError::backend_error(format!("Failed to read transit key metadata from Vault KV: {e}"))), - } + let path = path.as_str(); + self.run("vault_transit_read_metadata", OpClass::ReadIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + match kv2::read::(&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), + Err(e) => Err(AttemptError::from_vaultrs(e, |e| { + KmsError::backend_error(format!("Failed to read transit key metadata from Vault KV: {e}")) + })), + } + }) + .await } async fn write_metadata_to_kv(&self, key_id: &str, metadata: &TransitKeyMetadata) -> Result<()> { let path = self.metadata_key_path(key_id); + let path = path.as_str(); let persisted: TransitKeyMetadataPersisted = metadata.clone().into(); - 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}"))) + let persisted = &persisted; + // Single attempt: this is a whole-record overwrite without a CAS + // precondition, so a replay after a lost response could clobber a + // concurrent writer. + self.run("vault_transit_write_metadata", OpClass::MutatingNonIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + kv2::set(&vault.client, &self.metadata_kv_mount, path, persisted) + .await + .map(|_| ()) + .map_err(|e| { + AttemptError::from_vaultrs(e, |e| { + KmsError::backend_error(format!("Failed to write transit key metadata to Vault KV: {e}")) + }) + }) + }) + .await } 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 { - Ok(_) => Ok(()), - Err(vaultrs::error::ClientError::ResponseWrapError) - | Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(()), - Err(e) => Err(KmsError::backend_error(format!( - "Failed to delete transit key metadata from Vault KV: {e}" - ))), - } + let path = path.as_str(); + self.run("vault_transit_delete_metadata", OpClass::MutatingNonIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + match kv2::delete_metadata(&vault.client, &self.metadata_kv_mount, path).await { + // Metadata that is already gone is a completed delete. + Ok(_) + | Err(vaultrs::error::ClientError::ResponseWrapError) + | Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(()), + Err(e) => Err(AttemptError::from_vaultrs(e, |e| { + KmsError::backend_error(format!("Failed to delete transit key metadata from Vault KV: {e}")) + })), + } + }) + .await + } + + /// Flip `deletion_allowed` on the transit key so it can be deleted. + async fn allow_transit_key_deletion(&self, key_id: &str) -> Result<()> { + self.run("vault_transit_allow_deletion", OpClass::MutatingNonIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + let mut builder = UpdateKeyConfigurationRequestBuilder::default(); + builder.deletion_allowed(true); + key::update(&vault.client, &self.config.mount_path, key_id, Some(&mut builder)) + .await + .map(|_| ()) + .map_err(|e| { + AttemptError::from_vaultrs(e, |e| { + KmsError::backend_error(format!("Failed to allow deletion of Vault Transit key {key_id}: {e}")) + }) + }) + }) + .await + } + + /// Physically delete the transit key material in Vault. + async fn delete_transit_key(&self, key_id: &str) -> Result<()> { + self.run("vault_transit_delete_key", OpClass::MutatingNonIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + key::delete(&vault.client, &self.config.mount_path, key_id) + .await + .map(|_| ()) + .map_err(|e| { + AttemptError::from_vaultrs(e, |e| { + KmsError::backend_error(format!("Failed to delete Vault Transit key {key_id}: {e}")) + }) + }) + }) + .await } async fn get_key_metadata(&self, key_id: &str) -> Result { @@ -462,8 +583,40 @@ impl KmsClient for VaultTransitKmsClient { return Err(KmsError::unsupported_algorithm(algorithm)); } - if self.read_transit_key(key_id).await.is_ok() { - return Err(KmsError::key_already_exists(key_id)); + // Existence pre-check with read-confirm recovery: a create whose + // response was lost gets retried by callers, and used to be + // misreported as KeyAlreadyExists. Transit keys are always AES-256, + // so an existing enabled key of the default usage is exactly what + // this create would have produced; report it as the create result. + // Anything else keeps failing. A failed pre-check read must fail the + // create rather than fall through to re-creating over an unknown key. + match self.read_transit_key(key_id).await { + Ok(_) => { + let existing = self.get_key_metadata(key_id).await?; + return if existing.key_state == KeyState::Enabled && existing.key_usage == KeyUsage::EncryptDecrypt { + info!( + key_id, + "Vault Transit create found an identical enabled key; treating it as a recovered create" + ); + Ok(MasterKeyInfo { + key_id: key_id.to_string(), + version: existing.current_version, + algorithm: algorithm.to_string(), + usage: existing.key_usage, + status: KeyStatus::Active, + description: existing.description, + metadata: existing.tags.clone(), + created_at: existing.created_at, + rotated_at: None, + created_by: existing.created_by, + deletion_date: None, + }) + } else { + Err(KmsError::key_already_exists(key_id)) + }; + } + Err(KmsError::KeyNotFound { .. }) => {} + Err(error) => return Err(error), } self.create_transit_key(key_id).await?; @@ -497,9 +650,14 @@ impl KmsClient for VaultTransitKmsClient { } async fn list_keys(&self, request: &ListKeysRequest, _context: Option<&OperationContext>) -> Result { - 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}")))? + let all_keys = self + .run("vault_transit_list_keys", OpClass::ReadIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + key::list(&vault.client, &self.config.mount_path).await.map_err(|e| { + AttemptError::from_vaultrs(e, |e| KmsError::backend_error(format!("Failed to list Vault Transit keys: {e}"))) + }) + }) + .await? .keys; let mut filtered = Vec::new(); @@ -576,9 +734,20 @@ impl KmsClient for VaultTransitKmsClient { async fn rotate_key(&self, key_id: &str, _context: Option<&OperationContext>) -> Result { self.ensure_key_state_allows(key_id, StateGatedOperation::Rotate).await?; - 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}")))?; + // Single attempt, never retried: replaying a rotate whose response was + // lost would advance the key version once more per replay. + self.run("vault_transit_rotate_key", OpClass::MutatingNonIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + key::rotate(&vault.client, &self.config.mount_path, key_id) + .await + .map(|_| ()) + .map_err(|e| { + AttemptError::from_vaultrs(e, |e| { + KmsError::backend_error(format!("Failed to rotate Vault Transit key {key_id}: {e}")) + }) + }) + }) + .await?; let mut metadata = self.get_key_metadata(key_id).await?; metadata.current_version += 1; @@ -600,10 +769,16 @@ impl KmsClient for VaultTransitKmsClient { } async fn health_check(&self) -> Result<()> { - key::list(&self.vault()?.client, &self.config.mount_path) - .await - .map(|_| ()) - .map_err(|e| KmsError::backend_error(format!("Vault Transit health check failed: {e}"))) + self.run("vault_transit_health_check", OpClass::ReadIdempotent, move || async move { + let vault = self.vault().map_err(AttemptError::fatal)?; + key::list(&vault.client, &self.config.mount_path) + .await + .map(|_| ()) + .map_err(|e| { + AttemptError::from_vaultrs(e, |e| KmsError::backend_error(format!("Vault Transit health check failed: {e}"))) + }) + }) + .await } fn backend_info(&self) -> BackendInfo { @@ -660,8 +835,46 @@ impl VaultTransitKmsBackend { impl KmsBackend for VaultTransitKmsBackend { async fn create_key(&self, request: CreateKeyRequest) -> Result { let key_id = request.key_name.clone().unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - if self.client.read_transit_key(&key_id).await.is_ok() { - return Err(KmsError::key_already_exists(&key_id)); + + // Existence pre-check with read-confirm recovery: a create whose + // response was lost gets retried by callers, and used to be + // misreported as KeyAlreadyExists. If the stored record is exactly + // what this request would have written, report it as the create + // result; any divergence keeps failing so a create can never adopt or + // reshape a key it would not have produced. + match self.client.read_transit_key(&key_id).await { + Ok(_) => { + let existing = self.client.get_key_metadata(&key_id).await?; + let requested = TransitKeyMetadata::from_create_request(&request); + return if existing.key_state == KeyState::Enabled + && existing.key_usage == requested.key_usage + && existing.description == requested.description + && existing.tags == requested.tags + { + info!( + key_id, + "Vault Transit create found an identical enabled key; treating it as a recovered create" + ); + Ok(CreateKeyResponse { + key_id: key_id.clone(), + key_metadata: KeyMetadata { + key_id, + key_state: existing.key_state, + key_usage: existing.key_usage, + description: existing.description, + creation_date: existing.created_at, + deletion_date: existing.deletion_date, + origin: existing.origin, + key_manager: "VAULT_TRANSIT".to_string(), + tags: existing.tags, + }, + }) + } else { + Err(KmsError::key_already_exists(&key_id)) + }; + } + Err(KmsError::KeyNotFound { .. }) => {} + Err(error) => return Err(error), } self.client.create_transit_key(&key_id).await?; @@ -734,22 +947,9 @@ impl KmsBackend for VaultTransitKmsBackend { let deletion_date = if request.force_immediate.unwrap_or(false) { if key_metadata.key_state == KeyState::PendingDeletion { if !self.client.read_transit_key(&key_id).await?.deletion_allowed { - let mut update_builder = UpdateKeyConfigurationRequestBuilder::default(); - update_builder.deletion_allowed(true); - key::update( - &self.client.vault()?.client, - &self.client.config.mount_path, - &key_id, - Some(&mut update_builder), - ) - .await - .map_err(|e| { - KmsError::backend_error(format!("Failed to allow deletion of Vault Transit key {key_id}: {e}")) - })?; + self.client.allow_transit_key_deletion(&key_id).await?; } - 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_transit_key(&key_id).await?; self.client.delete_key_metadata(&key_id).await?; None } else { @@ -852,20 +1052,9 @@ impl KmsBackend for VaultTransitKmsBackend { } if !self.client.read_transit_key(key_id).await?.deletion_allowed { - let mut update_builder = UpdateKeyConfigurationRequestBuilder::default(); - update_builder.deletion_allowed(true); - key::update( - &self.client.vault()?.client, - &self.client.config.mount_path, - key_id, - Some(&mut update_builder), - ) - .await - .map_err(|e| KmsError::backend_error(format!("Failed to allow deletion of Vault Transit key {key_id}: {e}")))?; + self.client.allow_transit_key_deletion(key_id).await?; } - 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_transit_key(key_id).await?; self.client.delete_key_metadata(key_id).await?; Ok(ExpiredKeyRemoval::Removed) } @@ -874,10 +1063,167 @@ impl KmsBackend for VaultTransitKmsBackend { #[cfg(test)] mod tests { use super::*; + use crate::backends::scripted_vault::{ScriptedResponse, ScriptedVault}; use crate::config::{ DEFAULT_VAULT_TRANSIT_METADATA_KEY_PREFIX, DEFAULT_VAULT_TRANSIT_METADATA_KV_MOUNT, VaultAuthMethod, VaultTransitConfig, }; use crate::types::KeyStatus; + use vaultrs::api::transit::responses::{ReadKeyData, ReadKeyResponse}; + + async fn scripted_client(responses: Vec) -> (ScriptedVault, VaultTransitKmsClient) { + let vault = ScriptedVault::serve(responses).await; + let config = VaultTransitConfig { + address: vault.address.clone(), + ..test_vault_transit_config() + }; + let kms_config = KmsConfig { + timeout: Duration::from_secs(5), + retry_attempts: 3, + ..KmsConfig::default() + }; + let client = VaultTransitKmsClient::new(config, &kms_config) + .await + .expect("scripted Vault Transit client"); + (vault, client) + } + + /// KV2 read payload for a persisted transit metadata record. + fn metadata_read_data(metadata: &TransitKeyMetadata) -> serde_json::Value { + let persisted: TransitKeyMetadataPersisted = metadata.clone().into(); + serde_json::json!({ + "data": serde_json::to_value(&persisted).expect("serialize transit metadata"), + "metadata": { + "created_time": "2026-01-01T00:00:00Z", + "deletion_time": "", + "custom_metadata": null, + "destroyed": false, + "version": 1, + }, + }) + } + + /// Transit read-key payload for an existing symmetric key. + fn transit_key_read_data(key_id: &str) -> serde_json::Value { + let response = ReadKeyResponse { + key_type: KeyType::Aes256Gcm96, + deletion_allowed: false, + derived: false, + exportable: false, + allow_plaintext_backup: false, + keys: ReadKeyData::Symmetric(HashMap::from([("1".to_string(), 1_700_000_000_u64)])), + min_decryption_version: 1, + min_encryption_version: 0, + name: key_id.to_string(), + supports_encryption: true, + supports_decryption: true, + supports_derivation: false, + supports_signing: false, + imported: Some(false), + }; + serde_json::to_value(&response).expect("serialize transit key read response") + } + + #[tokio::test] + async fn wired_transit_encrypt_retries_transient_status() { + let metadata = TransitKeyMetadata::from_create_request(&CreateKeyRequest::default()); + let (vault, client) = scripted_client(vec![ + ScriptedResponse::ok(metadata_read_data(&metadata)), + ScriptedResponse::error(429, "throttled"), + ScriptedResponse::ok(serde_json::json!({ "ciphertext": "vault:v1:scripted" })), + ]) + .await; + + let response = client + .encrypt( + &EncryptRequest { + key_id: "wired-key".to_string(), + plaintext: b"plaintext".to_vec(), + encryption_context: HashMap::new(), + grant_tokens: Vec::new(), + }, + None, + ) + .await + .expect("encrypt must retry past a transient 429"); + assert_eq!(response.ciphertext, b"vault:v1:scripted".to_vec()); + + let requests = vault.requests(); + assert_eq!(requests.len(), 3, "metadata read plus two encrypt attempts: {requests:?}"); + assert_eq!(requests[1], "POST /v1/transit/encrypt/wired-key"); + assert_eq!(requests[2], "POST /v1/transit/encrypt/wired-key"); + } + + #[tokio::test] + async fn wired_transit_rotate_is_never_retried() { + let metadata = TransitKeyMetadata::from_create_request(&CreateKeyRequest::default()); + let (vault, client) = scripted_client(vec![ + ScriptedResponse::ok(metadata_read_data(&metadata)), + ScriptedResponse::error(503, "standby"), + ]) + .await; + + let error = client + .rotate_key("wired-key", None) + .await + .expect_err("the scripted 503 must fail the rotation"); + assert!(matches!(error, KmsError::BackendError { .. }), "got {error:?}"); + + let requests = vault.requests(); + assert_eq!(requests.len(), 2, "metadata read plus exactly one rotate attempt: {requests:?}"); + assert_eq!( + requests[1], "POST /v1/transit/keys/wired-key/rotate", + "a rotation must never be replayed: {requests:?}" + ); + } + + #[tokio::test] + async fn wired_transit_create_read_confirms_identical_existing_key() { + // The stored key and metadata are exactly what this create would have + // produced, so a retried create whose first response was lost recovers + // by reading them back instead of failing. + let metadata = TransitKeyMetadata::from_create_request(&CreateKeyRequest::default()); + let (vault, client) = scripted_client(vec![ + ScriptedResponse::ok(transit_key_read_data("wired-key")), + ScriptedResponse::ok(metadata_read_data(&metadata)), + ]) + .await; + + let recovered = client + .create_key("wired-key", "AES_256", None) + .await + .expect("an identical enabled key must read-confirm as a recovered create"); + assert_eq!(recovered.status, KeyStatus::Active); + + let requests = vault.requests(); + assert_eq!(requests.len(), 2, "read-confirm must be decided from reads alone: {requests:?}"); + assert!( + requests.iter().all(|line| line.starts_with("GET ")), + "a recovered create must not write anything: {requests:?}" + ); + } + + #[tokio::test] + async fn wired_transit_create_still_fails_on_mismatched_existing_key() { + let mut metadata = TransitKeyMetadata::from_create_request(&CreateKeyRequest::default()); + metadata.key_state = KeyState::Disabled; + let (vault, client) = scripted_client(vec![ + ScriptedResponse::ok(transit_key_read_data("wired-key")), + ScriptedResponse::ok(metadata_read_data(&metadata)), + ]) + .await; + + let error = client + .create_key("wired-key", "AES_256", None) + .await + .expect_err("a non-enabled existing key must keep failing the create"); + assert!(matches!(error, KmsError::KeyAlreadyExists { .. }), "got {error:?}"); + + let requests = vault.requests(); + assert!( + requests.iter().all(|line| line.starts_with("GET ")), + "the rejected create must not write anything: {requests:?}" + ); + } fn test_vault_transit_config() -> VaultTransitConfig { VaultTransitConfig { diff --git a/crates/kms/src/lib.rs b/crates/kms/src/lib.rs index e88369669..892ed5a35 100644 --- a/crates/kms/src/lib.rs +++ b/crates/kms/src/lib.rs @@ -73,9 +73,6 @@ pub mod deletion_worker; mod encryption; mod error; pub mod manager; -// The executor is wired into the Vault backends in a follow-up change; until -// then the module is only exercised by its own tests. -#[allow(dead_code)] mod policy; pub mod service; pub mod service_manager; diff --git a/crates/kms/src/policy.rs b/crates/kms/src/policy.rs index c2014b0e0..58411a51b 100644 --- a/crates/kms/src/policy.rs +++ b/crates/kms/src/policy.rs @@ -16,8 +16,8 @@ //! //! Vault-backed operations leave the process boundary, so every call needs a //! per-attempt timeout, a total operation deadline, and classification-driven -//! bounded retries. This module provides the engine only; the Vault backends -//! wire their call sites through [`execute`] in a follow-up change. +//! bounded retries. The Vault backends and the credential provider wire every +//! outbound `vaultrs` call through [`execute`]. //! //! Retry safety is driven by two orthogonal classifications: //! - [`OpClass`] states whether replaying the operation is safe at all. @@ -112,6 +112,32 @@ pub(crate) struct AttemptError { pub(crate) error: KmsError, } +impl AttemptError { + /// A failure that must never be retried, regardless of operation class. + pub(crate) fn fatal(error: KmsError) -> Self { + Self { + class: ErrorClass::Fatal, + error, + } + } + + /// Classify a `vaultrs` failure and map it onto a domain error. + /// + /// Classification reads the raw error before `map` consumes it, so call + /// sites keep their site-specific error mapping (404 to key-not-found and + /// so on) without losing the status code the retry decision needs. + pub(crate) fn from_vaultrs( + error: vaultrs::error::ClientError, + map: impl FnOnce(vaultrs::error::ClientError) -> KmsError, + ) -> Self { + let class = classify_vaultrs(&error); + Self { + class, + error: map(error), + } + } +} + /// Budgets applied by [`execute`]. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct RetryPolicy {