assay-workflow 0.2.0

//! Auth middleware + API-key helpers.
//!
//! `AuthMode`, `JwtConfig`, and `JwksCache` are defined in
//! `crate::auth_mode` (no crate-internal deps) and re-exported here so
//! the public path `assay_workflow::api::auth::AuthMode` keeps working.

use std::sync::Arc;

use axum::extract::{Request, State};
use axum::http::StatusCode;
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use axum::Json;
use sha2::{Digest, Sha256};
use tracing::{info, warn};

use crate::ctx::WorkflowCtx;
use crate::store::WorkflowStore;

// Re-export the config types from their cycle-free home.
pub use crate::auth_mode::{AuthMode, JwksCache, JwtConfig};

// ── Middleware ───────────────────────────────────────────────

/// Axum middleware that enforces authentication based on the configured mode.
///
/// When both JWT and API-key auth are enabled, dispatch is based on token shape:
/// if the Bearer token parses as a JWS header it takes the JWT path, otherwise
/// the API-key path. A semantically-invalid JWT (expired, forged signature, wrong
/// audience) is rejected and is *not* retried as an API key — a token that looks
/// like a JWT is treated as a JWT.
///
/// **Bootstrap window:** `POST /api/v1/api-keys` is accepted without a Bearer
/// token iff the `api_keys` table is empty. This is the only way a freshly
/// deployed server running in API-key or combined mode can receive its first
/// credential without operator shell access. The window closes the moment any
/// key exists.
pub async fn auth_middleware<S: WorkflowStore>(
    State(state): State<Arc<WorkflowCtx<S>>>,
    request: Request,
    next: Next,
) -> Response {
    let auth = &state.auth_mode;

    if !auth.is_enabled() {
        return next.run(request).await;
    }

    if is_bootstrap_request(&request) {
        match state.store().api_keys_empty().await {
            Ok(true) => {
                info!(
                    "Allowing unauthenticated POST /api/v1/api-keys — api_keys table is empty (bootstrap window)"
                );
                return next.run(request).await;
            }
            Ok(false) => {
                // Fall through to normal auth — bootstrap window is closed.
            }
            Err(e) => {
                warn!("api_keys_empty check failed: {e}");
                return (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    Json(serde_json::json!({"error": "auth bootstrap check failed"})),
                )
                    .into_response();
            }
        }
    }

    let token = match extract_bearer(&request) {
        Some(t) => t,
        None => return auth_error("Missing Authorization: Bearer <token>"),
    };

    if jsonwebtoken::decode_header(token).is_ok() {
        match &auth.jwt {
            Some(jwt) => {
                validate_jwt(
                    &jwt.issuer,
                    jwt.audience.as_deref(),
                    &jwt.jwks_cache,
                    request,
                    next,
                )
                .await
            }
            None => auth_error("JWT authentication is not enabled on this server"),
        }
    } else if auth.api_key {
        validate_api_key(state, request, next).await
    } else {
        auth_error("Token is not a valid JWT and API-key authentication is not enabled")
    }
}

async fn validate_api_key<S: WorkflowStore>(
    state: Arc<WorkflowCtx<S>>,
    request: Request,
    next: Next,
) -> Response {
    let token = match extract_bearer(&request) {
        Some(t) => t,
        None => return auth_error("Missing Authorization: Bearer <api-key>"),
    };

    let hash = hash_api_key(token);
    match state.store().validate_api_key(&hash).await {
        Ok(true) => next.run(request).await,
        Ok(false) => {
            warn!(
                "Invalid API key (prefix: {}...)",
                &token[..8.min(token.len())]
            );
            auth_error("Invalid API key")
        }
        Err(e) => {
            warn!("API key validation error: {e}");
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(serde_json::json!({"error": "auth check failed"})),
            )
                .into_response()
        }
    }
}

async fn validate_jwt(
    issuer: &str,
    audience: Option<&str>,
    jwks_cache: &JwksCache,
    request: Request,
    next: Next,
) -> Response {
    use jsonwebtoken::Validation;

    let token = match extract_bearer(&request) {
        Some(t) => t,
        None => return auth_error("Missing Authorization: Bearer <jwt>"),
    };

    // Decode header to get algorithm and kid
    let header = match jsonwebtoken::decode_header(token) {
        Ok(h) => h,
        Err(e) => {
            warn!("Invalid JWT header: {e}");
            return auth_error("Invalid JWT");
        }
    };

    // Find the decoding key from JWKS
    let decoding_key = match &header.kid {
        Some(kid) => match jwks_cache.find_key(kid).await {
            Ok(key) => key,
            Err(e) => {
                warn!("JWKS key lookup failed: {e}");
                return auth_error("JWT validation failed: key not found");
            }
        },
        None => match jwks_cache.find_any_key(header.alg).await {
            Ok(key) => key,
            Err(e) => {
                warn!("JWKS key lookup failed (no kid): {e}");
                return auth_error("JWT validation failed: no suitable key");
            }
        },
    };

    // Build validation rules
    let mut validation = Validation::new(header.alg);
    validation.set_issuer(&[issuer]);
    if let Some(aud) = audience {
        validation.set_audience(&[aud]);
    } else {
        validation.validate_aud = false;
    }

    // Validate signature + claims
    match jsonwebtoken::decode::<serde_json::Value>(token, &decoding_key, &validation) {
        Ok(_) => next.run(request).await,
        Err(e) => {
            warn!("JWT validation failed: {e}");
            auth_error(&format!("JWT validation failed: {e}"))
        }
    }
}

fn extract_bearer(request: &Request) -> Option<&str> {
    request
        .headers()
        .get("authorization")
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer "))
}

/// True iff the request is the bootstrap-window endpoint
/// (`POST /api/v1/api-keys`). The caller is still responsible for checking
/// that the `api_keys` table is empty before actually allowing unauth access.
fn is_bootstrap_request(request: &Request) -> bool {
    request.method() == axum::http::Method::POST
        && request.uri().path() == "/api/v1/api-keys"
}

fn auth_error(msg: &str) -> Response {
    (
        StatusCode::UNAUTHORIZED,
        Json(serde_json::json!({"error": msg})),
    )
        .into_response()
}

// ── API Key Helpers ─────────────────────────────────────────

/// Hash an API key with SHA-256 for storage/lookup.
pub fn hash_api_key(key: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(key.as_bytes());
    data_encoding::HEXLOWER.encode(&hasher.finalize())
}

/// Generate a new random API key (32 bytes, hex-encoded).
pub fn generate_api_key() -> String {
    use rand::Rng;
    let bytes: [u8; 32] = rand::rng().random();
    format!("assay_{}", data_encoding::HEXLOWER.encode(&bytes))
}

/// Extract the prefix (first 8 chars after "assay_") for display.
pub fn key_prefix(key: &str) -> String {
    let stripped = key.strip_prefix("assay_").unwrap_or(key);
    format!("assay_{}...", &stripped[..8.min(stripped.len())])
}