use crate::{
config::StorageConfig,
models::{ApiKey, SecureApiKey},
};
use async_trait::async_trait;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use thiserror::Error;
use tokio::fs;
use tracing::{debug, error, info, warn};
#[derive(Debug, Error)]
pub enum StorageError {
#[error("Storage error: {0}")]
General(String),
#[error("File I/O error: {0}")]
Io(#[from] std::io::Error),
#[error("Serialization error: {0}")]
Serialization(#[from] serde_json::Error),
#[error("Permission error: {0}")]
Permission(String),
#[error("Encryption error: {0}")]
Encryption(#[from] crate::crypto::encryption::EncryptionError),
}
#[async_trait]
pub trait StorageBackend: Send + Sync {
async fn load_keys(&self) -> Result<HashMap<String, ApiKey>, StorageError>;
async fn save_key(&self, key: &ApiKey) -> Result<(), StorageError>;
async fn delete_key(&self, key_id: &str) -> Result<(), StorageError>;
async fn save_all_keys(&self, keys: &HashMap<String, ApiKey>) -> Result<(), StorageError>;
}
pub async fn create_storage_backend(
config: &StorageConfig,
) -> Result<Arc<dyn StorageBackend>, StorageError> {
match config {
StorageConfig::File {
path,
file_permissions,
dir_permissions,
require_secure_filesystem,
enable_filesystem_monitoring,
} => {
let storage = FileStorage::new(
path.clone(),
*file_permissions,
*dir_permissions,
*require_secure_filesystem,
*enable_filesystem_monitoring,
)
.await?;
Ok(Arc::new(storage))
}
StorageConfig::Environment { prefix } => {
let storage = EnvironmentStorage::new(prefix.clone());
Ok(Arc::new(storage))
}
StorageConfig::Memory => {
let storage = MemoryStorage::new();
Ok(Arc::new(storage))
}
}
}
pub struct FileStorage {
path: PathBuf,
encryption_key: [u8; 32],
#[allow(dead_code)]
file_permissions: u32,
#[allow(dead_code)]
dir_permissions: u32,
#[allow(dead_code)]
require_secure_filesystem: bool,
enable_filesystem_monitoring: bool,
write_mutex: tokio::sync::Mutex<()>,
}
impl FileStorage {
pub async fn new(
path: PathBuf,
file_permissions: u32,
dir_permissions: u32,
require_secure_filesystem: bool,
enable_filesystem_monitoring: bool,
) -> Result<Self, StorageError> {
use crate::crypto::encryption::derive_encryption_key;
use crate::crypto::keys::generate_master_key;
if require_secure_filesystem {
Self::validate_filesystem_security(&path).await?;
}
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).await?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mut perms = fs::metadata(parent).await?.permissions();
perms.set_mode(dir_permissions); fs::set_permissions(parent, perms).await?;
Self::verify_directory_security(parent, dir_permissions).await?;
}
}
let master_key = generate_master_key().map_err(|e| StorageError::General(e.to_string()))?;
let encryption_key = derive_encryption_key(&master_key, "api-key-storage");
let storage = Self {
path,
encryption_key,
file_permissions,
dir_permissions,
require_secure_filesystem,
enable_filesystem_monitoring,
write_mutex: tokio::sync::Mutex::new(()),
};
if !storage.path.exists() {
storage.save_all_keys(&HashMap::new()).await?;
} else {
storage.ensure_secure_permissions().await?;
}
Ok(storage)
}
async fn ensure_secure_permissions(&self) -> Result<(), StorageError> {
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
if self.path.exists() {
let metadata = fs::metadata(&self.path).await?;
let mode = metadata.permissions().mode() & 0o777;
if mode != self.file_permissions {
warn!(
"Incorrect permissions on key file: {:o}, fixing to {:o}",
mode, self.file_permissions
);
let mut perms = metadata.permissions();
perms.set_mode(self.file_permissions);
fs::set_permissions(&self.path, perms).await?;
}
Self::verify_file_ownership(&self.path).await?;
}
}
Ok(())
}
#[allow(unused_variables)]
async fn validate_filesystem_security(path: &PathBuf) -> Result<(), StorageError> {
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
if let Some(parent) = path.parent() {
if parent.exists() {
let metadata = fs::metadata(parent).await?;
let _dev = metadata.dev();
if let Ok(mount_info) = fs::read_to_string("/proc/mounts").await {
let path_str = parent.to_string_lossy();
for line in mount_info.lines() {
let parts: Vec<&str> = line.split_whitespace().collect();
if parts.len() >= 3 {
let mount_point = parts[1];
let fs_type = parts[2];
if path_str.starts_with(mount_point) {
match fs_type {
"nfs" | "nfs4" | "cifs" | "smb" | "smbfs"
| "fuse.sshfs" => {
return Err(StorageError::Permission(format!(
"Storage path {} is on insecure network filesystem: {}",
path_str, fs_type
)));
}
_ => {}
}
}
}
}
}
}
}
}
Ok(())
}
#[allow(unused_variables)]
async fn verify_directory_security(
dir: &std::path::Path,
expected_perms: u32,
) -> Result<(), StorageError> {
#[cfg(unix)]
{
use std::os::unix::fs::{MetadataExt, PermissionsExt};
let metadata = fs::metadata(dir).await?;
let mode = metadata.permissions().mode() & 0o777;
if (mode & !expected_perms) != 0 {
return Err(StorageError::Permission(format!(
"Directory {} has insecure permissions: {:o} (expected: {:o})",
dir.display(),
mode,
expected_perms
)));
}
let current_uid = unsafe { libc::getuid() };
if metadata.uid() != current_uid {
return Err(StorageError::Permission(format!(
"Directory {} is not owned by current user (uid: {} vs {})",
dir.display(),
metadata.uid(),
current_uid
)));
}
}
Ok(())
}
#[allow(unused_variables)]
async fn verify_file_ownership(file: &std::path::Path) -> Result<(), StorageError> {
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
let metadata = fs::metadata(file).await?;
let current_uid = unsafe { libc::getuid() };
if metadata.uid() != current_uid {
return Err(StorageError::Permission(format!(
"File {} is not owned by current user (uid: {} vs {})",
file.display(),
metadata.uid(),
current_uid
)));
}
}
Ok(())
}
async fn save_secure_keys(
&self,
keys: &HashMap<String, SecureApiKey>,
) -> Result<(), StorageError> {
use crate::crypto::encryption::encrypt_data;
let content = serde_json::to_string_pretty(keys)?;
let encrypted_data = encrypt_data(content.as_bytes(), &self.encryption_key)?;
let encrypted_content = serde_json::to_string_pretty(&encrypted_data)?;
let temp_path = self.path.with_extension("tmp");
fs::write(&temp_path, encrypted_content).await?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mut perms = fs::metadata(&temp_path).await?.permissions();
perms.set_mode(self.file_permissions); fs::set_permissions(&temp_path, perms).await?;
}
fs::rename(&temp_path, &self.path).await?;
debug!("Saved {} keys to encrypted file storage", keys.len());
Ok(())
}
pub async fn create_backup(&self) -> Result<PathBuf, StorageError> {
if !self.path.exists() {
return Err(StorageError::General(
"Storage file does not exist".to_string(),
));
}
let timestamp = chrono::Utc::now().format("%Y%m%d_%H%M%S_%3f");
let backup_path = self
.path
.with_extension(format!("backup_{}.enc", timestamp));
fs::copy(&self.path, &backup_path).await?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mut perms = fs::metadata(&backup_path).await?.permissions();
perms.set_mode(self.file_permissions);
fs::set_permissions(&backup_path, perms).await?;
}
debug!("Created secure backup: {}", backup_path.display());
Ok(backup_path)
}
pub async fn restore_from_backup(&self, backup_path: &PathBuf) -> Result<(), StorageError> {
if !backup_path.exists() {
return Err(StorageError::General(
"Backup file does not exist".to_string(),
));
}
Self::verify_file_ownership(backup_path).await?;
let temp_path = self.path.with_extension("restore_tmp");
fs::copy(backup_path, &temp_path).await?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mut perms = fs::metadata(&temp_path).await?.permissions();
perms.set_mode(self.file_permissions);
fs::set_permissions(&temp_path, perms).await?;
}
fs::rename(&temp_path, &self.path).await?;
info!("Restored from backup: {}", backup_path.display());
Ok(())
}
pub async fn cleanup_backups(&self, keep_count: usize) -> Result<(), StorageError> {
if let Some(parent) = self.path.parent() {
let filename_stem = self
.path
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("keys");
let mut backups = Vec::new();
let mut entries = fs::read_dir(parent).await?;
while let Some(entry) = entries.next_entry().await? {
let path = entry.path();
if let Some(filename) = path.file_name().and_then(|n| n.to_str()) {
if filename.starts_with(&format!("{}.backup_", filename_stem)) {
if let Ok(metadata) = entry.metadata().await {
backups.push((
path,
metadata
.modified()
.unwrap_or(std::time::SystemTime::UNIX_EPOCH),
));
}
}
}
}
backups.sort_by(|a, b| b.1.cmp(&a.1));
for (backup_path, _) in backups.iter().skip(keep_count) {
if let Err(e) = fs::remove_file(backup_path).await {
warn!(
"Failed to remove old backup {}: {}",
backup_path.display(),
e
);
} else {
debug!("Removed old backup: {}", backup_path.display());
}
}
}
Ok(())
}
#[cfg(target_os = "linux")]
pub async fn start_filesystem_monitoring(&self) -> Result<(), StorageError> {
if !self.enable_filesystem_monitoring {
return Ok(());
}
use inotify::{Inotify, WatchMask};
let mut inotify = Inotify::init()
.map_err(|e| StorageError::General(format!("Failed to initialize inotify: {}", e)))?;
if let Some(parent) = self.path.parent() {
inotify
.watches()
.add(
parent,
WatchMask::MODIFY | WatchMask::ATTRIB | WatchMask::MOVED_TO | WatchMask::DELETE,
)
.map_err(|e| {
StorageError::General(format!("Failed to add inotify watch: {}", e))
})?;
info!("Started filesystem monitoring for: {}", parent.display());
let path = self.path.clone();
let file_permissions = self.file_permissions;
tokio::spawn(async move {
use tracing::{error, warn};
let mut buffer = [0; 1024];
loop {
match inotify.read_events_blocking(&mut buffer) {
Ok(events) => {
for event in events {
if let Some(name) = event.name {
if name.to_string_lossy().contains("keys") {
warn!(
"Detected unauthorized change to auth storage: {:?} (mask: {:?})",
name, event.mask
);
if path.exists() {
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
if let Ok(metadata) = std::fs::metadata(&path) {
let mode =
metadata.permissions().mode() & 0o777;
if mode != file_permissions {
error!(
"Security violation: File permissions changed from {:o} to {:o}",
file_permissions, mode
);
}
}
}
}
}
}
}
}
Err(e) => {
error!("Error reading inotify events: {}", e);
break;
}
}
}
});
}
Ok(())
}
#[cfg(not(target_os = "linux"))]
pub async fn start_filesystem_monitoring(&self) -> Result<(), StorageError> {
if self.enable_filesystem_monitoring {
warn!("Filesystem monitoring is only supported on Linux systems");
}
Ok(())
}
}
#[async_trait]
impl StorageBackend for FileStorage {
async fn load_keys(&self) -> Result<HashMap<String, ApiKey>, StorageError> {
use crate::crypto::encryption::decrypt_data;
self.ensure_secure_permissions().await?;
if !self.path.exists() {
return Ok(HashMap::new());
}
let content = fs::read(&self.path).await?;
if content.is_empty() {
return Ok(HashMap::new());
}
let decrypted_content = if let Ok(encrypted_data) = serde_json::from_slice(&content) {
let decrypted_bytes = decrypt_data(&encrypted_data, &self.encryption_key)?;
String::from_utf8(decrypted_bytes)
.map_err(|e| StorageError::General(format!("Invalid UTF-8: {}", e)))?
} else {
let plain_text = String::from_utf8(content)
.map_err(|e| StorageError::General(format!("Invalid UTF-8: {}", e)))?;
warn!("Found legacy plain text keys, converting to secure format");
let legacy_keys: HashMap<String, ApiKey> = serde_json::from_str(&plain_text)?;
let secure_keys: HashMap<String, SecureApiKey> = legacy_keys
.into_iter()
.map(|(id, key)| (id, key.to_secure_storage()))
.collect();
self.save_secure_keys(&secure_keys).await?;
plain_text
};
let secure_keys: HashMap<String, SecureApiKey> = serde_json::from_str(&decrypted_content)?;
let keys: HashMap<String, ApiKey> = secure_keys
.into_iter()
.map(|(id, secure_key)| (id, secure_key.to_api_key()))
.collect();
debug!("Loaded {} keys from encrypted file storage", keys.len());
Ok(keys)
}
async fn save_key(&self, key: &ApiKey) -> Result<(), StorageError> {
let _lock = self.write_mutex.lock().await;
let mut keys = self.load_keys().await?;
keys.insert(key.id.clone(), key.clone());
self.save_all_keys_internal(&keys).await
}
async fn delete_key(&self, key_id: &str) -> Result<(), StorageError> {
let _lock = self.write_mutex.lock().await;
let mut keys = self.load_keys().await?;
keys.remove(key_id);
self.save_all_keys_internal(&keys).await
}
async fn save_all_keys(&self, keys: &HashMap<String, ApiKey>) -> Result<(), StorageError> {
let _lock = self.write_mutex.lock().await;
self.save_all_keys_internal(keys).await
}
}
impl FileStorage {
async fn save_all_keys_internal(
&self,
keys: &HashMap<String, ApiKey>,
) -> Result<(), StorageError> {
let secure_keys: HashMap<String, SecureApiKey> = keys
.iter()
.map(|(id, key)| (id.clone(), key.to_secure_storage()))
.collect();
self.save_secure_keys(&secure_keys).await
}
}
pub struct EnvironmentStorage {
var_name: String,
}
impl EnvironmentStorage {
pub fn new(var_name: String) -> Self {
Self { var_name }
}
}
#[async_trait]
impl StorageBackend for EnvironmentStorage {
async fn load_keys(&self) -> Result<HashMap<String, ApiKey>, StorageError> {
match std::env::var(&self.var_name) {
Ok(content) => {
if content.trim().is_empty() {
return Ok(HashMap::new());
}
let keys: HashMap<String, ApiKey> = serde_json::from_str(&content)?;
debug!("Loaded {} keys from environment storage", keys.len());
Ok(keys)
}
Err(_) => {
debug!(
"Environment variable {} not found, returning empty keys",
self.var_name
);
Ok(HashMap::new())
}
}
}
async fn save_key(&self, key: &ApiKey) -> Result<(), StorageError> {
let mut keys = self.load_keys().await?;
keys.insert(key.id.clone(), key.clone());
self.save_all_keys(&keys).await
}
async fn delete_key(&self, key_id: &str) -> Result<(), StorageError> {
let mut keys = self.load_keys().await?;
keys.remove(key_id);
self.save_all_keys(&keys).await
}
async fn save_all_keys(&self, keys: &HashMap<String, ApiKey>) -> Result<(), StorageError> {
let content = serde_json::to_string(keys)?;
unsafe {
std::env::set_var(&self.var_name, content);
}
debug!("Saved {} keys to environment storage", keys.len());
Ok(())
}
}
pub struct MemoryStorage {
keys: tokio::sync::RwLock<HashMap<String, ApiKey>>,
}
impl MemoryStorage {
pub fn new() -> Self {
Self {
keys: tokio::sync::RwLock::new(HashMap::new()),
}
}
}
#[async_trait]
impl StorageBackend for MemoryStorage {
async fn load_keys(&self) -> Result<HashMap<String, ApiKey>, StorageError> {
let keys = self.keys.read().await;
debug!("Loaded {} keys from memory storage", keys.len());
Ok(keys.clone())
}
async fn save_key(&self, key: &ApiKey) -> Result<(), StorageError> {
let mut keys = self.keys.write().await;
keys.insert(key.id.clone(), key.clone());
debug!("Saved key {} to memory storage", key.id);
Ok(())
}
async fn delete_key(&self, key_id: &str) -> Result<(), StorageError> {
let mut keys = self.keys.write().await;
keys.remove(key_id);
debug!("Deleted key {} from memory storage", key_id);
Ok(())
}
async fn save_all_keys(&self, new_keys: &HashMap<String, ApiKey>) -> Result<(), StorageError> {
let mut keys = self.keys.write().await;
*keys = new_keys.clone();
debug!(
"Replaced all keys in memory storage with {} keys",
new_keys.len()
);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::models::{ApiKey, Role};
use chrono::{Duration, Utc};
use std::collections::HashMap;
use tempfile::TempDir;
use tokio::fs;
fn create_test_key(name: &str, role: Role) -> ApiKey {
ApiKey::new(
name.to_string(),
role,
Some(Utc::now() + Duration::days(30)),
vec!["127.0.0.1".to_string()],
)
}
fn create_test_keys() -> HashMap<String, ApiKey> {
let mut keys = HashMap::new();
let admin_key = create_test_key("admin-key", Role::Admin);
let operator_key = create_test_key("operator-key", Role::Operator);
let monitor_key = create_test_key("monitor-key", Role::Monitor);
keys.insert(admin_key.id.clone(), admin_key);
keys.insert(operator_key.id.clone(), operator_key);
keys.insert(monitor_key.id.clone(), monitor_key);
keys
}
#[test]
fn test_storage_error_display() {
let error = StorageError::General("test error".to_string());
assert_eq!(error.to_string(), "Storage error: test error");
let io_error = std::io::Error::new(std::io::ErrorKind::NotFound, "file not found");
let storage_error = StorageError::Io(io_error);
assert!(storage_error.to_string().contains("File I/O error"));
let perm_error = StorageError::Permission("access denied".to_string());
assert_eq!(perm_error.to_string(), "Permission error: access denied");
}
#[test]
fn test_storage_error_from_io_error() {
let io_error =
std::io::Error::new(std::io::ErrorKind::PermissionDenied, "permission denied");
let storage_error: StorageError = io_error.into();
match storage_error {
StorageError::Io(_) => (),
_ => panic!("Expected Io variant"),
}
}
#[test]
fn test_storage_error_from_serde_error() {
let serde_error = serde_json::from_str::<serde_json::Value>("invalid json").unwrap_err();
let storage_error: StorageError = serde_error.into();
match storage_error {
StorageError::Serialization(_) => (),
_ => panic!("Expected Serialization variant"),
}
}
mod memory_storage_tests {
use super::*;
#[tokio::test]
async fn test_memory_storage_new() {
let storage = MemoryStorage::new();
let keys = storage.load_keys().await.unwrap();
assert!(keys.is_empty());
}
#[tokio::test]
async fn test_memory_storage_save_and_load_key() {
let storage = MemoryStorage::new();
let test_key = create_test_key("test-key", Role::Operator);
storage.save_key(&test_key).await.unwrap();
let keys = storage.load_keys().await.unwrap();
assert_eq!(keys.len(), 1);
assert!(keys.contains_key(&test_key.id));
let loaded_key = &keys[&test_key.id];
assert_eq!(loaded_key.name, test_key.name);
assert_eq!(loaded_key.role, test_key.role);
}
#[tokio::test]
async fn test_memory_storage_save_multiple_keys() {
let storage = MemoryStorage::new();
let test_keys = create_test_keys();
for key in test_keys.values() {
storage.save_key(key).await.unwrap();
}
let loaded_keys = storage.load_keys().await.unwrap();
assert_eq!(loaded_keys.len(), test_keys.len());
for (id, key) in test_keys.iter() {
assert!(loaded_keys.contains_key(id));
assert_eq!(loaded_keys[id].name, key.name);
}
}
#[tokio::test]
async fn test_memory_storage_delete_key() {
let storage = MemoryStorage::new();
let test_key = create_test_key("test-key", Role::Monitor);
storage.save_key(&test_key).await.unwrap();
assert_eq!(storage.load_keys().await.unwrap().len(), 1);
storage.delete_key(&test_key.id).await.unwrap();
let keys = storage.load_keys().await.unwrap();
assert!(keys.is_empty());
}
#[tokio::test]
async fn test_memory_storage_delete_nonexistent_key() {
let storage = MemoryStorage::new();
storage.delete_key("nonexistent").await.unwrap();
assert!(storage.load_keys().await.unwrap().is_empty());
}
#[tokio::test]
async fn test_memory_storage_save_all_keys() {
let storage = MemoryStorage::new();
let test_keys = create_test_keys();
storage.save_all_keys(&test_keys).await.unwrap();
let loaded_keys = storage.load_keys().await.unwrap();
assert_eq!(loaded_keys.len(), test_keys.len());
for (id, key) in test_keys.iter() {
assert!(loaded_keys.contains_key(id));
assert_eq!(loaded_keys[id].name, key.name);
}
}
#[tokio::test]
async fn test_memory_storage_save_all_keys_replaces_existing() {
let storage = MemoryStorage::new();
let initial_keys = create_test_keys();
storage.save_all_keys(&initial_keys).await.unwrap();
assert_eq!(storage.load_keys().await.unwrap().len(), initial_keys.len());
let mut new_keys = HashMap::new();
let new_key = create_test_key("new-key", Role::Admin);
new_keys.insert(new_key.id.clone(), new_key);
storage.save_all_keys(&new_keys).await.unwrap();
let loaded_keys = storage.load_keys().await.unwrap();
assert_eq!(loaded_keys.len(), 1);
assert!(loaded_keys.contains_key(new_keys.keys().next().unwrap()));
}
#[tokio::test]
async fn test_memory_storage_concurrent_access() {
let storage = std::sync::Arc::new(MemoryStorage::new());
let mut handles = vec![];
for i in 0..10 {
let storage_clone = storage.clone();
let handle = tokio::spawn(async move {
let key = create_test_key(&format!("key-{}", i), Role::Operator);
storage_clone.save_key(&key).await.unwrap();
key.id
});
handles.push(handle);
}
let mut saved_ids = vec![];
for handle in handles {
saved_ids.push(handle.await.unwrap());
}
let keys = storage.load_keys().await.unwrap();
assert_eq!(keys.len(), 10);
for id in saved_ids {
assert!(keys.contains_key(&id));
}
}
}
mod environment_storage_tests {
use super::*;
#[tokio::test]
async fn test_environment_storage_new() {
let storage = EnvironmentStorage::new("TEST_MCP_KEYS".to_string());
unsafe {
std::env::remove_var("TEST_MCP_KEYS");
}
let keys = storage.load_keys().await.unwrap();
assert!(keys.is_empty());
}
#[tokio::test]
async fn test_environment_storage_save_and_load_key() {
let var_name = "TEST_MCP_KEYS_SAVE_LOAD";
unsafe {
std::env::remove_var(var_name);
}
let storage = EnvironmentStorage::new(var_name.to_string());
let test_key = create_test_key("env-test-key", Role::Monitor);
storage.save_key(&test_key).await.unwrap();
let keys = storage.load_keys().await.unwrap();
assert_eq!(keys.len(), 1);
assert!(keys.contains_key(&test_key.id));
assert!(std::env::var(var_name).is_ok());
unsafe {
std::env::remove_var(var_name);
}
}
#[tokio::test]
async fn test_environment_storage_multiple_keys() {
let var_name = "TEST_MCP_KEYS_MULTIPLE";
unsafe {
std::env::remove_var(var_name);
}
let storage = EnvironmentStorage::new(var_name.to_string());
let test_keys = create_test_keys();
storage.save_all_keys(&test_keys).await.unwrap();
let loaded_keys = storage.load_keys().await.unwrap();
assert_eq!(loaded_keys.len(), test_keys.len());
for (id, key) in test_keys.iter() {
assert!(loaded_keys.contains_key(id));
assert_eq!(loaded_keys[id].name, key.name);
}
unsafe {
std::env::remove_var(var_name);
}
}
#[tokio::test]
async fn test_environment_storage_delete_key() {
let var_name = "TEST_MCP_KEYS_DELETE";
unsafe {
std::env::remove_var(var_name);
}
let storage = EnvironmentStorage::new(var_name.to_string());
let test_keys = create_test_keys();
let key_to_delete = test_keys.values().next().unwrap().id.clone();
storage.save_all_keys(&test_keys).await.unwrap();
assert_eq!(storage.load_keys().await.unwrap().len(), test_keys.len());
storage.delete_key(&key_to_delete).await.unwrap();
let remaining_keys = storage.load_keys().await.unwrap();
assert_eq!(remaining_keys.len(), test_keys.len() - 1);
assert!(!remaining_keys.contains_key(&key_to_delete));
unsafe {
std::env::remove_var(var_name);
}
}
#[tokio::test]
async fn test_environment_storage_empty_content() {
let var_name = "TEST_MCP_KEYS_EMPTY";
unsafe {
std::env::set_var(var_name, "");
}
let storage = EnvironmentStorage::new(var_name.to_string());
let keys = storage.load_keys().await.unwrap();
assert!(keys.is_empty());
unsafe {
std::env::remove_var(var_name);
}
}
#[tokio::test]
async fn test_environment_storage_invalid_json() {
let var_name = "TEST_MCP_KEYS_INVALID";
unsafe {
std::env::set_var(var_name, "invalid json content");
}
let storage = EnvironmentStorage::new(var_name.to_string());
let result = storage.load_keys().await;
assert!(result.is_err());
match result.unwrap_err() {
StorageError::Serialization(_) => (),
_ => panic!("Expected serialization error"),
}
unsafe {
std::env::remove_var(var_name);
}
}
#[tokio::test]
async fn test_environment_storage_overwrite_existing() {
let var_name = "TEST_MCP_KEYS_OVERWRITE";
unsafe {
std::env::remove_var(var_name);
}
let storage = EnvironmentStorage::new(var_name.to_string());
let initial_keys = create_test_keys();
storage.save_all_keys(&initial_keys).await.unwrap();
let mut new_keys = HashMap::new();
let new_key = create_test_key("overwrite-key", Role::Admin);
new_keys.insert(new_key.id.clone(), new_key);
storage.save_all_keys(&new_keys).await.unwrap();
let loaded_keys = storage.load_keys().await.unwrap();
assert_eq!(loaded_keys.len(), 1);
assert!(loaded_keys.contains_key(new_keys.keys().next().unwrap()));
unsafe {
std::env::remove_var(var_name);
}
}
}
mod file_storage_tests {
use super::*;
async fn create_test_file_storage() -> (FileStorage, TempDir) {
unsafe {
std::env::set_var(
"PULSEENGINE_MCP_MASTER_KEY",
"l9EYbalIRp2CF35M4mKcWDqRvx3TFc7U4nX5zvQF56Q",
);
}
let temp_dir = TempDir::new().unwrap();
let storage_path = temp_dir.path().join("test_keys.enc");
let storage = FileStorage::new(
storage_path,
0o600,
0o700,
false, false, )
.await
.unwrap();
(storage, temp_dir)
}
#[tokio::test]
async fn test_file_storage_new() {
let (storage, _temp_dir) = create_test_file_storage().await;
let keys = storage.load_keys().await.unwrap();
assert!(keys.is_empty());
assert!(storage.path.exists());
}
#[tokio::test]
async fn test_file_storage_save_and_load_key() {
let (storage, _temp_dir) = create_test_file_storage().await;
let test_key = create_test_key("file-test-key", Role::Operator);
storage.save_key(&test_key).await.unwrap();
let keys = storage.load_keys().await.unwrap();
assert_eq!(keys.len(), 1);
assert!(keys.contains_key(&test_key.id));
let loaded_key = &keys[&test_key.id];
assert_eq!(loaded_key.name, test_key.name);
assert_eq!(loaded_key.role, test_key.role);
assert_eq!(loaded_key.key, "***redacted***");
}
#[tokio::test]
async fn test_file_storage_multiple_keys() {
let (storage, _temp_dir) = create_test_file_storage().await;
let test_keys = create_test_keys();
storage.save_all_keys(&test_keys).await.unwrap();
let loaded_keys = storage.load_keys().await.unwrap();
assert_eq!(loaded_keys.len(), test_keys.len());
for (id, key) in test_keys.iter() {
assert!(loaded_keys.contains_key(id));
assert_eq!(loaded_keys[id].name, key.name);
assert_eq!(loaded_keys[id].role, key.role);
}
}
#[tokio::test]
async fn test_file_storage_delete_key() {
let (storage, _temp_dir) = create_test_file_storage().await;
let test_keys = create_test_keys();
let key_to_delete = test_keys.values().next().unwrap().id.clone();
storage.save_all_keys(&test_keys).await.unwrap();
assert_eq!(storage.load_keys().await.unwrap().len(), test_keys.len());
storage.delete_key(&key_to_delete).await.unwrap();
let remaining_keys = storage.load_keys().await.unwrap();
assert_eq!(remaining_keys.len(), test_keys.len() - 1);
assert!(!remaining_keys.contains_key(&key_to_delete));
}
#[tokio::test]
async fn test_file_storage_encryption() {
let (storage, _temp_dir) = create_test_file_storage().await;
let test_key = create_test_key("encryption-test", Role::Admin);
storage.save_key(&test_key).await.unwrap();
let raw_content = fs::read(&storage.path).await.unwrap();
let raw_text = String::from_utf8_lossy(&raw_content);
assert!(!raw_text.contains(&test_key.name));
assert!(!raw_text.contains(&test_key.key));
let loaded_keys = storage.load_keys().await.unwrap();
assert_eq!(loaded_keys.len(), 1);
assert!(loaded_keys.contains_key(&test_key.id));
}
#[tokio::test]
async fn test_file_storage_empty_file() {
let temp_dir = TempDir::new().unwrap();
let storage_path = temp_dir.path().join("empty_keys.enc");
fs::write(&storage_path, "").await.unwrap();
let storage = FileStorage::new(storage_path, 0o600, 0o700, false, false)
.await
.unwrap();
let keys = storage.load_keys().await.unwrap();
assert!(keys.is_empty());
}
#[tokio::test]
async fn test_file_storage_nonexistent_file() {
let temp_dir = TempDir::new().unwrap();
let storage_path = temp_dir.path().join("nonexistent").join("keys.enc");
let storage = FileStorage::new(storage_path.clone(), 0o600, 0o700, false, false)
.await
.unwrap();
let keys = storage.load_keys().await.unwrap();
assert!(keys.is_empty());
assert!(storage_path.exists());
}
#[tokio::test]
#[allow(clippy::await_holding_lock)] #[ignore = "Flaky test - encryption race condition in CI environment"]
async fn test_file_storage_persistence() {
static TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
let _lock = TEST_LOCK.lock().await;
let original_master_key = std::env::var("PULSEENGINE_MCP_MASTER_KEY").ok();
unsafe {
std::env::set_var(
"PULSEENGINE_MCP_MASTER_KEY",
"l9EYbalIRp2CF35M4mKcWDqRvx3TFc7U4nX5zvQF56Q",
);
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
let temp_dir = TempDir::new().unwrap();
let storage_path = temp_dir.path().join("persistent_keys.enc");
let test_keys = create_test_keys();
{
let storage = FileStorage::new(storage_path.clone(), 0o600, 0o700, false, false)
.await
.unwrap();
storage.save_all_keys(&test_keys).await.unwrap();
}
assert!(storage_path.exists());
let file_metadata = std::fs::metadata(&storage_path).unwrap();
assert!(file_metadata.len() > 0);
{
let storage = FileStorage::new(storage_path, 0o600, 0o700, false, false)
.await
.unwrap();
let loaded_keys = storage.load_keys().await.unwrap();
assert_eq!(loaded_keys.len(), test_keys.len());
for (id, key) in test_keys.iter() {
assert!(loaded_keys.contains_key(id));
assert_eq!(loaded_keys[id].name, key.name);
}
}
unsafe {
match original_master_key {
Some(key) => std::env::set_var("PULSEENGINE_MCP_MASTER_KEY", key),
None => std::env::remove_var("PULSEENGINE_MCP_MASTER_KEY"),
}
}
}
#[tokio::test]
async fn test_file_storage_backup_and_restore() {
let (storage, _temp_dir) = create_test_file_storage().await;
let test_keys = create_test_keys();
storage.save_all_keys(&test_keys).await.unwrap();
let backup_path = storage.create_backup().await.unwrap();
assert!(backup_path.exists());
assert!(backup_path.to_string_lossy().contains("backup_"));
let mut modified_keys = HashMap::new();
let new_key = create_test_key("backup-test", Role::Monitor);
modified_keys.insert(new_key.id.clone(), new_key);
storage.save_all_keys(&modified_keys).await.unwrap();
assert_eq!(storage.load_keys().await.unwrap().len(), 1);
storage.restore_from_backup(&backup_path).await.unwrap();
let restored_keys = storage.load_keys().await.unwrap();
assert_eq!(restored_keys.len(), test_keys.len());
for id in test_keys.keys() {
assert!(restored_keys.contains_key(id));
}
}
#[tokio::test]
async fn test_file_storage_backup_nonexistent_storage() {
let temp_dir = TempDir::new().unwrap();
let storage_path = temp_dir.path().join("missing_keys.enc");
let storage = FileStorage::new(storage_path, 0o600, 0o700, false, false)
.await
.unwrap();
fs::remove_file(&storage.path).await.unwrap();
let result = storage.create_backup().await;
assert!(result.is_err());
match result.unwrap_err() {
StorageError::General(msg) => assert!(msg.contains("does not exist")),
_ => panic!("Expected general error"),
}
}
#[tokio::test]
async fn test_file_storage_restore_nonexistent_backup() {
let (storage, temp_dir) = create_test_file_storage().await;
let nonexistent_backup = temp_dir.path().join("nonexistent_backup.enc");
let result = storage.restore_from_backup(&nonexistent_backup).await;
assert!(result.is_err());
match result.unwrap_err() {
StorageError::General(msg) => assert!(msg.contains("does not exist")),
_ => panic!("Expected general error"),
}
}
#[tokio::test]
#[allow(clippy::await_holding_lock)] async fn test_file_storage_cleanup_backups() {
static TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
let _lock = TEST_LOCK.lock().await;
let original_master_key = {
let original = std::env::var("PULSEENGINE_MCP_MASTER_KEY").ok();
unsafe {
std::env::set_var(
"PULSEENGINE_MCP_MASTER_KEY",
"l9EYbalIRp2CF35M4mKcWDqRvx3TFc7U4nX5zvQF56Q",
);
}
original
};
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
let (storage, _temp_dir) = create_test_file_storage().await;
let test_key = create_test_key("cleanup-test", Role::Admin);
storage.save_key(&test_key).await.unwrap();
let mut backup_paths = vec![];
for _i in 0..5 {
let backup_path = storage.create_backup().await.unwrap();
backup_paths.push(backup_path);
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
}
for (i, path) in backup_paths.iter().enumerate() {
assert!(path.exists(), "Backup {} does not exist: {:?}", i, path);
}
storage.cleanup_backups(2).await.unwrap();
let parent = storage.path.parent().unwrap();
let mut remaining_backups = 0;
let mut entries = fs::read_dir(parent).await.unwrap();
while let Some(entry) = entries.next_entry().await.unwrap() {
if entry.file_name().to_string_lossy().contains("backup_") {
remaining_backups += 1;
}
}
assert_eq!(remaining_backups, 2);
unsafe {
match original_master_key {
Some(key) => std::env::set_var("PULSEENGINE_MCP_MASTER_KEY", key),
None => std::env::remove_var("PULSEENGINE_MCP_MASTER_KEY"),
}
}
}
#[cfg(unix)]
#[tokio::test]
async fn test_file_storage_permissions() {
use std::os::unix::fs::PermissionsExt;
let (storage, _temp_dir) = create_test_file_storage().await;
let test_key = create_test_key("perm-test", Role::Operator);
storage.save_key(&test_key).await.unwrap();
let metadata = fs::metadata(&storage.path).await.unwrap();
let mode = metadata.permissions().mode() & 0o777;
assert_eq!(mode, 0o600);
if let Some(parent) = storage.path.parent() {
let parent_metadata = fs::metadata(parent).await.unwrap();
let parent_mode = parent_metadata.permissions().mode() & 0o777;
assert_eq!(parent_mode, 0o700);
}
}
#[tokio::test]
#[allow(clippy::await_holding_lock)] async fn test_file_storage_atomic_operations() {
static TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
let _lock = TEST_LOCK.lock().await;
let original_master_key = {
let original = std::env::var("PULSEENGINE_MCP_MASTER_KEY").ok();
unsafe {
std::env::set_var(
"PULSEENGINE_MCP_MASTER_KEY",
"l9EYbalIRp2CF35M4mKcWDqRvx3TFc7U4nX5zvQF56Q",
);
}
original
};
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
let (storage, _temp_dir) = create_test_file_storage().await;
let initial_keys = create_test_keys();
storage.save_all_keys(&initial_keys).await.unwrap();
let storage_clone = std::sync::Arc::new(storage);
let mut handles = vec![];
for i in 0..5 {
let storage_ref = storage_clone.clone();
let handle = tokio::spawn(async move {
let key = create_test_key(&format!("concurrent-{}", i), Role::Monitor);
tokio::time::timeout(
std::time::Duration::from_secs(10),
storage_ref.save_key(&key),
)
.await
});
handles.push(handle);
}
for handle in handles {
match tokio::time::timeout(std::time::Duration::from_secs(15), handle).await {
Ok(result) => {
match result {
Ok(save_result) => {
let _ = save_result;
}
Err(_) => {
}
}
}
Err(_) => {
eprintln!(
"Concurrent operation timed out - acceptable on slower platforms"
);
}
}
}
let final_keys = storage_clone.load_keys().await.unwrap();
assert!(final_keys.len() >= initial_keys.len());
for id in initial_keys.keys() {
assert!(final_keys.contains_key(id));
}
unsafe {
match original_master_key {
Some(key) => std::env::set_var("PULSEENGINE_MCP_MASTER_KEY", key),
None => std::env::remove_var("PULSEENGINE_MCP_MASTER_KEY"),
}
}
}
}
mod storage_factory_tests {
use super::*;
use crate::config::StorageConfig;
#[tokio::test]
async fn test_create_memory_storage_backend() {
let config = StorageConfig::Memory;
let backend = create_storage_backend(&config).await.unwrap();
let test_key = create_test_key("memory-factory-test", Role::Admin);
backend.save_key(&test_key).await.unwrap();
let keys = backend.load_keys().await.unwrap();
assert_eq!(keys.len(), 1);
assert!(keys.contains_key(&test_key.id));
}
#[tokio::test]
async fn test_create_environment_storage_backend() {
let var_name = "TEST_FACTORY_ENV_STORAGE";
unsafe {
std::env::remove_var(var_name);
}
let config = StorageConfig::Environment {
prefix: var_name.to_string(),
};
let backend = create_storage_backend(&config).await.unwrap();
let test_key = create_test_key("env-factory-test", Role::Operator);
backend.save_key(&test_key).await.unwrap();
let keys = backend.load_keys().await.unwrap();
assert_eq!(keys.len(), 1);
assert!(keys.contains_key(&test_key.id));
unsafe {
std::env::remove_var(var_name);
}
}
#[tokio::test]
async fn test_create_file_storage_backend() {
let temp_dir = TempDir::new().unwrap();
let storage_path = temp_dir.path().join("factory_test_keys.enc");
let config = StorageConfig::File {
path: storage_path.clone(),
file_permissions: 0o600,
dir_permissions: 0o700,
require_secure_filesystem: false,
enable_filesystem_monitoring: false,
};
let backend = create_storage_backend(&config).await.unwrap();
let test_key = create_test_key("file-factory-test", Role::Monitor);
backend.save_key(&test_key).await.unwrap();
let keys = backend.load_keys().await.unwrap();
assert_eq!(keys.len(), 1);
assert!(keys.contains_key(&test_key.id));
assert!(storage_path.exists());
}
#[tokio::test]
async fn test_create_file_storage_backend_with_nested_path() {
let temp_dir = TempDir::new().unwrap();
let storage_path = temp_dir.path().join("nested").join("dirs").join("keys.enc");
let config = StorageConfig::File {
path: storage_path.clone(),
file_permissions: 0o600,
dir_permissions: 0o700,
require_secure_filesystem: false,
enable_filesystem_monitoring: false,
};
let backend = create_storage_backend(&config).await.unwrap();
assert!(storage_path.parent().unwrap().exists());
let test_key = create_test_key(
"nested-factory-test",
Role::Device {
allowed_devices: vec!["device1".to_string()],
},
);
backend.save_key(&test_key).await.unwrap();
let keys = backend.load_keys().await.unwrap();
assert_eq!(keys.len(), 1);
assert!(keys.contains_key(&test_key.id));
}
}
#[tokio::test]
async fn test_storage_backend_trait_object() {
let memory_storage: Box<dyn StorageBackend> = Box::new(MemoryStorage::new());
let env_storage: Box<dyn StorageBackend> =
Box::new(EnvironmentStorage::new("TEST_TRAIT_OBJECT".to_string()));
let storages: Vec<Box<dyn StorageBackend>> = vec![memory_storage, env_storage];
for (i, storage) in storages.into_iter().enumerate() {
let test_key = create_test_key(
&format!("trait-test-{}", i),
Role::Custom {
permissions: vec!["test:read".to_string()],
},
);
storage.save_key(&test_key).await.unwrap();
let keys = storage.load_keys().await.unwrap();
assert_eq!(keys.len(), 1);
assert!(keys.contains_key(&test_key.id));
}
unsafe {
std::env::remove_var("TEST_TRAIT_OBJECT");
}
}
#[tokio::test]
async fn test_secure_api_key_conversion() {
let original_key = create_test_key("conversion-test", Role::Admin);
let secure_key = original_key.to_secure_storage();
let restored_key = secure_key.to_api_key();
assert_eq!(restored_key.id, original_key.id);
assert_eq!(restored_key.name, original_key.name);
assert_eq!(restored_key.role, original_key.role);
assert_eq!(restored_key.created_at, original_key.created_at);
assert_eq!(restored_key.expires_at, original_key.expires_at);
assert_eq!(restored_key.ip_whitelist, original_key.ip_whitelist);
assert_eq!(restored_key.active, original_key.active);
assert_eq!(restored_key.usage_count, original_key.usage_count);
assert_eq!(restored_key.key, "***redacted***");
assert_ne!(restored_key.key, original_key.key);
assert_eq!(restored_key.secret_hash, original_key.secret_hash);
assert_eq!(restored_key.salt, original_key.salt);
}
}