guru-connect/server/src/api/auth_logout.rs

//! Logout and token revocation endpoints

use axum::{
    extract::{Request, State},
    http::{HeaderMap, StatusCode},
    Json,
};
use serde::Serialize;
use tracing::{info, warn};
use uuid::Uuid;

use crate::auth::AuthenticatedUser;
use crate::AppState;

use super::auth::ErrorResponse;

/// Extract JWT token from Authorization header
fn extract_token_from_headers(
    headers: &HeaderMap,
) -> Result<String, (StatusCode, Json<ErrorResponse>)> {
    let auth_header = headers
        .get("Authorization")
        .and_then(|v| v.to_str().ok())
        .ok_or_else(|| {
            (
                StatusCode::UNAUTHORIZED,
                Json(ErrorResponse {
                    error: "Missing Authorization header".to_string(),
                }),
            )
        })?;

    let token = auth_header.strip_prefix("Bearer ").ok_or_else(|| {
        (
            StatusCode::UNAUTHORIZED,
            Json(ErrorResponse {
                error: "Invalid Authorization format".to_string(),
            }),
        )
    })?;

    Ok(token.to_string())
}

/// Logout response
#[derive(Debug, Serialize)]
pub struct LogoutResponse {
    pub message: String,
}

/// POST /api/auth/logout - Revoke current token (logout)
///
/// Adds the user's current JWT token to the blacklist, effectively logging them out.
/// The token will no longer be valid for any requests.
pub async fn logout(
    State(state): State<AppState>,
    user: AuthenticatedUser,
    request: Request,
) -> Result<Json<LogoutResponse>, (StatusCode, Json<ErrorResponse>)> {
    // Extract token from headers
    let token = extract_token_from_headers(request.headers())?;

    // Add the login JWT to the blacklist.
    state.token_blacklist.revoke(&token).await;

    // Also revoke any outstanding session-scoped VIEWER tokens this user minted
    // (CRITICAL #2). The login-JWT blacklist alone leaves a viewer token minted
    // before logout valid for the rest of its 5-minute TTL, keeping a live
    // viewer/remote-control plane open after logout. `user.user_id` is the `sub`
    // the viewer tokens were registered under (the same claim stamped into them).
    // The viewer WS already blacklist-checks the exact token string, so adding
    // them here is sufficient — no WS change needed. take_for_user drains and
    // clears the registry entry; expired tokens are pruned (not returned).
    let viewer_tokens = state.viewer_tokens.take_for_user(&user.user_id);
    let revoked_viewer_count = viewer_tokens.len();
    for viewer_token in viewer_tokens {
        state.token_blacklist.revoke(&viewer_token).await;
    }

    info!(
        "User {} logged out (login token revoked, {} viewer token(s) revoked)",
        user.username, revoked_viewer_count
    );

    Ok(Json(LogoutResponse {
        message: "Logged out successfully".to_string(),
    }))
}

/// POST /api/auth/revoke-token - Revoke own token (same as logout)
///
/// Alias for logout endpoint for consistency with revocation terminology.
pub async fn revoke_own_token(
    State(state): State<AppState>,
    user: AuthenticatedUser,
    request: Request,
) -> Result<Json<LogoutResponse>, (StatusCode, Json<ErrorResponse>)> {
    logout(State(state), user, request).await
}

/// Revoke user request
#[derive(Debug, serde::Deserialize)]
pub struct RevokeUserRequest {
    pub user_id: Uuid,
}

/// POST /api/auth/admin/revoke-user - Admin endpoint to revoke all tokens for a user
///
/// WARNING: This currently only revokes the admin's own token as a demonstration.
/// Full implementation would require:
/// 1. Session tracking table to store active JWT tokens
/// 2. Query to find all tokens for the target user
/// 3. Add all found tokens to blacklist
///
/// For MVP, we're implementing the foundation but not the full user tracking.
pub async fn revoke_user_tokens(
    State(_state): State<AppState>,
    admin: AuthenticatedUser,
    Json(req): Json<RevokeUserRequest>,
) -> Result<Json<LogoutResponse>, (StatusCode, Json<ErrorResponse>)> {
    // Verify admin permission
    if !admin.is_admin() {
        return Err((
            StatusCode::FORBIDDEN,
            Json(ErrorResponse {
                error: "Admin access required".to_string(),
            }),
        ));
    }

    warn!(
        "Admin {} attempted to revoke tokens for user {} - NOT IMPLEMENTED (requires session tracking)",
        admin.username, req.user_id
    );

    // TODO: Implement session tracking
    // 1. Query active_sessions table for all tokens belonging to user_id
    // 2. Add each token to blacklist
    // 3. Delete session records from database

    Err((
        StatusCode::NOT_IMPLEMENTED,
        Json(ErrorResponse {
            error: "User token revocation not yet implemented - requires session tracking table"
                .to_string(),
        }),
    ))
}

/// GET /api/auth/blacklist/stats - Get blacklist statistics (admin only)
///
/// Returns information about the current token blacklist for monitoring.
#[derive(Debug, Serialize)]
pub struct BlacklistStatsResponse {
    pub revoked_tokens_count: usize,
}

pub async fn get_blacklist_stats(
    State(state): State<AppState>,
    admin: AuthenticatedUser,
) -> Result<Json<BlacklistStatsResponse>, (StatusCode, Json<ErrorResponse>)> {
    if !admin.is_admin() {
        return Err((
            StatusCode::FORBIDDEN,
            Json(ErrorResponse {
                error: "Admin access required".to_string(),
            }),
        ));
    }

    let count = state.token_blacklist.len().await;

    Ok(Json(BlacklistStatsResponse {
        revoked_tokens_count: count,
    }))
}

/// POST /api/auth/blacklist/cleanup - Clean up expired tokens from blacklist (admin only)
///
/// Removes expired tokens from the blacklist to prevent memory buildup.
#[derive(Debug, Serialize)]
pub struct CleanupResponse {
    pub removed_count: usize,
    pub remaining_count: usize,
}

pub async fn cleanup_blacklist(
    State(state): State<AppState>,
    admin: AuthenticatedUser,
) -> Result<Json<CleanupResponse>, (StatusCode, Json<ErrorResponse>)> {
    if !admin.is_admin() {
        return Err((
            StatusCode::FORBIDDEN,
            Json(ErrorResponse {
                error: "Admin access required".to_string(),
            }),
        ));
    }

    let removed = state
        .token_blacklist
        .cleanup_expired(&state.jwt_config)
        .await;
    let remaining = state.token_blacklist.len().await;

    info!(
        "Admin {} cleaned up blacklist: {} tokens removed, {} remaining",
        admin.username, removed, remaining
    );

    Ok(Json(CleanupResponse {
        removed_count: removed,
        remaining_count: remaining,
    }))
}

#[cfg(test)]
mod tests {
    use crate::auth::{JwtConfig, TokenBlacklist, ViewerAccess, ViewerTokenRegistry};
    use std::time::Duration;
    use uuid::Uuid;

    /// End-to-end (component-level) proof of CRITICAL #2: a viewer token minted
    /// under a user, then revoked at logout via the registry drain, is rejected
    /// by the SAME blacklist check the viewer WS runs (`is_revoked`).
    ///
    /// This exercises the real mint → register → logout-drain → blacklist path
    /// without the HTTP/DB plumbing of the full handler. Uses local component
    /// instances only (no process-global env), so it is parallel-safe.
    #[tokio::test]
    async fn logout_revokes_minted_viewer_token() {
        let jwt = JwtConfig::new("test-secret-at-least-32-chars-long!!".to_string(), 24);
        let registry = ViewerTokenRegistry::new();
        let blacklist = TokenBlacklist::new();

        let user_sub = Uuid::new_v4().to_string();
        let session_id = Uuid::new_v4();
        let tenant_id = Uuid::new_v4();

        // Mint a viewer token and register it under the user (as mint_viewer_token does).
        let viewer_token = jwt
            .create_viewer_token(&user_sub, session_id, tenant_id, ViewerAccess::Control)
            .unwrap();
        registry.register(&user_sub, &viewer_token, Duration::from_secs(300));

        // The viewer WS check (is_revoked) passes BEFORE logout — token is live.
        assert!(!blacklist.is_revoked(&viewer_token).await);

        // Logout drains the user's viewer tokens into the blacklist (handler logic).
        for tok in registry.take_for_user(&user_sub) {
            blacklist.revoke(&tok).await;
        }

        // After logout the same WS check now REJECTS the viewer token.
        assert!(blacklist.is_revoked(&viewer_token).await);
        // The token also remains a structurally valid viewer JWT (not expired),
        // proving revocation — not natural expiry — is what blocks it.
        assert!(jwt.validate_viewer_token(&viewer_token).is_ok());
    }

    /// A different user's logout must NOT revoke this user's viewer token.
    #[tokio::test]
    async fn logout_does_not_revoke_other_users_viewer_token() {
        let jwt = JwtConfig::new("test-secret-at-least-32-chars-long!!".to_string(), 24);
        let registry = ViewerTokenRegistry::new();
        let blacklist = TokenBlacklist::new();

        let user_a = Uuid::new_v4().to_string();
        let user_b = Uuid::new_v4().to_string();
        let session_id = Uuid::new_v4();
        let tenant_id = Uuid::new_v4();

        let token_b = jwt
            .create_viewer_token(&user_b, session_id, tenant_id, ViewerAccess::ViewOnly)
            .unwrap();
        registry.register(&user_b, &token_b, Duration::from_secs(300));

        // user_a logs out — drains only user_a's (empty) token set.
        for tok in registry.take_for_user(&user_a) {
            blacklist.revoke(&tok).await;
        }

        // user_b's token is untouched.
        assert!(!blacklist.is_revoked(&token_b).await);
    }
}