typesec_rbac/

engine.rs

//! RBAC policy engine — implements [`PolicyEngine`] for [`RbacPolicy`].

use std::collections::{HashMap, HashSet};

use glob::Pattern;
use tracing::debug;
use typesec_core::{
    ResourceId, SubjectId,
    policy::{PolicyEngine, PolicyResult},
};

use crate::model::RbacPolicy;

/// A compiled, fast-lookup RBAC engine.
///
/// After construction from an [`RbacPolicy`], the engine pre-computes:
/// - Effective permissions per role (with inheritance flattened).
/// - Subject → role mappings.
///
/// Every `check()` call does O(roles × patterns) work — fast enough for
/// the sizes of policies used in AI agent deployments.
pub struct RbacEngine {
    /// Subject → set of effective (permission, resource_pattern) pairs.
    subject_grants: HashMap<String, Vec<CompiledGrant>>,
    /// Glob subject pattern → set of effective grants.
    wildcard_subject_grants: Vec<(SubjectPattern, Vec<CompiledGrant>)>,
}

#[derive(Debug, Clone)]
struct Grant {
    permission: String,
    resource_patterns: Vec<String>,
}

/// A grant with its glob patterns validated and compiled once at load time.
///
/// Compiling here (rather than per `check()`) both surfaces pattern typos as
/// load errors — a malformed pattern would otherwise silently never match,
/// i.e. silently deny — and avoids re-parsing the glob on every check.
#[derive(Debug, Clone)]
struct CompiledGrant {
    permission: String,
    resource_patterns: Vec<ResourcePattern>,
}

#[derive(Debug, Clone)]
enum ResourcePattern {
    /// The literal `"*"` — matches every resource, including across `/`.
    Any,
    /// A compiled glob. Note glob `*` does not cross `/` separators:
    /// `reports/*` matches `reports/q1` but not `reports/2024/q1` (use
    /// `reports/**` for that).
    Glob(Pattern),
}

#[derive(Debug, Clone)]
enum SubjectPattern {
    /// The literal `"*"` — matches every subject.
    Any,
    /// A compiled glob.
    Glob(Pattern),
}

impl SubjectPattern {
    fn compile(pattern: &str) -> Result<Self, String> {
        if pattern == "*" {
            return Ok(Self::Any);
        }
        Pattern::new(pattern)
            .map(Self::Glob)
            .map_err(|e| format!("invalid subject pattern '{pattern}': {e}"))
    }

    fn matches(&self, subject: &str) -> bool {
        match self {
            Self::Any => true,
            Self::Glob(pattern) => pattern.matches(subject),
        }
    }
}

impl ResourcePattern {
    fn compile(pattern: &str) -> Result<Self, String> {
        if pattern == "*" {
            return Ok(Self::Any);
        }
        Pattern::new(pattern)
            .map(Self::Glob)
            .map_err(|e| format!("invalid resource pattern '{pattern}': {e}"))
    }

    fn matches(&self, resource: &str) -> bool {
        match self {
            Self::Any => true,
            Self::Glob(pattern) => pattern.matches(resource),
        }
    }
}

impl RbacEngine {
    /// Build an engine from a validated [`RbacPolicy`].
    ///
    /// Returns an error if the policy fails validation.
    pub fn new(policy: RbacPolicy) -> Result<Self, String> {
        policy.validate()?;

        // Step 1: flatten role inheritance into effective (permission, resources) pairs.
        let effective_roles: HashMap<String, Vec<Grant>> = {
            let mut map = HashMap::new();
            for role in &policy.roles {
                let grants = flatten_role(&role.name, &policy);
                map.insert(role.name.clone(), grants);
            }
            map
        };

        // Step 2: build subject → grants mapping, compiling patterns up front
        // so invalid globs fail the policy load instead of silently denying.
        let mut subject_grants: HashMap<String, Vec<CompiledGrant>> = HashMap::new();
        let mut wildcard_subject_grants: Vec<(SubjectPattern, Vec<CompiledGrant>)> = Vec::new();
        for assignment in &policy.assignments {
            let mut all_grants: Vec<CompiledGrant> = Vec::new();
            for role_name in &assignment.roles {
                if let Some(grants) = effective_roles.get(role_name) {
                    for grant in grants {
                        all_grants.push(CompiledGrant {
                            permission: grant.permission.clone(),
                            resource_patterns: grant
                                .resource_patterns
                                .iter()
                                .map(|p| ResourcePattern::compile(p))
                                .collect::<Result<_, _>>()?,
                        });
                    }
                }
            }
            if is_glob_pattern(&assignment.subject) {
                wildcard_subject_grants
                    .push((SubjectPattern::compile(&assignment.subject)?, all_grants));
            } else {
                subject_grants
                    .entry(assignment.subject.clone())
                    .or_default()
                    .extend(all_grants);
            }
        }

        Ok(Self {
            subject_grants,
            wildcard_subject_grants,
        })
    }

    /// Load an engine directly from a YAML string.
    pub fn from_yaml(yaml: &str) -> Result<Self, String> {
        let policy = RbacPolicy::from_yaml(yaml).map_err(|e| format!("YAML parse error: {e}"))?;
        Self::new(policy)
    }
}

impl PolicyEngine for RbacEngine {
    fn check(&self, subject: &SubjectId, action: &str, resource: &ResourceId) -> PolicyResult {
        let subject = subject.as_str();
        let resource = resource.as_str();
        debug!(subject, action, resource, "rbac check");

        let exact_grants = self.subject_grants.get(subject).into_iter().flatten();
        let wildcard_grants = self
            .wildcard_subject_grants
            .iter()
            .filter(|(pattern, _)| pattern.matches(subject))
            .flat_map(|(_, grants)| grants);

        let mut matched_subject = false;
        for grant in exact_grants.chain(wildcard_grants) {
            matched_subject = true;
            if grant.permission == action {
                for pattern in &grant.resource_patterns {
                    if pattern.matches(resource) {
                        return PolicyResult::Allow;
                    }
                }
            }
        }

        if !matched_subject {
            return PolicyResult::Deny(format!("no role assignments for subject '{subject}'"));
        }

        PolicyResult::Deny(format!(
            "no rule grants '{subject}' permission '{action}' on '{resource}'"
        ))
    }
}

fn is_glob_pattern(value: &str) -> bool {
    value.contains(['*', '?', '['])
}

/// Recursively flatten a role's permissions by resolving inheritance.
fn flatten_role(role_name: &str, policy: &RbacPolicy) -> Vec<Grant> {
    let mut seen = HashSet::new();
    flatten_role_inner(role_name, policy, &mut seen)
}

fn flatten_role_inner(
    role_name: &str,
    policy: &RbacPolicy,
    seen: &mut HashSet<String>,
) -> Vec<Grant> {
    if !seen.insert(role_name.to_owned()) {
        return vec![]; // cycle guard (already validated, but be safe)
    }

    let role = match policy.roles.iter().find(|r| r.name == role_name) {
        Some(r) => r,
        None => return vec![],
    };

    let mut grants: Vec<Grant> = Vec::new();

    // Own permissions.
    for perm in &role.permissions {
        grants.push(Grant {
            permission: perm.clone(),
            resource_patterns: role.resources.clone(),
        });
    }

    // Inherited permissions (recursive).
    for parent_name in &role.inherits {
        let inherited = flatten_role_inner(parent_name, policy, seen);
        grants.extend(inherited);
    }

    grants
}

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

    const YAML: &str = r#"
roles:
  - name: analyst
    permissions: [read, read_sensitive]
    resources: ["reports/*", "metrics/*"]
  - name: engineer
    permissions: [read, write, execute]
    resources: ["code/*", "infra/*"]
  - name: admin
    inherits: [analyst, engineer]
    permissions: [delete, delegate]
    resources: ["*"]

assignments:
  - subject: "agent:data-pipeline"
    roles: [analyst]
  - subject: "agent:deploy-bot"
    roles: [engineer]
  - subject: "agent:superuser"
    roles: [admin]
"#;

    fn engine() -> RbacEngine {
        RbacEngine::from_yaml(YAML).expect("engine build should succeed")
    }

    fn check(e: &RbacEngine, subject: &str, action: &str, resource: &str) -> PolicyResult {
        e.check(
            &SubjectId::from(subject),
            action,
            &ResourceId::from(resource),
        )
    }

    #[test]
    fn analyst_can_read_reports() {
        let e = engine();
        assert_eq!(
            check(&e, "agent:data-pipeline", "read", "reports/q1"),
            PolicyResult::Allow
        );
    }

    #[test]
    fn analyst_cannot_write() {
        let e = engine();
        assert!(matches!(
            check(&e, "agent:data-pipeline", "write", "reports/q1"),
            PolicyResult::Deny(_)
        ));
    }

    #[test]
    fn engineer_can_write_code() {
        let e = engine();
        assert_eq!(
            check(&e, "agent:deploy-bot", "write", "code/main.rs"),
            PolicyResult::Allow
        );
    }

    #[test]
    fn engineer_cannot_access_reports() {
        let e = engine();
        assert!(matches!(
            check(&e, "agent:deploy-bot", "read", "reports/q1"),
            PolicyResult::Deny(_)
        ));
    }

    #[test]
    fn admin_inherits_analyst_and_engineer() {
        let e = engine();
        // Inherited from analyst:
        assert_eq!(
            check(&e, "agent:superuser", "read_sensitive", "reports/q1"),
            PolicyResult::Allow
        );
        // Inherited from engineer:
        assert_eq!(
            check(&e, "agent:superuser", "execute", "code/deploy.sh"),
            PolicyResult::Allow
        );
        // Own permissions:
        assert_eq!(
            check(&e, "agent:superuser", "delete", "anything"),
            PolicyResult::Allow
        );
    }

    #[test]
    fn invalid_resource_pattern_fails_policy_load() {
        let yaml = r#"
roles:
  - name: broken
    permissions: [read]
    resources: ["reports/[unclosed"]

assignments:
  - subject: "agent:x"
    roles: [broken]
"#;
        let result = RbacEngine::from_yaml(yaml);
        assert!(
            result.is_err(),
            "malformed glob must fail at load, not silently deny"
        );
    }

    #[test]
    fn unknown_subject_is_denied() {
        let e = engine();
        assert!(matches!(
            check(&e, "agent:ghost", "read", "reports/q1"),
            PolicyResult::Deny(_)
        ));
    }

    #[test]
    fn wildcard_subject_assignment_matches_globbed_agents() {
        let yaml = r#"
roles:
  - name: deployer
    permissions: [execute]
    resources: ["infra/*"]

assignments:
  - subject: "agent:deploy-*"
    roles: [deployer]
"#;
        let e = RbacEngine::from_yaml(yaml).expect("engine build should succeed");
        assert_eq!(
            check(&e, "agent:deploy-prod", "execute", "infra/restart"),
            PolicyResult::Allow
        );
        assert!(matches!(
            check(&e, "agent:build-prod", "execute", "infra/restart"),
            PolicyResult::Deny(_)
        ));
    }

    #[test]
    fn exact_subject_and_wildcard_subject_grants_are_combined() {
        let yaml = r#"
roles:
  - name: reader
    permissions: [read]
    resources: ["reports/*"]
  - name: writer
    permissions: [write]
    resources: ["reports/*"]

assignments:
  - subject: "agent:report-*"
    roles: [reader]
  - subject: "agent:report-prod"
    roles: [writer]
"#;
        let e = RbacEngine::from_yaml(yaml).expect("engine build should succeed");
        assert_eq!(
            check(&e, "agent:report-prod", "read", "reports/q1"),
            PolicyResult::Allow
        );
        assert_eq!(
            check(&e, "agent:report-prod", "write", "reports/q1"),
            PolicyResult::Allow
        );
    }

    #[test]
    fn invalid_subject_pattern_fails_policy_load() {
        let yaml = r#"
roles:
  - name: reader
    permissions: [read]
    resources: ["*"]

assignments:
  - subject: "agent:[broken"
    roles: [reader]
"#;
        let result = RbacEngine::from_yaml(yaml);
        assert!(
            result.is_err(),
            "malformed subject glob must fail at load, not silently deny"
        );
    }

    #[test]
    fn cyclic_role_inheritance_fails_engine_construction() {
        let yaml = r#"
roles:
  - name: a
    inherits: [b]
    permissions: [read]
    resources: ["*"]
  - name: b
    inherits: [a]
    permissions: [write]
    resources: ["*"]

assignments:
  - subject: "agent:x"
    roles: [a]
"#;

        assert!(RbacEngine::from_yaml(yaml).is_err());
    }
}