agentsec_core/registry/

mod.rs

//! Known-good MCP server registry.
//!
//! A registry is a mapping `name → source URL` for MCP servers the user
//! has signalled are safe (because they came from a curated public list,
//! a personal allowlist, or a cached pull). It feeds the BlackList check
//! in [`crate::scan::unknown::classify`]:
//!
//! - **Exact match** of an installed MCP server name against a registry
//!   entry ⇒ `KnownGood`.
//! - **Levenshtein distance ≤ 2** to some entry (and not identical) ⇒
//!   `Typosquat` — high false-positive risk if the user named their own
//!   server something close, but worth surfacing for review.
//! - **No match** ⇒ `Unknown` — neutral verdict, just "AgentSec doesn't
//!   recognise this name".
//!
//! ## Source chain (demo-level)
//!
//! Phase 0 is a stub for the Service-side signature DB that Phase 1+
//! would build. For now the resolution order is:
//!
//! 1. **Cache** — `<paths.home>/registry.json`, if present.
//! 2. **Network** — fetch from [`DEFAULT_SOURCE_URL`] (or a caller-provided
//!    URL) when explicitly requested. This is opt-in to keep `scan` and
//!    `blacklist` runs deterministic by default.
//! 3. **Builtin fallback** — a hardcoded list of well-known
//!    `modelcontextprotocol/servers` entries. Always available, never
//!    fails.
//!
//! The wire format for the network endpoint is a flat JSON array of
//! `{"name": "...", "source": "..."}` objects.

use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use std::time::Duration;

use serde::{Deserialize, Serialize};

use crate::Paths;
use crate::error::{Error, Result};

/// Public demo endpoint. Open-source / read-only; the builtin fallback
/// keeps the binary working when this URL is unreachable.
pub const DEFAULT_SOURCE_URL: &str =
    "https://raw.githubusercontent.com/modelcontextprotocol/servers/main/.well-known/registry.json";

const FETCH_TIMEOUT_SECS: u64 = 10;
const USER_AGENT: &str = concat!("agentsec-registry/", env!("CARGO_PKG_VERSION"));

/// One registry entry: a known-good MCP server name and where it came from.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RegistryEntry {
    /// Server name as it would appear under `mcpServers.<name>` in
    /// `.mcp.json` or `.claude.json`.
    pub name: String,
    /// Free-form provenance string (URL, "builtin", filename, etc.). Echoed
    /// back through [`crate::scan::unknown::UnknownVerdict::reason`] so the
    /// user can audit why a name was approved.
    pub source: String,
}

/// In-memory known-good registry. Cheap to clone (entries live in a
/// `HashMap`, keys are short ASCII names).
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Registry {
    /// All entries keyed by `name` for O(1) exact-match lookup.
    pub entries: HashMap<String, RegistryEntry>,
}

impl Registry {
    /// Hardcoded fallback list — well-known `modelcontextprotocol/servers`
    /// entries. Always succeeds, never touches the network or filesystem.
    pub fn builtin() -> Self {
        const BUILTIN: &[(&str, &str)] = &[
            ("filesystem", "modelcontextprotocol/servers"),
            ("github", "modelcontextprotocol/servers"),
            ("gitlab", "modelcontextprotocol/servers"),
            ("git", "modelcontextprotocol/servers"),
            ("memory", "modelcontextprotocol/servers"),
            ("fetch", "modelcontextprotocol/servers"),
            ("puppeteer", "modelcontextprotocol/servers"),
            ("sequential-thinking", "modelcontextprotocol/servers"),
            ("everything", "modelcontextprotocol/servers"),
            ("brave-search", "modelcontextprotocol/servers"),
            ("google-maps", "modelcontextprotocol/servers"),
            ("slack", "modelcontextprotocol/servers"),
            ("sqlite", "modelcontextprotocol/servers"),
            ("postgres", "modelcontextprotocol/servers"),
            ("time", "modelcontextprotocol/servers"),
        ];
        let entries = BUILTIN
            .iter()
            .map(|(name, source)| {
                (
                    (*name).to_string(),
                    RegistryEntry {
                        name: (*name).to_string(),
                        source: (*source).to_string(),
                    },
                )
            })
            .collect();
        Self { entries }
    }

    /// Number of entries. Convenience for callers that want a one-line
    /// "registry has N known-good entries" message.
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// `true` if the registry has zero entries.
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Look up an exact name match.
    pub fn get(&self, name: &str) -> Option<&RegistryEntry> {
        self.entries.get(name)
    }

    /// Find the closest entry to `name` by Levenshtein distance, returning
    /// `Some((entry, distance))` only when the distance is in `1..=max`
    /// (zero would be an exact match — not a typosquat).
    pub fn closest(&self, name: &str, max: usize) -> Option<(&RegistryEntry, usize)> {
        let mut best: Option<(&RegistryEntry, usize)> = None;
        for entry in self.entries.values() {
            let d = levenshtein(name, &entry.name);
            if d == 0 || d > max {
                continue;
            }
            best = match best {
                Some((_, bd)) if bd <= d => best,
                _ => Some((entry, d)),
            };
        }
        best
    }

    /// Fetch the registry from a JSON endpoint serving a
    /// `[{"name":..., "source":...}, ...]` array.
    ///
    /// # Errors
    ///
    /// - [`crate::Error::Http`] on network / TLS / non-2xx response.
    /// - [`crate::Error::Json`] on payload shape mismatch.
    pub async fn fetch(url: &str) -> Result<Self> {
        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(FETCH_TIMEOUT_SECS))
            .user_agent(USER_AGENT)
            .build()?;
        let resp = client.get(url).send().await?;
        if !resp.status().is_success() {
            return Err(Error::Sanitize(format!(
                "registry fetch {url} returned HTTP {}",
                resp.status()
            )));
        }
        let entries: Vec<RegistryEntry> = resp.json().await?;
        Ok(Self::from_entries(entries))
    }

    /// Build a registry from an arbitrary entries list. Useful for tests
    /// and for parsing locally-managed allowlists.
    pub fn from_entries(entries: Vec<RegistryEntry>) -> Self {
        let entries = entries.into_iter().map(|e| (e.name.clone(), e)).collect();
        Self { entries }
    }

    /// Path of the local cache file under [`Paths::home`].
    pub fn cache_path(paths: &Paths) -> PathBuf {
        paths.home.join("registry.json")
    }

    /// Load the cached registry if [`Self::cache_path`] exists, else
    /// `Ok(None)`.
    ///
    /// # Errors
    ///
    /// - [`crate::Error::Io`] on read failure.
    /// - [`crate::Error::Json`] on parse failure.
    pub fn load_cached(paths: &Paths) -> Result<Option<Self>> {
        let path = Self::cache_path(paths);
        if !path.exists() {
            return Ok(None);
        }
        let body = fs::read_to_string(&path)?;
        let reg: Self = serde_json::from_str(&body)?;
        Ok(Some(reg))
    }

    /// Persist `self` to [`Self::cache_path`].
    ///
    /// # Errors
    ///
    /// - [`crate::Error::Io`] on dir / file write failure.
    /// - [`crate::Error::Json`] on serialization failure.
    pub fn save_cached(&self, paths: &Paths) -> Result<PathBuf> {
        fs::create_dir_all(&paths.home)?;
        let path = Self::cache_path(paths);
        let body = serde_json::to_string_pretty(self)?;
        fs::write(&path, body)?;
        Ok(path)
    }

    /// Resolve a registry through the source chain documented at the
    /// module level: cache → builtin. Network fetch is **not** triggered
    /// here — callers that want a fresh pull should call
    /// [`Self::fetch`] explicitly and then [`Self::save_cached`].
    pub fn load_or_builtin(paths: &Paths) -> Self {
        Self::load_cached(paths)
            .ok()
            .flatten()
            .unwrap_or_else(Self::builtin)
    }

    // ── per-user local allowlist ──────────────────────────────────────────

    /// Path of the per-user local allowlist under [`Paths::home`].
    pub fn local_path(paths: &Paths) -> PathBuf {
        paths.home.join("registry-local.json")
    }

    /// Load the local allowlist.  File absent → empty [`Registry`]; parse
    /// error → propagate [`crate::Error`].
    ///
    /// # Errors
    ///
    /// - [`crate::Error::Io`] on read failure (file exists but unreadable).
    /// - [`crate::Error::Json`] on parse failure.
    pub fn load_local(paths: &Paths) -> Result<Self> {
        let path = Self::local_path(paths);
        if !path.exists() {
            return Ok(Self::default());
        }
        let body = fs::read_to_string(&path)?;
        let reg: Self = serde_json::from_str(&body)?;
        Ok(reg)
    }

    /// Persist `self` to [`Self::local_path`].
    ///
    /// # Errors
    ///
    /// - [`crate::Error::Io`] on dir / file write failure.
    /// - [`crate::Error::Json`] on serialization failure.
    pub fn save_local(&self, paths: &Paths) -> Result<()> {
        fs::create_dir_all(&paths.home)?;
        let path = Self::local_path(paths);
        let body = serde_json::to_string_pretty(self)?;
        fs::write(&path, body)?;
        Ok(())
    }

    /// Add (or upsert) a single entry to the local allowlist.
    ///
    /// `source` defaults to `"local"` when `None`.
    ///
    /// # Errors
    ///
    /// Propagates I/O or JSON errors from [`Self::load_local`] /
    /// [`Self::save_local`].
    pub fn add_local(paths: &Paths, name: &str, source: Option<&str>) -> Result<()> {
        let mut reg = Self::load_local(paths)?;
        let entry = RegistryEntry {
            name: name.to_string(),
            source: source.unwrap_or("local").to_string(),
        };
        reg.entries.insert(name.to_string(), entry);
        reg.save_local(paths)
    }

    /// Remove a single entry from the local allowlist.  No-op when the
    /// entry is absent.
    ///
    /// # Errors
    ///
    /// Propagates I/O or JSON errors from [`Self::load_local`] /
    /// [`Self::save_local`].
    pub fn remove_local(paths: &Paths, name: &str) -> Result<()> {
        let mut reg = Self::load_local(paths)?;
        reg.entries.remove(name);
        reg.save_local(paths)
    }

    /// Resolve the three-source chain: `builtin < fetched-cache < local`.
    ///
    /// When the same name appears in multiple sources, the highest-priority
    /// source wins (local > fetched > builtin).
    ///
    /// # Errors
    ///
    /// Propagates I/O or JSON errors from loading the fetched cache or the
    /// local allowlist.  The builtin layer never fails.
    pub fn load_resolved(paths: &Paths) -> Result<Self> {
        // Start from the builtin base.
        let mut merged = Self::builtin();
        // Layer 2: fetched cache (overwrites builtin on name collision).
        if let Some(fetched) = Self::load_cached(paths)? {
            for (name, entry) in fetched.entries {
                merged.entries.insert(name, entry);
            }
        }
        // Layer 3: local allowlist (highest priority).
        let local = Self::load_local(paths)?;
        for (name, entry) in local.entries {
            merged.entries.insert(name, entry);
        }
        Ok(merged)
    }
}

/// Plain Levenshtein distance over byte-strings. Lowercases both inputs so
/// `"GitHub"` and `"github"` collide at distance 0. Names are ASCII in
/// practice, so byte-wise distance equals codepoint distance.
fn levenshtein(a: &str, b: &str) -> usize {
    let a = a.to_lowercase();
    let b = b.to_lowercase();
    let av: Vec<u8> = a.bytes().collect();
    let bv: Vec<u8> = b.bytes().collect();
    let (n, m) = (av.len(), bv.len());
    if n == 0 {
        return m;
    }
    if m == 0 {
        return n;
    }
    let mut prev: Vec<usize> = (0..=m).collect();
    let mut curr: Vec<usize> = vec![0; m + 1];
    for i in 1..=n {
        curr[0] = i;
        for j in 1..=m {
            let cost = usize::from(av[i - 1] != bv[j - 1]);
            curr[j] = (curr[j - 1] + 1) // insertion
                .min(prev[j] + 1) // deletion
                .min(prev[j - 1] + cost); // substitution
        }
        std::mem::swap(&mut prev, &mut curr);
    }
    prev[m]
}

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

    fn paths_for(tmp: &tempfile::TempDir) -> Paths {
        Paths {
            home: tmp.path().to_path_buf(),
            user_home: tmp.path().to_path_buf(),
        }
    }

    #[test]
    fn builtin_is_non_empty_and_contains_filesystem() {
        let r = Registry::builtin();
        assert!(!r.is_empty());
        assert!(r.get("filesystem").is_some());
    }

    #[test]
    fn exact_lookup_is_case_sensitive() {
        // Demo-level: name matching is byte-exact. `closest` lowercases,
        // but `get` does not — callers should canonicalise names before
        // exact lookup.
        let r = Registry::builtin();
        assert!(r.get("filesystem").is_some());
        assert!(r.get("FileSystem").is_none());
    }

    #[test]
    fn closest_finds_typosquats_within_distance() {
        let r = Registry::builtin();
        // "filesystme" → "filesystem" (transposition + missing letter,
        // distance 2).
        let (entry, d) = r.closest("filesystme", 2).expect("typosquat hit");
        assert_eq!(entry.name, "filesystem");
        assert!(d > 0 && d <= 2);
    }

    #[test]
    fn closest_skips_exact_matches() {
        let r = Registry::builtin();
        assert!(r.closest("filesystem", 2).is_none());
    }

    #[test]
    fn closest_returns_none_when_outside_distance() {
        let r = Registry::builtin();
        assert!(r.closest("completely-unrelated-name", 2).is_none());
    }

    #[test]
    fn cache_round_trip() {
        let tmp = tempfile::tempdir().unwrap();
        let paths = paths_for(&tmp);
        let original = Registry::builtin();
        let written = original.save_cached(&paths).unwrap();
        assert!(written.exists());
        let loaded = Registry::load_cached(&paths).unwrap().unwrap();
        assert_eq!(loaded.len(), original.len());
    }

    #[test]
    fn load_or_builtin_uses_cache_when_present() {
        let tmp = tempfile::tempdir().unwrap();
        let paths = paths_for(&tmp);
        // Save a 1-entry custom registry; load_or_builtin must surface it
        // rather than the builtin fallback.
        let custom = Registry::from_entries(vec![RegistryEntry {
            name: "my-internal-server".into(),
            source: "private allowlist".into(),
        }]);
        custom.save_cached(&paths).unwrap();
        let resolved = Registry::load_or_builtin(&paths);
        assert_eq!(resolved.len(), 1);
        assert!(resolved.get("my-internal-server").is_some());
        assert!(resolved.get("filesystem").is_none());
    }

    #[test]
    fn load_or_builtin_falls_back_when_no_cache() {
        let tmp = tempfile::tempdir().unwrap();
        let paths = paths_for(&tmp);
        let resolved = Registry::load_or_builtin(&paths);
        assert!(resolved.get("filesystem").is_some());
    }

    #[test]
    fn levenshtein_known_pairs() {
        assert_eq!(levenshtein("kitten", "sitting"), 3);
        assert_eq!(levenshtein("a", "a"), 0);
        assert_eq!(levenshtein("", "abc"), 3);
        assert_eq!(levenshtein("abc", ""), 3);
        assert_eq!(levenshtein("GitHub", "github"), 0); // case-insensitive
    }

    // ── local allowlist tests ─────────────────────────────────────────────

    #[test]
    fn load_resolved_empty_local_falls_back_to_builtin_only() {
        let tmp = tempfile::tempdir().unwrap();
        let paths = paths_for(&tmp);
        // No cache, no local file → must equal builtin.
        let resolved = Registry::load_resolved(&paths).unwrap();
        let builtin = Registry::builtin();
        assert_eq!(resolved.len(), builtin.len());
        assert!(resolved.get("filesystem").is_some());
    }

    #[test]
    fn load_resolved_local_overrides_builtin() {
        let tmp = tempfile::tempdir().unwrap();
        let paths = paths_for(&tmp);
        // Add a local entry with the same name as a builtin but different source.
        Registry::add_local(&paths, "filesystem", Some("my-fork")).unwrap();
        let resolved = Registry::load_resolved(&paths).unwrap();
        let entry = resolved.get("filesystem").unwrap();
        assert_eq!(entry.source, "my-fork");
    }

    #[test]
    fn load_resolved_local_overrides_fetched() {
        let tmp = tempfile::tempdir().unwrap();
        let paths = paths_for(&tmp);
        // Simulate a fetched cache entry.
        let fetched = Registry::from_entries(vec![RegistryEntry {
            name: "github".into(),
            source: "fetched-registry".into(),
        }]);
        fetched.save_cached(&paths).unwrap();
        // Local entry for the same name wins.
        Registry::add_local(&paths, "github", Some("personal-fork")).unwrap();
        let resolved = Registry::load_resolved(&paths).unwrap();
        let entry = resolved.get("github").unwrap();
        assert_eq!(entry.source, "personal-fork");
    }

    #[test]
    fn add_local_persists_across_reload() {
        let tmp = tempfile::tempdir().unwrap();
        let paths = paths_for(&tmp);
        Registry::add_local(&paths, "my-custom-server", None).unwrap();
        let loaded = Registry::load_local(&paths).unwrap();
        let entry = loaded.get("my-custom-server").unwrap();
        assert_eq!(entry.source, "local");
    }

    #[test]
    fn remove_local_removes_entry() {
        let tmp = tempfile::tempdir().unwrap();
        let paths = paths_for(&tmp);
        Registry::add_local(&paths, "to-remove", Some("test")).unwrap();
        assert!(
            Registry::load_local(&paths)
                .unwrap()
                .get("to-remove")
                .is_some()
        );
        Registry::remove_local(&paths, "to-remove").unwrap();
        assert!(
            Registry::load_local(&paths)
                .unwrap()
                .get("to-remove")
                .is_none()
        );
    }
}