fraiseql_auth/

state_store.rs

//! CSRF state store — trait definition and backends.
//!
//! Stores OAuth `state` parameters for the duration of an authorization flow and
//! removes them on first retrieval, preventing state-replay attacks.

use std::sync::Arc;

use async_trait::async_trait;
use dashmap::DashMap;

use crate::error::Result;

/// StateStore trait - implement this for different storage backends
///
/// Stores OAuth state parameters with expiration for CSRF protection.
/// In distributed deployments, use a persistent backend (Redis) instead of in-memory.
///
/// # Examples
///
/// Use in-memory store for single-instance deployments:
/// ```rust
/// use std::sync::Arc;
/// use fraiseql_auth::state_store::InMemoryStateStore;
/// let state_store = Arc::new(InMemoryStateStore::new());
/// ```
///
/// Use Redis for distributed deployments:
/// ```no_run
/// // Requires: live Redis server.
/// use std::sync::Arc;
/// # async fn example() -> fraiseql_auth::error::Result<()> {
/// # #[cfg(feature = "redis-rate-limiting")] {
/// use fraiseql_auth::state_store::RedisStateStore;
/// let state_store = Arc::new(RedisStateStore::new("redis://localhost:6379").await?);
/// # }
/// # Ok(())
/// # }
/// ```
// Reason: used as dyn Trait (Arc<dyn StateStore>); async_trait ensures Send bounds and
// dyn-compatibility async_trait: dyn-dispatch required; remove when RTN + Send is stable (RFC 3425)
#[async_trait]
pub trait StateStore: Send + Sync {
    /// Store a state value with provider and expiration
    ///
    /// # Arguments
    /// * `state` - The state parameter value
    /// * `provider` - OAuth provider name
    /// * `expiry_secs` - Unix timestamp when this state expires
    async fn store(&self, state: String, provider: String, expiry_secs: u64) -> Result<()>;

    /// Retrieve and remove a state value
    ///
    /// Returns (provider, expiry_secs) if state exists and is valid
    /// Returns error if state doesn't exist or is invalid
    async fn retrieve(&self, state: &str) -> Result<(String, u64)>;
}

/// In-memory state store using DashMap
///
/// **Warning**: Only suitable for single-instance deployments!
/// For distributed systems, use RedisStateStore instead.
///
/// # SECURITY
/// - Bounded to MAX_STATES entries to prevent unbounded memory growth
/// - Expired states are automatically cleaned up on store operations
/// - Implements LRU-like eviction when max capacity is reached
#[derive(Debug)]
pub struct InMemoryStateStore {
    // Map of state -> (provider, expiry_secs)
    pub(crate) states: Arc<DashMap<String, (String, u64)>>,
    // Maximum number of states to store (prevents memory exhaustion)
    max_states:        usize,
}

impl InMemoryStateStore {
    /// Default maximum number of states to store (10,000 states)
    /// At ~100 bytes per state, this limits memory to ~1 MB
    const MAX_STATES: usize = 10_000;

    /// Create a new in-memory state store with default limits
    #[must_use]
    pub fn new() -> Self {
        Self {
            states:     Arc::new(DashMap::new()),
            max_states: Self::MAX_STATES,
        }
    }

    /// Create a new in-memory state store with custom max size
    ///
    /// # Arguments
    /// * `max_states` - Maximum number of states to store
    #[must_use]
    pub fn with_max_states(max_states: usize) -> Self {
        Self {
            states:     Arc::new(DashMap::new()),
            max_states: max_states.max(1), // Ensure at least 1 state
        }
    }

    /// Clean up expired states and check capacity
    ///
    /// # SECURITY
    /// Called before inserting new states to:
    /// 1. Remove expired states (automatic cleanup)
    /// 2. Check if store is at capacity
    /// 3. Return eviction needed flag if cleanup doesn't free space
    fn cleanup_expired(&self) -> bool {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        // Remove all expired states
        self.states.retain(|_key, (_provider, expiry)| *expiry > now);

        // Return true if we're still over capacity after cleanup
        self.states.len() >= self.max_states
    }
}

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

// Reason: StateStore is defined with #[async_trait]; all implementations must match
// its transformed method signatures to satisfy the trait contract
// async_trait: dyn-dispatch required; remove when RTN + Send is stable (RFC 3425)
#[async_trait]
impl StateStore for InMemoryStateStore {
    async fn store(&self, state: String, provider: String, expiry_secs: u64) -> Result<()> {
        // SECURITY: Clean up expired states before inserting new one
        if self.cleanup_expired() {
            // Still over capacity after cleanup - reject to prevent memory exhaustion
            return Err(crate::error::AuthError::ConfigError {
                message: "State store at capacity, cannot store new state".to_string(),
            });
        }

        self.states.insert(state, (provider, expiry_secs));
        Ok(())
    }

    async fn retrieve(&self, state: &str) -> Result<(String, u64)> {
        let (_key, value) =
            self.states.remove(state).ok_or_else(|| crate::error::AuthError::InvalidState)?;
        Ok(value)
    }
}

/// Redis-backed state store for distributed deployments
///
/// Uses Redis to store OAuth state parameters, allowing state validation
/// across multiple server instances. Automatically expires states after TTL.
#[cfg(feature = "redis-rate-limiting")]
#[derive(Clone)]
pub struct RedisStateStore {
    client: redis::aio::ConnectionManager,
}

#[cfg(feature = "redis-rate-limiting")]
impl RedisStateStore {
    /// Create a new Redis state store
    ///
    /// # Arguments
    /// * `redis_url` - Connection string (e.g., "redis://localhost:6379")
    ///
    /// # Example
    /// ```no_run
    /// // Requires: live Redis server.
    /// # async fn example() -> fraiseql_auth::error::Result<()> {
    /// use fraiseql_auth::state_store::RedisStateStore;
    /// let store = RedisStateStore::new("redis://localhost:6379").await?;
    /// # Ok(())
    /// # }
    /// ```
    /// # Errors
    ///
    /// Returns [`AuthError::ConfigError`](crate::error::AuthError::ConfigError) if the Redis URL is
    /// invalid or if the connection manager cannot be established.
    pub async fn new(redis_url: &str) -> Result<Self> {
        let client =
            redis::Client::open(redis_url).map_err(|e| crate::error::AuthError::ConfigError {
                message: e.to_string(),
            })?;

        let connection_manager = client.get_connection_manager().await.map_err(|e| {
            crate::error::AuthError::ConfigError {
                message: e.to_string(),
            }
        })?;

        Ok(Self {
            client: connection_manager,
        })
    }

    /// Get Redis key for state
    fn state_key(state: &str) -> String {
        format!("oauth:state:{}", state)
    }
}

#[cfg(feature = "redis-rate-limiting")]
// Reason: StateStore is defined with #[async_trait]; all implementations must match
// its transformed method signatures to satisfy the trait contract
// async_trait: dyn-dispatch required; remove when RTN + Send is stable (RFC 3425)
#[async_trait]
impl StateStore for RedisStateStore {
    async fn store(&self, state: String, provider: String, expiry_secs: u64) -> Result<()> {
        use redis::AsyncCommands;

        let key = Self::state_key(&state);
        let ttl = expiry_secs
            .saturating_sub(
                std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_secs(),
            )
            .max(1); // Minimum 1 second TTL

        let mut conn = self.client.clone();
        let _: () = conn.set_ex(&key, &provider, ttl).await.map_err(|e| {
            crate::error::AuthError::ConfigError {
                message: e.to_string(),
            }
        })?;

        Ok(())
    }

    async fn retrieve(&self, state: &str) -> Result<(String, u64)> {
        use redis::AsyncCommands;

        let key = Self::state_key(state);
        let mut conn = self.client.clone();

        // SECURITY: Use GETDEL (atomic get-and-delete, Redis ≥6.2) to prevent the
        // GET+DEL race condition where two concurrent requests could both read the
        // same state token before either deletes it, enabling replay attacks.
        let provider: Option<String> =
            conn.get_del(&key).await.map_err(|e| crate::error::AuthError::ConfigError {
                message: e.to_string(),
            })?;

        let provider = provider.ok_or(crate::error::AuthError::InvalidState)?;

        // Return current time as expiry (it was already validated by Redis TTL)
        let expiry_secs = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        Ok((provider, expiry_secs))
    }
}