From 4ce0e280f2dfa6a1b43d2399a044e5a14335c935 Mon Sep 17 00:00:00 2001 From: Zhengchao An Date: Sat, 1 Aug 2026 12:30:35 +0800 Subject: [PATCH] feat(kms): add a synthetic encrypt-decrypt probe worker (#5543) --- crates/kms/src/lib.rs | 2 + crates/kms/src/probe.rs | 894 ++++++++++++++++++++++++++++++ crates/kms/src/service_manager.rs | 78 ++- 3 files changed, 971 insertions(+), 3 deletions(-) create mode 100644 crates/kms/src/probe.rs diff --git a/crates/kms/src/lib.rs b/crates/kms/src/lib.rs index a52ef3d52..6e9ff1c37 100644 --- a/crates/kms/src/lib.rs +++ b/crates/kms/src/lib.rs @@ -75,6 +75,7 @@ mod encryption; mod error; pub mod manager; mod policy; +pub mod probe; pub mod service; pub mod service_manager; mod time_serde; @@ -93,6 +94,7 @@ pub use deletion_worker::DeletionReferenceChecker; pub use encryption::is_data_key_envelope; pub use error::{KmsError, KmsUnavailableError, Result}; pub use manager::KmsManager; +pub use probe::{ProbeFailureKind, ProbeResult, ProbeStatus}; pub use service::{DataKey, ObjectEncryptionService}; pub use service_manager::{ KmsServiceManager, KmsServiceStatus, KmsStartOutcome, get_global_encryption_service, get_global_kms_service_manager, diff --git a/crates/kms/src/probe.rs b/crates/kms/src/probe.rs new file mode 100644 index 000000000..3ec8f8b03 --- /dev/null +++ b/crates/kms/src/probe.rs @@ -0,0 +1,894 @@ +// 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. + +//! Background synthetic probe for the configured KMS backend. +//! +//! A health check that only reports "the backend answered" cannot distinguish a +//! usable KMS from one that is reachable but no longer returns material that +//! round-trips. Each probe round therefore generates a data key under a +//! dedicated key, decrypts the returned ciphertext, and compares the plaintext, +//! so a backend that answers without being usable becomes visible before object +//! traffic depends on it. +//! +//! The probe key uses a fixed reserved id ([`PROBE_KEY_ID`]) and is created only +//! when missing, so every node of a deployment converges on the same key and a +//! create that loses the race against another node is a success, not an error. +//! +//! Backends that cannot host such a key — single-key read-only backends, or +//! backends without data key generation — are reported as +//! [`ProbeResult::Unsupported`] exactly once and then left alone: an +//! unsupported backend is a configuration fact, not an outage, and must never +//! raise failure metrics or alerts. +//! +//! Results are published two ways: as a lock-free [`ProbeStatus`] snapshot that +//! readiness reporting can read without touching the backend, and as +//! `rustfs_kms_probe_*` metrics. As everywhere else in this crate, metric labels +//! carry only static outcome classes — never key identifiers, key material, or +//! ciphertext. + +use crate::backends::KmsBackend; +use crate::error::KmsError; +use crate::types::{CreateKeyRequest, DecryptRequest, DescribeKeyRequest, GenerateDataKeyRequest, KeySpec, KeyUsage}; +use arc_swap::ArcSwap; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; +use subtle::ConstantTimeEq; +use tokio::time::Instant; +use tokio_util::sync::CancellationToken; +use tracing::{debug, info, warn}; +use zeroize::Zeroize; + +/// Reserved id of the key every probe round runs through. +/// +/// Fixed rather than per-node or per-process: all nodes must converge on one +/// key instead of littering the backend with probe keys, and the reserved +/// `rustfs-internal-` prefix marks it as RustFS-owned for operators reviewing +/// the key list. +pub const PROBE_KEY_ID: &str = "rustfs-internal-kms-probe"; + +/// Description stored on the probe key when this node creates it. +const PROBE_KEY_DESCRIPTION: &str = "RustFS internal KMS health probe key; created and used only by the synthetic probe"; + +/// Default interval between probe rounds. +pub const DEFAULT_PROBE_INTERVAL: Duration = Duration::from_secs(60); + +/// Floor for the configured interval. Probing faster buys no additional signal +/// and multiplies backend load by the size of the deployment. +pub const MIN_PROBE_INTERVAL: Duration = Duration::from_secs(5); + +/// Probe interval in whole seconds; `0` disables the probe entirely. +pub const ENV_KMS_PROBE_INTERVAL_SECS: &str = "RUSTFS_KMS_PROBE_INTERVAL_SECS"; + +/// Result of the most recent probe round. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProbeResult { + /// No round has completed yet for this service version. + Pending, + /// The full generate/decrypt/compare cycle succeeded. + Success, + /// The backend cannot host the probe key; no further rounds run and no + /// failures are recorded. + Unsupported, + /// The round failed with the given classification. + Failure(ProbeFailureKind), +} + +impl ProbeResult { + /// Static metric label for this result. + fn as_label(self) -> &'static str { + match self { + ProbeResult::Pending => "pending", + ProbeResult::Success => "success", + ProbeResult::Unsupported => "unsupported", + ProbeResult::Failure(_) => "failure", + } + } +} + +/// Stage of the round trip that failed, used as a static metric label. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProbeFailureKind { + /// The dedicated probe key could not be described or created. + KeyProvisioning, + /// Data key generation failed. + Generate, + /// The freshly generated ciphertext could not be decrypted. + Decrypt, + /// Decryption succeeded but returned material that differs from what was + /// generated — the backend answers without round-tripping. + Mismatch, +} + +impl ProbeFailureKind { + /// Static metric label for this failure class. + fn as_label(self) -> &'static str { + match self { + ProbeFailureKind::KeyProvisioning => "key_provisioning", + ProbeFailureKind::Generate => "generate", + ProbeFailureKind::Decrypt => "decrypt", + ProbeFailureKind::Mismatch => "mismatch", + } + } +} + +/// Snapshot of what the probe has observed so far. +/// +/// Consumed by readiness reporting, which owns the staleness and consecutive +/// failure thresholds; the probe itself only records what happened. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProbeStatus { + /// Result of the most recent completed round. + pub result: ProbeResult, + /// When the most recent round completed. + pub last_round_at: Option, + /// When the most recent successful round completed. + pub last_success_at: Option, + /// Wall-clock timestamp of the most recent success, in Unix seconds, for + /// display and for the exported gauge. Use [`ProbeStatus::last_success_age`] + /// for staleness decisions: it is monotonic and immune to clock steps. + pub last_success_unix_secs: Option, + /// Failed rounds since the last success; reset by any non-failure result. + pub consecutive_failures: u32, +} + +impl ProbeStatus { + /// Initial snapshot published before the first round completes. + fn pending() -> Self { + Self { + result: ProbeResult::Pending, + last_round_at: None, + last_success_at: None, + last_success_unix_secs: None, + consecutive_failures: 0, + } + } + + /// How long ago the most recent round completed. + pub fn last_round_age(&self) -> Option { + self.last_round_at.map(|at| at.elapsed()) + } + + /// How long ago the most recent successful round completed. + pub fn last_success_age(&self) -> Option { + self.last_success_at.map(|at| at.elapsed()) + } +} + +/// Lock-free publication slot for [`ProbeStatus`]. +/// +/// Written by the single worker task and read by anyone; readers never block +/// the probe and the probe never blocks a readiness handler. +struct ProbeState { + status: ArcSwap, +} + +impl ProbeState { + fn new() -> Self { + Self { + status: ArcSwap::from_pointee(ProbeStatus::pending()), + } + } + + fn status(&self) -> Arc { + self.status.load_full() + } + + /// Fold one round result into the published snapshot. + /// + /// Safe as a plain load-then-store because the owning worker task is the + /// only writer. + fn publish(&self, result: ProbeResult) { + let previous = self.status.load(); + let now = Instant::now(); + let (last_success_at, last_success_unix_secs) = match result { + ProbeResult::Success => (Some(now), Some(jiff::Timestamp::now().as_second())), + _ => (previous.last_success_at, previous.last_success_unix_secs), + }; + let consecutive_failures = match result { + ProbeResult::Failure(_) => previous.consecutive_failures.saturating_add(1), + _ => 0, + }; + self.status.store(Arc::new(ProbeStatus { + result, + last_round_at: Some(now), + last_success_at, + last_success_unix_secs, + consecutive_failures, + })); + record_state(result, last_success_unix_secs, consecutive_failures); + } +} + +/// Owner of one service version's probe worker. +/// +/// Mirrors the deletion worker's ownership model: cancelling stops the loop, +/// and dropping the handle — a service version replaced without an explicit +/// stop — cancels as a safety net. +pub struct ProbeHandle { + state: Arc, + cancel: CancellationToken, + task: std::sync::Mutex>>, +} + +impl ProbeHandle { + /// Latest published snapshot. + pub fn status(&self) -> Arc { + self.state.status() + } + + /// Stop the worker. The task observes the cancelled token on its next poll. + pub(crate) fn shutdown(&self) { + self.cancel.cancel(); + if let Ok(mut task) = self.task.lock() { + drop(task.take()); + } + } + + #[cfg(test)] + pub(crate) fn is_cancelled(&self) -> bool { + self.cancel.is_cancelled() + } +} + +impl Drop for ProbeHandle { + fn drop(&mut self) { + self.cancel.cancel(); + } +} + +/// Start the probe for a service version. +/// +/// Returns `None` when the backend advertises no data key round trip, or when +/// the operator disabled the probe; both cases leave no background task behind. +pub(crate) fn spawn_probe_worker(backend: Arc) -> Option { + let capabilities = backend.capabilities(); + if !capabilities.generate_data_key || !capabilities.decrypt { + return None; + } + let interval = configured_interval()?; + let state = Arc::new(ProbeState::new()); + let cancel = CancellationToken::new(); + let task = ProbeWorker::new(backend, state.clone(), interval).spawn(cancel.clone()); + Some(ProbeHandle { + state, + cancel, + task: std::sync::Mutex::new(Some(task)), + }) +} + +/// Probe interval from the environment, or `None` when the probe is disabled. +fn configured_interval() -> Option { + parse_interval(std::env::var(ENV_KMS_PROBE_INTERVAL_SECS).ok().as_deref()) +} + +/// Interpret the configured interval: unset or unparsable falls back to the +/// default, `0` disables the probe, and anything below [`MIN_PROBE_INTERVAL`] +/// is raised to it so a misconfigured value cannot turn the probe into a load +/// generator against the backend. +fn parse_interval(value: Option<&str>) -> Option { + let Some(value) = value else { + return Some(DEFAULT_PROBE_INTERVAL); + }; + let Ok(seconds) = value.trim().parse::() else { + warn!( + variable = ENV_KMS_PROBE_INTERVAL_SECS, + "ignoring unparsable KMS probe interval; falling back to the default" + ); + return Some(DEFAULT_PROBE_INTERVAL); + }; + if seconds == 0 { + return None; + } + Some(Duration::from_secs(seconds).max(MIN_PROBE_INTERVAL)) +} + +/// Encryption context bound to every probe data key. +/// +/// Static and free of per-node values, so material generated by one node stays +/// decryptable by every other node running the same probe. +fn probe_encryption_context() -> HashMap { + HashMap::from([("rustfs-purpose".to_string(), "kms-probe".to_string())]) +} + +/// Outcome of making sure the dedicated probe key exists. +enum KeyProvisioning { + Ready, + Unsupported, + Failed, +} + +pub(crate) struct ProbeWorker { + backend: Arc, + state: Arc, + interval: Duration, +} + +impl ProbeWorker { + fn new(backend: Arc, state: Arc, interval: Duration) -> Self { + Self { + backend, + state, + interval, + } + } + + fn spawn(self, cancel: CancellationToken) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { self.run(cancel).await }) + } + + async fn run(self, cancel: CancellationToken) { + describe_metrics(); + let mut ticker = tokio::time::interval(self.interval); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + _ = cancel.cancelled() => { + debug!("KMS probe worker stopped"); + return; + } + _ = ticker.tick() => {} + } + let started = Instant::now(); + let result = self.round().await; + record_round(result, started.elapsed()); + self.state.publish(result); + if result == ProbeResult::Unsupported { + // A capability gap does not change while a service version + // lives, so retrying it forever would only produce noise. + info!("KMS backend cannot host the synthetic probe key; probe disabled for this service version"); + return; + } + } + } + + /// Run one round. Exposed separately so tests can drive the round trip + /// deterministically without the interval loop. + pub(crate) async fn round(&self) -> ProbeResult { + let capabilities = self.backend.capabilities(); + if !capabilities.generate_data_key || !capabilities.decrypt { + return ProbeResult::Unsupported; + } + match self.ensure_probe_key().await { + KeyProvisioning::Ready => {} + KeyProvisioning::Unsupported => return ProbeResult::Unsupported, + KeyProvisioning::Failed => return ProbeResult::Failure(ProbeFailureKind::KeyProvisioning), + } + + let mut generated = match self + .backend + .generate_data_key(GenerateDataKeyRequest { + key_id: PROBE_KEY_ID.to_string(), + key_spec: KeySpec::Aes256, + encryption_context: probe_encryption_context(), + }) + .await + { + Ok(generated) => generated, + Err(error) => { + warn!(%error, "KMS probe could not generate a data key"); + return ProbeResult::Failure(ProbeFailureKind::Generate); + } + }; + + let decrypted = self + .backend + .decrypt(DecryptRequest { + ciphertext: generated.ciphertext_blob.clone(), + encryption_context: probe_encryption_context(), + grant_tokens: Vec::new(), + }) + .await; + + let result = match decrypted { + Ok(mut response) => { + let round_trips = bool::from(response.plaintext.ct_eq(&generated.plaintext_key)); + response.plaintext.zeroize(); + if round_trips { + debug!("KMS probe round trip succeeded"); + ProbeResult::Success + } else { + // The backend answered both calls but the material does not + // survive the round trip: report it as loudly as an outage. + warn!("KMS probe decrypted material does not match the generated data key"); + ProbeResult::Failure(ProbeFailureKind::Mismatch) + } + } + Err(error) => { + warn!(%error, "KMS probe could not decrypt its generated data key"); + ProbeResult::Failure(ProbeFailureKind::Decrypt) + } + }; + generated.plaintext_key.zeroize(); + result + } + + /// Create the probe key if it is missing. + /// + /// Idempotent by construction: an existing key is used as is, and a create + /// that loses the race against another node reports the key as ready + /// instead of failing the round. + async fn ensure_probe_key(&self) -> KeyProvisioning { + match self + .backend + .describe_key(DescribeKeyRequest { + key_id: PROBE_KEY_ID.to_string(), + }) + .await + { + Ok(_) => return KeyProvisioning::Ready, + Err(KmsError::KeyNotFound { .. }) => {} + Err(error) if is_capability_gap(&error) => return KeyProvisioning::Unsupported, + Err(error) => { + warn!(%error, "KMS probe could not describe its probe key"); + return KeyProvisioning::Failed; + } + } + + match self + .backend + .create_key(CreateKeyRequest { + key_name: Some(PROBE_KEY_ID.to_string()), + key_usage: KeyUsage::EncryptDecrypt, + description: Some(PROBE_KEY_DESCRIPTION.to_string()), + ..Default::default() + }) + .await + { + Ok(_) => { + info!(key_id = PROBE_KEY_ID, "created the KMS synthetic probe key"); + KeyProvisioning::Ready + } + // Another node created the key between the describe and the create. + Err(KmsError::KeyAlreadyExists { .. }) => KeyProvisioning::Ready, + Err(error) if is_capability_gap(&error) => KeyProvisioning::Unsupported, + Err(error) => { + warn!(%error, "KMS probe could not create its probe key"); + KeyProvisioning::Failed + } + } + } +} + +/// Whether an error means "this backend will never host a probe key". +/// +/// Read-only single-key backends reject key creation outright, and backends +/// without the operation report an unsupported capability. Both are properties +/// of the configuration rather than symptoms of an outage, so they stop the +/// probe instead of counting as failures. +fn is_capability_gap(error: &KmsError) -> bool { + matches!(error, KmsError::UnsupportedCapability { .. } | KmsError::InvalidOperation { .. }) +} + +// --------------------------------------------------------------------------- +// Metrics +// +// Label values are exclusively static outcome classes; the probe key id is a +// compile-time constant and still never becomes a label, so the metric shape +// stays independent of how many keys or nodes exist. +// --------------------------------------------------------------------------- + +/// Counter: probe rounds completed, by `result`. +const METRIC_PROBE_ROUNDS_TOTAL: &str = "rustfs_kms_probe_rounds_total"; +/// Counter: failed probe rounds, by `failure_kind`. +const METRIC_PROBE_FAILURES_TOTAL: &str = "rustfs_kms_probe_failures_total"; +/// Histogram: wall-clock duration of one probe round, in seconds, by `result`. +const METRIC_PROBE_DURATION_SECONDS: &str = "rustfs_kms_probe_duration_seconds"; +/// Gauge: Unix timestamp of the most recent successful round. +const METRIC_PROBE_LAST_SUCCESS_TIMESTAMP: &str = "rustfs_kms_probe_last_success_timestamp_seconds"; +/// Gauge: failed rounds since the last success. +const METRIC_PROBE_CONSECUTIVE_FAILURES: &str = "rustfs_kms_probe_consecutive_failures"; + +/// Register metric descriptions once per process. +fn describe_metrics() { + static DESCRIBE: std::sync::Once = std::sync::Once::new(); + DESCRIBE.call_once(|| { + metrics::describe_counter!(METRIC_PROBE_ROUNDS_TOTAL, "Total synthetic KMS probe rounds completed, by result"); + metrics::describe_counter!( + METRIC_PROBE_FAILURES_TOTAL, + "Total failed synthetic KMS probe rounds, by the round-trip stage that failed" + ); + metrics::describe_histogram!( + METRIC_PROBE_DURATION_SECONDS, + "Wall-clock duration of one synthetic KMS probe round, in seconds" + ); + metrics::describe_gauge!( + METRIC_PROBE_LAST_SUCCESS_TIMESTAMP, + "Unix timestamp of the most recent successful synthetic KMS probe round" + ); + metrics::describe_gauge!( + METRIC_PROBE_CONSECUTIVE_FAILURES, + "Synthetic KMS probe rounds that have failed since the last success" + ); + }); +} + +/// Record one completed round. +/// +/// An unsupported backend is deliberately counted as its own result and never +/// as a failure, so alerts on the failure counter stay silent for deployments +/// whose backend cannot host the probe. +fn record_round(result: ProbeResult, elapsed: Duration) { + metrics::counter!(METRIC_PROBE_ROUNDS_TOTAL, "result" => result.as_label()).increment(1); + if let ProbeResult::Failure(kind) = result { + metrics::counter!(METRIC_PROBE_FAILURES_TOTAL, "failure_kind" => kind.as_label()).increment(1); + } + metrics::histogram!(METRIC_PROBE_DURATION_SECONDS, "result" => result.as_label()).record(elapsed.as_secs_f64()); +} + +/// Mirror the published snapshot into the gauges. +fn record_state(result: ProbeResult, last_success_unix_secs: Option, consecutive_failures: u32) { + if result == ProbeResult::Success + && let Some(timestamp) = last_success_unix_secs + { + // Only ever moved forward by a success, so a failing probe leaves the + // gauge at the last known good time and its age keeps growing. + metrics::gauge!(METRIC_PROBE_LAST_SUCCESS_TIMESTAMP).set(timestamp as f64); + } + metrics::gauge!(METRIC_PROBE_CONSECUTIVE_FAILURES).set(f64::from(consecutive_failures)); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::backends::local::LocalKmsBackend; + use crate::backends::static_kms::StaticKmsBackend; + use crate::config::KmsConfig; + use crate::types::{ + CancelKeyDeletionRequest, CancelKeyDeletionResponse, CreateKeyResponse, DecryptResponse, DeleteKeyRequest, + DeleteKeyResponse, DescribeKeyResponse, EncryptRequest, EncryptResponse, GenerateDataKeyResponse, KeyMetadata, KeyState, + ListKeysRequest, ListKeysResponse, + }; + use async_trait::async_trait; + use base64::Engine as _; + use metrics_util::MetricKind; + use metrics_util::debugging::{DebugValue, DebuggingRecorder}; + use std::future::Future; + + async fn local_backend(temp_dir: &tempfile::TempDir) -> Arc { + let config = KmsConfig::local(temp_dir.path().to_path_buf()).with_insecure_development_defaults(); + Arc::new(LocalKmsBackend::new(config).await.expect("local backend should build")) + } + + async fn static_backend() -> Arc { + let config = + KmsConfig::static_kms("static-key".to_string(), base64::engine::general_purpose::STANDARD.encode([0x42u8; 32])); + Arc::new(StaticKmsBackend::new(config).await.expect("static backend should build")) + } + + fn worker(backend: Arc) -> ProbeWorker { + ProbeWorker::new(backend, Arc::new(ProbeState::new()), DEFAULT_PROBE_INTERVAL) + } + + /// Backend that answers every probe call but corrupts the round trip in a + /// configurable way, so each failure classification can be provoked without + /// a real backend outage. + struct BrokenBackend { + fault: Fault, + } + + #[derive(Clone, Copy)] + enum Fault { + Generate, + Decrypt, + /// Decrypt succeeds but returns material other than the generated key. + Mismatch, + } + + impl BrokenBackend { + fn arc(fault: Fault) -> Arc { + Arc::new(Self { fault }) + } + + fn probe_key_metadata() -> KeyMetadata { + KeyMetadata { + key_id: PROBE_KEY_ID.to_string(), + key_state: KeyState::Enabled, + key_usage: KeyUsage::EncryptDecrypt, + description: None, + creation_date: jiff::Zoned::now(), + deletion_date: None, + origin: "test".to_string(), + key_manager: "test".to_string(), + tags: HashMap::new(), + } + } + } + + #[async_trait] + impl KmsBackend for BrokenBackend { + async fn create_key(&self, _request: CreateKeyRequest) -> crate::error::Result { + unimplemented!("the probe key always exists on this backend") + } + + async fn encrypt(&self, _request: EncryptRequest) -> crate::error::Result { + unimplemented!("not exercised by probe tests") + } + + async fn decrypt(&self, _request: DecryptRequest) -> crate::error::Result { + match self.fault { + Fault::Decrypt => Err(KmsError::backend_error("decrypt unavailable")), + Fault::Mismatch => Ok(DecryptResponse { + plaintext: vec![0xffu8; 32], + key_id: PROBE_KEY_ID.to_string(), + encryption_algorithm: None, + }), + Fault::Generate => unreachable!("generation fails before decrypt is reached"), + } + } + + async fn generate_data_key(&self, _request: GenerateDataKeyRequest) -> crate::error::Result { + match self.fault { + Fault::Generate => Err(KmsError::backend_error("data key generation unavailable")), + _ => Ok(GenerateDataKeyResponse { + key_id: PROBE_KEY_ID.to_string(), + plaintext_key: vec![0x11u8; 32], + ciphertext_blob: vec![0x22u8; 48], + }), + } + } + + async fn describe_key(&self, _request: DescribeKeyRequest) -> crate::error::Result { + Ok(DescribeKeyResponse { + key_metadata: Self::probe_key_metadata(), + }) + } + + async fn list_keys(&self, _request: ListKeysRequest) -> crate::error::Result { + unimplemented!("not exercised by probe tests") + } + + async fn delete_key(&self, _request: DeleteKeyRequest) -> crate::error::Result { + unimplemented!("not exercised by probe tests") + } + + async fn cancel_key_deletion( + &self, + _request: CancelKeyDeletionRequest, + ) -> crate::error::Result { + unimplemented!("not exercised by probe tests") + } + + async fn health_check(&self) -> crate::error::Result { + Ok(true) + } + } + + #[tokio::test] + async fn round_creates_the_probe_key_once_and_round_trips_it() { + let temp_dir = tempfile::tempdir().expect("temp dir"); + let backend = local_backend(&temp_dir).await; + let worker = worker(backend.clone()); + + assert_eq!(worker.round().await, ProbeResult::Success); + let created = backend + .describe_key(DescribeKeyRequest { + key_id: PROBE_KEY_ID.to_string(), + }) + .await + .expect("the first round must create the probe key"); + assert_eq!(created.key_metadata.key_id, PROBE_KEY_ID); + + // A second round reuses the key instead of failing on the existing one, + // which is also what a second node in the deployment would do. + assert_eq!(worker.round().await, ProbeResult::Success); + assert_eq!( + backend + .list_keys(ListKeysRequest { + limit: Some(100), + marker: None, + usage_filter: None, + status_filter: None, + }) + .await + .expect("list keys") + .keys + .iter() + .filter(|key| key.key_id == PROBE_KEY_ID) + .count(), + 1, + "probe key creation must be idempotent" + ); + } + + #[test] + fn read_only_backend_is_reported_unsupported_without_failures() { + let (snapshot, result) = record_metrics(|| { + Box::pin(async { + let worker = worker(static_backend().await); + let result = worker.round().await; + record_round(result, Duration::from_millis(1)); + result + }) + }); + + assert_eq!(result, ProbeResult::Unsupported); + assert_eq!(counter_value(&snapshot, METRIC_PROBE_ROUNDS_TOTAL, &[("result", "unsupported")]), 1); + assert_eq!( + counter_value(&snapshot, METRIC_PROBE_FAILURES_TOTAL, &[]), + 0, + "an unsupported backend must never raise probe failures" + ); + } + + #[tokio::test] + async fn each_round_trip_stage_gets_its_own_failure_classification() { + for (fault, expected) in [ + (Fault::Generate, ProbeFailureKind::Generate), + (Fault::Decrypt, ProbeFailureKind::Decrypt), + (Fault::Mismatch, ProbeFailureKind::Mismatch), + ] { + assert_eq!(worker(BrokenBackend::arc(fault)).round().await, ProbeResult::Failure(expected)); + } + } + + #[test] + fn metrics_and_status_record_success_then_failure() { + let (snapshot, status) = record_metrics(|| { + Box::pin(async { + let temp_dir = tempfile::tempdir().expect("temp dir"); + let state = Arc::new(ProbeState::new()); + let healthy = ProbeWorker::new(local_backend(&temp_dir).await, state.clone(), DEFAULT_PROBE_INTERVAL); + let round = healthy.round().await; + record_round(round, Duration::from_millis(10)); + state.publish(round); + + let broken = ProbeWorker::new(BrokenBackend::arc(Fault::Mismatch), state.clone(), DEFAULT_PROBE_INTERVAL); + for _ in 0..2 { + let round = broken.round().await; + record_round(round, Duration::from_millis(10)); + state.publish(round); + } + state.status() + }) + }); + + assert_eq!(status.result, ProbeResult::Failure(ProbeFailureKind::Mismatch)); + assert_eq!(status.consecutive_failures, 2); + assert!( + status.last_success_at.is_some() && status.last_success_unix_secs.is_some(), + "a failing probe must keep reporting when it last succeeded" + ); + + assert_eq!(counter_value(&snapshot, METRIC_PROBE_ROUNDS_TOTAL, &[("result", "success")]), 1); + assert_eq!(counter_value(&snapshot, METRIC_PROBE_ROUNDS_TOTAL, &[("result", "failure")]), 2); + assert_eq!(counter_value(&snapshot, METRIC_PROBE_FAILURES_TOTAL, &[("failure_kind", "mismatch")]), 2); + assert_eq!(histogram_values(&snapshot, METRIC_PROBE_DURATION_SECONDS, &[]).len(), 3); + assert!( + gauge_value(&snapshot, METRIC_PROBE_LAST_SUCCESS_TIMESTAMP).is_some_and(|timestamp| timestamp > 0.0), + "the last success gauge must carry a wall-clock timestamp" + ); + assert_eq!(gauge_value(&snapshot, METRIC_PROBE_CONSECUTIVE_FAILURES), Some(2.0)); + } + + #[tokio::test(start_paused = true)] + async fn worker_publishes_status_and_stops_on_cancel() { + let temp_dir = tempfile::tempdir().expect("temp dir"); + let state = Arc::new(ProbeState::new()); + assert_eq!(state.status().result, ProbeResult::Pending); + + let cancel = CancellationToken::new(); + let task = ProbeWorker::new(local_backend(&temp_dir).await, state.clone(), DEFAULT_PROBE_INTERVAL).spawn(cancel.clone()); + + // The paused clock auto-advances through the worker's interval ticks. + let mut observed = None; + for _ in 0..100 { + tokio::time::sleep(Duration::from_secs(1)).await; + if state.status().result != ProbeResult::Pending { + observed = Some(state.status()); + break; + } + } + let observed = observed.expect("the worker loop must publish a round result"); + assert_eq!(observed.result, ProbeResult::Success); + assert_eq!(observed.consecutive_failures, 0); + assert!(observed.last_success_at.is_some()); + + cancel.cancel(); + task.await.expect("worker task must stop after cancellation"); + } + + #[tokio::test(start_paused = true)] + async fn unsupported_backend_stops_the_worker_without_cancellation() { + let state = Arc::new(ProbeState::new()); + let cancel = CancellationToken::new(); + let task = ProbeWorker::new(static_backend().await, state.clone(), DEFAULT_PROBE_INTERVAL).spawn(cancel.clone()); + + task.await.expect("worker must exit on its own for an unsupported backend"); + assert_eq!(state.status().result, ProbeResult::Unsupported); + assert!(!cancel.is_cancelled()); + } + + #[test] + fn interval_configuration_defaults_clamps_and_disables() { + assert_eq!(parse_interval(None), Some(DEFAULT_PROBE_INTERVAL)); + assert_eq!(parse_interval(Some("not-a-number")), Some(DEFAULT_PROBE_INTERVAL)); + assert_eq!(parse_interval(Some(" 120 ")), Some(Duration::from_secs(120))); + assert_eq!(parse_interval(Some("1")), Some(MIN_PROBE_INTERVAL)); + assert_eq!(parse_interval(Some("0")), None, "zero must disable the probe"); + } + + // -- Metric assertion helpers, mirroring the policy engine's tests -------- + + type MetricEntry = ( + metrics_util::CompositeKey, + Option, + Option, + DebugValue, + ); + + /// Run `test` on a paused current-thread runtime under a debugging recorder + /// and return one snapshot of everything it emitted. + fn record_metrics(test: impl FnOnce() -> std::pin::Pin>>) -> (Vec, Out) { + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + let out = metrics::with_local_recorder(&recorder, || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_time() + .start_paused(true) + .build() + .expect("current-thread runtime must build"); + runtime.block_on(test()) + }); + (snapshotter.snapshot().into_vec(), out) + } + + fn labels_match(key: &metrics::Key, labels: &[(&str, &str)]) -> bool { + labels + .iter() + .all(|(label, expected)| key.labels().any(|l| l.key() == *label && l.value() == *expected)) + } + + fn counter_value(snapshot: &[MetricEntry], name: &str, labels: &[(&str, &str)]) -> u64 { + snapshot + .iter() + .filter_map(|(composite, _unit, _description, value)| { + let matches = composite.kind() == MetricKind::Counter + && composite.key().name() == name + && labels_match(composite.key(), labels); + match (matches, value) { + (true, DebugValue::Counter(count)) => Some(*count), + _ => None, + } + }) + .sum() + } + + fn histogram_values(snapshot: &[MetricEntry], name: &str, labels: &[(&str, &str)]) -> Vec { + snapshot + .iter() + .filter_map(|(composite, _unit, _description, value)| { + let matches = composite.kind() == MetricKind::Histogram + && composite.key().name() == name + && labels_match(composite.key(), labels); + match (matches, value) { + (true, DebugValue::Histogram(values)) => Some(values), + _ => None, + } + }) + .flatten() + .map(|value| value.into_inner()) + .collect() + } + + fn gauge_value(snapshot: &[MetricEntry], name: &str) -> Option { + snapshot.iter().find_map(|(composite, _unit, _description, value)| { + let matches = composite.kind() == MetricKind::Gauge && composite.key().name() == name; + match (matches, value) { + (true, DebugValue::Gauge(gauge)) => Some(gauge.into_inner()), + _ => None, + } + }) + } +} diff --git a/crates/kms/src/service_manager.rs b/crates/kms/src/service_manager.rs index 5acd1a15a..77104d815 100644 --- a/crates/kms/src/service_manager.rs +++ b/crates/kms/src/service_manager.rs @@ -21,6 +21,7 @@ use crate::config::{BackendConfig, KmsConfig}; use crate::deletion_worker::{DeletionReferenceChecker, DeletionWorker}; use crate::error::{KmsError, Result}; use crate::manager::KmsManager; +use crate::probe::{ProbeHandle, ProbeStatus, spawn_probe_worker}; use crate::service::ObjectEncryptionService; use arc_swap::ArcSwap; use sha2::{Digest, Sha256}; @@ -121,13 +122,23 @@ struct ServiceVersion { /// Background deletion worker owned by this service version, if the /// backend supports deletion scheduling deletion_worker: Option>, + /// Background synthetic probe owned by this service version, if the + /// backend can round-trip a data key and the probe is enabled + probe_worker: Option>, } impl ServiceVersion { - fn shutdown_deletion_worker(&self) { + /// Stop the background workers this version owns. + /// + /// The credential renewal task is recycled separately because its shutdown + /// is asynchronous. + fn shutdown_background_workers(&self) { if let Some(worker) = &self.deletion_worker { worker.shutdown(); } + if let Some(probe) = &self.probe_worker { + probe.shutdown(); + } } } @@ -399,7 +410,7 @@ impl KmsServiceManager { // Note: Existing Arc references will keep the service alive until operations complete let state = self.state.load_full(); if let Some(current) = state.current_service.as_ref() { - current.shutdown_deletion_worker(); + current.shutdown_background_workers(); } self.state.store(Arc::new(RuntimeState { config: state.config.clone(), @@ -534,6 +545,19 @@ impl KmsServiceManager { self.state.load().current_service.as_ref().map(|sv| sv.version) } + /// Latest synthetic probe snapshot of the running service version. + /// + /// `None` when KMS is not running, when the backend cannot round-trip a + /// data key, or when the probe is disabled. The read is lock-free and + /// performs no backend I/O, so a health endpoint can consult it without + /// turning an external KMS hiccup into probe pressure of its own; the + /// staleness and consecutive-failure thresholds belong to the caller. + pub fn probe_status(&self) -> Option> { + let state = self.state.load(); + let service_version = state.current_service.as_ref()?; + Some(service_version.probe_worker.as_ref()?.status()) + } + /// Health check for the KMS service pub async fn health_check(&self) -> Result { let checked_state = self.state.load_full(); @@ -620,6 +644,7 @@ impl KmsServiceManager { manager: kms_manager, credential_task, deletion_worker: None, + probe_worker: None, }) } @@ -633,9 +658,13 @@ impl KmsServiceManager { fn publish_running(&self, config: KmsConfig, mut service_version: ServiceVersion) { if let Some(previous) = self.state.load().current_service.as_ref() { - previous.shutdown_deletion_worker(); + previous.shutdown_background_workers(); } service_version.deletion_worker = self.spawn_deletion_worker(&config, &service_version); + // Started at publish time for the same reason as the deletion worker: + // a start or reconfigure candidate that never becomes current must not + // leak a running task or publish probe results nobody asked for. + service_version.probe_worker = spawn_probe_worker(service_version.manager.backend()).map(Arc::new); self.state.store(Arc::new(RuntimeState { config: Some(config), status: KmsServiceStatus::Running, @@ -1010,6 +1039,49 @@ mod tests { ); } + #[tokio::test] + async fn probe_worker_follows_the_service_lifecycle() { + use tempfile::TempDir; + + let key_dir = TempDir::new().expect("create local KMS directory"); + let mut config = KmsConfig::local(key_dir.path().to_path_buf()); + config.allow_insecure_dev_defaults = true; + let manager = KmsServiceManager::new(); + manager.configure(config).await.expect("configure local KMS"); + manager.start().await.expect("start local KMS"); + + let first_probe = manager + .state + .load() + .current_service + .as_ref() + .expect("running service") + .probe_worker + .clone() + .expect("a backend with a data key round trip must run a probe"); + assert!(!first_probe.is_cancelled()); + assert!(manager.probe_status().is_some(), "a running service must publish probe status"); + + // Replacing the service version replaces (and cancels) its probe. + manager.restart().await.expect("restart"); + assert!(first_probe.is_cancelled(), "replaced version's probe must be cancelled"); + let second_probe = manager + .state + .load() + .current_service + .as_ref() + .expect("running service") + .probe_worker + .clone() + .expect("restarted service must run a fresh probe"); + assert!(!second_probe.is_cancelled()); + + // Stopping the service stops its probe and withdraws the snapshot. + manager.stop().await.expect("stop"); + assert!(second_probe.is_cancelled(), "stop must cancel the probe worker"); + assert!(manager.probe_status().is_none(), "a stopped service must publish no status"); + } + #[tokio::test] async fn reconfigure_allows_safe_local_runtime_settings_only() { use tempfile::TempDir;