//! 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 //! (`__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>, } #[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 { 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 { self.inner.collection("findings") } pub fn sbom_entries(&self) -> Collection { self.inner.collection("sbom_entries") } pub fn dast_findings(&self) -> Collection { self.inner.collection("dast_findings") } pub fn dast_scan_runs(&self) -> Collection { self.inner.collection("dast_scan_runs") } pub fn pentest_sessions(&self) -> Collection { self.inner.collection("pentest_sessions") } pub fn attack_chain_nodes(&self) -> Collection { self.inner.collection("attack_chain_nodes") } pub fn pentest_messages(&self) -> Collection { self.inner.collection("pentest_messages") } }