Skip to main content

oxigdal_security/encryption/
key_management.rs

1//! Key management and rotation.
2
3use crate::encryption::{AtRestEncryptor, EncryptionAlgorithm};
4use crate::error::{Result, SecurityError};
5use chrono::{DateTime, Utc};
6use dashmap::DashMap;
7use serde::{Deserialize, Serialize};
8use std::sync::Arc;
9use uuid::Uuid;
10
11/// Key metadata.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct KeyMetadata {
14    /// Key ID.
15    pub key_id: String,
16    /// Algorithm.
17    pub algorithm: EncryptionAlgorithm,
18    /// Created at.
19    pub created_at: DateTime<Utc>,
20    /// Expires at.
21    pub expires_at: Option<DateTime<Utc>>,
22    /// Rotation period in days.
23    pub rotation_period_days: Option<u32>,
24    /// Whether the key is active.
25    pub active: bool,
26    /// Key version.
27    pub version: u32,
28}
29
30impl KeyMetadata {
31    /// Create new key metadata.
32    pub fn new(
33        key_id: String,
34        algorithm: EncryptionAlgorithm,
35        rotation_period_days: Option<u32>,
36    ) -> Self {
37        let expires_at =
38            rotation_period_days.map(|days| Utc::now() + chrono::Duration::days(days as i64));
39
40        Self {
41            key_id,
42            algorithm,
43            created_at: Utc::now(),
44            expires_at,
45            rotation_period_days,
46            active: true,
47            version: 1,
48        }
49    }
50
51    /// Check if the key is expired.
52    pub fn is_expired(&self) -> bool {
53        self.expires_at.is_some_and(|exp| Utc::now() > exp)
54    }
55
56    /// Check if the key needs rotation (within 7 days of expiration).
57    pub fn needs_rotation(&self) -> bool {
58        self.expires_at
59            .is_some_and(|exp| Utc::now() + chrono::Duration::days(7) > exp)
60    }
61}
62
63/// Key manager for storing and rotating encryption keys.
64pub struct KeyManager {
65    keys: Arc<DashMap<String, (Vec<u8>, KeyMetadata)>>,
66    current_key_id: Arc<parking_lot::RwLock<Option<String>>>,
67}
68
69impl KeyManager {
70    /// Create a new key manager.
71    pub fn new() -> Self {
72        Self {
73            keys: Arc::new(DashMap::new()),
74            current_key_id: Arc::new(parking_lot::RwLock::new(None)),
75        }
76    }
77
78    /// Generate a new key.
79    pub fn generate_key(
80        &self,
81        algorithm: EncryptionAlgorithm,
82        rotation_period_days: Option<u32>,
83    ) -> Result<String> {
84        let key_id = Uuid::new_v4().to_string();
85        let key = AtRestEncryptor::generate_key(algorithm);
86        let metadata = KeyMetadata::new(key_id.clone(), algorithm, rotation_period_days);
87
88        self.keys.insert(key_id.clone(), (key, metadata));
89
90        // Set as current key if no current key exists
91        {
92            let mut current = self.current_key_id.write();
93            if current.is_none() {
94                *current = Some(key_id.clone());
95            }
96        }
97
98        Ok(key_id)
99    }
100
101    /// Add an existing key.
102    pub fn add_key(
103        &self,
104        key_id: String,
105        key: Vec<u8>,
106        algorithm: EncryptionAlgorithm,
107        rotation_period_days: Option<u32>,
108    ) -> Result<()> {
109        let metadata = KeyMetadata::new(key_id.clone(), algorithm, rotation_period_days);
110        self.keys.insert(key_id, (key, metadata));
111        Ok(())
112    }
113
114    /// Get a key by ID.
115    pub fn get_key(&self, key_id: &str) -> Result<(Vec<u8>, KeyMetadata)> {
116        self.keys
117            .get(key_id)
118            .map(|entry| entry.value().clone())
119            .ok_or_else(|| SecurityError::key_management(format!("Key not found: {}", key_id)))
120    }
121
122    /// Get the current active key.
123    pub fn get_current_key(&self) -> Result<(String, Vec<u8>, KeyMetadata)> {
124        let current_id = self
125            .current_key_id
126            .read()
127            .clone()
128            .ok_or_else(|| SecurityError::key_management("No current key set"))?;
129
130        let (key, metadata) = self.get_key(&current_id)?;
131        Ok((current_id, key, metadata))
132    }
133
134    /// Set the current active key.
135    pub fn set_current_key(&self, key_id: String) -> Result<()> {
136        // Verify key exists
137        if !self.keys.contains_key(&key_id) {
138            return Err(SecurityError::key_management(format!(
139                "Key not found: {}",
140                key_id
141            )));
142        }
143
144        let mut current = self.current_key_id.write();
145        *current = Some(key_id);
146        Ok(())
147    }
148
149    /// Rotate the current key.
150    pub fn rotate_key(&self) -> Result<String> {
151        let (current_id, _, metadata) = self.get_current_key()?;
152
153        // Generate new key with same parameters
154        let new_key_id = self.generate_key(metadata.algorithm, metadata.rotation_period_days)?;
155
156        // Deactivate old key
157        if let Some(mut entry) = self.keys.get_mut(&current_id) {
158            entry.value_mut().1.active = false;
159        }
160
161        // Set new key as current
162        self.set_current_key(new_key_id.clone())?;
163
164        Ok(new_key_id)
165    }
166
167    /// List all keys.
168    pub fn list_keys(&self) -> Vec<(String, KeyMetadata)> {
169        self.keys
170            .iter()
171            .map(|entry| (entry.key().clone(), entry.value().1.clone()))
172            .collect()
173    }
174
175    /// List expired keys.
176    pub fn list_expired_keys(&self) -> Vec<(String, KeyMetadata)> {
177        self.keys
178            .iter()
179            .filter(|entry| entry.value().1.is_expired())
180            .map(|entry| (entry.key().clone(), entry.value().1.clone()))
181            .collect()
182    }
183
184    /// List keys that need rotation.
185    pub fn list_keys_needing_rotation(&self) -> Vec<(String, KeyMetadata)> {
186        self.keys
187            .iter()
188            .filter(|entry| entry.value().1.needs_rotation())
189            .map(|entry| (entry.key().clone(), entry.value().1.clone()))
190            .collect()
191    }
192
193    /// Delete a key.
194    pub fn delete_key(&self, key_id: &str) -> Result<()> {
195        // Prevent deleting current key
196        {
197            let current = self.current_key_id.read();
198            if current.as_ref().is_some_and(|id| id == key_id) {
199                return Err(SecurityError::key_management("Cannot delete current key"));
200            }
201        }
202
203        self.keys
204            .remove(key_id)
205            .ok_or_else(|| SecurityError::key_management(format!("Key not found: {}", key_id)))?;
206
207        Ok(())
208    }
209
210    /// Create an encryptor for a specific key.
211    pub fn create_encryptor(&self, key_id: &str) -> Result<AtRestEncryptor> {
212        let (key, metadata) = self.get_key(key_id)?;
213
214        if metadata.is_expired() {
215            return Err(SecurityError::key_management(format!(
216                "Key expired: {}",
217                key_id
218            )));
219        }
220
221        AtRestEncryptor::new(metadata.algorithm, key, key_id.to_string())
222    }
223
224    /// Create an encryptor for the current key.
225    pub fn create_current_encryptor(&self) -> Result<AtRestEncryptor> {
226        let (key_id, key, metadata) = self.get_current_key()?;
227
228        if metadata.is_expired() {
229            return Err(SecurityError::key_management("Current key expired"));
230        }
231
232        AtRestEncryptor::new(metadata.algorithm, key, key_id)
233    }
234
235    /// Get the number of keys.
236    pub fn key_count(&self) -> usize {
237        self.keys.len()
238    }
239
240    /// Clear all keys.
241    pub fn clear(&self) {
242        self.keys.clear();
243        *self.current_key_id.write() = None;
244    }
245}
246
247impl Default for KeyManager {
248    fn default() -> Self {
249        Self::new()
250    }
251}
252
253#[cfg(test)]
254mod tests {
255    use super::*;
256
257    #[test]
258    fn test_key_generation() {
259        let manager = KeyManager::new();
260        let key_id = manager
261            .generate_key(EncryptionAlgorithm::Aes256Gcm, Some(365))
262            .expect("Failed to generate key");
263
264        assert!(!key_id.is_empty());
265
266        let (key, metadata) = manager.get_key(&key_id).expect("Failed to get key");
267        assert_eq!(key.len(), 32);
268        assert_eq!(metadata.algorithm, EncryptionAlgorithm::Aes256Gcm);
269        assert!(metadata.active);
270        assert!(!metadata.is_expired());
271    }
272
273    #[test]
274    fn test_current_key() {
275        let manager = KeyManager::new();
276        let key_id = manager
277            .generate_key(EncryptionAlgorithm::Aes256Gcm, Some(365))
278            .expect("Failed to generate key");
279
280        let (current_id, _, _) = manager
281            .get_current_key()
282            .expect("Failed to get current key");
283        assert_eq!(current_id, key_id);
284    }
285
286    #[test]
287    fn test_key_rotation() {
288        let manager = KeyManager::new();
289        let old_key_id = manager
290            .generate_key(EncryptionAlgorithm::Aes256Gcm, Some(365))
291            .expect("Failed to generate key");
292
293        let new_key_id = manager.rotate_key().expect("Failed to rotate key");
294        assert_ne!(old_key_id, new_key_id);
295
296        let (current_id, _, _) = manager
297            .get_current_key()
298            .expect("Failed to get current key");
299        assert_eq!(current_id, new_key_id);
300
301        let (_, old_metadata) = manager.get_key(&old_key_id).expect("Failed to get old key");
302        assert!(!old_metadata.active);
303    }
304
305    #[test]
306    fn test_list_keys() {
307        let manager = KeyManager::new();
308        manager
309            .generate_key(EncryptionAlgorithm::Aes256Gcm, Some(365))
310            .expect("Failed to generate key");
311        manager
312            .generate_key(EncryptionAlgorithm::ChaCha20Poly1305, Some(365))
313            .expect("Failed to generate key");
314
315        let keys = manager.list_keys();
316        assert_eq!(keys.len(), 2);
317    }
318
319    #[test]
320    fn test_delete_key() {
321        let manager = KeyManager::new();
322        let key_id1 = manager
323            .generate_key(EncryptionAlgorithm::Aes256Gcm, Some(365))
324            .expect("Failed to generate key");
325        let key_id2 = manager
326            .generate_key(EncryptionAlgorithm::Aes256Gcm, Some(365))
327            .expect("Failed to generate key");
328
329        // Cannot delete current key
330        assert!(manager.delete_key(&key_id1).is_err());
331
332        // Can delete non-current key
333        assert!(manager.delete_key(&key_id2).is_ok());
334        assert_eq!(manager.key_count(), 1);
335    }
336
337    #[test]
338    fn test_create_encryptor() {
339        let manager = KeyManager::new();
340        let key_id = manager
341            .generate_key(EncryptionAlgorithm::Aes256Gcm, Some(365))
342            .expect("Failed to generate key");
343
344        let encryptor = manager
345            .create_encryptor(&key_id)
346            .expect("Failed to create encryptor");
347
348        assert_eq!(encryptor.algorithm(), EncryptionAlgorithm::Aes256Gcm);
349        assert_eq!(encryptor.key_id(), key_id);
350    }
351}