Skip to main content

kaccy_db/security/
encryption.rs

1use aes_gcm::{
2    aead::{Aead, KeyInit, OsRng},
3    Aes256Gcm, Key, Nonce,
4};
5use argon2::password_hash::{rand_core::RngCore, SaltString};
6use argon2::{Argon2, PasswordHash, PasswordHasher as Argon2Hasher, PasswordVerifier};
7use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
8use std::collections::HashMap;
9use std::sync::RwLock;
10
11/// Encryption provider for column-level encryption
12#[derive(Debug)]
13pub struct EncryptionProvider {
14    /// Active encryption keys (key_id -> key)
15    keys: RwLock<HashMap<String, Vec<u8>>>,
16    /// Current active key ID
17    active_key_id: RwLock<String>,
18}
19
20impl EncryptionProvider {
21    /// Create a new encryption provider with a master key
22    pub fn new(master_key: &[u8]) -> Result<Self, EncryptionError> {
23        if master_key.len() != 32 {
24            return Err(EncryptionError::InvalidKeySize {
25                expected: 32,
26                actual: master_key.len(),
27            });
28        }
29
30        let mut keys = HashMap::new();
31        let key_id = "master".to_string();
32        keys.insert(key_id.clone(), master_key.to_vec());
33
34        Ok(Self {
35            keys: RwLock::new(keys),
36            active_key_id: RwLock::new(key_id),
37        })
38    }
39
40    /// Generate a new random encryption key
41    pub fn generate_key() -> Vec<u8> {
42        let mut key = vec![0u8; 32];
43        OsRng.fill_bytes(&mut key);
44        key
45    }
46
47    /// Add a new encryption key for key rotation
48    pub fn add_key(&self, key_id: String, key: Vec<u8>) -> Result<(), EncryptionError> {
49        if key.len() != 32 {
50            return Err(EncryptionError::InvalidKeySize {
51                expected: 32,
52                actual: key.len(),
53            });
54        }
55
56        let mut keys = self.keys.write().map_err(|_| EncryptionError::LockError)?;
57        keys.insert(key_id, key);
58        Ok(())
59    }
60
61    /// Set the active key ID for new encryptions
62    pub fn set_active_key(&self, key_id: String) -> Result<(), EncryptionError> {
63        let keys = self.keys.read().map_err(|_| EncryptionError::LockError)?;
64        if !keys.contains_key(&key_id) {
65            return Err(EncryptionError::KeyNotFound(key_id));
66        }
67
68        let mut active = self
69            .active_key_id
70            .write()
71            .map_err(|_| EncryptionError::LockError)?;
72        *active = key_id;
73        Ok(())
74    }
75
76    /// Encrypt data using the active key
77    pub fn encrypt(&self, plaintext: &[u8]) -> Result<EncryptedData, EncryptionError> {
78        let active_key_id = self
79            .active_key_id
80            .read()
81            .map_err(|_| EncryptionError::LockError)?;
82        self.encrypt_with_key(&active_key_id, plaintext)
83    }
84
85    /// Encrypt data with a specific key
86    pub fn encrypt_with_key(
87        &self,
88        key_id: &str,
89        plaintext: &[u8],
90    ) -> Result<EncryptedData, EncryptionError> {
91        let keys = self.keys.read().map_err(|_| EncryptionError::LockError)?;
92        let key = keys
93            .get(key_id)
94            .ok_or_else(|| EncryptionError::KeyNotFound(key_id.to_string()))?;
95
96        // Generate random nonce
97        let mut nonce_bytes = [0u8; 12];
98        OsRng.fill_bytes(&mut nonce_bytes);
99        let nonce = Nonce::from_slice(&nonce_bytes);
100
101        // Encrypt
102        let key_obj = Key::<Aes256Gcm>::from_slice(key);
103        let cipher = Aes256Gcm::new(key_obj);
104        let ciphertext = cipher
105            .encrypt(nonce, plaintext)
106            .map_err(|e| EncryptionError::EncryptionFailed(e.to_string()))?;
107
108        Ok(EncryptedData {
109            key_id: key_id.to_string(),
110            nonce: nonce_bytes.to_vec(),
111            ciphertext,
112        })
113    }
114
115    /// Decrypt data
116    pub fn decrypt(&self, encrypted: &EncryptedData) -> Result<Vec<u8>, EncryptionError> {
117        let keys = self.keys.read().map_err(|_| EncryptionError::LockError)?;
118        let key = keys
119            .get(&encrypted.key_id)
120            .ok_or_else(|| EncryptionError::KeyNotFound(encrypted.key_id.clone()))?;
121
122        if encrypted.nonce.len() != 12 {
123            return Err(EncryptionError::InvalidNonceSize {
124                expected: 12,
125                actual: encrypted.nonce.len(),
126            });
127        }
128
129        let nonce = Nonce::from_slice(&encrypted.nonce);
130        let key_obj = Key::<Aes256Gcm>::from_slice(key);
131        let cipher = Aes256Gcm::new(key_obj);
132
133        cipher
134            .decrypt(nonce, encrypted.ciphertext.as_ref())
135            .map_err(|e| EncryptionError::DecryptionFailed(e.to_string()))
136    }
137
138    /// Encrypt a string value
139    pub fn encrypt_string(&self, plaintext: &str) -> Result<String, EncryptionError> {
140        let encrypted = self.encrypt(plaintext.as_bytes())?;
141        Ok(encrypted.to_base64())
142    }
143
144    /// Decrypt a string value
145    pub fn decrypt_string(&self, ciphertext: &str) -> Result<String, EncryptionError> {
146        let encrypted = EncryptedData::from_base64(ciphertext)?;
147        let plaintext = self.decrypt(&encrypted)?;
148        String::from_utf8(plaintext).map_err(|e| EncryptionError::InvalidUtf8(e.to_string()))
149    }
150
151    /// Re-encrypt data with a new key (for key rotation)
152    pub fn reencrypt(
153        &self,
154        encrypted: &EncryptedData,
155        new_key_id: &str,
156    ) -> Result<EncryptedData, EncryptionError> {
157        let plaintext = self.decrypt(encrypted)?;
158        self.encrypt_with_key(new_key_id, &plaintext)
159    }
160}
161
162/// Encrypted data structure
163#[derive(Debug, Clone)]
164pub struct EncryptedData {
165    /// Key ID used for encryption
166    pub key_id: String,
167    /// Nonce used for encryption
168    pub nonce: Vec<u8>,
169    /// Encrypted ciphertext
170    pub ciphertext: Vec<u8>,
171}
172
173impl EncryptedData {
174    /// Serialize to base64-encoded string format: key_id:nonce:ciphertext
175    pub fn to_base64(&self) -> String {
176        format!(
177            "{}:{}:{}",
178            self.key_id,
179            BASE64.encode(&self.nonce),
180            BASE64.encode(&self.ciphertext)
181        )
182    }
183
184    /// Deserialize from base64-encoded string
185    pub fn from_base64(encoded: &str) -> Result<Self, EncryptionError> {
186        let parts: Vec<&str> = encoded.split(':').collect();
187        if parts.len() != 3 {
188            return Err(EncryptionError::InvalidFormat);
189        }
190
191        let key_id = parts[0].to_string();
192        let nonce = BASE64
193            .decode(parts[1])
194            .map_err(|e| EncryptionError::Base64Error(e.to_string()))?;
195        let ciphertext = BASE64
196            .decode(parts[2])
197            .map_err(|e| EncryptionError::Base64Error(e.to_string()))?;
198
199        Ok(Self {
200            key_id,
201            nonce,
202            ciphertext,
203        })
204    }
205}
206
207/// Key rotation manager
208#[derive(Debug)]
209pub struct KeyRotationManager {
210    provider: EncryptionProvider,
211}
212
213impl KeyRotationManager {
214    /// Create a new key rotation manager
215    pub fn new(provider: EncryptionProvider) -> Self {
216        Self { provider }
217    }
218
219    /// Rotate to a new key
220    pub fn rotate_key(&self, new_key_id: String, new_key: Vec<u8>) -> Result<(), EncryptionError> {
221        // Add the new key
222        self.provider.add_key(new_key_id.clone(), new_key)?;
223
224        // Set it as active
225        self.provider.set_active_key(new_key_id)?;
226
227        Ok(())
228    }
229
230    /// Re-encrypt data with the current active key
231    pub fn reencrypt_data(
232        &self,
233        old_encrypted: &EncryptedData,
234    ) -> Result<EncryptedData, EncryptionError> {
235        let active_key_id = self
236            .provider
237            .active_key_id
238            .read()
239            .map_err(|_| EncryptionError::LockError)?;
240        self.provider.reencrypt(old_encrypted, &active_key_id)
241    }
242}
243
244/// Password hashing service using Argon2
245#[derive(Debug, Clone)]
246pub struct PasswordHashingService {
247    argon2: Argon2<'static>,
248}
249
250impl Default for PasswordHashingService {
251    fn default() -> Self {
252        Self::new()
253    }
254}
255
256impl PasswordHashingService {
257    /// Create a new password hasher
258    pub fn new() -> Self {
259        Self {
260            argon2: Argon2::default(),
261        }
262    }
263
264    /// Hash a password
265    pub fn hash_password(&self, password: &str) -> Result<String, EncryptionError> {
266        let salt = SaltString::generate(&mut OsRng);
267
268        let password_hash = self
269            .argon2
270            .hash_password(password.as_bytes(), &salt)
271            .map_err(|e| EncryptionError::HashingFailed(e.to_string()))?;
272
273        Ok(password_hash.to_string())
274    }
275
276    /// Verify a password against a hash
277    pub fn verify_password(&self, password: &str, hash: &str) -> Result<bool, EncryptionError> {
278        let parsed_hash =
279            PasswordHash::new(hash).map_err(|e| EncryptionError::InvalidHash(e.to_string()))?;
280
281        match self
282            .argon2
283            .verify_password(password.as_bytes(), &parsed_hash)
284        {
285            Ok(()) => Ok(true),
286            Err(argon2::password_hash::Error::Password) => Ok(false),
287            Err(e) => Err(EncryptionError::VerificationFailed(e.to_string())),
288        }
289    }
290}
291
292/// Errors that can occur during encryption operations
293#[derive(Debug, Clone, thiserror::Error)]
294pub enum EncryptionError {
295    /// The provided key has the wrong byte length.
296    #[error("Invalid key size: expected {expected}, got {actual}")]
297    InvalidKeySize {
298        /// Expected key length in bytes.
299        expected: usize,
300        /// Actual key length provided.
301        actual: usize,
302    },
303
304    /// The nonce provided for decryption has the wrong length.
305    #[error("Invalid nonce size: expected {expected}, got {actual}")]
306    InvalidNonceSize {
307        /// Expected nonce length in bytes.
308        expected: usize,
309        /// Actual nonce length found.
310        actual: usize,
311    },
312
313    /// The requested key ID does not exist in the key store.
314    #[error("Key not found: {0}")]
315    KeyNotFound(String),
316
317    /// The AEAD encryption operation failed.
318    #[error("Encryption failed: {0}")]
319    EncryptionFailed(String),
320
321    /// The AEAD decryption operation failed (e.g., authentication tag mismatch).
322    #[error("Decryption failed: {0}")]
323    DecryptionFailed(String),
324
325    /// Decrypted bytes are not valid UTF-8.
326    #[error("Invalid UTF-8: {0}")]
327    InvalidUtf8(String),
328
329    /// The serialized encrypted-data string has an invalid format.
330    #[error("Invalid format")]
331    InvalidFormat,
332
333    /// Base64 decoding failed.
334    #[error("Base64 error: {0}")]
335    Base64Error(String),
336
337    /// A read/write lock was poisoned.
338    #[error("Lock error")]
339    LockError,
340
341    /// Argon2 password hashing failed.
342    #[error("Hashing failed: {0}")]
343    HashingFailed(String),
344
345    /// The stored password hash string is malformed.
346    #[error("Invalid hash: {0}")]
347    InvalidHash(String),
348
349    /// Password verification encountered an unexpected error.
350    #[error("Verification failed: {0}")]
351    VerificationFailed(String),
352}
353
354#[cfg(test)]
355mod tests {
356    use super::*;
357
358    #[test]
359    fn test_generate_key() {
360        let key = EncryptionProvider::generate_key();
361        assert_eq!(key.len(), 32);
362    }
363
364    #[test]
365    fn test_encrypt_decrypt() {
366        let key = EncryptionProvider::generate_key();
367        let provider = EncryptionProvider::new(&key).unwrap();
368
369        let plaintext = b"Hello, World!";
370        let encrypted = provider.encrypt(plaintext).unwrap();
371        let decrypted = provider.decrypt(&encrypted).unwrap();
372
373        assert_eq!(plaintext, decrypted.as_slice());
374    }
375
376    #[test]
377    fn test_encrypt_decrypt_string() {
378        let key = EncryptionProvider::generate_key();
379        let provider = EncryptionProvider::new(&key).unwrap();
380
381        let plaintext = "Hello, World!";
382        let encrypted = provider.encrypt_string(plaintext).unwrap();
383        let decrypted = provider.decrypt_string(&encrypted).unwrap();
384
385        assert_eq!(plaintext, decrypted);
386    }
387
388    #[test]
389    fn test_encrypted_data_base64() {
390        let data = EncryptedData {
391            key_id: "test".to_string(),
392            nonce: vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
393            ciphertext: vec![13, 14, 15, 16],
394        };
395
396        let encoded = data.to_base64();
397        let decoded = EncryptedData::from_base64(&encoded).unwrap();
398
399        assert_eq!(data.key_id, decoded.key_id);
400        assert_eq!(data.nonce, decoded.nonce);
401        assert_eq!(data.ciphertext, decoded.ciphertext);
402    }
403
404    #[test]
405    fn test_key_rotation() {
406        let key1 = EncryptionProvider::generate_key();
407        let provider = EncryptionProvider::new(&key1).unwrap();
408
409        let plaintext = b"Secret data";
410        let encrypted1 = provider.encrypt(plaintext).unwrap();
411
412        // Add a new key and rotate
413        let key2 = EncryptionProvider::generate_key();
414        provider.add_key("key2".to_string(), key2).unwrap();
415        provider.set_active_key("key2".to_string()).unwrap();
416
417        // Re-encrypt with new key
418        let encrypted2 = provider.reencrypt(&encrypted1, "key2").unwrap();
419
420        // Verify both can be decrypted
421        let decrypted1 = provider.decrypt(&encrypted1).unwrap();
422        let decrypted2 = provider.decrypt(&encrypted2).unwrap();
423
424        assert_eq!(plaintext, decrypted1.as_slice());
425        assert_eq!(plaintext, decrypted2.as_slice());
426        assert_eq!(encrypted2.key_id, "key2");
427    }
428
429    #[test]
430    fn test_invalid_key_size() {
431        let short_key = vec![0u8; 16]; // Too short
432        let result = EncryptionProvider::new(&short_key);
433        assert!(result.is_err());
434    }
435
436    #[test]
437    fn test_key_not_found() {
438        let key = EncryptionProvider::generate_key();
439        let provider = EncryptionProvider::new(&key).unwrap();
440
441        let result = provider.encrypt_with_key("nonexistent", b"data");
442        assert!(result.is_err());
443    }
444
445    #[test]
446    fn test_password_hasher() {
447        let hasher = PasswordHashingService::new();
448        let password = "my_secure_password";
449
450        let hash = hasher.hash_password(password).unwrap();
451        assert!(hasher.verify_password(password, &hash).unwrap());
452        assert!(!hasher.verify_password("wrong_password", &hash).unwrap());
453    }
454
455    #[test]
456    fn test_password_hashing_produces_different_hashes() {
457        let hasher = PasswordHashingService::new();
458        let password = "same_password";
459
460        let hash1 = hasher.hash_password(password).unwrap();
461        let hash2 = hasher.hash_password(password).unwrap();
462
463        // Different salts should produce different hashes
464        assert_ne!(hash1, hash2);
465
466        // But both should verify correctly
467        assert!(hasher.verify_password(password, &hash1).unwrap());
468        assert!(hasher.verify_password(password, &hash2).unwrap());
469    }
470
471    #[test]
472    fn test_key_rotation_manager() {
473        let key1 = EncryptionProvider::generate_key();
474        let provider = EncryptionProvider::new(&key1).unwrap();
475        let manager = KeyRotationManager::new(provider);
476
477        let plaintext = b"Test data";
478        let encrypted = manager.provider.encrypt(plaintext).unwrap();
479
480        // Rotate key
481        let key2 = EncryptionProvider::generate_key();
482        manager.rotate_key("new_key".to_string(), key2).unwrap();
483
484        // Re-encrypt with new key
485        let reencrypted = manager.reencrypt_data(&encrypted).unwrap();
486        let decrypted = manager.provider.decrypt(&reencrypted).unwrap();
487
488        assert_eq!(plaintext, decrypted.as_slice());
489        assert_eq!(reencrypted.key_id, "new_key");
490    }
491}