eric / rustwarden

/// Authentication Required Middleware
/// Feature 004: UI Post-Auth
///
/// Middleware to protect routes that require authentication.
/// Redirects to login page if user is not authenticated.
use axum::{
    extract::{Request, State},
    http::Uri,
    middleware::Next,
    response::{IntoResponse, Response},
};
use sea_orm::DatabaseConnection;

use std::sync::Arc;

use crate::auth::jwt::{Claims, JwtService};
use crate::auth::jwt_extractor;
use crate::config::auth::AuthConfig;
use crate::error::AppError;
use crate::services::jwt_blacklist_service::JwtBlacklistService;

/// Extract user claims from JWT token in request
///
/// Helper function to extract user claims from JWT token (cookie or header).
/// Returns None if not authenticated or error.
///
/// # Arguments
///
/// * `db` - Database connection (for blacklist check)
/// * `request` - The HTTP request
///
/// # Returns
///
/// * `Some(Claims)` if authenticated and token valid
/// * `None` if not authenticated or error
pub async fn extract_user_from_jwt(db: &DatabaseConnection, request: &Request) -> Option<Claims> {
    extract_user_from_jwt_impl(db, request).await
}

/// Internal implementation of extract_user_from_jwt
async fn extract_user_from_jwt_impl(db: &DatabaseConnection, request: &Request) -> Option<Claims> {
    // Load auth configuration from request extensions (injected at startup)
    let auth_config = match request.extensions().get::<Arc<AuthConfig>>() {
        Some(config) => config.clone(),
        None => {
            tracing::warn!("AuthConfig extension missing, falling back to from_env");
            match AuthConfig::from_env() {
                Ok(config) => Arc::new(config),
                Err(e) => {
                    tracing::warn!("Failed to load auth config: {}", e);
                    return None;
                }
            }
        }
    };

    // Extract JWT from request
    let jwt_token = match jwt_extractor::extract_jwt_from_request(request) {
        Ok(token) => token,
        Err(_) => return None,
    };

    // Check if token is blacklisted
    if JwtBlacklistService::is_blacklisted(db, &jwt_token)
        .await
        .is_ok_and(|b| b)
    {
        tracing::debug!("JWT token is blacklisted");
        return None;
    }

    // Validate JWT token
    match JwtService::validate_token(&auth_config, &jwt_token) {
        Ok(claims) => {
            tracing::debug!("User extracted from JWT: {}", claims.sub);
            Some(claims)
        }
        Err(e) => {
            tracing::debug!("Failed to validate JWT: {}", e);
            None
        }
    }
}

/// Middleware that requires authentication
///
/// Checks if the user has a valid JWT token (from cookie or Authorization header).
/// If authenticated, allows the request to proceed.
/// If not authenticated, redirects to `/login` with a redirect parameter.
///
/// # Arguments
///
/// * `State(db)` - Database connection (for blacklist check)
/// * `uri` - Original request URI (for redirect after login)
/// * `request` - The HTTP request
/// * `next` - Next middleware or handler
///
/// # Returns
///
/// * `200 OK` if authenticated (passes to handler)
/// * `302 Found` redirects to `/login` if not authenticated
///
/// # Example
///
/// ```rust,ignore
/// use axum::Router;
/// use axum::routing::get;
/// use axum::middleware;
///
/// let protected_routes = Router::new()
///     .route("/home", get(home_handler))
///     .layer(middleware::from_fn(auth_required));
/// ```
/// Internal implementation that takes db as parameter
/// Made public for use in main.rs with from_fn closure
pub async fn auth_required_impl(
    db: &DatabaseConnection,
    uri: Uri,
    request: Request,
    next: Next,
) -> Result<Response, AppError> {
    tracing::debug!(
        "auth_required middleware: checking authentication for {}",
        uri
    );

    // Load auth configuration from request extensions (injected at startup)
    let auth_config = request
        .extensions()
        .get::<Arc<AuthConfig>>()
        .cloned()
        .ok_or_else(|| {
            tracing::error!("AuthConfig extension missing from request");
            AppError::Config("Auth configuration not available".to_string())
        })?;

    // Extract JWT from request (cookie or header)
    let jwt_token = match jwt_extractor::extract_jwt_from_request(&request) {
        Ok(token) => token,
        Err(_) => {
            tracing::warn!("Unauthenticated access attempt to: {} (no token)", uri);
            return Ok(crate::middleware::auth_service::auth_failure_response(
                &auth_config,
                &uri,
            ));
        }
    };

    // Check if token is blacklisted
    if JwtBlacklistService::is_blacklisted(db, &jwt_token).await? {
        tracing::warn!("Blacklisted JWT token used for: {}", uri);
        return Ok(crate::middleware::auth_service::auth_failure_response(
            &auth_config,
            &uri,
        ));
    }

    // Validate JWT token
    let claims = match JwtService::validate_token(&auth_config, &jwt_token) {
        Ok(claims) => claims,
        Err(e) => {
            tracing::warn!("Invalid JWT token for {}: {}", uri, e);
            return Ok(crate::middleware::auth_service::auth_failure_response(
                &auth_config,
                &uri,
            ));
        }
    };

    tracing::debug!("User authenticated: {} (from JWT)", claims.sub);

    // Store claims in request extensions for handlers to use
    let (mut parts, body) = request.into_parts();
    parts.extensions.insert(claims.clone());
    let request = Request::from_parts(parts, body);

    // User is authenticated, proceed to handler
    Ok(next.run(request).await)
}

/// Middleware that requires authentication (for use with from_fn_with_state)
///
/// Checks if the user has a valid JWT token (from cookie or Authorization header).
/// If authenticated, allows the request to proceed.
/// If not authenticated, redirects to `/login` with a redirect parameter.
///
/// # Note
///
/// This function signature is compatible with `axum::middleware::from_fn_with_state`.
/// The URI is extracted from `request.uri()` internally.
/// Returns `Response` directly (not `Result`) as required by `from_fn_with_state`.
/// Errors are converted to responses using `IntoResponse`.
pub async fn auth_required(
    State(db): State<DatabaseConnection>,
    request: Request,
    next: Next,
) -> Response {
    // Extract URI from request (as per Axum middleware best practices)
    let uri = request.uri().clone();

    // Call the implementation and convert Result to Response
    match auth_required_impl(&db, uri, request, next).await {
        Ok(response) => response,
        Err(e) => {
            tracing::error!("Auth middleware error: {}", e);
            e.into_response()
        }
    }
}

/// Optional authentication middleware
///
/// Similar to `auth_required` but doesn't redirect if not authenticated.
/// Useful for routes that are accessible to both authenticated and
/// non-authenticated users.
///
/// # Returns
///
/// Always allows the request to proceed, whether authenticated or not.
pub async fn optional_auth(
    State(db): State<DatabaseConnection>,
    request: Request,
    next: Next,
) -> Response {
    tracing::debug!("optional_auth middleware: checking authentication");

    // Try to extract and validate JWT (but don't fail if missing)
    if let Ok(jwt_token) = jwt_extractor::extract_jwt_from_request(&request) {
        if let Some(auth_config) = request.extensions().get::<Arc<AuthConfig>>().cloned() {
            if JwtBlacklistService::is_blacklisted(&db, &jwt_token)
                .await
                .is_ok_and(|b| !b)
            {
                if let Ok(claims) = JwtService::validate_token(&auth_config, &jwt_token) {
                    tracing::debug!("User authenticated: {} (from JWT)", claims.sub);
                }
            }
        }
    }

    // Always proceed to handler
    next.run(request).await
}