AkurAI-Drive / backend/src/oauth.rs

Latest change f891f45338fe5e5d14370df77cc8001dceb06426 - Drive MCP: recover /mcp tools and make it an IdP OAuth protected resource by AkurAI Agent

//! IdP access-token validation for the MCP endpoint.
//!
//! The AkurAI IdP issues EdDSA (Ed25519) JWT access tokens to the public PKCE
//! client `akurai-agent-akurai-drive`, with `aud` = [client id, "akurai-drive"].
//! They are verified locally against the IdP JWKS (fetched over the same-box
//! back channel and cached); no token is ever stored.

use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use ed25519_dalek::{Signature, Verifier, VerifyingKey};
use serde_json::Value;
use std::sync::{LazyLock, Mutex};
use std::time::{Duration, Instant};

use crate::config;

const JWKS_TTL: Duration = Duration::from_secs(600);

/// The user a valid IdP access token was issued to.
#[derive(Debug, Clone, PartialEq)]
pub struct TokenGrant {
    pub sub: String,
    pub tenant_id: Option<String>,
    pub scopes: Vec<String>,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum TokenError {
    Malformed,
    UnknownKey,
    BadSignature,
    WrongIssuer,
    Expired,
    WrongAudience,
    NotAUserToken,
}

/// Does the `aud` claim (string or array) contain `audience`?
pub fn audience_contains(aud: Option<&Value>, audience: &str) -> bool {
    match aud {
        Some(Value::String(s)) => s == audience,
        Some(Value::Array(items)) => items.iter().any(|v| v.as_str() == Some(audience)),
        _ => false,
    }
}

fn decode_json(segment: &str) -> Result<Value, TokenError> {
    let bytes = URL_SAFE_NO_PAD
        .decode(segment)
        .map_err(|_| TokenError::Malformed)?;
    serde_json::from_slice(&bytes).map_err(|_| TokenError::Malformed)
}

/// The Ed25519 public key for `kid` in a JWKS document, if present and usable.
fn find_key(jwks: &Value, kid: &str) -> Option<VerifyingKey> {
    let jwk = jwks
        .get("keys")?
        .as_array()?
        .iter()
        .find(|k| k.get("kid").and_then(Value::as_str) == Some(kid))?;
    if jwk.get("kty").and_then(Value::as_str) != Some("OKP")
        || jwk.get("crv").and_then(Value::as_str) != Some("Ed25519")
    {
        return None;
    }
    let x = URL_SAFE_NO_PAD.decode(jwk.get("x")?.as_str()?).ok()?;
    VerifyingKey::from_bytes(&x.try_into().ok()?).ok()
}

/// Verify a compact JWS access token and its claims. `now` is Unix seconds.
pub fn validate_access_token(
    token: &str,
    jwks: &Value,
    issuer: &str,
    accepted_audiences: &[&str],
    now: i64,
) -> Result<TokenGrant, TokenError> {
    let mut parts = token.split('.');
    let (Some(h), Some(p), Some(s), None) =
        (parts.next(), parts.next(), parts.next(), parts.next())
    else {
        return Err(TokenError::Malformed);
    };
    let header = decode_json(h)?;
    if header.get("alg").and_then(Value::as_str) != Some("EdDSA") {
        return Err(TokenError::Malformed);
    }
    let kid = header
        .get("kid")
        .and_then(Value::as_str)
        .ok_or(TokenError::Malformed)?;
    let key = find_key(jwks, kid).ok_or(TokenError::UnknownKey)?;
    let sig_bytes = URL_SAFE_NO_PAD
        .decode(s)
        .map_err(|_| TokenError::Malformed)?;
    let signature = Signature::from_slice(&sig_bytes).map_err(|_| TokenError::Malformed)?;
    key.verify(format!("{h}.{p}").as_bytes(), &signature)
        .map_err(|_| TokenError::BadSignature)?;

    let claims = decode_json(p)?;
    if claims.get("iss").and_then(Value::as_str) != Some(issuer.trim_end_matches('/')) {
        return Err(TokenError::WrongIssuer);
    }
    let exp = claims.get("exp").and_then(Value::as_i64).unwrap_or(0);
    let nbf = claims
        .get("nbf")
        .and_then(Value::as_i64)
        .unwrap_or(i64::MIN);
    if exp <= now || nbf > now {
        return Err(TokenError::Expired);
    }
    if !accepted_audiences
        .iter()
        .any(|a| audience_contains(claims.get("aud"), a))
    {
        return Err(TokenError::WrongAudience);
    }
    // client_credentials tokens carry `client_id` and have no Drive owner.
    let sub = claims.get("sub").and_then(Value::as_str).unwrap_or("");
    if sub.is_empty() || claims.get("client_id").is_some() || accepted_audiences.contains(&sub) {
        return Err(TokenError::NotAUserToken);
    }
    let tenant_id = ["tenant_id", "organization_id"]
        .iter()
        .find_map(|c| claims.get(*c).and_then(Value::as_str))
        .filter(|t| !t.is_empty())
        .map(str::to_string);
    let scopes = claims
        .get("scope")
        .and_then(Value::as_str)
        .unwrap_or("")
        .split_whitespace()
        .map(str::to_string)
        .collect();
    Ok(TokenGrant {
        sub: sub.to_string(),
        tenant_id,
        scopes,
    })
}

static JWKS_CACHE: LazyLock<Mutex<Option<(Instant, Value)>>> = LazyLock::new(|| Mutex::new(None));

/// IdP JWKS over the plain-HTTP same-box back channel, cached for `JWKS_TTL`.
fn fetch_jwks(force: bool) -> Result<Value, String> {
    if !force
        && let Some((at, jwks)) = JWKS_CACHE.lock().unwrap().as_ref()
        && at.elapsed() < JWKS_TTL
    {
        return Ok(jwks.clone());
    }
    let url = format!("{}/jwks", config::get().idp_internal_url);
    let body = ureq::get(&url)
        .timeout(Duration::from_secs(10))
        .call()
        .map_err(|e| format!("JWKS request failed: {e}"))?
        .into_string()
        .map_err(|e| format!("JWKS read failed: {e}"))?;
    let jwks: Value =
        serde_json::from_str(&body).map_err(|e| format!("invalid JWKS response: {e}"))?;
    *JWKS_CACHE.lock().unwrap() = Some((Instant::now(), jwks.clone()));
    Ok(jwks)
}

/// Validate an IdP access token for this app (blocking: may fetch the JWKS).
/// Refetches the JWKS once on an unknown `kid` to follow key rotation.
pub fn authenticate(token: &str) -> Result<TokenGrant, TokenError> {
    let cfg = config::get();
    let resource = crate::mcp::resource_url(&cfg.drive_public_url);
    let accepted = [cfg.mcp_audience.as_str(), resource.as_str()];
    let now = chrono::Utc::now().timestamp();
    let check = |jwks: &Value| validate_access_token(token, jwks, &cfg.idp_issuer, &accepted, now);
    let jwks = fetch_jwks(false).map_err(|e| {
        tracing::warn!("{e}");
        TokenError::UnknownKey
    })?;
    match check(&jwks) {
        Err(TokenError::UnknownKey) => {
            let jwks = fetch_jwks(true).map_err(|e| {
                tracing::warn!("{e}");
                TokenError::UnknownKey
            })?;
            check(&jwks)
        }
        other => other,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ed25519_dalek::{Signer, SigningKey};
    use serde_json::json;

    const ISS: &str = "https://auth.olibuijr.com";
    const NOW: i64 = 1_800_000_000;
    const AUD: [&str; 1] = ["akurai-drive"];

    fn key() -> SigningKey {
        SigningKey::from_bytes(&[7u8; 32])
    }

    fn jwks() -> Value {
        json!({"keys": [{
            "kty": "OKP", "crv": "Ed25519", "alg": "EdDSA", "kid": "k1",
            "x": URL_SAFE_NO_PAD.encode(key().verifying_key().to_bytes()),
        }]})
    }

    fn sign(header: Value, claims: Value, signer: &SigningKey) -> String {
        let h = URL_SAFE_NO_PAD.encode(header.to_string());
        let p = URL_SAFE_NO_PAD.encode(claims.to_string());
        let sig = signer.sign(format!("{h}.{p}").as_bytes());
        format!("{h}.{p}.{}", URL_SAFE_NO_PAD.encode(sig.to_bytes()))
    }

    fn claims() -> Value {
        json!({
            "iss": ISS, "sub": "user-1", "exp": NOW + 60, "iat": NOW,
            "aud": ["akurai-agent-akurai-drive", "akurai-drive"],
            "tenant_id": "t1", "scope": "openid profile email",
        })
    }

    fn token(claims: Value) -> String {
        sign(
            json!({"alg": "EdDSA", "kid": "k1", "typ": "JWT"}),
            claims,
            &key(),
        )
    }

    fn check(token: &str) -> Result<TokenGrant, TokenError> {
        validate_access_token(token, &jwks(), ISS, &AUD, NOW)
    }

    #[test]
    fn accepts_idp_user_token_with_array_audience() {
        let grant = check(&token(claims())).unwrap();
        assert_eq!(grant.sub, "user-1");
        assert_eq!(grant.tenant_id.as_deref(), Some("t1"));
        assert_eq!(grant.scopes, ["openid", "profile", "email"]);
    }

    #[test]
    fn accepts_string_audience() {
        let mut c = claims();
        c["aud"] = json!("akurai-drive");
        assert!(check(&token(c)).is_ok());
    }

    #[test]
    fn rejects_other_audience() {
        let mut c = claims();
        c["aud"] = json!(["akurai-agent-akurai-mail", "akurai-mail"]);
        assert_eq!(check(&token(c)), Err(TokenError::WrongAudience));
    }

    #[test]
    fn rejects_expired_and_not_yet_valid() {
        let mut c = claims();
        c["exp"] = json!(NOW);
        assert_eq!(check(&token(c)), Err(TokenError::Expired));
        let mut c = claims();
        c["nbf"] = json!(NOW + 5);
        assert_eq!(check(&token(c)), Err(TokenError::Expired));
    }

    #[test]
    fn rejects_wrong_issuer() {
        let mut c = claims();
        c["iss"] = json!("https://evil.example");
        assert_eq!(check(&token(c)), Err(TokenError::WrongIssuer));
    }

    #[test]
    fn rejects_forged_signature_and_unknown_kid() {
        let forged = sign(
            json!({"alg": "EdDSA", "kid": "k1"}),
            claims(),
            &SigningKey::from_bytes(&[9u8; 32]),
        );
        assert_eq!(check(&forged), Err(TokenError::BadSignature));
        let other_kid = sign(json!({"alg": "EdDSA", "kid": "k2"}), claims(), &key());
        assert_eq!(check(&other_kid), Err(TokenError::UnknownKey));
    }

    #[test]
    fn rejects_non_eddsa_and_garbage() {
        let none = sign(json!({"alg": "none", "kid": "k1"}), claims(), &key());
        assert_eq!(check(&none), Err(TokenError::Malformed));
        assert_eq!(check("not-a-jwt"), Err(TokenError::Malformed));
        assert_eq!(check("a.b.c.d"), Err(TokenError::Malformed));
    }

    #[test]
    fn rejects_client_credentials_tokens() {
        let mut c = claims();
        c["client_id"] = json!("some-service");
        assert_eq!(check(&token(c)), Err(TokenError::NotAUserToken));
        let mut c = claims();
        c["sub"] = json!("akurai-drive");
        assert_eq!(check(&token(c)), Err(TokenError::NotAUserToken));
    }

    #[test]
    fn audience_contains_handles_shapes() {
        assert!(audience_contains(Some(&json!("a")), "a"));
        assert!(audience_contains(Some(&json!(["x", "a"])), "a"));
        assert!(!audience_contains(Some(&json!(["x"])), "a"));
        assert!(!audience_contains(Some(&json!(1)), "a"));
        assert!(!audience_contains(None, "a"));
    }
}