diff --git a/plugins/rust/python-package/rate_limiter/README.md b/plugins/rust/python-package/rate_limiter/README.md index 825c3350..723adaf7 100644 --- a/plugins/rust/python-package/rate_limiter/README.md +++ b/plugins/rust/python-package/rate_limiter/README.md @@ -39,6 +39,11 @@ If any configured dimension is exceeded, the plugin returns a violation with HTT # Redis options (required when backend: redis) redis_url: "redis://redis:6379/0" redis_key_prefix: "rl" + + # Backend failure policy (default: "open" — fail-open) + # "closed" — return HTTP 503 BACKEND_UNAVAILABLE violation when the + # backend can't be reached (correctness over availability) + fail_mode: "open" ``` ### Configuration reference @@ -52,8 +57,13 @@ If any configured dimension is exceeded, the plugin returns a violation with HTT | `backend` | string | `"memory"` | `"memory"` or `"redis"` | | `redis_url` | string | `null` | Redis connection URL (required when `backend: redis`) | | `redis_key_prefix` | string | `"rl"` | Prefix for all Redis keys | +| `fail_mode` | string | `"open"` | Behaviour when the backend can't be reached: `"open"` allows the request through, `"closed"` blocks with a 503 `BACKEND_UNAVAILABLE` violation | + +**Rate string format:** `"/"` where unit is `s`/`sec`/`second`, `m`/`min`/`minute`, or `h`/`hr`/`hour`. Malformed strings raise `ValueError` at startup. Counts above `1_000_000` are rejected as a sanity ceiling — anything higher is almost certainly a misconfig or a denial-of-service vector against the memory backend. -**Rate string format:** `"/"` where unit is `s`/`sec`/`second`, `m`/`min`/`minute`, or `h`/`hr`/`hour`. Malformed strings raise `ValueError` at startup. +**Unknown config keys** (e.g. a typo like `redis_ur`) are logged at `WARN` at engine init alongside the accepted-key list, instead of being silently ignored. + +**Invalid `fail_mode` values** (e.g. `"clsoed"`) are logged at `WARN` and fall back to `"open"` so an operator's typo surfaces instead of silently disabling the hardening they asked for. **Omitting a dimension** (e.g. no `by_tenant`) means that dimension is unlimited — no counter is tracked for it. @@ -107,10 +117,26 @@ Each identity (user, tenant, tool) has a bucket that holds up to `count` tokens. - `token_bucket`: atomic Lua script — reads `{tokens, last_refill}` hash, refills proportionally, consumes 1 token, writes back — one round-trip, no race condition - All gateway instances share the same counter — the configured limit is the true cluster-wide limit - Requires `redis_url` to be set -- If Redis is unavailable, the plugin fails open — the request is allowed through without rate limiting. This is a deliberate design choice: an infrastructure failure must never block legitimate traffic. Operators should monitor for rate-limiter error logs and treat them as high-priority alerts +- **Backend failure policy** is governed by `fail_mode`: + - `"open"` (default) — the request is allowed through without rate limiting. Availability over correctness; an infrastructure failure must never block legitimate traffic. Operators should monitor for rate-limiter error logs and treat them as high-priority alerts. + - `"closed"` — the request is blocked with a `PluginViolation` (code `BACKEND_UNAVAILABLE`, HTTP 503, `Retry-After: 1`). Correctness over availability; pick this when a failed rate-limit check is less acceptable than a brief outage. **Multi-instance deployment (important):** The `memory` backend is local to a single gateway instance — rate limit counters are not shared across replicas. For multi-instance deployments (e.g., behind nginx or on OpenShift with multiple gateway pods), always use `backend: redis` to ensure rate limits are enforced correctly across all instances. +### Tenant-scoped Redis key layout + +When the plugin context carries a `tenant_id`, every dimension key is prefixed with it so counters are isolated per tenant: + +``` +rl:{tenant_id}:user:{email}:{window_seconds} +rl:{tenant_id}:tenant:{tenant_id}:{window_seconds} +rl:{tenant_id}:tool:{tool_name}:{window_seconds} +``` + +When `tenant_id` is absent (single-tenant deployments), the prefix is omitted and keys revert to the pre-tenant-scoping layout (`rl:user:{email}:{window}`), so single-tenant behaviour is unchanged. + +**Upgrade note:** the first deploy of the tenant-scoping change causes counters under `rl:user:*` / `rl:tool:*` to be orphaned while new writes land at `rl:{tenant}:user:*`. Counters effectively reset once for all in-flight windows — non-event for typical second/minute windows. + ## Examples ### Single-instance (default config) @@ -170,6 +196,15 @@ config: In `permissive` mode the plugin records violations and emits `X-RateLimit-*` headers but does not block requests. Useful for baselining traffic before switching to `enforce`. +## Lifecycle + +The plugin participates in the plugin manager's lifecycle contract: + +- `async def initialize(self)` — invoked once when the plugin manager constructs the plugin. Logs one `INFO` record naming the active backend (`memory` / `redis`). +- `async def shutdown(self)` — invoked when the plugin manager tears the plugin down (runtime disable, re-instantiation after a config change). Releases backend-held resources — specifically, drops the Rust core's cached Redis multiplexed connection and the SCRIPT LOAD SHA cache. In-flight requests already hold their own clones of the connection and remain valid; the cached reference is replaced on the next request. + +Without `shutdown`, the cached Redis connection would leak across plugin re-instantiation, producing connection churn on the server. + ## Limitations | Limitation | Severity | Status | diff --git a/plugins/rust/python-package/rate_limiter/cpex_rate_limiter/rate_limiter.py b/plugins/rust/python-package/rate_limiter/cpex_rate_limiter/rate_limiter.py index 1cd89c2a..7a71b353 100644 --- a/plugins/rust/python-package/rate_limiter/cpex_rate_limiter/rate_limiter.py +++ b/plugins/rust/python-package/rate_limiter/cpex_rate_limiter/rate_limiter.py @@ -3,6 +3,8 @@ from __future__ import annotations +import logging + try: from mcpgateway.plugins.framework import Plugin, PromptPrehookResult, ToolPreInvokeResult except ModuleNotFoundError: @@ -45,6 +47,7 @@ class RateLimiterConfig: "backend", "redis_url", "redis_key_prefix", + "fail_mode", ) def __init__(self, **overrides) -> None: @@ -54,6 +57,9 @@ def __init__(self, **overrides) -> None: setattr(self, field, config.get(field)) +_logger = logging.getLogger(__name__) + + class RateLimiterPlugin(Plugin): """Gateway-facing Plugin subclass that delegates behavior to Rust.""" @@ -61,13 +67,36 @@ def __init__(self, config) -> None: super().__init__(config) self._core = RateLimiterPluginCore(config.config or {}) + async def initialize(self) -> None: + """Lifecycle hook: called once when the plugin manager constructs us.""" + cfg = self.config.config or {} + backend = cfg.get("backend", "memory") + _logger.info("rate limiter initialized: backend=%s", backend) + + async def shutdown(self) -> None: + """Lifecycle hook: release Rust-held resources (e.g. Redis connection). + + The plugin manager calls this on disable and on re-instantiation. + Without it, the cached Redis connection leaks until the plugin + instance is garbage-collected. + """ + try: + self._core.shutdown() + except Exception: + _logger.exception("rate limiter shutdown: core.shutdown() raised") + async def prompt_pre_fetch(self, payload, context): + # The Rust core handles fail_mode policy internally (open vs closed) + # and logs backend errors via log_exception. The except here is a + # final safety net for the unlikely case that a non-backend bug in + # the core escapes as a Python exception. try: result = self._core.prompt_pre_fetch(payload, context) if hasattr(result, "__await__"): return await result return result except Exception: + _logger.warning("rate limiter prompt_pre_fetch: unexpected core error; allowing request", exc_info=True) return PromptPrehookResult() async def tool_pre_invoke(self, payload, context): @@ -77,6 +106,7 @@ async def tool_pre_invoke(self, payload, context): return await result return result except Exception: + _logger.warning("rate limiter tool_pre_invoke: unexpected core error; allowing request", exc_info=True) return ToolPreInvokeResult() diff --git a/plugins/rust/python-package/rate_limiter/cpex_rate_limiter/rate_limiter_rust/__init__.pyi b/plugins/rust/python-package/rate_limiter/cpex_rate_limiter/rate_limiter_rust/__init__.pyi index b3dfb8cb..55f7e79f 100644 --- a/plugins/rust/python-package/rate_limiter/cpex_rate_limiter/rate_limiter_rust/__init__.pyi +++ b/plugins/rust/python-package/rate_limiter/cpex_rate_limiter/rate_limiter_rust/__init__.pyi @@ -102,8 +102,15 @@ class RateLimiterEngine: - `backend`: `"memory"` (default) or `"redis"` - `redis_url`: required when `backend = "redis"` - `redis_key_prefix`: key namespace prefix (default `"rl"`) + - `fail_mode`: `"open"` (default) or `"closed"` — handled by the + plugin shim, but accepted here so it doesn't trip the unknown-key + warning below. + + Any other key in the dict is logged at WARN so misspellings (e.g. + `redis_ur` instead of `redis_url`) surface visibly instead of being + silently ignored. """ - def check(self, user: builtins.str, tenant: typing.Optional[builtins.str], tool: builtins.str, now_unix: builtins.int, include_retry_after: builtins.bool) -> tuple[builtins.bool, dict, dict]: + def check(self, user: builtins.str, tenant: typing.Optional[builtins.str], tool: builtins.str, now_unix: builtins.int, include_retry_after: builtins.bool, context_prefix: typing.Optional[builtins.str]) -> tuple[builtins.bool, dict, dict]: r""" High-level check: builds dimension keys internally, evaluates, and returns pre-built Python dicts for headers and metadata. @@ -113,13 +120,18 @@ class RateLimiterEngine: Returns `(allowed, headers_dict, meta_dict)`. + When `context_prefix` is provided (e.g. a team/tenant ID), it is + prepended to every dimension key so that separate plugin instances + for different tenants use isolated Redis counters instead of sharing + a single key namespace. + **Note:** The Redis backend arm uses `block_on()` on a dedicated Tokio runtime, which would deadlock if called from within a Tokio context. The Python wrapper routes Redis to `check_async()` instead; this sync path is intended for the memory backend. The `debug_assert` below guards against accidental misuse. """ - def check_async(self, user: builtins.str, tenant: typing.Optional[builtins.str], tool: builtins.str, now_unix: builtins.int, include_retry_after: builtins.bool) -> typing.Any: + def check_async(self, user: builtins.str, tenant: typing.Optional[builtins.str], tool: builtins.str, now_unix: builtins.int, include_retry_after: builtins.bool, context_prefix: typing.Optional[builtins.str]) -> typing.Any: r""" Async variant of `check()` for Redis-backed deployments. @@ -129,6 +141,12 @@ class RateLimiterEngine: @typing.final class RateLimiterPluginCore: def __new__(cls, config: dict) -> RateLimiterPluginCore: ... + def shutdown(self) -> None: + r""" + Release backend-held resources (e.g. the cached Redis multiplexed + connection). Called by the Python shim's `shutdown()` when the plugin + framework tears the plugin down. + """ def prompt_pre_fetch(self, payload: typing.Any, context: typing.Any) -> typing.Any: ... def tool_pre_invoke(self, payload: typing.Any, context: typing.Any) -> typing.Any: ... diff --git a/plugins/rust/python-package/rate_limiter/src/config.rs b/plugins/rust/python-package/rate_limiter/src/config.rs index 5eb7b01c..6b990b6f 100644 --- a/plugins/rust/python-package/rate_limiter/src/config.rs +++ b/plugins/rust/python-package/rate_limiter/src/config.rs @@ -31,6 +31,10 @@ pub enum ConfigError { InvalidRateString(String), #[error("rate count must be > 0, got {0}")] ZeroCount(u64), + #[error( + "rate count {0} exceeds the sanity ceiling of {MAX_RATE_COUNT}; if you really need this, raise the limit deliberately" + )] + CountAboveCeiling(u64), #[error("{field}: {message}")] FieldError { field: String, message: String }, #[error( @@ -39,6 +43,11 @@ pub enum ConfigError { InvalidAlgorithm(String), } +/// Upper bound for any single rate-limit count. Above this, configurations +/// are almost certainly typos or denial-of-service vectors against the +/// memory backend (which allocates per dimension key). +pub const MAX_RATE_COUNT: u64 = 1_000_000; + /// Parse a rate string like `"30/m"`, `"100/s"`, `"1000/h"`. /// /// Accepted units (case-insensitive): `s`, `sec`, `second`, `m`, `min`, @@ -57,6 +66,9 @@ pub fn parse_rate(s: &str) -> Result { if count == 0 { return Err(ConfigError::ZeroCount(count)); } + if count > MAX_RATE_COUNT { + return Err(ConfigError::CountAboveCeiling(count)); + } let window_secs: u64 = match unit_str.trim().to_ascii_lowercase().as_str() { "s" | "sec" | "second" => 1, @@ -198,6 +210,25 @@ mod tests { assert!(parse_rate("0/s").is_err()); } + // --- Bounds: reject absurd counts so misconfig can't allocate huge + // windows, overflow internal math, or starve other memory. + // See config_hardening_bounds_rationale in the README. + + #[test] + fn parse_rate_rejects_count_above_upper_bound() { + assert!( + parse_rate("99999999/h").is_err(), + "parse_rate must reject counts above the sanity ceiling" + ); + } + + #[test] + fn parse_rate_accepts_reasonable_large_count() { + // 100000/h is plausibly a real quota — must still parse. + let r = parse_rate("100000/h").unwrap(); + assert_eq!(r.count, 100_000); + } + // --- Algorithm::from_str --- #[test] diff --git a/plugins/rust/python-package/rate_limiter/src/engine.rs b/plugins/rust/python-package/rate_limiter/src/engine.rs index 56025949..a418338d 100644 --- a/plugins/rust/python-package/rate_limiter/src/engine.rs +++ b/plugins/rust/python-package/rate_limiter/src/engine.rs @@ -77,6 +77,16 @@ impl RateLimiterEngine { pub fn uses_async_backend(&self) -> bool { matches!(self.backend, EngineBackend::Redis(_)) } + + /// Release backend-held resources. For Redis, this drops the cached + /// multiplexed connection so the server can close the socket; in-flight + /// requests that already cloned the handle remain valid. Memory backend + /// has no external resources and is a no-op. + pub fn shutdown(&self) { + if let EngineBackend::Redis(redis) = &self.backend { + redis.shutdown(); + } + } } #[gen_stub_pymethods] @@ -91,8 +101,17 @@ impl RateLimiterEngine { /// - `backend`: `"memory"` (default) or `"redis"` /// - `redis_url`: required when `backend = "redis"` /// - `redis_key_prefix`: key namespace prefix (default `"rl"`) + /// - `fail_mode`: `"open"` (default) or `"closed"` — handled by the + /// plugin shim, but accepted here so it doesn't trip the unknown-key + /// warning below. + /// + /// Any other key in the dict is logged at WARN so misspellings (e.g. + /// `redis_ur` instead of `redis_url`) surface visibly instead of being + /// silently ignored. #[new] pub fn new(config: &Bound<'_, PyDict>) -> PyResult { + warn_on_unknown_config_keys(config); + let by_user: Option = config.get_item("by_user")?.and_then(|v| v.extract().ok()); let by_tenant: Option = config.get_item("by_tenant")?.and_then(|v| v.extract().ok()); @@ -161,11 +180,17 @@ impl RateLimiterEngine { /// /// Returns `(allowed, headers_dict, meta_dict)`. /// + /// When `context_prefix` is provided (e.g. a team/tenant ID), it is + /// prepended to every dimension key so that separate plugin instances + /// for different tenants use isolated Redis counters instead of sharing + /// a single key namespace. + /// /// **Note:** The Redis backend arm uses `block_on()` on a dedicated Tokio /// runtime, which would deadlock if called from within a Tokio context. /// The Python wrapper routes Redis to `check_async()` instead; this sync /// path is intended for the memory backend. The `debug_assert` below /// guards against accidental misuse. + #[allow(clippy::too_many_arguments)] pub fn check<'py>( &self, py: Python<'py>, @@ -174,13 +199,14 @@ impl RateLimiterEngine { tool: &str, now_unix: i64, include_retry_after: bool, + context_prefix: Option<&str>, ) -> PyResult<(bool, Bound<'py, PyDict>, Bound<'py, PyDict>)> { if matches!(self.backend, EngineBackend::Redis(_)) { return Err(pyo3::exceptions::PyRuntimeError::new_err( "check() must not be called with the Redis backend — use check_async() instead", )); } - let checks = self.build_checks(user, tenant, tool); + let checks = self.build_checks(user, tenant, tool, context_prefix); if checks.is_empty() { let headers = PyDict::new(py); let meta = PyDict::new(py); @@ -216,13 +242,14 @@ impl RateLimiterEngine { let eval = EvalResult::from_dims(&dim_results); let headers = build_headers_dict(py, &eval, include_retry_after)?; - let meta = build_meta_dict(py, &eval, now_unix)?; + let meta = build_meta_dict(py, &eval, now_unix, Some(user), tenant)?; Ok((eval.allowed, headers, meta)) } /// Async variant of `check()` for Redis-backed deployments. /// /// Returns an awaitable that resolves to `(allowed, headers_dict, meta_dict)`. + #[allow(clippy::too_many_arguments)] pub fn check_async<'py>( &self, py: Python<'py>, @@ -231,8 +258,9 @@ impl RateLimiterEngine { tool: &str, now_unix: i64, include_retry_after: bool, + context_prefix: Option<&str>, ) -> PyResult> { - let checks = self.build_checks(user, tenant, tool); + let checks = self.build_checks(user, tenant, tool, context_prefix); if checks.is_empty() { return future_into_py(py, async move { Python::attach(|py| -> PyResult> { @@ -255,6 +283,10 @@ impl RateLimiterEngine { let backend = self.backend.clone(); let algorithm = self.config.algorithm; let clock = Arc::clone(&self.clock); + // Capture identity as owned for the async move so build_meta_dict + // can surface tenant_id/user_id in violation details (G7). + let user_owned = user.to_string(); + let tenant_owned = tenant.map(|s| s.to_string()); future_into_py(py, async move { let dim_results: Vec = @@ -284,7 +316,13 @@ impl RateLimiterEngine { let eval = EvalResult::from_dims(&dim_results); Python::attach(|py| -> PyResult> { let headers = build_headers_dict(py, &eval, include_retry_after)?; - let meta = build_meta_dict(py, &eval, now_unix)?; + let meta = build_meta_dict( + py, + &eval, + now_unix, + Some(user_owned.as_str()), + tenant_owned.as_deref(), + )?; let tup = pyo3::types::PyTuple::new( py, [ @@ -306,28 +344,83 @@ impl RateLimiterEngine { impl RateLimiterEngine { /// Build dimension checks from engine config. /// Mirrors Python `_build_rust_checks()` but runs in Rust. + /// + /// When `context_prefix` is `Some("team_a")` the keys become + /// `team_a:user:alice`, `team_a:tenant:team_a`, `team_a:tool:search` + /// instead of the unprefixed versions. This prevents counter collisions + /// across plugin-manager contexts (teams) that share the same Redis. fn build_checks( &self, user: &str, tenant: Option<&str>, tool: &str, + context_prefix: Option<&str>, ) -> Vec<(String, u64, u64)> { let mut checks = Vec::with_capacity(3); + let pfx = context_prefix.unwrap_or(""); if let Some(ref rl) = self.config.by_user { - checks.push((format!("user:{}", user), rl.count, rl.window_nanos)); + let key = if pfx.is_empty() { + format!("user:{}", user) + } else { + format!("{}:user:{}", pfx, user) + }; + checks.push((key, rl.count, rl.window_nanos)); } if let (Some(t), Some(rl)) = (tenant, &self.config.by_tenant) { - checks.push((format!("tenant:{}", t), rl.count, rl.window_nanos)); + let key = if pfx.is_empty() { + format!("tenant:{}", t) + } else { + format!("{}:tenant:{}", pfx, t) + }; + checks.push((key, rl.count, rl.window_nanos)); } // Tool names are already normalised (lowercase) in EngineConfig at init time. // The caller passes the already-lowercased tool name from Python. if let Some(rl) = self.config.by_tool.get(tool) { - checks.push((format!("tool:{}", tool), rl.count, rl.window_nanos)); + let key = if pfx.is_empty() { + format!("tool:{}", tool) + } else { + format!("{}:tool:{}", pfx, tool) + }; + checks.push((key, rl.count, rl.window_nanos)); } checks } } +/// Emit a single WARN-level log listing config keys we don't recognise. +/// Catches misspellings (e.g. ``redis_ur`` instead of ``redis_url``) that +/// would otherwise silently default and surprise the operator at runtime. +fn warn_on_unknown_config_keys(config: &Bound<'_, PyDict>) { + const KNOWN: &[&str] = &[ + "by_user", + "by_tenant", + "by_tool", + "algorithm", + "backend", + "redis_url", + "redis_key_prefix", + "fail_mode", + ]; + let mut unknown: Vec = Vec::new(); + for (key, _) in config.iter() { + let Ok(name) = key.extract::() else { + continue; + }; + if !KNOWN.contains(&name.as_str()) { + unknown.push(name); + } + } + if !unknown.is_empty() { + unknown.sort(); + warn!( + "rate limiter: unknown config key(s): {}; expected one of: {}", + unknown.join(", "), + KNOWN.join(", "), + ); + } +} + /// Build HTTP rate-limit headers dict — mirrors Python `_make_headers()`. fn build_headers_dict<'py>( py: Python<'py>, @@ -349,10 +442,18 @@ fn build_headers_dict<'py>( } /// Build metadata dict — mirrors Python `_rust_to_plugin_meta()`. +/// +/// When a request is blocked, identity (`user_id` / `tenant_id`) is +/// surfaced in the meta dict so the resulting `PluginViolation.details` +/// carries enough context for downstream debugging (G7). Identity is +/// intentionally NOT attached on allowed responses to avoid widening +/// identity exposure to every metadata consumer on the hot path. fn build_meta_dict<'py>( py: Python<'py>, eval: &EvalResult, now_unix: i64, + user: Option<&str>, + tenant: Option<&str>, ) -> PyResult> { let meta = PyDict::new(py); let reset_in = eval @@ -361,6 +462,14 @@ fn build_meta_dict<'py>( meta.set_item("limited", true)?; meta.set_item("remaining", eval.remaining)?; meta.set_item("reset_in", reset_in)?; + if !eval.allowed { + if let Some(u) = user.filter(|s| !s.is_empty()) { + meta.set_item("user_id", u)?; + } + if let Some(t) = tenant.filter(|s| !s.is_empty()) { + meta.set_item("tenant_id", t)?; + } + } let has_violated = !eval.violated_dimensions.is_empty(); let has_allowed = !eval.allowed_dimensions.is_empty(); @@ -559,4 +668,75 @@ mod tests { let result = engine.evaluate_many(checks(), 1_000_000).unwrap(); assert!(!result.allowed); // tenant exhausted → blocked } + + // --- Context prefix: tenant-scoped key isolation --- + + #[test] + fn build_checks_without_prefix_produces_unprefixed_keys() { + let (engine, _handle) = engine_with_fake_clock(Some("10/s"), Algorithm::FixedWindow); + let checks = engine.build_checks("alice", Some("acme"), "search", None); + let keys: Vec<&str> = checks.iter().map(|(k, _, _)| k.as_str()).collect(); + assert!(keys.contains(&"user:alice")); + assert!(keys.contains(&"tool:search")); + } + + #[test] + fn build_checks_with_prefix_prepends_to_all_keys() { + let (engine, _handle) = engine_with_fake_clock(Some("10/s"), Algorithm::FixedWindow); + let checks = engine.build_checks("alice", Some("acme"), "search", Some("team_a")); + let keys: Vec<&str> = checks.iter().map(|(k, _, _)| k.as_str()).collect(); + assert!(keys.contains(&"team_a:user:alice"), "keys: {:?}", keys); + assert!(keys.contains(&"team_a:tool:search"), "keys: {:?}", keys); + } + + #[test] + fn build_checks_with_prefix_includes_tenant_dimension() { + init_python(); + let (clock, _handle) = FakeClock::new(1_000_000); + let cfg = EngineConfig { + by_user: Some(crate::config::parse_rate("10/s").unwrap()), + by_tenant: Some(crate::config::parse_rate("100/s").unwrap()), + by_tool: HashMap::new(), + algorithm: Algorithm::FixedWindow, + }; + let engine = RateLimiterEngine::new_with_clock(cfg, Arc::new(clock)); + let checks = engine.build_checks("alice", Some("acme"), "search", Some("team_a")); + let keys: Vec<&str> = checks.iter().map(|(k, _, _)| k.as_str()).collect(); + assert!(keys.contains(&"team_a:user:alice"), "keys: {:?}", keys); + assert!(keys.contains(&"team_a:tenant:acme"), "keys: {:?}", keys); + } + + #[test] + fn different_prefixes_produce_isolated_counters() { + let (engine, _handle) = engine_with_fake_clock(Some("2/s"), Algorithm::FixedWindow); + // Exhaust limit for team_a + let checks_a = || engine.build_checks("alice", None, "search", Some("team_a")); + let _ = engine.evaluate_many(checks_a(), 1_000_000).unwrap(); + let _ = engine.evaluate_many(checks_a(), 1_000_000).unwrap(); + let result_a = engine.evaluate_many(checks_a(), 1_000_000).unwrap(); + assert!( + !result_a.allowed, + "team_a should be blocked after 2 requests" + ); + + // team_b should still be allowed — different prefix, different counters + let checks_b = || engine.build_checks("alice", None, "search", Some("team_b")); + let result_b = engine.evaluate_many(checks_b(), 1_000_000).unwrap(); + assert!( + result_b.allowed, + "team_b should be allowed — isolated counter" + ); + } + + #[test] + fn empty_prefix_matches_no_prefix_behavior() { + let (engine, _handle) = engine_with_fake_clock(Some("10/s"), Algorithm::FixedWindow); + let checks_none = engine.build_checks("alice", None, "search", None); + let checks_empty = engine.build_checks("alice", None, "search", Some("")); + // Both should produce the same unprefixed keys + assert_eq!(checks_none.len(), checks_empty.len()); + for ((k1, _, _), (k2, _, _)) in checks_none.iter().zip(checks_empty.iter()) { + assert_eq!(k1, k2); + } + } } diff --git a/plugins/rust/python-package/rate_limiter/src/lib.rs b/plugins/rust/python-package/rate_limiter/src/lib.rs index 40da9503..67855275 100644 --- a/plugins/rust/python-package/rate_limiter/src/lib.rs +++ b/plugins/rust/python-package/rate_limiter/src/lib.rs @@ -31,6 +31,7 @@ fn compat_default_config(py: Python<'_>) -> PyResult> { defaults.set_item("backend", "memory")?; defaults.set_item("redis_url", py.None())?; defaults.set_item("redis_key_prefix", "rl")?; + defaults.set_item("fail_mode", "open")?; Ok(defaults.unbind()) } diff --git a/plugins/rust/python-package/rate_limiter/src/plugin.rs b/plugins/rust/python-package/rate_limiter/src/plugin.rs index bf968da0..87573cb3 100644 --- a/plugins/rust/python-package/rate_limiter/src/plugin.rs +++ b/plugins/rust/python-package/rate_limiter/src/plugin.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use cpex_framework_bridge::{build_framework_object, default_result}; +use log::warn; use pyo3::prelude::*; use pyo3::types::{PyAny, PyDict, PyModule, PyTuple}; use pyo3_async_runtimes::tokio::{future_into_py, into_future}; @@ -21,6 +22,10 @@ const LOGGER_NAME: &str = "cpex_rate_limiter.rate_limiter"; pub struct RateLimiterPluginCore { engine: Arc, use_async: bool, + /// When true, backend failures produce a BACKEND_UNAVAILABLE violation + /// (fail-closed) instead of the default allow result (fail-open). Read + /// from the `fail_mode` config key at init time. + fail_closed: bool, } #[gen_stub_pymethods] @@ -29,12 +34,21 @@ impl RateLimiterPluginCore { #[new] pub fn new(config: &Bound<'_, PyDict>) -> PyResult { let engine = Arc::new(RateLimiterEngine::new(config)?); + let fail_closed = parse_fail_mode(config)?; Ok(Self { use_async: engine.uses_async_backend(), engine, + fail_closed, }) } + /// Release backend-held resources (e.g. the cached Redis multiplexed + /// connection). Called by the Python shim's `shutdown()` when the plugin + /// framework tears the plugin down. + pub fn shutdown(&self) { + self.engine.shutdown(); + } + pub fn prompt_pre_fetch<'py>( &self, py: Python<'py>, @@ -47,8 +61,18 @@ impl RateLimiterPluginCore { .trim() .to_ascii_lowercase(); let (user, tenant) = extract_request_context(context)?; + // Use tenant_id as the context prefix so that each team's rate limit + // counters are isolated in Redis. Without this, all teams share keys. + let context_prefix = tenant.as_deref(); + let fail_closed = self.fail_closed; if !self.use_async { - return match evaluate_sync_request(&self.engine, &user, tenant.as_deref(), &prompt) { + return match evaluate_sync_request( + &self.engine, + &user, + tenant.as_deref(), + &prompt, + context_prefix, + ) { Ok((allowed, headers, meta)) => Ok(build_prehook_result( py, "PromptPrehookResult", @@ -58,18 +82,27 @@ impl RateLimiterPluginCore { )? .into_bound(py)), Err(_err) => { - log_exception( - py, - "RateLimiterPlugin.prompt_pre_fetch error; allowing request", - )?; - Ok(default_result(py, "PromptPrehookResult")?.into_bound(py)) + log_exception(py, error_log_message("prompt_pre_fetch", fail_closed))?; + Ok( + backend_error_result(py, "PromptPrehookResult", fail_closed)? + .into_bound(py), + ) } }; } let engine = Arc::clone(&self.engine); + let context_prefix_owned = context_prefix.map(|s| s.to_string()); future_into_py(py, async move { - match evaluate_async_request(&engine, &user, tenant.as_deref(), &prompt).await { + match evaluate_async_request( + &engine, + &user, + tenant.as_deref(), + &prompt, + context_prefix_owned.as_deref(), + ) + .await + { Ok((allowed, headers, meta)) => Python::attach(|py| { build_prehook_result( py, @@ -80,11 +113,8 @@ impl RateLimiterPluginCore { ) }), Err(_err) => Python::attach(|py| { - log_exception( - py, - "RateLimiterPlugin.prompt_pre_fetch error; allowing request", - )?; - default_result(py, "PromptPrehookResult") + log_exception(py, error_log_message("prompt_pre_fetch", fail_closed))?; + backend_error_result(py, "PromptPrehookResult", fail_closed) }), } }) @@ -102,8 +132,16 @@ impl RateLimiterPluginCore { .trim() .to_ascii_lowercase(); let (user, tenant) = extract_request_context(context)?; + let context_prefix = tenant.as_deref(); + let fail_closed = self.fail_closed; if !self.use_async { - return match evaluate_sync_request(&self.engine, &user, tenant.as_deref(), &tool) { + return match evaluate_sync_request( + &self.engine, + &user, + tenant.as_deref(), + &tool, + context_prefix, + ) { Ok((allowed, headers, meta)) => Ok(build_prehook_result( py, "ToolPreInvokeResult", @@ -113,18 +151,27 @@ impl RateLimiterPluginCore { )? .into_bound(py)), Err(_err) => { - log_exception( - py, - "RateLimiterPlugin.tool_pre_invoke error; allowing request", - )?; - Ok(default_result(py, "ToolPreInvokeResult")?.into_bound(py)) + log_exception(py, error_log_message("tool_pre_invoke", fail_closed))?; + Ok( + backend_error_result(py, "ToolPreInvokeResult", fail_closed)? + .into_bound(py), + ) } }; } let engine = Arc::clone(&self.engine); + let context_prefix_owned = context_prefix.map(|s| s.to_string()); future_into_py(py, async move { - match evaluate_async_request(&engine, &user, tenant.as_deref(), &tool).await { + match evaluate_async_request( + &engine, + &user, + tenant.as_deref(), + &tool, + context_prefix_owned.as_deref(), + ) + .await + { Ok((allowed, headers, meta)) => Python::attach(|py| { build_prehook_result( py, @@ -135,22 +182,99 @@ impl RateLimiterPluginCore { ) }), Err(_err) => Python::attach(|py| { - log_exception( - py, - "RateLimiterPlugin.tool_pre_invoke error; allowing request", - )?; - default_result(py, "ToolPreInvokeResult") + log_exception(py, error_log_message("tool_pre_invoke", fail_closed))?; + backend_error_result(py, "ToolPreInvokeResult", fail_closed) }), } }) } } +/// Parse the ``fail_mode`` config key into a ``fail_closed`` bool. +/// +/// Accepted values (case-insensitive, trimmed): ``"open"`` and ``"closed"``. +/// An absent key, an explicit ``None``, or an empty string all resolve to +/// fail-open (the safe default for backwards compatibility). Any other +/// value — including typos like ``"clsoed"`` and non-string types — is +/// logged at WARN and falls through to fail-open rather than silently +/// disabling the fail-closed hardening the operator asked for. +fn parse_fail_mode(config: &Bound<'_, PyDict>) -> PyResult { + let item = match config.get_item("fail_mode")? { + None => return Ok(false), + Some(v) if v.is_none() => return Ok(false), + Some(v) => v, + }; + + match item.extract::() { + Ok(s) => { + let trimmed = s.trim(); + if trimmed.is_empty() { + return Ok(false); + } + match trimmed.to_ascii_lowercase().as_str() { + "open" => Ok(false), + "closed" => Ok(true), + _ => { + warn!( + "rate limiter: unknown fail_mode={:?}; expected \"open\" or \"closed\"; defaulting to \"open\"", + trimmed, + ); + Ok(false) + } + } + } + Err(_) => { + // Non-string value (dict, int, list, ...). Stringify for the log + // so the operator sees what was actually passed. + let repr = item + .repr() + .map(|r| r.to_string()) + .unwrap_or_else(|_| "".into()); + warn!( + "rate limiter: fail_mode must be a string (\"open\" or \"closed\"); got {}; defaulting to \"open\"", + repr, + ); + Ok(false) + } + } +} + +fn error_log_message(hook: &str, fail_closed: bool) -> &'static str { + match (hook, fail_closed) { + ("prompt_pre_fetch", false) => { + "RateLimiterPlugin.prompt_pre_fetch error; allowing request (fail_mode=open)" + } + ("prompt_pre_fetch", true) => { + "RateLimiterPlugin.prompt_pre_fetch error; blocking request (fail_mode=closed)" + } + ("tool_pre_invoke", false) => { + "RateLimiterPlugin.tool_pre_invoke error; allowing request (fail_mode=open)" + } + ("tool_pre_invoke", true) => { + "RateLimiterPlugin.tool_pre_invoke error; blocking request (fail_mode=closed)" + } + _ => "RateLimiterPlugin hook error", + } +} + +fn backend_error_result( + py: Python<'_>, + class_name: &str, + fail_closed: bool, +) -> PyResult> { + if fail_closed { + build_backend_unavailable_result(py, class_name) + } else { + default_result(py, class_name) + } +} + fn evaluate_sync_request( engine: &RateLimiterEngine, user: &str, tenant: Option<&str>, tool_or_prompt: &str, + context_prefix: Option<&str>, ) -> PyResult<(bool, Py, Py)> { let now_unix = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -158,8 +282,15 @@ fn evaluate_sync_request( .as_secs() as i64; Python::attach(|py| { - let (allowed, headers, meta) = - engine.check(py, user, tenant, tool_or_prompt, now_unix, true)?; + let (allowed, headers, meta) = engine.check( + py, + user, + tenant, + tool_or_prompt, + now_unix, + true, + context_prefix, + )?; Ok((allowed, headers.unbind(), meta.unbind())) }) } @@ -169,6 +300,7 @@ async fn evaluate_async_request( user: &str, tenant: Option<&str>, tool_or_prompt: &str, + context_prefix: Option<&str>, ) -> PyResult<(bool, Py, Py)> { let now_unix = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -176,7 +308,15 @@ async fn evaluate_async_request( .as_secs() as i64; let awaitable = Python::attach(|py| { engine - .check_async(py, user, tenant, tool_or_prompt, now_unix, true) + .check_async( + py, + user, + tenant, + tool_or_prompt, + now_unix, + true, + context_prefix, + ) .map(|awaitable| awaitable.unbind()) })?; await_async_tuple(awaitable).await @@ -241,6 +381,56 @@ fn build_prehook_result( ) } +/// Build a prehook/tool result that carries a BACKEND_UNAVAILABLE violation. +/// Used when `fail_mode=closed` and the Rust engine could not evaluate the +/// rate limit (e.g. Redis unreachable). The result blocks the request with +/// HTTP 503 and a Retry-After hint. +fn build_backend_unavailable_result(py: Python<'_>, class_name: &str) -> PyResult> { + let headers = PyDict::new(py); + headers.set_item("Retry-After", "1")?; + let violation = build_framework_object( + py, + "PluginViolation", + [ + ( + "reason", + "Rate limiter backend unavailable" + .into_pyobject(py)? + .into_any() + .unbind(), + ), + ( + "description", + "Rate limiter backend unavailable; failing closed per fail_mode=closed" + .into_pyobject(py)? + .into_any() + .unbind(), + ), + ( + "code", + "BACKEND_UNAVAILABLE".into_pyobject(py)?.into_any().unbind(), + ), + ("details", PyDict::new(py).into_any().unbind()), + ( + "http_status_code", + 503i32.into_pyobject(py)?.into_any().unbind(), + ), + ("http_headers", headers.into_any().unbind()), + ], + )?; + build_framework_object( + py, + class_name, + [ + ( + "continue_processing", + false.into_pyobject(py)?.to_owned().into_any().unbind(), + ), + ("violation", violation), + ], + ) +} + fn build_violation( py: Python<'_>, meta: &Bound<'_, PyDict>, diff --git a/plugins/rust/python-package/rate_limiter/src/redis_backend.rs b/plugins/rust/python-package/rate_limiter/src/redis_backend.rs index f94d3119..9f968281 100644 --- a/plugins/rust/python-package/rate_limiter/src/redis_backend.rs +++ b/plugins/rust/python-package/rate_limiter/src/redis_backend.rs @@ -233,6 +233,13 @@ impl RedisRateLimiter { *self.script_sha.lock() = None; } + /// Drop the cached multiplexed connection and script SHA so the server + /// can close the socket. In-flight requests hold their own clones and + /// remain valid. Called from `RateLimiterEngine::shutdown()`. + pub fn shutdown(&self) { + self.reset_connection(); + } + /// Return the batch Lua script for the active algorithm. fn batch_script(&self) -> &'static str { match self.algorithm { diff --git a/plugins/rust/python-package/rate_limiter/tests/integration/test_rate_limiter.py b/plugins/rust/python-package/rate_limiter/tests/integration/test_rate_limiter.py index 6e5a827a..148b4909 100644 --- a/plugins/rust/python-package/rate_limiter/tests/integration/test_rate_limiter.py +++ b/plugins/rust/python-package/rate_limiter/tests/integration/test_rate_limiter.py @@ -1041,6 +1041,35 @@ async def _flush_redis(redis_url: str) -> None: await client.aclose() +async def _keys_in_redis(redis_url: str, pattern: str) -> list[str]: + """Return all Redis keys matching the pattern. Used by isolation tests.""" + # Third-Party + import redis.asyncio as aioredis # noqa: PLC0415 + + client = aioredis.from_url(redis_url, decode_responses=True) + try: + return sorted(await client.keys(pattern)) + finally: + await client.aclose() + + +async def _count_redis_clients(redis_url: str) -> int: + """Return the number of connected clients reported by Redis CLIENT LIST. + + Includes the monitoring client we open here, so callers use this as a + relative measure (delta before/after an action), not an absolute count. + """ + # Third-Party + import redis.asyncio as aioredis # noqa: PLC0415 + + client = aioredis.from_url(redis_url, decode_responses=True) + try: + listing = await client.execute_command("CLIENT", "LIST") + return len([line for line in listing.splitlines() if line.strip()]) + finally: + await client.aclose() + + class TestRedisBackendIntegration: """End-to-end integration tests for the Redis backend. @@ -1256,3 +1285,296 @@ async def test_redis_token_bucket_refills_over_time(self, redis_url_for_integrat result = await plugin.tool_pre_invoke(payload, ctx) assert result.violation is None, "After tokens refill over time, token_bucket Redis backend must allow requests again" + + +class TestRedisTenantIsolation: + """Tenant-scoped Redis key isolation (G2). + + Without tenant-scoped keys, the same user hitting two different teams + through the same Redis would share a single ``rl:user:alice:*`` counter — + Team A's strict limit would poison Team B. The fix prefixes dimension + keys with the tenant id so counters are isolated per tenant. + """ + + @pytest.mark.asyncio + async def test_same_user_different_tenants_isolated_in_redis(self, redis_url_for_integration): + """Same user, two tenant ids, one Redis → independent per-user counters.""" + await _flush_redis(redis_url_for_integration) + + plugin = _make_redis_plugin(redis_url_for_integration, limit="3/s") + payload = ToolPreInvokePayload(name="tool", arguments={}) + ctx_team_a = PluginContext(global_context=GlobalContext(request_id="r1", user="alice", tenant_id="team_a")) + ctx_team_b = PluginContext(global_context=GlobalContext(request_id="r2", user="alice", tenant_id="team_b")) + + # Exhaust alice's limit under team_a. + for _ in range(3): + result = await plugin.tool_pre_invoke(payload, ctx_team_a) + assert result.violation is None, "team_a: requests within limit must be allowed" + blocked = await plugin.tool_pre_invoke(payload, ctx_team_a) + assert blocked.violation is not None, "team_a: 4th request must be blocked" + + # Same user under team_b must not be affected. + result = await plugin.tool_pre_invoke(payload, ctx_team_b) + assert result.violation is None, "team_b: alice must have an independent counter — team_a's limit must not bleed across tenants" + + # Key-shape proof: the two tenants occupy disjoint key namespaces in Redis. + keys_team_a = await _keys_in_redis(redis_url_for_integration, "rl:team_a:user:alice:*") + keys_team_b = await _keys_in_redis(redis_url_for_integration, "rl:team_b:user:alice:*") + assert keys_team_a, f"expected team_a-prefixed key, got none. keys_team_a={keys_team_a}" + assert keys_team_b, f"expected team_b-prefixed key, got none. keys_team_b={keys_team_b}" + assert set(keys_team_a).isdisjoint(set(keys_team_b)), "team_a and team_b key sets must not overlap" + + +class TestRedisLifecycle: + """Plugin-manager lifecycle compliance (G3, G10). + + When the plugin framework disables a plugin it calls ``await plugin.shutdown()``; + when it re-enables, a fresh instance is constructed. A compliant Redis-backed + plugin must release its cached connection on shutdown so the old instance + doesn't leak sockets while the new one opens its own. + """ + + @pytest.mark.asyncio + async def test_initialize_logs_backend(self, redis_url_for_integration, caplog): + """plugin.initialize() emits a log record identifying the active backend.""" + import logging # noqa: PLC0415 + + plugin = _make_redis_plugin(redis_url_for_integration, limit="5/s") + + with caplog.at_level(logging.INFO, logger="cpex_rate_limiter.rate_limiter"): + await plugin.initialize() + + matches = [r for r in caplog.records if "initialized" in r.getMessage() and "redis" in r.getMessage()] + assert matches, ( + "plugin.initialize() must log a record mentioning the backend so operators can confirm the plugin is live — " + f"captured records: {[r.getMessage() for r in caplog.records]}" + ) + + @pytest.mark.asyncio + async def test_shutdown_releases_redis_connection(self, redis_url_for_integration): + """After plugin.shutdown(), the Redis connection cached by the Rust core is dropped.""" + await _flush_redis(redis_url_for_integration) + + plugin = _make_redis_plugin(redis_url_for_integration, limit="5/s") + ctx = PluginContext(global_context=GlobalContext(request_id="r1", user="alice")) + payload = ToolPreInvokePayload(name="tool", arguments={}) + + # Warm the connection — the Rust core opens Redis lazily on first request. + await plugin.tool_pre_invoke(payload, ctx) + + clients_before = await _count_redis_clients(redis_url_for_integration) + + await plugin.shutdown() + + # Allow the TCP close to propagate to the server's client list. + await asyncio.sleep(0.2) + clients_after = await _count_redis_clients(redis_url_for_integration) + + assert clients_after < clients_before, ( + "plugin.shutdown() must release the Rust core's cached Redis connection — " + f"expected fewer clients after shutdown, got before={clients_before} after={clients_after}" + ) + + +def _make_redis_plugin_with_config(redis_url: str, extra_config: dict) -> RateLimiterPlugin: + """Redis-backed plugin with caller-supplied extra config keys (e.g. fail_mode).""" + base = { + "by_user": "3/s", + "backend": "redis", + "redis_url": redis_url, + "algorithm": "fixed_window", + } + base.update(extra_config) + return RateLimiterPlugin( + PluginConfig( + name="RateLimiter", + kind="cpex_rate_limiter.rate_limiter.RateLimiterPlugin", + hooks=["tool_pre_invoke"], + priority=100, + config=base, + ) + ) + + +# A Redis URL that parses fine but points at a port where nothing listens. +# Used by fail-mode tests to trigger a backend error deterministically +# without depending on Docker being slow or flaky. +_DEAD_REDIS_URL = "redis://127.0.0.1:1/15" + + +class TestRedisFailModeAndViolationContext: + """Fail-mode policy and violation context (G4, G7, G8, G14). + + When the Redis backend is unreachable, the plugin must either fail open + (default) or fail closed (opt-in via fail_mode="closed") — and in both + cases an operator-visible log record must describe what happened. When + the plugin blocks a request the violation must carry enough context + (tenant_id, user_id) for downstream debugging. + """ + + @pytest.mark.asyncio + async def test_redis_unreachable_default_fail_open_logs_warning(self, caplog): + """Unreachable backend + default fail_mode: request is allowed AND a WARNING is logged.""" + import logging # noqa: PLC0415 + + plugin = RateLimiterPlugin( + PluginConfig( + name="RateLimiter", + kind="cpex_rate_limiter.rate_limiter.RateLimiterPlugin", + hooks=["tool_pre_invoke"], + priority=100, + config={"by_user": "3/s", "backend": "redis", "redis_url": _DEAD_REDIS_URL}, + ) + ) + ctx = PluginContext(global_context=GlobalContext(request_id="r1", user="alice")) + payload = ToolPreInvokePayload(name="tool", arguments={}) + + with caplog.at_level(logging.WARNING, logger="cpex_rate_limiter.rate_limiter"): + result = await plugin.tool_pre_invoke(payload, ctx) + + assert result.violation is None, "default fail_mode=open: unreachable backend must not block" + warnings = [r for r in caplog.records if r.levelno >= logging.WARNING] + assert warnings, ( + "backend failures must be logged at WARNING or higher so operators notice silent fail-open — " + f"captured records: {[(r.levelname, r.getMessage()) for r in caplog.records]}" + ) + + @pytest.mark.asyncio + async def test_redis_unreachable_fail_mode_closed_blocks(self): + """fail_mode=closed: unreachable backend blocks the request with BACKEND_UNAVAILABLE.""" + plugin = _make_redis_plugin_with_config(_DEAD_REDIS_URL, {"fail_mode": "closed"}) + ctx = PluginContext(global_context=GlobalContext(request_id="r1", user="alice")) + payload = ToolPreInvokePayload(name="tool", arguments={}) + + result = await plugin.tool_pre_invoke(payload, ctx) + + assert result.violation is not None, "fail_mode=closed must block when the backend is unreachable" + assert result.violation.code == "BACKEND_UNAVAILABLE", ( + f"expected code=BACKEND_UNAVAILABLE, got {result.violation.code!r}" + ) + + @pytest.mark.asyncio + async def test_violation_details_includes_tenant_and_user(self, redis_url_for_integration): + """Blocked requests carry tenant_id + user_id in violation.details for downstream debugging.""" + await _flush_redis(redis_url_for_integration) + + plugin = _make_redis_plugin(redis_url_for_integration, limit="2/s") + ctx = PluginContext( + global_context=GlobalContext(request_id="r1", user="alice@example.com", tenant_id="team_a") + ) + payload = ToolPreInvokePayload(name="tool", arguments={}) + + await plugin.tool_pre_invoke(payload, ctx) + await plugin.tool_pre_invoke(payload, ctx) + blocked = await plugin.tool_pre_invoke(payload, ctx) + + assert blocked.violation is not None, "3rd request must be blocked" + details = blocked.violation.details or {} + assert details.get("tenant_id") == "team_a", ( + f"violation.details must carry tenant_id; got details={details!r}" + ) + assert details.get("user_id") == "alice@example.com", ( + f"violation.details must carry user_id; got details={details!r}" + ) + + @pytest.mark.asyncio + async def test_invalid_fail_mode_logs_warning_and_defaults_open(self, redis_url_for_integration, caplog): + """Typos like 'clsoed' must WARN and default to fail-open, not silently disable. + + Before the strict parser, anything that didn't case-insensitively + equal 'closed' silently became fail-open — including obvious + typos — which undermined the point of having a fail-closed knob. + Operators get no signal that their configured policy isn't the + one being applied. After the fix, unknown values are rejected + with a WARN log so the typo surfaces at init time. + """ + import logging # noqa: PLC0415 + + with caplog.at_level(logging.WARNING): + plugin = _make_redis_plugin_with_config(redis_url_for_integration, {"fail_mode": "clsoed"}) + + warnings = [ + r for r in caplog.records + if r.levelno >= logging.WARNING + and "fail_mode" in r.getMessage() + and "clsoed" in r.getMessage() + ] + assert warnings, ( + "invalid fail_mode must emit a WARN naming both the field and the bad value; " + f"captured records: {[(r.levelname, r.getMessage()) for r in caplog.records]}" + ) + + # Behaviour: fall back to fail-open (not fail-closed, not crash). + ctx = PluginContext(global_context=GlobalContext(request_id="r1", user="alice")) + payload = ToolPreInvokePayload(name="tool", arguments={}) + # Deliberately hit the happy path — limit is 3/s, one request stays well under. + result = await plugin.tool_pre_invoke(payload, ctx) + assert result.violation is None, "invalid fail_mode must default to fail-open, not block" + + @pytest.mark.asyncio + async def test_allowed_request_metadata_does_not_carry_identity(self, redis_url_for_integration): + """Allowed responses must NOT carry user_id / tenant_id in metadata. + + Identity fields belong on the block path (violation.details) so + operators can see who triggered a 429. Exposing them on every + allowed response widens identity leak surface to downstream + consumers that inspect plugin metadata and bloats the hot path + with data no caller needs. + """ + await _flush_redis(redis_url_for_integration) + + plugin = _make_redis_plugin(redis_url_for_integration, limit="5/s") + ctx = PluginContext( + global_context=GlobalContext(request_id="r1", user="alice@example.com", tenant_id="team_a") + ) + payload = ToolPreInvokePayload(name="tool", arguments={}) + + result = await plugin.tool_pre_invoke(payload, ctx) + + assert result.violation is None, "request under limit must be allowed" + metadata = result.metadata or {} + assert "user_id" not in metadata, ( + f"allowed response must not carry user_id in metadata; got metadata={metadata!r}" + ) + assert "tenant_id" not in metadata, ( + f"allowed response must not carry tenant_id in metadata; got metadata={metadata!r}" + ) + + +class TestConfigHardening: + """Config hardening (G13, G15). + + The plugin must reject obviously-misconfigured rate strings and warn + when the operator passes a key the engine doesn't understand (typos + surface visibly instead of silently being ignored). + """ + + @pytest.mark.asyncio + async def test_unknown_config_key_emits_warning(self, redis_url_for_integration, caplog): + """Misspelled config keys (e.g. 'redis_ur') log a WARNING so the operator notices.""" + import logging # noqa: PLC0415 + + with caplog.at_level(logging.WARNING): + RateLimiterPlugin( + PluginConfig( + name="RateLimiter", + kind="cpex_rate_limiter.rate_limiter.RateLimiterPlugin", + hooks=["tool_pre_invoke"], + priority=100, + config={ + "by_user": "3/s", + "backend": "redis", + "redis_url": redis_url_for_integration, + "redis_ur": "typo-key", # misspelled + }, + ) + ) + + warnings = [ + r for r in caplog.records + if r.levelno >= logging.WARNING and "redis_ur" in r.getMessage() + ] + assert warnings, ( + "engine must warn on unknown config keys so misspellings are visible — " + f"captured records: {[(r.levelname, r.getMessage()) for r in caplog.records]}" + ) diff --git a/plugins/rust/python-package/rate_limiter/tests/mcpgateway_mock/plugins/framework.py b/plugins/rust/python-package/rate_limiter/tests/mcpgateway_mock/plugins/framework.py index ef2b553e..463190ce 100644 --- a/plugins/rust/python-package/rate_limiter/tests/mcpgateway_mock/plugins/framework.py +++ b/plugins/rust/python-package/rate_limiter/tests/mcpgateway_mock/plugins/framework.py @@ -103,3 +103,9 @@ def priority(self) -> int: @property def mode(self) -> str: return self._config.mode + + async def initialize(self) -> None: + """Lifecycle hook. Mirrors mcpgateway.plugins.framework.base.Plugin.""" + + async def shutdown(self) -> None: + """Lifecycle hook. Mirrors mcpgateway.plugins.framework.base.Plugin.""" diff --git a/plugins/rust/python-package/rate_limiter/tests/test_config.py b/plugins/rust/python-package/rate_limiter/tests/test_config.py index 887cea0f..d22dd8c6 100644 --- a/plugins/rust/python-package/rate_limiter/tests/test_config.py +++ b/plugins/rust/python-package/rate_limiter/tests/test_config.py @@ -106,3 +106,38 @@ def test_redis_with_url_constructs_successfully(self): redis_url="redis://localhost:6379/0", )) assert plugin is not None + + +class TestFailModePublicSurface: + """fail_mode must be reachable through the advertised Python compat API. + + G6 of the review feedback: the Rust core supports fail_mode, but the + public compat layer was dropping it, so callers following the + RateLimiterConfig helper pattern couldn't actually enable fail-closed + behaviour. + """ + + def test_compat_default_config_includes_fail_mode(self): + """compat_default_config() must list fail_mode among its keys.""" + from cpex_rate_limiter.rate_limiter_rust import compat_default_config # noqa: PLC0415 + + defaults = compat_default_config() + assert "fail_mode" in defaults, ( + f"compat_default_config() must include 'fail_mode'; got keys={sorted(defaults.keys())}" + ) + + def test_rate_limiter_config_preserves_fail_mode(self): + """RateLimiterConfig(fail_mode=...) must round-trip through the __slots__.""" + cfg = RateLimiterConfig(fail_mode="closed") + assert getattr(cfg, "fail_mode", None) == "closed", ( + "RateLimiterConfig must expose fail_mode via its attribute surface" + ) + + def test_rate_limiter_config_fail_mode_defaults_to_open(self): + """When fail_mode is not passed, RateLimiterConfig exposes the safe default.""" + cfg = RateLimiterConfig() + # Default should be the string "open" (mirroring the Rust-side fallback), + # not None — this is what operators read when inspecting the config object. + assert getattr(cfg, "fail_mode", "__missing__") == "open", ( + "RateLimiterConfig.fail_mode must default to 'open' for fail-open behaviour" + )