AkurAI-IDP / src/lib/clients.rs

Latest change fce146c19e4ee44321ed964db4111dca7d75cc8b - recover: MCP OAuth public clients, RFC 8414 metadata, /api/health (0.8.4) by Ólafur Búi Ólafsson

//! OAuth client policy: public (PKCE-only) clients, token audiences and the
//! built-in MCP clients that let AkurAI agents reach AkurAI app MCP servers.

use rusqlite::{Connection, OptionalExtension};

use crate::lib::crypto::{constant_time_equal, sha256};

pub const AUTH_METHOD_NONE: &str = "none";

/// Loopback redirect URIs used by agent MCP clients. Agents run on Titan and
/// complete the OAuth code flow on a fixed local port.
const AGENT_REDIRECT_URIS: &str =
    r#"["http://localhost:47840/callback","http://127.0.0.1:47840/callback"]"#;
const AGENT_GRANT_TYPES: &str = r#"["authorization_code","refresh_token"]"#;
const APP_SCOPES: &str = r#"["openid","profile","email","groups","offline_access"]"#;
const MAIL_SCOPES: &str =
    r#"["openid","email","offline_access","mail.read","mail.draft","mail.send"]"#;

/// (client id, display name, scopes, audience) for every AkurAI app MCP server.
const AGENT_MCP_CLIENTS: &[(&str, &str, &str, &str)] = &[
    (
        "akurai-agent-mail",
        "AkurAI Agent Mail",
        MAIL_SCOPES,
        "https://mail.olibuijr.com/mcp",
    ),
    (
        "akurai-agent-notes",
        "AkurAI Agent Notes",
        APP_SCOPES,
        "akurai-notes",
    ),
    (
        "akurai-agent-passvault",
        "AkurAI Agent PassVault",
        APP_SCOPES,
        "akurai-passvault",
    ),
    (
        "akurai-agent-tasks",
        "AkurAI Agent Tasks",
        APP_SCOPES,
        "akurai-tasks",
    ),
    (
        "akurai-agent-crm",
        "AkurAI Agent CRM",
        APP_SCOPES,
        "akurai-crm",
    ),
    (
        "akurai-agent-build",
        "AkurAI Agent Build",
        APP_SCOPES,
        "akurai-build",
    ),
    (
        "akurai-agent-drive",
        "AkurAI Agent Drive",
        APP_SCOPES,
        "akurai-drive",
    ),
];

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientPolicy {
    pub tenant_id: String,
    pub auth_method: String,
    pub platform_wide: bool,
    pub audience: String,
}

impl ClientPolicy {
    pub fn is_public(&self) -> bool {
        self.auth_method == AUTH_METHOD_NONE
    }

    /// Platform-wide clients serve every organization; others only their own.
    pub fn allows_organization(&self, organization_id: &str) -> bool {
        self.platform_wide || self.tenant_id == organization_id
    }
}

pub fn load_policy(conn: &Connection, client_id: &str) -> Option<ClientPolicy> {
    conn.query_row(
        "SELECT tenant_id, token_endpoint_auth_method, platform_wide, audience
         FROM clients WHERE id = ?1",
        [client_id],
        |row| {
            Ok(ClientPolicy {
                tenant_id: row.get(0)?,
                auth_method: row.get(1)?,
                platform_wide: row.get::<_, i64>(2)? != 0,
                audience: row.get(3)?,
            })
        },
    )
    .optional()
    .ok()
    .flatten()
}

/// Authenticate a client at the token endpoint. Public clients carry no
/// secret; their proof of possession is the PKCE verifier checked by the caller.
pub fn authenticate(conn: &Connection, client_id: &str, secret: Option<&str>) -> bool {
    let Some(policy) = load_policy(conn, client_id) else {
        return false;
    };
    if policy.is_public() {
        return true;
    }
    let Some(secret) = secret.filter(|s| !s.is_empty()) else {
        return false;
    };
    conn.query_row(
        "SELECT client_secret_hash FROM clients WHERE id = ?1",
        [client_id],
        |row| row.get::<_, String>(0),
    )
    .map(|hash| !hash.is_empty() && constant_time_equal(&sha256(secret), &hash))
    .unwrap_or(false)
}

/// Access-token audience for a client; falls back to the client id.
pub fn audience(conn: &Connection, client_id: &str) -> String {
    load_policy(conn, client_id)
        .map(|p| p.audience)
        .filter(|a| !a.trim().is_empty())
        .unwrap_or_else(|| client_id.to_string())
}

/// Register the built-in agent MCP clients under the oldest organization.
/// Existing rows are never overwritten.
pub fn register_builtin_clients(conn: &Connection, now: i64) -> rusqlite::Result<()> {
    let Some(organization_id) = conn
        .query_row(
            "SELECT id FROM organizations ORDER BY created_at, id LIMIT 1",
            [],
            |row| row.get::<_, String>(0),
        )
        .optional()?
    else {
        return Ok(());
    };
    for (id, name, scopes, audience) in AGENT_MCP_CLIENTS {
        conn.execute(
            "INSERT OR IGNORE INTO clients
             (id, tenant_id, name, client_secret_hash, redirect_uris, grant_types,
              scopes, first_party, token_endpoint_auth_method, platform_wide, audience, created_at)
             VALUES (?1, ?2, ?3, '', ?4, ?5, ?6, 1, 'none', 0, ?7, ?8)",
            rusqlite::params![
                id,
                organization_id,
                name,
                AGENT_REDIRECT_URIS,
                AGENT_GRANT_TYPES,
                scopes,
                audience,
                now
            ],
        )?;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::db::schema::create_tables;

    fn fixture() -> Connection {
        let conn = Connection::open_in_memory().expect("in-memory db");
        create_tables(&conn);
        conn.execute(
            "INSERT INTO tenants (id, name, slug, created_at) VALUES ('org-a', 'A', 'a', 1)",
            [],
        )
        .expect("tenant");
        create_tables(&conn);
        conn.execute(
            "INSERT INTO clients (id, tenant_id, name, client_secret_hash, redirect_uris,
             grant_types, scopes, created_at)
             VALUES ('conf', 'org-a', 'Confidential', ?1, '[]', '[]', '[]', 1)",
            [sha256("s3cret")],
        )
        .expect("confidential client");
        conn
    }

    #[test]
    fn confidential_client_requires_matching_secret() {
        let conn = fixture();
        assert!(authenticate(&conn, "conf", Some("s3cret")));
        assert!(!authenticate(&conn, "conf", Some("wrong")));
        assert!(!authenticate(&conn, "conf", None));
    }

    #[test]
    fn legacy_clients_default_to_secret_auth_and_own_audience() {
        let conn = fixture();
        let policy = load_policy(&conn, "conf").expect("policy");
        assert_eq!(policy.auth_method, "client_secret_basic");
        assert!(!policy.is_public());
        assert_eq!(audience(&conn, "conf"), "conf");
    }

    #[test]
    fn builtin_agent_clients_are_public_and_bound_to_app_audiences() {
        let conn = fixture();
        register_builtin_clients(&conn, 5).expect("register");
        let mail = load_policy(&conn, "akurai-agent-mail").expect("mail client");
        assert!(mail.is_public());
        assert_eq!(mail.tenant_id, "org-a");
        assert_eq!(
            audience(&conn, "akurai-agent-mail"),
            "https://mail.olibuijr.com/mcp"
        );
        assert_eq!(audience(&conn, "akurai-agent-notes"), "akurai-notes");
        assert!(authenticate(&conn, "akurai-agent-notes", None));
    }

    #[test]
    fn only_akurai_agent_clients_are_registered() {
        let conn = fixture();
        register_builtin_clients(&conn, 5).expect("register");
        let foreign: i64 = conn
            .query_row(
                "SELECT COUNT(*) FROM clients WHERE token_endpoint_auth_method = 'none'
                 AND id NOT LIKE 'akurai-agent-%'",
                [],
                |row| row.get(0),
            )
            .expect("count");
        assert_eq!(foreign, 0);
    }

    #[test]
    fn registration_never_overwrites_existing_rows() {
        let conn = fixture();
        register_builtin_clients(&conn, 5).expect("register");
        conn.execute(
            "UPDATE clients SET name = 'Custom' WHERE id = 'akurai-agent-mail'",
            [],
        )
        .expect("rename");
        register_builtin_clients(&conn, 9).expect("re-register");
        let name: String = conn
            .query_row(
                "SELECT name FROM clients WHERE id = 'akurai-agent-mail'",
                [],
                |r| r.get(0),
            )
            .expect("name");
        assert_eq!(name, "Custom");
    }

    #[test]
    fn platform_wide_clients_allow_any_organization() {
        let own = ClientPolicy {
            tenant_id: "org-a".into(),
            auth_method: "none".into(),
            platform_wide: false,
            audience: "x".into(),
        };
        assert!(own.allows_organization("org-a"));
        assert!(!own.allows_organization("org-b"));
        let wide = ClientPolicy {
            platform_wide: true,
            ..own
        };
        assert!(wide.allows_organization("org-b"));
    }

    #[test]
    fn unknown_client_fails_authentication() {
        let conn = fixture();
        assert!(!authenticate(&conn, "missing", None));
    }
}