Skip to main content

fakecloud_kms/
state.rs

1use std::collections::HashMap;
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: HashMap<String, KmsKey>,
19    pub aliases: HashMap<String, KmsAlias>,
20    pub grants: Vec<KmsGrant>,
21    pub custom_key_stores: HashMap<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}
30
31fn default_master_key_bytes() -> Vec<u8> {
32    use aes_gcm::aead::rand_core::RngCore;
33    use aes_gcm::aead::OsRng;
34    let mut bytes = vec![0u8; 32];
35    OsRng.fill_bytes(&mut bytes);
36    bytes
37}
38
39impl KmsState {
40    pub fn new(account_id: &str, region: &str) -> Self {
41        Self {
42            account_id: account_id.to_string(),
43            region: region.to_string(),
44            keys: HashMap::new(),
45            aliases: HashMap::new(),
46            grants: Vec::new(),
47            custom_key_stores: HashMap::new(),
48            master_key_bytes: default_master_key_bytes(),
49        }
50    }
51
52    pub fn reset(&mut self) {
53        self.keys.clear();
54        self.aliases.clear();
55        self.grants.clear();
56        self.custom_key_stores.clear();
57        // Keep the master key across resets so ciphertexts still decrypt.
58    }
59}
60
61#[derive(Clone, serde::Serialize, serde::Deserialize)]
62pub struct KmsKey {
63    pub key_id: String,
64    pub arn: String,
65    pub creation_date: f64,
66    pub description: String,
67    pub enabled: bool,
68    pub key_usage: String,
69    pub key_spec: String,
70    pub key_manager: String,
71    pub key_state: String,
72    pub deletion_date: Option<f64>,
73    pub tags: HashMap<String, String>,
74    pub policy: String,
75    pub key_rotation_enabled: bool,
76    pub origin: String,
77    pub multi_region: bool,
78    pub rotations: Vec<KeyRotation>,
79    pub signing_algorithms: Option<Vec<String>>,
80    pub encryption_algorithms: Option<Vec<String>>,
81    pub mac_algorithms: Option<Vec<String>>,
82    pub custom_key_store_id: Option<String>,
83    pub imported_key_material: bool,
84    /// Raw bytes of imported key material (used as AES key for encrypt/decrypt).
85    pub imported_material_bytes: Option<Vec<u8>>,
86    /// Deterministic seed for the key (used for DeriveSharedSecret).
87    pub private_key_seed: Vec<u8>,
88    pub primary_region: Option<String>,
89}
90
91#[derive(Clone, serde::Serialize, serde::Deserialize)]
92pub struct KmsAlias {
93    pub alias_name: String,
94    pub alias_arn: String,
95    pub target_key_id: String,
96    pub creation_date: f64,
97}
98
99#[derive(Clone, serde::Serialize, serde::Deserialize)]
100pub struct KmsGrant {
101    pub grant_id: String,
102    pub grant_token: String,
103    pub key_id: String,
104    pub grantee_principal: String,
105    pub retiring_principal: Option<String>,
106    pub operations: Vec<String>,
107    pub constraints: Option<serde_json::Value>,
108    pub name: Option<String>,
109    pub creation_date: f64,
110}
111
112#[derive(Clone, serde::Serialize, serde::Deserialize)]
113pub struct KeyRotation {
114    pub key_id: String,
115    pub rotation_date: f64,
116    pub rotation_type: String,
117}
118
119#[derive(Clone, serde::Serialize, serde::Deserialize)]
120pub struct CustomKeyStore {
121    pub custom_key_store_id: String,
122    pub custom_key_store_name: String,
123    pub custom_key_store_type: String,
124    pub cloud_hsm_cluster_id: Option<String>,
125    pub trust_anchor_certificate: Option<String>,
126    pub connection_state: String,
127    pub creation_date: f64,
128    pub xks_proxy_uri_endpoint: Option<String>,
129    pub xks_proxy_uri_path: Option<String>,
130    pub xks_proxy_vpc_endpoint_service_name: Option<String>,
131    pub xks_proxy_connectivity: Option<String>,
132}
133
134/// On-disk snapshot envelope for KMS state. Versioned so format
135/// changes fail loudly on upgrade.
136#[derive(Clone, serde::Serialize, serde::Deserialize)]
137pub struct KmsSnapshot {
138    pub schema_version: u32,
139    #[serde(default)]
140    pub accounts: Option<fakecloud_core::multi_account::MultiAccountState<KmsState>>,
141    #[serde(default)]
142    pub state: Option<KmsState>,
143}
144
145pub const KMS_SNAPSHOT_SCHEMA_VERSION: u32 = 2;
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    #[test]
152    fn new_has_empty_collections() {
153        let state = KmsState::new("123456789012", "us-east-1");
154        assert_eq!(state.account_id, "123456789012");
155        assert_eq!(state.region, "us-east-1");
156        assert!(state.keys.is_empty());
157        assert!(state.aliases.is_empty());
158        assert!(state.grants.is_empty());
159        assert!(state.custom_key_stores.is_empty());
160    }
161
162    #[test]
163    fn reset_clears_collections() {
164        let mut state = KmsState::new("123456789012", "us-east-1");
165        state.aliases.insert(
166            "alias/test".to_string(),
167            KmsAlias {
168                alias_name: "alias/test".to_string(),
169                alias_arn: "arn".to_string(),
170                target_key_id: "k".to_string(),
171                creation_date: 0.0,
172            },
173        );
174        assert!(!state.aliases.is_empty());
175        state.reset();
176        assert!(state.aliases.is_empty());
177    }
178}