rocketmq_auth/authorization/strategy/

stateful_authorization_strategy.rs

// Copyright 2023 The RocketMQ Rust Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Stateful authorization strategy with caching.
//!
//! This strategy caches authorization decisions for connection-based requests,
//! providing better performance for repeated authorization checks on the same
//! channel/connection.

use std::any::Any;
use std::collections::HashMap;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::sync::RwLock;
use std::sync::RwLockWriteGuard;
use std::time::Duration;
use std::time::Instant;

use rocketmq_security_api::AuthorizationDecision;
#[cfg(test)]
use rocketmq_security_api::AuthorizationDenial;
use tracing::debug;

use crate::authorization::context::default_authorization_context::DefaultAuthorizationContext;
use crate::authorization::strategy::abstract_authorization_strategy::AbstractAuthorizationStrategy;
use crate::authorization::strategy::abstract_authorization_strategy::AuthorizationStrategy;
use crate::authorization::strategy::evaluate_base_authorization;
use crate::authorization::strategy::AuthorizationFuture;
use crate::config::AuthConfig;
use crate::AuthMetrics;
#[cfg(test)]
use crate::AuthServiceError;
use crate::AuthServiceResult;

type AuthCache = HashMap<String, CachedAuthResult>;

/// Cached authorization result.
#[derive(Clone, Debug)]
struct CachedAuthResult {
    /// The final policy decision. Operational failures are never cached.
    decision: AuthorizationDecision,
    /// Timestamp when this entry was cached
    cached_at: Instant,
}

impl CachedAuthResult {
    fn new(decision: AuthorizationDecision) -> Self {
        Self {
            decision,
            cached_at: Instant::now(),
        }
    }

    fn is_expired(&self, ttl: Duration) -> bool {
        self.cached_at.elapsed() > ttl
    }
}

fn cacheable_decision(
    result: &AuthServiceResult<AuthorizationDecision>,
    cache_negative_result: bool,
) -> Option<AuthorizationDecision> {
    match result {
        Ok(AuthorizationDecision::Allow) => Some(AuthorizationDecision::Allow),
        Ok(decision @ AuthorizationDecision::Deny(_)) if cache_negative_result => Some(*decision),
        Ok(AuthorizationDecision::Deny(_)) | Err(_) => None,
    }
}

/// Stateful authorization strategy with caching.
///
/// This strategy caches authorization decisions based on channel ID, reducing
/// the overhead of repeated authorization checks for the same connection.
///
/// # Cache Key
///
/// The cache key is built from:
/// - Channel ID
/// - Subject key
/// - Resource key
/// - Actions (comma-separated)
/// - Source IP
///
/// Format: `{channel_id}#{subject_key}#{resource_key}#{actions}#{source_ip}`
///
/// # Cache Behavior
///
/// - **Cache Hit**: Returns cached result immediately
/// - **Cache Miss**: Evaluates authorization and caches the result
/// - **TTL**: Cached entries expire after configured duration
/// - **No Channel ID**: Bypasses cache and evaluates directly
///
/// # Use Cases
///
/// - Connection-based protocols (e.g., long-lived TCP connections)
/// - High-throughput scenarios with repeated authorization checks
/// - Systems where authorization policies are relatively stable
///
/// # Trade-offs
///
/// **Advantages:**
/// - Significantly reduced latency for cached requests
/// - Lower load on authorization provider
/// - Better performance for connection-based protocols
///
/// **Disadvantages:**
/// - Potential delay in policy updates (up to TTL)
/// - Higher memory footprint
/// - Cache consistency considerations
///
/// # Examples
///
/// ```rust,ignore
/// use rocketmq_auth::StatefulAuthorizationStrategy;
/// use rocketmq_auth::AuthConfig;
///
/// let config = AuthConfig::default();
/// let strategy = StatefulAuthorizationStrategy::new(config, None)?;
///
/// // First evaluation: queries provider
/// let context = DefaultAuthorizationContext::of(subject, resource, action, ip);
/// context.set_channel_id("channel-123");
/// match strategy.evaluate(&context).await? { // Cache miss
///     AuthorizationDecision::Allow => proceed(),
///     AuthorizationDecision::Deny(reason) => reject(reason),
/// }
///
/// // Second evaluation: uses cache
/// match strategy.evaluate(&context).await? { // Cache hit
///     AuthorizationDecision::Allow => proceed(),
///     AuthorizationDecision::Deny(reason) => reject(reason),
/// }
/// ```
pub struct StatefulAuthorizationStrategy {
    /// Base authorization strategy
    base: AbstractAuthorizationStrategy,

    /// Authorization cache (key -> result)
    auth_cache: Arc<RwLock<AuthCache>>,

    /// Cache TTL
    cache_ttl: Duration,

    /// Maximum cache size
    cache_max_size: usize,

    /// Request counter for periodic cleanup
    request_counter: AtomicUsize,

    /// Shared ACL snapshot generation. A successful ACL reload advances this value.
    acl_generation: Arc<AtomicU64>,

    /// Last generation observed by this strategy, used to drop old entries promptly.
    last_seen_acl_generation: AtomicU64,
    cache_negative_result: bool,
    metrics: AuthMetrics,
}

impl StatefulAuthorizationStrategy {
    /// Creates a new stateful authorization strategy with caching.
    ///
    /// # Arguments
    ///
    /// * `auth_config` - Authorization configuration
    /// * `metadata_service` - Optional metadata service for provider initialization
    ///
    /// # Returns
    ///
    /// A new `StatefulAuthorizationStrategy` instance.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let config = AuthConfig::default();
    /// let strategy = StatefulAuthorizationStrategy::new(config, None)?;
    /// ```
    pub fn new(
        auth_config: AuthConfig,
        metadata_service: Option<Box<dyn Any + Send + Sync>>,
    ) -> AuthServiceResult<Self> {
        let acl_generation = metadata_service
            .as_ref()
            .and_then(|service| service.downcast_ref::<crate::ProviderRegistry>())
            .map(crate::ProviderRegistry::acl_generation_counter)
            .unwrap_or_default();
        Self::new_with_acl_generation(auth_config, metadata_service, acl_generation)
    }

    /// Creates a new stateful authorization strategy bound to a shared ACL generation counter.
    pub fn new_with_acl_generation(
        auth_config: AuthConfig,
        metadata_service: Option<Box<dyn Any + Send + Sync>>,
        acl_generation: Arc<AtomicU64>,
    ) -> AuthServiceResult<Self> {
        Self::new_with_acl_generation_and_metrics(auth_config, metadata_service, acl_generation, AuthMetrics::default())
    }

    pub fn new_with_acl_generation_and_metrics(
        auth_config: AuthConfig,
        metadata_service: Option<Box<dyn Any + Send + Sync>>,
        acl_generation: Arc<AtomicU64>,
        metrics: AuthMetrics,
    ) -> AuthServiceResult<Self> {
        let cache_ttl = Duration::from_secs(auth_config.stateful_authorization_cache_expired_second as u64);
        let cache_max_size = auth_config.stateful_authorization_cache_max_num as usize;
        let cache_negative_result = auth_config.stateful_authorization_cache_negative_enable;

        let base = AbstractAuthorizationStrategy::new(auth_config, metadata_service)?;
        let initial_generation = acl_generation.load(Ordering::Acquire);

        debug!(
            "StatefulAuthorizationStrategy initialized with TTL: {:?}, max size: {}",
            cache_ttl, cache_max_size
        );

        Ok(Self {
            base,
            auth_cache: Arc::new(RwLock::new(HashMap::new())),
            cache_ttl,
            cache_max_size,
            request_counter: AtomicUsize::new(0),
            acl_generation,
            last_seen_acl_generation: AtomicU64::new(initial_generation),
            cache_negative_result,
            metrics,
        })
    }

    /// Builds a cache key from the authorization context.
    ///
    /// Format: `{generation}#{channel_id}#{subject_key}#{resource_key}#{actions}#{source_ip}`
    fn build_key(&self, context: &DefaultAuthorizationContext) -> String {
        let generation = self.refresh_acl_generation();
        let channel_id = context.channel_id().unwrap_or("");
        let subject_key = context.subject().map(|s| s.subject_key()).unwrap_or_default();
        let resource_key = context.resource().map(|r| format!("{:?}", r)).unwrap_or_default();
        let actions = context
            .actions()
            .iter()
            .map(|a| format!("{:?}", a))
            .collect::<Vec<_>>()
            .join(",");
        let source_ip = context.source_ip().unwrap_or("");

        format!(
            "{}#{}#{}#{}#{}#{}",
            generation, channel_id, subject_key, resource_key, actions, source_ip
        )
    }

    fn refresh_acl_generation(&self) -> u64 {
        let current = self.acl_generation.load(Ordering::Acquire);
        let previous = self.last_seen_acl_generation.load(Ordering::Acquire);
        if previous != current
            && self
                .last_seen_acl_generation
                .compare_exchange(previous, current, Ordering::AcqRel, Ordering::Acquire)
                .is_ok()
        {
            self.clear_cache();
            self.metrics.record_cache_invalidation();
        }
        current
    }

    fn recover_poisoned_cache(&self) {
        let mut cache = match self.auth_cache.write() {
            Ok(cache) => cache,
            Err(poisoned) => poisoned.into_inner(),
        };
        cache.clear();
        self.auth_cache.clear_poison();
        debug!("Authorization cache lock poisoned; cache cleared");
    }

    fn write_cache_recovering(&self) -> RwLockWriteGuard<'_, AuthCache> {
        match self.auth_cache.write() {
            Ok(cache) => cache,
            Err(poisoned) => {
                let mut cache = poisoned.into_inner();
                cache.clear();
                self.auth_cache.clear_poison();
                debug!("Authorization cache write lock poisoned; cache cleared before reuse");
                cache
            }
        }
    }

    /// Evicts expired entries from the cache.
    fn evict_expired(&self) {
        let mut cache = self.write_cache_recovering();
        let before_size = cache.len();

        cache.retain(|_, result| !result.is_expired(self.cache_ttl));

        let evicted = before_size - cache.len();
        if evicted > 0 {
            debug!("Evicted {} expired cache entries", evicted);
        }
    }

    /// Evicts oldest entries if cache is over capacity.
    fn evict_if_full(&self) {
        let mut cache = self.write_cache_recovering();

        if cache.len() >= self.cache_max_size {
            // Simple eviction: remove first 10% of entries
            let to_remove = (self.cache_max_size / 10).max(1);
            let keys_to_remove: Vec<String> = cache.keys().take(to_remove).cloned().collect();

            for key in keys_to_remove {
                cache.remove(&key);
            }

            debug!("Evicted {} entries to maintain cache size", to_remove);
        }
    }

    /// Gets a reference to the base strategy.
    pub fn base(&self) -> &AbstractAuthorizationStrategy {
        &self.base
    }

    /// Gets the current cache size.
    pub fn cache_size(&self) -> usize {
        match self.auth_cache.read() {
            Ok(cache) => cache.len(),
            Err(poisoned) => {
                drop(poisoned);
                self.recover_poisoned_cache();
                0
            }
        }
    }

    /// Clears the entire cache.
    pub fn clear_cache(&self) {
        let mut cache = self.write_cache_recovering();
        cache.clear();
        debug!("Authorization cache cleared");
    }
}

impl AuthorizationStrategy for StatefulAuthorizationStrategy {
    /// Evaluates authorization with caching support.
    ///
    /// If the context has a channel ID, the result is cached. Subsequent
    /// requests with the same cache key will use the cached result until
    /// it expires.
    ///
    /// # Arguments
    ///
    /// * `context` - The authorization context to evaluate
    ///
    /// # Returns
    ///
    /// * `Ok(AuthorizationDecision)` for a final allow or deny
    /// * `Err(AuthServiceError)` if evaluation cannot make a decision
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let strategy = StatefulAuthorizationStrategy::new(config, None)?;
    /// let mut context = DefaultAuthorizationContext::of(subject, resource, action, ip);
    /// context.set_channel_id("channel-123");
    /// match strategy.evaluate(&context).await? {
    ///     AuthorizationDecision::Allow => proceed(),
    ///     AuthorizationDecision::Deny(reason) => reject(reason),
    /// }
    /// ```
    fn evaluate<'a>(&'a self, context: &'a DefaultAuthorizationContext) -> AuthorizationFuture<'a> {
        Box::pin(async move {
            let _operation = self.base.enter_operation()?;
            // If no channel ID, bypass cache
            if context.channel_id().is_none_or(str::is_empty) {
                debug!("No channel ID, bypassing cache");
                return evaluate_base_authorization(&self.base, context).await;
            }

            let cache_key = self.build_key(context);

            // Check cache
            {
                match self.auth_cache.read() {
                    Ok(cache) => {
                        if let Some(cached_result) = cache.get(&cache_key) {
                            if !cached_result.is_expired(self.cache_ttl) {
                                debug!("Authorization cache hit");
                                self.metrics.record_cache_hit();
                                return Ok(cached_result.decision);
                            } else {
                                debug!("Authorization cache entry expired");
                            }
                        }
                    }
                    Err(poisoned) => {
                        drop(poisoned);
                        self.recover_poisoned_cache();
                        debug!("Authorization cache unavailable after poison recovery");
                    }
                }
            }

            // Cache miss - evaluate and cache result
            debug!("Authorization cache miss");
            self.metrics.record_cache_miss();

            let result = evaluate_base_authorization(&self.base, context).await;

            // Cache granted results by default. Denied results are cached only when
            // explicitly enabled because ACL updates must take effect conservatively.
            if let Some(decision) = cacheable_decision(&result, self.cache_negative_result) {
                self.evict_if_full();

                let mut cache = self.write_cache_recovering();
                cache.insert(cache_key.clone(), CachedAuthResult::new(decision));
                debug!("Authorization decision cached");
            }

            // Periodic cleanup every 100 requests
            let counter = self.request_counter.fetch_add(1, Ordering::Relaxed);
            if counter.is_multiple_of(100) {
                self.evict_expired();
            }

            result
        })
    }
}

#[cfg(test)]
mod tests {
    use cheetah_string::CheetahString;
    use rocketmq_security_api::Action;

    use super::*;
    use crate::authentication::enums::subject_type::SubjectType;
    use crate::authorization::model::resource::Resource;

    fn create_test_config() -> AuthConfig {
        AuthConfig {
            config_name: CheetahString::from("test"),
            cluster_name: CheetahString::from("test-cluster"),
            auth_config_path: CheetahString::from("/tmp/auth"),
            acl_file: CheetahString::new(),
            acl_file_watch_enabled: false,
            acl_file_watch_interval_millis: 5_000,
            authentication_enabled: false,
            authentication_provider: CheetahString::new(),
            authentication_metadata_provider: CheetahString::new(),
            authentication_strategy: CheetahString::new(),
            authentication_whitelist: CheetahString::new(),
            init_authentication_user: CheetahString::new(),
            inner_client_authentication_credentials: CheetahString::new(),
            signature_algorithm: Default::default(),
            request_timestamp_expired_millis: 0,
            authorization_enabled: true,
            authorization_provider: CheetahString::new(),
            authorization_metadata_provider: CheetahString::new(),
            authorization_strategy: CheetahString::new(),
            authorization_whitelist: CheetahString::new(),
            migrate_auth_from_v1_enabled: false,
            user_cache_max_num: 1000,
            user_cache_expired_second: 300,
            user_cache_refresh_second: 60,
            acl_cache_max_num: 1000,
            acl_cache_expired_second: 300,
            acl_cache_refresh_second: 60,
            stateful_authentication_cache_max_num: 1000,
            stateful_authentication_cache_expired_second: 300,
            stateful_authorization_cache_max_num: 100,
            stateful_authorization_cache_expired_second: 60,
            stateful_authorization_cache_negative_enable: false,
            ..AuthConfig::default()
        }
    }

    #[test]
    fn test_stateful_strategy_creation() {
        let config = create_test_config();
        let strategy = StatefulAuthorizationStrategy::new(config, None);
        assert!(strategy.is_ok());
    }

    #[test]
    fn test_cache_key_building() {
        let config = create_test_config();
        let strategy = StatefulAuthorizationStrategy::new(config, None).unwrap();

        let mut context = DefaultAuthorizationContext::default();
        context.set_channel_id("ch-123".to_string());
        context.set_source_ip("192.168.0.1".to_string());

        let key = strategy.build_key(&context);
        assert!(key.starts_with("0#"));
        assert!(key.contains("ch-123"));
        assert!(key.contains("192.168.0.1"));
    }

    #[test]
    fn test_cache_key_includes_acl_generation() {
        let config = create_test_config();
        let acl_generation = Arc::new(AtomicU64::new(11));
        let strategy =
            StatefulAuthorizationStrategy::new_with_acl_generation(config, None, acl_generation.clone()).unwrap();

        let mut context = DefaultAuthorizationContext::default();
        context.set_channel_id("ch-123".to_string());
        context.set_source_ip("192.168.0.1".to_string());

        let key = strategy.build_key(&context);
        assert!(key.starts_with("11#ch-123#"));

        acl_generation.store(12, Ordering::Release);
        let key = strategy.build_key(&context);
        assert!(key.starts_with("12#ch-123#"));
    }

    #[tokio::test]
    async fn test_no_channel_id_bypasses_cache() {
        let mut config = create_test_config();
        config.authorization_enabled = false;

        let strategy = StatefulAuthorizationStrategy::new(config, None).unwrap();
        let context = DefaultAuthorizationContext::default();

        let result = strategy.evaluate(&context).await;
        assert!(result.is_ok());
        assert_eq!(strategy.cache_size(), 0); // Nothing cached
    }

    #[test]
    fn test_clear_cache() {
        let config = create_test_config();
        let strategy = StatefulAuthorizationStrategy::new(config, None).unwrap();

        strategy.clear_cache();
        assert_eq!(strategy.cache_size(), 0);
    }

    #[tokio::test]
    async fn test_acl_generation_change_invalidates_cached_authorization_result() {
        let mut config = create_test_config();
        config.authorization_enabled = false;
        let acl_generation = Arc::new(AtomicU64::new(0));
        let strategy =
            StatefulAuthorizationStrategy::new_with_acl_generation(config, None, acl_generation.clone()).unwrap();
        let mut context = DefaultAuthorizationContext::default();
        context.set_channel_id("ch-123".to_string());
        context.set_source_ip("192.168.0.1".to_string());

        assert_eq!(strategy.evaluate(&context).await.unwrap(), AuthorizationDecision::Allow);
        assert_eq!(strategy.cache_size(), 1);

        acl_generation.fetch_add(1, Ordering::AcqRel);
        assert_eq!(strategy.evaluate(&context).await.unwrap(), AuthorizationDecision::Allow);
        assert_eq!(strategy.cache_size(), 1);

        let cache = strategy.auth_cache.read().unwrap();
        assert!(cache.keys().all(|key| key.starts_with("1#")));
    }

    #[tokio::test]
    async fn poisoned_authorization_cache_is_cleared_and_evaluation_continues() {
        let mut config = create_test_config();
        config.authorization_enabled = false;
        let strategy = StatefulAuthorizationStrategy::new(config, None).unwrap();
        let cache = strategy.auth_cache.clone();

        let poison_result = std::thread::spawn(move || {
            let _guard = cache.write().unwrap();
            panic!("poison authorization cache for test");
        })
        .join();
        assert!(poison_result.is_err());

        let mut context = DefaultAuthorizationContext::default();
        context.set_channel_id("ch-poison".to_string());

        assert_eq!(strategy.evaluate(&context).await.unwrap(), AuthorizationDecision::Allow);
        assert_eq!(strategy.cache_size(), 1);
        assert!(strategy.auth_cache.read().is_ok());
    }

    #[tokio::test(flavor = "current_thread")]
    async fn evaluate_cache_miss_inside_current_thread_runtime_without_nested_block_on_panic() {
        let mut config = create_test_config();
        config.authorization_enabled = false;
        let strategy = StatefulAuthorizationStrategy::new(config, None).unwrap();
        let mut context = DefaultAuthorizationContext::default();
        context.set_channel_id("ch-current-thread".to_string());
        assert_eq!(strategy.evaluate(&context).await.unwrap(), AuthorizationDecision::Allow);

        assert_eq!(strategy.cache_size(), 1);
    }

    #[tokio::test]
    async fn denied_authorization_is_not_cached_by_default() {
        let config = create_test_config();
        let strategy = StatefulAuthorizationStrategy::new(config, None).unwrap();
        let mut context = DefaultAuthorizationContext::of(
            "alice",
            SubjectType::User,
            Resource::of_topic("TopicA"),
            Action::Pub,
            "127.0.0.1",
        );
        context.set_channel_id("ch-denied-default".to_string());

        assert_eq!(
            strategy.evaluate(&context).await.unwrap(),
            AuthorizationDecision::Deny(AuthorizationDenial::SubjectUnknown)
        );
        assert_eq!(strategy.cache_size(), 0);
    }

    #[tokio::test]
    async fn denied_authorization_can_be_cached_when_enabled() {
        let mut config = create_test_config();
        config.stateful_authorization_cache_negative_enable = true;
        let strategy = StatefulAuthorizationStrategy::new(config, None).unwrap();
        let mut context = DefaultAuthorizationContext::of(
            "alice",
            SubjectType::User,
            Resource::of_topic("TopicA"),
            Action::Pub,
            "127.0.0.1",
        );
        context.set_channel_id("ch-denied-cached".to_string());

        assert_eq!(
            strategy.evaluate(&context).await.unwrap(),
            AuthorizationDecision::Deny(AuthorizationDenial::SubjectUnknown)
        );
        assert_eq!(strategy.cache_size(), 1);
    }

    #[test]
    fn operational_errors_are_never_cacheable() {
        let result = Err(AuthServiceError::invalid_context("broken context".to_owned()));

        assert_eq!(cacheable_decision(&result, false), None);
        assert_eq!(cacheable_decision(&result, true), None);
    }
}