Skip to main content

fakecloud_kms/
state.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3
4use parking_lot::RwLock;
5
6pub type SharedKmsState = Arc<RwLock<fakecloud_core::multi_account::MultiAccountState<KmsState>>>;
7
8impl fakecloud_core::multi_account::AccountState for KmsState {
9    fn new_for_account(account_id: &str, region: &str, _endpoint: &str) -> Self {
10        Self::new(account_id, region)
11    }
12}
13
14#[derive(Clone, serde::Serialize, serde::Deserialize)]
15pub struct KmsState {
16    pub account_id: String,
17    pub region: String,
18    pub keys: BTreeMap<String, KmsKey>,
19    pub aliases: BTreeMap<String, KmsAlias>,
20    pub grants: Vec<KmsGrant>,
21    pub custom_key_stores: BTreeMap<String, CustomKeyStore>,
22    /// Per-account master key bytes (32 bytes for AES-256-GCM) used to
23    /// wrap plaintext in the AWS-shaped ciphertext blob format. Generated
24    /// lazily on first encrypt and persisted alongside the rest of the
25    /// state so that ciphertexts produced before a server restart still
26    /// decrypt afterwards.
27    #[serde(default = "default_master_key_bytes")]
28    pub master_key_bytes: Vec<u8>,
29    /// In-flight RSA wrapping keypairs handed out by GetParametersForImport
30    /// and consumed by ImportKeyMaterial to RSA-OAEP-unwrap the encrypted
31    /// key material. Keyed by the import token bytes returned to the
32    /// caller; entries are removed after a successful import.
33    #[serde(default)]
34    pub import_wrapping_keys: BTreeMap<String, ImportWrapEntry>,
35}
36
37#[derive(Clone, serde::Serialize, serde::Deserialize)]
38pub struct ImportWrapEntry {
39    /// PKCS#8 DER-encoded RSA-2048 private key half of the wrapping
40    /// keypair. The corresponding SubjectPublicKeyInfo DER was returned
41    /// to the caller in `GetParametersForImport.PublicKey` so they can
42    /// RSA-OAEP-encrypt the key material under it.
43    pub private_key_der: Vec<u8>,
44    /// CMK key id this token is bound to. ImportKeyMaterial rejects the
45    /// token when it doesn't match the CMK in the request.
46    pub key_id: String,
47}
48
49fn default_master_key_bytes() -> Vec<u8> {
50    use aes_gcm::aead::rand_core::RngCore;
51    use aes_gcm::aead::OsRng;
52    let mut bytes = vec![0u8; 32];
53    OsRng.fill_bytes(&mut bytes);
54    bytes
55}
56
57impl KmsState {
58    pub fn new(account_id: &str, region: &str) -> Self {
59        Self {
60            account_id: account_id.to_string(),
61            region: region.to_string(),
62            keys: BTreeMap::new(),
63            aliases: BTreeMap::new(),
64            grants: Vec::new(),
65            custom_key_stores: BTreeMap::new(),
66            master_key_bytes: default_master_key_bytes(),
67            import_wrapping_keys: BTreeMap::new(),
68        }
69    }
70
71    pub fn reset(&mut self) {
72        self.keys.clear();
73        self.aliases.clear();
74        self.grants.clear();
75        self.custom_key_stores.clear();
76        // Keep the master key across resets so ciphertexts still decrypt.
77    }
78}
79
80#[derive(Clone, serde::Serialize, serde::Deserialize)]
81pub struct KmsKey {
82    pub key_id: String,
83    pub arn: String,
84    pub creation_date: f64,
85    pub description: String,
86    pub enabled: bool,
87    pub key_usage: String,
88    pub key_spec: String,
89    pub key_manager: String,
90    pub key_state: String,
91    pub deletion_date: Option<f64>,
92    pub tags: BTreeMap<String, String>,
93    pub policy: String,
94    pub key_rotation_enabled: bool,
95    /// Customer-specified rotation cadence from EnableKeyRotation
96    /// (`RotationPeriodInDays`, 90..2560). `None` means AWS's 365-day
97    /// default. Echoed by GetKeyRotationStatus.
98    #[serde(default)]
99    pub rotation_period_in_days: Option<i32>,
100    pub origin: String,
101    pub multi_region: bool,
102    pub rotations: Vec<KeyRotation>,
103    pub signing_algorithms: Option<Vec<String>>,
104    pub encryption_algorithms: Option<Vec<String>>,
105    pub mac_algorithms: Option<Vec<String>>,
106    pub custom_key_store_id: Option<String>,
107    pub imported_key_material: bool,
108    /// Raw bytes of imported key material (used as AES key for encrypt/decrypt).
109    pub imported_material_bytes: Option<Vec<u8>>,
110    /// Deterministic seed for the key (used for DeriveSharedSecret).
111    pub private_key_seed: Vec<u8>,
112    pub primary_region: Option<String>,
113    /// PKCS#8 DER-encoded private key, populated for asymmetric specs
114    /// (RSA_2048/3072/4096, ECC_*) at CreateKey time. None for
115    /// symmetric / HMAC specs.
116    #[serde(default, skip_serializing_if = "Option::is_none")]
117    pub asymmetric_private_key_der: Option<Vec<u8>>,
118    /// SubjectPublicKeyInfo DER-encoded public key. Returned by
119    /// GetPublicKey verbatim.
120    #[serde(default, skip_serializing_if = "Option::is_none")]
121    pub asymmetric_public_key_der: Option<Vec<u8>>,
122}
123
124#[derive(Clone, serde::Serialize, serde::Deserialize)]
125pub struct KmsAlias {
126    pub alias_name: String,
127    pub alias_arn: String,
128    pub target_key_id: String,
129    pub creation_date: f64,
130}
131
132#[derive(Clone, serde::Serialize, serde::Deserialize)]
133pub struct KmsGrant {
134    pub grant_id: String,
135    pub grant_token: String,
136    pub key_id: String,
137    pub grantee_principal: String,
138    pub retiring_principal: Option<String>,
139    pub operations: Vec<String>,
140    pub constraints: Option<serde_json::Value>,
141    pub name: Option<String>,
142    pub creation_date: f64,
143}
144
145#[derive(Clone, serde::Serialize, serde::Deserialize)]
146pub struct KeyRotation {
147    pub key_id: String,
148    pub rotation_date: f64,
149    pub rotation_type: String,
150}
151
152#[derive(Clone, serde::Serialize, serde::Deserialize)]
153pub struct CustomKeyStore {
154    pub custom_key_store_id: String,
155    pub custom_key_store_name: String,
156    pub custom_key_store_type: String,
157    pub cloud_hsm_cluster_id: Option<String>,
158    pub trust_anchor_certificate: Option<String>,
159    pub connection_state: String,
160    pub creation_date: f64,
161    pub xks_proxy_uri_endpoint: Option<String>,
162    pub xks_proxy_uri_path: Option<String>,
163    pub xks_proxy_vpc_endpoint_service_name: Option<String>,
164    pub xks_proxy_connectivity: Option<String>,
165}
166
167/// On-disk snapshot envelope for KMS state. Versioned so format
168/// changes fail loudly on upgrade.
169#[derive(Clone, serde::Serialize, serde::Deserialize)]
170pub struct KmsSnapshot {
171    pub schema_version: u32,
172    #[serde(default)]
173    pub accounts: Option<fakecloud_core::multi_account::MultiAccountState<KmsState>>,
174    #[serde(default)]
175    pub state: Option<KmsState>,
176}
177
178pub const KMS_SNAPSHOT_SCHEMA_VERSION: u32 = 2;
179
180#[cfg(test)]
181mod tests {
182    use super::*;
183
184    #[test]
185    fn new_has_empty_collections() {
186        let state = KmsState::new("123456789012", "us-east-1");
187        assert_eq!(state.account_id, "123456789012");
188        assert_eq!(state.region, "us-east-1");
189        assert!(state.keys.is_empty());
190        assert!(state.aliases.is_empty());
191        assert!(state.grants.is_empty());
192        assert!(state.custom_key_stores.is_empty());
193    }
194
195    #[test]
196    fn reset_clears_collections() {
197        let mut state = KmsState::new("123456789012", "us-east-1");
198        state.aliases.insert(
199            "alias/test".to_string(),
200            KmsAlias {
201                alias_name: "alias/test".to_string(),
202                alias_arn: "arn".to_string(),
203                target_key_id: "k".to_string(),
204                creation_date: 0.0,
205            },
206        );
207        assert!(!state.aliases.is_empty());
208        state.reset();
209        assert!(state.aliases.is_empty());
210    }
211}