fix(kms): key usage/state/spec validation + real ECDH DeriveSharedSecret (bug-hunt) (#2280)

Author: vieiralucas

crates/fakecloud-e2e/tests/kms.rs

@@ -390,12 +390,31 @@ async fn kms_derive_shared_secret() {
         .unwrap();
     let key_id = resp.key_metadata().unwrap().key_id().to_string();
 
-    let fake_pub = Blob::new(vec![0x04; 65]); // Fake uncompressed EC point
+    // Real ECDH requires a valid SubjectPublicKeyInfo for the counterparty;
+    // use a second KEY_AGREEMENT key's public key.
+    let peer = client
+        .create_key()
+        .key_usage(aws_sdk_kms::types::KeyUsageType::KeyAgreement)
+        .key_spec(aws_sdk_kms::types::KeySpec::EccNistP256)
+        .send()
+        .await
+        .unwrap();
+    let peer_id = peer.key_metadata().unwrap().key_id().to_string();
+    let peer_pub = client
+        .get_public_key()
+        .key_id(&peer_id)
+        .send()
+        .await
+        .unwrap()
+        .public_key()
+        .unwrap()
+        .clone();
+
     let result = client
         .derive_shared_secret()
         .key_id(&key_id)
         .key_agreement_algorithm(aws_sdk_kms::types::KeyAgreementAlgorithmSpec::Ecdh)
-        .public_key(fake_pub)
+        .public_key(peer_pub)
         .send()
         .await
         .unwrap();
@@ -1151,3 +1170,124 @@ async fn kms_generate_data_key_binds_encryption_context() {
         "decrypt with a different encryption context must fail"
     );
 }
+
+#[tokio::test]
+async fn kms_generate_data_key_rejects_non_encrypt_key() {
+    // Bug-hunt 1.18: GenerateDataKey must reject a key whose usage isn't
+    // ENCRYPT_DECRYPT (InvalidKeyUsageException), not silently succeed.
+    let server = helpers::TestServer::start().await;
+    let client = server.kms_client().await;
+
+    let key = client
+        .create_key()
+        .key_usage(aws_sdk_kms::types::KeyUsageType::SignVerify)
+        .key_spec(aws_sdk_kms::types::KeySpec::Rsa2048)
+        .send()
+        .await
+        .unwrap();
+    let key_id = key.key_metadata().unwrap().key_id().to_string();
+
+    let result = client
+        .generate_data_key()
+        .key_id(&key_id)
+        .key_spec(aws_sdk_kms::types::DataKeySpec::Aes256)
+        .send()
+        .await;
+    assert!(
+        result.is_err(),
+        "GenerateDataKey on a SIGN_VERIFY key must fail"
+    );
+}
+
+#[tokio::test]
+async fn kms_create_key_rejects_incompatible_spec_usage() {
+    // Bug-hunt 1.20: an HMAC spec with ENCRYPT_DECRYPT usage is unusable and
+    // must be rejected at CreateKey.
+    let server = helpers::TestServer::start().await;
+    let client = server.kms_client().await;
+
+    let result = client
+        .create_key()
+        .key_spec(aws_sdk_kms::types::KeySpec::Hmac256)
+        .key_usage(aws_sdk_kms::types::KeyUsageType::EncryptDecrypt)
+        .send()
+        .await;
+    assert!(
+        result.is_err(),
+        "HMAC + ENCRYPT_DECRYPT is incompatible and must be rejected"
+    );
+}
+
+#[tokio::test]
+async fn kms_enable_key_rejects_pending_deletion() {
+    // Bug-hunt 1.19: EnableKey must not resurrect a key scheduled for deletion.
+    let server = helpers::TestServer::start().await;
+    let client = server.kms_client().await;
+
+    let key = client.create_key().send().await.unwrap();
+    let key_id = key.key_metadata().unwrap().key_id().to_string();
+
+    client
+        .schedule_key_deletion()
+        .key_id(&key_id)
+        .pending_window_in_days(7)
+        .send()
+        .await
+        .unwrap();
+
+    let result = client.enable_key().key_id(&key_id).send().await;
+    assert!(
+        result.is_err(),
+        "EnableKey on a PendingDeletion key must fail (KMSInvalidStateException)"
+    );
+}
+
+#[tokio::test]
+async fn kms_derive_shared_secret_converges() {
+    // Bug-hunt 1.21: two parties running DeriveSharedSecret with each other's
+    // public key must derive the SAME secret (real ECDH), which the prior
+    // SHA-256(seed || peer) construction did not.
+    let server = helpers::TestServer::start().await;
+    let client = server.kms_client().await;
+
+    async fn make_agreement_key(
+        c: &aws_sdk_kms::Client,
+    ) -> (String, aws_sdk_kms::primitives::Blob) {
+        let k = c
+            .create_key()
+            .key_spec(aws_sdk_kms::types::KeySpec::EccNistP256)
+            .key_usage(aws_sdk_kms::types::KeyUsageType::KeyAgreement)
+            .send()
+            .await
+            .unwrap();
+        let id = k.key_metadata().unwrap().key_id().to_string();
+        let pk = c.get_public_key().key_id(&id).send().await.unwrap();
+        let spki = pk.public_key().unwrap().clone();
+        (id, spki)
+    }
+    let (id_a, pub_a) = make_agreement_key(&client).await;
+    let (id_b, pub_b) = make_agreement_key(&client).await;
+
+    let secret_a = client
+        .derive_shared_secret()
+        .key_id(&id_a)
+        .key_agreement_algorithm(aws_sdk_kms::types::KeyAgreementAlgorithmSpec::Ecdh)
+        .public_key(pub_b)
+        .send()
+        .await
+        .unwrap();
+    let secret_b = client
+        .derive_shared_secret()
+        .key_id(&id_b)
+        .key_agreement_algorithm(aws_sdk_kms::types::KeyAgreementAlgorithmSpec::Ecdh)
+        .public_key(pub_a)
+        .send()
+        .await
+        .unwrap();
+
+    assert_eq!(
+        secret_a.shared_secret().unwrap().as_ref(),
+        secret_b.shared_secret().unwrap().as_ref(),
+        "both parties must derive identical shared secrets"
+    );
+}

crates/fakecloud-kms/Cargo.toml

@@ -29,7 +29,7 @@ aes-gcm = "0.10"
 rsa = { workspace = true }
 rand = "0.8"
 signature = "2"
-p256 = { version = "0.13", features = ["ecdsa", "pkcs8"] }
-p384 = { version = "0.13", features = ["ecdsa", "pkcs8"] }
+p256 = { version = "0.13", features = ["ecdsa", "ecdh", "pkcs8"] }
+p384 = { version = "0.13", features = ["ecdsa", "ecdh", "pkcs8"] }
 k256 = { version = "0.13", features = ["ecdsa", "pkcs8"] }
 p521 = { version = "0.13", features = ["ecdsa", "pkcs8"] }

crates/fakecloud-kms/src/helpers.rs

@@ -540,6 +540,78 @@ pub(crate) fn require_usable_key_state(key: &KmsKey) -> Result<(), AwsServiceErr
     ))
 }
 
+/// The KeyUsage values compatible with a given KeySpec, per AWS KMS. An
+/// empty slice means the spec is unknown (leave validation to later paths).
+fn allowed_usages_for_key_spec(key_spec: &str) -> &'static [&'static str] {
+    match key_spec {
+        "SYMMETRIC_DEFAULT" => &["ENCRYPT_DECRYPT"],
+        s if s.starts_with("HMAC_") => &["GENERATE_VERIFY_MAC"],
+        s if s.starts_with("RSA_") => &["ENCRYPT_DECRYPT", "SIGN_VERIFY"],
+        // NIST curves support signing and key agreement.
+        "ECC_NIST_P256" | "ECC_NIST_P384" | "ECC_NIST_P521" => &["SIGN_VERIFY", "KEY_AGREEMENT"],
+        // secp256k1 supports signing only.
+        "ECC_SECG_P256K1" => &["SIGN_VERIFY"],
+        // SM2 (China regions) supports all three.
+        "SM2" => &["SIGN_VERIFY", "ENCRYPT_DECRYPT", "KEY_AGREEMENT"],
+        _ => &[],
+    }
+}
+
+/// CreateKey must reject a KeyUsage that is incompatible with the KeySpec
+/// (e.g. HMAC_256 + ENCRYPT_DECRYPT), which would otherwise create a key that
+/// no crypto operation can use. AWS returns ValidationException.
+pub(crate) fn validate_key_spec_usage(
+    key_spec: &str,
+    key_usage: &str,
+) -> Result<(), AwsServiceError> {
+    let allowed = allowed_usages_for_key_spec(key_spec);
+    if allowed.is_empty() || allowed.contains(&key_usage) {
+        return Ok(());
+    }
+    Err(AwsServiceError::aws_error(
+        StatusCode::BAD_REQUEST,
+        "ValidationException",
+        format!("KeyUsage {key_usage} is not compatible with KeySpec {key_spec}"),
+    ))
+}
+
+/// EnableKey / DisableKey may only transition a key that is currently
+/// `Enabled` or `Disabled`. Any other lifecycle state
+/// (`PendingDeletion`, `PendingImport`, `Unavailable`, ...) is a
+/// `KMSInvalidStateException`; in particular a scheduled-for-deletion key must
+/// not be resurrected without CancelKeyDeletion.
+pub(crate) fn require_enable_disable_transition(key: &KmsKey) -> Result<(), AwsServiceError> {
+    if key.key_state == "Enabled" || key.key_state == "Disabled" {
+        return Ok(());
+    }
+    Err(AwsServiceError::aws_error(
+        StatusCode::BAD_REQUEST,
+        "KMSInvalidStateException",
+        format!(
+            "Key '{}' is not in a state that allows this operation (current state: {})",
+            key.arn, key.key_state
+        ),
+    ))
+}
+
+/// GenerateDataKey / GenerateDataKeyPair (and their `WithoutPlaintext`
+/// variants) require an `ENCRYPT_DECRYPT` key. A SIGN_VERIFY /
+/// GENERATE_VERIFY_MAC / KEY_AGREEMENT key must be rejected with
+/// `InvalidKeyUsageException`, matching the Encrypt/ReEncrypt paths.
+pub(crate) fn require_key_usage_encrypt_decrypt(key: &KmsKey) -> Result<(), AwsServiceError> {
+    if key.key_usage != "ENCRYPT_DECRYPT" {
+        return Err(AwsServiceError::aws_error(
+            StatusCode::BAD_REQUEST,
+            "InvalidKeyUsageException",
+            format!(
+                "The operation failed because the KMS key {} is not enabled for the requested operation. The key usage must be ENCRYPT_DECRYPT but is {}.",
+                key.arn, key.key_usage
+            ),
+        ));
+    }
+    Ok(())
+}
+
 pub(crate) fn validate_key_usage_signing(
     key: &KmsKey,
     resolved: &str,
@@ -923,3 +995,28 @@ pub(crate) fn action_matches(policy_action: &str, requested_action: &str) -> boo
     }
     false
 }
+
+#[cfg(test)]
+mod key_spec_usage_tests {
+    use super::validate_key_spec_usage;
+
+    #[test]
+    fn compatible_combos_pass() {
+        assert!(validate_key_spec_usage("SYMMETRIC_DEFAULT", "ENCRYPT_DECRYPT").is_ok());
+        assert!(validate_key_spec_usage("HMAC_256", "GENERATE_VERIFY_MAC").is_ok());
+        assert!(validate_key_spec_usage("RSA_2048", "SIGN_VERIFY").is_ok());
+        assert!(validate_key_spec_usage("RSA_2048", "ENCRYPT_DECRYPT").is_ok());
+        assert!(validate_key_spec_usage("ECC_NIST_P256", "KEY_AGREEMENT").is_ok());
+        assert!(validate_key_spec_usage("ECC_NIST_P256", "SIGN_VERIFY").is_ok());
+    }
+
+    #[test]
+    fn incompatible_combos_rejected() {
+        // HMAC key can't encrypt/decrypt.
+        assert!(validate_key_spec_usage("HMAC_256", "ENCRYPT_DECRYPT").is_err());
+        // Symmetric key can't sign.
+        assert!(validate_key_spec_usage("SYMMETRIC_DEFAULT", "SIGN_VERIFY").is_err());
+        // secp256k1 supports signing only.
+        assert!(validate_key_spec_usage("ECC_SECG_P256K1", "KEY_AGREEMENT").is_err());
+    }
+}

crates/fakecloud-kms/src/service.rs

@@ -546,6 +546,9 @@ impl KmsService {
 
     fn create_key(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
         let input = CreateKeyInput::from_body(&req.json_body())?;
+        // Reject incompatible KeySpec/KeyUsage combinations (e.g. an HMAC spec
+        // with ENCRYPT_DECRYPT) before minting a key no operation can use.
+        validate_key_spec_usage(&input.key_spec, &input.key_usage)?;
 
         let mut accounts = self.state.write();
         let state = accounts.get_or_create(&req.account_id);
@@ -825,6 +828,11 @@ impl KmsService {
                 "Key state became inconsistent",
             )
         })?;
+        // EnableKey is only valid from Enabled/Disabled. A key that is
+        // PendingDeletion / PendingImport / Unavailable must not be silently
+        // resurrected to Enabled — AWS returns KMSInvalidStateException (you
+        // must CancelKeyDeletion / import material first).
+        require_enable_disable_transition(key)?;
         key.enabled = true;
         key.key_state = "Enabled".to_string();
 
@@ -844,6 +852,7 @@ impl KmsService {
                 "Key state became inconsistent",
             )
         })?;
+        require_enable_disable_transition(key)?;
         key.enabled = false;
         key.key_state = "Disabled".to_string();
 
@@ -1209,6 +1218,9 @@ impl KmsService {
                 "Key state became inconsistent",
             )
         })?;
+        // Rotation cannot be toggled on a key pending deletion / import;
+        // AWS rejects with KMSInvalidStateException (Disabled -> DisabledException).
+        require_usable_key_state(key)?;
         // RotationPeriodInDays is optional; AWS validates 90..=2560 and
         // defaults to 365. Persist it so GetKeyRotationStatus echoes it back
         // instead of dropping it (bug-audit 2026-06-20, 1.24).
@@ -1251,6 +1263,7 @@ impl KmsService {
                 "Key state became inconsistent",
             )
         })?;
+        require_usable_key_state(key)?;
         key.key_rotation_enabled = false;
 
         Ok(AwsResponse::json(StatusCode::OK, "{}"))
@@ -1269,6 +1282,7 @@ impl KmsService {
                 "Key state became inconsistent",
             )
         })?;
+        require_usable_key_state(key)?;
 
         let rotation = KeyRotation {
             key_id: key.key_id.clone(),