AkurAI-Notes / crates/notes-mcp/src/oauth.rs

Latest change 228bc96dddc8e5490cfc2ed27705c0d1dbbfaf83 - MCP: OAuth protected resource for the AkurAI IdP by Ólafur Búi Ólafsson

//! OAuth 2.0 protected-resource support for the MCP endpoint.
//!
//! The AkurAI IdP issues EdDSA access tokens whose `aud` array contains the
//! app audience (`akurai-notes`). Tokens are verified locally against the
//! IdP JWKS, then the caller's email is resolved through `/userinfo` and
//! checked against the Notes allow list.

use jsonwebtoken::jwk::{AlgorithmParameters, EllipticCurve, JwkSet, KeyAlgorithm};
use jsonwebtoken::{decode, decode_header, Algorithm, DecodingKey, Validation};
use serde::Deserialize;
use serde_json::{json, Value};
use std::collections::HashSet;
use std::sync::Mutex;
use std::time::{Duration, Instant};

const RESOURCE_NAME: &str = "AkurAI Notes MCP";
const SCOPES: [&str; 5] = ["openid", "profile", "email", "groups", "offline_access"];
const JWKS_TTL: Duration = Duration::from_secs(600);
pub const METADATA_PATH: &str = "/.well-known/oauth-protected-resource";

pub struct OAuthConfig {
    /// Public origin of the app, e.g. `https://akurai-notes.olibuijr.com`.
    pub public_url: String,
    pub issuer: String,
    /// Audience the IdP puts in tokens for this app.
    pub audience: String,
    /// Emails allowed to use Notes.
    pub allow: Vec<String>,
}

impl OAuthConfig {
    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 resource_metadata_url(&self) -> String {
        format!("{}{METADATA_PATH}/mcp", self.base())
    }

    /// 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 == self.base())
    }
}

pub fn is_metadata_path(path: &str) -> bool {
    path == METADATA_PATH || path == format!("{METADATA_PATH}/mcp")
}

pub fn protected_resource_metadata(cfg: &OAuthConfig) -> Value {
    json!({
        "resource": cfg.resource_url(),
        "authorization_servers": [cfg.issuer()],
        "scopes_supported": SCOPES,
        "bearer_methods_supported": ["header"],
        "resource_name": RESOURCE_NAME,
    })
}

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

// ---------------------------------------------------------------------------
// Access-token validation
// ---------------------------------------------------------------------------

#[derive(Deserialize)]
struct AccessClaims {
    sub: String,
    #[serde(default)]
    aud: Value,
    #[serde(default)]
    scope: Option<String>,
    #[serde(default)]
    client_id: Option<String>,
}

#[derive(Debug, Clone, PartialEq)]
pub struct TokenGrant {
    pub sub: String,
    pub scopes: HashSet<String>,
}

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

pub fn validate_access_token(
    token: &str,
    jwks: &JwkSet,
    issuer: &str,
    audience: &str,
) -> Result<TokenGrant, &'static str> {
    let header = decode_header(token).map_err(|_| "access token is not a JWT")?;
    if header.alg != Algorithm::EdDSA {
        return Err("unexpected JWT algorithm");
    }
    let kid = header.kid.as_deref().ok_or("missing JWT key id")?;
    let jwk = jwks.find(kid).ok_or("unknown JWT key id")?;
    if jwk.common.key_algorithm != Some(KeyAlgorithm::EdDSA)
        || !matches!(
            &jwk.algorithm,
            AlgorithmParameters::OctetKeyPair(p) if p.curve == EllipticCurve::Ed25519
        )
    {
        return Err("unsupported JWKS key");
    }
    let key = DecodingKey::from_jwk(jwk).map_err(|_| "unsupported JWKS key")?;
    let mut validation = Validation::new(Algorithm::EdDSA);
    validation.set_issuer(&[issuer.trim_end_matches('/')]);
    validation.validate_aud = false;
    validation.leeway = 0;
    let claims = decode::<AccessClaims>(token, &key, &validation)
        .map_err(|_| "JWT signature or claims validation failed")?
        .claims;
    if !audience_contains(&claims.aud, audience) {
        return Err("access token audience mismatch");
    }
    // client_credentials tokens name the client, not a user.
    if claims.client_id.is_some() || claims.sub.is_empty() {
        return Err("invalid user access token");
    }
    let scopes = claims
        .scope
        .unwrap_or_default()
        .split_whitespace()
        .map(str::to_string)
        .collect();
    Ok(TokenGrant {
        sub: claims.sub,
        scopes,
    })
}

/// Allowed caller email from the IdP userinfo response for the same token.
pub fn identity_from_userinfo(
    userinfo: &Value,
    sub: &str,
    allow: &[String],
) -> Result<String, &'static str> {
    if userinfo.get("sub").and_then(Value::as_str) != Some(sub) {
        return Err("userinfo identity mismatch");
    }
    if userinfo.get("email_verified").and_then(Value::as_bool) != Some(true) {
        return Err("verified email is required");
    }
    let email = userinfo
        .get("email")
        .and_then(Value::as_str)
        .map(|e| e.trim().to_lowercase())
        .ok_or("userinfo omitted email")?;
    if !allow.iter().any(|a| a.trim().eq_ignore_ascii_case(&email)) {
        return Err("IDP user is not allowed to access AkurAI-Notes");
    }
    Ok(email)
}

// ---------------------------------------------------------------------------
// IdP access
// ---------------------------------------------------------------------------

pub trait IdpClient: Send + Sync {
    /// Signing keys; `force` bypasses any cache (key rotation).
    fn jwks(&self, force: bool) -> Result<JwkSet, &'static str>;
    fn userinfo(&self, token: &str) -> Result<Value, &'static str>;
}

/// Talks to the IdP over its loopback-internal URL.
pub struct HttpIdp {
    internal_url: String,
    http: reqwest::blocking::Client,
    cache: Mutex<Option<(Instant, JwkSet)>>,
}

impl HttpIdp {
    pub fn new(internal_url: &str) -> Result<Self, String> {
        let http = reqwest::blocking::Client::builder()
            .timeout(Duration::from_secs(10))
            .build()
            .map_err(|e| format!("IDP http client: {e}"))?;
        Ok(Self {
            internal_url: internal_url.trim_end_matches('/').to_string(),
            http,
            cache: Mutex::new(None),
        })
    }
}

impl IdpClient for HttpIdp {
    fn jwks(&self, force: bool) -> Result<JwkSet, &'static str> {
        let mut cache = self.cache.lock().map_err(|_| "JWKS cache poisoned")?;
        if let Some((at, jwks)) = cache.as_ref() {
            if !force && at.elapsed() < JWKS_TTL {
                return Ok(jwks.clone());
            }
        }
        let jwks = self
            .http
            .get(format!("{}/jwks", self.internal_url))
            .send()
            .and_then(|r| r.error_for_status())
            .map_err(|_| "JWKS request failed")?
            .json::<JwkSet>()
            .map_err(|_| "invalid JWKS response")?;
        *cache = Some((Instant::now(), jwks.clone()));
        Ok(jwks)
    }

    fn userinfo(&self, token: &str) -> Result<Value, &'static str> {
        let resp = self
            .http
            .get(format!("{}/userinfo", self.internal_url))
            .bearer_auth(token)
            .send()
            .map_err(|_| "userinfo request failed")?;
        if !resp.status().is_success() {
            return Err("inactive IDP identity");
        }
        resp.json::<Value>()
            .map_err(|_| "invalid userinfo response")
    }
}

pub struct OAuth {
    pub config: OAuthConfig,
    idp: Box<dyn IdpClient>,
}

impl OAuth {
    pub fn new(config: OAuthConfig, idp: Box<dyn IdpClient>) -> Self {
        Self { config, idp }
    }

    /// Validate an IdP access token and return the allowed caller's email.
    pub fn authenticate(&self, token: &str) -> Result<String, &'static str> {
        let cfg = &self.config;
        let validate =
            |jwks: &JwkSet| validate_access_token(token, jwks, &cfg.issuer, &cfg.audience);
        let grant = match validate(&self.idp.jwks(false)?) {
            Err("unknown JWT key id") => validate(&self.idp.jwks(true)?)?,
            other => other?,
        };
        identity_from_userinfo(&self.idp.userinfo(token)?, &grant.sub, &cfg.allow)
    }
}