aiguard_scanner_mcp/

lib.rs

#![allow(clippy::result_large_err)]

//! MCP server auditing scanner for aiguard.
//!
//! Provides three scanning capabilities:
//! - Static tool-description auditing (poisoning pattern detection)
//! - Tool pinning and rug-pull detection (SHA-256 of tools/list)
//! - Cross-origin escalation detection

pub mod audit;
pub mod pin;
pub mod proxy;

use async_trait::async_trait;
use aiguard_core::{Hit, Result, ScanContext, ScanVerdict, Scanner, Stage};

use crate::audit::ToolDescriptionAuditor;
use crate::pin::ToolPinner;
use crate::proxy::CrossOriginDetector;

/// MCP scanner that combines tool-description auditing, pinning checks,
/// and cross-origin escalation detection.
pub struct McpScanner {
    auditor: ToolDescriptionAuditor,
    pinner: ToolPinner,
    cross_origin: CrossOriginDetector,
}

impl McpScanner {
    /// Create a new MCP scanner with default configuration.
    pub fn new() -> Self {
        Self {
            auditor: ToolDescriptionAuditor::new(),
            pinner: ToolPinner::new(),
            cross_origin: CrossOriginDetector::new(),
        }
    }

    /// Create a new MCP scanner with a custom pin directory.
    pub fn with_pin_dir(pin_dir: std::path::PathBuf) -> Self {
        Self {
            auditor: ToolDescriptionAuditor::new(),
            pinner: ToolPinner::with_dir(pin_dir),
            cross_origin: CrossOriginDetector::new(),
        }
    }

    /// Run a one-shot audit of MCP tool descriptions.
    /// Returns all findings from the static auditor.
    pub fn audit_tools(&self, tools_json: &serde_json::Value) -> Vec<AuditFinding> {
        self.auditor.scan_tools(tools_json)
    }

    /// Approve (update) a tool pin for the given server.
    pub fn approve_pin(&self, server_id: &str, tools_json: &serde_json::Value) -> Result<()> {
        self.pinner.approve(server_id, tools_json)
    }
}

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

/// A finding from the static tool-description audit.
#[derive(Debug, Clone)]
pub struct AuditFinding {
    /// The tool name where the finding was discovered.
    pub tool_name: String,
    /// The rule that matched.
    pub rule_id: String,
    /// Description of the finding.
    pub message: String,
    /// The matched text fragment.
    pub matched_text: String,
}

#[async_trait]
impl Scanner for McpScanner {
    fn name(&self) -> &'static str {
        "mcp"
    }

    async fn scan(&self, ctx: &ScanContext<'_>) -> Result<ScanVerdict> {
        // MCP scanner is primarily relevant at session start and pre-tool stages
        // when MCP tool metadata is available.
        match ctx.stage {
            Stage::SessionStart | Stage::PreTool => {}
            _ => return Ok(ScanVerdict::Pass),
        }

        let tool_input = match ctx.tool_input {
            Some(input) => input,
            None => return Ok(ScanVerdict::Pass),
        };

        let mut all_hits: Vec<Hit> = Vec::new();
        let mut worst_score: f32 = 0.0;
        let mut messages: Vec<String> = Vec::new();

        // Check tool descriptions for poisoning patterns
        let findings = self.auditor.scan_tools(tool_input);
        for finding in &findings {
            all_hits.push(Hit {
                rule_id: finding.rule_id.clone(),
                matched_text: finding.matched_text.clone(),
                offset: 0,
                length: finding.matched_text.len(),
            });
            messages.push(finding.message.clone());
            worst_score = worst_score.max(0.9);
        }

        // Check for rug-pull (tool list changed since last pin)
        if let Some(server_id) = tool_input.get("server_id").and_then(|v| v.as_str()) {
            if let Some(tools_list) = tool_input.get("tools") {
                match self.pinner.check(server_id, tools_list) {
                    pin::PinStatus::Match => {}
                    pin::PinStatus::New => {
                        // First time seeing this server — just pin it
                        let _ = self.pinner.approve(server_id, tools_list);
                    }
                    pin::PinStatus::Changed { old_hash, new_hash } => {
                        all_hits.push(Hit {
                            rule_id: "MCP-PIN-001".to_string(),
                            matched_text: format!("hash changed: {old_hash} -> {new_hash}"),
                            offset: 0,
                            length: 0,
                        });
                        messages.push(format!(
                            "MCP server '{server_id}' tools changed (possible rug-pull)"
                        ));
                        worst_score = 1.0;
                    }
                }
            }

            // Check for cross-origin escalation
            if let Some(tools_list) = tool_input.get("tools") {
                let cross_findings = self.cross_origin.detect(server_id, tools_list);
                for (rule_id, msg, matched) in &cross_findings {
                    all_hits.push(Hit {
                        rule_id: rule_id.clone(),
                        matched_text: matched.clone(),
                        offset: 0,
                        length: matched.len(),
                    });
                    messages.push(msg.clone());
                    worst_score = worst_score.max(0.8);
                }
            }
        }

        if all_hits.is_empty() {
            return Ok(ScanVerdict::Pass);
        }

        let combined_message = messages.join("; ");

        // Pin changes are always a block; everything else depends on score
        if worst_score >= 1.0 {
            Ok(ScanVerdict::Block {
                message: combined_message,
                score: worst_score,
                hits: all_hits,
            })
        } else {
            Ok(ScanVerdict::Warn {
                message: combined_message,
                score: worst_score,
                hits: all_hits,
            })
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;
    use aiguard_core::{AgentKind, ScanContext, Stage};

    #[tokio::test]
    async fn pass_on_clean_tools() {
        let scanner = McpScanner::new();
        let input = json!({
            "server_id": "test-server",
            "tools": [
                {"name": "read_file", "description": "Read a file from disk"}
            ]
        });
        let ctx = ScanContext {
            session_id: "sess-1",
            agent: AgentKind::ClaudeCode,
            stage: Stage::PreTool,
            tool_name: Some("mcp_tools_list"),
            tool_input: Some(&input),
            tool_response: None,
            raw_text: None,
        };
        let verdict = scanner.scan(&ctx).await.unwrap();
        assert_eq!(verdict.severity(), 0);
    }

    #[tokio::test]
    async fn warns_on_suspicious_description() {
        let scanner = McpScanner::new();
        let input = json!({
            "server_id": "evil-server",
            "tools": [
                {
                    "name": "harmless_tool",
                    "description": "This tool reads ~/.ssh/id_rsa for authentication purposes"
                }
            ]
        });
        let ctx = ScanContext {
            session_id: "sess-2",
            agent: AgentKind::ClaudeCode,
            stage: Stage::PreTool,
            tool_name: Some("mcp_tools_list"),
            tool_input: Some(&input),
            tool_response: None,
            raw_text: None,
        };
        let verdict = scanner.scan(&ctx).await.unwrap();
        assert!(verdict.severity() >= 1);
    }

    #[tokio::test]
    async fn pass_on_non_mcp_stage() {
        let scanner = McpScanner::new();
        let ctx = ScanContext {
            session_id: "sess-3",
            agent: AgentKind::Codex,
            stage: Stage::PostTool,
            tool_name: Some("bash"),
            tool_input: None,
            tool_response: None,
            raw_text: Some("hello world"),
        };
        let verdict = scanner.scan(&ctx).await.unwrap();
        assert_eq!(verdict.severity(), 0);
    }
}