kv_core/
encryption.rs

1//! Encryption layer for the KV service
2//! 
3//! Provides AES-256-GCM encryption for data at rest with key derivation
4//! using HKDF and secure key management.
5
6use aes_gcm::{
7    aead::{Aead, KeyInit},
8    Aes256Gcm, Nonce, Key,
9};
10use hkdf::Hkdf;
11use sha2::Sha256;
12use base64::{Engine as _, engine::general_purpose};
13use getrandom;
14use std::collections::HashMap;
15
16use crate::{KVError, KVResult, DatabaseId};
17
18/// Encryption key for a specific database
19#[derive(Debug, Clone)]
20pub struct DatabaseKey {
21    pub database_id: DatabaseId,
22    pub key: Key<Aes256Gcm>,
23    pub nonce: [u8; 12], // GCM nonce size
24}
25
26impl DatabaseKey {
27    /// Generate a new database key from master key
28    /// 
29    /// # Errors
30    /// Returns error if master key is too short or key derivation fails
31    pub fn derive(master_key: &[u8], database_id: DatabaseId) -> KVResult<Self> {
32        if master_key.len() < 32 {
33            return Err(KVError::Encryption("Master key must be at least 32 bytes".to_string()));
34        }
35
36        // Use HKDF to derive database-specific key
37        let hk = Hkdf::<Sha256>::new(None, master_key);
38        let mut derived_key = [0u8; 32];
39        let info = format!("kv-database-{database_id}");
40        hk.expand(info.as_bytes(), &mut derived_key)
41            .map_err(|e| KVError::Encryption(format!("Key derivation failed: {e}")))?;
42
43        // Generate random nonce for this database
44        let mut nonce = [0u8; 12];
45        getrandom::fill(&mut nonce).map_err(|e| KVError::Encryption(format!("Failed to generate random nonce: {e}")))?;
46
47        let key = Key::<Aes256Gcm>::from_slice(&derived_key);
48        
49        Ok(Self {
50            database_id,
51            key: *key,
52            nonce,
53        })
54    }
55
56    /// Encrypt data using this database key
57    /// 
58    /// # Errors
59    /// Returns error if encryption fails
60    pub fn encrypt(&self, data: &[u8]) -> KVResult<Vec<u8>> {
61        let cipher = Aes256Gcm::new(&self.key);
62        let nonce = Nonce::from_slice(&self.nonce);
63        
64        cipher.encrypt(nonce, data)
65            .map_err(|e| KVError::Encryption(format!("Encryption failed: {e}")))
66    }
67
68    /// Decrypt data using this database key
69    /// 
70    /// # Errors
71    /// Returns error if decryption fails
72    pub fn decrypt(&self, encrypted_data: &[u8]) -> KVResult<Vec<u8>> {
73        let cipher = Aes256Gcm::new(&self.key);
74        let nonce = Nonce::from_slice(&self.nonce);
75        
76        cipher.decrypt(nonce, encrypted_data)
77            .map_err(|e| KVError::Encryption(format!("Decryption failed: {e}")))
78    }
79}
80
81/// Key manager for handling encryption keys
82pub struct KeyManager {
83    master_key: Vec<u8>,
84    database_keys: HashMap<DatabaseId, DatabaseKey>,
85}
86
87impl KeyManager {
88    /// Create a new key manager with master key
89    /// 
90    /// # Errors
91    /// Returns error if key generation or decoding fails
92    pub fn new(master_key: &str) -> KVResult<Self> {
93        let master_key = if master_key.is_empty() {
94            // Generate a new master key
95            let mut key = [0u8; 32];
96            getrandom::fill(&mut key).map_err(|e| KVError::Encryption(format!("Failed to generate random key: {e}")))?;
97            key.to_vec()
98        } else {
99            // Decode base64 master key
100            general_purpose::STANDARD.decode(master_key)
101                .map_err(|e| KVError::Encryption(format!("Invalid base64 master key: {e}")))?
102        };
103
104        Ok(Self {
105            master_key,
106            database_keys: HashMap::new(),
107        })
108    }
109
110    /// Get or create encryption key for a database
111    /// 
112    /// # Errors
113    /// Returns error if key derivation fails
114    pub fn get_database_key(&mut self, database_id: DatabaseId) -> KVResult<&DatabaseKey> {
115        if !self.database_keys.contains_key(&database_id) {
116            let db_key = DatabaseKey::derive(&self.master_key, database_id)?;
117            self.database_keys.insert(database_id, db_key);
118        }
119        
120        self.database_keys.get(&database_id)
121            .ok_or_else(|| KVError::Encryption("Failed to get database key".to_string()))
122    }
123
124    /// Get master key as base64 string (for configuration)
125    #[must_use] 
126    pub fn master_key_base64(&self) -> String {
127        general_purpose::STANDARD.encode(&self.master_key)
128    }
129
130    /// Rotate master key (generates new key and invalidates all database keys)
131    /// 
132    /// # Errors
133    /// Returns error if new key generation fails
134    pub fn rotate_master_key(&mut self) -> KVResult<String> {
135        // Generate new master key
136        let mut new_master_key = [0u8; 32];
137        getrandom::fill(&mut new_master_key).map_err(|e| KVError::Encryption(format!("Failed to generate random key: {e}")))?;
138        
139        // Update master key
140        self.master_key = new_master_key.to_vec();
141        
142        // Clear all database keys (they will be regenerated on next access)
143        self.database_keys.clear();
144        
145        Ok(self.master_key_base64())
146    }
147
148    /// Encrypt data for a specific database
149    /// 
150    /// # Errors
151    /// Returns error if encryption fails
152    pub fn encrypt(&mut self, database_id: DatabaseId, data: &[u8]) -> KVResult<Vec<u8>> {
153        let db_key = self.get_database_key(database_id)?;
154        db_key.encrypt(data)
155    }
156
157    /// Decrypt data for a specific database
158    /// 
159    /// # Errors
160    /// Returns error if decryption fails
161    pub fn decrypt(&mut self, database_id: DatabaseId, encrypted_data: &[u8]) -> KVResult<Vec<u8>> {
162        let db_key = self.get_database_key(database_id)?;
163        db_key.decrypt(encrypted_data)
164    }
165}
166
167/// Encrypted storage wrapper
168#[derive(Debug, Clone)]
169pub struct EncryptedData {
170    pub database_id: DatabaseId,
171    pub encrypted_data: Vec<u8>,
172    pub nonce: [u8; 12],
173}
174
175impl EncryptedData {
176    /// Create encrypted data from raw data
177    /// 
178    /// # Errors
179    /// Returns error if encryption fails
180    pub fn encrypt(data: &[u8], database_id: DatabaseId, key_manager: &mut KeyManager) -> KVResult<Self> {
181        let encrypted_data = key_manager.encrypt(database_id, data)?;
182        
183        Ok(Self {
184            database_id,
185            encrypted_data,
186            nonce: [0u8; 12], // Will be set by the database key
187        })
188    }
189
190    /// Decrypt this data
191    /// 
192    /// # Errors
193    /// Returns error if decryption fails
194    pub fn decrypt(&self, key_manager: &mut KeyManager) -> KVResult<Vec<u8>> {
195        key_manager.decrypt(self.database_id, &self.encrypted_data)
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202
203    #[test]
204    fn test_key_derivation() {
205        // Create a 32-byte master key
206        let master_key = b"0123456789abcdef0123456789abcdef"; // 32 bytes
207        let db_key1 = DatabaseKey::derive(master_key, 0).unwrap();
208        let db_key2 = DatabaseKey::derive(master_key, 1).unwrap();
209        
210        // Different databases should have different keys
211        assert_ne!(db_key1.key, db_key2.key);
212    }
213
214    #[test]
215    fn test_encryption_decryption() {
216        // Use empty string to generate a random master key
217        let mut key_manager = KeyManager::new("").unwrap();
218        
219        let data = b"Hello, encrypted world!";
220        let encrypted = key_manager.encrypt(0, data).unwrap();
221        let decrypted = key_manager.decrypt(0, &encrypted).unwrap();
222        
223        assert_eq!(data, &decrypted[..]);
224    }
225
226    #[test]
227    fn test_key_rotation() {
228        let mut key_manager = KeyManager::new("").unwrap();
229        let old_master_key = key_manager.master_key_base64();
230        
231        let new_master_key = key_manager.rotate_master_key().unwrap();
232        
233        assert_ne!(old_master_key, new_master_key);
234        assert!(key_manager.database_keys.is_empty());
235    }
236
237    #[test]
238    fn test_encrypted_data_wrapper() {
239        let mut key_manager = KeyManager::new("").unwrap();
240        let data = b"Test data for encryption wrapper";
241        
242        let encrypted = EncryptedData::encrypt(data, 0, &mut key_manager).unwrap();
243        let decrypted = encrypted.decrypt(&mut key_manager).unwrap();
244        
245        assert_eq!(data, &decrypted[..]);
246    }
247}
kv_core/encryption.rs

kv_core/
encryption.rs
