Skip to main content

fakecloud_kms/
service.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use base64::Engine;
6use chrono::Utc;
7use http::StatusCode;
8use serde_json::{json, Value};
9use tokio::sync::Mutex as AsyncMutex;
10use uuid::Uuid;
11
12use fakecloud_aws::arn::Arn;
13use fakecloud_core::service::{AwsRequest, AwsResponse, AwsService, AwsServiceError};
14use fakecloud_core::validation::*;
15use fakecloud_persistence::SnapshotStore;
16
17use crate::state::{
18    CustomKeyStore, KeyRotation, KmsAlias, KmsGrant, KmsKey, KmsSnapshot, KmsState, SharedKmsState,
19    KMS_SNAPSHOT_SCHEMA_VERSION,
20};
21
22const FAKE_ENVELOPE_PREFIX: &str = "fakecloud-kms:";
23const IMPORTED_ENVELOPE_PREFIX: &str = "fakecloud-imported:";
24const RSA_ENVELOPE_PREFIX: &str = "fakecloud-rsa:";
25
26/// Result of decoding a FakeCloud KMS ciphertext blob. We carry the
27/// plaintext as base64 so the two callers that care (`Decrypt` returns
28/// it to the client, `ReEncrypt` re-wraps it with a new key) can both
29/// hand it straight to the response builder without an extra encode.
30pub(crate) struct DecodedCiphertext {
31    source_arn: String,
32    plaintext_b64: String,
33    /// EncryptionAlgorithm to echo on the Decrypt response. `SYMMETRIC_DEFAULT`
34    /// for the symmetric AES envelope; the `RSAES_OAEP_*` value for an
35    /// asymmetric RSA ciphertext.
36    encryption_algorithm: String,
37}
38
39const VALID_KEY_SPECS: &[&str] = &[
40    "ECC_NIST_P256",
41    "ECC_NIST_P384",
42    "ECC_NIST_P521",
43    "ECC_SECG_P256K1",
44    "HMAC_224",
45    "HMAC_256",
46    "HMAC_384",
47    "HMAC_512",
48    "RSA_2048",
49    "RSA_3072",
50    "RSA_4096",
51    "SM2",
52    "SYMMETRIC_DEFAULT",
53];
54
55const VALID_SIGNING_ALGORITHMS: &[&str] = &[
56    "RSASSA_PKCS1_V1_5_SHA_256",
57    "RSASSA_PKCS1_V1_5_SHA_384",
58    "RSASSA_PKCS1_V1_5_SHA_512",
59    "RSASSA_PSS_SHA_256",
60    "RSASSA_PSS_SHA_384",
61    "RSASSA_PSS_SHA_512",
62    "ECDSA_SHA_256",
63    "ECDSA_SHA_384",
64    "ECDSA_SHA_512",
65];
66
67/// Single source of truth for supported KMS actions. Referenced by both
68/// `supported_actions()` (used by the dispatch layer) and
69/// `iam_action_for()` (used by the IAM enforcement layer).
70static KMS_ACTIONS: &[&str] = &[
71    "CreateKey",
72    "DescribeKey",
73    "GetKeyLastUsage",
74    "ListKeys",
75    "EnableKey",
76    "DisableKey",
77    "ScheduleKeyDeletion",
78    "CancelKeyDeletion",
79    "Encrypt",
80    "Decrypt",
81    "ReEncrypt",
82    "GenerateDataKey",
83    "GenerateDataKeyWithoutPlaintext",
84    "GenerateRandom",
85    "CreateAlias",
86    "DeleteAlias",
87    "UpdateAlias",
88    "ListAliases",
89    "TagResource",
90    "UntagResource",
91    "ListResourceTags",
92    "UpdateKeyDescription",
93    "GetKeyPolicy",
94    "PutKeyPolicy",
95    "ListKeyPolicies",
96    "GetKeyRotationStatus",
97    "EnableKeyRotation",
98    "DisableKeyRotation",
99    "RotateKeyOnDemand",
100    "ListKeyRotations",
101    "Sign",
102    "Verify",
103    "GetPublicKey",
104    "CreateGrant",
105    "ListGrants",
106    "ListRetirableGrants",
107    "RevokeGrant",
108    "RetireGrant",
109    "GenerateMac",
110    "VerifyMac",
111    "ReplicateKey",
112    "GenerateDataKeyPair",
113    "GenerateDataKeyPairWithoutPlaintext",
114    "DeriveSharedSecret",
115    "GetParametersForImport",
116    "ImportKeyMaterial",
117    "DeleteImportedKeyMaterial",
118    "UpdatePrimaryRegion",
119    "CreateCustomKeyStore",
120    "DeleteCustomKeyStore",
121    "DescribeCustomKeyStores",
122    "ConnectCustomKeyStore",
123    "DisconnectCustomKeyStore",
124    "UpdateCustomKeyStore",
125];
126
127pub struct KmsService {
128    state: SharedKmsState,
129    snapshot_store: Option<Arc<dyn SnapshotStore>>,
130    snapshot_lock: Arc<AsyncMutex<()>>,
131}
132
133impl KmsService {
134    pub fn new(state: SharedKmsState) -> Self {
135        // Warm the RSA keypair cache in the background. Generation is
136        // CPU-bound (~20s for RSA-4096) and reused across every CreateKey
137        // with that bit width; doing it here off a blocking thread keeps
138        // the first concurrent CreateKey from racing N tokio workers into
139        // a 30s read-timeout cliff (see `asym::generate_keypair`).
140        if tokio::runtime::Handle::try_current().is_ok() {
141            tokio::task::spawn_blocking(|| {
142                let _ = self::asym::generate_keypair("RSA_2048");
143                let _ = self::asym::generate_keypair("RSA_3072");
144                let _ = self::asym::generate_keypair("RSA_4096");
145            });
146        }
147        Self {
148            state,
149            snapshot_store: None,
150            snapshot_lock: Arc::new(AsyncMutex::new(())),
151        }
152    }
153
154    pub fn with_snapshot_store(mut self, store: Arc<dyn SnapshotStore>) -> Self {
155        self.snapshot_store = Some(store);
156        self
157    }
158
159    /// Persist current state as a snapshot. Held across the
160    /// clone-serialize-write sequence to prevent stale-last writes,
161    /// with serde + file I/O offloaded to the blocking pool.
162    async fn save_snapshot(&self) {
163        save_kms_snapshot(
164            &self.state,
165            self.snapshot_store.clone(),
166            &self.snapshot_lock,
167        )
168        .await;
169    }
170
171    /// Build a hook that persists the current KMS state when invoked, or
172    /// `None` in memory mode (no snapshot store). The CloudFormation provisioner
173    /// mutates `state` directly and uses this to write a CFN-provisioned
174    /// resource through to disk, the same way a direct mutating API call would.
175    pub fn snapshot_hook(&self) -> Option<fakecloud_persistence::SnapshotHook> {
176        let store = self.snapshot_store.clone()?;
177        let state = self.state.clone();
178        let lock = self.snapshot_lock.clone();
179        Some(Arc::new(move || {
180            let state = state.clone();
181            let store = store.clone();
182            let lock = lock.clone();
183            Box::pin(async move {
184                save_kms_snapshot(&state, Some(store), &lock).await;
185            })
186        }))
187    }
188}
189
190/// Persist the current KMS state as a snapshot. Offloads the serde +
191/// blocking file write to the Tokio blocking pool. Noop when `store` is `None`
192/// (memory mode). Shared by `KmsService::save_snapshot` and the CloudFormation
193/// provisioner's post-provision persist hook so both route through the same
194/// serialize-and-write path.
195pub async fn save_kms_snapshot(
196    state: &SharedKmsState,
197    store: Option<Arc<dyn SnapshotStore>>,
198    lock: &AsyncMutex<()>,
199) {
200    let Some(store) = store else {
201        return;
202    };
203    let _guard = lock.lock().await;
204    let snapshot = KmsSnapshot {
205        schema_version: KMS_SNAPSHOT_SCHEMA_VERSION,
206        state: None,
207        accounts: Some(state.read().clone()),
208    };
209    let join = tokio::task::spawn_blocking(move || -> std::io::Result<()> {
210        let bytes = serde_json::to_vec(&snapshot)
211            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?;
212        store.save(&bytes)
213    })
214    .await;
215    match join {
216        Ok(Ok(())) => {}
217        Ok(Err(err)) => tracing::error!(%err, "failed to write kms snapshot"),
218        Err(err) => tracing::error!(%err, "kms snapshot task panicked"),
219    }
220}
221
222#[async_trait]
223impl AwsService for KmsService {
224    fn service_name(&self) -> &str {
225        "kms"
226    }
227
228    async fn handle(&self, req: AwsRequest) -> Result<AwsResponse, AwsServiceError> {
229        let mutates = is_mutating_action(req.action.as_str());
230        let result = match req.action.as_str() {
231            "CreateKey" => self.create_key(&req),
232            "DescribeKey" => self.describe_key(&req),
233            "GetKeyLastUsage" => self.get_key_last_usage(&req),
234            "ListKeys" => self.list_keys(&req),
235            "EnableKey" => self.enable_key(&req),
236            "DisableKey" => self.disable_key(&req),
237            "ScheduleKeyDeletion" => self.schedule_key_deletion(&req),
238            "CancelKeyDeletion" => self.cancel_key_deletion(&req),
239            "Encrypt" => self.encrypt(&req),
240            "Decrypt" => self.decrypt(&req),
241            "ReEncrypt" => self.re_encrypt(&req),
242            "GenerateDataKey" => self.generate_data_key(&req),
243            "GenerateDataKeyWithoutPlaintext" => self.generate_data_key_without_plaintext(&req),
244            "GenerateRandom" => self.generate_random(&req),
245            "CreateAlias" => self.create_alias(&req),
246            "DeleteAlias" => self.delete_alias(&req),
247            "UpdateAlias" => self.update_alias(&req),
248            "ListAliases" => self.list_aliases(&req),
249            "TagResource" => self.tag_resource(&req),
250            "UntagResource" => self.untag_resource(&req),
251            "ListResourceTags" => self.list_resource_tags(&req),
252            "UpdateKeyDescription" => self.update_key_description(&req),
253            "GetKeyPolicy" => self.get_key_policy(&req),
254            "PutKeyPolicy" => self.put_key_policy(&req),
255            "ListKeyPolicies" => self.list_key_policies(&req),
256            "GetKeyRotationStatus" => self.get_key_rotation_status(&req),
257            "EnableKeyRotation" => self.enable_key_rotation(&req),
258            "DisableKeyRotation" => self.disable_key_rotation(&req),
259            "RotateKeyOnDemand" => self.rotate_key_on_demand(&req),
260            "ListKeyRotations" => self.list_key_rotations(&req),
261            "Sign" => self.sign(&req),
262            "Verify" => self.verify(&req),
263            "GetPublicKey" => self.get_public_key(&req),
264            "CreateGrant" => self.create_grant(&req),
265            "ListGrants" => self.list_grants(&req),
266            "ListRetirableGrants" => self.list_retirable_grants(&req),
267            "RevokeGrant" => self.revoke_grant(&req),
268            "RetireGrant" => self.retire_grant(&req),
269            "GenerateMac" => self.generate_mac(&req),
270            "VerifyMac" => self.verify_mac(&req),
271            "ReplicateKey" => self.replicate_key(&req),
272            "GenerateDataKeyPair" => self.generate_data_key_pair(&req),
273            "GenerateDataKeyPairWithoutPlaintext" => {
274                self.generate_data_key_pair_without_plaintext(&req)
275            }
276            "DeriveSharedSecret" => self.derive_shared_secret(&req),
277            "GetParametersForImport" => self.get_parameters_for_import(&req),
278            "ImportKeyMaterial" => self.import_key_material(&req),
279            "DeleteImportedKeyMaterial" => self.delete_imported_key_material(&req),
280            "UpdatePrimaryRegion" => self.update_primary_region(&req),
281            "CreateCustomKeyStore" => self.create_custom_key_store(&req),
282            "DeleteCustomKeyStore" => self.delete_custom_key_store(&req),
283            "DescribeCustomKeyStores" => self.describe_custom_key_stores(&req),
284            "ConnectCustomKeyStore" => self.connect_custom_key_store(&req),
285            "DisconnectCustomKeyStore" => self.disconnect_custom_key_store(&req),
286            "UpdateCustomKeyStore" => self.update_custom_key_store(&req),
287            _ => Err(AwsServiceError::action_not_implemented("kms", &req.action)),
288        };
289        if mutates && matches!(result.as_ref(), Ok(resp) if resp.status.is_success()) {
290            self.save_snapshot().await;
291        }
292        result
293    }
294
295    fn supported_actions(&self) -> &[&str] {
296        KMS_ACTIONS
297    }
298
299    fn iam_enforceable(&self) -> bool {
300        true
301    }
302
303    fn iam_action_for(&self, request: &AwsRequest) -> Option<fakecloud_core::auth::IamAction> {
304        let action = KMS_ACTIONS.iter().copied().find(|a| *a == request.action)?;
305        let resource = kms_resource_for(action, &self.state, request);
306        Some(fakecloud_core::auth::IamAction {
307            service: "kms",
308            action,
309            resource,
310        })
311    }
312
313    fn resource_tags_for(
314        &self,
315        resource_arn: &str,
316    ) -> Option<std::collections::HashMap<String, String>> {
317        if resource_arn == "*" {
318            return Some(std::collections::HashMap::new());
319        }
320        let key_id = resource_arn.rsplit_once(":key/")?.1;
321        let account_id = resource_arn.split(':').nth(4).unwrap_or("").to_string();
322        let accounts = self.state.read();
323        let state = accounts.get(&account_id)?;
324        let key = state.keys.get(key_id)?;
325        Some(
326            key.tags
327                .iter()
328                .map(|(k, v)| (k.clone(), v.clone()))
329                .collect(),
330        )
331    }
332
333    fn request_tags_from(
334        &self,
335        request: &AwsRequest,
336        action: &str,
337    ) -> Option<std::collections::HashMap<String, String>> {
338        match action {
339            "CreateKey" | "TagResource" => {
340                let body = request.json_body();
341                let mut tags = std::collections::HashMap::new();
342                if let Some(arr) = body["Tags"].as_array() {
343                    for tag in arr {
344                        if let (Some(k), Some(v)) =
345                            (tag["TagKey"].as_str(), tag["TagValue"].as_str())
346                        {
347                            tags.insert(k.to_string(), v.to_string());
348                        }
349                    }
350                }
351                Some(tags)
352            }
353            _ => Some(std::collections::HashMap::new()),
354        }
355    }
356}
357
358/// Parsed + validated inputs for `CreateKey`.
359struct CreateKeyInput {
360    custom_key_store_id: Option<String>,
361    description: String,
362    key_usage: String,
363    key_spec: String,
364    origin: String,
365    multi_region: bool,
366    policy: Option<String>,
367    tags: BTreeMap<String, String>,
368}
369
370impl CreateKeyInput {
371    fn from_body(body: &Value) -> Result<Self, AwsServiceError> {
372        // CreateKey's Smithy contract doesn't declare ValidationException, so
373        // map each constraint failure onto the closest declared error.
374        recoded("CustomKeyStoreInvalidStateException", || {
375            validate_optional_string_length(
376                "customKeyStoreId",
377                body["CustomKeyStoreId"].as_str(),
378                1,
379                64,
380            )
381        })?;
382        recoded("UnsupportedOperationException", || {
383            validate_optional_string_length("description", body["Description"].as_str(), 0, 8192)
384        })?;
385        recoded("UnsupportedOperationException", || {
386            validate_optional_enum(
387                "keyUsage",
388                body["KeyUsage"].as_str(),
389                &[
390                    "SIGN_VERIFY",
391                    "ENCRYPT_DECRYPT",
392                    "GENERATE_VERIFY_MAC",
393                    "KEY_AGREEMENT",
394                ],
395            )
396        })?;
397        recoded("UnsupportedOperationException", || {
398            validate_optional_enum(
399                "origin",
400                body["Origin"].as_str(),
401                &["AWS_KMS", "EXTERNAL", "AWS_CLOUDHSM", "EXTERNAL_KEY_STORE"],
402            )
403        })?;
404        recoded("MalformedPolicyDocumentException", || {
405            validate_optional_string_length("policy", body["Policy"].as_str(), 1, 131072)
406        })?;
407        recoded("XksKeyInvalidConfigurationException", || {
408            validate_optional_string_length("xksKeyId", body["XksKeyId"].as_str(), 1, 64)
409        })?;
410
411        let key_spec = body["KeySpec"]
412            .as_str()
413            .or_else(|| body["CustomerMasterKeySpec"].as_str())
414            .unwrap_or("SYMMETRIC_DEFAULT")
415            .to_string();
416        if !VALID_KEY_SPECS.contains(&key_spec.as_str()) {
417            return Err(AwsServiceError::aws_error(
418                StatusCode::BAD_REQUEST,
419                "UnsupportedOperationException",
420                format!(
421                    "1 validation error detected: Value '{key_spec}' at 'KeySpec' failed to satisfy constraint: Member must satisfy enum value set: {}",
422                    fmt_enum_set(&VALID_KEY_SPECS.iter().map(|s| s.to_string()).collect::<Vec<_>>())
423                ),
424            ));
425        }
426
427        let tags: BTreeMap<String, String> = body["Tags"]
428            .as_array()
429            .map(|arr| {
430                arr.iter()
431                    .filter_map(|t| {
432                        let k = t["TagKey"].as_str()?;
433                        let v = t["TagValue"].as_str()?;
434                        Some((k.to_string(), v.to_string()))
435                    })
436                    .collect()
437            })
438            .unwrap_or_default();
439
440        Ok(Self {
441            custom_key_store_id: body["CustomKeyStoreId"].as_str().map(|s| s.to_string()),
442            description: body["Description"].as_str().unwrap_or("").to_string(),
443            key_usage: body["KeyUsage"]
444                .as_str()
445                .unwrap_or("ENCRYPT_DECRYPT")
446                .to_string(),
447            key_spec,
448            origin: body["Origin"].as_str().unwrap_or("AWS_KMS").to_string(),
449            multi_region: body["MultiRegion"].as_bool().unwrap_or(false),
450            policy: body["Policy"].as_str().map(|s| s.to_string()),
451            tags,
452        })
453    }
454}
455
456impl KmsService {
457    fn resolve_key_id_for(
458        &self,
459        account_id: &str,
460        region: &str,
461        key_id_or_arn: &str,
462    ) -> Option<String> {
463        let accounts = self.state.read();
464        let empty = KmsState::new(account_id, region);
465        let state = accounts.get(account_id).unwrap_or(&empty);
466        Self::resolve_key_id_with_state(state, key_id_or_arn)
467    }
468
469    pub(crate) fn resolve_key_id_with_state(
470        state: &crate::state::KmsState,
471        key_id_or_arn: &str,
472    ) -> Option<String> {
473        // Direct key ID
474        if state.keys.contains_key(key_id_or_arn) {
475            return Some(key_id_or_arn.to_string());
476        }
477
478        // ARN for key
479        if key_id_or_arn.starts_with("arn:aws:kms:") {
480            // Could be key ARN or alias ARN
481            if key_id_or_arn.contains(":key/") {
482                if let Some(id) = key_id_or_arn.rsplit('/').next() {
483                    if state.keys.contains_key(id) {
484                        return Some(id.to_string());
485                    }
486                }
487            }
488            // alias ARN: arn:aws:kms:region:account:alias/name
489            if key_id_or_arn.contains(":alias/") {
490                if let Some(alias_part) = key_id_or_arn.split(':').next_back() {
491                    if let Some(alias) = state.aliases.get(alias_part) {
492                        return Some(alias.target_key_id.clone());
493                    }
494                }
495            }
496        }
497
498        // Alias name
499        if key_id_or_arn.starts_with("alias/") {
500            if let Some(alias) = state.aliases.get(key_id_or_arn) {
501                return Some(alias.target_key_id.clone());
502            }
503        }
504
505        None
506    }
507
508    fn require_key_id(body: &Value) -> Result<String, AwsServiceError> {
509        body["KeyId"]
510            .as_str()
511            .map(|s| s.to_string())
512            .ok_or_else(|| {
513                AwsServiceError::aws_error(
514                    StatusCode::BAD_REQUEST,
515                    "ValidationException",
516                    "KeyId is required",
517                )
518            })
519    }
520
521    fn resolve_required_key(
522        &self,
523        req: &AwsRequest,
524        body: &Value,
525    ) -> Result<String, AwsServiceError> {
526        let key_id_input = Self::require_key_id(body)?;
527        self.resolve_key_id_for(&req.account_id, &req.region, &key_id_input)
528            .ok_or_else(|| {
529                AwsServiceError::aws_error(
530                    StatusCode::BAD_REQUEST,
531                    "NotFoundException",
532                    format!("Key '{key_id_input}' does not exist"),
533                )
534            })
535    }
536
537    fn create_key(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
538        let input = CreateKeyInput::from_body(&req.json_body())?;
539
540        let mut accounts = self.state.write();
541        let state = accounts.get_or_create(&req.account_id);
542
543        let key_id = if input.multi_region {
544            format!("mrk-{}", Uuid::new_v4().as_simple())
545        } else {
546            Uuid::new_v4().to_string()
547        };
548
549        let arn = format!(
550            "arn:aws:kms:{}:{}:key/{}",
551            state.region, state.account_id, key_id
552        );
553        let now = Utc::now().timestamp() as f64;
554
555        let signing_algs = if input.key_usage == "SIGN_VERIFY" {
556            signing_algorithms_for_key_spec(&input.key_spec)
557        } else {
558            None
559        };
560        let encryption_algs = encryption_algorithms_for_key(&input.key_usage, &input.key_spec);
561        let mac_algs = if input.key_usage == "GENERATE_VERIFY_MAC" {
562            mac_algorithms_for_key_spec(&input.key_spec)
563        } else {
564            None
565        };
566
567        let key_policy = input
568            .policy
569            .unwrap_or_else(|| default_key_policy(&state.account_id));
570
571        let mut asym_priv: Option<Vec<u8>> = None;
572        let mut asym_pub: Option<Vec<u8>> = None;
573        if let Some((p, k)) = asym::generate_keypair(&input.key_spec).map_err(|e| {
574            AwsServiceError::aws_error(
575                StatusCode::INTERNAL_SERVER_ERROR,
576                "KMSInternalException",
577                format!("failed to generate asymmetric key: {e}"),
578            )
579        })? {
580            asym_priv = Some(p);
581            asym_pub = Some(k);
582        } else if let Some((p, k)) = asym_ecdsa::generate_keypair(&input.key_spec).map_err(|e| {
583            AwsServiceError::aws_error(
584                StatusCode::INTERNAL_SERVER_ERROR,
585                "KMSInternalException",
586                format!("failed to generate ecdsa key: {e}"),
587            )
588        })? {
589            asym_priv = Some(p);
590            asym_pub = Some(k);
591        }
592
593        // Refuse asymmetric specs we cannot really generate keys for
594        // rather than store a no-DER key that would later fall through
595        // to a fake-bytes Sign/Verify path. SM2 currently has no
596        // pure-Rust impl wired in.
597        let is_asymmetric = input.key_spec.starts_with("ECC_")
598            || input.key_spec.starts_with("RSA_")
599            || input.key_spec == "SM2";
600        if is_asymmetric && asym_priv.is_none() {
601            return Err(AwsServiceError::aws_error(
602                StatusCode::BAD_REQUEST,
603                "UnsupportedOperationException",
604                format!(
605                    "KeySpec '{}' is not supported by this fakecloud build; \
606                     no fake-signature fallback is provided",
607                    input.key_spec
608                ),
609            ));
610        }
611
612        let key = KmsKey {
613            key_id: key_id.clone(),
614            arn: arn.clone(),
615            creation_date: now,
616            description: input.description,
617            enabled: true,
618            key_usage: input.key_usage,
619            key_spec: input.key_spec,
620            key_manager: "CUSTOMER".to_string(),
621            key_state: "Enabled".to_string(),
622            deletion_date: None,
623            tags: input.tags,
624            policy: key_policy,
625            key_rotation_enabled: false,
626            rotation_period_in_days: None,
627            origin: input.origin,
628            multi_region: input.multi_region,
629            rotations: Vec::new(),
630            signing_algorithms: signing_algs,
631            encryption_algorithms: encryption_algs,
632            mac_algorithms: mac_algs,
633            custom_key_store_id: input.custom_key_store_id,
634            imported_key_material: false,
635            imported_material_bytes: None,
636            private_key_seed: rand_bytes(32),
637            primary_region: None,
638            asymmetric_private_key_der: asym_priv,
639            asymmetric_public_key_der: asym_pub,
640        };
641
642        let metadata = key_metadata_json(&key, &state.account_id);
643        state.keys.insert(key_id, key);
644
645        Ok(AwsResponse::json(
646            StatusCode::OK,
647            serde_json::to_string(&json!({ "KeyMetadata": metadata })).unwrap(),
648        ))
649    }
650
651    fn describe_key(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
652        let body = req.json_body();
653        let key_id_input = body["KeyId"].as_str().ok_or_else(|| {
654            AwsServiceError::aws_error(
655                StatusCode::BAD_REQUEST,
656                "ValidationException",
657                "KeyId is required",
658            )
659        })?;
660
661        let accounts = self.state.read();
662        let empty = KmsState::new(&req.account_id, &req.region);
663        let state = accounts.get(&req.account_id).unwrap_or(&empty);
664
665        // Check key policy for Deny rules
666        let resolved = Self::resolve_key_id_with_state(state, key_id_input).ok_or_else(|| {
667            AwsServiceError::aws_error(
668                StatusCode::BAD_REQUEST,
669                "NotFoundException",
670                format!("Key '{key_id_input}' does not exist"),
671            )
672        })?;
673
674        let key = state.keys.get(&resolved).ok_or_else(|| {
675            AwsServiceError::aws_error(
676                StatusCode::BAD_REQUEST,
677                "NotFoundException",
678                format!("Key '{key_id_input}' does not exist"),
679            )
680        })?;
681
682        // Check policy for Deny on DescribeKey
683        check_policy_deny(key, "kms:DescribeKey")?;
684
685        let metadata = key_metadata_json(key, &state.account_id);
686        Ok(AwsResponse::json(
687            StatusCode::OK,
688            serde_json::to_string(&json!({ "KeyMetadata": metadata })).unwrap(),
689        ))
690    }
691
692    fn get_key_last_usage(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
693        let body = req.json_body();
694        let key_id_input = body["KeyId"].as_str().ok_or_else(|| {
695            AwsServiceError::aws_error(
696                StatusCode::BAD_REQUEST,
697                "ValidationException",
698                "KeyId is required",
699            )
700        })?;
701
702        let accounts = self.state.read();
703        let empty = KmsState::new(&req.account_id, &req.region);
704        let state = accounts.get(&req.account_id).unwrap_or(&empty);
705
706        let resolved = Self::resolve_key_id_with_state(state, key_id_input).ok_or_else(|| {
707            AwsServiceError::aws_error(
708                StatusCode::BAD_REQUEST,
709                "NotFoundException",
710                format!("Key '{key_id_input}' does not exist"),
711            )
712        })?;
713        let key = state.keys.get(&resolved).ok_or_else(|| {
714            AwsServiceError::aws_error(
715                StatusCode::BAD_REQUEST,
716                "NotFoundException",
717                format!("Key '{key_id_input}' does not exist"),
718            )
719        })?;
720
721        // KMS started tracking on the key's creation date. We don't yet
722        // record per-op timestamps, so KeyLastUsage is omitted — the AWS
723        // spec explicitly allows this when no tracked op has run.
724        Ok(AwsResponse::json(
725            StatusCode::OK,
726            serde_json::to_string(&json!({
727                "KeyId": key.key_id,
728                "KeyCreationDate": key.creation_date,
729                "TrackingStartDate": key.creation_date,
730            }))
731            .unwrap(),
732        ))
733    }
734
735    fn list_keys(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
736        let body = req.json_body();
737
738        // ListKeys only declares InvalidMarkerException / KMSInternal /
739        // DependencyTimeout; map both Limit and Marker shape failures onto
740        // InvalidMarkerException so the probe sees a declared error.
741        recoded("InvalidMarkerException", || {
742            validate_optional_json_range("limit", &body["Limit"], 1, 1000)
743        })?;
744        recoded("InvalidMarkerException", || {
745            validate_optional_string_length("marker", body["Marker"].as_str(), 1, 320)
746        })?;
747
748        let limit = body["Limit"].as_i64().unwrap_or(1000) as usize;
749        let marker = body["Marker"].as_str();
750
751        let accounts = self.state.read();
752        let empty = KmsState::new(&req.account_id, &req.region);
753        let state = accounts.get(&req.account_id).unwrap_or(&empty);
754        let all_keys: Vec<Value> = state
755            .keys
756            .values()
757            .map(|k| {
758                json!({
759                    "KeyId": k.key_id,
760                    "KeyArn": k.arn,
761                })
762            })
763            .collect();
764
765        let start = if let Some(m) = marker {
766            all_keys
767                .iter()
768                .position(|k| k["KeyId"].as_str() == Some(m))
769                .map(|pos| pos + 1)
770                .unwrap_or(0)
771        } else {
772            0
773        };
774
775        let page = &all_keys[start..all_keys.len().min(start + limit)];
776        let truncated = start + limit < all_keys.len();
777
778        let mut result = json!({
779            "Keys": page,
780            "Truncated": truncated,
781        });
782
783        if truncated {
784            if let Some(last) = page.last() {
785                result["NextMarker"] = last["KeyId"].clone();
786            }
787        }
788
789        Ok(AwsResponse::json(
790            StatusCode::OK,
791            serde_json::to_string(&result).unwrap(),
792        ))
793    }
794
795    fn enable_key(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
796        let body = req.json_body();
797        let resolved = self.resolve_required_key(req, &body)?;
798
799        let mut accounts = self.state.write();
800        let state = accounts.get_or_create(&req.account_id);
801        let key = state.keys.get_mut(&resolved).ok_or_else(|| {
802            AwsServiceError::aws_error(
803                StatusCode::INTERNAL_SERVER_ERROR,
804                "KMSInternalException",
805                "Key state became inconsistent",
806            )
807        })?;
808        key.enabled = true;
809        key.key_state = "Enabled".to_string();
810
811        Ok(AwsResponse::json(StatusCode::OK, "{}"))
812    }
813
814    fn disable_key(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
815        let body = req.json_body();
816        let resolved = self.resolve_required_key(req, &body)?;
817
818        let mut accounts = self.state.write();
819        let state = accounts.get_or_create(&req.account_id);
820        let key = state.keys.get_mut(&resolved).ok_or_else(|| {
821            AwsServiceError::aws_error(
822                StatusCode::INTERNAL_SERVER_ERROR,
823                "KMSInternalException",
824                "Key state became inconsistent",
825            )
826        })?;
827        key.enabled = false;
828        key.key_state = "Disabled".to_string();
829
830        Ok(AwsResponse::json(StatusCode::OK, "{}"))
831    }
832
833    fn schedule_key_deletion(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
834        let body = req.json_body();
835        let resolved = self.resolve_required_key(req, &body)?;
836        let pending_days = body["PendingWindowInDays"].as_i64().unwrap_or(30);
837
838        let mut accounts = self.state.write();
839        let state = accounts.get_or_create(&req.account_id);
840        let key = state.keys.get_mut(&resolved).ok_or_else(|| {
841            AwsServiceError::aws_error(
842                StatusCode::INTERNAL_SERVER_ERROR,
843                "KMSInternalException",
844                "Key state became inconsistent",
845            )
846        })?;
847        let deletion_date =
848            Utc::now().timestamp() as f64 + (pending_days as f64 * 24.0 * 60.0 * 60.0);
849        key.key_state = "PendingDeletion".to_string();
850        key.enabled = false;
851        key.deletion_date = Some(deletion_date);
852
853        Ok(AwsResponse::json(
854            StatusCode::OK,
855            serde_json::to_string(&json!({
856                "KeyId": key.key_id,
857                "DeletionDate": deletion_date,
858                "KeyState": "PendingDeletion",
859                "PendingWindowInDays": pending_days,
860            }))
861            .unwrap(),
862        ))
863    }
864
865    fn cancel_key_deletion(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
866        let body = req.json_body();
867        let resolved = self.resolve_required_key(req, &body)?;
868
869        let mut accounts = self.state.write();
870        let state = accounts.get_or_create(&req.account_id);
871        let key = state.keys.get_mut(&resolved).ok_or_else(|| {
872            AwsServiceError::aws_error(
873                StatusCode::INTERNAL_SERVER_ERROR,
874                "KMSInternalException",
875                "Key state became inconsistent",
876            )
877        })?;
878        key.key_state = "Disabled".to_string();
879        key.deletion_date = None;
880
881        Ok(AwsResponse::json(
882            StatusCode::OK,
883            serde_json::to_string(&json!({
884                "KeyId": key.key_id,
885            }))
886            .unwrap(),
887        ))
888    }
889
890    fn tag_resource(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
891        let body = req.json_body();
892        let key_id = Self::require_key_id(&body)?;
893
894        let resolved = self
895            .resolve_key_id_for(&req.account_id, &req.region, &key_id)
896            .ok_or_else(|| {
897                AwsServiceError::aws_error(
898                    StatusCode::BAD_REQUEST,
899                    "NotFoundException",
900                    format!("Invalid keyId {key_id}"),
901                )
902            })?;
903
904        let mut accounts = self.state.write();
905        let state = accounts.get_or_create(&req.account_id);
906        let key = state.keys.get_mut(&resolved).ok_or_else(|| {
907            AwsServiceError::aws_error(
908                StatusCode::INTERNAL_SERVER_ERROR,
909                "KMSInternalException",
910                "Key state became inconsistent",
911            )
912        })?;
913
914        fakecloud_core::tags::apply_tags(&mut key.tags, &body, "Tags", "TagKey", "TagValue")
915            .map_err(|f| {
916                AwsServiceError::aws_error(
917                    StatusCode::BAD_REQUEST,
918                    "ValidationException",
919                    format!("{f} must be a list"),
920                )
921            })?;
922
923        Ok(AwsResponse::json(StatusCode::OK, "{}"))
924    }
925
926    fn untag_resource(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
927        let body = req.json_body();
928        let key_id = Self::require_key_id(&body)?;
929
930        let resolved = self
931            .resolve_key_id_for(&req.account_id, &req.region, &key_id)
932            .ok_or_else(|| {
933                AwsServiceError::aws_error(
934                    StatusCode::BAD_REQUEST,
935                    "NotFoundException",
936                    format!("Invalid keyId {key_id}"),
937                )
938            })?;
939
940        let mut accounts = self.state.write();
941        let state = accounts.get_or_create(&req.account_id);
942        let key = state.keys.get_mut(&resolved).ok_or_else(|| {
943            AwsServiceError::aws_error(
944                StatusCode::INTERNAL_SERVER_ERROR,
945                "KMSInternalException",
946                "Key state became inconsistent",
947            )
948        })?;
949
950        fakecloud_core::tags::remove_tags(&mut key.tags, &body, "TagKeys").map_err(|f| {
951            AwsServiceError::aws_error(
952                StatusCode::BAD_REQUEST,
953                "ValidationException",
954                format!("{f} must be a list"),
955            )
956        })?;
957
958        Ok(AwsResponse::json(StatusCode::OK, "{}"))
959    }
960
961    fn list_resource_tags(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
962        let body = req.json_body();
963        let key_id = Self::require_key_id(&body)?;
964
965        let resolved = self
966            .resolve_key_id_for(&req.account_id, &req.region, &key_id)
967            .ok_or_else(|| {
968                AwsServiceError::aws_error(
969                    StatusCode::BAD_REQUEST,
970                    "NotFoundException",
971                    format!("Invalid keyId {key_id}"),
972                )
973            })?;
974
975        let accounts = self.state.read();
976        let empty = KmsState::new(&req.account_id, &req.region);
977        let state = accounts.get(&req.account_id).unwrap_or(&empty);
978        let key = state.keys.get(&resolved).ok_or_else(|| {
979            AwsServiceError::aws_error(
980                StatusCode::INTERNAL_SERVER_ERROR,
981                "KMSInternalException",
982                "Key state became inconsistent",
983            )
984        })?;
985        let tags = fakecloud_core::tags::tags_to_json(&key.tags, "TagKey", "TagValue");
986
987        Ok(AwsResponse::json(
988            StatusCode::OK,
989            serde_json::to_string(&json!({
990                "Tags": tags,
991                "Truncated": false,
992            }))
993            .unwrap(),
994        ))
995    }
996
997    fn update_key_description(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
998        let body = req.json_body();
999        let resolved = self.resolve_required_key(req, &body)?;
1000        let description = body["Description"].as_str().unwrap_or("").to_string();
1001
1002        let mut accounts = self.state.write();
1003        let state = accounts.get_or_create(&req.account_id);
1004        let key = state.keys.get_mut(&resolved).ok_or_else(|| {
1005            AwsServiceError::aws_error(
1006                StatusCode::INTERNAL_SERVER_ERROR,
1007                "KMSInternalException",
1008                "Key state became inconsistent",
1009            )
1010        })?;
1011        key.description = description;
1012
1013        Ok(AwsResponse::json(StatusCode::OK, "{}"))
1014    }
1015
1016    fn get_key_policy(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
1017        let body = req.json_body();
1018        let key_id = Self::require_key_id(&body)?;
1019
1020        // For key policy operations, aliases should not work
1021        if key_id.starts_with("alias/") {
1022            return Err(AwsServiceError::aws_error(
1023                StatusCode::BAD_REQUEST,
1024                "NotFoundException",
1025                format!("Invalid keyId {key_id}"),
1026            ));
1027        }
1028
1029        let resolved = self
1030            .resolve_key_id_for(&req.account_id, &req.region, &key_id)
1031            .ok_or_else(|| {
1032                AwsServiceError::aws_error(
1033                    StatusCode::BAD_REQUEST,
1034                    "NotFoundException",
1035                    format!("Key '{key_id}' does not exist"),
1036                )
1037            })?;
1038
1039        let accounts = self.state.read();
1040        let empty = KmsState::new(&req.account_id, &req.region);
1041        let state = accounts.get(&req.account_id).unwrap_or(&empty);
1042        let key = state.keys.get(&resolved).ok_or_else(|| {
1043            AwsServiceError::aws_error(
1044                StatusCode::INTERNAL_SERVER_ERROR,
1045                "KMSInternalException",
1046                "Key state became inconsistent",
1047            )
1048        })?;
1049
1050        Ok(AwsResponse::json(
1051            StatusCode::OK,
1052            serde_json::to_string(&json!({
1053                "Policy": key.policy,
1054            }))
1055            .unwrap(),
1056        ))
1057    }
1058
1059    fn put_key_policy(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
1060        let body = req.json_body();
1061        let key_id = Self::require_key_id(&body)?;
1062
1063        // For key policy operations, aliases should not work
1064        if key_id.starts_with("alias/") {
1065            return Err(AwsServiceError::aws_error(
1066                StatusCode::BAD_REQUEST,
1067                "NotFoundException",
1068                format!("Invalid keyId {key_id}"),
1069            ));
1070        }
1071
1072        let resolved = self
1073            .resolve_key_id_for(&req.account_id, &req.region, &key_id)
1074            .ok_or_else(|| {
1075                AwsServiceError::aws_error(
1076                    StatusCode::BAD_REQUEST,
1077                    "NotFoundException",
1078                    format!("Key '{key_id}' does not exist"),
1079                )
1080            })?;
1081
1082        let policy = body["Policy"].as_str().unwrap_or("").to_string();
1083
1084        let mut accounts = self.state.write();
1085        let state = accounts.get_or_create(&req.account_id);
1086        let key = state.keys.get_mut(&resolved).ok_or_else(|| {
1087            AwsServiceError::aws_error(
1088                StatusCode::INTERNAL_SERVER_ERROR,
1089                "KMSInternalException",
1090                "Key state became inconsistent",
1091            )
1092        })?;
1093        key.policy = policy;
1094
1095        Ok(AwsResponse::json(StatusCode::OK, "{}"))
1096    }
1097
1098    fn list_key_policies(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
1099        let body = req.json_body();
1100        let _resolved = self.resolve_required_key(req, &body)?;
1101
1102        Ok(AwsResponse::json(
1103            StatusCode::OK,
1104            serde_json::to_string(&json!({
1105                "PolicyNames": ["default"],
1106                "Truncated": false,
1107            }))
1108            .unwrap(),
1109        ))
1110    }
1111
1112    fn get_key_rotation_status(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
1113        let body = req.json_body();
1114        let key_id = Self::require_key_id(&body)?;
1115
1116        // Real KMS resolves alias/* and alias-ARNs identically to a key id.
1117        let resolved = self
1118            .resolve_key_id_for(&req.account_id, &req.region, &key_id)
1119            .ok_or_else(|| {
1120                AwsServiceError::aws_error(
1121                    StatusCode::BAD_REQUEST,
1122                    "NotFoundException",
1123                    format!("Key '{key_id}' does not exist"),
1124                )
1125            })?;
1126
1127        let accounts = self.state.read();
1128        let empty = KmsState::new(&req.account_id, &req.region);
1129        let state = accounts.get(&req.account_id).unwrap_or(&empty);
1130        let key = state.keys.get(&resolved).ok_or_else(|| {
1131            AwsServiceError::aws_error(
1132                StatusCode::INTERNAL_SERVER_ERROR,
1133                "KMSInternalException",
1134                "Key state became inconsistent",
1135            )
1136        })?;
1137
1138        let mut result = json!({
1139            "KeyRotationEnabled": key.key_rotation_enabled,
1140        });
1141        // AWS echoes the configured cadence (default 365) only while rotation
1142        // is enabled.
1143        if key.key_rotation_enabled {
1144            result["RotationPeriodInDays"] = json!(key.rotation_period_in_days.unwrap_or(365));
1145        }
1146
1147        Ok(AwsResponse::json(
1148            StatusCode::OK,
1149            serde_json::to_string(&result).unwrap(),
1150        ))
1151    }
1152
1153    fn enable_key_rotation(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
1154        let body = req.json_body();
1155        let key_id = Self::require_key_id(&body)?;
1156
1157        // Real KMS resolves alias/* and alias-ARNs identically to a key id
1158        // here. Earlier code rejected `alias/*` outright, breaking IaC
1159        // configs that reference keys by alias.
1160        let resolved = self
1161            .resolve_key_id_for(&req.account_id, &req.region, &key_id)
1162            .ok_or_else(|| {
1163                AwsServiceError::aws_error(
1164                    StatusCode::BAD_REQUEST,
1165                    "NotFoundException",
1166                    format!("Key '{key_id}' does not exist"),
1167                )
1168            })?;
1169
1170        let mut accounts = self.state.write();
1171        let state = accounts.get_or_create(&req.account_id);
1172        let key = state.keys.get_mut(&resolved).ok_or_else(|| {
1173            AwsServiceError::aws_error(
1174                StatusCode::INTERNAL_SERVER_ERROR,
1175                "KMSInternalException",
1176                "Key state became inconsistent",
1177            )
1178        })?;
1179        // RotationPeriodInDays is optional; AWS validates 90..=2560 and
1180        // defaults to 365. Persist it so GetKeyRotationStatus echoes it back
1181        // instead of dropping it (bug-audit 2026-06-20, 1.24).
1182        if let Some(period) = body.get("RotationPeriodInDays").and_then(|v| v.as_i64()) {
1183            if !(90..=2560).contains(&period) {
1184                return Err(AwsServiceError::aws_error(
1185                    StatusCode::BAD_REQUEST,
1186                    "ValidationException",
1187                    format!("RotationPeriodInDays must be between 90 and 2560, got {period}"),
1188                ));
1189            }
1190            key.rotation_period_in_days = Some(period as i32);
1191        }
1192        key.key_rotation_enabled = true;
1193
1194        Ok(AwsResponse::json(StatusCode::OK, "{}"))
1195    }
1196
1197    fn disable_key_rotation(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
1198        let body = req.json_body();
1199        let key_id = Self::require_key_id(&body)?;
1200
1201        // Real KMS resolves alias/* and alias-ARNs identically to a key id.
1202        let resolved = self
1203            .resolve_key_id_for(&req.account_id, &req.region, &key_id)
1204            .ok_or_else(|| {
1205                AwsServiceError::aws_error(
1206                    StatusCode::BAD_REQUEST,
1207                    "NotFoundException",
1208                    format!("Key '{key_id}' does not exist"),
1209                )
1210            })?;
1211
1212        let mut accounts = self.state.write();
1213        let state = accounts.get_or_create(&req.account_id);
1214        let key = state.keys.get_mut(&resolved).ok_or_else(|| {
1215            AwsServiceError::aws_error(
1216                StatusCode::INTERNAL_SERVER_ERROR,
1217                "KMSInternalException",
1218                "Key state became inconsistent",
1219            )
1220        })?;
1221        key.key_rotation_enabled = false;
1222
1223        Ok(AwsResponse::json(StatusCode::OK, "{}"))
1224    }
1225
1226    fn rotate_key_on_demand(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
1227        let body = req.json_body();
1228        let resolved = self.resolve_required_key(req, &body)?;
1229
1230        let mut accounts = self.state.write();
1231        let state = accounts.get_or_create(&req.account_id);
1232        let key = state.keys.get_mut(&resolved).ok_or_else(|| {
1233            AwsServiceError::aws_error(
1234                StatusCode::INTERNAL_SERVER_ERROR,
1235                "KMSInternalException",
1236                "Key state became inconsistent",
1237            )
1238        })?;
1239
1240        let rotation = KeyRotation {
1241            key_id: key.key_id.clone(),
1242            rotation_date: Utc::now().timestamp() as f64,
1243            rotation_type: "ON_DEMAND".to_string(),
1244        };
1245        key.rotations.push(rotation);
1246
1247        Ok(AwsResponse::json(
1248            StatusCode::OK,
1249            serde_json::to_string(&json!({
1250                "KeyId": key.key_id,
1251            }))
1252            .unwrap(),
1253        ))
1254    }
1255
1256    fn list_key_rotations(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
1257        let body = req.json_body();
1258        let resolved = self.resolve_required_key(req, &body)?;
1259        validate_optional_json_range("limit", &body["Limit"], 1, 1000)?;
1260        let limit = body["Limit"].as_i64().unwrap_or(1000) as usize;
1261        let marker = body["Marker"].as_str();
1262
1263        let accounts = self.state.read();
1264        let empty = KmsState::new(&req.account_id, &req.region);
1265        let state = accounts.get(&req.account_id).unwrap_or(&empty);
1266        let key = state.keys.get(&resolved).ok_or_else(|| {
1267            AwsServiceError::aws_error(
1268                StatusCode::INTERNAL_SERVER_ERROR,
1269                "KMSInternalException",
1270                "Key state became inconsistent",
1271            )
1272        })?;
1273
1274        let start_index = if let Some(marker) = marker {
1275            marker.parse::<usize>().unwrap_or(0)
1276        } else {
1277            0
1278        };
1279
1280        let rotations: Vec<Value> = key
1281            .rotations
1282            .iter()
1283            .skip(start_index)
1284            .take(limit)
1285            .map(|r| {
1286                json!({
1287                    "KeyId": r.key_id,
1288                    "RotationDate": r.rotation_date,
1289                    "RotationType": r.rotation_type,
1290                })
1291            })
1292            .collect();
1293
1294        let total_after_start = key.rotations.len().saturating_sub(start_index);
1295        let truncated = total_after_start > limit;
1296
1297        let mut response = json!({
1298            "Rotations": rotations,
1299            "Truncated": truncated,
1300        });
1301
1302        if truncated {
1303            response["NextMarker"] = json!((start_index + limit).to_string());
1304        }
1305
1306        Ok(AwsResponse::json(
1307            StatusCode::OK,
1308            serde_json::to_string(&response).unwrap(),
1309        ))
1310    }
1311
1312    fn replicate_key(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
1313        let body = req.json_body();
1314        let key_id = Self::require_key_id(&body)?;
1315        let replica_region = body["ReplicaRegion"].as_str().unwrap_or("").to_string();
1316
1317        let resolved = self
1318            .resolve_key_id_for(&req.account_id, &req.region, &key_id)
1319            .ok_or_else(|| {
1320                AwsServiceError::aws_error(
1321                    StatusCode::BAD_REQUEST,
1322                    "NotFoundException",
1323                    format!("Key '{key_id}' does not exist"),
1324                )
1325            })?;
1326
1327        let mut accounts = self.state.write();
1328        let state = accounts.get_or_create(&req.account_id);
1329
1330        // Clone the source key once and drop the borrow — the replica reuses
1331        // every field except the region-dependent ones.
1332        let source_key = state
1333            .keys
1334            .get(&resolved)
1335            .ok_or_else(|| {
1336                AwsServiceError::aws_error(
1337                    StatusCode::INTERNAL_SERVER_ERROR,
1338                    "KMSInternalException",
1339                    "Key state became inconsistent",
1340                )
1341            })?
1342            .clone();
1343        let account_id = state.account_id.clone();
1344        let source_region = state.region.clone();
1345
1346        let replica_arn = format!(
1347            "arn:aws:kms:{}:{}:key/{}",
1348            replica_region, account_id, source_key.key_id
1349        );
1350
1351        let metadata = json!({
1352            "KeyId": source_key.key_id,
1353            "Arn": replica_arn,
1354            "AWSAccountId": account_id,
1355            "CreationDate": source_key.creation_date,
1356            "Description": source_key.description,
1357            "Enabled": source_key.enabled,
1358            "KeyUsage": source_key.key_usage,
1359            "KeySpec": source_key.key_spec,
1360            "CustomerMasterKeySpec": source_key.key_spec,
1361            "KeyManager": source_key.key_manager,
1362            "KeyState": source_key.key_state,
1363            "Origin": source_key.origin,
1364            "MultiRegion": true,
1365            "MultiRegionConfiguration": {
1366                "MultiRegionKeyType": "REPLICA",
1367                "PrimaryKey": {
1368                    "Arn": source_key.arn,
1369                    "Region": source_region,
1370                },
1371                "ReplicaKeys": [],
1372            },
1373        });
1374
1375        let replica_storage_key = format!("{}:{}", replica_region, source_key.key_id);
1376        let source_policy = source_key.policy.clone();
1377        let replica_key = KmsKey {
1378            arn: replica_arn,
1379            deletion_date: None,
1380            key_rotation_enabled: false,
1381            rotation_period_in_days: None,
1382            multi_region: true,
1383            rotations: Vec::new(),
1384            custom_key_store_id: None,
1385            imported_key_material: false,
1386            imported_material_bytes: None,
1387            private_key_seed: rand_bytes(32),
1388            primary_region: None,
1389            ..source_key
1390        };
1391
1392        state.keys.insert(replica_storage_key, replica_key);
1393
1394        Ok(AwsResponse::json(
1395            StatusCode::OK,
1396            serde_json::to_string(&json!({
1397                "ReplicaKeyMetadata": metadata,
1398                "ReplicaPolicy": source_policy,
1399            }))
1400            .unwrap(),
1401        ))
1402    }
1403
1404    fn update_primary_region(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
1405        let body = req.json_body();
1406        let key_id = Self::require_key_id(&body)?;
1407        let primary_region = body["PrimaryRegion"]
1408            .as_str()
1409            .ok_or_else(|| {
1410                AwsServiceError::aws_error(
1411                    StatusCode::BAD_REQUEST,
1412                    "ValidationException",
1413                    "PrimaryRegion is required",
1414                )
1415            })?
1416            .to_string();
1417
1418        let resolved = self
1419            .resolve_key_id_for(&req.account_id, &req.region, &key_id)
1420            .ok_or_else(|| {
1421                AwsServiceError::aws_error(
1422                    StatusCode::BAD_REQUEST,
1423                    "NotFoundException",
1424                    format!("Key '{key_id}' does not exist"),
1425                )
1426            })?;
1427
1428        let mut accounts = self.state.write();
1429        let state = accounts.get_or_create(&req.account_id);
1430        let account_id = state.account_id.clone();
1431        let key = state.keys.get_mut(&resolved).ok_or_else(|| {
1432            AwsServiceError::aws_error(
1433                StatusCode::BAD_REQUEST,
1434                "NotFoundException",
1435                format!("Key '{key_id}' does not exist"),
1436            )
1437        })?;
1438
1439        if !key.multi_region {
1440            return Err(AwsServiceError::aws_error(
1441                StatusCode::BAD_REQUEST,
1442                "UnsupportedOperationException",
1443                format!("Key '{}' is not a multi-Region key", key.arn),
1444            ));
1445        }
1446        key.primary_region = Some(primary_region.clone());
1447        // Update the ARN to reflect the new region
1448        key.arn = format!(
1449            "arn:aws:kms:{}:{}:key/{}",
1450            primary_region, account_id, key.key_id
1451        );
1452
1453        Ok(AwsResponse::json(StatusCode::OK, "{}"))
1454    }
1455}
1456
1457#[path = "asym.rs"]
1458pub(crate) mod asym;
1459#[path = "asym_ecdsa.rs"]
1460pub(crate) mod asym_ecdsa;
1461#[path = "mac.rs"]
1462pub(crate) mod mac;
1463#[path = "service_aliases.rs"]
1464mod service_aliases;
1465#[path = "service_crypto.rs"]
1466mod service_crypto;
1467#[path = "service_custom_store.rs"]
1468mod service_custom_store;
1469#[path = "service_grants.rs"]
1470mod service_grants;
1471
1472#[path = "helpers.rs"]
1473mod helpers;
1474pub(crate) use helpers::*;
1475
1476#[path = "provisioner.rs"]
1477pub mod provisioner;
1478
1479#[cfg(test)]
1480#[path = "service_tests.rs"]
1481mod tests;