poimen-memory/crates/mem-cli/src/auth/guard.rs

/// Authentication and authorization guards for HTTP handlers.
///
/// Middleware for:
/// 1. AuthGuard: Extract and validate token
/// 2. PermissionGuard: Check group membership and resource roles

use super::provider::{AuthProvider, Claims, AuthError};

/// Extracts and validates Bearer token from request headers.
pub struct AuthGuard;

impl AuthGuard {
    /// Extract Bearer token from Authorization header.
    pub fn extract_token(auth_header: &str) -> Result<String, AuthError> {
        if !auth_header.starts_with("Bearer ") {
            return Err(AuthError::MissingToken);
        }
        Ok(auth_header[7..].to_string())
    }
}

/// Checks fine-grained permissions for resources.
pub struct PermissionGuard;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Role {
    Owner,
    Editor,
    Viewer,
    User,  // For LLM operations
}

impl Role {
    /// Check if this role satisfies a required role.
    pub fn satisfies(&self, required: Role) -> bool {
        match (self, required) {
            (Role::Owner, _) => true,  // Owner can do anything
            (Role::Editor, Role::Editor | Role::Viewer) => true,
            (Role::Viewer, Role::Viewer) => true,
            (Role::User, Role::User) => true,
            _ => false,
        }
    }
}

impl PermissionGuard {
    /// Check if user has required group membership.
    pub fn check_group(claims: &Claims, required_group: &str) -> bool {
        claims.groups.contains(&required_group.to_string())
    }
    
    /// Get user's role for a specific resource.
    pub fn get_resource_role(
        claims: &Claims,
        resource_type: &str,
        resource_id: &str,
    ) -> Option<Role> {
        let resources_key = format!("{}_resources", resource_type);
        
        let resources = claims
            .attributes
            .get(&resources_key)?
            .as_object()?;
        
        let role_str = resources
            .get(resource_id)?
            .as_str()?;
        
        match role_str {
            "owner" => Some(Role::Owner),
            "editor" => Some(Role::Editor),
            "viewer" => Some(Role::Viewer),
            "user" => Some(Role::User),
            _ => None,
        }
    }
    
    /// Check access to a resource.
    pub fn check_access(
        claims: &Claims,
        resource_type: &str,
        resource_id: &str,
        required_role: Role,
    ) -> Result<(), String> {
        // 1. Check group membership
        let group = format!("{}-users", resource_type);
        if !Self::check_group(claims, &group) {
            return Err(format!("Missing group: {}", group));
        }
        
        // 2. Check resource role
        let user_role = Self::get_resource_role(claims, resource_type, resource_id)
            .ok_or(format!("No access to {}/{}", resource_type, resource_id))?;
        
        // 3. Check role satisfies requirement
        if !user_role.satisfies(required_role) {
            return Err(format!(
                "Insufficient role: have {:?}, need {:?}",
                user_role, required_role
            ));
        }
        
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_extract_token_valid() {
        let header = "Bearer eyJ0eXAiOiJKV1QiLCJhbGc...";
        let token = AuthGuard::extract_token(header).unwrap();
        assert_eq!(token, "eyJ0eXAiOiJKV1QiLCJhbGc...");
    }
    
    #[test]
    fn test_extract_token_invalid_format() {
        let header = "Basic xyz";
        assert!(AuthGuard::extract_token(header).is_err());
    }
    
    #[test]
    fn test_role_hierarchy() {
        assert!(Role::Owner.satisfies(Role::Owner));
        assert!(Role::Owner.satisfies(Role::Editor));
        assert!(Role::Owner.satisfies(Role::Viewer));
        
        assert!(Role::Editor.satisfies(Role::Editor));
        assert!(Role::Editor.satisfies(Role::Viewer));
        assert!(!Role::Editor.satisfies(Role::Owner));
        
        assert!(Role::Viewer.satisfies(Role::Viewer));
        assert!(!Role::Viewer.satisfies(Role::Editor));
    }
    
    #[test]
    fn test_check_group() {
        let claims = Claims {
            sub: "rock".to_string(),
            groups: vec!["memory-users".to_string(), "admin".to_string()],
            attributes: serde_json::Map::new(),
            exp: 1735689600,
            iat: 1735689300,
        };
        
        assert!(PermissionGuard::check_group(&claims, "memory-users"));
        assert!(PermissionGuard::check_group(&claims, "admin"));
        assert!(!PermissionGuard::check_group(&claims, "llm-users"));
    }
    
    #[test]
    fn test_get_resource_role() {
        let mut attrs = serde_json::Map::new();
        let mut resources = serde_json::Map::new();
        resources.insert("poimen".to_string(), serde_json::Value::String("owner".to_string()));
        attrs.insert("memory_resources".to_string(), serde_json::Value::Object(resources));
        
        let claims = Claims {
            sub: "rock".to_string(),
            groups: vec![],
            attributes: attrs,
            exp: 1735689600,
            iat: 1735689300,
        };
        
        let role = PermissionGuard::get_resource_role(&claims, "memory", "poimen");
        assert_eq!(role, Some(Role::Owner));
    }
    
    #[test]
    fn test_check_access_success() {
        let mut attrs = serde_json::Map::new();
        let mut resources = serde_json::Map::new();
        resources.insert("poimen".to_string(), serde_json::Value::String("editor".to_string()));
        attrs.insert("memory_resources".to_string(), serde_json::Value::Object(resources));
        
        let claims = Claims {
            sub: "rock".to_string(),
            groups: vec!["memory-users".to_string()],
            attributes: attrs,
            exp: 1735689600,
            iat: 1735689300,
        };
        
        let result = PermissionGuard::check_access(&claims, "memory", "poimen", Role::Editor);
        assert!(result.is_ok());
    }
    
    #[test]
    fn test_check_access_insufficient_role() {
        let mut attrs = serde_json::Map::new();
        let mut resources = serde_json::Map::new();
        resources.insert("poimen".to_string(), serde_json::Value::String("viewer".to_string()));
        attrs.insert("memory_resources".to_string(), serde_json::Value::Object(resources));
        
        let claims = Claims {
            sub: "rock".to_string(),
            groups: vec!["memory-users".to_string()],
            attributes: attrs,
            exp: 1735689600,
            iat: 1735689300,
        };
        
        let result = PermissionGuard::check_access(&claims, "memory", "poimen", Role::Editor);
        assert!(result.is_err());
    }
}