guru-connect/server/src/auth/token_blacklist.rs

//! Token blacklist for JWT revocation
//!
//! Provides in-memory token blacklist for immediate revocation of JWTs.
//! Tokens are automatically cleaned up after expiration.

use std::collections::HashSet;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info};

/// Token blacklist for revocation
///
/// Maintains a set of revoked token signatures. When a token is revoked
/// (e.g., on logout or admin action), it's added to this blacklist and
/// all subsequent validation attempts will fail.
#[derive(Clone)]
pub struct TokenBlacklist {
    /// Set of revoked token strings
    tokens: Arc<RwLock<HashSet<String>>>,
}

impl TokenBlacklist {
    /// Create a new empty blacklist
    pub fn new() -> Self {
        Self {
            tokens: Arc::new(RwLock::new(HashSet::new())),
        }
    }

    /// Add a token to the blacklist (revoke it)
    ///
    /// # Arguments
    /// * `token` - The full JWT token string to revoke
    ///
    /// # Example
    /// ```rust
    /// blacklist.revoke("eyJ...").await;
    /// ```
    pub async fn revoke(&self, token: &str) {
        let mut tokens = self.tokens.write().await;
        let was_new = tokens.insert(token.to_string());

        if was_new {
            debug!(
                "Token revoked and added to blacklist (length: {})",
                token.len()
            );
        }
    }

    /// Check if a token has been revoked
    ///
    /// # Arguments
    /// * `token` - The JWT token string to check
    ///
    /// # Returns
    /// `true` if the token is in the blacklist (revoked), `false` otherwise
    pub async fn is_revoked(&self, token: &str) -> bool {
        let tokens = self.tokens.read().await;
        tokens.contains(token)
    }

    /// Get the number of tokens currently in the blacklist
    pub async fn len(&self) -> usize {
        let tokens = self.tokens.read().await;
        tokens.len()
    }

    /// Check if the blacklist is empty
    pub async fn is_empty(&self) -> bool {
        let tokens = self.tokens.read().await;
        tokens.is_empty()
    }

    /// Remove expired tokens from blacklist (cleanup)
    ///
    /// This should be called periodically to prevent memory buildup.
    /// Tokens that can no longer be validated (expired) are removed.
    ///
    /// # Arguments
    /// * `jwt_config` - JWT configuration for validating token expiration
    ///
    /// # Returns
    /// Number of tokens removed from blacklist
    pub async fn cleanup_expired(&self, jwt_config: &super::JwtConfig) -> usize {
        let mut tokens = self.tokens.write().await;
        let original_len = tokens.len();

        // Remove tokens that fail validation (expired)
        tokens.retain(|token| {
            // If token is expired (validation fails), remove it from blacklist
            jwt_config.validate_token(token).is_ok()
        });

        let removed = original_len - tokens.len();

        if removed > 0 {
            info!(
                "Cleaned {} expired tokens from blacklist ({} remaining)",
                removed,
                tokens.len()
            );
        }

        removed
    }

    /// Clear all tokens from the blacklist
    ///
    /// WARNING: This removes all revoked tokens. Use with caution.
    pub async fn clear(&self) {
        let mut tokens = self.tokens.write().await;
        let count = tokens.len();
        tokens.clear();
        info!("Cleared {} tokens from blacklist", count);
    }
}

impl Default for TokenBlacklist {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_revoke_and_check() {
        let blacklist = TokenBlacklist::new();
        let token = "test.token.here";

        assert!(!blacklist.is_revoked(token).await);

        blacklist.revoke(token).await;

        assert!(blacklist.is_revoked(token).await);
        assert_eq!(blacklist.len().await, 1);
    }

    #[tokio::test]
    async fn test_multiple_revocations() {
        let blacklist = TokenBlacklist::new();

        blacklist.revoke("token1").await;
        blacklist.revoke("token2").await;
        blacklist.revoke("token3").await;

        assert_eq!(blacklist.len().await, 3);
        assert!(blacklist.is_revoked("token1").await);
        assert!(blacklist.is_revoked("token2").await);
        assert!(blacklist.is_revoked("token3").await);
        assert!(!blacklist.is_revoked("token4").await);
    }

    #[tokio::test]
    async fn test_clear() {
        let blacklist = TokenBlacklist::new();

        blacklist.revoke("token1").await;
        blacklist.revoke("token2").await;

        assert_eq!(blacklist.len().await, 2);

        blacklist.clear().await;

        assert_eq!(blacklist.len().await, 0);
        assert!(blacklist.is_empty().await);
    }
}