feat(kms): enforce shared key state machine across backends (#5489)

* feat(kms): enforce shared key state machine across backends Unify the key state x operation matrix behind a single gate in backends/mod.rs and wire it into the Local, Vault KV2 and Vault Transit backends: Disabled keys reject encryption, data key generation and rotation while still allowing decryption and lifecycle recovery; PendingDeletion keys reject everything except decryption and cancellation (including repeated deletion scheduling); cancellation now requires an actual pending deletion everywhere. This closes the missing gates on KV2 encrypt/generate and Local generate_data_key, and stops enable_key from silently reverting a pending deletion. Decryption is deliberately left ungated in Disabled/PendingDeletion — an explicit, documented and tested deviation from AWS KMS, since gating it would break reads of existing objects the moment a key is disabled. Add shared contract tests driving the full matrix offline for Local (and via ignored tests against a live Vault for KV2/Transit), a stateless contract for Static, an SSE-shaped regression proving existing envelopes stay decryptable after disable, and a pin on the known-risk Enabled default of Transit's synthesized metadata fallback. Refs rustfs/backlog#1571 (part of rustfs/backlog#1562) * feat(kms): persist deletion deadlines and run a restartable deletion worker (#5491)
2026-08-30 00:47:13 +00:00 · 2026-07-31 07:24:39 +08:00
parent 3921336b23
commit 2e29c330a9
10 changed files with 1394 additions and 38 deletions
@@ -18,7 +18,10 @@ use crate::backends::vault_credentials::{
    CredentialTaskHandle, VaultClientHandle, VaultConnectionSettings, VaultCredentialPolicy, VaultCredentialProvider,
    token_source_for,
 };
-use crate::backends::{BackendCapabilities, BackendInfo, KmsBackend, KmsClient};
+use crate::backends::{
+    BackendCapabilities, BackendInfo, ExpiredKeyRemoval, KmsBackend, KmsClient, StateGatedOperation, ensure_key_state_permits,
+    ensure_key_status_permits,
+};
 use crate::config::{KmsConfig, VaultConfig};
 use crate::encryption::{AesDekCrypto, DataKeyEnvelope, DekCrypto, generate_key_material};
 use crate::error::{KmsError, Result};
@@ -64,6 +67,10 @@ struct VaultKeyData {
    metadata: HashMap<String, String>,
    /// Key tags
    tags: HashMap<String, String>,
+    /// Scheduled deletion deadline; absent on records written before deadline
+    /// persistence landed, so it must stay optional for backward compatibility.
+    #[serde(default)]
+    deletion_date: Option<Zoned>,
    /// Encrypted key material (base64 encoded)
    encrypted_key_material: String,
    /// Version that pre-versioning envelopes (no `master_key_version`) resolve to.
@@ -381,6 +388,7 @@ impl VaultKmsClient {
            description: request.description.clone(),
            metadata: existing_key_data.metadata.clone(),
            tags: request.tags.clone(),
+            deletion_date: existing_key_data.deletion_date.clone(),
            encrypted_key_material: existing_key_data.encrypted_key_material.clone(), // Preserve the key material
            baseline_version: existing_key_data.baseline_version,
        };
@@ -472,6 +480,9 @@ impl KmsClient for VaultKmsClient {
    async fn generate_data_key(&self, request: &GenerateKeyRequest, _context: Option<&OperationContext>) -> Result<DataKeyInfo> {
        debug!("Generating data key for master key: {}", request.master_key_id);

+        let key_data = self.get_key_data(&request.master_key_id).await?;
+        ensure_key_status_permits(&request.master_key_id, &key_data.status, StateGatedOperation::GenerateDataKey)?;
+
        // Generate random data key material using the existing method
        let plaintext_key = generate_key_material(&request.key_spec)?;

@@ -510,8 +521,9 @@ impl KmsClient for VaultKmsClient {
    async fn encrypt(&self, request: &EncryptRequest, _context: Option<&OperationContext>) -> Result<EncryptResponse> {
        debug!("Encrypting data with key: {}", request.key_id);

-        // Get the master key
+        // Get the master key and verify its state allows encryption
        let key_data = self.get_key_data(&request.key_id).await?;
+        ensure_key_status_permits(&request.key_id, &key_data.status, StateGatedOperation::Encrypt)?;
        let key_material = self.decrypt_key_material(&key_data.encrypted_key_material).await?;

        // For simplicity, we'll use a basic encryption approach
@@ -593,6 +605,7 @@ impl KmsClient for VaultKmsClient {
            description: None,
            metadata: HashMap::new(),
            tags: HashMap::new(),
+            deletion_date: None,
            encrypted_key_material: encrypted_material,
            baseline_version: None,
        };
@@ -611,6 +624,7 @@ impl KmsClient for VaultKmsClient {
            created_at: key_data.created_at,
            rotated_at: None,
            created_by: None,
+            deletion_date: None,
        };

        debug!(key_id, "Vault KMS master key created");
@@ -678,6 +692,7 @@ impl KmsClient for VaultKmsClient {
        debug!("Enabling key: {}", key_id);

        let mut key_data = self.get_key_data(key_id).await?;
+        ensure_key_status_permits(key_id, &key_data.status, StateGatedOperation::Enable)?;
        key_data.status = KeyStatus::Active;
        self.store_key_data(key_id, &key_data).await?;

@@ -689,6 +704,7 @@ impl KmsClient for VaultKmsClient {
        debug!("Disabling key: {}", key_id);

        let mut key_data = self.get_key_data(key_id).await?;
+        ensure_key_status_permits(key_id, &key_data.status, StateGatedOperation::Disable)?;
        key_data.status = KeyStatus::Disabled;
        self.store_key_data(key_id, &key_data).await?;

@@ -699,13 +715,15 @@ impl KmsClient for VaultKmsClient {
    async fn schedule_key_deletion(
        &self,
        key_id: &str,
-        _pending_window_days: u32,
+        pending_window_days: u32,
        _context: Option<&OperationContext>,
    ) -> Result<()> {
        debug!("Scheduling key deletion: {}", key_id);

        let mut key_data = self.get_key_data(key_id).await?;
+        ensure_key_status_permits(key_id, &key_data.status, StateGatedOperation::ScheduleDeletion)?;
        key_data.status = KeyStatus::PendingDeletion;
+        key_data.deletion_date = Some(Zoned::now() + Duration::from_secs(pending_window_days as u64 * 86400));
        self.store_key_data(key_id, &key_data).await?;

        debug!(key_id, "Vault KMS key deletion scheduled");
@@ -716,7 +734,11 @@ impl KmsClient for VaultKmsClient {
        debug!("Canceling key deletion: {}", key_id);

        let mut key_data = self.get_key_data(key_id).await?;
+        if key_data.status != KeyStatus::PendingDeletion {
+            return Err(KmsError::invalid_key_state(format!("Key {key_id} is not pending deletion")));
+        }
        key_data.status = KeyStatus::Active;
+        key_data.deletion_date = None;
        self.store_key_data(key_id, &key_data).await?;

        debug!(key_id, "Vault KMS key deletion canceled");
@@ -818,6 +840,7 @@ impl KmsClient for VaultKmsClient {
            created_at: key_data.created_at.clone(),
            rotated_at: Some(Zoned::now()),
            created_by: None,
+            deletion_date: key_data.deletion_date.clone(),
        })
    }

@@ -860,6 +883,12 @@ pub struct VaultKmsBackend {
 }

 impl VaultKmsBackend {
+    /// Lifecycle driver for the shared state-machine contract tests.
+    #[cfg(test)]
+    pub(crate) fn lifecycle_client(&self) -> &VaultKmsClient {
+        &self.client
+    }
+
    /// Create a new VaultKmsBackend
    pub async fn new(config: KmsConfig) -> Result<Self> {
        config.validate()?;
@@ -905,6 +934,7 @@ impl VaultKmsBackend {
            KeyState::Unavailable => KeyStatus::Deleted,
            KeyState::PendingImport => KeyStatus::Disabled, // Treat as disabled until import completes
        };
+        key_data.deletion_date = metadata.deletion_date.clone();

        // Update the key data in Vault storage
        self.client.store_key_data(key_id, &key_data).await?;
@@ -1004,7 +1034,7 @@ impl KmsBackend for VaultKmsBackend {
            key_usage: key_info.usage,
            description: key_info.description,
            creation_date: key_info.created_at,
-            deletion_date: None,
+            deletion_date: key_data.deletion_date.clone(),
            origin: "VAULT".to_string(),
            key_manager: "VAULT".to_string(),
            tags: key_data.tags,
@@ -1033,8 +1063,17 @@ impl KmsBackend for VaultKmsBackend {
        };

        let deletion_date = if request.force_immediate.unwrap_or(false) {
-            // Check if key is already in PendingDeletion state
-            if key_metadata.key_state == KeyState::PendingDeletion {
+            // Check if key is already in PendingDeletion state (or a tombstone
+            // left by a crashed removal, which may simply be completed)
+            if key_metadata.key_state == KeyState::PendingDeletion || key_metadata.key_state == KeyState::Unavailable {
+                // Tombstone first: mark the record Deleted before removing it,
+                // so a crash between the two steps leaves a key that is already
+                // unusable and whose removal can simply be re-run.
+                if key_metadata.key_state == KeyState::PendingDeletion {
+                    let mut key_data = self.client.get_key_data(key_id).await?;
+                    key_data.status = KeyStatus::Deleted;
+                    self.client.store_key_data(key_id, &key_data).await?;
+                }
                // Force immediate deletion: physically delete the key from Vault storage
                self.client.delete_key(key_id).await?;

@@ -1052,6 +1091,8 @@ impl KmsBackend for VaultKmsBackend {
            }
        } else {
            // Schedule for deletion (default 30 days)
+            ensure_key_state_permits(key_id, &key_metadata.key_state, StateGatedOperation::ScheduleDeletion)?;
+
            let days = request.pending_window_in_days.unwrap_or(30);
            if !(7..=30).contains(&days) {
                return Err(crate::error::KmsError::invalid_parameter(
@@ -1120,6 +1161,43 @@ impl KmsBackend for VaultKmsBackend {
            .with_schedule_deletion(true)
            .with_physical_delete(true)
    }
+
+    async fn remove_expired_key(&self, key_id: &str, now: &Zoned) -> Result<ExpiredKeyRemoval> {
+        // Vault KV2 offers no compare-and-swap here, so a cancellation racing
+        // the read below can still lose; the window is a single read-write
+        // gap and the sweep re-reads on every pass.
+        let mut key_data = match self.client.get_key_data(key_id).await {
+            Ok(key_data) => key_data,
+            Err(KmsError::KeyNotFound { .. }) => return Ok(ExpiredKeyRemoval::Removed),
+            Err(error) => return Err(error),
+        };
+        match key_data.status {
+            // Tombstone left by a crashed removal: complete it.
+            KeyStatus::Deleted => {}
+            KeyStatus::PendingDeletion => {
+                match &key_data.deletion_date {
+                    Some(deadline) if deadline <= now => {}
+                    // Not yet due, or a legacy record without a persisted
+                    // deadline — never auto-remove those.
+                    _ => return Ok(ExpiredKeyRemoval::NotExpired),
+                }
+                // Tombstone first: mark the record Deleted before removing it,
+                // so a crash between the two steps leaves a key that is
+                // already unusable and whose removal can simply be re-run.
+                key_data.status = KeyStatus::Deleted;
+                self.client.store_key_data(key_id, &key_data).await?;
+            }
+            KeyStatus::Active | KeyStatus::Disabled => return Ok(ExpiredKeyRemoval::StateChanged),
+        }
+
+        match self.client.delete_key(key_id).await {
+            Ok(()) | Err(KmsError::KeyNotFound { .. }) => {
+                debug!(key_id, "Vault KV2 expired key removed");
+                Ok(ExpiredKeyRemoval::Removed)
+            }
+            Err(error) => Err(error),
+        }
+    }
 }

 #[cfg(test)]
@@ -1284,6 +1362,7 @@ mod tests {
            tags: HashMap::new(),
            encrypted_key_material: general_purpose::STANDARD.encode([0x42u8; 32]),
            baseline_version: Some(1),
+            deletion_date: None,
        };

        let mut value = serde_json::to_value(&key_data).expect("serialize key data");
@@ -1642,4 +1721,41 @@ mod tests {
            "cancel_key_deletion must persist Active status to Vault, not only mutate the response"
        );
    }
+
+    /// The persisted KV2 record round-trips its deletion deadline, and records
+    /// written before the field existed keep deserializing (as None). A revert
+    /// of deadline persistence turns this test red.
+    #[test]
+    fn vault_key_data_deletion_date_round_trips_and_stays_backward_compatible() {
+        let deadline = Zoned::now() + Duration::from_secs(7 * 86400);
+        let key_data = VaultKeyData {
+            algorithm: "AES_256".to_string(),
+            usage: KeyUsage::EncryptDecrypt,
+            created_at: Zoned::now(),
+            status: KeyStatus::PendingDeletion,
+            version: 1,
+            description: None,
+            metadata: HashMap::new(),
+            tags: HashMap::new(),
+            deletion_date: Some(deadline.clone()),
+            encrypted_key_material: "material".to_string(),
+            baseline_version: None,
+        };
+
+        let mut value = serde_json::to_value(&key_data).expect("serialize");
+        let restored: VaultKeyData = serde_json::from_value(value.clone()).expect("round trip");
+        assert_eq!(
+            restored.deletion_date.as_ref().map(Zoned::timestamp),
+            Some(deadline.timestamp()),
+            "deletion deadline must survive the KV2 round trip"
+        );
+
+        value
+            .as_object_mut()
+            .expect("record must be a JSON object")
+            .remove("deletion_date")
+            .expect("current records must carry the field");
+        let legacy: VaultKeyData = serde_json::from_value(value).expect("legacy record must deserialize");
+        assert!(legacy.deletion_date.is_none());
+    }
 }