lean_ctx_server/server/

role_guard.rs

// SPDX-License-Identifier: Apache-2.0
//! Role-based tool access guard for the MCP server pipeline.
//!
//! Checks the active role's tool policy before dispatching a tool call.
//! A denial becomes a tool error (`is_error`) carrying the reason, and its
//! audit entry names the agent that made the call.

use rmcp::model::{CallToolResult, ContentBlock};

use lean_ctx::core::roles;

pub struct RoleCheckResult {
    pub blocked: bool,
    pub role_name: String,
    pub message: Option<String>,
}

/// [`check_tool_access_as`] for this process's agent identity
/// (`LEAN_CTX_AGENT_ID`, else `local`), for callers without an MCP session.
pub fn check_tool_access(tool_name: &str) -> RoleCheckResult {
    check_tool_access_as(
        tool_name,
        lean_ctx::core::agent_identity::current_agent_id(),
    )
}

/// Check `tool_name` against the active role; a denial is audited for
/// `agent_id`.
pub fn check_tool_access_as(tool_name: &str, agent_id: &str) -> RoleCheckResult {
    let role_name = roles::active_role_name();
    let role = roles::active_role();

    if tool_name == "ctx_session" || tool_name == "ctx" {
        return RoleCheckResult {
            blocked: false,
            role_name,
            message: None,
        };
    }

    if !role.is_tool_allowed(tool_name) {
        lean_ctx::core::events::emit_policy_violation(
            &role_name,
            tool_name,
            "tool not allowed by role policy",
        );
        lean_ctx::core::audit_trail::record(lean_ctx::core::audit_trail::AuditEntryData {
            agent_id: agent_id.to_string(),
            tool: tool_name.to_string(),
            action: None,
            input_hash: String::new(),
            output_tokens: 0,
            role: role_name.clone(),
            event_type: lean_ctx::core::audit_trail::AuditEventType::ToolDenied,
        });
        let denied_msg = format!(
            "[ROLE DENIED] Tool '{}' is not allowed for role '{}' ({}).\n\
             Allowed tools: {}\n\
             Use `ctx_session` with action `role` to switch roles.",
            tool_name,
            role_name,
            role.role.description,
            if role.tools.allowed.is_empty() || role.tools.allowed.iter().any(|a| a == "*") {
                "* (all except denied)".to_string()
            } else {
                role.tools.allowed.join(", ")
            }
        );
        return RoleCheckResult {
            blocked: true,
            role_name,
            message: Some(denied_msg),
        };
    }

    if is_shell_tool(tool_name) && !role.is_shell_allowed() {
        lean_ctx::core::events::emit_policy_violation(
            &role_name,
            tool_name,
            &format!("shell denied by policy: {}", role.role.shell_policy),
        );
        lean_ctx::core::audit_trail::record(lean_ctx::core::audit_trail::AuditEntryData {
            agent_id: agent_id.to_string(),
            tool: tool_name.to_string(),
            action: None,
            input_hash: String::new(),
            output_tokens: 0,
            role: role_name.clone(),
            event_type: lean_ctx::core::audit_trail::AuditEventType::ToolDenied,
        });
        let msg = format!(
            "[ROLE DENIED] Shell access denied for role '{}'. Shell policy: {}.",
            role_name, role.role.shell_policy
        );
        return RoleCheckResult {
            blocked: true,
            role_name,
            message: Some(msg),
        };
    }

    let cap_result = lean_ctx::core::capabilities::check_capabilities(&role_name, tool_name);
    if !cap_result.allowed {
        let missing_names: Vec<&str> = cap_result
            .missing
            .iter()
            .map(lean_ctx::core::capabilities::Capability::display_name)
            .collect();
        lean_ctx::core::events::emit_policy_violation(
            &role_name,
            tool_name,
            &format!("missing capabilities: {}", missing_names.join(", ")),
        );
        let msg = format!(
            "[CAPABILITY DENIED] Tool '{}' requires capabilities [{}] which role '{}' does not grant.",
            tool_name,
            missing_names.join(", "),
            role_name
        );
        return RoleCheckResult {
            blocked: true,
            role_name,
            message: Some(msg),
        };
    }

    RoleCheckResult {
        blocked: false,
        role_name,
        message: None,
    }
}

/// The denial as a tool error the agent can read; `None` when allowed.
pub fn into_call_tool_result(check: &RoleCheckResult) -> Option<CallToolResult> {
    if check.blocked {
        Some(CallToolResult::error(vec![ContentBlock::text(
            check.message.as_deref().unwrap_or("Blocked by role policy"),
        )]))
    } else {
        None
    }
}

fn is_shell_tool(name: &str) -> bool {
    matches!(name, "ctx_shell" | "ctx_execute")
}

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

    #[test]
    fn session_tool_always_allowed() {
        let result = check_tool_access("ctx_session");
        assert!(!result.blocked);
    }

    #[test]
    fn meta_tool_always_allowed() {
        let result = check_tool_access("ctx");
        assert!(!result.blocked);
    }

    #[test]
    fn coder_role_allows_all() {
        // Other tests change the global role; hold the guard so none can
        // switch it between the set and the check below.
        let _role = roles::lock_global_role();
        let _ = roles::set_active_role("coder");
        let result = check_tool_access("ctx_edit");
        assert!(
            !result.blocked,
            "coder role should allow ctx_edit, active={}",
            result.role_name
        );
        assert_eq!(result.role_name, "coder");
    }

    // A denial was a successful result audited for agent "unknown".
    #[test]
    fn a_denial_is_a_tool_error_audited_for_the_calling_agent() {
        let _data_dir = lean_ctx::core::data_dir::isolated_data_dir();
        let reviewer = roles::load_role("reviewer").expect("built-in reviewer role");
        let check = roles::with_test_active_role(reviewer, || {
            check_tool_access_as("ctx_edit", "role-guard-test-agent")
        });
        assert!(check.blocked);
        let denied = into_call_tool_result(&check).expect("a denial");
        assert_eq!(denied.is_error, Some(true));
        assert!(
            lean_ctx::core::audit_trail::load_recent(50)
                .iter()
                .any(|entry| {
                    entry.agent_id == "role-guard-test-agent"
                        && entry.tool == "ctx_edit"
                        && matches!(
                            entry.event_type,
                            lean_ctx::core::audit_trail::AuditEventType::ToolDenied
                        )
                })
        );
    }

    #[test]
    fn ctx_call_is_not_exempt_from_guard() {
        let _role = roles::lock_global_role();
        let _ = roles::set_active_role("coder");
        let result = check_tool_access("ctx_call");
        assert!(
            !result.blocked,
            "ctx_call itself should be allowed for coder"
        );
    }
}