AkurAI-IDP / src/routes/revoke.rs

Latest change 26d50ee76eb4c6c0aaca914e91a507fae2a243a6 - feat: add tenant-safe organization identity by Ólafur Búi Ólafsson

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

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

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

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

async fn revoke_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: RevokeRequest =
        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.into_response(),
    };
    // RFC 7009 token_type_hint is advisory.  Always try this client's
    // refresh-token namespace; JWT access tokens have no per-token revocation
    // record and therefore remain a non-enumerating 200 response.
    let result = with_db(|conn| {
        revoke_refresh_token(conn, &token, client_id, params.token_type_hint.as_deref())
    });
    if result.is_err() {
        return StatusCode::INTERNAL_SERVER_ERROR.into_response();
    }
    // User access tokens are invalidated by authorization-version or
    // membership changes; machine JWTs by client lifecycle changes or their
    // short expiry.  This endpoint does not claim per-token JWT revocation.
    StatusCode::OK.into_response()
}

fn revoke_refresh_token(
    conn: &rusqlite::Connection,
    token: &str,
    client_id: &str,
    _token_type_hint: Option<&str>,
) -> rusqlite::Result<usize> {
    conn.execute(
        "UPDATE refresh_tokens SET revoked = 1
         WHERE token_hash = ?1 AND client_id = ?2 AND revoked = 0",
        rusqlite::params![sha256(token), client_id],
    )
}

fn verify_client_secret_with_conn(
    conn: &rusqlite::Connection,
    client_id: &str,
    secret: Option<&str>,
) -> bool {
    let Some(secret) = secret.filter(|value| !value.is_empty()) else {
        return false;
    };
    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: String| constant_time_equal(&sha256(secret), &hash))
}

fn verify_client_secret(client_id: &str, secret: Option<&str>) -> bool {
    with_db(|conn| verify_client_secret_with_conn(conn, client_id, secret))
}

fn parse_form_urlencoded(bytes: &[u8]) -> RevokeRequest {
    let s = std::str::from_utf8(bytes).unwrap_or("");
    let mut params = RevokeRequest::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),
    }
}

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

    fn fixture() -> (Connection, String) {
        let conn = Connection::open_in_memory().expect("in-memory db");
        create_tables(&conn);
        conn.execute(
            "INSERT INTO tenants (id, name, slug, created_at)
             VALUES ('tenant-a', 'Tenant A', 'tenant-a', 1)",
            [],
        )
        .expect("tenant");
        conn.execute(
            "INSERT INTO users
             (id, tenant_id, email, password_hash, created_at, updated_at)
             VALUES ('user-a', 'tenant-a', 'user@example.test', 'hash', 1, 1)",
            [],
        )
        .expect("user");
        conn.execute(
            "INSERT INTO clients
             (id, tenant_id, name, client_secret_hash, redirect_uris, grant_types, scopes, created_at)
             VALUES ('client-a', 'tenant-a', 'Client A', ?1, '', 'refresh_token', 'openid', 1)",
            [sha256("secret-a")],
        )
        .expect("client");
        let token = "refresh-token-a".to_string();
        conn.execute(
            "INSERT INTO refresh_tokens
             (token_hash, client_id, user_id, scopes, expires_at)
             VALUES (?1, 'client-a', 'user-a', 'openid', 9999999999)",
            [sha256(&token)],
        )
        .expect("refresh token");
        (conn, token)
    }

    fn revoked(conn: &Connection, token: &str) -> i64 {
        conn.query_row(
            "SELECT revoked FROM refresh_tokens WHERE token_hash = ?1",
            [sha256(token)],
            |row| row.get(0),
        )
        .expect("revocation row")
    }

    #[test]
    fn absent_refresh_token_hint_still_revokes_once() {
        let (conn, token) = fixture();
        assert_eq!(
            revoke_refresh_token(&conn, &token, "client-a", None).expect("revoke"),
            1
        );
        assert_eq!(revoked(&conn, &token), 1);
        assert_eq!(
            revoke_refresh_token(&conn, &token, "client-a", None).expect("revoke again"),
            0
        );
    }

    #[test]
    fn refresh_token_hint_revokes_refresh_token() {
        let (conn, token) = fixture();
        revoke_refresh_token(&conn, &token, "client-a", Some("refresh_token")).expect("revoke");
        assert_eq!(revoked(&conn, &token), 1);
    }

    #[test]
    fn access_token_hint_is_advisory_and_does_not_claim_jwt_revocation() {
        let (conn, token) = fixture();
        revoke_refresh_token(&conn, &token, "client-a", Some("access_token")).expect("revoke");
        assert_eq!(revoked(&conn, &token), 1);
    }

    #[test]
    fn unsupported_hint_still_revokes_refresh_token() {
        let (conn, token) = fixture();
        revoke_refresh_token(&conn, &token, "client-a", Some("id_token")).expect("revoke");
        assert_eq!(revoked(&conn, &token), 1);
    }

    #[test]
    fn wrong_client_cannot_revoke_refresh_token() {
        let (conn, token) = fixture();
        assert_eq!(
            revoke_refresh_token(&conn, &token, "other-client", Some("refresh_token"))
                .expect("revoke"),
            0
        );
        assert_eq!(revoked(&conn, &token), 0);
    }

    #[test]
    fn fake_client_cannot_authenticate() {
        let (conn, _token) = fixture();
        assert!(!verify_client_secret_with_conn(
            &conn,
            "fake-client",
            Some("secret-a")
        ));
    }
}