AkurAI-CRM / crates/crm-mcp/src/oauth.rs

Latest change 46f8d3437352b9043c8406756a07d2c52c892d19 - Gate MCP merge tool on IdP role; send app.js 401s through login by Ólafur Búi Ólafsson

//! The MCP endpoint as an OAuth 2.0 protected resource of the AkurAI IdP.
//!
//! - RFC 9728 metadata at `/.well-known/oauth-protected-resource[/mcp]`.
//! - Bearer access tokens are EdDSA (Ed25519) JWTs verified locally against
//!   the IdP JWKS. `aud` (string or array) must contain the app audience
//!   (`akurai-crm`); agent tokens carry `["akurai-agent-crm","akurai-crm"]`.
//! - Only user tokens with a tenant are accepted; `client_credentials`
//!   tokens have no user and are refused.

use akurai_json::Value;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine as _;
use crm_api::session::Role;
use ring::signature::{UnparsedPublicKey, ED25519};
use std::sync::Mutex;
use std::time::{Duration, Instant};

pub const DEFAULT_PUBLIC_URL: &str = "https://akurai-crm.olibuijr.com";
pub const DEFAULT_ISSUER: &str = "https://auth.olibuijr.com";
pub const AUDIENCE: &str = "akurai-crm";
const RESOURCE_NAME: &str = "AkurAI-CRM MCP";
const SCOPES: [&str; 5] = ["openid", "profile", "email", "groups", "offline_access"];
const JWKS_TTL: Duration = Duration::from_secs(600);
const JWKS_MIN_REFRESH: Duration = Duration::from_secs(30);

#[derive(Debug, Clone)]
pub struct ResourceConfig {
    pub public_url: String,
    pub issuer: String,
    pub audience: String,
}

impl ResourceConfig {
    pub fn from_env() -> Self {
        let env = |name: &str, default: &str| {
            std::env::var(name)
                .ok()
                .filter(|v| !v.trim().is_empty())
                .unwrap_or_else(|| default.to_string())
        };
        Self {
            public_url: env("AKURAI_PUBLIC_URL", DEFAULT_PUBLIC_URL),
            issuer: env("AKURAI_IDP_ISSUER", DEFAULT_ISSUER),
            audience: AUDIENCE.to_string(),
        }
    }

    fn base(&self) -> &str {
        self.public_url.trim_end_matches('/')
    }

    fn issuer(&self) -> &str {
        self.issuer.trim_end_matches('/')
    }

    pub fn resource_url(&self) -> String {
        format!("{}/mcp", self.base())
    }

    pub fn metadata_url(&self) -> String {
        format!("{}/.well-known/oauth-protected-resource/mcp", self.base())
    }

    pub fn metadata_json(&self) -> String {
        let strings = |items: &[&str]| {
            Value::Array(items.iter().map(|s| Value::Str(s.to_string())).collect())
        };
        Value::Object(vec![
            ("resource".into(), Value::Str(self.resource_url())),
            ("authorization_servers".into(), strings(&[self.issuer()])),
            ("scopes_supported".into(), strings(&SCOPES)),
            ("bearer_methods_supported".into(), strings(&["header"])),
            ("resource_name".into(), Value::Str(RESOURCE_NAME.into())),
        ])
        .to_json()
    }

    /// `WWW-Authenticate` value; `invalid` marks a presented-but-rejected token.
    pub fn challenge(&self, invalid: bool) -> String {
        let url = self.metadata_url();
        if invalid {
            format!("Bearer error=\"invalid_token\", resource_metadata=\"{url}\"")
        } else {
            format!("Bearer resource_metadata=\"{url}\"")
        }
    }

    /// Browsers must be same-origin; non-browser clients send no Origin.
    pub fn origin_allowed(&self, origin: Option<&str>) -> bool {
        origin.is_none_or(|o| o.trim_end_matches('/') == self.base())
    }
}

/// Ed25519 verification keys from a JWKS document, by `kid`.
#[derive(Debug, Clone, Default)]
pub struct Jwks {
    keys: Vec<(String, Vec<u8>)>,
}

impl Jwks {
    pub fn parse(json: &str) -> Result<Self, &'static str> {
        let doc = akurai_json::parse(json).map_err(|_| "invalid JWKS")?;
        let Some(Value::Array(entries)) = doc.get("keys") else {
            return Err("JWKS has no keys");
        };
        let keys = entries
            .iter()
            .filter(|k| {
                str_field(k, "kty") == Some("OKP") && str_field(k, "crv") == Some("Ed25519")
            })
            .filter_map(|k| {
                let kid = str_field(k, "kid")?.to_string();
                let x = URL_SAFE_NO_PAD.decode(str_field(k, "x")?).ok()?;
                Some((kid, x))
            })
            .collect();
        Ok(Self { keys })
    }

    fn find(&self, kid: &str) -> Option<&[u8]> {
        self.keys
            .iter()
            .find(|(k, _)| k == kid)
            .map(|(_, x)| x.as_slice())
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct Grant {
    pub sub: String,
    pub tenant: String,
    pub scopes: Vec<String>,
    /// Organization role from the `role` claim; missing/unknown is member.
    pub role: Role,
}

fn str_field<'a>(value: &'a Value, key: &str) -> Option<&'a str> {
    value.get(key).and_then(Value::as_str)
}

fn decode_segment(segment: &str) -> Result<Value, &'static str> {
    let bytes = URL_SAFE_NO_PAD
        .decode(segment)
        .map_err(|_| "JWT segment is not base64url")?;
    let text = std::str::from_utf8(&bytes).map_err(|_| "JWT segment is not UTF-8")?;
    akurai_json::parse(text).map_err(|_| "JWT segment is not JSON")
}

fn audience_contains(aud: Option<&Value>, audience: &str) -> bool {
    match aud {
        Some(Value::Str(s)) => s == audience,
        Some(Value::Array(items)) => items.iter().any(|v| v.as_str() == Some(audience)),
        _ => false,
    }
}

pub fn validate_access_token(
    token: &str,
    jwks: &Jwks,
    cfg: &ResourceConfig,
    now: i64,
) -> Result<Grant, &'static str> {
    let mut parts = token.split('.');
    let (Some(header), Some(payload), Some(signature), None) =
        (parts.next(), parts.next(), parts.next(), parts.next())
    else {
        return Err("access token is not a JWT");
    };
    let header = decode_segment(header)?;
    if str_field(&header, "alg") != Some("EdDSA") {
        return Err("unexpected JWT algorithm");
    }
    let kid = str_field(&header, "kid").ok_or("missing JWT key id")?;
    let key = jwks.find(kid).ok_or("unknown JWT key id")?;
    let signature = URL_SAFE_NO_PAD
        .decode(signature)
        .map_err(|_| "JWT signature is not base64url")?;
    let (signing_input, _) = token.rsplit_once('.').ok_or("access token is not a JWT")?;
    UnparsedPublicKey::new(&ED25519, key)
        .verify(signing_input.as_bytes(), &signature)
        .map_err(|_| "JWT signature validation failed")?;

    let claims = decode_segment(payload)?;
    if str_field(&claims, "iss").map(|i| i.trim_end_matches('/')) != Some(cfg.issuer()) {
        return Err("access token issuer mismatch");
    }
    match claims.get("exp").and_then(Value::as_i64) {
        Some(exp) if exp > now => {}
        _ => return Err("access token has expired"),
    }
    if claims
        .get("nbf")
        .and_then(Value::as_i64)
        .is_some_and(|nbf| nbf > now)
    {
        return Err("access token is not yet valid");
    }
    if !audience_contains(claims.get("aud"), &cfg.audience) {
        return Err("access token audience mismatch");
    }
    let sub = str_field(&claims, "sub").unwrap_or("");
    // client_credentials tokens name the client as subject; they have no user.
    if sub.is_empty()
        || claims.get("client_id").is_some()
        || str_field(&claims, "token_type") == Some("client_credentials")
        || audience_contains(claims.get("aud"), sub)
    {
        return Err("invalid user access token");
    }
    let tenant = str_field(&claims, "tenant_id")
        .or_else(|| str_field(&claims, "organization_id"))
        .filter(|t| !t.is_empty())
        .ok_or("access token has no tenant")?;
    let scopes = str_field(&claims, "scope")
        .unwrap_or("")
        .split_whitespace()
        .map(str::to_string)
        .collect();
    Ok(Grant {
        sub: sub.to_string(),
        tenant: tenant.to_string(),
        scopes,
        role: Role::parse(str_field(&claims, "role").unwrap_or("")),
    })
}

pub type JwksSource = Box<dyn Fn() -> Result<Jwks, &'static str> + Send + Sync>;

/// Looks up a user's IdP groups for an access token. IdP access tokens carry
/// no `groups` claim, so production asks the IdP's `/userinfo`.
pub type GroupSource = Box<dyn Fn(&str) -> Result<Vec<String>, &'static str> + Send + Sync>;

pub const GROUP_DENIED: &str = crm_api::auth::idp_client::GROUP_DENIED;

struct GroupPolicy {
    allowed: Vec<String>,
    source: GroupSource,
}

/// Validates access tokens with a cached JWKS, refetching on expiry or on an
/// unknown key id (key rotation), rate-limited.
pub struct TokenVerifier {
    cfg: ResourceConfig,
    source: JwksSource,
    cache: Mutex<Option<(Instant, Jwks)>>,
    groups: Option<GroupPolicy>,
}

impl TokenVerifier {
    pub fn new(cfg: ResourceConfig, source: JwksSource) -> Self {
        Self {
            cfg,
            source,
            cache: Mutex::new(None),
            groups: None,
        }
    }

    /// Require membership in one of `allowed` (no-op when empty).
    pub fn with_allowed_groups(mut self, allowed: Vec<String>, source: GroupSource) -> Self {
        self.groups = (!allowed.is_empty()).then_some(GroupPolicy { allowed, source });
        self
    }

    /// Verifier that fetches `<issuer>/jwks` over HTTPS and, when
    /// `AKURAI_IDP_ALLOWED_GROUPS` is set, checks groups via the IdP's
    /// internal `/userinfo`.
    pub fn from_env() -> Result<Self, String> {
        use crm_api::auth::idp_client;
        let cfg = ResourceConfig::from_env();
        let url = format!("{}/jwks", cfg.issuer());
        let verifier = Self::new(cfg, Box::new(move || fetch_jwks(&url)));
        let allowed = idp_client::allowed_groups_from_env();
        if allowed.is_empty() {
            return Ok(verifier);
        }
        let internal = idp_client::internal_url_from_env()?;
        Ok(verifier.with_allowed_groups(
            allowed,
            Box::new(move |token| {
                idp_client::fetch_userinfo(&internal, token)
                    .map(|info| info.groups)
                    .map_err(|_| "userinfo lookup failed")
            }),
        ))
    }

    pub fn config(&self) -> &ResourceConfig {
        &self.cfg
    }

    pub fn verify(&self, token: &str) -> Result<Grant, &'static str> {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs() as i64;
        let jwks = self.jwks(false)?;
        let grant = match validate_access_token(token, &jwks, &self.cfg, now) {
            Err("unknown JWT key id") => {
                validate_access_token(token, &self.jwks(true)?, &self.cfg, now)
            }
            result => result,
        }?;
        if let Some(policy) = &self.groups {
            let groups = (policy.source)(token)?;
            if !crm_api::auth::idp_client::groups_allowed(&policy.allowed, &groups) {
                return Err(GROUP_DENIED);
            }
        }
        Ok(grant)
    }

    fn jwks(&self, refresh: bool) -> Result<Jwks, &'static str> {
        let mut cache = self.cache.lock().map_err(|_| "JWKS cache poisoned")?;
        if let Some((at, jwks)) = cache.as_ref() {
            let fresh = at.elapsed() < JWKS_TTL;
            if (fresh && !refresh) || at.elapsed() < JWKS_MIN_REFRESH {
                return Ok(jwks.clone());
            }
        }
        let jwks = (self.source)()?;
        *cache = Some((Instant::now(), jwks.clone()));
        Ok(jwks)
    }
}

fn fetch_jwks(url: &str) -> Result<Jwks, &'static str> {
    let body = reqwest::blocking::Client::builder()
        .timeout(Duration::from_secs(10))
        .build()
        .map_err(|_| "JWKS client configuration failed")?
        .get(url)
        .send()
        .and_then(|r| r.error_for_status())
        .and_then(|r| r.text())
        .map_err(|_| "JWKS request failed")?;
    Jwks::parse(&body)
}