AkurAI-Build / src/oauth.rs

Latest change fcd6ef8973fbf7b205d6da53e253f715b0a44d71 - MCP: restrict IdP tokens to verified admins from AKURAI_BUILD_ADMIN_EMAILS by Ólafur Búi Ólafsson

//! OAuth 2.0 protected-resource support for the HTTP MCP endpoint (`POST /mcp`).
//!
//! The AkurAI IdP issues EdDSA JWT access tokens to the public PKCE client
//! `akurai-agent-akurai-build`; `aud` is an array that also names this app's
//! audience. Tokens are verified locally against the IdP JWKS (cached, and
//! refetched once on an unknown `kid` to follow key rotation). Only user
//! tokens bound to a tenant are accepted. Metadata follows RFC 9728.

use std::{
    io::Write,
    process::{Command, Stdio},
    sync::Mutex,
    time::{Duration, Instant},
};

use jsonwebtoken::{
    Algorithm, DecodingKey, Validation, decode, decode_header,
    jwk::{AlgorithmParameters, EllipticCurve, JwkSet},
};
use serde::Deserialize;
use serde_json::{Value, json};

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

/// Why a presented bearer token was rejected. Never includes token material.
pub type Rejection = &'static str;

/// Outcome of a failed authentication, mapped to 401 or 403 by the caller.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Denial {
    /// The token is not a valid IdP user access token for this resource.
    Invalid(Rejection),
    /// The token is valid but its user is not an allowlisted admin.
    Forbidden(Rejection),
}

/// Claims of a validated access token, before the admin check.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Grant {
    pub subject: String,
    pub tenant: String,
}

/// An allowlisted admin authenticated through the IdP.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Identity {
    pub subject: String,
    pub tenant: String,
    pub email: String,
}

type UserinfoFetch = Box<dyn Fn(&str) -> Result<Value, Rejection> + Send + Sync>;

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

pub struct ResourceServer {
    issuer: String,
    audience: String,
    public_origin: String,
    jwks: Mutex<Option<(Instant, JwkSet)>>,
    admins: Vec<String>,
    userinfo: UserinfoFetch,
}

impl ResourceServer {
    pub fn new(issuer: &str, audience: &str, public_origin: &str) -> Self {
        Self {
            issuer: issuer.trim_end_matches('/').to_owned(),
            audience: audience.to_owned(),
            public_origin: public_origin.trim_end_matches('/').to_owned(),
            jwks: Mutex::new(None),
            admins: Vec::new(),
            userinfo: Box::new(|_| Err("userinfo endpoint is not configured")),
        }
    }

    /// Emails allowed to use the MCP endpoint; empty denies every IdP user.
    pub fn with_admins(mut self, admins: Vec<String>) -> Self {
        self.admins = admins.iter().map(|email| normalize_email(email)).collect();
        self
    }

    /// Fetches IdP userinfo for a token with curl (token passed on stdin).
    pub fn with_userinfo_url(self, url: String) -> Self {
        self.with_userinfo(move |token| fetch_userinfo(&url, token))
    }

    pub fn with_userinfo(
        mut self,
        fetch: impl Fn(&str) -> Result<Value, Rejection> + Send + Sync + 'static,
    ) -> Self {
        self.userinfo = Box::new(fetch);
        self
    }

    /// Preloads the key set; used by tests and to avoid a first-call fetch.
    pub fn with_jwks(self, jwks: JwkSet) -> Self {
        if let Ok(mut cache) = self.jwks.lock() {
            *cache = Some((Instant::now(), jwks));
        }
        self
    }

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

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

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

    /// `WWW-Authenticate` value, with an RFC 6750 error code when a token
    /// was presented and refused.
    pub fn challenge(&self, error: Option<&str>) -> String {
        let url = self.metadata_url();
        match error {
            Some(error) => format!("Bearer error=\"{error}\", resource_metadata=\"{url}\""),
            None => format!("Bearer resource_metadata=\"{url}\""),
        }
    }

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

    /// Blocking: may fetch the JWKS and the IdP userinfo with curl.
    pub fn authenticate(&self, token: &str) -> Result<Identity, Denial> {
        let grant = self.grant(token).map_err(Denial::Invalid)?;
        let userinfo = (self.userinfo)(token).map_err(Denial::Invalid)?;
        let email =
            admin_identity(&userinfo, &grant.subject, &self.admins).map_err(Denial::Forbidden)?;
        Ok(Identity {
            subject: grant.subject,
            tenant: grant.tenant,
            email,
        })
    }

    fn grant(&self, token: &str) -> Result<Grant, Rejection> {
        match validate(token, &self.keys(false)?, &self.issuer, &self.audience) {
            Err(UNKNOWN_KID) => validate(token, &self.keys(true)?, &self.issuer, &self.audience),
            other => other,
        }
    }

    fn keys(&self, refresh: bool) -> Result<JwkSet, Rejection> {
        if !refresh
            && let Ok(cache) = self.jwks.lock()
            && let Some((fetched, jwks)) = cache.as_ref()
            && fetched.elapsed() < JWKS_TTL
        {
            return Ok(jwks.clone());
        }
        let jwks = fetch_jwks(&self.issuer)?;
        if let Ok(mut cache) = self.jwks.lock() {
            *cache = Some((Instant::now(), jwks.clone()));
        }
        Ok(jwks)
    }
}

fn fetch_jwks(issuer: &str) -> Result<JwkSet, Rejection> {
    let output = Command::new("curl")
        .args(["--fail", "--silent", "--show-error", "--max-time", "10"])
        .arg(format!("{issuer}/jwks"))
        .output()
        .map_err(|_| "JWKS request failed")?;
    if !output.status.success() {
        return Err("JWKS request failed");
    }
    serde_json::from_slice(&output.stdout).map_err(|_| "invalid JWKS response")
}

fn fetch_userinfo(url: &str, token: &str) -> Result<Value, Rejection> {
    let mut child = Command::new("curl")
        .args(["--fail", "--silent", "--max-time", "10", "--header", "@-"])
        .arg(url)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .spawn()
        .map_err(|_| "userinfo request failed")?;
    child
        .stdin
        .take()
        .ok_or("userinfo request failed")?
        .write_all(format!("Authorization: Bearer {token}\n").as_bytes())
        .map_err(|_| "userinfo request failed")?;
    let output = child
        .wait_with_output()
        .map_err(|_| "userinfo request failed")?;
    if !output.status.success() {
        return Err("inactive IdP identity");
    }
    serde_json::from_slice(&output.stdout).map_err(|_| "invalid userinfo response")
}

fn normalize_email(email: &str) -> String {
    email.trim().to_lowercase()
}

/// The verified, allowlisted email userinfo reports for `subject`.
pub fn admin_identity(
    userinfo: &Value,
    subject: &str,
    admins: &[String],
) -> Result<String, Rejection> {
    if userinfo.get("sub").and_then(Value::as_str) != Some(subject) {
        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(normalize_email)
        .filter(|email| !email.is_empty())
        .ok_or("userinfo has no email")?;
    if !admins.contains(&email) {
        return Err("user is not an AkurAI Build admin");
    }
    Ok(email)
}

const UNKNOWN_KID: Rejection = "unknown JWT key id";

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

pub fn validate(
    token: &str,
    jwks: &JwkSet,
    issuer: &str,
    audience: &str,
) -> Result<Grant, Rejection> {
    let header = decode_header(token).map_err(|_| "invalid JWT header")?;
    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_KID)?;
    if !matches!(
        &jwk.algorithm,
        AlgorithmParameters::OctetKeyPair(key) if key.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]);
    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("token audience does not include this resource");
    }
    let client_token =
        claims.client_id.is_some() || claims.token_type.as_deref() == Some("client_credentials");
    if claims.sub.is_empty() || client_token {
        return Err("a user access token is required");
    }
    let tenant = claims
        .tenant_id
        .filter(|tenant| !tenant.is_empty())
        .ok_or("token has no tenant")?;
    Ok(Grant {
        subject: claims.sub,
        tenant,
    })
}

#[cfg(test)]
pub(crate) mod testing {
    use jsonwebtoken::{EncodingKey, Header, encode, jwk::JwkSet};
    use ring::{
        rand::SystemRandom,
        signature::{Ed25519KeyPair, KeyPair},
    };
    use serde_json::{Value, json};

    pub const ISSUER: &str = "https://auth.example";
    pub const KID: &str = "test-key";

    /// Ephemeral Ed25519 signer standing in for the IdP.
    pub struct TestIdp {
        pkcs8: Vec<u8>,
        public: Vec<u8>,
    }

    impl TestIdp {
        pub fn new() -> Self {
            let pkcs8 = Ed25519KeyPair::generate_pkcs8(&SystemRandom::new())
                .expect("generate key")
                .as_ref()
                .to_vec();
            let public = Ed25519KeyPair::from_pkcs8(&pkcs8)
                .expect("parse key")
                .public_key()
                .as_ref()
                .to_vec();
            Self { pkcs8, public }
        }

        pub fn jwks(&self) -> JwkSet {
            let x = base64url(&self.public);
            serde_json::from_value(json!({"keys": [{
                "kty": "OKP", "crv": "Ed25519", "alg": "EdDSA", "use": "sig", "kid": KID, "x": x
            }]}))
            .expect("jwks")
        }

        pub fn sign(&self, kid: &str, claims: &Value) -> String {
            let mut header = Header::new(jsonwebtoken::Algorithm::EdDSA);
            header.kid = Some(kid.to_owned());
            encode(&header, claims, &EncodingKey::from_ed_der(&self.pkcs8)).expect("sign")
        }
    }

    pub fn user_claims(aud: Value) -> Value {
        json!({
            "iss": ISSUER,
            "sub": "user-1",
            "aud": aud,
            "exp": 4_102_444_800_u64,
            "iat": 1_700_000_000_u64,
            "tenant_id": "tenant-1",
            "organization_id": "tenant-1",
            "scope": "openid profile email groups offline_access",
        })
    }

    fn base64url(bytes: &[u8]) -> String {
        const TABLE: &[u8; 64] =
            b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
        let mut out = String::new();
        for chunk in bytes.chunks(3) {
            let value = chunk.iter().enumerate().fold(0_u32, |acc, (i, byte)| {
                acc | (u32::from(*byte) << (16 - 8 * i))
            });
            for index in 0..=chunk.len() {
                out.push(TABLE[((value >> (18 - 6 * index)) & 63) as usize] as char);
            }
        }
        out
    }
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::testing::{ISSUER, KID, TestIdp, user_claims};
    use super::*;

    const AUD: [&str; 2] = ["akurai-agent-akurai-build", "akurai-build"];

    fn server(idp: &TestIdp) -> ResourceServer {
        ResourceServer::new(ISSUER, "akurai-build", "https://build.example/")
            .with_jwks(idp.jwks())
            .with_admins(vec!["admin@example.com".to_owned()])
            .with_userinfo(|_| Ok(userinfo("user-1", "admin@example.com", true)))
    }

    #[test]
    fn metadata_names_resource_and_issuer() {
        let server = ResourceServer::new(
            "https://auth.example/",
            "akurai-build",
            "https://build.example",
        );
        let metadata = server.metadata();
        assert_eq!(metadata["resource"], "https://build.example/mcp");
        assert_eq!(
            metadata["authorization_servers"],
            json!(["https://auth.example"])
        );
        assert_eq!(metadata["bearer_methods_supported"], json!(["header"]));
    }

    #[test]
    fn challenge_points_at_metadata() {
        let server = ResourceServer::new(ISSUER, "akurai-build", "https://build.example");
        let url = "https://build.example/.well-known/oauth-protected-resource/mcp";
        assert_eq!(
            server.challenge(None),
            format!("Bearer resource_metadata=\"{url}\"")
        );
        assert_eq!(
            server.challenge(Some("invalid_token")),
            format!("Bearer error=\"invalid_token\", resource_metadata=\"{url}\"")
        );
    }

    #[test]
    fn origin_must_match_when_present() {
        let server = ResourceServer::new(ISSUER, "akurai-build", "https://build.example");
        assert!(server.origin_allowed(None));
        assert!(server.origin_allowed(Some("https://build.example")));
        assert!(!server.origin_allowed(Some("https://evil.example")));
    }

    #[test]
    fn audience_matches_string_or_array() {
        assert!(audience_contains(&json!("akurai-build"), "akurai-build"));
        assert!(audience_contains(&json!(AUD), "akurai-build"));
        assert!(!audience_contains(&json!(["akurai-notes"]), "akurai-build"));
        assert!(!audience_contains(&json!(null), "akurai-build"));
    }

    #[test]
    fn accepts_idp_user_token_with_array_audience() {
        let idp = TestIdp::new();
        let token = idp.sign(KID, &user_claims(json!(AUD)));
        assert_eq!(
            server(&idp).authenticate(&token),
            Ok(Identity {
                subject: "user-1".into(),
                tenant: "tenant-1".into(),
                email: "admin@example.com".into(),
            })
        );
    }

    #[test]
    fn rejects_other_audience() {
        let idp = TestIdp::new();
        let token = idp.sign(
            KID,
            &user_claims(json!(["akurai-agent-akurai-notes", "akurai-notes"])),
        );
        assert!(server(&idp).authenticate(&token).is_err());
    }

    #[test]
    fn rejects_wrong_issuer_expiry_and_signature() {
        let idp = TestIdp::new();
        let mut claims = user_claims(json!(AUD));
        claims["iss"] = json!("https://evil.example");
        assert!(server(&idp).authenticate(&idp.sign(KID, &claims)).is_err());

        let mut claims = user_claims(json!(AUD));
        claims["exp"] = json!(1_000_u64);
        assert!(server(&idp).authenticate(&idp.sign(KID, &claims)).is_err());

        let forged = TestIdp::new().sign(KID, &user_claims(json!(AUD)));
        assert!(server(&idp).authenticate(&forged).is_err());
        assert!(server(&idp).authenticate("not-a-jwt").is_err());
    }

    #[test]
    fn rejects_client_credentials_and_tenantless_tokens() {
        let idp = TestIdp::new();
        let mut claims = user_claims(json!(AUD));
        claims["client_id"] = json!("akurai-agent-akurai-build");
        assert!(server(&idp).authenticate(&idp.sign(KID, &claims)).is_err());

        let mut claims = user_claims(json!(AUD));
        claims["token_type"] = json!("client_credentials");
        assert!(server(&idp).authenticate(&idp.sign(KID, &claims)).is_err());

        let mut claims = user_claims(json!(AUD));
        claims.as_object_mut().expect("object").remove("tenant_id");
        assert!(server(&idp).authenticate(&idp.sign(KID, &claims)).is_err());
    }

    #[test]
    fn subject_prefix_is_not_treated_as_a_client() {
        let idp = TestIdp::new();
        let mut claims = user_claims(json!(AUD));
        claims["sub"] = json!("akurai-agent-someone");
        let token = idp.sign(KID, &claims);
        assert!(validate(&token, &idp.jwks(), ISSUER, "akurai-build").is_ok());
    }

    fn userinfo(sub: &str, email: &str, verified: bool) -> Value {
        json!({"sub": sub, "email": email, "email_verified": verified})
    }

    #[test]
    fn admin_identity_requires_matching_verified_allowlisted_email() {
        let admins = vec!["admin@example.com".to_owned()];
        assert_eq!(
            admin_identity(
                &userinfo("user-1", "Admin@Example.com", true),
                "user-1",
                &admins
            ),
            Ok("admin@example.com".to_owned())
        );
        for (info, reason) in [
            (
                userinfo("user-2", "admin@example.com", true),
                "subject mismatch",
            ),
            (userinfo("user-1", "admin@example.com", false), "unverified"),
            (
                userinfo("user-1", "member@example.com", true),
                "not allowlisted",
            ),
            (json!({"sub": "user-1", "email_verified": true}), "no email"),
        ] {
            assert!(
                admin_identity(&info, "user-1", &admins).is_err(),
                "{reason}"
            );
        }
    }

    #[test]
    fn admin_identity_fails_closed_on_empty_allowlist() {
        let info = userinfo("user-1", "admin@example.com", true);
        assert!(admin_identity(&info, "user-1", &[]).is_err());
    }

    #[test]
    fn authenticate_forbids_non_admins_after_valid_token() {
        let idp = TestIdp::new();
        let token = idp.sign(KID, &user_claims(json!(AUD)));
        let admins = || vec!["admin@example.com".to_owned()];
        let member = server(&idp)
            .with_admins(admins())
            .with_userinfo(|_| Ok(userinfo("user-1", "member@example.com", true)));
        assert!(matches!(
            member.authenticate(&token),
            Err(Denial::Forbidden(_))
        ));
        let unreachable = server(&idp)
            .with_admins(admins())
            .with_userinfo(|_| Err("userinfo request failed"));
        assert!(matches!(
            unreachable.authenticate(&token),
            Err(Denial::Invalid(_))
        ));
    }

    #[test]
    fn validate_reports_unknown_kid_for_rotation() {
        let idp = TestIdp::new();
        let token = idp.sign("rotated", &user_claims(json!(AUD)));
        assert_eq!(
            validate(&token, &idp.jwks(), ISSUER, "akurai-build"),
            Err(UNKNOWN_KID)
        );
    }
}