Rust Security Best Practices¶

// rust/01-input-validation.rs

/*
Input validation patterns and best practices
*/

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use serde::{Deserialize, Serialize};
use regex::Regex;

/// Input validator.
pub struct InputValidator {
    patterns: HashMap<String, Regex>,
    max_lengths: HashMap<String, usize>,
    allowed_chars: HashMap<String, Regex>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationResult {
    pub is_valid: bool,
    pub errors: Vec<String>,
    pub sanitized_value: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationRule {
    pub field_name: String,
    pub required: bool,
    pub min_length: Option<usize>,
    pub max_length: Option<usize>,
    pub pattern: Option<String>,
    pub allowed_values: Option<Vec<String>>,
    pub custom_validator: Option<String>,
}

impl InputValidator {
    pub fn new() -> Self {
        let mut patterns = HashMap::new();
        let mut max_lengths = HashMap::new();
        let mut allowed_chars = HashMap::new();

        // Email pattern
        patterns.insert("email".to_string(), Regex::new(r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$").unwrap());

        // Phone pattern
        patterns.insert("phone".to_string(), Regex::new(r"^\+?[1-9]\d{1,14}$").unwrap());

        // URL pattern
        patterns.insert("url".to_string(), Regex::new(r"^https?://[^\s/$.?#].[^\s]*$").unwrap());

        // Alphanumeric pattern
        patterns.insert("alphanumeric".to_string(), Regex::new(r"^[a-zA-Z0-9]+$").unwrap());

        // Set max lengths
        max_lengths.insert("email".to_string(), 254);
        max_lengths.insert("phone".to_string(), 15);
        max_lengths.insert("url".to_string(), 2048);
        max_lengths.insert("name".to_string(), 100);
        max_lengths.insert("description".to_string(), 1000);

        // Set allowed characters
        allowed_chars.insert("name".to_string(), Regex::new(r"^[a-zA-Z\s\-'\.]+$").unwrap());
        allowed_chars.insert("description".to_string(), Regex::new(r"^[a-zA-Z0-9\s\-'\.\,\!\?]+$").unwrap());

        Self {
            patterns,
            max_lengths,
            allowed_chars,
        }
    }

    /// Validate a single field.
    pub fn validate_field(&self, field_name: &str, value: &str, rules: &ValidationRule) -> ValidationResult {
        let mut errors = Vec::new();
        let mut sanitized_value = Some(value.to_string());

        // Check if required
        if rules.required && value.trim().is_empty() {
            errors.push(format!("Field '{}' is required", field_name));
            return ValidationResult {
                is_valid: false,
                errors,
                sanitized_value: None,
            };
        }

        // Skip validation for empty optional fields
        if !rules.required && value.trim().is_empty() {
            return ValidationResult {
                is_valid: true,
                errors: Vec::new(),
                sanitized_value: Some(String::new()),
            };
        }

        let trimmed_value = value.trim();

        // Check minimum length
        if let Some(min_length) = rules.min_length {
            if trimmed_value.len() < min_length {
                errors.push(format!("Field '{}' must be at least {} characters long", field_name, min_length));
            }
        }

        // Check maximum length
        if let Some(max_length) = rules.max_length {
            if trimmed_value.len() > max_length {
                errors.push(format!("Field '{}' must be no more than {} characters long", field_name, max_length));
            }
        }

        // Check pattern
        if let Some(pattern) = &rules.pattern {
            if let Ok(regex) = Regex::new(pattern) {
                if !regex.is_match(trimmed_value) {
                    errors.push(format!("Field '{}' does not match required pattern", field_name));
                }
            }
        }

        // Check allowed values
        if let Some(allowed_values) = &rules.allowed_values {
            if !allowed_values.contains(&trimmed_value.to_string()) {
                errors.push(format!("Field '{}' must be one of: {}", field_name, allowed_values.join(", ")));
            }
        }

        // Check against predefined patterns
        if let Some(pattern_name) = self.get_pattern_name(field_name) {
            if let Some(pattern) = self.patterns.get(pattern_name) {
                if !pattern.is_match(trimmed_value) {
                    errors.push(format!("Field '{}' does not match required format", field_name));
                }
            }
        }

        // Check against max lengths
        if let Some(max_length) = self.max_lengths.get(field_name) {
            if trimmed_value.len() > *max_length {
                errors.push(format!("Field '{}' exceeds maximum length of {}", field_name, max_length));
            }
        }

        // Check allowed characters
        if let Some(allowed_chars) = self.allowed_chars.get(field_name) {
            if !allowed_chars.is_match(trimmed_value) {
                errors.push(format!("Field '{}' contains invalid characters", field_name));
            }
        }

        // Sanitize value
        if errors.is_empty() {
            sanitized_value = Some(self.sanitize_value(trimmed_value));
        }

        ValidationResult {
            is_valid: errors.is_empty(),
            errors,
            sanitized_value,
        }
    }

    /// Validate multiple fields.
    pub fn validate_fields(&self, data: &HashMap<String, String>, rules: &[ValidationRule]) -> HashMap<String, ValidationResult> {
        let mut results = HashMap::new();

        for rule in rules {
            let value = data.get(&rule.field_name).map(|s| s.as_str()).unwrap_or("");
            let result = self.validate_field(&rule.field_name, value, rule);
            results.insert(rule.field_name.clone(), result);
        }

        results
    }

    /// Get pattern name for field.
    fn get_pattern_name(&self, field_name: &str) -> Option<&str> {
        match field_name {
            "email" | "email_address" => Some("email"),
            "phone" | "phone_number" => Some("phone"),
            "url" | "website" => Some("url"),
            _ => None,
        }
    }

    /// Sanitize input value.
    fn sanitize_value(&self, value: &str) -> String {
        // Remove null bytes
        let sanitized = value.replace('\0', "");

        // Trim whitespace
        let trimmed = sanitized.trim();

        // Remove control characters
        let cleaned = trimmed.chars()
            .filter(|c| !c.is_control() || *c == '\n' || *c == '\r' || *c == '\t')
            .collect::<String>();

        cleaned
    }

    /// Validate JSON input.
    pub fn validate_json(&self, json: &str) -> Result<serde_json::Value, String> {
        // Parse JSON
        let value: serde_json::Value = serde_json::from_str(json)
            .map_err(|e| format!("Invalid JSON: {}", e))?;

        // Check for dangerous keys
        if let Some(obj) = value.as_object() {
            for key in obj.keys() {
                if key.starts_with("__") || key.contains("..") || key.contains("$") {
                    return Err(format!("Dangerous key detected: {}", key));
                }
            }
        }

        // Check for dangerous values
        if let Some(str_value) = value.as_str() {
            if str_value.contains("<script>") || str_value.contains("javascript:") {
                return Err("Dangerous content detected".to_string());
            }
        }

        Ok(value)
    }

    /// Validate file upload.
    pub fn validate_file_upload(&self, filename: &str, content_type: &str, size: u64) -> Result<(), String> {
        // Check file extension
        let allowed_extensions = ["jpg", "jpeg", "png", "gif", "pdf", "txt", "doc", "docx"];
        if let Some(extension) = filename.split('.').last() {
            if !allowed_extensions.contains(&extension.to_lowercase().as_str()) {
                return Err(format!("File type '{}' not allowed", extension));
            }
        }

        // Check content type
        let allowed_types = ["image/jpeg", "image/png", "image/gif", "application/pdf", "text/plain"];
        if !allowed_types.contains(&content_type) {
            return Err(format!("Content type '{}' not allowed", content_type));
        }

        // Check file size (10MB limit)
        let max_size = 10 * 1024 * 1024;
        if size > max_size {
            return Err(format!("File size {} exceeds maximum allowed size {}", size, max_size));
        }

        // Check filename for dangerous characters
        if filename.contains("..") || filename.contains("/") || filename.contains("\\") {
            return Err("Filename contains dangerous characters".to_string());
        }

        Ok(())
    }
}

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

    #[test]
    fn test_input_validation() {
        let validator = InputValidator::new();

        let rules = vec![
            ValidationRule {
                field_name: "email".to_string(),
                required: true,
                min_length: Some(5),
                max_length: Some(254),
                pattern: None,
                allowed_values: None,
                custom_validator: None,
            },
            ValidationRule {
                field_name: "name".to_string(),
                required: true,
                min_length: Some(2),
                max_length: Some(100),
                pattern: None,
                allowed_values: None,
                custom_validator: None,
            },
        ];

        let mut data = HashMap::new();
        data.insert("email".to_string(), "test@example.com".to_string());
        data.insert("name".to_string(), "John Doe".to_string());

        let results = validator.validate_fields(&data, &rules);

        assert!(results.get("email").unwrap().is_valid);
        assert!(results.get("name").unwrap().is_valid);
    }

    #[test]
    fn test_json_validation() {
        let validator = InputValidator::new();

        let valid_json = r#"{"name": "John", "age": 30}"#;
        let result = validator.validate_json(valid_json);
        assert!(result.is_ok());

        let dangerous_json = r#"{"__proto__": {"isAdmin": true}}"#;
        let result = validator.validate_json(dangerous_json);
        assert!(result.is_err());
    }

    #[test]
    fn test_file_upload_validation() {
        let validator = InputValidator::new();

        let result = validator.validate_file_upload("test.jpg", "image/jpeg", 1024);
        assert!(result.is_ok());

        let result = validator.validate_file_upload("test.exe", "application/octet-stream", 1024);
        assert!(result.is_err());
    }
}

// rust/02-auth-authorization.rs

/*
Authentication and authorization patterns
*/

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::RwLock;
use serde::{Deserialize, Serialize};
use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation};
use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier};
use argon2::password_hash::{rand_core::OsRng, SaltString};

/// Authentication manager.
pub struct AuthManager {
    jwt_secret: String,
    jwt_expiry: Duration,
    password_hasher: Argon2<'static>,
    user_sessions: Arc<RwLock<HashMap<String, UserSession>>>,
    rate_limiter: Arc<RwLock<HashMap<String, RateLimitInfo>>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct User {
    pub id: String,
    pub username: String,
    pub email: String,
    pub password_hash: String,
    pub roles: Vec<String>,
    pub permissions: Vec<String>,
    pub is_active: bool,
    pub created_at: Instant,
    pub last_login: Option<Instant>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserSession {
    pub user_id: String,
    pub session_id: String,
    pub created_at: Instant,
    pub expires_at: Instant,
    pub ip_address: String,
    pub user_agent: String,
    pub is_active: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RateLimitInfo {
    pub attempts: u32,
    pub last_attempt: Instant,
    pub blocked_until: Option<Instant>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoginRequest {
    pub username: String,
    pub password: String,
    pub ip_address: String,
    pub user_agent: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoginResponse {
    pub token: String,
    pub refresh_token: String,
    pub expires_in: u64,
    pub user: User,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JwtClaims {
    pub sub: String, // user id
    pub username: String,
    pub roles: Vec<String>,
    pub permissions: Vec<String>,
    pub iat: u64, // issued at
    pub exp: u64, // expires at
    pub jti: String, // jwt id
}

impl AuthManager {
    pub fn new(jwt_secret: String, jwt_expiry: Duration) -> Self {
        Self {
            jwt_secret,
            jwt_expiry,
            password_hasher: Argon2::default(),
            user_sessions: Arc::new(RwLock::new(HashMap::new())),
            rate_limiter: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Register a new user.
    pub async fn register_user(&self, username: String, email: String, password: String) -> Result<User, String> {
        // Check if user already exists
        if self.user_exists(&username, &email).await {
            return Err("User already exists".to_string());
        }

        // Hash password
        let password_hash = self.hash_password(&password)?;

        // Create user
        let user = User {
            id: uuid::Uuid::new_v4().to_string(),
            username,
            email,
            password_hash,
            roles: vec!["user".to_string()],
            permissions: vec!["read".to_string()],
            is_active: true,
            created_at: Instant::now(),
            last_login: None,
        };

        // Store user (in real implementation, store in database)
        self.store_user(&user).await?;

        Ok(user)
    }

    /// Authenticate user.
    pub async fn authenticate(&self, request: LoginRequest) -> Result<LoginResponse, String> {
        // Check rate limiting
        if self.is_rate_limited(&request.ip_address).await {
            return Err("Too many login attempts".to_string());
        }

        // Find user
        let user = self.find_user_by_username(&request.username).await
            .ok_or("Invalid credentials")?;

        // Verify password
        if !self.verify_password(&request.password, &user.password_hash)? {
            self.record_failed_attempt(&request.ip_address).await;
            return Err("Invalid credentials".to_string());
        }

        // Check if user is active
        if !user.is_active {
            return Err("Account is disabled".to_string());
        }

        // Generate JWT token
        let token = self.generate_jwt_token(&user)?;
        let refresh_token = self.generate_refresh_token(&user)?;

        // Create session
        let session = UserSession {
            user_id: user.id.clone(),
            session_id: uuid::Uuid::new_v4().to_string(),
            created_at: Instant::now(),
            expires_at: Instant::now() + self.jwt_expiry,
            ip_address: request.ip_address,
            user_agent: request.user_agent,
            is_active: true,
        };

        // Store session
        self.store_session(&session).await?;

        // Update last login
        self.update_last_login(&user.id).await?;

        // Clear rate limiting
        self.clear_rate_limit(&request.ip_address).await;

        Ok(LoginResponse {
            token,
            refresh_token,
            expires_in: self.jwt_expiry.as_secs(),
            user,
        })
    }

    /// Verify JWT token.
    pub async fn verify_token(&self, token: &str) -> Result<JwtClaims, String> {
        let key = DecodingKey::from_secret(self.jwt_secret.as_ref());
        let validation = Validation::new(Algorithm::HS256);

        let token_data = decode::<JwtClaims>(token, &key, &validation)
            .map_err(|e| format!("Invalid token: {}", e))?;

        // Check if token is expired
        let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs();
        if token_data.claims.exp < now {
            return Err("Token expired".to_string());
        }

        // Check if session is still active
        if !self.is_session_active(&token_data.claims.jti).await {
            return Err("Session expired".to_string());
        }

        Ok(token_data.claims)
    }

    /// Check if user has permission.
    pub async fn has_permission(&self, user_id: &str, permission: &str) -> bool {
        if let Some(user) = self.find_user_by_id(user_id).await {
            user.permissions.contains(&permission.to_string())
        } else {
            false
        }
    }

    /// Check if user has role.
    pub async fn has_role(&self, user_id: &str, role: &str) -> bool {
        if let Some(user) = self.find_user_by_id(user_id).await {
            user.roles.contains(&role.to_string())
        } else {
            false
        }
    }

    /// Logout user.
    pub async fn logout(&self, session_id: &str) -> Result<(), String> {
        let mut sessions = self.user_sessions.write().await;
        sessions.remove(session_id);
        Ok(())
    }

    /// Hash password.
    fn hash_password(&self, password: &str) -> Result<String, String> {
        let salt = SaltString::generate(&mut OsRng);
        let password_hash = self.password_hasher
            .hash_password(password.as_bytes(), &salt)
            .map_err(|e| format!("Failed to hash password: {}", e))?;

        Ok(password_hash.to_string())
    }

    /// Verify password.
    fn verify_password(&self, password: &str, hash: &str) -> Result<bool, String> {
        let parsed_hash = PasswordHash::new(hash)
            .map_err(|e| format!("Invalid password hash: {}", e))?;

        Ok(Argon2::default().verify_password(password.as_bytes(), &parsed_hash).is_ok())
    }

    /// Generate JWT token.
    fn generate_jwt_token(&self, user: &User) -> Result<String, String> {
        let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs();
        let exp = now + self.jwt_expiry.as_secs();

        let claims = JwtClaims {
            sub: user.id.clone(),
            username: user.username.clone(),
            roles: user.roles.clone(),
            permissions: user.permissions.clone(),
            iat: now,
            exp,
            jti: uuid::Uuid::new_v4().to_string(),
        };

        let header = Header::new(Algorithm::HS256);
        let key = EncodingKey::from_secret(self.jwt_secret.as_ref());

        encode(&header, &claims, &key)
            .map_err(|e| format!("Failed to generate token: {}", e))
    }

    /// Generate refresh token.
    fn generate_refresh_token(&self, user: &User) -> Result<String, String> {
        // In a real implementation, you would generate a secure refresh token
        Ok(uuid::Uuid::new_v4().to_string())
    }

    /// Check if user exists.
    async fn user_exists(&self, username: &str, email: &str) -> bool {
        // In a real implementation, check database
        false
    }

    /// Store user.
    async fn store_user(&self, user: &User) -> Result<(), String> {
        // In a real implementation, store in database
        Ok(())
    }

    /// Find user by username.
    async fn find_user_by_username(&self, username: &str) -> Option<User> {
        // In a real implementation, query database
        None
    }

    /// Find user by ID.
    async fn find_user_by_id(&self, user_id: &str) -> Option<User> {
        // In a real implementation, query database
        None
    }

    /// Store session.
    async fn store_session(&self, session: &UserSession) -> Result<(), String> {
        let mut sessions = self.user_sessions.write().await;
        sessions.insert(session.session_id.clone(), session.clone());
        Ok(())
    }

    /// Check if session is active.
    async fn is_session_active(&self, session_id: &str) -> bool {
        let sessions = self.user_sessions.read().await;
        if let Some(session) = sessions.get(session_id) {
            session.is_active && session.expires_at > Instant::now()
        } else {
            false
        }
    }

    /// Update last login.
    async fn update_last_login(&self, user_id: &str) -> Result<(), String> {
        // In a real implementation, update database
        Ok(())
    }

    /// Check rate limiting.
    async fn is_rate_limited(&self, ip_address: &str) -> bool {
        let rate_limits = self.rate_limiter.read().await;
        if let Some(info) = rate_limits.get(ip_address) {
            if let Some(blocked_until) = info.blocked_until {
                if blocked_until > Instant::now() {
                    return true;
                }
            }

            // Check if too many attempts in short time
            if info.attempts >= 5 && info.last_attempt.elapsed() < Duration::from_minutes(15) {
                return true;
            }
        }

        false
    }

    /// Record failed attempt.
    async fn record_failed_attempt(&self, ip_address: &str) {
        let mut rate_limits = self.rate_limiter.write().await;
        let info = rate_limits.entry(ip_address.to_string()).or_insert(RateLimitInfo {
            attempts: 0,
            last_attempt: Instant::now(),
            blocked_until: None,
        });

        info.attempts += 1;
        info.last_attempt = Instant::now();

        if info.attempts >= 5 {
            info.blocked_until = Some(Instant::now() + Duration::from_minutes(15));
        }
    }

    /// Clear rate limit.
    async fn clear_rate_limit(&self, ip_address: &str) {
        let mut rate_limits = self.rate_limiter.write().await;
        rate_limits.remove(ip_address);
    }
}

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

    #[tokio::test]
    async fn test_auth_manager() {
        let auth_manager = AuthManager::new("secret".to_string(), Duration::from_hours(1));

        let user = auth_manager.register_user(
            "testuser".to_string(),
            "test@example.com".to_string(),
            "password123".to_string(),
        ).await.unwrap();

        assert_eq!(user.username, "testuser");
        assert_eq!(user.email, "test@example.com");
    }

    #[tokio::test]
    async fn test_authentication() {
        let auth_manager = AuthManager::new("secret".to_string(), Duration::from_hours(1));

        let request = LoginRequest {
            username: "testuser".to_string(),
            password: "password123".to_string(),
            ip_address: "127.0.0.1".to_string(),
            user_agent: "test".to_string(),
        };

        // This would fail in real implementation without proper user setup
        let result = auth_manager.authenticate(request).await;
        assert!(result.is_err());
    }
}

// rust/03-data-encryption.rs

/*
Data encryption patterns and best practices
*/

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use serde::{Deserialize, Serialize};
use aes_gcm::{Aes256Gcm, Key, Nonce};
use aes_gcm::aead::{Aead, NewAead};
use chacha20poly1305::{ChaCha20Poly1305, Key as ChaChaKey, Nonce as ChaChaNonce};
use chacha20poly1305::aead::{Aead as ChaChaAead, NewAead as ChaChaNewAead};

/// Encryption manager.
pub struct EncryptionManager {
    aes_key: Key<Aes256Gcm>,
    chacha_key: ChaChaKey<ChaCha20Poly1305>,
    key_rotation_interval: Duration,
    last_key_rotation: Instant,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptedData {
    pub data: Vec<u8>,
    pub algorithm: String,
    pub key_id: String,
    pub nonce: Vec<u8>,
    pub tag: Vec<u8>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptionConfig {
    pub algorithm: EncryptionAlgorithm,
    pub key_size: usize,
    pub nonce_size: usize,
    pub tag_size: usize,
    pub key_rotation_interval: Duration,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum EncryptionAlgorithm {
    Aes256Gcm,
    ChaCha20Poly1305,
}

impl EncryptionManager {
    pub fn new() -> Self {
        let aes_key = Key::<Aes256Gcm>::from_slice(b"01234567890123456789012345678901");
        let chacha_key = ChaChaKey::<ChaCha20Poly1305>::from_slice(b"01234567890123456789012345678901");

        Self {
            aes_key: *aes_key,
            chacha_key: *chacha_key,
            key_rotation_interval: Duration::from_days(30),
            last_key_rotation: Instant::now(),
        }
    }

    /// Encrypt data using AES-256-GCM.
    pub fn encrypt_aes(&self, data: &[u8]) -> Result<EncryptedData, String> {
        let cipher = Aes256Gcm::new(&self.aes_key);
        let nonce = Nonce::from_slice(b"012345678901"); // 12 bytes for GCM

        let encrypted_data = cipher.encrypt(nonce, data)
            .map_err(|e| format!("Encryption failed: {}", e))?;

        Ok(EncryptedData {
            data: encrypted_data,
            algorithm: "AES-256-GCM".to_string(),
            key_id: "default".to_string(),
            nonce: nonce.to_vec(),
            tag: Vec::new(), // GCM includes tag in encrypted data
        })
    }

    /// Decrypt data using AES-256-GCM.
    pub fn decrypt_aes(&self, encrypted_data: &EncryptedData) -> Result<Vec<u8>, String> {
        let cipher = Aes256Gcm::new(&self.aes_key);
        let nonce = Nonce::from_slice(&encrypted_data.nonce);

        let decrypted_data = cipher.decrypt(nonce, &encrypted_data.data)
            .map_err(|e| format!("Decryption failed: {}", e))?;

        Ok(decrypted_data)
    }

    /// Encrypt data using ChaCha20-Poly1305.
    pub fn encrypt_chacha(&self, data: &[u8]) -> Result<EncryptedData, String> {
        let cipher = ChaCha20Poly1305::new(&self.chacha_key);
        let nonce = ChaChaNonce::from_slice(b"012345678901"); // 12 bytes for ChaCha20-Poly1305

        let encrypted_data = cipher.encrypt(nonce, data)
            .map_err(|e| format!("Encryption failed: {}", e))?;

        Ok(EncryptedData {
            data: encrypted_data,
            algorithm: "ChaCha20-Poly1305".to_string(),
            key_id: "default".to_string(),
            nonce: nonce.to_vec(),
            tag: Vec::new(), // ChaCha20-Poly1305 includes tag in encrypted data
        })
    }

    /// Decrypt data using ChaCha20-Poly1305.
    pub fn decrypt_chacha(&self, encrypted_data: &EncryptedData) -> Result<Vec<u8>, String> {
        let cipher = ChaCha20Poly1305::new(&self.chacha_key);
        let nonce = ChaChaNonce::from_slice(&encrypted_data.nonce);

        let decrypted_data = cipher.decrypt(nonce, &encrypted_data.data)
            .map_err(|e| format!("Decryption failed: {}", e))?;

        Ok(decrypted_data)
    }

    /// Encrypt data with automatic algorithm selection.
    pub fn encrypt(&self, data: &[u8], algorithm: EncryptionAlgorithm) -> Result<EncryptedData, String> {
        match algorithm {
            EncryptionAlgorithm::Aes256Gcm => self.encrypt_aes(data),
            EncryptionAlgorithm::ChaCha20Poly1305 => self.encrypt_chacha(data),
        }
    }

    /// Decrypt data with automatic algorithm detection.
    pub fn decrypt(&self, encrypted_data: &EncryptedData) -> Result<Vec<u8>, String> {
        match encrypted_data.algorithm.as_str() {
            "AES-256-GCM" => self.decrypt_aes(encrypted_data),
            "ChaCha20-Poly1305" => self.decrypt_chacha(encrypted_data),
            _ => Err("Unsupported encryption algorithm".to_string()),
        }
    }

    /// Encrypt sensitive fields in a struct.
    pub fn encrypt_struct<T>(&self, data: &T, fields: &[String]) -> Result<HashMap<String, EncryptedData>, String> 
    where
        T: Serialize,
    {
        let json = serde_json::to_value(data)
            .map_err(|e| format!("Serialization failed: {}", e))?;

        let mut encrypted_fields = HashMap::new();

        if let Some(obj) = json.as_object() {
            for field in fields {
                if let Some(value) = obj.get(field) {
                    let field_data = serde_json::to_vec(value)
                        .map_err(|e| format!("Field serialization failed: {}", e))?;

                    let encrypted = self.encrypt(&field_data, EncryptionAlgorithm::Aes256Gcm)?;
                    encrypted_fields.insert(field.clone(), encrypted);
                }
            }
        }

        Ok(encrypted_fields)
    }

    /// Decrypt sensitive fields in a struct.
    pub fn decrypt_struct<T>(&self, encrypted_fields: &HashMap<String, EncryptedData>) -> Result<HashMap<String, serde_json::Value>, String> 
    where
        T: Deserialize<'static>,
    {
        let mut decrypted_fields = HashMap::new();

        for (field, encrypted_data) in encrypted_fields {
            let decrypted_data = self.decrypt(encrypted_data)?;
            let value: serde_json::Value = serde_json::from_slice(&decrypted_data)
                .map_err(|e| format!("Deserialization failed: {}", e))?;

            decrypted_fields.insert(field.clone(), value);
        }

        Ok(decrypted_fields)
    }

    /// Check if key rotation is needed.
    pub fn needs_key_rotation(&self) -> bool {
        self.last_key_rotation.elapsed() >= self.key_rotation_interval
    }

    /// Rotate encryption keys.
    pub fn rotate_keys(&mut self) -> Result<(), String> {
        // In a real implementation, generate new keys and update key store
        self.last_key_rotation = Instant::now();
        Ok(())
    }

    /// Generate secure random key.
    pub fn generate_key(&self, size: usize) -> Result<Vec<u8>, String> {
        use rand::RngCore;

        let mut key = vec![0u8; size];
        rand::thread_rng().fill_bytes(&mut key);
        Ok(key)
    }

    /// Generate secure random nonce.
    pub fn generate_nonce(&self, size: usize) -> Result<Vec<u8>, String> {
        use rand::RngCore;

        let mut nonce = vec![0u8; size];
        rand::thread_rng().fill_bytes(&mut nonce);
        Ok(nonce)
    }
}

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

    #[test]
    fn test_aes_encryption() {
        let manager = EncryptionManager::new();
        let data = b"Hello, World!";

        let encrypted = manager.encrypt_aes(data).unwrap();
        let decrypted = manager.decrypt_aes(&encrypted).unwrap();

        assert_eq!(decrypted, data);
    }

    #[test]
    fn test_chacha_encryption() {
        let manager = EncryptionManager::new();
        let data = b"Hello, World!";

        let encrypted = manager.encrypt_chacha(data).unwrap();
        let decrypted = manager.decrypt_chacha(&encrypted).unwrap();

        assert_eq!(decrypted, data);
    }

    #[test]
    fn test_automatic_encryption() {
        let manager = EncryptionManager::new();
        let data = b"Hello, World!";

        let encrypted = manager.encrypt(data, EncryptionAlgorithm::Aes256Gcm).unwrap();
        let decrypted = manager.decrypt(&encrypted).unwrap();

        assert_eq!(decrypted, data);
    }
}