rust-memory-service/src/encryption/mod.rs

//! Encryption Module
//!
//! Per-user encryption using AES-256-GCM.

use aes_gcm::{
    aead::{Aead, KeyInit, OsRng},
    Aes256Gcm, Nonce,
};
use sha2::{Sha256, Digest};
use rand::RngCore;
use thiserror::Error;

/// Encryption errors
#[derive(Debug, Error)]
pub enum EncryptionError {
    #[error("Encryption failed: {0}")]
    Encryption(String),

    #[error("Decryption failed: {0}")]
    Decryption(String),

    #[error("Invalid key: {0}")]
    InvalidKey(String),

    #[error("Invalid ciphertext")]
    InvalidCiphertext,
}

/// Per-user encryptor using AES-256-GCM
pub struct UserEncryptor {
    cipher: Aes256Gcm,
    user_id: String,
}

impl UserEncryptor {
    /// Create a new encryptor for a user
    ///
    /// Derives a user-specific key using HKDF-SHA256(master_key, user_id)
    pub fn new(user_id: &str, master_key: &[u8]) -> Result<Self, EncryptionError> {
        if master_key.len() < 32 {
            return Err(EncryptionError::InvalidKey(
                "Master key must be at least 32 bytes".to_string(),
            ));
        }

        // Derive user-specific key
        let mut hasher = Sha256::new();
        hasher.update(master_key);
        hasher.update(user_id.as_bytes());
        let key = hasher.finalize();

        let cipher = Aes256Gcm::new_from_slice(&key)
            .map_err(|e| EncryptionError::InvalidKey(e.to_string()))?;

        Ok(Self {
            cipher,
            user_id: user_id.to_string(),
        })
    }

    /// Encrypt data
    ///
    /// Returns nonce + ciphertext
    pub fn encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>, EncryptionError> {
        // Generate random nonce
        let mut nonce_bytes = [0u8; 12];
        OsRng.fill_bytes(&mut nonce_bytes);
        let nonce = Nonce::from_slice(&nonce_bytes);

        // Encrypt
        let ciphertext = self
            .cipher
            .encrypt(nonce, plaintext)
            .map_err(|e| EncryptionError::Encryption(e.to_string()))?;

        // Prepend nonce to ciphertext
        let mut result = Vec::with_capacity(12 + ciphertext.len());
        result.extend_from_slice(&nonce_bytes);
        result.extend_from_slice(&ciphertext);

        Ok(result)
    }

    /// Decrypt data
    ///
    /// Expects nonce + ciphertext format
    pub fn decrypt(&self, ciphertext: &[u8]) -> Result<Vec<u8>, EncryptionError> {
        if ciphertext.len() < 12 {
            return Err(EncryptionError::InvalidCiphertext);
        }

        // Extract nonce
        let (nonce_bytes, encrypted) = ciphertext.split_at(12);
        let nonce = Nonce::from_slice(nonce_bytes);

        // Decrypt
        self.cipher
            .decrypt(nonce, encrypted)
            .map_err(|e| EncryptionError::Decryption(e.to_string()))
    }

    /// Get the user ID this encryptor is for
    pub fn user_id(&self) -> &str {
        &self.user_id
    }
}

/// Encryption manager for managing per-user encryptors
pub struct EncryptionManager {
    master_key: Vec<u8>,
}

impl EncryptionManager {
    /// Create a new encryption manager
    pub fn new(master_key: Vec<u8>) -> Result<Self, EncryptionError> {
        if master_key.len() < 32 {
            return Err(EncryptionError::InvalidKey(
                "Master key must be at least 32 bytes".to_string(),
            ));
        }

        Ok(Self { master_key })
    }

    /// Create a new encryption manager from hex-encoded key
    pub fn from_hex(hex_key: &str) -> Result<Self, EncryptionError> {
        let key = hex::decode(hex_key)
            .map_err(|e| EncryptionError::InvalidKey(e.to_string()))?;
        Self::new(key)
    }

    /// Get an encryptor for a specific user
    pub fn for_user(&self, user_id: &str) -> Result<UserEncryptor, EncryptionError> {
        UserEncryptor::new(user_id, &self.master_key)
    }

    /// Encrypt data for a user
    pub fn encrypt(&self, user_id: &str, plaintext: &[u8]) -> Result<Vec<u8>, EncryptionError> {
        let encryptor = self.for_user(user_id)?;
        encryptor.encrypt(plaintext)
    }

    /// Decrypt data for a user
    pub fn decrypt(&self, user_id: &str, ciphertext: &[u8]) -> Result<Vec<u8>, EncryptionError> {
        let encryptor = self.for_user(user_id)?;
        encryptor.decrypt(ciphertext)
    }
}

/// Generate a random 32-byte master key
pub fn generate_master_key() -> Vec<u8> {
    let mut key = vec![0u8; 32];
    OsRng.fill_bytes(&mut key);
    key
}

/// Generate a master key and return as hex string
pub fn generate_master_key_hex() -> String {
    hex::encode(generate_master_key())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_encrypt_decrypt() {
        let master_key = generate_master_key();
        let encryptor = UserEncryptor::new("user-1", &master_key).unwrap();

        let plaintext = b"Hello, World!";
        let ciphertext = encryptor.encrypt(plaintext).unwrap();
        let decrypted = encryptor.decrypt(&ciphertext).unwrap();

        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_different_users_different_keys() {
        let master_key = generate_master_key();
        let enc1 = UserEncryptor::new("user-1", &master_key).unwrap();
        let enc2 = UserEncryptor::new("user-2", &master_key).unwrap();

        let plaintext = b"Secret data";
        let ciphertext = enc1.encrypt(plaintext).unwrap();

        // User 2 should not be able to decrypt User 1's data
        assert!(enc2.decrypt(&ciphertext).is_err());
    }

    #[test]
    fn test_encryption_manager() {
        let manager = EncryptionManager::new(generate_master_key()).unwrap();

        let plaintext = b"Sensitive information";
        let ciphertext = manager.encrypt("user-1", plaintext).unwrap();
        let decrypted = manager.decrypt("user-1", &ciphertext).unwrap();

        assert_eq!(decrypted, plaintext);
    }
}