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-31 09:18:28 +00:00 · 2026-07-31 07:24:39 +08:00
parent 3921336b23
commit 2e29c330a9
10 changed files with 1394 additions and 38 deletions
@@ -14,7 +14,9 @@

 //! Local file-based KMS backend implementation

-use crate::backends::{BackendCapabilities, BackendInfo, KmsBackend, KmsClient};
+use crate::backends::{
+    BackendCapabilities, BackendInfo, ExpiredKeyRemoval, KmsBackend, KmsClient, StateGatedOperation, ensure_key_status_permits,
+};
 use crate::config::KmsConfig;
 use crate::config::LocalConfig;
 use crate::encryption::{AesDekCrypto, DataKeyEnvelope, DekCrypto, generate_key_material};
@@ -425,6 +427,10 @@ struct StoredMasterKey {
    #[serde(with = "crate::time_serde::option_zoned")]
    rotated_at: Option<Zoned>,
    created_by: Option<String>,
+    /// Scheduled deletion deadline; absent on records written before deadline
+    /// persistence landed, so it must stay optional for backward compatibility.
+    #[serde(default, with = "crate::time_serde::option_zoned")]
+    deletion_date: Option<Zoned>,
    /// Encrypted key material (32 bytes encoded in base64 for AES-256)
    encrypted_key_material: String,
    /// Nonce used for encryption
@@ -770,6 +776,7 @@ impl LocalKmsClient {
            created_at: stored_key.created_at,
            rotated_at: stored_key.rotated_at,
            created_by: stored_key.created_by,
+            deletion_date: stored_key.deletion_date,
        })
    }

@@ -843,6 +850,7 @@ impl LocalKmsClient {
            created_at: master_key.created_at.clone(),
            rotated_at: master_key.rotated_at.clone(),
            created_by: master_key.created_by.clone(),
+            deletion_date: master_key.deletion_date.clone(),
            encrypted_key_material,
            nonce,
            at_rest_protection,
@@ -931,9 +939,12 @@ impl LocalKmsClient {

 #[async_trait]
 impl KmsClient for LocalKmsClient {
-    async fn generate_data_key(&self, request: &GenerateKeyRequest, _context: Option<&OperationContext>) -> Result<DataKeyInfo> {
+    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_info = self.describe_key(&request.master_key_id, context).await?;
+        ensure_key_status_permits(&request.master_key_id, &key_info.status, StateGatedOperation::GenerateDataKey)?;
+
        // Generate random data key material
        let key_length = match request.key_spec.as_str() {
            "AES_256" => 32,
@@ -972,14 +983,9 @@ impl KmsClient for LocalKmsClient {
    async fn encrypt(&self, request: &EncryptRequest, context: Option<&OperationContext>) -> Result<EncryptResponse> {
        debug!("Encrypting data with key: {}", request.key_id);

-        // Verify key exists and is active
+        // Verify key exists and its state allows encryption
        let key_info = self.describe_key(&request.key_id, context).await?;
-        if key_info.status != KeyStatus::Active {
-            return Err(KmsError::invalid_operation(format!(
-                "Key {} is not active (status: {:?})",
-                request.key_id, key_info.status
-            )));
-        }
+        ensure_key_status_permits(&request.key_id, &key_info.status, StateGatedOperation::Encrypt)?;

        let (ciphertext, _nonce) = self.encrypt_with_master_key(&request.key_id, &request.plaintext).await?;

@@ -1110,6 +1116,7 @@ impl KmsClient for LocalKmsClient {

        let _write_guard = self.lock_key_for_write(key_id).await;
        let mut master_key = self.load_master_key(key_id).await?;
+        ensure_key_status_permits(key_id, &master_key.status, StateGatedOperation::Enable)?;
        master_key.status = KeyStatus::Active;

        // Preserve the existing key material. Regenerating it on a pure status change would
@@ -1127,6 +1134,7 @@ impl KmsClient for LocalKmsClient {

        let _write_guard = self.lock_key_for_write(key_id).await;
        let mut master_key = self.load_master_key(key_id).await?;
+        ensure_key_status_permits(key_id, &master_key.status, StateGatedOperation::Disable)?;
        master_key.status = KeyStatus::Disabled;

        // Preserve the existing key material (see enable_key): a status change must never
@@ -1141,14 +1149,16 @@ impl KmsClient for LocalKmsClient {
    async fn schedule_key_deletion(
        &self,
        key_id: &str,
-        _pending_window_days: u32,
+        pending_window_days: u32,
        _context: Option<&OperationContext>,
    ) -> Result<()> {
        debug!("Scheduling deletion for key: {}", key_id);

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

        // Preserve the existing key material (see enable_key): scheduling deletion must not
        // regenerate the master key, or cancelling the deletion later would recover a key that
@@ -1165,7 +1175,11 @@ impl KmsClient for LocalKmsClient {

        let _write_guard = self.lock_key_for_write(key_id).await;
        let mut master_key = self.load_master_key(key_id).await?;
+        if master_key.status != KeyStatus::PendingDeletion {
+            return Err(KmsError::invalid_key_state(format!("Key {key_id} is not pending deletion")));
+        }
        master_key.status = KeyStatus::Active;
+        master_key.deletion_date = None;

        // Preserve the existing key material (see enable_key): cancelling deletion must recover
        // the ORIGINAL key, not mint a new one that cannot decrypt existing data.
@@ -1215,6 +1229,12 @@ pub struct LocalKmsBackend {
 }

 impl LocalKmsBackend {
+    /// Lifecycle driver for the shared state-machine contract tests.
+    #[cfg(test)]
+    pub(crate) fn lifecycle_client(&self) -> &LocalKmsClient {
+        &self.client
+    }
+
    /// Create a new LocalKmsBackend
    pub async fn new(config: KmsConfig) -> Result<Self> {
        config.validate()?;
@@ -1328,6 +1348,11 @@ impl KmsBackend for LocalKmsBackend {

    async fn describe_key(&self, request: DescribeKeyRequest) -> Result<DescribeKeyResponse> {
        let key_info = self.client.describe_key(&request.key_id, None).await?;
+        let deletion_date = if key_info.status == KeyStatus::PendingDeletion {
+            self.client.load_master_key(&request.key_id).await?.deletion_date
+        } else {
+            None
+        };

        let metadata = KeyMetadata {
            key_id: key_info.key_id,
@@ -1340,7 +1365,7 @@ impl KmsBackend for LocalKmsBackend {
            key_usage: key_info.usage,
            description: key_info.description,
            creation_date: key_info.created_at,
-            deletion_date: None,
+            deletion_date,
            origin: "KMS".to_string(),
            key_manager: "CUSTOMER".to_string(),
            tags: key_info.tags,
@@ -1371,7 +1396,22 @@ impl KmsBackend for LocalKmsBackend {
            .map_err(|_| KmsError::key_not_found(format!("Key {key_id} not found")))?;

        let (deletion_date_str, deletion_date_dt) = if request.force_immediate.unwrap_or(false) {
-            // For immediate deletion, actually delete the key from filesystem
+            // Tombstone first: mark the record Deleted before removing the
+            // file, so a crash between the two steps leaves a key that is
+            // already unusable and whose removal can simply be re-run.
+            match self.client.decode_stored_key(key_id).await {
+                Ok((_stored, key_material)) => {
+                    let mut tombstone = master_key.clone();
+                    tombstone.status = KeyStatus::Deleted;
+                    tombstone.deletion_date = Some(Zoned::now());
+                    self.client.save_master_key(&tombstone, &key_material).await?;
+                }
+                Err(error) => {
+                    // A record whose material can no longer be decoded cannot be
+                    // re-encrypted into a tombstone; proceed with the removal.
+                    warn!(key_id, %error, "skipping tombstone for undecodable key record");
+                }
+            }
            let key_path = self.client.master_key_path(key_id)?;
            durable_file::remove_durably(key_path)
                .await
@@ -1399,6 +1439,8 @@ impl KmsBackend for LocalKmsBackend {
            });
        } else {
            // Schedule for deletion (default 30 days)
+            ensure_key_status_permits(key_id, &master_key.status, StateGatedOperation::ScheduleDeletion)?;
+
            let days = request.pending_window_in_days.unwrap_or(30);
            if !(7..=30).contains(&days) {
                return Err(KmsError::invalid_parameter("pending_window_in_days must be between 7 and 30".to_string()));
@@ -1406,6 +1448,7 @@ impl KmsBackend for LocalKmsBackend {

            let deletion_date = Zoned::now() + Duration::from_secs(days as u64 * 86400);
            master_key.status = KeyStatus::PendingDeletion;
+            master_key.deletion_date = Some(deletion_date.clone());

            (Some(deletion_date.to_string()), Some(deletion_date))
        };
@@ -1459,6 +1502,7 @@ impl KmsBackend for LocalKmsBackend {

        // Cancel the deletion by resetting the state
        master_key.status = KeyStatus::Active;
+        master_key.deletion_date = None;

        // Save the updated key to disk - this is the missing critical step!
        // Preserve existing key material instead of generating new one
@@ -1496,13 +1540,55 @@ impl KmsBackend for LocalKmsBackend {
    fn capabilities(&self) -> BackendCapabilities {
        // Rotation stays unadvertised until historical key versions can be
        // retained (see LocalKmsClient::rotate_key); without version history
-        // there is also no versioning capability. Deletion deadlines are not
-        // yet persisted across restarts, but scheduling itself is supported.
+        // there is also no versioning capability.
        BackendCapabilities::minimal()
            .with_enable_disable(true)
            .with_schedule_deletion(true)
            .with_physical_delete(true)
    }
+
+    async fn remove_expired_key(&self, key_id: &str, now: &Zoned) -> Result<ExpiredKeyRemoval> {
+        // The per-key write lock serializes this against a concurrent
+        // cancellation, closing the check-then-remove race.
+        let _write_guard = self.client.lock_key_for_write(key_id).await;
+
+        if !fs::try_exists(self.client.master_key_path(key_id)?).await? {
+            return Ok(ExpiredKeyRemoval::Removed);
+        }
+        let master_key = self.client.load_master_key(key_id).await?;
+        match master_key.status {
+            // Tombstone left by a crashed removal: complete it.
+            KeyStatus::Deleted => {}
+            KeyStatus::PendingDeletion => {
+                match &master_key.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 (see delete_key): a crash between the state
+                // write and the file removal must leave an unusable record.
+                match self.client.decode_stored_key(key_id).await {
+                    Ok((_stored, key_material)) => {
+                        let mut tombstone = master_key.clone();
+                        tombstone.status = KeyStatus::Deleted;
+                        tombstone.deletion_date = Some(now.clone());
+                        self.client.save_master_key(&tombstone, &key_material).await?;
+                    }
+                    Err(error) => {
+                        warn!(key_id, %error, "skipping tombstone for undecodable key record");
+                    }
+                }
+            }
+            KeyStatus::Active | KeyStatus::Disabled => return Ok(ExpiredKeyRemoval::StateChanged),
+        }
+
+        durable_file::remove_durably(self.client.master_key_path(key_id)?)
+            .await
+            .map_err(|e| KmsError::internal_error(format!("Failed to delete key file: {e}")))?;
+        debug!(key_id, "Local KMS expired key removed");
+        Ok(ExpiredKeyRemoval::Removed)
+    }
 }

 #[cfg(test)]
@@ -2617,9 +2703,16 @@ mod tests {
            client.schedule_key_deletion(key_id, 7, None),
            client.enable_key(key_id, None),
        );
-        disable.expect("disable");
-        schedule.expect("schedule deletion");
-        enable.expect("enable");
+        // The per-key lock serializes the three transitions in an arbitrary
+        // order, and the state gate may legitimately reject a transition that
+        // lost the race (e.g. enable after deletion was scheduled). Any other
+        // error kind would still mean corrupted storage.
+        for result in [disable, schedule, enable] {
+            match result {
+                Ok(()) | Err(KmsError::InvalidOperation { .. }) => {}
+                Err(other) => panic!("concurrent transition must only fail with a state rejection, got {other:?}"),
+            }
+        }

        // Whatever the serialization order, the file must be one writer's
        // complete output with the original material intact.
@@ -2634,4 +2727,67 @@ mod tests {
            "concurrent status updates must never lose or regenerate key material"
        );
    }
+
+    /// Records written before deadline persistence landed have no
+    /// deletion_date field and must keep deserializing (as None).
+    #[tokio::test]
+    async fn stored_master_key_without_deletion_date_still_deserializes() {
+        let (client, _temp_dir) = create_test_client().await;
+        client.create_key("legacy-key", "AES_256", None).await.expect("create key");
+
+        let path = client.master_key_path("legacy-key").expect("key path");
+        let bytes = fs::read(&path).await.expect("read stored key");
+        let mut value: serde_json::Value = serde_json::from_slice(&bytes).expect("stored key must be JSON");
+        value
+            .as_object_mut()
+            .expect("stored key must be a JSON object")
+            .remove("deletion_date")
+            .expect("current records must carry the field");
+
+        let stored: StoredMasterKey = serde_json::from_value(value).expect("legacy record must deserialize");
+        assert!(stored.deletion_date.is_none());
+    }
+
+    #[tokio::test]
+    async fn remove_expired_key_completes_a_tombstone_and_stays_idempotent() {
+        let temp_dir = TempDir::new().expect("temp dir");
+        let config = KmsConfig::local(temp_dir.path().to_path_buf()).with_insecure_development_defaults();
+        let backend = LocalKmsBackend::new(config).await.expect("backend");
+        let created = backend
+            .create_key(CreateKeyRequest {
+                key_name: Some("tombstoned-key".to_string()),
+                key_usage: KeyUsage::EncryptDecrypt,
+                ..Default::default()
+            })
+            .await
+            .expect("create key");
+        let key_id = created.key_id;
+
+        // Craft the state a removal crashed in: tombstone written, file not
+        // yet removed.
+        let client = backend.lifecycle_client();
+        let (_stored, key_material) = client.decode_stored_key(&key_id).await.expect("decode stored key");
+        let mut tombstone = client.load_master_key(&key_id).await.expect("load key");
+        tombstone.status = KeyStatus::Deleted;
+        tombstone.deletion_date = Some(Zoned::now());
+        client
+            .save_master_key(&tombstone, &key_material)
+            .await
+            .expect("write tombstone");
+
+        // The sweep primitive completes the crashed removal...
+        let outcome = backend
+            .remove_expired_key(&key_id, &Zoned::now())
+            .await
+            .expect("tombstone completion");
+        assert_eq!(outcome, crate::backends::ExpiredKeyRemoval::Removed);
+        assert!(!client.master_key_path(&key_id).expect("key path").exists());
+
+        // ...and stays idempotent once the key is gone.
+        let outcome = backend
+            .remove_expired_key(&key_id, &Zoned::now())
+            .await
+            .expect("repeat removal");
+        assert_eq!(outcome, crate::backends::ExpiredKeyRemoval::Removed);
+    }
 }