Mercury/src/auth/middleware.rs

use axum::{
    extract::{Request, State},
    http::{header, HeaderMap, HeaderValue, Method, StatusCode},
    middleware::Next,
    response::{IntoResponse, Response},
};

use crate::{
    auth::{decode_jwt, resolve_api_key, Claims},
    state::AppState,
};

/// Resolves a Bearer token to Claims, trying JWT then API key.
/// Inserts Claims into request extensions on success so downstream
/// middleware can reuse them without an additional DB round-trip.
async fn authenticate(token: &str, state: &AppState, req: &mut Request) -> Option<Claims> {
    if let Some(existing) = req.extensions().get::<Claims>().cloned() {
        return Some(existing);
    }
    let claims = if let Ok(c) = decode_jwt(token, &state.config.jwt_secret) {
        c
    } else {
        resolve_api_key(token, &state.pool).await?
    };
    req.extensions_mut().insert(claims.clone());
    Some(claims)
}

pub async fn blacklist_layer(
    State(state): State<AppState>,
    mut req: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    let method = req.method().as_str().to_uppercase();
    let path = req.uri().path().to_string();
    let caller_mask = if let Some(token) = extract_bearer(&req) {
        authenticate(&token, &state, &mut req)
            .await
            .map(|c| c.permissions_mask())
            .unwrap_or(0)
    } else {
        0
    };
    if state
        .blacklist_cache
        .is_blocked(&method, &path, caller_mask)
        .await
    {
        return Err(StatusCode::FORBIDDEN);
    }
    Ok(next.run(req).await)
}

pub fn extract_bearer(req: &Request) -> Option<String> {
    req.headers()
        .get("authorization")
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer "))
        .map(|s| s.to_string())
}

pub async fn require_auth(
    State(state): State<AppState>,
    mut req: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    let token = extract_bearer(&req).ok_or(StatusCode::UNAUTHORIZED)?;
    authenticate(&token, &state, &mut req)
        .await
        .ok_or(StatusCode::UNAUTHORIZED)?;
    Ok(next.run(req).await)
}

pub async fn require_super_admin(
    State(state): State<AppState>,
    mut req: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    let token = extract_bearer(&req).ok_or(StatusCode::UNAUTHORIZED)?;
    let claims = authenticate(&token, &state, &mut req)
        .await
        .ok_or(StatusCode::UNAUTHORIZED)?;
    if !claims.has_permission(crate::auth::permissions::SUPER_ADMIN) {
        return Err(StatusCode::FORBIDDEN);
    }
    Ok(next.run(req).await)
}

pub async fn require_admin_query(
    State(state): State<AppState>,
    mut req: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    let token = extract_bearer(&req).ok_or(StatusCode::UNAUTHORIZED)?;
    let claims = authenticate(&token, &state, &mut req)
        .await
        .ok_or(StatusCode::UNAUTHORIZED)?;
    if !claims.has_permission(crate::auth::permissions::ADMIN_QUERY) {
        return Err(StatusCode::FORBIDDEN);
    }
    Ok(next.run(req).await)
}

pub async fn require_admin_cache(
    State(state): State<AppState>,
    mut req: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    let token = extract_bearer(&req).ok_or(StatusCode::UNAUTHORIZED)?;
    let claims = authenticate(&token, &state, &mut req)
        .await
        .ok_or(StatusCode::UNAUTHORIZED)?;
    if !claims.has_permission(crate::auth::permissions::ADMIN_CACHE) {
        return Err(StatusCode::FORBIDDEN);
    }
    Ok(next.run(req).await)
}

pub async fn cors_layer(State(state): State<AppState>, req: Request, next: Next) -> Response {
    let is_preflight = req.method() == Method::OPTIONS;
    let origin_header = req.headers().get(header::ORIGIN).cloned();

    enum CorsDecision {
        None,
        Wildcard,
        Specific(HeaderValue),
    }

    let decision = {
        let guard = state.cors_cache.inner.read().await;
        if guard.wildcard {
            CorsDecision::Wildcard
        } else if let Some(origin) = origin_header.as_ref() {
            if guard.origins.contains(origin) {
                CorsDecision::Specific(origin.clone())
            } else {
                CorsDecision::None
            }
        } else {
            CorsDecision::None
        }
    };

    let (cors_origin, vary) = match &decision {
        CorsDecision::None => (None, false),
        CorsDecision::Wildcard => (Some(HeaderValue::from_static("*")), false),
        CorsDecision::Specific(v) => (Some(v.clone()), true),
    };

    if is_preflight {
        let mut headers = HeaderMap::new();
        if let Some(origin) = cors_origin {
            headers.insert(header::ACCESS_CONTROL_ALLOW_ORIGIN, origin);
            headers.insert(
                header::ACCESS_CONTROL_ALLOW_METHODS,
                HeaderValue::from_static("GET, POST, PUT, DELETE, OPTIONS"),
            );
            headers.insert(
                header::ACCESS_CONTROL_ALLOW_HEADERS,
                HeaderValue::from_static("content-type, authorization"),
            );
            if vary {
                headers.insert(header::VARY, HeaderValue::from_static("Origin"));
            }
        }
        return (StatusCode::NO_CONTENT, headers).into_response();
    }

    let mut response = next.run(req).await;
    if let Some(origin) = cors_origin {
        response
            .headers_mut()
            .insert(header::ACCESS_CONTROL_ALLOW_ORIGIN, origin);
        if vary {
            response
                .headers_mut()
                .append(header::VARY, HeaderValue::from_static("Origin"));
        }
    }
    response
}