compliance-scanner-agent/compliance-mcp/src/database.rs

//! Per-tenant Mongo broker for the MCP server.
//!
//! Mirror of the agent's `compliance_agent::database::DatabasePool` —
//! duplicated here rather than lifted into `compliance-core` to keep
//! this PR focused. If a third consumer ever needs it, lift then.
//!
//! Bearer tokens (validated by the auth middleware) carry a tenant_id
//! and the handler resolves the per-tenant database via
//! [`DatabasePool::for_tenant_id`]. The admin database
//! (`<db_prefix>__admin`) holds the cross-tenant `mcp_tokens`
//! collection that the middleware queries on every request.

use std::sync::Arc;

use dashmap::DashMap;
use mongodb::{bson::doc, Client, Collection};
use sha2::{Digest, Sha256};

use compliance_core::models::*;

/// 63-byte Mongo db-name cap; same invariant as the agent's pool.
const MAX_DB_NAME_LEN: usize = 63;
/// 16-byte SHA-256 truncation, hex-encoded → 32 chars.
const HASH_HEX_LEN: usize = 32;
const MAX_PREFIX_LEN: usize = MAX_DB_NAME_LEN - 1 - HASH_HEX_LEN;

#[derive(Clone, Debug)]
pub struct DatabasePool {
    client: Client,
    db_prefix: String,
    /// Tenants we've handed out a [`Database`] for. The MCP server
    /// doesn't ensure indexes (the agent owns that side of the
    /// schema), so the marker exists only to satisfy the parallel
    /// shape — current code never reads it.
    #[allow(dead_code)]
    seen: Arc<DashMap<String, ()>>,
}

#[derive(Debug, thiserror::Error)]
pub enum DbError {
    #[error("db_prefix '{prefix}' is {len} chars; max is {max} so the hash-fallback DB name fits Mongo's 63-byte cap")]
    PrefixTooLong {
        prefix: String,
        len: usize,
        max: usize,
    },
    #[error(transparent)]
    Mongo(#[from] mongodb::error::Error),
}

impl DatabasePool {
    pub async fn connect(uri: &str, db_prefix: &str) -> Result<Self, DbError> {
        if db_prefix.len() > MAX_PREFIX_LEN {
            return Err(DbError::PrefixTooLong {
                prefix: db_prefix.to_string(),
                len: db_prefix.len(),
                max: MAX_PREFIX_LEN,
            });
        }
        let client = Client::with_uri_str(uri).await?;
        client
            .database("admin")
            .run_command(doc! { "ping": 1 })
            .await?;
        tracing::info!(
            "MCP MongoDB cluster reachable; per-tenant pool ready (db prefix '{db_prefix}')"
        );
        Ok(Self {
            client,
            db_prefix: db_prefix.to_string(),
            seen: Arc::new(DashMap::new()),
        })
    }

    /// Read-only handle to the tenant's database. No indexes are
    /// ensured here — the agent owns writes, MCP only reads.
    pub fn for_tenant_id(&self, tenant_id: &str) -> Database {
        let db_name = self.tenant_db_name(tenant_id);
        self.seen.insert(tenant_id.to_string(), ());
        Database::new(self.client.database(&db_name))
    }

    /// Cross-tenant admin DB — holds the `mcp_tokens` collection that
    /// the auth middleware queries to map bearer → tenant_id.
    pub fn admin_db(&self) -> mongodb::Database {
        self.client.database(&format!("{}__admin", self.db_prefix))
    }

    pub fn tenant_db_name(&self, tenant_id: &str) -> String {
        let sanitized = sanitize_tenant_id(tenant_id);
        let natural = format!("{}_{}", self.db_prefix, sanitized);
        if natural.len() <= MAX_DB_NAME_LEN {
            natural
        } else {
            let mut h = Sha256::new();
            h.update(tenant_id.as_bytes());
            let digest = h.finalize();
            let suffix = hex::encode(&digest[..HASH_HEX_LEN / 2]);
            format!("{}_{}", self.db_prefix, suffix)
        }
    }
}

fn sanitize_tenant_id(tenant_id: &str) -> String {
    tenant_id
        .chars()
        .map(|c| match c {
            '/' | '\\' | '.' | '"' | '$' | ' ' | '\0' => '_',
            c => c,
        })
        .collect()
}

/// Typed accessors for the MCP-readable collections in a tenant DB.
/// Matches the agent's `Database` shape but only exposes what the MCP
/// tool handlers actually need.
#[derive(Clone, Debug)]
pub struct Database {
    inner: mongodb::Database,
}

impl Database {
    pub(crate) fn new(inner: mongodb::Database) -> Self {
        Self { inner }
    }

    pub fn findings(&self) -> Collection<Finding> {
        self.inner.collection("findings")
    }

    pub fn sbom_entries(&self) -> Collection<SbomEntry> {
        self.inner.collection("sbom_entries")
    }

    pub fn dast_findings(&self) -> Collection<DastFinding> {
        self.inner.collection("dast_findings")
    }

    pub fn dast_scan_runs(&self) -> Collection<DastScanRun> {
        self.inner.collection("dast_scan_runs")
    }

    pub fn pentest_sessions(&self) -> Collection<PentestSession> {
        self.inner.collection("pentest_sessions")
    }

    pub fn attack_chain_nodes(&self) -> Collection<AttackChainNode> {
        self.inner.collection("attack_chain_nodes")
    }

    pub fn pentest_messages(&self) -> Collection<PentestMessage> {
        self.inner.collection("pentest_messages")
    }
}