use crate::config::{ConfigChangeRecord, ConfigHistory, FlowControlConfig};
use crate::decision_chain::DecisionChain;
use crate::error::Decision;
use crate::error::LimiteronError;
#[cfg(feature = "fallback")]
use crate::error::RejectionMetadata;
#[cfg(feature = "fallback")]
use crate::fallback::FallbackManager;
use crate::i18n::t;
#[cfg(feature = "fallback")]
use crate::l1_cache::IslandFallbackStrategy;
#[cfg(feature = "fallback")]
use crate::l1_cache::IslandModeConfig;
use crate::l1_cache::{CacheableDecision, L1Cache, L1CacheConfig, RateLimitCacheKey};
#[cfg(any(feature = "prometheus", feature = "telemetry"))]
fn decision_outcome(decision: &Decision) -> &'static str {
match decision {
Decision::Allowed(_) => "allowed",
Decision::Rejected(_) => "rejected",
Decision::Banned(_) => "banned",
}
}
fn log_fingerprint(key: &str) -> String {
format!("<{} chars>", key.len())
}
use crate::matchers::{IdentifierExtractor, RequestContext, RuleMatcher};
use crate::rules::{RuleBuilder, StatsManager, StatsSnapshot};
#[cfg(all(feature = "ban-manager", not(feature = "parallel-checker")))]
use crate::error::BanInfo;
#[cfg(all(feature = "ban-manager", not(feature = "parallel-checker")))]
use crate::storage::BanTarget;
use crate::storage::{BanStorage, Storage};
use dashmap::DashMap;
#[cfg(feature = "parallel-checker")]
use log::warn;
use log::{debug, error, info, trace};
use std::sync::Arc;
use tokio::sync::RwLock;
#[cfg(feature = "ban-manager")]
use crate::BanSource;
#[cfg(feature = "ban-manager")]
use crate::ban::BanManager;
#[cfg(feature = "circuit-breaker")]
use crate::circuit::CircuitBreaker;
#[cfg(feature = "audit-log")]
use crate::logging::AuditLogger;
#[cfg(any(
feature = "parallel-checker",
feature = "ban-manager",
feature = "multi-tenant"
))]
use crate::matchers::Identifier;
#[cfg(feature = "prometheus")]
use crate::telemetry::Metrics;
#[cfg(feature = "telemetry")]
use crate::telemetry::Tracer;
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
pub struct GovernorStats {
pub total_requests: u64,
pub allowed_requests: u64,
pub rejected_requests: u64,
pub banned_requests: u64,
pub error_count: u64,
pub last_updated: Option<chrono::DateTime<chrono::Utc>>,
}
impl From<StatsSnapshot> for GovernorStats {
fn from(snapshot: StatsSnapshot) -> Self {
Self {
total_requests: snapshot.total_requests,
allowed_requests: snapshot.allowed_requests,
rejected_requests: snapshot.rejected_requests,
banned_requests: snapshot.banned_requests,
error_count: snapshot.error_count,
last_updated: snapshot.last_updated,
}
}
}
#[cfg(feature = "fallback")]
pub(crate) struct ConservativeQuotaCounter {
window_start_secs: u64,
count: u64,
}
#[cfg(feature = "fallback")]
impl ConservativeQuotaCounter {
fn try_acquire(&mut self, max_requests: u64, window_secs: u64, now_secs: u64) -> bool {
let window_secs = window_secs.max(1);
if now_secs >= self.window_start_secs.saturating_add(window_secs) {
self.window_start_secs = now_secs;
self.count = 0;
}
if self.count >= max_requests {
return false;
}
self.count += 1;
true
}
}
pub struct Governor {
config: Arc<RwLock<FlowControlConfig>>,
storage: Arc<dyn Storage>,
ban_storage: Arc<dyn BanStorage>,
quota_storage: Option<Arc<dyn crate::storage::QuotaStorage>>,
#[cfg(feature = "fallback")]
conservative_quota: parking_lot::Mutex<Option<ConservativeQuotaCounter>>,
#[cfg(feature = "ban-manager")]
ban_manager: Arc<BanManager>,
#[cfg(feature = "parallel-checker")]
parallel_ban_checker: Arc<crate::storage::ParallelBanChecker>,
decision_chain: Arc<RwLock<DecisionChain>>,
rule_matcher: Arc<RwLock<RuleMatcher>>,
rule_chains: Arc<RwLock<DashMap<String, DecisionChain>>>,
identifier_extractor: Arc<dyn IdentifierExtractor>,
#[cfg(feature = "circuit-breaker")]
circuit_breaker: Arc<CircuitBreaker>,
#[cfg(feature = "audit-log")]
audit_logger: Arc<RwLock<Option<Arc<AuditLogger>>>>,
config_history: Arc<RwLock<ConfigHistory>>,
stats: StatsManager,
l1_cache: L1Cache<CacheableDecision>,
l1_cache_enabled: std::sync::atomic::AtomicBool,
#[cfg(feature = "fallback")]
fallback_manager: Option<Arc<FallbackManager>>,
#[cfg(feature = "event-system")]
event_emitter: Option<Arc<crate::events::EventEmitter>>,
#[cfg(feature = "prometheus")]
metrics: Option<Arc<Metrics>>,
#[cfg(feature = "prometheus")]
per_rule_metrics: bool,
#[cfg(feature = "telemetry")]
tracer: Option<Arc<Tracer>>,
#[cfg(feature = "multi-tenant")]
tenant_resolver: Option<Arc<dyn crate::tenant::TenantResolver>>,
shutdown_token: tokio_util::sync::CancellationToken,
config_watcher_tokens: parking_lot::Mutex<Vec<tokio_util::sync::CancellationToken>>,
shutdown_snapshot_dir: Option<std::path::PathBuf>,
is_shutdown: std::sync::atomic::AtomicBool,
#[cfg(feature = "prometheus")]
exported_rejections: parking_lot::Mutex<ExportedRejections>,
clock: Arc<dyn crate::clock::Clock>,
}
#[cfg(feature = "prometheus")]
#[derive(Default)]
struct ExportedRejections {
per_rule: std::collections::HashMap<String, std::collections::HashMap<String, u64>>,
}
#[cfg(feature = "prometheus")]
impl ExportedRejections {
fn clear(&mut self) {
self.per_rule.clear();
}
}
#[derive(Clone, Default)]
#[allow(clippy::type_complexity)]
pub struct GovernorBuilder {
config: Option<FlowControlConfig>,
storage: Option<Arc<dyn Storage>>,
ban_storage: Option<Arc<dyn BanStorage>>,
quota_storage: Option<Arc<dyn crate::storage::QuotaStorage>>,
identifier_extractor: Option<Arc<dyn crate::matchers::IdentifierExtractor>>,
#[cfg(feature = "circuit-breaker")]
circuit_breaker: Option<Arc<CircuitBreaker>>,
#[cfg(feature = "audit-log")]
audit_logger: Option<Arc<crate::logging::AuditLogger>>,
#[cfg(feature = "prometheus")]
metrics: Option<Arc<Metrics>>,
#[cfg(feature = "prometheus")]
per_rule_metrics: bool,
#[cfg(feature = "telemetry")]
tracer: Option<Arc<Tracer>>,
#[cfg(feature = "parallel-checker")]
parallel_ban_checker: Option<Arc<crate::storage::ParallelBanChecker>>,
l1_cache_config: Option<L1CacheConfig>,
l1_cache_enabled: bool,
#[cfg(feature = "fallback")]
fallback_manager: Option<Arc<FallbackManager>>,
#[cfg(feature = "event-system")]
event_emitter: Option<Arc<crate::events::EventEmitter>>,
custom_matcher_registry: Option<Arc<crate::matchers::custom::CustomMatcherRegistry>>,
#[cfg(feature = "multi-tenant")]
tenant_resolver: Option<Arc<dyn crate::tenant::TenantResolver>>,
shutdown_snapshot_dir: Option<std::path::PathBuf>,
clock: Option<Arc<dyn crate::clock::Clock>>,
}
impl GovernorBuilder {
pub fn new() -> Self {
Self {
config: None,
storage: None,
ban_storage: None,
quota_storage: None,
identifier_extractor: None,
#[cfg(feature = "circuit-breaker")]
circuit_breaker: None,
#[cfg(feature = "audit-log")]
audit_logger: None,
#[cfg(feature = "prometheus")]
metrics: None,
#[cfg(feature = "prometheus")]
per_rule_metrics: true,
#[cfg(feature = "telemetry")]
tracer: None,
#[cfg(feature = "parallel-checker")]
parallel_ban_checker: None,
l1_cache_config: None,
l1_cache_enabled: true, #[cfg(feature = "fallback")]
fallback_manager: None,
#[cfg(feature = "event-system")]
event_emitter: None,
custom_matcher_registry: None,
#[cfg(feature = "multi-tenant")]
tenant_resolver: None,
shutdown_snapshot_dir: None,
clock: None,
}
}
#[cfg(feature = "multi-tenant")]
pub fn with_tenant_resolver(
mut self,
resolver: Arc<dyn crate::tenant::TenantResolver>,
) -> Self {
self.tenant_resolver = Some(resolver);
self
}
pub fn with_shutdown_snapshot_dir(mut self, dir: impl Into<std::path::PathBuf>) -> Self {
self.shutdown_snapshot_dir = Some(dir.into());
self
}
pub fn with_custom_matcher_registry(
mut self,
registry: Arc<crate::matchers::custom::CustomMatcherRegistry>,
) -> Self {
self.custom_matcher_registry = Some(registry);
self
}
pub fn with_config(mut self, config: FlowControlConfig) -> Self {
self.config = Some(config);
self
}
pub fn with_storage(mut self, storage: Arc<dyn Storage>) -> Self {
self.storage = Some(storage);
self
}
pub fn with_ban_storage(mut self, ban_storage: Arc<dyn BanStorage>) -> Self {
self.ban_storage = Some(ban_storage);
self
}
pub fn with_quota_storage(
mut self,
quota_storage: Arc<dyn crate::storage::QuotaStorage>,
) -> Self {
self.quota_storage = Some(quota_storage);
self
}
pub fn with_clock(mut self, clock: Arc<dyn crate::clock::Clock>) -> Self {
self.clock = Some(clock);
self
}
pub fn with_identifier_extractor(
mut self,
extractor: Arc<dyn crate::matchers::IdentifierExtractor>,
) -> Self {
self.identifier_extractor = Some(extractor);
self
}
#[cfg(feature = "circuit-breaker")]
pub fn with_circuit_breaker(mut self, circuit_breaker: Arc<CircuitBreaker>) -> Self {
self.circuit_breaker = Some(circuit_breaker);
self
}
#[cfg(feature = "audit-log")]
pub fn with_audit_logger(mut self, audit_logger: Arc<crate::logging::AuditLogger>) -> Self {
self.audit_logger = Some(audit_logger);
self
}
#[cfg(feature = "prometheus")]
pub fn with_metrics(mut self, metrics: Arc<Metrics>) -> Self {
self.metrics = Some(metrics);
self
}
#[cfg(feature = "prometheus")]
pub fn with_per_rule_metrics(mut self, enabled: bool) -> Self {
self.per_rule_metrics = enabled;
self
}
#[cfg(feature = "telemetry")]
pub fn with_tracer(mut self, tracer: Arc<Tracer>) -> Self {
self.tracer = Some(tracer);
self
}
pub fn with_l1_cache_config(mut self, config: L1CacheConfig) -> Self {
self.l1_cache_config = Some(config);
self
}
pub fn with_l1_cache_enabled(mut self, enabled: bool) -> Self {
self.l1_cache_enabled = enabled;
self
}
#[cfg(feature = "fallback")]
pub fn with_fallback_manager(mut self, fallback_manager: Arc<FallbackManager>) -> Self {
self.fallback_manager = Some(fallback_manager);
self
}
#[cfg(feature = "event-system")]
pub fn with_event_emitter(mut self, emitter: Arc<crate::events::EventEmitter>) -> Self {
self.event_emitter = Some(emitter);
self
}
pub async fn build(self) -> Result<Governor, LimiteronError> {
let config = self.config.unwrap_or_default();
config.validate().map_err(LimiteronError::ConfigError)?;
let storage = self
.storage
.ok_or_else(|| LimiteronError::DependencyError("storage is required".to_string()))?;
let ban_storage = self.ban_storage.ok_or_else(|| {
LimiteronError::DependencyError("ban_storage is required".to_string())
})?;
#[cfg(feature = "ban-manager")]
let ban_manager = {
use crate::ban::{BanManager, BanManagerConfig};
let config = BanManagerConfig::default();
BanManager::with_dependencies(ban_storage.clone(), config)
.await
.map(Arc::new)?
};
#[cfg(feature = "parallel-checker")]
let parallel_ban_checker = self.parallel_ban_checker.unwrap_or_else(|| {
Arc::new(crate::storage::ParallelBanChecker::new(ban_manager.clone()))
});
let identifier_extractor = if let Some(extractor) = self.identifier_extractor {
extractor
} else {
Arc::new(
crate::matchers::CompositeExtractor::builder()
.add_extractor(Box::new(crate::matchers::UserIdExtractor::from_header(
"X-User-Id",
)))
.add_extractor(Box::new(crate::matchers::IpExtractor::builder().build()))
.add_extractor(Box::new(crate::matchers::ApiKeyExtractor::from_header(
"X-API-Key",
)))
.build(),
)
};
let rules = if let Some(registry) = &self.custom_matcher_registry {
RuleBuilder::build_rules_with_registry(&config, registry).await?
} else {
RuleBuilder::build_rules(&config)?
};
let rule_matcher = Arc::new(tokio::sync::RwLock::new(
crate::matchers::RuleMatcher::with_dependencies(rules),
));
let decision_chain = Arc::new(tokio::sync::RwLock::new(
crate::decision_chain::DecisionChain::with_dependencies(vec![]),
));
#[cfg(feature = "circuit-breaker")]
let circuit_breaker = self.circuit_breaker.unwrap_or_else(|| {
Arc::new(CircuitBreaker::with_dependencies(
crate::circuit::CircuitBreakerConfig::default(),
))
});
#[cfg(feature = "audit-log")]
let audit_logger = self
.audit_logger
.map(|logger| Arc::new(tokio::sync::RwLock::new(Some(logger))))
.unwrap_or_else(|| Arc::new(tokio::sync::RwLock::new(None)));
let clock: Arc<dyn crate::clock::Clock> = self
.clock
.clone()
.unwrap_or_else(|| Arc::new(crate::clock::SystemClock));
let rule_chains_map = RuleBuilder::build_rule_chains_with_quota_storage(
&config,
self.quota_storage.as_ref(),
clock.clone(),
)?;
let rule_chains = Arc::new(tokio::sync::RwLock::new(rule_chains_map));
let l1_cache_config = self
.l1_cache_config
.unwrap_or_else(|| L1CacheConfig::new(std::time::Duration::from_secs(60), 10_000));
let l1_cache = L1Cache::with_config(l1_cache_config).await.map_err(|e| {
LimiteronError::DependencyError(format!("Failed to create L1Cache: {}", e))
})?;
let l1_cache_enabled = self.l1_cache_enabled;
#[cfg(feature = "fallback")]
let fallback_manager = self.fallback_manager.clone();
#[cfg(feature = "fallback")]
if let Some(ref fm) = self.fallback_manager {
let l1_cache_ref = l1_cache.clone();
fm.register_island_mode_callback(std::sync::Arc::new(move |is_island| {
if is_island {
let island_config =
IslandModeConfig::new(IslandFallbackStrategy::LocalDecision);
l1_cache_ref.enable_island_mode(island_config);
} else {
l1_cache_ref.disable_island_mode();
}
}))
.await;
log::info!(
target: "governor",
"{}",
t("governor-island-callback-registered", &[])
);
}
Ok(Governor {
config: Arc::new(tokio::sync::RwLock::new(config)),
storage,
ban_storage,
quota_storage: self.quota_storage,
#[cfg(feature = "fallback")]
conservative_quota: parking_lot::Mutex::new(None),
#[cfg(feature = "ban-manager")]
ban_manager,
#[cfg(feature = "parallel-checker")]
parallel_ban_checker,
#[cfg(feature = "circuit-breaker")]
circuit_breaker,
decision_chain,
rule_matcher,
rule_chains,
identifier_extractor,
#[cfg(feature = "audit-log")]
audit_logger,
config_history: Arc::new(tokio::sync::RwLock::new(ConfigHistory::new(100))),
stats: StatsManager::new(),
l1_cache,
l1_cache_enabled: std::sync::atomic::AtomicBool::new(l1_cache_enabled),
#[cfg(feature = "fallback")]
fallback_manager,
#[cfg(feature = "event-system")]
event_emitter: self.event_emitter,
#[cfg(feature = "prometheus")]
metrics: self.metrics,
#[cfg(feature = "prometheus")]
per_rule_metrics: self.per_rule_metrics,
#[cfg(feature = "telemetry")]
tracer: self.tracer,
#[cfg(feature = "multi-tenant")]
tenant_resolver: self.tenant_resolver,
shutdown_token: tokio_util::sync::CancellationToken::new(),
config_watcher_tokens: parking_lot::Mutex::new(Vec::new()),
shutdown_snapshot_dir: self.shutdown_snapshot_dir.clone(),
is_shutdown: std::sync::atomic::AtomicBool::new(false),
#[cfg(feature = "prometheus")]
exported_rejections: parking_lot::Mutex::new(ExportedRejections::default()),
clock,
})
}
}
impl Governor {
pub fn builder() -> GovernorBuilder {
GovernorBuilder::new()
}
pub async fn new() -> Result<Self, LimiteronError> {
use crate::config::{Action, ActionConfig, LimiterConfig, Matcher, Rule};
use crate::storage::{MemoryBanStorage, MemoryStorage};
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let config = FlowControlConfig {
version: "0.1.0".to_string(),
global: crate::config::GlobalConfig::default(),
rules: vec![Rule {
id: "default".to_string(),
name: "Default out-of-the-box rule".to_string(),
priority: 100,
matchers: vec![Matcher::User {
user_ids: vec!["*".to_string()],
}],
limiters: vec![LimiterConfig::TokenBucket {
capacity: 100,
refill_rate: 10,
}],
action: ActionConfig {
on_exceed: Action::Reject,
ban: None,
},
}],
};
Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
}
#[allow(clippy::too_many_arguments)]
pub async fn with_dependencies(
config: Arc<tokio::sync::RwLock<FlowControlConfig>>,
storage: Arc<dyn Storage>,
ban_storage: Arc<dyn BanStorage>,
identifier_extractor: Arc<dyn crate::matchers::IdentifierExtractor>,
rule_matcher: Arc<tokio::sync::RwLock<crate::matchers::RuleMatcher>>,
decision_chain: Arc<tokio::sync::RwLock<crate::decision_chain::DecisionChain>>,
rule_chains: Arc<
tokio::sync::RwLock<DashMap<String, crate::decision_chain::DecisionChain>>,
>,
#[cfg(feature = "circuit-breaker")] circuit_breaker: Arc<CircuitBreaker>,
) -> Result<Self, LimiteronError> {
#[cfg(feature = "ban-manager")]
use crate::ban::{BanManager, BanManagerConfig};
#[cfg(feature = "ban-manager")]
let ban_manager: Arc<BanManager> =
BanManager::with_dependencies(ban_storage.clone(), BanManagerConfig::default())
.await
.map(Arc::new)?;
#[cfg(feature = "parallel-checker")]
let parallel_ban_checker: Arc<crate::storage::ParallelBanChecker> =
Arc::new(crate::storage::ParallelBanChecker::new(ban_manager.clone()));
Ok(Self {
config,
storage,
ban_storage,
quota_storage: None,
#[cfg(feature = "fallback")]
conservative_quota: parking_lot::Mutex::new(None),
#[cfg(feature = "ban-manager")]
ban_manager,
#[cfg(feature = "parallel-checker")]
parallel_ban_checker,
#[cfg(feature = "circuit-breaker")]
circuit_breaker,
decision_chain,
rule_matcher,
rule_chains,
identifier_extractor,
#[cfg(feature = "audit-log")]
audit_logger: Arc::new(tokio::sync::RwLock::new(None)),
config_history: Arc::new(tokio::sync::RwLock::new(ConfigHistory::new(100))),
stats: StatsManager::new(),
l1_cache: L1Cache::new().await.map_err(|e| {
LimiteronError::ConfigError(format!("Failed to create L1Cache: {}", e))
})?,
l1_cache_enabled: std::sync::atomic::AtomicBool::new(true),
#[cfg(feature = "fallback")]
fallback_manager: None,
#[cfg(feature = "event-system")]
event_emitter: None,
#[cfg(feature = "prometheus")]
metrics: None,
#[cfg(feature = "prometheus")]
per_rule_metrics: true,
#[cfg(feature = "telemetry")]
tracer: None,
#[cfg(feature = "multi-tenant")]
tenant_resolver: None,
shutdown_token: tokio_util::sync::CancellationToken::new(),
config_watcher_tokens: parking_lot::Mutex::new(Vec::new()),
shutdown_snapshot_dir: None,
is_shutdown: std::sync::atomic::AtomicBool::new(false),
#[cfg(feature = "prometheus")]
exported_rejections: parking_lot::Mutex::new(ExportedRejections::default()),
clock: Arc::new(crate::clock::SystemClock),
})
}
pub async fn with_storage(
config: FlowControlConfig,
storage: Arc<dyn Storage>,
ban_storage: Arc<dyn BanStorage>,
#[cfg(feature = "prometheus")] metrics: Option<Arc<Metrics>>,
#[cfg(feature = "telemetry")] tracer: Option<Arc<Tracer>>,
) -> Result<Self, LimiteronError> {
let builder = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.with_identifier_extractor(Arc::new(
crate::matchers::CompositeExtractor::builder()
.add_extractor(Box::new(crate::matchers::UserIdExtractor::from_header(
"X-User-Id",
)))
.add_extractor(Box::new(crate::matchers::IpExtractor::builder().build()))
.add_extractor(Box::new(crate::matchers::ApiKeyExtractor::from_header(
"X-API-Key",
)))
.build(),
));
#[cfg(feature = "prometheus")]
let builder = if let Some(m) = metrics {
builder.with_metrics(m)
} else {
builder
};
#[cfg(feature = "telemetry")]
let builder = if let Some(t) = tracer {
builder.with_tracer(t)
} else {
builder
};
builder.build().await
}
pub async fn from_config_file<P: AsRef<std::path::Path>>(
config_path: P,
storage: Arc<dyn Storage>,
ban_storage: Arc<dyn BanStorage>,
) -> Result<Self, LimiteronError> {
let config = crate::config::ConfigLoader::load_from_file(config_path)?;
Self::create_with_config(config, storage, ban_storage).await
}
pub async fn from_config_with_env<P: AsRef<std::path::Path>>(
config_path: P,
storage: Arc<dyn Storage>,
ban_storage: Arc<dyn BanStorage>,
) -> Result<Self, LimiteronError> {
let config = crate::ConfigLoader::load_from_file_with_env(config_path)?;
Self::create_with_config(config, storage, ban_storage).await
}
async fn create_with_config(
config: FlowControlConfig,
storage: Arc<dyn Storage>,
ban_storage: Arc<dyn BanStorage>,
) -> Result<Self, LimiteronError> {
Self::with_storage(
config,
storage,
ban_storage,
#[cfg(feature = "prometheus")]
None,
#[cfg(feature = "telemetry")]
None,
)
.await
}
pub async fn check(&self, context: &RequestContext) -> Result<Decision, LimiteronError> {
#[cfg(feature = "prometheus")]
let start_time = std::time::Instant::now();
#[cfg(feature = "telemetry")]
let check_span = self.tracer.as_ref().map(|t| t.start_span("governor_check"));
#[cfg(feature = "telemetry")]
let result = self.check_inner(context, check_span.as_ref()).await;
#[cfg(not(feature = "telemetry"))]
let result = self.check_inner(context).await;
#[cfg(feature = "audit-log")]
if let Ok(decision) = &result
&& let Some(logger) = self.audit_logger.read().await.as_ref()
{
let (kind, reason) = match decision {
Decision::Allowed(_) => ("allowed".to_string(), String::new()),
Decision::Rejected(m) => ("rejected".to_string(), m.reason.clone()),
Decision::Banned(b) => ("banned".to_string(), b.reason().to_string()),
};
let identifier = self
.identifier_extractor
.extract(context)
.map(|i| i.key())
.unwrap_or_default();
let logger = logger.clone();
tokio::spawn(async move {
logger.log_decision(identifier, kind, reason, None).await;
});
}
#[cfg(feature = "prometheus")]
if let Some(ref metrics) = self.metrics {
let duration = start_time.elapsed();
let allowed = matches!(&result, Ok(Decision::Allowed(_)));
metrics.record_check(duration, allowed);
if matches!(&result, Ok(Decision::Banned(_))) {
metrics.record_ban();
}
if result.is_err() {
metrics.record_error("check_error");
}
}
#[cfg(feature = "telemetry")]
if let Some(s) = check_span {
s.finish();
}
result
}
async fn check_inner(
&self,
context: &RequestContext,
#[cfg(feature = "telemetry")] check_span: Option<&crate::telemetry::Span>,
) -> Result<Decision, LimiteronError> {
#[cfg(feature = "fallback")]
if let Some(ref fallback_mgr) = self.fallback_manager {
#[cfg(feature = "telemetry")]
let result = self
.check_with_fallback(context, fallback_mgr, check_span)
.await;
#[cfg(not(feature = "telemetry"))]
let result = self.check_with_fallback(context, fallback_mgr).await;
return result;
}
#[cfg(feature = "telemetry")]
let result = self.check_internal(context, check_span).await;
#[cfg(not(feature = "telemetry"))]
let result = self.check_internal(context).await;
result
}
async fn check_internal(
&self,
context: &RequestContext,
#[cfg(feature = "telemetry")] check_span: Option<&crate::telemetry::Span>,
) -> Result<Decision, LimiteronError> {
self.stats.increment_total();
let identifier = self.identifier_extractor.extract(context).ok_or_else(|| {
LimiteronError::ConfigError("Failed to extract identifier".to_string())
})?;
trace!(
"Extracted identifier: {}",
log_fingerprint(&identifier.key())
);
#[cfg(feature = "multi-tenant")]
let identifier = self.tenant_scoped_identifier(context, identifier);
let matcher_guard = self.rule_matcher.read().await;
let chains_guard = self.rule_chains.read().await;
let matched_rules = {
#[allow(clippy::disallowed_methods)]
matcher_guard
.match_all(context)
.into_iter()
.cloned()
.collect::<Vec<_>>()
};
#[cfg(feature = "parallel-checker")]
{
let ban_target = identifier.to_ban_target();
if let Some(target) = ban_target {
#[cfg(feature = "circuit-breaker")]
let ban_info = match self
.circuit_breaker
.execute(|| self.parallel_ban_checker.check_single_target(&target))
.await
{
Ok(info) => info,
Err(e) => {
warn!(
"{}",
t(
"governor-ban-check-circuit-open",
&[("error", e.to_string())],
)
);
None
}
};
#[cfg(not(feature = "circuit-breaker"))]
let ban_info = self
.parallel_ban_checker
.check_single_target(&target)
.await?;
if let Some(info) = ban_info {
warn!(
"{}",
t(
"governor-request-banned",
&[
("user", log_fingerprint(&identifier.key()).to_string(),),
("reason", info.reason().to_string()),
],
)
);
self.stats.increment_banned();
return Ok(Decision::Banned(info));
}
}
}
if self.is_l1_cache_enabled() && !matched_rules.is_empty() {
let cache_key = self.build_cache_key_multi(&identifier, &matched_rules);
if let Ok(Some(cached_decision)) = self.l1_cache.get(&cache_key).await {
let decision = cached_decision.to_decision();
if !matches!(decision, Decision::Allowed(_)) {
trace!(
"L1 cache hit (reject decision): key={}",
log_fingerprint(&cache_key)
);
#[cfg(feature = "prometheus")]
if let Some(ref metrics) = self.metrics {
if self.per_rule_metrics
&& let Some(rule_id) = cached_decision.rule_id.as_deref()
{
metrics.record_rule_outcome(rule_id, decision_outcome(&decision));
} else {
metrics.record_negative_cache_hit();
}
}
#[cfg(feature = "telemetry")]
if let Some(span) = check_span {
if let Some(rule_id) = cached_decision.rule_id.as_deref() {
span.set_attribute("rule.id", rule_id);
}
span.set_attribute("rule.outcome", decision_outcome(&decision));
}
self.update_stats_for_decision(&Result::Ok(decision.clone()));
return Ok(decision);
}
trace!(
"L1 cache hit but allow decision, ignoring and re-running full check: key={}",
log_fingerprint(&cache_key)
);
}
}
if matched_rules.is_empty() {
let result = self.decision_chain.read().await.check().await;
self.update_stats_for_decision(&result);
return result;
}
let rule_chains = &*chains_guard;
for rule in &matched_rules {
if let Some(chain) = rule_chains.get(&rule.id) {
let result = chain.check().await;
#[cfg(any(feature = "prometheus", feature = "telemetry"))]
let outcome: &str = match &result {
Ok(decision) => decision_outcome(decision),
Err(_) => "error",
};
#[cfg(feature = "prometheus")]
if self.per_rule_metrics
&& let Some(ref metrics) = self.metrics
{
metrics.record_rule_outcome(&rule.id, outcome);
if chain.node_rejections_version() != chain.exported_rejections_version() {
self.export_chain_rejections(&chain, &rule.id, metrics);
}
}
match result {
Ok(Decision::Allowed(_)) => {
continue;
}
_ => {
self.update_stats_for_decision(&result);
#[cfg(feature = "telemetry")]
if let Some(span) = check_span {
span.set_attribute("rule.id", &rule.id);
span.set_attribute("rule.outcome", outcome);
}
if self.is_l1_cache_enabled()
&& let Ok(ref decision) = result
{
let cache_key = self.build_cache_key_multi(&identifier, &matched_rules);
let mut cacheable = CacheableDecision::from_decision(decision);
cacheable.rule_id = Some(rule.id.clone());
let _ = self.l1_cache.set(cache_key, cacheable).await;
trace!("L1 cache updated: decision=rejected");
}
#[cfg(feature = "event-system")]
{
if let Ok(ref decision) = result {
let decision_str = match decision {
Decision::Banned(_) => "Banned",
Decision::Rejected(_) => "Denied",
_ => "Allowed",
};
self.emit_rate_limit_event(
&identifier.key(),
&rule.id,
decision_str,
)
.await;
}
}
return result;
}
}
}
}
self.stats.increment_allowed();
let decision = Decision::allowed_default();
#[cfg(feature = "telemetry")]
if let Some(span) = check_span {
span.set_attribute("rule.outcome", "allowed");
}
Ok(decision)
}
#[cfg(feature = "prometheus")]
fn export_chain_rejections(&self, chain: &DecisionChain, rule_id: &str, metrics: &Metrics) {
let mut exported = self.exported_rejections.lock();
if chain.exported_rejections_version() >= chain.node_rejections_version() {
return;
}
let nodes = match exported.per_rule.get_mut(rule_id) {
Some(nodes) => nodes,
None => exported.per_rule.entry(rule_id.to_owned()).or_default(),
};
let version = chain.for_each_node_rejection(|node, count| match nodes.get_mut(node) {
Some(last) => {
let delta = count.saturating_sub(*last);
if delta > 0 {
metrics.record_limiter_rejections(rule_id, node, delta);
}
*last = count;
}
None => {
nodes.insert(node.to_owned(), count);
metrics.record_limiter_rejections(rule_id, node, count);
}
});
chain.set_exported_rejections_version(version);
}
#[cfg(feature = "fallback")]
async fn check_with_fallback(
&self,
context: &RequestContext,
fallback_mgr: &Arc<FallbackManager>,
#[cfg(feature = "telemetry")] check_span: Option<&crate::telemetry::Span>,
) -> Result<Decision, LimiteronError> {
let context_clone = context.clone();
fallback_mgr
.execute_with_fallback(
crate::fallback::ComponentType::Redis,
|| async {
#[cfg(feature = "telemetry")]
let result = self.check_internal(&context_clone, check_span).await;
#[cfg(not(feature = "telemetry"))]
let result = self.check_internal(&context_clone).await;
#[cfg(feature = "prometheus")]
self.note_degraded_check(&result);
result
},
|| async {
self.check_l1_cache_only(&context_clone).await
},
)
.await
}
#[cfg(all(feature = "prometheus", feature = "fallback"))]
fn note_degraded_check(&self, result: &Result<Decision, LimiteronError>) {
if let Err(LimiteronError::StorageError(_)) = result
&& let Some(ref metrics) = self.metrics
{
metrics.record_degraded();
}
}
#[cfg(feature = "fallback")]
async fn check_l1_cache_only(
&self,
context: &RequestContext,
) -> Result<Decision, LimiteronError> {
if !self.is_l1_cache_enabled() {
return Err(LimiteronError::LimitError(t(
"governor-storage-failure-no-l1",
&[],
)));
}
let identifier = self.identifier_extractor.extract(context).ok_or_else(|| {
LimiteronError::ConfigError("Failed to extract identifier".to_string())
})?;
#[cfg(feature = "multi-tenant")]
let identifier = self.tenant_scoped_identifier(context, identifier);
let matched_rules = {
let matcher = self.rule_matcher.read().await;
#[allow(clippy::disallowed_methods)]
matcher
.match_all(context)
.into_iter()
.cloned()
.collect::<Vec<_>>()
};
if matched_rules.is_empty() {
return Ok(Decision::allowed_default());
}
let cache_key = self.build_cache_key_multi(&identifier, &matched_rules);
let mut valid_cached: Option<Decision> = None;
if let Ok(Some(cached_decision)) = self.l1_cache.get(&cache_key).await {
match cached_decision.to_decision_strict() {
Ok(decision) => {
trace!(
"island mode - L1 cache hit: key={}",
log_fingerprint(&cache_key)
);
valid_cached = Some(decision);
}
Err(e) => {
log::warn!(
"island mode - corrupted cache entry: {e}; falling back to island strategy"
);
self.stats.increment_error();
}
}
}
if let Some(decision) = valid_cached {
self.update_stats_for_decision(&Result::Ok(decision.clone()));
return Ok(decision);
}
{
if self.l1_cache.is_island_mode() {
if let Some(config) = self.l1_cache.island_config() {
match config.fallback_strategy {
IslandFallbackStrategy::AllowAll => {
log::warn!(
target: "governor",
"{}",
t("governor-island-allow-all", &[])
);
let decision = Decision::allowed_default();
self.update_stats_for_decision(&Ok(decision.clone()));
Ok(decision)
}
IslandFallbackStrategy::RejectAll => {
log::warn!(
target: "governor",
"{}",
t("governor-island-reject-all", &[])
);
self.stats.increment_error();
Err(LimiteronError::LimitError(t(
"governor-island-reject-storage-failure",
&[],
)))
}
IslandFallbackStrategy::LocalDecision => {
log::warn!(
target: "governor",
"{}",
t("governor-island-l1-miss-conservative", &[])
);
let decision = Decision::allowed_default();
self.update_stats_for_decision(&Ok(decision.clone()));
Ok(decision)
}
IslandFallbackStrategy::ConservativeQuota {
max_requests,
window_secs,
} => {
let max_requests = u64::from(max_requests);
log::warn!(
target: "governor",
"{}",
t(
"governor-island-conservative-quota",
&[
("max", max_requests.to_string()),
("window", window_secs.to_string()),
],
)
);
let now_secs = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
let mut guard = self.conservative_quota.lock();
let counter = guard.get_or_insert_with(|| ConservativeQuotaCounter {
window_start_secs: now_secs,
count: 0,
});
let decision =
if counter.try_acquire(max_requests, window_secs, now_secs) {
Decision::allowed_default()
} else {
Decision::Rejected(RejectionMetadata {
reason: t(
"governor-island-conservative-quota",
&[
("max", max_requests.to_string()),
("window", window_secs.to_string()),
],
),
retry_after: window_secs.max(1),
limit: max_requests,
reset_at: now_secs + window_secs.max(1),
})
};
drop(guard);
self.update_stats_for_decision(&Ok(decision.clone()));
Ok(decision)
}
}
} else {
let decision = Decision::allowed_default();
self.update_stats_for_decision(&Ok(decision.clone()));
Ok(decision)
}
} else {
self.stats.increment_error();
Err(LimiteronError::LimitError(t(
"governor-storage-failure-cache-miss",
&[],
)))
}
}
}
fn build_cache_key_multi(
&self,
identifier: &crate::matchers::Identifier,
rules: &[crate::Rule],
) -> String {
let rule_part = rules
.iter()
.map(|r| r.id.as_str())
.collect::<Vec<_>>()
.join("|");
match identifier {
crate::matchers::Identifier::UserId(user_id) => {
RateLimitCacheKey::user_rate_limit(user_id, &rule_part)
}
crate::matchers::Identifier::Ip(ip) => RateLimitCacheKey::ip_rate_limit(ip, &rule_part),
crate::matchers::Identifier::ApiKey(api_key) => {
RateLimitCacheKey::api_key_rate_limit(api_key, &rule_part)
}
_ => RateLimitCacheKey::generic(&identifier.key(), &rule_part),
}
}
fn update_stats_for_decision(&self, result: &Result<Decision, LimiteronError>) {
match result {
Ok(Decision::Allowed(_)) => {
self.stats.increment_allowed();
}
Ok(Decision::Banned(_)) => {
self.stats.increment_banned();
}
Ok(Decision::Rejected(_)) => {
self.stats.increment_rejected();
}
Err(_) => {
self.stats.increment_error();
}
}
}
#[cfg(feature = "event-system")]
async fn emit_rate_limit_event(&self, key: &str, rule_id: &str, decision: &str) {
if let Some(ref emitter) = self.event_emitter {
let event = crate::events::Event::new(crate::events::EventType::RateLimitTriggered {
key: key.to_string(),
rule_id: rule_id.to_string(),
decision: decision.to_string(),
});
if let Err(e) = emitter.emit(event).await {
error!(
"{}",
t(
"governor-emit-rate-limit-event-failed",
&[("error", e.to_string())]
)
);
}
}
}
#[cfg(feature = "parallel-checker")]
pub async fn check_resource_parallel(
&self,
resource: &str,
) -> Result<Decision, LimiteronError> {
let ban_info = self
.parallel_ban_checker
.check_user_banned(resource)
.await?;
match ban_info {
Some(info) => {
warn!(
"{}",
t(
"governor-resource-banned",
&[
("resource", resource.to_string()),
("reason", info.reason().to_string()),
],
)
);
Ok(Decision::Banned(info))
}
None => Ok(Decision::allowed_default()),
}
}
#[cfg(not(feature = "parallel-checker"))]
pub async fn check_resource_parallel(
&self,
_resource: &str,
) -> Result<Decision, LimiteronError> {
#[cfg(feature = "ban-manager")]
{
let target = BanTarget::UserId(_resource.to_string());
let ban_record = self.ban_manager.is_banned(&target).await?;
if let Some(record) = ban_record {
return Ok(Decision::Banned(BanInfo::new(
record.reason,
record.expires_at,
record.ban_times,
)));
}
Ok(Decision::allowed_default())
}
#[cfg(not(feature = "ban-manager"))]
{
Err(LimiteronError::ConfigError(t(
"governor-resource-ban-check-unavailable",
&[],
)))
}
}
#[cfg(feature = "ban-manager")]
pub async fn ban_identifier(
&self,
identifier: &Identifier,
reason: &str,
source: Option<BanSource>,
) -> Result<(), LimiteronError> {
debug!(
"Ban user: {} reason: {}",
log_fingerprint(&identifier.key()),
reason
);
let ban_target = identifier.to_ban_target();
if let Some(target) = ban_target {
let ban_source = source.unwrap_or(BanSource::Manual {
operator: "unknown".to_string(),
});
self.ban_manager
.create_ban(
target,
reason.to_string(),
ban_source,
serde_json::json!({}),
None,
)
.await?;
info!(
"{}",
t(
"governor-user-banned",
&[(
"user",
crate::logging::redact_user_id(Some(identifier.key().as_ref())).to_string(),
)],
)
);
} else {
return Err(LimiteronError::ValidationError(
"Unsupported identifier type".to_string(),
));
}
Ok(())
}
#[cfg(feature = "ban-manager")]
pub async fn unban_identifier(&self, identifier: &Identifier) -> Result<(), LimiteronError> {
debug!("Unban user: {}", log_fingerprint(&identifier.key()));
let ban_target = identifier.to_ban_target();
if let Some(target) = ban_target {
let unbanned = self
.ban_manager
.delete_ban(&target, "unknown".to_string())
.await?;
if unbanned {
info!(
"{}",
t(
"governor-user-unbanned",
&[(
"user",
crate::logging::redact_user_id(Some(identifier.key().as_ref()))
.to_string(),
)],
)
);
}
Ok(())
} else {
Err(LimiteronError::ValidationError(
"Unsupported identifier type".to_string(),
))
}
}
pub async fn get_config_history(&self) -> Vec<ConfigChangeRecord> {
self.config_history.read().await.get_records().to_vec()
}
#[cfg(feature = "multi-tenant")]
pub fn resolve_tenant(&self, context: &RequestContext) -> Option<crate::tenant::Namespace> {
self.tenant_resolver
.as_ref()
.and_then(|resolver| resolver.resolve(context))
}
#[cfg(feature = "multi-tenant")]
pub fn decision_key(&self, context: &RequestContext, identifier: &Identifier) -> String {
match self.resolve_tenant(context) {
Some(ns) if !ns.is_default() => ns.qualify_key(&identifier.key()),
_ => identifier.key(),
}
}
#[cfg(feature = "multi-tenant")]
fn tenant_scoped_identifier(
&self,
context: &RequestContext,
identifier: Identifier,
) -> Identifier {
let Some(ns) = self.resolve_tenant(context) else {
return identifier;
};
if ns.is_default() {
return identifier;
}
let qualified = ns.qualify_key(identifier.as_str());
match identifier {
Identifier::UserId(_) => Identifier::UserId(qualified),
Identifier::Ip(_) => Identifier::Ip(qualified),
Identifier::Mac(_) => Identifier::Mac(qualified),
Identifier::ApiKey(_) => Identifier::ApiKey(qualified),
Identifier::DeviceId(_) => Identifier::DeviceId(qualified),
}
}
#[cfg(all(feature = "multi-tenant", feature = "ban-manager"))]
pub async fn ban_identifier_for_namespace(
&self,
namespace: &crate::tenant::Namespace,
identifier: &Identifier,
reason: &str,
duration: Option<std::time::Duration>,
) -> Result<(), LimiteronError> {
let Some(target) = identifier.to_ban_target() else {
return Err(LimiteronError::ValidationError(
"Unsupported identifier type".to_string(),
));
};
let target = if namespace.is_default() {
target
} else {
crate::storage::qualify_ban_target(&target, namespace).unwrap_or(target)
};
let now = chrono::Utc::now();
let duration = duration.unwrap_or(std::time::Duration::from_secs(3600));
let record = crate::storage::BanRecord {
target,
ban_times: 1,
duration,
banned_at: now,
expires_at: now
+ chrono::Duration::from_std(duration).unwrap_or(chrono::Duration::hours(1)),
is_manual: true,
reason: reason.to_string(),
};
self.ban_storage
.save(&record)
.await
.map_err(LimiteronError::StorageError)?;
let message_key = if namespace.is_default() {
"governor-global-ban-applied"
} else {
"governor-tenant-ban-applied"
};
info!(
"{}",
t(
message_key,
&[
("namespace", namespace.to_string()),
("key", log_fingerprint(&identifier.key()).to_string(),),
],
)
);
Ok(())
}
#[cfg(all(feature = "multi-tenant", feature = "ban-manager"))]
pub async fn is_identifier_banned(
&self,
context: &RequestContext,
identifier: &Identifier,
) -> Result<Option<crate::storage::BanRecord>, LimiteronError> {
let Some(target) = identifier.to_ban_target() else {
return Ok(None);
};
if let Some(ns) = self.resolve_tenant(context)
&& !ns.is_default()
&& let Some(scoped) = crate::storage::qualify_ban_target(&target, &ns)
&& let Some(record) = self.ban_storage.is_banned(&scoped).await?
{
return Ok(Some(record));
}
Ok(self.ban_storage.is_banned(&target).await?)
}
pub async fn stop_config_watcher(&self) -> Result<(), LimiteronError> {
for token in self.config_watcher_tokens.lock().drain(..) {
token.cancel();
}
info!("{}", t("governor-config-watcher-stopped", &[]));
Ok(())
}
pub async fn manual_config_check(&self) -> Result<bool, LimiteronError> {
info!("{}", t("governor-manual-config-check", &[]));
let _config = self.config.read().await;
Ok(true)
}
pub async fn stats(&self) -> GovernorStats {
self.stats.snapshot().into()
}
pub async fn decision_chain_stats(&self) -> crate::decision_chain::ChainStats {
self.decision_chain.read().await.stats().await
}
pub async fn rule_matcher_stats(&self) -> crate::matchers::MatcherStats {
self.rule_matcher.read().await.stats()
}
pub async fn reset_stats(&self) {
info!("{}", t("governor-stats-reset", &[]));
self.decision_chain.write().await.reset_stats().await;
self.rule_matcher.write().await.reset_stats();
self.stats.reset();
self.l1_cache.reset_stats();
}
pub fn config_handle(&self) -> Arc<RwLock<FlowControlConfig>> {
self.config.clone()
}
pub async fn apply_config(
&self,
new_config: FlowControlConfig,
) -> Result<ConfigApplyReport, LimiteronError> {
new_config.validate().map_err(LimiteronError::ConfigError)?;
let rules = RuleBuilder::build_rules(&new_config)?;
let new_matcher = RuleMatcher::with_dependencies(rules);
let new_chains = RuleBuilder::build_rule_chains_with_quota_storage(
&new_config,
self.quota_storage.as_ref(),
Arc::clone(&self.clock),
)?;
let new_hash = new_config.compute_hash();
let new_version = new_config.version.clone();
let rule_count = new_config.rules.len();
let (old_version, old_hash) = {
let mut matcher_guard = self.rule_matcher.write().await;
let mut chains_guard = self.rule_chains.write().await;
let mut cfg_guard = self.config.write().await;
*matcher_guard = new_matcher;
*chains_guard = new_chains;
let old = std::mem::replace(&mut *cfg_guard, new_config);
#[cfg(feature = "prometheus")]
self.exported_rejections.lock().clear();
let old_hash = old.compute_hash();
let old_version = old.version;
(old_version, old_hash)
};
if let Err(e) = self.l1_cache.clear().await {
log::warn!("apply_config: L1 cache clear failed: {e}");
}
self.config_history
.write()
.await
.add_record(ConfigChangeRecord {
timestamp: chrono::Utc::now(),
old_version: Some(old_version.clone()),
new_version: new_version.clone(),
old_hash: Some(old_hash),
new_hash: new_hash.clone(),
source: crate::config::ChangeSource::Api,
changes: vec![format!("rules={rule_count} (hot-reload)")],
});
info!(
"apply_config: hot-reloaded config {} → {} (rules={})",
old_version, new_version, rule_count
);
Ok(ConfigApplyReport {
applied: true,
old_version,
new_version,
config_hash: new_hash,
rule_count,
applied_at: chrono::Utc::now(),
})
}
pub fn enable_l1_cache(&self) {
self.l1_cache_enabled
.store(true, std::sync::atomic::Ordering::Release);
info!("{}", t("governor-l1-cache-enabled", &[]));
}
pub fn disable_l1_cache(&self) {
self.l1_cache_enabled
.store(false, std::sync::atomic::Ordering::Release);
info!("{}", t("governor-l1-cache-disabled", &[]));
}
pub fn is_l1_cache_enabled(&self) -> bool {
self.l1_cache_enabled
.load(std::sync::atomic::Ordering::Acquire)
}
pub async fn clear_l1_cache(&self) {
let _ = self.l1_cache.clear().await;
info!("{}", t("governor-l1-cache-cleared", &[]));
}
pub async fn evict_expired_l1_cache(&self) -> usize {
let evicted = self.l1_cache.evict_expired().await.unwrap_or(0);
if evicted > 0 {
debug!("L1 cache evicted {} expired entries", evicted);
}
evicted
}
pub async fn invalidate_l1_cache(&self, identifier: &str) {
let _ = self
.l1_cache
.invalidate_by_prefix(&format!("rl:user:{}:", identifier))
.await;
let _ = self
.l1_cache
.invalidate_by_prefix(&format!("rl:ip:{}:", identifier))
.await;
let _ = self
.l1_cache
.invalidate_by_prefix(&format!("rl:apikey:{}:", identifier))
.await;
let _ = self
.l1_cache
.invalidate_by_prefix(&format!("rl:generic:{}:", identifier))
.await;
let _ = self
.l1_cache
.invalidate(&RateLimitCacheKey::ban_check(identifier))
.await;
debug!("invalidated L1 cache for identifier {}", identifier);
}
pub async fn invalidate_rule_cache(&self, rule_id: &str) {
let _ = self
.l1_cache
.invalidate_containing(&format!(":{}", rule_id))
.await;
debug!("invalidated L1 cache for rule {}", rule_id);
}
pub async fn l1_cache_size(&self) -> usize {
self.l1_cache.len().await.unwrap_or(0)
}
#[cfg(feature = "audit-log")]
pub async fn set_audit_logger(&self, audit_logger: Arc<AuditLogger>) {
let mut logger = self.audit_logger.write().await;
*logger = Some(audit_logger);
info!("{}", t("governor-audit-logger-set", &[]));
}
#[cfg(feature = "audit-log")]
pub async fn audit_logger(&self) -> Option<Arc<AuditLogger>> {
let guard = self.audit_logger.read().await;
guard.as_ref().cloned()
}
pub async fn introspect(&self) -> IntrospectionSnapshot {
let config = self.config.read().await;
let config_version = config.version.clone();
let rules = config
.rules
.iter()
.map(|r| RuleIntrospection {
id: r.id.clone(),
name: r.name.clone(),
priority: r.priority,
matcher_count: r.matchers.len(),
limiters: r.limiters.iter().map(limiter_kind).collect(),
ban_on_exceed: r.action.ban.is_some(),
})
.collect();
drop(config);
let rule_chains = self.rule_chains.read().await;
let mut chains: Vec<ChainIntrospection> = Vec::with_capacity(rule_chains.len());
for entry in rule_chains.iter() {
let chain = entry.value();
let stats = chain.stats().await;
chains.push(ChainIntrospection {
rule_id: entry.key().clone(),
node_count: chain.node_count(),
enabled_node_count: chain.enabled_node_count(),
total_checks: stats.total_checks,
allowed_count: stats.allowed_count,
rejected_count: stats.rejected_count,
error_count: stats.error_count,
node_rejections: stats.node_rejections,
});
}
drop(rule_chains);
chains.sort_by(|a, b| a.rule_id.cmp(&b.rule_id));
IntrospectionSnapshot {
config_version,
rules,
chains,
stats: self.stats().await,
l1_cache: L1Introspection {
enabled: self.is_l1_cache_enabled(),
size: self.l1_cache_size().await,
},
health: HealthIntrospection::from(&self.health_status().await),
#[cfg(feature = "multi-tenant")]
multi_tenant_enabled: self.tenant_resolver.is_some(),
is_shutdown: self.is_shutdown(),
}
}
pub async fn health_check(&self) -> Result<(), LimiteronError> {
info!("{}", t("governor-health-check", &[]));
let status = self.health_status().await;
if status.healthy() {
Ok(())
} else {
let mut issues = Vec::new();
if !status.storage_healthy {
issues.push("storage".to_string());
}
if !status.ban_storage_healthy {
issues.push("ban_storage".to_string());
}
if !status.cache_healthy {
issues.push("cache".to_string());
}
if !status.background_tasks_alive {
issues.push("background_tasks".to_string());
}
Err(LimiteronError::StorageError(
crate::error::StorageError::ConnectionError(format!(
"Components unhealthy: {}",
issues.join(", ")
)),
))
}
}
pub async fn health_status(&self) -> HealthStatus {
let storage_healthy = self.storage.get("__health_probe__").await.is_ok();
let ban_storage_healthy = self.ban_storage.list_bans(false, 0, 1).await.is_ok();
let _config_guard = self.config.read().await;
let cache_healthy = true;
#[cfg(feature = "ban-manager")]
{
let _ = self.ban_manager.get_config().await;
}
#[cfg(feature = "circuit-breaker")]
{
let _ = self.circuit_breaker.get_state().await;
}
#[cfg(feature = "audit-log")]
{
let _ = self.audit_logger().await;
}
HealthStatus {
storage_healthy,
ban_storage_healthy,
cache_healthy,
background_tasks_alive: !self.is_shutdown.load(std::sync::atomic::Ordering::SeqCst),
}
}
pub async fn shutdown(&self) -> Result<(), LimiteronError> {
let already_shutdown = self
.is_shutdown
.compare_exchange(
false,
true,
std::sync::atomic::Ordering::SeqCst,
std::sync::atomic::Ordering::SeqCst,
)
.is_err();
if already_shutdown {
info!("{}", t("governor-already-shutdown", &[]));
return Ok(());
}
info!("{}", t("governor-shutdown-started", &[]));
if let Err(e) = self.stop_config_watcher().await {
error!("failed to stop config watcher during shutdown: {}", e);
}
self.shutdown_token.cancel();
self.persist_shutdown_snapshot();
#[cfg(feature = "ban-manager")]
self.ban_manager.stop_auto_unban_task().await;
#[cfg(feature = "audit-log")]
self.detach_audit_logger().await;
self.clear_l1_cache().await;
info!("{}", t("governor-shutdown-complete", &[]));
Ok(())
}
fn persist_shutdown_snapshot(&self) {
let Some(dir) = self.shutdown_snapshot_dir.as_ref() else {
return;
};
let snapshot = serde_json::json!({
"captured_at": chrono::Utc::now().to_rfc3339(),
"stats": GovernorStats::from(self.stats.snapshot()),
});
let payload = match serde_json::to_vec_pretty(&snapshot) {
Ok(payload) => payload,
Err(e) => {
error!("failed to serialize shutdown snapshot: {}", e);
return;
}
};
if let Err(e) = std::fs::create_dir_all(dir) {
error!("failed to create snapshot dir {}: {}", dir.display(), e);
return;
}
let target = dir.join("governor-shutdown-snapshot.json");
let mut tmp = dir.join("governor-shutdown-snapshot.json.tmp");
match Self::create_private_tmp(&tmp, &payload) {
Ok(()) => {}
Err(first_err) => {
tmp = dir.join(format!(
"governor-shutdown-snapshot.json.tmp.{}",
std::process::id()
));
if let Err(e) = Self::create_private_tmp(&tmp, &payload) {
error!(
"failed to write shutdown snapshot to {}: first={} retry={}",
tmp.display(),
first_err,
e
);
return;
}
}
}
if let Err(e) = std::fs::rename(&tmp, &target) {
error!(
"failed to finalize shutdown snapshot at {}: {}",
target.display(),
e
);
}
}
fn create_private_tmp(tmp: &std::path::Path, payload: &[u8]) -> std::io::Result<()> {
use std::io::Write;
let mut file = std::fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(tmp)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
file.set_permissions(std::fs::Permissions::from_mode(0o600))?;
}
file.write_all(payload)?;
file.sync_all().ok();
Ok(())
}
#[cfg(feature = "audit-log")]
async fn detach_audit_logger(&self) {
let logger = self.audit_logger.write().await.take();
if let Some(logger) = logger.and_then(|a| Arc::try_unwrap(a).ok()) {
logger.shutdown().await;
}
}
pub fn shutdown_token(&self) -> &tokio_util::sync::CancellationToken {
&self.shutdown_token
}
pub fn register_config_watcher_token(&self, token: tokio_util::sync::CancellationToken) {
self.config_watcher_tokens.lock().push(token);
}
pub fn is_shutdown(&self) -> bool {
self.is_shutdown.load(std::sync::atomic::Ordering::SeqCst)
}
}
impl Drop for Governor {
fn drop(&mut self) {
self.is_shutdown
.store(true, std::sync::atomic::Ordering::SeqCst);
self.shutdown_token.cancel();
for token in self.config_watcher_tokens.lock().drain(..) {
token.cancel();
}
}
}
#[derive(Debug, Clone)]
pub struct HealthStatus {
pub storage_healthy: bool,
pub ban_storage_healthy: bool,
pub cache_healthy: bool,
pub background_tasks_alive: bool,
}
impl HealthStatus {
pub fn healthy(&self) -> bool {
self.storage_healthy
&& self.ban_storage_healthy
&& self.cache_healthy
&& self.background_tasks_alive
}
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct RuleIntrospection {
pub id: String,
pub name: String,
pub priority: u16,
pub matcher_count: usize,
pub limiters: Vec<String>,
pub ban_on_exceed: bool,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ChainIntrospection {
pub rule_id: String,
pub node_count: usize,
pub enabled_node_count: usize,
pub total_checks: u64,
pub allowed_count: u64,
pub rejected_count: u64,
pub error_count: u64,
pub node_rejections: Vec<(String, u64)>,
}
#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)]
pub struct L1Introspection {
pub enabled: bool,
pub size: usize,
}
#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)]
pub struct HealthIntrospection {
pub storage_healthy: bool,
pub ban_storage_healthy: bool,
pub cache_healthy: bool,
pub background_tasks_alive: bool,
}
impl From<&HealthStatus> for HealthIntrospection {
fn from(s: &HealthStatus) -> Self {
Self {
storage_healthy: s.storage_healthy,
ban_storage_healthy: s.ban_storage_healthy,
cache_healthy: s.cache_healthy,
background_tasks_alive: s.background_tasks_alive,
}
}
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ConfigApplyReport {
pub applied: bool,
pub old_version: String,
pub new_version: String,
pub config_hash: String,
pub rule_count: usize,
pub applied_at: chrono::DateTime<chrono::Utc>,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct IntrospectionSnapshot {
pub config_version: String,
pub rules: Vec<RuleIntrospection>,
pub chains: Vec<ChainIntrospection>,
pub stats: GovernorStats,
pub l1_cache: L1Introspection,
pub health: HealthIntrospection,
#[cfg(feature = "multi-tenant")]
pub multi_tenant_enabled: bool,
pub is_shutdown: bool,
}
fn limiter_kind(l: &crate::config::LimiterConfig) -> String {
use crate::config::LimiterConfig as L;
match l {
L::TokenBucket { .. } => "token_bucket".to_string(),
L::SlidingWindow { .. } => "sliding_window".to_string(),
L::FixedWindow { .. } => "fixed_window".to_string(),
L::Quota { .. } => "quota".to_string(),
L::Concurrency { .. } => "concurrency".to_string(),
L::LeakyBucket { .. } => "leaky_bucket".to_string(),
L::SlidingWindowLog { .. } => "sliding_window_log".to_string(),
L::PriorityQueue { .. } => "priority_queue".to_string(),
L::AdmissionControl { .. } => "admission_control".to_string(),
L::Custom { name, .. } => format!("custom:{name}"),
}
}
#[cfg(test)]
#[allow(clippy::field_reassign_with_default)]
mod governor_construction_tests {
use super::*;
use crate::config::{Action, ActionConfig, FlowControlConfig, LimiterConfig, Matcher, Rule};
use crate::error::StorageError;
use crate::storage::{BanHistory, BanRecord, BanTarget, MemoryBanStorage, MemoryStorage};
use async_trait::async_trait;
fn create_valid_test_config() -> FlowControlConfig {
FlowControlConfig {
version: "0.1.0".to_string(),
global: crate::config::types::GlobalConfig::default(),
rules: vec![Rule {
id: "test_rule".to_string(),
name: "Test Rule".to_string(),
priority: 100,
matchers: vec![Matcher::User {
user_ids: vec!["*".to_string()],
}],
limiters: vec![LimiterConfig::TokenBucket {
capacity: 100,
refill_rate: 10,
}],
action: ActionConfig {
on_exceed: Action::Reject,
ban: None,
},
}],
}
}
#[tokio::test]
async fn test_governor_new_with_valid_config() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 0);
assert_eq!(stats.allowed_requests, 0);
}
#[tokio::test]
async fn test_governor_builder_with_storage() {
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let config = create_valid_test_config();
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 0);
}
#[tokio::test]
async fn test_governor_with_storage_explicit() {
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let config = create_valid_test_config();
#[cfg(feature = "prometheus")]
let metrics: Option<Arc<Metrics>> = None;
#[cfg(feature = "telemetry")]
let tracer: Option<Arc<Tracer>> = None;
let governor = Governor::with_storage(
config,
storage,
ban_storage,
#[cfg(feature = "prometheus")]
metrics,
#[cfg(feature = "telemetry")]
tracer,
)
.await
.expect("Governor creation should succeed");
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 0);
}
#[cfg(feature = "ban-manager")]
#[tokio::test]
async fn test_governor_with_memory_ban_storage() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let result = governor.health_check().await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_governor_check_with_valid_request() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let mut ctx = RequestContext::default();
ctx.user_id = Some("test_user".to_string());
let result = governor.check(&ctx).await;
let _ = result;
}
#[tokio::test]
async fn test_governor_stats_after_requests() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let mut ctx = RequestContext::default();
ctx.user_id = Some("test_user".to_string());
for _ in 0..5 {
let _ = governor.check(&ctx).await;
}
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 5);
}
#[tokio::test]
async fn test_governor_health_check() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let result = governor.health_check().await;
assert!(
result.is_ok(),
"health_check should succeed with healthy MemoryStorage: {:?}",
result.err()
);
}
#[tokio::test]
async fn test_governor_health_status_real() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let status = governor.health_status().await;
assert!(
status.storage_healthy,
"storage_healthy should be true for MemoryStorage"
);
assert!(
status.ban_storage_healthy,
"ban_storage_healthy should be true for MemoryBanStorage"
);
assert!(status.cache_healthy, "cache_healthy should be true");
assert!(
status.background_tasks_alive,
"background_tasks_alive should be true"
);
assert!(status.healthy(), "all components should be healthy");
}
#[tokio::test]
async fn test_shutdown_idempotent() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
assert!(
!governor.is_shutdown(),
"Governor should not be shutdown initially"
);
let result1 = governor.shutdown().await;
assert!(
result1.is_ok(),
"first shutdown should succeed: {:?}",
result1.err()
);
assert!(
governor.is_shutdown(),
"Governor should be shutdown after first call"
);
let result2 = governor.shutdown().await;
assert!(
result2.is_ok(),
"second shutdown should be idempotent: {:?}",
result2.err()
);
assert!(governor.is_shutdown(), "Governor should still be shutdown");
}
#[tokio::test]
async fn test_shutdown_affects_health_status() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let status_before = governor.health_status().await;
assert!(
status_before.background_tasks_alive,
"background_tasks_alive should be true before shutdown"
);
governor.shutdown().await.unwrap();
let status_after = governor.health_status().await;
assert!(
!status_after.background_tasks_alive,
"background_tasks_alive should be false after shutdown"
);
}
#[tokio::test]
async fn test_shutdown_token_cancellation() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let token = governor.shutdown_token();
assert!(
!token.is_cancelled(),
"token should not be cancelled initially"
);
governor.shutdown().await.unwrap();
assert!(
token.is_cancelled(),
"token should be cancelled after shutdown"
);
}
#[tokio::test]
async fn test_shutdown_persists_stats_snapshot() {
let dir = tempfile::tempdir().expect("tempdir should be created");
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.with_shutdown_snapshot_dir(dir.path())
.build()
.await
.expect("Governor build should succeed");
governor.stats.increment_total();
governor.stats.increment_total();
governor.stats.increment_allowed();
governor.stats.increment_rejected();
governor.shutdown().await.unwrap();
let snapshot_path = dir.path().join("governor-shutdown-snapshot.json");
let raw = std::fs::read_to_string(&snapshot_path)
.expect("shutdown snapshot should be written to the configured dir");
let parsed: serde_json::Value =
serde_json::from_str(&raw).expect("snapshot should be valid JSON");
let stats = &parsed["stats"];
assert_eq!(
stats["total_requests"].as_u64(),
Some(2),
"total_requests should be captured in snapshot"
);
assert_eq!(
stats["allowed_requests"].as_u64(),
Some(1),
"allowed_requests should be captured in snapshot"
);
assert_eq!(
stats["rejected_requests"].as_u64(),
Some(1),
"rejected_requests should be captured in snapshot"
);
let captured_at_valid = parsed["captured_at"]
.as_str()
.is_some_and(|ts| chrono::DateTime::parse_from_rfc3339(ts).is_ok());
assert!(
captured_at_valid,
"captured_at should be a valid RFC3339 timestamp"
);
}
#[tokio::test]
async fn test_shutdown_without_snapshot_dir_writes_nothing() {
let dir = tempfile::tempdir().expect("tempdir should be created");
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
governor.shutdown().await.unwrap();
assert!(
dir.path()
.read_dir()
.expect("dir should exist")
.next()
.is_none(),
"no snapshot file should be written when snapshot dir is not configured"
);
}
#[tokio::test]
async fn test_shutdown_creates_missing_snapshot_dir() {
let dir = tempfile::tempdir().expect("tempdir should be created");
let nested = dir.path().join("metrics").join("snapshots");
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.with_shutdown_snapshot_dir(&nested)
.build()
.await
.expect("Governor build should succeed");
governor.shutdown().await.unwrap();
assert!(
nested.join("governor-shutdown-snapshot.json").exists(),
"missing snapshot dir should be created and snapshot written"
);
}
#[cfg(unix)]
#[tokio::test]
async fn test_shutdown_snapshot_tmp_conflict_falls_back() {
let dir = tempfile::tempdir().expect("tempdir should be created");
let tmp_path = dir.path().join("governor-shutdown-snapshot.json.tmp");
std::fs::write(&tmp_path, b"occupied").expect("placeholder should be written");
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.with_shutdown_snapshot_dir(dir.path())
.build()
.await
.expect("Governor build should succeed");
governor.shutdown().await.unwrap();
let snapshot_path = dir.path().join("governor-shutdown-snapshot.json");
assert!(
snapshot_path.exists(),
"snapshot should be written via pid-suffixed tmp despite conflict"
);
let occupied = std::fs::read(&tmp_path).expect("placeholder should remain");
assert_eq!(
occupied, b"occupied",
"conflicting file must not be touched"
);
}
#[tokio::test]
async fn test_shutdown_cancels_registered_config_watcher() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let watcher_token = tokio_util::sync::CancellationToken::new();
governor.register_config_watcher_token(watcher_token.clone());
governor.shutdown().await.unwrap();
assert!(
watcher_token.is_cancelled(),
"registered config watcher token should be cancelled by shutdown"
);
}
#[tokio::test]
async fn test_shutdown_cancels_all_registered_watchers() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let first = tokio_util::sync::CancellationToken::new();
let second = tokio_util::sync::CancellationToken::new();
let third = tokio_util::sync::CancellationToken::new();
governor.register_config_watcher_token(first.clone());
governor.register_config_watcher_token(second.clone());
governor.register_config_watcher_token(third.clone());
governor.shutdown().await.unwrap();
assert!(
first.is_cancelled(),
"first watcher token should be cancelled"
);
assert!(
second.is_cancelled(),
"second watcher token should be cancelled"
);
assert!(
third.is_cancelled(),
"third watcher token should be cancelled"
);
}
#[tokio::test]
async fn test_drop_cancels_all_registered_watchers() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let first = tokio_util::sync::CancellationToken::new();
let second = tokio_util::sync::CancellationToken::new();
{
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
governor.register_config_watcher_token(first.clone());
governor.register_config_watcher_token(second.clone());
}
assert!(
first.is_cancelled(),
"first watcher token cancelled on drop"
);
assert!(
second.is_cancelled(),
"second watcher token cancelled on drop"
);
}
#[tokio::test]
async fn test_stop_config_watcher_cancels_registered_token() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let watcher_token = tokio_util::sync::CancellationToken::new();
governor.register_config_watcher_token(watcher_token.clone());
governor.stop_config_watcher().await.unwrap();
assert!(
watcher_token.is_cancelled(),
"stop_config_watcher should cancel the registered token"
);
governor.stop_config_watcher().await.unwrap();
}
#[tokio::test]
async fn test_drop_cancels_shutdown_tokens() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let watcher_token = tokio_util::sync::CancellationToken::new();
let shutdown_token;
{
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
governor.register_config_watcher_token(watcher_token.clone());
shutdown_token = governor.shutdown_token().clone();
assert!(!shutdown_token.is_cancelled());
}
assert!(
shutdown_token.is_cancelled(),
"shutdown token should be cancelled on drop"
);
assert!(
watcher_token.is_cancelled(),
"registered config watcher token should be cancelled on drop"
);
}
#[tokio::test]
async fn test_drop_after_explicit_shutdown_is_noop() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let shutdown_token;
{
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
shutdown_token = governor.shutdown_token().clone();
governor.shutdown().await.unwrap();
assert!(shutdown_token.is_cancelled());
}
assert!(
shutdown_token.is_cancelled(),
"drop after explicit shutdown should keep token cancelled"
);
}
#[cfg(feature = "ban-manager")]
#[tokio::test]
async fn test_shutdown_stops_auto_unban_task() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
assert!(
governor.ban_manager.is_auto_unban_running().await,
"auto-unban task should be running before shutdown"
);
governor.shutdown().await.unwrap();
assert!(
!governor.ban_manager.is_auto_unban_running().await,
"auto-unban task should be stopped after shutdown"
);
}
#[cfg(feature = "audit-log")]
#[tokio::test]
async fn test_shutdown_detaches_audit_logger() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let audit_logger = Arc::new(
crate::logging::AuditLogger::new(crate::logging::AuditLogConfig::default()).await,
);
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.with_audit_logger(audit_logger)
.build()
.await
.expect("Governor build should succeed");
assert!(
governor.audit_logger.read().await.is_some(),
"audit logger should be attached before shutdown"
);
governor.shutdown().await.unwrap();
assert!(
governor.audit_logger.read().await.is_none(),
"audit logger should be detached after shutdown"
);
}
#[cfg(feature = "audit-log")]
#[tokio::test]
async fn test_shutdown_detaches_audit_logger_with_external_ref() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let external_ref = Arc::new(
crate::logging::AuditLogger::new(crate::logging::AuditLogConfig::default()).await,
);
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.with_audit_logger(external_ref.clone())
.build()
.await
.expect("Governor build should succeed");
governor.shutdown().await.unwrap();
assert!(
governor.audit_logger.read().await.is_none(),
"audit logger slot should be cleared even when ownership is shared"
);
drop(external_ref);
}
struct FailingStorage;
#[async_trait]
impl Storage for FailingStorage {
async fn get(&self, _key: &str) -> Result<Option<String>, StorageError> {
Err(StorageError::ConnectionError(
"mock storage failure".to_string(),
))
}
async fn set(
&self,
_key: &str,
_value: &str,
_ttl: Option<u64>,
) -> Result<(), StorageError> {
Err(StorageError::ConnectionError(
"mock storage failure".to_string(),
))
}
async fn delete(&self, _key: &str) -> Result<(), StorageError> {
Err(StorageError::ConnectionError(
"mock storage failure".to_string(),
))
}
}
struct FailingBanStorage;
#[async_trait]
impl BanStorage for FailingBanStorage {
async fn is_banned(&self, _target: &BanTarget) -> Result<Option<BanRecord>, StorageError> {
Err(StorageError::ConnectionError(
"mock ban storage failure".to_string(),
))
}
async fn save(&self, _record: &BanRecord) -> Result<(), StorageError> {
Err(StorageError::ConnectionError(
"mock ban storage failure".to_string(),
))
}
async fn get_history(
&self,
_target: &BanTarget,
) -> Result<Option<BanHistory>, StorageError> {
Err(StorageError::ConnectionError(
"mock ban storage failure".to_string(),
))
}
async fn increment_ban_times(&self, _target: &BanTarget) -> Result<u64, StorageError> {
Err(StorageError::ConnectionError(
"mock ban storage failure".to_string(),
))
}
async fn get_ban_times(&self, _target: &BanTarget) -> Result<u64, StorageError> {
Err(StorageError::ConnectionError(
"mock ban storage failure".to_string(),
))
}
async fn remove_ban(&self, _target: &BanTarget) -> Result<(), StorageError> {
Err(StorageError::ConnectionError(
"mock ban storage failure".to_string(),
))
}
async fn cleanup_expired_bans(&self) -> Result<u64, StorageError> {
Err(StorageError::ConnectionError(
"mock ban storage failure".to_string(),
))
}
async fn list_bans(
&self,
_active_only: bool,
_offset: u64,
_limit: u64,
) -> Result<Vec<BanRecord>, StorageError> {
Err(StorageError::ConnectionError(
"mock ban storage failure".to_string(),
))
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
#[tokio::test]
async fn test_health_check_returns_err_when_storage_unhealthy() {
let storage: Arc<dyn Storage> = Arc::new(FailingStorage);
let ban_storage: Arc<dyn BanStorage> = Arc::new(FailingBanStorage);
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let result = governor.health_check().await;
assert!(
result.is_err(),
"health_check should return Err when storage is unhealthy"
);
}
#[tokio::test]
async fn test_health_check_returns_err_after_shutdown() {
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let result_before = governor.health_check().await;
assert!(
result_before.is_ok(),
"health_check should pass before shutdown"
);
governor.shutdown().await.unwrap();
let result_after = governor.health_check().await;
assert!(
result_after.is_err(),
"health_check should return Err after shutdown (background_tasks not alive)"
);
}
#[tokio::test]
async fn test_health_check_error_message_contains_component_names() {
let storage: Arc<dyn Storage> = Arc::new(FailingStorage);
let ban_storage: Arc<dyn BanStorage> = Arc::new(FailingBanStorage);
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let err = governor.health_check().await.unwrap_err();
let err_msg = match err {
LimiteronError::StorageError(StorageError::ConnectionError(msg)) => msg,
other => panic!("expected StorageError::ConnectionError, got {other:?}"),
};
assert!(
err_msg.contains("storage"),
"error message should mention 'storage' component: {err_msg}"
);
assert!(
err_msg.contains("ban_storage"),
"error message should mention 'ban_storage' component: {err_msg}"
);
assert!(
err_msg.contains("Components unhealthy"),
"error message should have descriptive prefix: {err_msg}"
);
}
#[tokio::test]
async fn test_health_status_reports_unhealthy_components() {
let storage: Arc<dyn Storage> = Arc::new(FailingStorage);
let ban_storage: Arc<dyn BanStorage> = Arc::new(FailingBanStorage);
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let status = governor.health_status().await;
assert!(
!status.storage_healthy,
"storage_healthy should be false with FailingStorage"
);
assert!(
!status.ban_storage_healthy,
"ban_storage_healthy should be false with FailingBanStorage"
);
assert!(
status.cache_healthy,
"cache_healthy should always be true (in-memory, cannot fail)"
);
assert!(
status.background_tasks_alive,
"background_tasks_alive should be true before shutdown"
);
assert!(
!status.healthy(),
"overall healthy() should be false when any component is down"
);
}
#[test]
fn test_health_status_healthy_all_true() {
let status = HealthStatus {
storage_healthy: true,
ban_storage_healthy: true,
cache_healthy: true,
background_tasks_alive: true,
};
assert!(status.healthy(), "all true should be healthy");
}
#[test]
fn test_health_status_healthy_any_false() {
let cases = [
(
"storage",
HealthStatus {
storage_healthy: false,
ban_storage_healthy: true,
cache_healthy: true,
background_tasks_alive: true,
},
),
(
"ban_storage",
HealthStatus {
storage_healthy: true,
ban_storage_healthy: false,
cache_healthy: true,
background_tasks_alive: true,
},
),
(
"cache",
HealthStatus {
storage_healthy: true,
ban_storage_healthy: true,
cache_healthy: false,
background_tasks_alive: true,
},
),
(
"background_tasks",
HealthStatus {
storage_healthy: true,
ban_storage_healthy: true,
cache_healthy: true,
background_tasks_alive: false,
},
),
];
for (name, status) in cases {
assert!(
!status.healthy(),
"healthy() should be false when {name} is false"
);
}
}
#[tokio::test]
async fn test_governor_new_with_defaults() {
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 0);
}
#[test]
fn test_governor_builder_validation_empty_config() {
let result = std::thread::spawn(|| {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
Governor::builder()
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
})
});
let result = result.join().unwrap();
assert!(result.is_err());
}
#[tokio::test]
async fn test_governor_new_works_with_builtin_default_rule() {
let governor = Governor::new().await.expect("new() should succeed");
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 0);
let ctx = RequestContext::new().with_header("X-User-Id", "out-of-box-user");
for _ in 0..100 {
let result = governor.check(&ctx).await.expect("check should not error");
assert!(matches!(result, Decision::Allowed(_)));
}
let result = governor.check(&ctx).await.expect("check should not error");
assert!(
matches!(result, Decision::Rejected(_)),
"the 101st request must be limited by the built-in fallback rule"
);
}
#[tokio::test]
async fn test_governor_builder_with_l1_cache_config() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
use crate::l1_cache::L1CacheConfig;
let l1_config = L1CacheConfig {
default_ttl: std::time::Duration::from_secs(60),
max_size: 10000,
..Default::default()
};
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.with_l1_cache_config(l1_config)
.build()
.await
.expect("Governor build should succeed");
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 0);
}
#[tokio::test]
async fn test_governor_config_validation_missing_rules() {
let config = FlowControlConfig {
version: "0.1.0".to_string(),
global: crate::config::types::GlobalConfig::default(),
rules: vec![],
};
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let result = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_governor_multiple_requests_stats() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed");
let mut ctx = RequestContext::default();
ctx.user_id = Some("user_a".to_string());
for _ in 0..3 {
let _ = governor.check(&ctx).await;
}
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 3);
let mut ctx_b = RequestContext::default();
ctx_b.user_id = Some("user_b".to_string());
let _ = governor.check(&ctx_b).await;
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 4);
}
#[tokio::test]
async fn test_governor_with_storage_reuse() {
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let governor = Governor::with_storage(
config,
storage.clone(),
ban_storage.clone(),
#[cfg(feature = "prometheus")]
None,
#[cfg(feature = "telemetry")]
None,
)
.await
.expect("Governor creation should succeed");
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 0);
}
#[tokio::test]
async fn test_builder_missing_storage() {
let result = Governor::builder()
.with_config(create_valid_test_config())
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await;
let err = result
.err()
.expect("expected DependencyError for missing storage");
match err {
LimiteronError::DependencyError(msg) => assert!(msg.contains("storage")),
_ => panic!("expected DependencyError"),
}
}
#[tokio::test]
async fn test_builder_missing_ban_storage() {
let result = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.build()
.await;
let err = result
.err()
.expect("expected DependencyError for missing ban_storage");
match err {
LimiteronError::DependencyError(msg) => assert!(msg.contains("ban_storage")),
_ => panic!("expected DependencyError"),
}
}
#[tokio::test]
async fn test_check_identifier_extraction_failure() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let ctx = RequestContext::default();
let result = governor.check(&ctx).await;
match result {
Err(LimiteronError::ConfigError(msg)) => {
assert!(msg.contains("Failed to extract identifier"))
}
other => panic!("Expected ConfigError, got: {:?}", other),
}
}
#[tokio::test]
async fn test_check_total_incremented_on_extraction_failure() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let ctx = RequestContext::default();
let _ = governor.check(&ctx).await;
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 1);
assert_eq!(stats.allowed_requests, 0);
assert_eq!(stats.rejected_requests, 0);
assert_eq!(stats.banned_requests, 0);
}
fn create_non_matching_config() -> FlowControlConfig {
FlowControlConfig {
version: "0.1.0".to_string(),
global: crate::config::types::GlobalConfig::default(),
rules: vec![Rule {
id: "specific_rule".to_string(),
name: "Specific Rule".to_string(),
priority: 100,
matchers: vec![Matcher::User {
user_ids: vec!["specific_user".to_string()],
}],
limiters: vec![LimiterConfig::TokenBucket {
capacity: 100,
refill_rate: 10,
}],
action: ActionConfig {
on_exceed: Action::Reject,
ban: None,
},
}],
}
}
#[tokio::test]
async fn test_check_no_matched_rules_falls_to_default_chain() {
let config = create_non_matching_config();
let governor = Governor::builder()
.with_config(config)
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let mut ctx = RequestContext::default();
ctx.client_ip = Some("192.168.1.1".to_string());
let result = governor.check(&ctx).await;
match result {
Ok(Decision::Allowed(_)) => {}
other => panic!("Expected Allowed via default chain, got: {:?}", other),
}
}
#[tokio::test]
async fn test_check_no_matched_rules_updates_stats() {
let config = create_non_matching_config();
let governor = Governor::builder()
.with_config(config)
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let mut ctx = RequestContext::default();
ctx.client_ip = Some("192.168.1.1".to_string());
let _ = governor.check(&ctx).await;
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 1);
assert_eq!(stats.allowed_requests, 1);
}
#[tokio::test]
async fn test_check_rule_allowed() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let mut ctx = RequestContext::default();
ctx.client_ip = Some("10.0.0.1".to_string());
let result = governor.check(&ctx).await;
match result {
Ok(Decision::Allowed(_)) => {}
other => panic!("Expected Allowed, got: {:?}", other),
}
}
#[tokio::test]
async fn test_check_allowed_increments_stats() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let mut ctx = RequestContext::default();
ctx.client_ip = Some("10.0.0.2".to_string());
let _ = governor.check(&ctx).await;
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 1);
assert_eq!(stats.allowed_requests, 1);
}
fn create_small_capacity_config() -> FlowControlConfig {
FlowControlConfig {
version: "0.1.0".to_string(),
global: crate::config::types::GlobalConfig::default(),
rules: vec![Rule {
id: "small_bucket".to_string(),
name: "Small Bucket Rule".to_string(),
priority: 100,
matchers: vec![Matcher::User {
user_ids: vec!["*".to_string()],
}],
limiters: vec![LimiterConfig::TokenBucket {
capacity: 3,
refill_rate: 1,
}],
action: ActionConfig {
on_exceed: Action::Reject,
ban: None,
},
}],
}
}
#[tokio::test]
async fn test_check_rule_accepted_then_rejected() {
let config = create_small_capacity_config();
let governor = Governor::builder()
.with_config(config)
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let mut ctx = RequestContext::default();
ctx.client_ip = Some("10.0.0.3".to_string());
for _ in 0..3 {
let result = governor.check(&ctx).await;
assert!(
result.is_ok(),
"Expected Ok within capacity, got: {:?}",
result
);
}
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 3);
assert_eq!(stats.allowed_requests, 3);
}
#[tokio::test]
async fn test_check_rejected_updates_rejected_stats() {
let config = create_small_capacity_config();
let governor = Governor::builder()
.with_config(config)
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let mut ctx = RequestContext::default();
ctx.client_ip = Some("10.0.0.4".to_string());
for _ in 0..4 {
let _ = governor.check(&ctx).await;
}
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 4);
}
#[tokio::test]
async fn test_check_multiple_users_isolated() {
let config = create_small_capacity_config();
let governor = Governor::builder()
.with_config(config)
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let mut ctx_a = RequestContext::default();
ctx_a.client_ip = Some("10.0.0.10".to_string());
let mut ctx_b = RequestContext::default();
ctx_b.client_ip = Some("10.0.0.11".to_string());
for _ in 0..3 {
let _ = governor.check(&ctx_a).await;
}
governor.check(&ctx_a).await.expect("check should succeed");
governor.check(&ctx_b).await.expect("check should succeed");
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 5);
}
#[tokio::test]
async fn test_l1_cache_enable_disable() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
assert!(governor.is_l1_cache_enabled());
governor.disable_l1_cache();
assert!(!governor.is_l1_cache_enabled());
governor.enable_l1_cache();
assert!(governor.is_l1_cache_enabled());
}
#[tokio::test]
async fn test_l1_cache_disabled_does_not_cache() {
let config = create_small_capacity_config();
let governor = Governor::builder()
.with_config(config)
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_l1_cache_enabled(false)
.build()
.await
.expect("Governor build should succeed");
assert!(!governor.is_l1_cache_enabled());
let mut ctx = RequestContext::default();
ctx.client_ip = Some("10.0.0.5".to_string());
for _ in 0..3 {
let result = governor.check(&ctx).await;
assert!(result.is_ok());
}
let result = governor.check(&ctx).await;
match result {
Ok(Decision::Rejected(_)) => {}
other => panic!(
"Expected Rejected after exhaust with L1 disabled, got: {:?}",
other
),
}
}
#[tokio::test]
async fn test_l1_cache_clear_and_size() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let size = governor.l1_cache_size().await;
assert_eq!(size, 0);
let _ = governor.clear_l1_cache().await;
assert_eq!(governor.l1_cache_size().await, 0);
}
#[tokio::test]
async fn test_evict_expired_l1_cache() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let evicted = governor.evict_expired_l1_cache().await;
assert_eq!(evicted, 0);
}
#[tokio::test]
async fn test_reset_stats() {
let config = create_small_capacity_config();
let governor = Governor::builder()
.with_config(config)
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let mut ctx = RequestContext::default();
ctx.client_ip = Some("10.0.0.6".to_string());
let _ = governor.check(&ctx).await;
let stats_before = governor.stats().await;
assert_eq!(stats_before.total_requests, 1);
assert_eq!(stats_before.allowed_requests, 1);
governor.reset_stats().await;
let stats_after = governor.stats().await;
assert_eq!(stats_after.total_requests, 0);
assert_eq!(stats_after.allowed_requests, 0);
assert_eq!(stats_after.rejected_requests, 0);
assert_eq!(stats_after.banned_requests, 0);
assert_eq!(stats_after.error_count, 0);
}
#[tokio::test]
async fn test_decision_chain_stats() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let stats = governor.decision_chain_stats().await;
assert_eq!(stats.allowed_count, 0);
assert_eq!(stats.rejected_count, 0);
assert_eq!(stats.error_count, 0);
}
#[tokio::test]
async fn test_rule_matcher_stats() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let stats = governor.rule_matcher_stats().await;
assert_eq!(stats.total_matches, 0);
assert_eq!(stats.total_mismatches, 0);
}
#[tokio::test]
async fn test_get_config_history() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let history = governor.get_config_history().await;
assert!(history.is_empty());
}
#[tokio::test]
async fn test_manual_config_check_returns_ok() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let result = governor.manual_config_check().await;
assert!(result.is_ok());
match result {
Ok(val) => assert!(val),
_ => unreachable!(),
}
}
#[tokio::test]
async fn test_stop_config_watcher() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let result = governor.stop_config_watcher().await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_health_check_returns_ok() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let result = governor.health_check().await;
assert!(result.is_ok());
}
#[cfg(feature = "ban-manager")]
#[tokio::test]
async fn test_ban_identifier_user_id() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let identifier = crate::matchers::Identifier::UserId("test_user".to_string());
let result = governor
.ban_identifier(&identifier, "test ban reason", None)
.await;
assert!(
result.is_ok(),
"ban_identifier should succeed: {:?}",
result
);
}
#[cfg(feature = "ban-manager")]
#[tokio::test]
async fn test_ban_identifier_ip() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let identifier = crate::matchers::Identifier::Ip("192.168.1.1".to_string());
let result = governor
.ban_identifier(&identifier, "test IP ban", None)
.await;
assert!(
result.is_ok(),
"ban_identifier for IP should succeed: {:?}",
result
);
}
#[cfg(feature = "ban-manager")]
#[tokio::test]
async fn test_ban_identifier_unsupported_type() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let identifier = crate::matchers::Identifier::ApiKey("test_key".to_string());
let result = governor.ban_identifier(&identifier, "test ban", None).await;
match result {
Err(LimiteronError::ValidationError(msg)) => {
assert!(msg.contains("Unsupported identifier type"))
}
other => panic!("Expected ValidationError, got: {:?}", other),
}
}
#[cfg(feature = "ban-manager")]
#[tokio::test]
async fn test_unban_identifier() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let identifier = crate::matchers::Identifier::UserId("unban_test_user".to_string());
let ban_result = governor.ban_identifier(&identifier, "temp ban", None).await;
assert!(ban_result.is_ok(), "ban should succeed");
let unban_result = governor.unban_identifier(&identifier).await;
assert!(
unban_result.is_ok(),
"unban should succeed: {:?}",
unban_result
);
}
#[cfg(feature = "ban-manager")]
#[tokio::test]
async fn test_unban_identifier_unsupported_type() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let identifier = crate::matchers::Identifier::ApiKey("test_key".to_string());
let result = governor.unban_identifier(&identifier).await;
match result {
Err(LimiteronError::ValidationError(msg)) => {
assert!(msg.contains("Unsupported identifier type"))
}
other => panic!("Expected ValidationError, got: {:?}", other),
}
}
#[tokio::test]
async fn test_governor_builder_full_chain() {
use crate::l1_cache::L1CacheConfig;
let config = create_valid_test_config();
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let l1_config = L1CacheConfig {
default_ttl: std::time::Duration::from_secs(30),
max_size: 5000,
..Default::default()
};
let governor = Governor::builder()
.with_config(config)
.with_storage(storage)
.with_ban_storage(ban_storage)
.with_l1_cache_config(l1_config)
.with_l1_cache_enabled(true)
.with_identifier_extractor(Arc::new(
crate::matchers::CompositeExtractor::builder()
.add_extractor(Box::new(crate::matchers::UserIdExtractor::from_header(
"X-User-Id",
)))
.build(),
))
.build()
.await
.expect("Governor build with full chain should succeed");
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 0);
assert!(governor.is_l1_cache_enabled());
}
#[tokio::test]
async fn test_invalidate_l1_cache() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
governor.invalidate_l1_cache("test_user").await;
governor.invalidate_rule_cache("test_rule").await;
assert_eq!(governor.l1_cache_size().await, 0);
}
#[test]
fn test_governor_stats_from_snapshot() {
let snapshot = crate::rules::StatsSnapshot {
total_requests: 100,
allowed_requests: 80,
rejected_requests: 15,
banned_requests: 3,
error_count: 2,
last_updated: Some(chrono::Utc::now()),
};
let stats: GovernorStats = snapshot.into();
assert_eq!(stats.total_requests, 100);
assert_eq!(stats.allowed_requests, 80);
assert_eq!(stats.rejected_requests, 15);
assert_eq!(stats.banned_requests, 3);
assert_eq!(stats.error_count, 2);
assert!(stats.last_updated.is_some());
}
#[test]
fn test_governor_stats_default_impl() {
let stats = GovernorStats::default();
assert_eq!(stats.total_requests, 0);
assert_eq!(stats.allowed_requests, 0);
assert_eq!(stats.rejected_requests, 0);
assert_eq!(stats.banned_requests, 0);
assert_eq!(stats.error_count, 0);
assert!(stats.last_updated.is_none());
}
#[test]
fn test_governor_stats_debug_format() {
let stats = GovernorStats::default();
let debug = format!("{:?}", stats);
assert!(debug.contains("total_requests"));
assert!(debug.contains("allowed_requests"));
assert!(debug.contains("rejected_requests"));
assert!(debug.contains("banned_requests"));
assert!(debug.contains("error_count"));
assert!(debug.contains("last_updated"));
}
#[test]
fn test_governor_stats_clone() {
let mut stats = GovernorStats::default();
stats.total_requests = 42;
let cloned = stats.clone();
assert_eq!(cloned.total_requests, 42);
}
#[tokio::test]
async fn test_l1_cache_stats() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let stats = governor.l1_cache.stats().await;
assert_eq!(stats.total_lookups, 0);
assert_eq!(stats.hits, 0);
assert_eq!(stats.misses, 0);
assert_eq!(stats.current_size, 0);
}
#[tokio::test]
async fn test_with_dependencies_constructor() {
let config = Arc::new(tokio::sync::RwLock::new(create_valid_test_config()));
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let extractor: Arc<dyn crate::matchers::IdentifierExtractor> = Arc::new(
crate::matchers::CompositeExtractor::builder()
.add_extractor(Box::new(crate::matchers::UserIdExtractor::from_header(
"X-User-Id",
)))
.build(),
);
let config_value = config.read().await.clone();
let rules = RuleBuilder::build_rules(&config_value).expect("build_rules should succeed");
let rule_matcher = Arc::new(tokio::sync::RwLock::new(
crate::matchers::RuleMatcher::with_dependencies(rules),
));
let decision_chain = Arc::new(tokio::sync::RwLock::new(
crate::decision_chain::DecisionChain::with_dependencies(vec![]),
));
let chain_map = RuleBuilder::build_rule_chains(&config_value)
.expect("build_rule_chains should succeed");
let rule_chains = Arc::new(tokio::sync::RwLock::new(chain_map));
#[cfg(feature = "circuit-breaker")]
let circuit_breaker = {
use crate::circuit::{CircuitBreaker, CircuitBreakerConfig};
Arc::new(CircuitBreaker::with_dependencies(
CircuitBreakerConfig::default(),
))
};
let governor = Governor::with_dependencies(
config,
storage,
ban_storage,
extractor,
rule_matcher,
decision_chain,
rule_chains,
#[cfg(feature = "circuit-breaker")]
circuit_breaker,
)
.await
.expect("with_dependencies should succeed");
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 0);
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_build_cache_key_all_types() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let rule = |id: &str| crate::Rule {
id: id.to_string(),
name: id.to_string(),
priority: 1,
condition: Box::new(crate::MatchCondition::User(vec![])),
enabled: true,
};
let rules = vec![rule("rule_1")];
let user_id = crate::matchers::Identifier::UserId("user123".to_string());
let key = governor.build_cache_key_multi(&user_id, &rules);
assert_eq!(key, "rl:user:user123:rule_1");
let ip = crate::matchers::Identifier::Ip("192.168.1.1".to_string());
let key = governor.build_cache_key_multi(&ip, &rules);
assert_eq!(key, "rl:ip:192.168.1.1:rule_1");
let api_key = crate::matchers::Identifier::ApiKey("key123".to_string());
let key = governor.build_cache_key_multi(&api_key, &rules);
assert_eq!(key, "rl:apikey:key123:rule_1");
let mac = crate::matchers::Identifier::Mac("AA:BB:CC:DD:EE:FF".to_string());
let key = governor.build_cache_key_multi(&mac, &rules);
assert_eq!(key, "rl:generic:mac:AA:BB:CC:DD:EE:FF:rule_1");
let device = crate::matchers::Identifier::DeviceId("device-001".to_string());
let key = governor.build_cache_key_multi(&device, &rules);
assert_eq!(key, "rl:generic:device_id:device-001:rule_1");
}
#[tokio::test]
async fn test_build_cache_key_multi_all_types() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("Governor build should succeed");
let rule = |id: &str| crate::Rule {
id: id.to_string(),
name: id.to_string(),
priority: 1,
condition: Box::new(crate::MatchCondition::User(vec![])),
enabled: true,
};
let rules = vec![rule("rule_a"), rule("rule_b")];
let user_id = crate::matchers::Identifier::UserId("user123".to_string());
let key = governor.build_cache_key_multi(&user_id, &rules);
assert_eq!(key, "rl:user:user123:rule_a|rule_b");
let ip = crate::matchers::Identifier::Ip("192.168.1.1".to_string());
let key = governor.build_cache_key_multi(&ip, &rules);
assert_eq!(key, "rl:ip:192.168.1.1:rule_a|rule_b");
}
#[tokio::test]
async fn test_governor_from_config_file_json() {
let config_json = r#"{
"version": "0.1.0",
"global": {
"storage": "memory",
"cache": "memory",
"metrics": "none"
},
"rules": [{
"id": "test_rule",
"name": "Test Rule",
"priority": 100,
"matchers": [{"type": "User", "user_ids": ["*"]}],
"limiters": [{"type": "TokenBucket", "capacity": 100, "refill_rate": 10}],
"action": {"on_exceed": "reject"}
}]
}"#;
let dir = std::env::temp_dir().join(format!("limiteron_cfg_{}", std::process::id()));
let _ = std::fs::create_dir_all(&dir);
let config_path = dir.join("test_config.json");
std::fs::write(&config_path, config_json).expect("write config file");
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let result = Governor::from_config_file(&config_path, storage, ban_storage).await;
assert!(
result.is_ok(),
"from_config_file failed: {:?}",
result.err()
);
let _ = std::fs::remove_file(&config_path);
let _ = std::fs::remove_dir(&dir);
}
#[tokio::test]
async fn test_governor_from_config_with_env_json() {
let config_json = r#"{
"version": "0.1.0",
"global": {
"storage": "memory",
"cache": "memory",
"metrics": "none"
},
"rules": [{
"id": "test_rule_env",
"name": "Test Rule Env",
"priority": 100,
"matchers": [{"type": "User", "user_ids": ["*"]}],
"limiters": [{"type": "TokenBucket", "capacity": 100, "refill_rate": 10}],
"action": {"on_exceed": "reject"}
}]
}"#;
let dir = std::env::temp_dir().join(format!("limiteron_cfg_env_{}", std::process::id()));
let _ = std::fs::create_dir_all(&dir);
let config_path = dir.join("test_env_config.json");
std::fs::write(&config_path, config_json).expect("write config file");
let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn BanStorage> = Arc::new(MemoryBanStorage::new());
let result = Governor::from_config_with_env(&config_path, storage, ban_storage).await;
assert!(
result.is_ok(),
"from_config_with_env failed: {:?}",
result.err()
);
let _ = std::fs::remove_file(&config_path);
let _ = std::fs::remove_dir(&dir);
}
#[cfg(feature = "parallel-checker")]
#[tokio::test]
async fn test_check_resource_parallel_banned() {
use crate::matchers::Identifier;
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let user_id = Identifier::UserId("banned_resource_user".to_string());
governor
.ban_identifier(&user_id, "test ban for parallel", None)
.await
.expect("ban should succeed");
let result = governor
.check_resource_parallel("banned_resource_user")
.await;
match result {
Ok(Decision::Banned(_)) => {}
other => panic!("Expected Banned for banned user, got: {:?}", other),
}
}
#[tokio::test]
async fn test_health_check_with_all_features() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let result = governor.health_check().await;
assert!(result.is_ok(), "health_check should succeed: {:?}", result);
}
}
#[cfg(test)]
#[allow(clippy::field_reassign_with_default, dead_code, unused_imports)]
mod governor_feature_gated_tests {
use super::*;
use crate::storage::{MemoryBanStorage, MemoryStorage};
use std::sync::Arc;
use std::time::Duration;
fn create_valid_test_config() -> FlowControlConfig {
FlowControlConfig {
version: "0.1.0".to_string(),
global: crate::config::types::GlobalConfig::default(),
rules: vec![crate::config::types::Rule {
id: "test_rule".to_string(),
name: "Test Rule".to_string(),
priority: 100,
matchers: vec![crate::config::types::Matcher::User {
user_ids: vec!["*".to_string()],
}],
limiters: vec![crate::config::types::LimiterConfig::TokenBucket {
capacity: 100,
refill_rate: 10,
}],
action: crate::config::types::ActionConfig {
on_exceed: crate::config::types::Action::Reject,
ban: None,
},
}],
}
}
#[cfg(feature = "circuit-breaker")]
#[tokio::test]
async fn test_builder_with_circuit_breaker() {
use crate::circuit::{CircuitBreaker, CircuitBreakerConfig};
let cb = Arc::new(CircuitBreaker::with_dependencies(
CircuitBreakerConfig::default(),
));
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_circuit_breaker(cb)
.build()
.await
.expect("build with circuit_breaker should succeed");
assert!(governor.health_check().await.is_ok());
}
#[cfg(feature = "audit-log")]
#[tokio::test]
async fn test_builder_with_audit_logger() {
use crate::logging::AuditLogger;
let logger = Arc::new(AuditLogger::default().await);
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_audit_logger(logger)
.build()
.await
.expect("build with audit_logger should succeed");
assert!(governor.health_check().await.is_ok());
}
#[cfg(feature = "prometheus")]
#[tokio::test]
async fn test_builder_with_metrics() {
use crate::telemetry::Metrics;
let metrics = Arc::new(Metrics::new());
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_metrics(metrics.clone())
.build()
.await
.expect("build with metrics should succeed");
assert_eq!(metrics.requests_total.get(), 0.0);
let mut ctx = RequestContext::default();
ctx.user_id = Some("test_user".to_string());
let _ = governor.check(&ctx).await;
assert!(
metrics.requests_total.get() >= 1.0,
"metrics.requests_total should be recorded after check()"
);
}
#[cfg(feature = "prometheus")]
#[tokio::test]
async fn test_per_rule_metrics_wiring() {
use crate::telemetry::Metrics;
let metrics = Arc::new(Metrics::new());
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_metrics(metrics.clone())
.build()
.await
.expect("build should succeed");
let ctx = RequestContext::new()
.with_header("X-User-Id", "u_per_rule")
.with_method("GET");
for _ in 0..3 {
let _ = governor.check(&ctx).await;
}
let output = metrics.gather();
assert!(
output.contains("flowguard_rule_checks_total"),
"dimension metrics must be registered in the registry: {output}"
);
assert!(
output.contains("rule=\"test_rule\""),
"there should be a test_rule label series: {output}"
);
assert!(
output.contains("outcome=\"allowed\""),
"there should be an allowed outcome series: {output}"
);
}
#[cfg(feature = "prometheus")]
#[tokio::test]
async fn test_negative_cache_hit_counts_per_rule_metrics() {
use crate::telemetry::Metrics;
let config = FlowControlConfig {
version: "0.1.0".to_string(),
global: crate::config::types::GlobalConfig::default(),
rules: vec![crate::config::types::Rule {
id: "test_rule".to_string(),
name: "Test Rule".to_string(),
priority: 100,
matchers: vec![crate::config::types::Matcher::User {
user_ids: vec!["*".to_string()],
}],
limiters: vec![crate::config::types::LimiterConfig::TokenBucket {
capacity: 1,
refill_rate: 1,
}],
action: crate::config::types::ActionConfig {
on_exceed: crate::config::types::Action::Reject,
ban: None,
},
}],
};
let metrics = Arc::new(Metrics::new());
let governor = Governor::builder()
.with_config(config)
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_metrics(metrics.clone())
.build()
.await
.expect("build should succeed");
let ctx = RequestContext::new()
.with_header("X-User-Id", "u_neg_cache")
.with_method("GET");
for _ in 0..3 {
let _ = governor.check(&ctx).await;
}
let output = metrics.gather();
assert!(
output
.contains("flowguard_rule_checks_total{outcome=\"rejected\",rule=\"test_rule\"} 2"),
"rejections from negative-cache hits must be counted into the per-rule series per the decision rule (1 on-chain + 1 cache hit): {output}"
);
assert_eq!(
metrics.negative_cache_hits_total.get(),
0,
"a rule-attributable hit must not be counted twice into negative_cache_hits_total"
);
}
#[cfg(feature = "prometheus")]
#[tokio::test]
async fn test_per_rule_metrics_disabled_opt_out() {
use crate::telemetry::Metrics;
let metrics = Arc::new(Metrics::new());
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_metrics(metrics.clone())
.with_per_rule_metrics(false)
.build()
.await
.expect("build should succeed");
let ctx = RequestContext::new()
.with_header("X-User-Id", "u_opt_out")
.with_method("GET");
let _ = governor.check(&ctx).await;
let output = metrics.gather();
assert!(
!output.contains("flowguard_rule_checks_total{"),
"there should be no rule dimension series after shutdown: {output}"
);
assert!(metrics.requests_total.get() >= 1.0);
}
#[cfg(feature = "prometheus")]
#[tokio::test]
async fn test_incremental_export_idempotent_and_exact() {
use crate::telemetry::Metrics;
let metrics = Arc::new(Metrics::new());
let governor = Governor::builder()
.with_config(FlowControlConfig {
version: "0.1.0".to_string(),
global: crate::config::types::GlobalConfig::default(),
rules: vec![crate::config::types::Rule {
id: "test_rule".to_string(),
name: "Test Rule".to_string(),
priority: 100,
matchers: vec![crate::config::types::Matcher::User {
user_ids: vec!["*".to_string()],
}],
limiters: vec![crate::config::types::LimiterConfig::TokenBucket {
capacity: 1,
refill_rate: 1,
}],
action: crate::config::types::ActionConfig {
on_exceed: crate::config::types::Action::Reject,
ban: None,
},
}],
})
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_metrics(metrics.clone())
.build()
.await
.expect("build should succeed");
let chain = governor
.rule_chains
.read()
.await
.get("test_rule")
.map(|entry| entry.value().clone())
.expect("test_rule chain should exist");
let exported_total = |m: &Metrics| -> u64 {
use prometheus::core::Collector;
m.rule_limiter_rejections_total
.collect()
.iter()
.map(|family| {
family
.get_metric()
.iter()
.map(|metric| metric.get_counter().value() as u64)
.sum::<u64>()
})
.sum()
};
let _ = chain.check().await.unwrap();
assert_eq!(
chain.stats_sync().rejected_count,
0,
"the first check must exhaust the tokens and allow"
);
chain.check().await.unwrap();
chain.check().await.unwrap();
let actual = chain.stats_sync().node_rejections[0].1;
governor.export_chain_rejections(&chain, "test_rule", &metrics);
assert_eq!(
exported_total(&metrics),
actual,
"the first export must exactly equal the true reject count"
);
assert_eq!(
chain.exported_rejections_version(),
chain.node_rejections_version(),
"the exported ledger version must catch up with the chain version"
);
governor.export_chain_rejections(&chain, "test_rule", &metrics);
assert_eq!(
exported_total(&metrics),
actual,
"repeated exports of the same version must not double count"
);
chain.check().await.unwrap();
let actual = chain.stats_sync().node_rejections[0].1;
governor.export_chain_rejections(&chain, "test_rule", &metrics);
assert_eq!(
exported_total(&metrics),
actual,
"the next export must exactly cover the new rejections"
);
}
#[cfg(feature = "prometheus")]
#[tokio::test]
async fn test_concurrent_exports_never_double_count() {
use crate::telemetry::Metrics;
let metrics = Arc::new(Metrics::new());
let governor = Arc::new(
Governor::builder()
.with_config(FlowControlConfig {
version: "0.1.0".to_string(),
global: crate::config::types::GlobalConfig::default(),
rules: vec![crate::config::types::Rule {
id: "test_rule".to_string(),
name: "Test Rule".to_string(),
priority: 100,
matchers: vec![crate::config::types::Matcher::User {
user_ids: vec!["*".to_string()],
}],
limiters: vec![crate::config::types::LimiterConfig::TokenBucket {
capacity: 1,
refill_rate: 1,
}],
action: crate::config::types::ActionConfig {
on_exceed: crate::config::types::Action::Reject,
ban: None,
},
}],
})
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_metrics(metrics.clone())
.build()
.await
.expect("build should succeed"),
);
let chain = governor
.rule_chains
.read()
.await
.get("test_rule")
.map(|entry| entry.value().clone())
.expect("test_rule chain should exist");
let _ = chain.check().await.unwrap();
let tasks = (0..8)
.map(|_| {
let governor = governor.clone();
let chain = chain.clone();
let metrics = metrics.clone();
tokio::spawn(async move {
for _ in 0..50 {
let _ = chain.check().await;
governor.export_chain_rejections(&chain, "test_rule", &metrics);
}
})
})
.collect::<Vec<_>>();
for task in tasks {
task.await.unwrap();
}
let actual: u64 = chain
.stats_sync()
.node_rejections
.iter()
.map(|(_, count)| count)
.sum();
let exported: u64 = {
use prometheus::core::Collector;
metrics
.rule_limiter_rejections_total
.collect()
.iter()
.map(|family| {
family
.get_metric()
.iter()
.map(|metric| metric.get_counter().value() as u64)
.sum::<u64>()
})
.sum()
};
assert_eq!(
exported, actual,
"concurrent exports must not double count or miss"
);
}
#[cfg(all(test, feature = "prometheus", feature = "telemetry", feature = "otlp"))]
#[tokio::test]
async fn test_negative_cache_hit_without_rule_id_records_outcome() {
use crate::l1_cache::CacheableDecision;
use crate::matchers::Identifier;
use crate::telemetry::Metrics;
use crate::telemetry::otlp::{InMemoryTransport, OtlpSpanExporter};
let transport = Arc::new(InMemoryTransport::new());
let exporter = OtlpSpanExporter::new(
"limiteron-test",
transport.clone(),
"http://mock-collector:4318/v1/traces",
);
let (sink, worker) = OtlpSpanExporter::spawn_worker(exporter, 64);
tokio::spawn(worker.run());
let metrics = Arc::new(Metrics::new());
let governor = Governor::builder()
.with_config(FlowControlConfig {
version: "0.1.0".to_string(),
global: crate::config::types::GlobalConfig::default(),
rules: vec![crate::config::types::Rule {
id: "test_rule".to_string(),
name: "Test Rule".to_string(),
priority: 100,
matchers: vec![crate::config::types::Matcher::User {
user_ids: vec!["*".to_string()],
}],
limiters: vec![crate::config::types::LimiterConfig::TokenBucket {
capacity: 100,
refill_rate: 10,
}],
action: crate::config::types::ActionConfig {
on_exceed: crate::config::types::Action::Reject,
ban: None,
},
}],
})
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_metrics(metrics.clone())
.with_tracer(Arc::new(Tracer::with_otlp_sink(true, Some(Arc::new(sink)))))
.build()
.await
.expect("build should succeed");
let ctx = RequestContext::new()
.with_header("X-User-Id", "u_noattr")
.with_method("GET");
let matched_rules: Vec<_> = governor
.rule_matcher
.read()
.await
.match_all(&ctx)
.into_iter()
.cloned()
.collect();
assert!(!matched_rules.is_empty(), "the rule should match");
let identifier = Identifier::UserId("u_noattr".to_string());
let cache_key = governor.build_cache_key_multi(&identifier, &matched_rules);
governor
.l1_cache
.set(cache_key, CacheableDecision::rejected("cached deny"))
.await
.expect("seed cache should succeed");
let decision = governor.check(&ctx).await.expect("check should succeed");
assert!(
matches!(decision, Decision::Rejected(_)),
"should hit the cached rejection: {decision:?}"
);
assert_eq!(
metrics.negative_cache_hits_total.get(),
1,
"an unattributable hit should fall back to an independent count"
);
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
let span = loop {
let requests = transport.requests();
if let Some(request) = requests.first() {
break request.1["resourceSpans"][0]["scopeSpans"][0]["spans"][0].clone();
}
assert!(
tokio::time::Instant::now() < deadline,
"no OTLP export received within 5s"
);
tokio::time::sleep(Duration::from_millis(10)).await;
};
let attribute = |key: &str| -> Option<String> {
span["attributes"]
.as_array()?
.iter()
.find(|a| a["key"] == key)
.map(|a| {
a["value"]["stringValue"]
.as_str()
.unwrap_or_default()
.to_string()
})
};
assert_eq!(
attribute("rule.id"),
None,
"an unattributable hit must not fabricate rule.id"
);
assert_eq!(
attribute("rule.outcome").as_deref(),
Some("rejected"),
"an unattributable hit must keep the result signal"
);
}
#[cfg(all(feature = "prometheus", feature = "fallback"))]
#[tokio::test]
async fn test_degraded_metrics_wiring() {
use crate::fallback::FallbackManager;
use crate::telemetry::Metrics;
use oxcache::Cache;
let metrics = Arc::new(Metrics::new());
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_metrics(metrics.clone())
.build()
.await
.expect("build should succeed");
let ctx = RequestContext::new()
.with_header("X-User-Id", "u_degraded")
.with_method("GET");
let _ = governor.check(&ctx).await;
assert_eq!(
metrics.degraded_checks_total.get(),
0,
"success on the main path must not count as degraded"
);
let cache: Cache<String, String> = Cache::builder()
.capacity(10000)
.ttl(Duration::from_secs(60))
.build()
.await
.unwrap();
let metrics_fb = Arc::new(Metrics::new());
let governor_fb = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_metrics(metrics_fb.clone())
.with_fallback_manager(Arc::new(FallbackManager::new(Arc::new(cache))))
.build()
.await
.expect("build with fallback should succeed");
let ctx_no_id = RequestContext::new().with_method("GET");
let _ = governor_fb.check(&ctx_no_id).await;
assert_eq!(
metrics_fb.degraded_checks_total.get(),
0,
"a config-class error triggering degradation must not count into degraded_checks_total"
);
}
#[cfg(all(feature = "prometheus", feature = "fallback"))]
#[tokio::test]
async fn test_degraded_error_classification() {
use crate::error::StorageError;
use crate::telemetry::Metrics;
let metrics = Arc::new(Metrics::new());
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_metrics(metrics.clone())
.build()
.await
.expect("build should succeed");
governor.note_degraded_check(&Err(LimiteronError::StorageError(
StorageError::ConnectionError("redis down".to_string()),
)));
assert_eq!(
metrics.degraded_checks_total.get(),
1,
"a storage-class error must count as degraded"
);
governor.note_degraded_check(&Err(LimiteronError::ConfigError(
"no identifier".to_string(),
)));
governor.note_degraded_check(&Ok(Decision::allowed_default()));
assert_eq!(
metrics.degraded_checks_total.get(),
1,
"config-class errors and successful results must not count as degraded"
);
}
#[cfg(feature = "prometheus")]
#[tokio::test]
async fn test_per_rule_limiter_rejections_reset_on_hot_reload() {
use crate::telemetry::Metrics;
let make_config = || FlowControlConfig {
version: "0.1.0".to_string(),
global: crate::config::types::GlobalConfig::default(),
rules: vec![crate::config::types::Rule {
id: "test_rule".to_string(),
name: "Test Rule".to_string(),
priority: 100,
matchers: vec![crate::config::types::Matcher::User {
user_ids: vec!["*".to_string()],
}],
limiters: vec![crate::config::types::LimiterConfig::TokenBucket {
capacity: 1,
refill_rate: 1,
}],
action: crate::config::types::ActionConfig {
on_exceed: crate::config::types::Action::Reject,
ban: None,
},
}],
};
let metrics = Arc::new(Metrics::new());
let governor = Governor::builder()
.with_config(make_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_metrics(metrics.clone())
.build()
.await
.expect("build should succeed");
let ctx = RequestContext::new()
.with_header("X-User-Id", "u_hot_reload")
.with_method("GET");
for _ in 0..3 {
let _ = governor.check(&ctx).await;
}
let counter = metrics
.rule_limiter_rejections_total
.with_label_values(&["test_rule", "test_rule_limiter_0"]);
assert_eq!(
counter.get(),
1,
"1 on-chain rejection from the old chain must export as 1"
);
governor
.apply_config(make_config())
.await
.expect("hot reload should succeed");
for _ in 0..2 {
let _ = governor.check(&ctx).await;
}
assert_eq!(
counter.get(),
2,
"after a hot swap, the new chain's rejections must export immediately (the snapshot is cleared with the swap)"
);
}
#[cfg(feature = "telemetry")]
#[tokio::test]
async fn test_builder_with_tracer() {
use crate::telemetry::Tracer;
let tracer = Arc::new(Tracer::new(true));
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_tracer(tracer.clone())
.build()
.await
.expect("build with tracer should succeed");
assert!(tracer.is_enabled());
let mut ctx = RequestContext::default();
ctx.user_id = Some("test_user".to_string());
let _ = governor.check(&ctx).await;
}
#[cfg(feature = "prometheus")]
#[tokio::test]
async fn test_with_storage_forwards_metrics() {
use crate::telemetry::Metrics;
let metrics = Arc::new(Metrics::new());
#[cfg(feature = "telemetry")]
let tracer: Option<Arc<crate::telemetry::Tracer>> = None;
let governor = Governor::with_storage(
create_valid_test_config(),
Arc::new(MemoryStorage::new()),
Arc::new(MemoryBanStorage::new()),
Some(metrics.clone()),
#[cfg(feature = "telemetry")]
tracer,
)
.await
.expect("with_storage should succeed");
let mut ctx = RequestContext::default();
ctx.user_id = Some("test_user".to_string());
let _ = governor.check(&ctx).await;
assert!(
metrics.requests_total.get() >= 1.0,
"with_storage should forward metrics to Governor (was silently dropped before fix)"
);
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_builder_with_fallback_manager() {
use crate::fallback::FallbackManager;
use oxcache::Cache;
let cache: Cache<String, String> = Cache::builder()
.capacity(10000)
.ttl(Duration::from_secs(60))
.build()
.await
.unwrap();
let fm = Arc::new(FallbackManager::new(Arc::new(cache)));
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_fallback_manager(fm)
.build()
.await
.expect("build with fallback_manager should succeed");
let stats = governor.stats().await;
assert_eq!(stats.total_requests, 0);
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_check_with_fallback_manager() {
use crate::fallback::FallbackManager;
use oxcache::Cache;
let cache: Cache<String, String> = Cache::builder()
.capacity(10000)
.ttl(Duration::from_secs(60))
.build()
.await
.unwrap();
let fm = Arc::new(FallbackManager::new(Arc::new(cache)));
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_fallback_manager(fm)
.build()
.await
.expect("build should succeed");
let mut ctx = RequestContext::default();
ctx.client_ip = Some("10.0.0.100".to_string());
let result = governor.check(&ctx).await;
assert!(
result.is_ok(),
"check with fallback should succeed: {:?}",
result
);
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_check_with_fallback_no_identifier() {
use crate::fallback::FallbackManager;
use oxcache::Cache;
let cache: Cache<String, String> = Cache::builder()
.capacity(10000)
.ttl(Duration::from_secs(60))
.build()
.await
.unwrap();
let fm = Arc::new(FallbackManager::new(Arc::new(cache)));
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_fallback_manager(fm)
.build()
.await
.expect("build should succeed");
let ctx = RequestContext::default();
let result = governor.check(&ctx).await;
assert!(result.is_err());
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_check_with_fallback_l1_cache_disabled() {
use crate::fallback::FallbackManager;
use oxcache::Cache;
let cache: Cache<String, String> = Cache::builder()
.capacity(10000)
.ttl(Duration::from_secs(60))
.build()
.await
.unwrap();
let fm = Arc::new(FallbackManager::new(Arc::new(cache)));
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_l1_cache_enabled(false)
.with_fallback_manager(fm)
.build()
.await
.expect("build should succeed");
let ctx = RequestContext::default();
let result = governor.check(&ctx).await;
assert!(result.is_err());
}
#[cfg(feature = "parallel-checker")]
#[tokio::test]
async fn test_check_parallel_ban_checker_path() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let mut ctx = RequestContext::default();
ctx.client_ip = Some("10.0.0.200".to_string());
let result = governor.check(&ctx).await;
assert!(
result.is_ok(),
"check with parallel ban checker should succeed: {:?}",
result
);
}
#[cfg(feature = "parallel-checker")]
#[tokio::test]
async fn test_check_resource_parallel_unbanned() {
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let result = governor.check_resource_parallel("test_user_123").await;
match result {
Ok(Decision::Allowed(_)) => {}
other => panic!("Expected Allowed for unbanned user, got: {:?}", other),
}
}
#[cfg(feature = "audit-log")]
#[tokio::test]
async fn test_audit_logger_set_and_get() {
use crate::logging::AuditLogger;
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let initial = governor.audit_logger().await;
assert!(initial.is_none());
let logger = Arc::new(AuditLogger::default().await);
governor.set_audit_logger(logger.clone()).await;
let retrieved = governor.audit_logger().await;
assert!(retrieved.is_some(), "audit_logger should be Some after set");
}
fn create_ip_test_config() -> FlowControlConfig {
FlowControlConfig {
version: "0.1.0".to_string(),
global: crate::config::types::GlobalConfig::default(),
rules: vec![crate::config::types::Rule {
id: "ip_rule".to_string(),
name: "IP Rule".to_string(),
priority: 100,
matchers: vec![crate::config::types::Matcher::Ip {
ip_ranges: vec!["0.0.0.0/0".to_string()],
}],
limiters: vec![crate::config::types::LimiterConfig::TokenBucket {
capacity: 100,
refill_rate: 10,
}],
action: crate::config::types::ActionConfig {
on_exceed: crate::config::types::Action::Reject,
ban: None,
},
}],
}
}
fn create_user_test_config(user_ids: Vec<String>) -> FlowControlConfig {
FlowControlConfig {
version: "0.1.0".to_string(),
global: crate::config::types::GlobalConfig::default(),
rules: vec![crate::config::types::Rule {
id: "user_rule".to_string(),
name: "User Rule".to_string(),
priority: 100,
matchers: vec![crate::config::types::Matcher::User { user_ids }],
limiters: vec![crate::config::types::LimiterConfig::TokenBucket {
capacity: 100,
refill_rate: 10,
}],
action: crate::config::types::ActionConfig {
on_exceed: crate::config::types::Action::Reject,
ban: None,
},
}],
}
}
fn create_ip_request_context(ip: &str) -> RequestContext {
let mut ctx = RequestContext::default();
ctx.client_ip = Some(ip.to_string());
ctx
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_check_l1_cache_only_disabled_returns_err() {
let governor = Governor::builder()
.with_config(create_ip_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_l1_cache_enabled(false)
.build()
.await
.expect("build should succeed");
let ctx = create_ip_request_context("10.0.0.1");
let result = governor.check_l1_cache_only(&ctx).await;
assert!(result.is_err());
match result.unwrap_err() {
LimiteronError::LimitError(msg) => {
assert!(
msg.contains("L1 缓存未启用") || msg.contains("L1 cache not enabled"),
"unexpected: {msg}"
);
}
other => panic!("expected LimitError, got: {:?}", other),
}
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_check_l1_cache_only_identifier_extraction_failure() {
let governor = Governor::builder()
.with_config(create_ip_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let ctx = RequestContext::default();
let result = governor.check_l1_cache_only(&ctx).await;
assert!(result.is_err());
match result.unwrap_err() {
LimiteronError::ConfigError(msg) => {
assert!(
msg.contains("Failed to extract identifier"),
"unexpected: {}",
msg
);
}
other => panic!("expected ConfigError, got: {:?}", other),
}
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_check_l1_cache_only_no_matched_rules() {
let governor = Governor::builder()
.with_config(create_user_test_config(vec!["specific_user".to_string()]))
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let ctx = create_ip_request_context("10.0.0.1");
let result = governor.check_l1_cache_only(&ctx).await;
assert!(result.is_ok(), "expected Ok, got: {:?}", result.err());
match result.unwrap() {
Decision::Allowed(_) => {}
other => panic!("expected Allowed, got: {:?}", other),
}
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_check_l1_cache_only_cache_miss_no_island() {
let governor = Governor::builder()
.with_config(create_ip_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
assert!(!governor.l1_cache.is_island_mode());
let ctx = create_ip_request_context("10.0.0.1");
let result = governor.check_l1_cache_only(&ctx).await;
assert!(result.is_err());
match result.unwrap_err() {
LimiteronError::LimitError(msg) => {
assert!(
msg.contains("降级缓存未命中") || msg.contains("degraded cache miss"),
"unexpected: {msg}"
);
}
other => panic!("expected LimitError, got: {:?}", other),
}
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_check_l1_cache_only_island_allow_all() {
let governor = Governor::builder()
.with_config(create_ip_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let island_config = IslandModeConfig::new(IslandFallbackStrategy::AllowAll);
governor.l1_cache.enable_island_mode(island_config);
assert!(governor.l1_cache.is_island_mode());
let ctx = create_ip_request_context("10.0.0.1");
let result = governor.check_l1_cache_only(&ctx).await;
assert!(result.is_ok(), "expected Ok, got: {:?}", result.err());
match result.unwrap() {
Decision::Allowed(_) => {}
other => panic!("expected Allowed, got: {:?}", other),
}
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_check_l1_cache_only_island_reject_all() {
let governor = Governor::builder()
.with_config(create_ip_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let island_config = IslandModeConfig::new(IslandFallbackStrategy::RejectAll);
governor.l1_cache.enable_island_mode(island_config);
let ctx = create_ip_request_context("10.0.0.1");
let result = governor.check_l1_cache_only(&ctx).await;
assert!(result.is_err());
match result.unwrap_err() {
LimiteronError::LimitError(msg) => {
assert!(
msg.contains("拒绝请求") || msg.contains("reject"),
"unexpected: {msg}"
);
}
other => panic!("expected LimitError, got: {:?}", other),
}
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_check_l1_cache_only_island_local_decision() {
let governor = Governor::builder()
.with_config(create_ip_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let island_config = IslandModeConfig::new(IslandFallbackStrategy::LocalDecision);
governor.l1_cache.enable_island_mode(island_config);
let ctx = create_ip_request_context("10.0.0.1");
let result = governor.check_l1_cache_only(&ctx).await;
assert!(result.is_ok(), "expected Ok, got: {:?}", result.err());
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_check_l1_cache_only_island_conservative_quota() {
let governor = Governor::builder()
.with_config(create_ip_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let island_config = IslandModeConfig::new(IslandFallbackStrategy::ConservativeQuota {
max_requests: 10,
window_secs: 60,
});
governor.l1_cache.enable_island_mode(island_config);
let ctx = create_ip_request_context("10.0.0.1");
let result = governor.check_l1_cache_only(&ctx).await;
assert!(result.is_ok(), "expected Ok, got: {:?}", result.err());
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_check_l1_cache_only_cache_hit() {
let governor = Governor::builder()
.with_config(create_ip_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let ctx = create_ip_request_context("10.0.0.1");
let _ = governor.check(&ctx).await;
let miss = governor.check_l1_cache_only(&ctx).await;
assert!(miss.is_err(), "expected miss, got: {:?}", miss);
use crate::l1_cache::CacheableDecision;
let cache_key = "rl:ip:10.0.0.1:ip_rule".to_string();
let rejected = CacheableDecision::rejected("rate limit exceeded");
governor
.l1_cache
.set(cache_key, rejected)
.await
.expect("cache set should succeed");
let result = governor.check_l1_cache_only(&ctx).await;
assert!(result.is_ok(), "expected Ok, got: {:?}", result.err());
match result.unwrap() {
Decision::Rejected(_) => {}
other => panic!("expected Rejected, got: {:?}", other),
}
}
#[cfg(feature = "fallback")]
#[tokio::test]
async fn test_check_l1_cache_only_cached_banned_updates_stats() {
use crate::error::BanInfo;
use crate::l1_cache::CacheableDecision;
let governor = Governor::builder()
.with_config(create_ip_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let ctx = create_ip_request_context("10.0.0.50");
let cache_key = "rl:ip:10.0.0.50:ip_rule".to_string();
let ban_info = BanInfo::new("test ban".to_string(), chrono::Utc::now(), 1);
let banned_decision = CacheableDecision::banned(&ban_info);
governor
.l1_cache
.set(cache_key, banned_decision)
.await
.expect("cache set should succeed");
let result = governor.check_l1_cache_only(&ctx).await;
assert!(result.is_ok(), "expected Ok, got: {:?}", result.err());
match result.unwrap() {
Decision::Banned(_) => {}
other => panic!("expected Banned, got: {:?}", other),
}
}
#[tokio::test]
async fn test_update_stats_for_decision_err_branch() {
let governor = Governor::builder()
.with_config(create_ip_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.build()
.await
.expect("build should succeed");
let err_result: Result<Decision, LimiteronError> =
Err(LimiteronError::LimitError("test error".to_string()));
governor.update_stats_for_decision(&err_result);
}
#[cfg(feature = "metrics")]
#[tokio::test]
async fn test_metrics_feature_records_check_metrics() {
use crate::telemetry::Metrics;
let metrics = Arc::new(Metrics::new());
let governor = Governor::builder()
.with_config(create_valid_test_config())
.with_storage(Arc::new(MemoryStorage::new()))
.with_ban_storage(Arc::new(MemoryBanStorage::new()))
.with_metrics(metrics.clone())
.build()
.await
.expect("build with metrics should succeed");
let ctx = RequestContext {
user_id: Some("metrics_test_user".to_string()),
ip: Some("127.0.0.1".to_string()),
path: "/test".to_string(),
method: "GET".to_string(),
..Default::default()
};
let _ = governor.check(&ctx).await;
assert!(
metrics.requests_total.get() > 0.0,
"metrics should record at least one check via requests_total counter"
);
}
#[cfg(all(test, feature = "fallback"))]
mod conservative_quota_counter_tests {
use super::*;
#[test]
fn test_conservative_quota_counter_bounds_window() {
let mut c = ConservativeQuotaCounter {
window_start_secs: 100,
count: 0,
};
for i in 0..5u64 {
assert!(
c.try_acquire(5, 60, 100 + i),
"request {} should be allowed",
i + 1
);
}
assert!(
!c.try_acquire(5, 60, 105),
"over-limit within the window must be rejected"
);
assert!(
c.try_acquire(5, 60, 160),
"the window flip should reset and allow"
);
assert!(
c.try_acquire(5, 60, 120),
"a clock turn-back does not reset the window, capacity stays bounded (2/5)"
);
}
}
}
#[cfg(all(test, feature = "multi-tenant"))]
mod governor_tenant_ban_tests {
use super::{Governor, RequestContext};
use super::{create_valid_test_config_for_tenant, limiter_kind};
use crate::config::{LimiterConfig, QuotaType};
use crate::storage::{MemoryBanStorage, MemoryStorage};
use std::sync::Arc;
async fn make_governor() -> Governor {
let storage: Arc<dyn crate::storage::Storage> = Arc::new(MemoryStorage::new());
let ban_storage: Arc<dyn crate::storage::BanStorage> = Arc::new(MemoryBanStorage::new());
Governor::builder()
.with_config(create_valid_test_config_for_tenant())
.with_storage(storage)
.with_ban_storage(ban_storage)
.build()
.await
.expect("Governor build should succeed")
}
#[test]
fn test_limiter_kind_covers_every_variant() {
let cases: Vec<(LimiterConfig, &str)> = vec![
(
LimiterConfig::TokenBucket {
capacity: 10,
refill_rate: 1,
},
"token_bucket",
),
(
LimiterConfig::SlidingWindow {
window_size: "1m".into(),
max_requests: 5,
},
"sliding_window",
),
(
LimiterConfig::FixedWindow {
window_size: "1m".into(),
max_requests: 5,
},
"fixed_window",
),
(
LimiterConfig::Quota {
quota_type: QuotaType::Count,
limit: 5,
window: "1d".into(),
alert_threshold: None,
overdraft: None,
},
"quota",
),
(
LimiterConfig::Concurrency { max_concurrent: 2 },
"concurrency",
),
(
LimiterConfig::LeakyBucket {
capacity: 4,
leak_rate: 2,
},
"leaky_bucket",
),
(
LimiterConfig::SlidingWindowLog {
window_size: "1m".into(),
max_requests: 3,
},
"sliding_window_log",
),
(
LimiterConfig::PriorityQueue {
window_size: "1m".into(),
total_per_window: 10,
level_weights: vec![1],
default_priority: None,
},
"priority_queue",
),
(
LimiterConfig::AdmissionControl {
max_concurrent: 2,
max_per_second: 5,
},
"admission_control",
),
(
LimiterConfig::Custom {
name: "mylim".into(),
config: serde_json::Value::Null,
},
"custom:mylim",
),
];
let mut kinds: Vec<String> = Vec::new();
for (cfg, expected) in &cases {
let got = limiter_kind(cfg);
assert_eq!(got, *expected, "limiter_kind mapping drift");
kinds.push(got);
}
let mut uniq = kinds.clone();
uniq.sort();
uniq.dedup();
assert_eq!(
uniq.len(),
cases.len(),
"limiter type labels must be mutually distinct"
);
}
#[cfg(feature = "ban-manager")]
#[tokio::test]
async fn test_ban_identifier_namespace_scoping() {
use super::LimiteronError;
use crate::matchers::Identifier;
use crate::tenant::Namespace;
use std::time::Duration;
let governor = make_governor().await;
let ns = Namespace::new("t1", "prod");
let user = Identifier::UserId("u1".to_string());
let ctx = RequestContext::default();
assert_eq!(
governor.tenant_scoped_identifier(&ctx, user.clone()).key(),
user.key()
);
let qualified = ns.qualify_key("u1");
assert_ne!(qualified, "u1");
assert!(
qualified.contains("t1"),
"the qualified key must carry the tenant id: {qualified}"
);
governor
.ban_identifier_for_namespace(&ns, &user, "abuse", None)
.await
.expect("UserId is bannable");
governor
.ban_identifier_for_namespace(
&Namespace::default(),
&user,
"global",
Some(Duration::from_secs(60)),
)
.await
.expect("default namespace path");
for id in [
Identifier::ApiKey("k".to_string()),
Identifier::DeviceId("d".to_string()),
] {
let err = governor
.ban_identifier_for_namespace(&ns, &id, "x", None)
.await
.expect_err("unsupported identifier must fail loudly");
assert!(
matches!(err, LimiteronError::ValidationError(_)),
"should be ValidationError: {err:?}"
);
}
assert!(
governor
.is_identifier_banned(&ctx, &Identifier::UserId("nobody".to_string()))
.await
.expect("ban lookup")
.is_none()
);
}
#[tokio::test]
async fn test_resolve_tenant_without_resolver_is_none() {
let governor = make_governor().await;
assert!(
governor
.resolve_tenant(&RequestContext::default())
.is_none()
);
}
}
#[cfg(all(test, feature = "multi-tenant"))]
fn create_valid_test_config_for_tenant() -> crate::config::FlowControlConfig {
use crate::config::types::{Action, ActionConfig, GlobalConfig, Matcher, Rule};
crate::config::FlowControlConfig {
version: "0.1.0".to_string(),
global: GlobalConfig::default(),
rules: vec![Rule {
id: "tenant_ban_rule".to_string(),
name: "Tenant Ban Rule".to_string(),
priority: 100,
matchers: vec![Matcher::User {
user_ids: vec!["*".to_string()],
}],
limiters: vec![limiter_cfg_token_bucket()],
action: ActionConfig {
on_exceed: Action::Reject,
ban: None,
},
}],
}
}
#[cfg(all(test, feature = "multi-tenant"))]
fn limiter_cfg_token_bucket() -> crate::config::LimiterConfig {
crate::config::LimiterConfig::TokenBucket {
capacity: 100,
refill_rate: 10,
}
}