AkurAI-Notes / crates/notes-mcp/tests/oauth_test.rs

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

use akurai_http::{Request, Response};
use akurai_storage::BTree;
use jsonwebtoken::jwk::JwkSet;
use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
use notes_mcp::oauth::{
    audience_contains, identity_from_userinfo, protected_resource_metadata, validate_access_token,
    www_authenticate, IdpClient, OAuth, OAuthConfig,
};
use notes_mcp::server::{McpConfig, McpServer};
use serde_json::{json, Value};
use std::sync::{Arc, Mutex};
use std::time::{SystemTime, UNIX_EPOCH};

const ISS: &str = "https://auth.olibuijr.com";
const AUD: &str = "akurai-notes";
const PUBLIC: &str = "https://akurai-notes.olibuijr.com";
const STATIC_TOKEN: &str = "static-token";

fn jwks() -> JwkSet {
    serde_json::from_value(json!({"keys": [{
        "kty": "OKP", "use": "sig", "crv": "Ed25519", "kid": "fixture", "alg": "EdDSA",
        "x": "2-Jj2UvNCvQiUPNYRgSi0cJSPiJI6Rs6D0UTeEpQVj8"
    }]}))
    .unwrap()
}

fn sign_with_kid(claims: Value, kid: &str) -> String {
    let mut header = Header::new(Algorithm::EdDSA);
    header.kid = Some(kid.into());
    let key = EncodingKey::from_ed_pem(
        b"-----BEGIN PRIVATE KEY-----\n\
          MC4CAQAwBQYDK2VwBCIEIGrD/e7uKYqSY4twDEsRfMMuLSrODf14dpTiTK6K1YI0\n\
          -----END PRIVATE KEY-----\n",
    )
    .unwrap();
    encode(&header, &claims, &key).unwrap()
}

fn sign(claims: Value) -> String {
    sign_with_kid(claims, "fixture")
}

fn now() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs()
}

fn user_claims() -> Value {
    json!({"iss": ISS, "sub": "user-1", "aud": ["akurai-agent-akurai-notes", AUD],
           "exp": now() + 300, "scope": "openid profile email", "tenant_id": "t1",
           "organization_id": "o1"})
}

fn config() -> OAuthConfig {
    OAuthConfig {
        public_url: format!("{PUBLIC}/"),
        issuer: ISS.into(),
        audience: AUD.into(),
        allow: vec!["owner@example.com".into()],
    }
}

struct FakeIdp {
    userinfo: Value,
}

impl IdpClient for FakeIdp {
    fn jwks(&self, _force: bool) -> Result<JwkSet, &'static str> {
        Ok(jwks())
    }
    fn userinfo(&self, _token: &str) -> Result<Value, &'static str> {
        Ok(self.userinfo.clone())
    }
}

fn owner_userinfo() -> Value {
    json!({"sub": "user-1", "email": "Owner@Example.com", "email_verified": true})
}

fn server_with(userinfo: Value) -> McpServer {
    let dir = std::env::temp_dir().join(format!("notes-oauth-test-{}", std::process::id()));
    let _ = std::fs::create_dir_all(&dir);
    let db_path = dir.join(format!(
        "test-{}.db",
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos()
    ));
    let db = BTree::open(&db_path).unwrap();
    McpServer::with_db(
        McpConfig {
            host: "127.0.0.1".into(),
            port: 0,
            db: db_path.display().to_string(),
            token: STATIC_TOKEN.into(),
            embed_url: None,
            embed_model: None,
            embed_api_key: None,
            owner_sub: "owner".into(),
        },
        Arc::new(Mutex::new(db)),
    )
    .with_oauth(OAuth::new(config(), Box::new(FakeIdp { userinfo })))
}

fn request(head: &str) -> Request {
    let mut req = Request::parse_head(head.as_bytes()).unwrap();
    req.body = br#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#.to_vec();
    req
}

fn mcp_post(extra: &str) -> Request {
    request(&format!("POST /mcp HTTP/1.1\r\n{extra}\r\n"))
}

fn header<'a>(resp: &'a Response, name: &str) -> Option<&'a str> {
    resp.headers
        .iter()
        .find(|(k, _)| k.eq_ignore_ascii_case(name))
        .map(|(_, v)| v.as_str())
}

// --- pure helpers --------------------------------------------------------

#[test]
fn metadata_names_resource_and_idp() {
    let m = protected_resource_metadata(&config());
    assert_eq!(m["resource"], "https://akurai-notes.olibuijr.com/mcp");
    assert_eq!(m["authorization_servers"], json!([ISS]));
    assert_eq!(m["bearer_methods_supported"], json!(["header"]));
    assert!(m["scopes_supported"]
        .as_array()
        .unwrap()
        .contains(&json!("openid")));
}

#[test]
fn challenge_points_at_metadata() {
    assert_eq!(
        www_authenticate(&config(), false),
        "Bearer resource_metadata=\"https://akurai-notes.olibuijr.com/.well-known/oauth-protected-resource/mcp\""
    );
    assert!(www_authenticate(&config(), true).starts_with("Bearer error=\"invalid_token\", "));
}

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

#[test]
fn accepts_idp_user_token_with_app_audience() {
    let grant = validate_access_token(&sign(user_claims()), &jwks(), ISS, AUD).unwrap();
    assert_eq!(grant.sub, "user-1");
    assert!(grant.scopes.contains("email"));
}

#[test]
fn rejects_wrong_audience_issuer_expiry_and_key() {
    let mut c = user_claims();
    c["aud"] = json!(["akurai-agent-other", "other-app"]);
    assert!(validate_access_token(&sign(c), &jwks(), ISS, AUD).is_err());

    let mut c = user_claims();
    c["iss"] = json!("https://evil.test");
    assert!(validate_access_token(&sign(c), &jwks(), ISS, AUD).is_err());

    let mut c = user_claims();
    c["exp"] = json!(now() - 10);
    assert!(validate_access_token(&sign(c), &jwks(), ISS, AUD).is_err());

    assert_eq!(
        validate_access_token(&sign_with_kid(user_claims(), "other"), &jwks(), ISS, AUD),
        Err("unknown JWT key id")
    );
    assert!(validate_access_token("not.a.jwt", &jwks(), ISS, AUD).is_err());
}

#[test]
fn rejects_client_credentials_tokens() {
    let mut c = user_claims();
    c["client_id"] = json!("akurai-agent-akurai-notes");
    assert!(validate_access_token(&sign(c), &jwks(), ISS, AUD).is_err());
}

#[test]
fn userinfo_identity_must_match_and_be_allowed() {
    let allow = vec!["owner@example.com".to_string()];
    assert_eq!(
        identity_from_userinfo(&owner_userinfo(), "user-1", &allow).unwrap(),
        "owner@example.com"
    );
    assert!(identity_from_userinfo(&owner_userinfo(), "user-2", &allow).is_err());
    let stranger = json!({"sub": "user-1", "email": "x@example.com", "email_verified": true});
    assert!(identity_from_userinfo(&stranger, "user-1", &allow).is_err());
    let unverified =
        json!({"sub": "user-1", "email": "owner@example.com", "email_verified": false});
    assert!(identity_from_userinfo(&unverified, "user-1", &allow).is_err());
}

// --- HTTP surface --------------------------------------------------------

#[test]
fn serves_protected_resource_metadata_on_both_paths() {
    let server = server_with(owner_userinfo());
    for path in [
        "/.well-known/oauth-protected-resource",
        "/.well-known/oauth-protected-resource/mcp",
    ] {
        let resp = server.handle_http(&request(&format!("GET {path} HTTP/1.1\r\n\r\n")));
        assert_eq!(resp.status, 200, "{path}");
        let body: Value = serde_json::from_slice(&resp.body).unwrap();
        assert_eq!(body["resource"], "https://akurai-notes.olibuijr.com/mcp");
    }
}

#[test]
fn missing_token_gets_resource_metadata_challenge() {
    let resp = server_with(owner_userinfo()).handle_http(&mcp_post(""));
    assert_eq!(resp.status, 401);
    assert_eq!(
        header(&resp, "WWW-Authenticate"),
        Some(www_authenticate(&config(), false).as_str())
    );
}

#[test]
fn rejected_token_gets_invalid_token_challenge() {
    let resp =
        server_with(owner_userinfo()).handle_http(&mcp_post("Authorization: Bearer nope\r\n"));
    assert_eq!(resp.status, 401);
    assert_eq!(
        header(&resp, "WWW-Authenticate"),
        Some(www_authenticate(&config(), true).as_str())
    );
}

#[test]
fn static_token_still_works() {
    let resp = server_with(owner_userinfo()).handle_http(&mcp_post(&format!(
        "Authorization: Bearer {STATIC_TOKEN}\r\n"
    )));
    assert_eq!(resp.status, 200);
}

#[test]
fn idp_token_for_allowed_user_is_accepted() {
    let token = sign(user_claims());
    let resp = server_with(owner_userinfo())
        .handle_http(&mcp_post(&format!("Authorization: Bearer {token}\r\n")));
    assert_eq!(resp.status, 200);
    assert!(String::from_utf8(resp.body).unwrap().contains("list_notes"));
}

#[test]
fn idp_token_for_disallowed_user_is_rejected() {
    let token = sign(user_claims());
    let stranger = json!({"sub": "user-1", "email": "x@example.com", "email_verified": true});
    let resp =
        server_with(stranger).handle_http(&mcp_post(&format!("Authorization: Bearer {token}\r\n")));
    assert_eq!(resp.status, 401);
}

#[test]
fn foreign_browser_origin_is_forbidden() {
    let resp = server_with(owner_userinfo()).handle_http(&mcp_post(&format!(
        "Authorization: Bearer {STATIC_TOKEN}\r\nOrigin: https://evil.test\r\n"
    )));
    assert_eq!(resp.status, 403);
    let same = server_with(owner_userinfo()).handle_http(&mcp_post(&format!(
        "Authorization: Bearer {STATIC_TOKEN}\r\nOrigin: {PUBLIC}\r\n"
    )));
    assert_eq!(same.status, 200);
}