AkurAI-IDP / src/routes/introspect.rs

Latest change 65d9e4103c7a6a1601ba4b287022e495bfac7f18 - tokens: aud lists client id and resource audience when they differ by Ólafur Búi Ólafsson

use axum::{
    Json, Router, extract::Request, http::StatusCode, response::IntoResponse, routing::post,
};
use serde::Deserialize;
use serde_json::json;

use crate::db::with_db;
use crate::lib::crypto::{constant_time_equal, sha256};
use crate::lib::jwt::verify_active_token;

pub fn router() -> Router {
    Router::new().route("/introspect", post(introspect_endpoint))
}

#[derive(Deserialize, Default)]
struct IntrospectRequest {
    token: Option<String>,
    token_type_hint: Option<String>,
    client_id: Option<String>,
    #[allow(dead_code)]
    client_secret: Option<String>,
}

async fn introspect_endpoint(request: Request) -> impl IntoResponse {
    let (basic_id, basic_secret) = extract_basic_auth(&request);
    let body_bytes = match axum::body::to_bytes(request.into_body(), 1024 * 64).await {
        Ok(b) => b,
        Err(_) => {
            return (
                StatusCode::BAD_REQUEST,
                Json(json!({"error": "invalid_request"})),
            )
                .into_response();
        }
    };
    let mut params: IntrospectRequest =
        serde_json::from_slice(&body_bytes).unwrap_or_else(|_| parse_form_urlencoded(&body_bytes));
    if let Some(id) = basic_id {
        params.client_id = Some(id);
        params.client_secret = basic_secret;
    }
    let client_id = match params.client_id.as_deref() {
        Some(id) if !id.is_empty() => id,
        _ => {
            return (
                StatusCode::UNAUTHORIZED,
                Json(json!({"error": "invalid_client"})),
            )
                .into_response();
        }
    };
    if !verify_client_secret(client_id, params.client_secret.as_deref()) {
        return (
            StatusCode::UNAUTHORIZED,
            Json(json!({"error": "invalid_client"})),
        )
            .into_response();
    }
    let token = match params.token {
        Some(t) if !t.is_empty() => t,
        _ => return (StatusCode::OK, Json(json!({"active": false}))).into_response(),
    };
    let hint = params.token_type_hint.as_deref();

    if hint != Some("access_token") {
        let token_hash = sha256(&token);
        let result = with_db(|conn| {
            conn.query_row(
                "SELECT rt.client_id, rt.user_id, rt.scopes, rt.expires_at, rt.revoked,
                        rt.organization_id, u.email
                 FROM refresh_tokens rt JOIN users u ON u.id = rt.user_id
                 WHERE rt.token_hash = ?1 AND rt.client_id = ?2",
                rusqlite::params![token_hash, client_id],
                |row| {
                    let expires_at: i64 = row.get(3)?;
                    let revoked: bool = row.get::<_, i64>(4)? != 0;
                    let now = now_secs();
                    Ok(json!({
                        "active": !revoked && expires_at > now,
                        "token_type": "refresh_token",
                        "client_id": row.get::<_, String>(0)?,
                        "sub": row.get::<_, String>(1)?,
                        "scope": row.get::<_, String>(2)?,
                        "exp": expires_at,
                        "email": row.get::<_, String>(6)?,
                        "organization_id": row.get::<_, Option<String>>(5)?,
                    }))
                },
            )
            .ok()
        });
        if let Some(resp) = result {
            return (StatusCode::OK, Json(resp)).into_response();
        }
    }

    let claims = match verify_active_token(&token) {
        Some(claims) => claims,
        None => return (StatusCode::OK, Json(json!({"active": false}))).into_response(),
    };
    if !crate::lib::clients::audience_contains(claims.get("aud"), client_id) {
        return (StatusCode::OK, Json(json!({"active": false}))).into_response();
    }
    let now = now_secs();
    let exp = claims.get("exp").and_then(|v| v.as_i64()).unwrap_or(0);
    let mut active = exp > now;
    if active && claims.get("token_type").and_then(|v| v.as_str()) == Some("client_credentials") {
        active = with_db(|conn| {
            conn.query_row(
                "SELECT COUNT(*) FROM clients WHERE id = ?1 AND tenant_id = ?2",
                rusqlite::params![
                    client_id,
                    claims
                        .get("organization_id")
                        .and_then(|v| v.as_str())
                        .unwrap_or("")
                ],
                |row| row.get::<_, i64>(0),
            )
            .map(|count| count == 1)
            .unwrap_or(false)
        });
    } else {
        let claims_version = claims.get("claims_version").and_then(|v| v.as_i64());
        let authorization_version = claims.get("authorization_version").and_then(|v| v.as_i64());
        let user_id = claims.get("sub").and_then(|v| v.as_str()).unwrap_or("");
        let organization_id = claims
            .get("organization_id")
            .or_else(|| claims.get("tenant_id"))
            .and_then(|v| v.as_str())
            .unwrap_or("");
        active = active
            && claims_version == Some(crate::lib::organizations::CLAIMS_VERSION)
            && with_db(|conn| {
                crate::lib::organizations::active_membership(conn, user_id, organization_id)
                    .ok()
                    .flatten()
                    .is_some_and(|m| Some(m.authorization_version) == authorization_version)
            });
    }
    let mut result = json!({"active": active, "token_type": "access_token"});
    if let Some(obj) = claims.as_object() {
        for (key, value) in obj {
            result[key] = value.clone();
        }
    }
    result["active"] = json!(active);
    (StatusCode::OK, Json(result)).into_response()
}

fn now_secs() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs() as i64
}

fn verify_client_secret(client_id: &str, secret: Option<&str>) -> bool {
    let Some(secret) = secret.filter(|value| !value.is_empty()) else {
        return false;
    };
    with_db(|conn| {
        let stored: Option<String> = conn
            .query_row(
                "SELECT client_secret_hash FROM clients WHERE id = ?1",
                [client_id],
                |row| row.get(0),
            )
            .ok();
        stored.is_some_and(|hash| constant_time_equal(&sha256(secret), &hash))
    })
}

fn parse_form_urlencoded(bytes: &[u8]) -> IntrospectRequest {
    let s = std::str::from_utf8(bytes).unwrap_or("");
    let mut params = IntrospectRequest::default();
    for pair in s.split('&') {
        if let Some((key, value)) = pair.split_once('=') {
            let value = urldecode(value);
            match key {
                "token" => params.token = Some(value),
                "token_type_hint" => params.token_type_hint = Some(value),
                "client_id" => params.client_id = Some(value),
                "client_secret" => params.client_secret = Some(value),
                _ => {}
            }
        }
    }
    params
}

fn urldecode(s: &str) -> String {
    let s = s.replace('+', " ");
    let mut result = Vec::with_capacity(s.len());
    let mut chars = s.bytes();
    while let Some(b) = chars.next() {
        if b == b'%' {
            let hi = chars.next().unwrap_or(b'0');
            let lo = chars.next().unwrap_or(b'0');
            let hex = [hi, lo];
            if let Ok(val) = u8::from_str_radix(std::str::from_utf8(&hex).unwrap_or("00"), 16) {
                result.push(val);
            }
        } else {
            result.push(b);
        }
    }
    String::from_utf8(result).unwrap_or_default()
}

fn extract_basic_auth(request: &Request) -> (Option<String>, Option<String>) {
    let header = match request.headers().get(axum::http::header::AUTHORIZATION) {
        Some(h) => h,
        None => return (None, None),
    };
    let value = match header.to_str().ok() {
        Some(v) => v,
        None => return (None, None),
    };
    let encoded = match value.strip_prefix("Basic ") {
        Some(e) => e,
        None => return (None, None),
    };
    use base64::Engine;
    let decoded = match base64::engine::general_purpose::STANDARD.decode(encoded) {
        Ok(d) => d,
        Err(_) => return (None, None),
    };
    let decoded_str = match String::from_utf8(decoded) {
        Ok(s) => s,
        Err(_) => return (None, None),
    };
    match decoded_str.split_once(':') {
        Some((id, secret)) => (Some(id.to_string()), Some(secret.to_string())),
        None => (None, None),
    }
}