rust_rbac/storage/

memory.rs

use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock};
use async_trait::async_trait;

use crate::error::RbacError;
use crate::models::{Permission, Role};
use crate::storage::traits::RbacStorage;

/// In-memory implementation of RbacStorage
pub struct MemoryStorage {
    permissions: Arc<RwLock<HashMap<String, Permission>>>,
    roles: Arc<RwLock<HashMap<String, Role>>>,
    role_permissions: Arc<RwLock<HashMap<String, HashSet<String>>>>,
    subject_roles: Arc<RwLock<HashMap<String, HashSet<String>>>>,
    subject_permissions: Arc<RwLock<HashMap<String, HashSet<String>>>>,
}

impl MemoryStorage {
    pub fn new() -> Self {
        Self {
            permissions: Arc::new(RwLock::new(HashMap::new())),
            roles: Arc::new(RwLock::new(HashMap::new())),
            role_permissions: Arc::new(RwLock::new(HashMap::new())),
            subject_roles: Arc::new(RwLock::new(HashMap::new())),
            subject_permissions: Arc::new(RwLock::new(HashMap::new())),
        }
    }
}

impl Default for MemoryStorage {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl RbacStorage for MemoryStorage {
    async fn create_permission(&self, permission: &Permission) -> Result<(), RbacError> {
        let mut permissions = self.permissions.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on permissions".to_string())
        })?;
        
        if permissions.contains_key(&permission.name) {
            return Err(RbacError::PermissionAlreadyExists(permission.name.clone()));
        }
        
        permissions.insert(permission.name.clone(), permission.clone());
        Ok(())
    }
    
    async fn get_permission(&self, name: &str) -> Result<Option<Permission>, RbacError> {
        let permissions = self.permissions.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on permissions".to_string())
        })?;
        
        Ok(permissions.get(name).cloned())
    }
    
    async fn delete_permission(&self, name: &str) -> Result<(), RbacError> {
        let mut permissions = self.permissions.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on permissions".to_string())
        })?;
        
        if permissions.remove(name).is_none() {
            return Err(RbacError::PermissionNotFound(name.to_string()));
        }
        
        // Remove this permission from all roles
        let mut role_permissions = self.role_permissions.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on role_permissions".to_string())
        })?;
        
        for permissions in role_permissions.values_mut() {
            permissions.remove(name);
        }
        
        // Remove this permission from all subjects
        let mut subject_permissions = self.subject_permissions.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on subject_permissions".to_string())
        })?;
        
        for permissions in subject_permissions.values_mut() {
            permissions.remove(name);
        }
        
        Ok(())
    }
    
    async fn create_role(&self, role: &Role) -> Result<(), RbacError> {
        let mut roles = self.roles.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on roles".to_string())
        })?;
        
        if roles.contains_key(&role.name) {
            return Err(RbacError::RoleAlreadyExists(role.name.clone()));
        }
        
        roles.insert(role.name.clone(), role.clone());
        
        // Initialize empty permission set for this role
        let mut role_permissions = self.role_permissions.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on role_permissions".to_string())
        })?;
        
        role_permissions.insert(role.name.clone(), HashSet::new());
        
        Ok(())
    }
    
    async fn get_role(&self, name: &str) -> Result<Option<Role>, RbacError> {
        let roles = self.roles.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on roles".to_string())
        })?;
        
        Ok(roles.get(name).cloned())
    }
    
    async fn delete_role(&self, name: &str) -> Result<(), RbacError> {
        let mut roles = self.roles.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on roles".to_string())
        })?;
        
        if roles.remove(name).is_none() {
            return Err(RbacError::RoleNotFound(name.to_string()));
        }
        
        // Remove this role's permissions
        let mut role_permissions = self.role_permissions.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on role_permissions".to_string())
        })?;
        
        role_permissions.remove(name);
        
        // Remove this role from all subjects
        let mut subject_roles = self.subject_roles.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on subject_roles".to_string())
        })?;
        
        for roles in subject_roles.values_mut() {
            roles.remove(name);
        }
        
        Ok(())
    }
    
    async fn assign_permission_to_role(&self, permission_name: &str, role_name: &str) -> Result<(), RbacError> {
        // Check if permission exists
        let permissions = self.permissions.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on permissions".to_string())
        })?;
        
        if !permissions.contains_key(permission_name) {
            return Err(RbacError::PermissionNotFound(permission_name.to_string()));
        }
        
        // Check if role exists
        let roles = self.roles.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on roles".to_string())
        })?;
        
        if !roles.contains_key(role_name) {
            return Err(RbacError::RoleNotFound(role_name.to_string()));
        }
        
        // Assign permission to role
        let mut role_permissions = self.role_permissions.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on role_permissions".to_string())
        })?;
        
        role_permissions
            .entry(role_name.to_string())
            .or_insert_with(HashSet::new)
            .insert(permission_name.to_string());
        
        Ok(())
    }
    
    async fn remove_permission_from_role(&self, permission_name: &str, role_name: &str) -> Result<(), RbacError> {
        // Check if role exists
        let roles = self.roles.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on roles".to_string())
        })?;
        
        if !roles.contains_key(role_name) {
            return Err(RbacError::RoleNotFound(role_name.to_string()));
        }
        
        // Remove permission from role
        let mut role_permissions = self.role_permissions.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on role_permissions".to_string())
        })?;
        
        if let Some(permissions) = role_permissions.get_mut(role_name) {
            permissions.remove(permission_name);
        }
        
        Ok(())
    }
    
    async fn get_permissions_for_role(&self, role_name: &str) -> Result<Vec<Permission>, RbacError> {
        // Check if role exists
        let roles = self.roles.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on roles".to_string())
        })?;
        
        if !roles.contains_key(role_name) {
            return Err(RbacError::RoleNotFound(role_name.to_string()));
        }
        
        // Get permissions for role
        let role_permissions = self.role_permissions.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on role_permissions".to_string())
        })?;
        
        let permissions_read = self.permissions.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on permissions".to_string())
        })?;
        
        let permission_names = role_permissions.get(role_name).map_or_else(
            || Vec::new(),
            |names| names.iter().cloned().collect::<Vec<_>>()
        );
        
        let mut result = Vec::new();
        for name in permission_names {
            if let Some(permission) = permissions_read.get(&name) {
                result.push(permission.clone());
            }
        }
        
        Ok(result)
    }
    
    async fn assign_role_to_subject(&self, role_name: &str, subject_id: &str) -> Result<(), RbacError> {
        // Check if role exists
        let roles = self.roles.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on roles".to_string())
        })?;
        
        if !roles.contains_key(role_name) {
            return Err(RbacError::RoleNotFound(role_name.to_string()));
        }
        
        // Assign role to subject
        let mut subject_roles = self.subject_roles.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on subject_roles".to_string())
        })?;
        
        subject_roles
            .entry(subject_id.to_string())
            .or_insert_with(HashSet::new)
            .insert(role_name.to_string());
        
        Ok(())
    }
    
    async fn remove_role_from_subject(&self, role_name: &str, subject_id: &str) -> Result<(), RbacError> {
        // Remove role from subject
        let mut subject_roles = self.subject_roles.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on subject_roles".to_string())
        })?;
        
        if let Some(roles) = subject_roles.get_mut(subject_id) {
            roles.remove(role_name);
        }
        
        Ok(())
    }
    
    async fn get_roles_for_subject(&self, subject_id: &str) -> Result<Vec<Role>, RbacError> {
        // Get roles for subject
        let subject_roles = self.subject_roles.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on subject_roles".to_string())
        })?;
        
        let roles_read = self.roles.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on roles".to_string())
        })?;
        
        let role_names = subject_roles.get(subject_id).map_or_else(
            || Vec::new(),
            |names| names.iter().cloned().collect::<Vec<_>>()
        );
        
        let mut result = Vec::new();
        for name in role_names {
            if let Some(role) = roles_read.get(&name) {
                result.push(role.clone());
            }
        }
        
        Ok(result)
    }
    
    async fn assign_permission_to_subject(&self, permission_name: &str, subject_id: &str) -> Result<(), RbacError> {
        // Check if permission exists
        let permissions = self.permissions.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on permissions".to_string())
        })?;
        
        if !permissions.contains_key(permission_name) {
            return Err(RbacError::PermissionNotFound(permission_name.to_string()));
        }
        
        // Assign permission to subject
        let mut subject_permissions = self.subject_permissions.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on subject_permissions".to_string())
        })?;
        
        subject_permissions
            .entry(subject_id.to_string())
            .or_insert_with(HashSet::new)
            .insert(permission_name.to_string());
        
        Ok(())
    }
    
    async fn remove_permission_from_subject(&self, permission_name: &str, subject_id: &str) -> Result<(), RbacError> {
        // Remove permission from subject
        let mut subject_permissions = self.subject_permissions.write().map_err(|_| {
            RbacError::StorageError("Failed to acquire write lock on subject_permissions".to_string())
        })?;
        
        if let Some(permissions) = subject_permissions.get_mut(subject_id) {
            permissions.remove(permission_name);
        }
        
        Ok(())
    }
    
    async fn get_direct_permissions_for_subject(&self, subject_id: &str) -> Result<Vec<Permission>, RbacError> {
        // Get direct permissions for subject
        let subject_permissions = self.subject_permissions.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on subject_permissions".to_string())
        })?;
        
        let permissions_read = self.permissions.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on permissions".to_string())
        })?;
        
        let permission_names = subject_permissions.get(subject_id).map_or_else(
            || Vec::new(),
            |names| names.iter().cloned().collect::<Vec<_>>()
        );
        
        let mut result = Vec::new();
        for name in permission_names {
            if let Some(permission) = permissions_read.get(&name) {
                result.push(permission.clone());
            }
        }
        
        Ok(result)
    }
    
    async fn subject_has_permission(&self, subject_id: &str, permission_name: &str) -> Result<bool, RbacError> {
        // Check direct permissions
        let subject_permissions = self.subject_permissions.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on subject_permissions".to_string())
        })?;
        
        if let Some(permissions) = subject_permissions.get(subject_id) {
            if permissions.contains(permission_name) {
                return Ok(true);
            }
        }
        
        // Check permissions via roles
        let subject_roles = self.subject_roles.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on subject_roles".to_string())
        })?;
        
        let role_permissions = self.role_permissions.read().map_err(|_| {
            RbacError::StorageError("Failed to acquire read lock on role_permissions".to_string())
        })?;
        
        if let Some(roles) = subject_roles.get(subject_id) {
            for role_name in roles {
                if let Some(permissions) = role_permissions.get(role_name) {
                    if permissions.contains(permission_name) {
                        return Ok(true);
                    }
                }
            }
        }
        
        Ok(false)
    }
}