From a54d856eaefec826eb51369ecfd6fab90cbd0d86 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Thu, 16 Jul 2026 10:49:42 +0300 Subject: [PATCH 01/55] feat(shield): GCRA rate limiting with pluggable keys and limits Replace the fixed-window limiter with an embedded, config-driven GCRA shaper that adds no blocking latency to the request path. - Local GCRA shaper (one timestamp per key): bursts up to `burst`, then throttles to `rate`, with no fixed-window boundary burst. - Keys: client IP, a header value, or a validated JWT claim, each with an IP fallback. Two-phase: IP/header rules run before auth (shed floods pre-verification), claim-keyed rules run after auth using the verified, un-forgeable principal. Phase is derived from the key. - Limit resolution chain: JWT claim (tier name to a profile, or explicit numbers), external service (cached and refreshed off the request path), rule profile, then default. Tiers are named config data. - Optional async cross-instance reconciliation over a shared store (feature `redis`): delta push / aggregate pull on an interval with a sliding-window counter, converging on an approximate fleet-wide limit. A store outage degrades to per-instance limiting, never failing requests. - Emit draft-ietf RateLimit-Limit / -Remaining / -Reset headers, plus Retry-After on 429. Auth now attaches the verified claims as a typed request extension for the post-auth phase. CI gains a Redis service for the reconciliation test. BREAKING CHANGE: the `shield` config schema is replaced. `endpoint_classes`, `identifier_endpoints`, and `window_secs` are removed in favour of `profiles`, `rules` (pattern + key + profile), `default_profile`, `jwt_limits`, `limit_service`, and `sync`. Closes #68 --- .github/workflows/ci.yml | 14 + README.md | 86 +++++- src/auth/forward.rs | 4 +- src/auth/jwks.rs | 2 +- src/auth/mod.rs | 24 +- src/config.rs | 195 ++++++++++--- src/lib.rs | 28 +- src/shield/gcra.rs | 242 ++++++++++++++++ src/shield/global.rs | 369 ++++++++++++++++++++++++ src/shield/matcher.rs | 222 +++++++++++---- src/shield/mod.rs | 591 +++++++++++++++------------------------ src/shield/resolve.rs | 306 ++++++++++++++++++++ src/shield/store.rs | 357 +++++++---------------- src/shield/tests.rs | 339 ++++++++++++++++++++++ src/shield/window.rs | 64 +++++ 15 files changed, 2110 insertions(+), 733 deletions(-) create mode 100644 src/shield/gcra.rs create mode 100644 src/shield/global.rs create mode 100644 src/shield/resolve.rs create mode 100644 src/shield/tests.rs create mode 100644 src/shield/window.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 12a993d..e040f3d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,6 +18,17 @@ jobs: quality-checks: name: Quality Checks (${{ matrix.backend.name }}) runs-on: ubuntu-latest + services: + # Shared store for the rate-limit cross-instance reconciliation test. + redis: + image: redis:7 + ports: + - 6379:6379 + options: >- + --health-cmd "redis-cli ping" + --health-interval 10s + --health-timeout 5s + --health-retries 5 strategy: fail-fast: false matrix: @@ -53,6 +64,9 @@ jobs: run: cargo build --release ${{ matrix.backend.flags }} - name: Test + env: + # Exercises the rate-limit reconciliation against the Redis service. + SHIELD_REDIS_TEST_URL: redis://127.0.0.1:6379/ run: cargo test ${{ matrix.backend.flags }} - name: Publish dry-run diff --git a/README.md b/README.md index d0bd61e..266e1d9 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ Works with **any** gRPC service via proto descriptor files. No code generation, - **Health endpoints** `/health/live`, `/health/ready` (upstream gRPC health probe), `/health/startup` - **Prometheus metrics** at `/metrics` - **CORS** with a configurable origin allow-list -- **Rate limiting (Shield)**: per-client endpoint classes + per-identifier limits, in-process by default or Redis-backed (feature `redis`) for multi-instance +- **Rate limiting (Shield)**: local GCRA shaper (no blocking latency) keyed by client IP, header, or validated JWT claim; named limit tiers as config data; optional async cross-instance reconciliation (feature `redis`) for an approximate fleet-wide limit - **JWT auth**: validate `Bearer` tokens via an Ed25519 PEM key or JWKS auto-discovery, enforce per-route `require_auth` / `required_roles`, and forward claims as headers - **OIDC discovery**: serve `/.well-known/openid-configuration` and a JWKS endpoint (Ed25519) built from config, to front an identity provider - **Forward-auth**: a verification endpoint (`/auth/verify`) for a fronting proxy (nginx `auth_request`, Traefik `forwardAuth`) to delegate auth, returning the verified identity as headers @@ -108,25 +108,38 @@ streaming: sse_keep_alive_secs: 15 # Rate limiting (Shield) +# +# Every decision is made locally with a GCRA shaper (no blocking latency). +# Named profiles define tiers as data; rules bind a path glob to a key and a +# limit. Rules keyed by a JWT claim run after auth (so the claim is verified); +# the rest run before auth so anonymous floods are shed cheaply. shield: enabled: true - window_secs: 60 # default window for bare counts like "20" - # Optional: shared counters across replicas (needs the `redis` build feature). - # Omit for an in-process per-replica store. - # redis_url: "redis://127.0.0.1/" # CIDR ranges of trusted proxies/LBs. X-Forwarded-For is honored only from # these peers; set this behind a load balancer for correct per-client limits. trusted_proxies: ["10.0.0.0/8"] - # Classify endpoints by glob pattern → class → rate (limited per client IP) - endpoint_classes: + # Limit tiers: sustained rate ("N/unit" or a bare count = per minute) + burst + # (max back-to-back requests; defaults to one window of the rate). + profiles: + anon: { rate: "60/min", burst: 20 } + premium: { rate: "1000/min", burst: 100 } + # Applied when a matched rule resolves no other limit. + default_profile: "anon" + rules: + # Anonymous heavy endpoints, keyed by client IP (runs before auth). - pattern: "/api/v1/heavy-*" - class: "heavy" - rate: "10/min" - # Per-identifier limits keyed by a request body field - identifier_endpoints: - - path: "/api/v1/login" - body_field: "email" - rate: "5/min" + key: { type: ip } + profile: "anon" + # Per-principal limit keyed by a validated JWT claim (runs after auth). + - pattern: "/api/v1/**" + key: { type: jwt_claim, claim: "sub" } + # Optional: resolve a key's limit from the JWT itself (tier name → a profile, + # or explicit ratelimit_rpm / ratelimit_burst claims). + # jwt_limits: { tier_claim: "ratelimit_tier" } + # Optional: async cross-instance reconciliation via a shared store for an + # approximate fleet-wide limit (needs the `redis` build feature). The request + # path never blocks on it; a store outage degrades to per-instance limits. + # sync: { redis_url: "redis://127.0.0.1/", interval_ms: 500 } # JWT auth auth: @@ -165,6 +178,51 @@ buf build -o my-service.descriptor.bin protoc --descriptor_set_out=my-service.descriptor.bin --include_imports *.proto ``` +## Rate limiting + +Shield is an embedded, config-driven limiter designed for a data plane: every +decision is made in-process by a GCRA shaper, so it adds no blocking latency to +the request path. GCRA (a token-bucket equivalent storing one timestamp per key) +lets legitimate bursts through up to a configured `burst` while throttling +sustained abuse to the `rate`, with no fixed-window boundary burst. + +**Keying and phases.** A rule keys on the client IP, a header value (API key), +or a validated JWT claim. IP/header rules run *before* auth so anonymous floods +are shed before any signature verification; claim-keyed rules run *after* auth so +they use the verified, un-forgeable principal. The phase is derived from the key; +there is no phase setting to misconfigure. Every key falls back to the client IP +when its value is absent, so a limit can't be dodged by omitting a header or +staying anonymous. + +**Limit sources.** A key's `{rate, burst}` resolves in order: the JWT itself +(a `ratelimit_tier` claim naming a profile, or explicit `ratelimit_rpm` / +`ratelimit_burst`), then an external service (cached and refreshed in the +background, never blocking), then the rule's pinned profile, then the default. +Tier-name indirection lets you retune the numbers in config without re-issuing +tokens or changing the service. + +**Deployment modes.** + +- **Local (default).** No shared store. Each instance enforces the limit + independently, so the fleet-wide effect is roughly `N × rate` for `N` + instances. Zero dependencies, lowest latency. Set per-instance limits with + that multiplier in mind. +- **Reconciled (`sync` + `redis` feature).** Instances asynchronously push their + deltas to a shared store and pull the aggregate on an interval, converging on + an approximate fleet-wide limit. The configured `rate` is then the *fleet* + budget. The request path still never blocks on the store; if the store is + unreachable, instances degrade to local limiting rather than failing requests. + +**Sizing the overshoot.** In reconciled mode the aggregate lags by up to one +`sync.interval_ms`, so the fleet can briefly overshoot the budget by about +`(N - 1) × rate × interval`. Shorter intervals tighten the bound at the cost of +more store traffic; the default (500 ms) suits per-minute limits. The global +view uses a sliding-window counter, so there is no boundary burst on top of this +lag. + +See the `shield:` block under [Configuration](#configuration) for the full +schema. + ## Library Usage ```rust diff --git a/src/auth/forward.rs b/src/auth/forward.rs index 582a973..9aa13bc 100644 --- a/src/auth/forward.rs +++ b/src/auth/forward.rs @@ -67,7 +67,9 @@ impl ForwardAuth { match self.auth.decide(headers, &path, &method).await { // 200 carries the verified claim headers for the fronting proxy to // copy upstream. - AuthDecision::Allow(claim_headers) => (StatusCode::OK, claim_headers).into_response(), + AuthDecision::Allow(claim_headers, _) => { + (StatusCode::OK, claim_headers).into_response() + } AuthDecision::Unauthenticated(msg) => self.deny(msg), AuthDecision::Forbidden(msg) => forbidden(msg), } diff --git a/src/auth/jwks.rs b/src/auth/jwks.rs index 342a129..6f148cf 100644 --- a/src/auth/jwks.rs +++ b/src/auth/jwks.rs @@ -38,7 +38,7 @@ const JWKS_HTTP_TIMEOUT: Duration = Duration::from_secs(5); /// Uses the pure-Rust `ring` provider (installed per-config, not process-global) /// and bundles Mozilla's root store via `webpki-roots`, so the binary needs no /// system CA bundle and works in musl / scratch / distroless images. -fn build_tls_config() -> rustls::ClientConfig { +pub(crate) fn build_tls_config() -> rustls::ClientConfig { let mut roots = rustls::RootCertStore::empty(); roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); rustls::ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) diff --git a/src/auth/mod.rs b/src/auth/mod.rs index eaf0642..b10c454 100644 --- a/src/auth/mod.rs +++ b/src/auth/mod.rs @@ -120,14 +120,23 @@ impl Auth { /// The outcome of an auth check for a request. pub(crate) enum AuthDecision { - /// Allowed; forward these (verified) claim headers to the upstream. - Allow(HeaderMap), + /// Allowed; forward these (verified) claim headers to the upstream, and the + /// verified claims themselves (`None` for anonymous access) for downstream + /// consumers such as per-principal rate limiting. + Allow(HeaderMap, Option), /// Rejected: no/invalid credentials (HTTP 401). Unauthenticated(&'static str), /// Rejected: authenticated but lacking a required role (HTTP 403). Forbidden(&'static str), } +/// Verified JWT claims attached to the request by the auth middleware. Present +/// only when a valid token was supplied, and set exclusively from a verified +/// token (never from client input), so downstream consumers may safely key +/// security decisions (e.g. rate limits) on it. +#[derive(Clone)] +pub(crate) struct ValidatedClaims(pub(crate) std::sync::Arc); + impl Auth { /// Evaluate auth for a request: validate the bearer token, apply the route /// policy, and render the claim headers to forward. This is the single @@ -168,7 +177,7 @@ impl Auth { if let Some(claims) = &claims { inject_claim_headers(&mut claim_headers, claims, &self.claims_headers); } - AuthDecision::Allow(claim_headers) + AuthDecision::Allow(claim_headers, claims) } } @@ -189,11 +198,18 @@ pub async fn middleware( match auth.decide(request.headers(), &path, &method).await { AuthDecision::Unauthenticated(msg) => unauthorized(msg), AuthDecision::Forbidden(msg) => forbidden(msg), - AuthDecision::Allow(claim_headers) => { + AuthDecision::Allow(claim_headers, claims) => { let dst = request.headers_mut(); for (name, value) in &claim_headers { dst.insert(name.clone(), value.clone()); } + // Expose the verified claims to inner layers (e.g. per-principal rate + // limiting) as a typed extension a client cannot forge. + if let Some(claims) = claims { + request + .extensions_mut() + .insert(ValidatedClaims(std::sync::Arc::new(claims))); + } next.run(request).await } } diff --git a/src/config.rs b/src/config.rs index 249db5c..54287f0 100644 --- a/src/config.rs +++ b/src/config.rs @@ -369,58 +369,176 @@ fn default_authz_timeout_ms() -> u64 { } /// Shield (rate limiting) configuration. +/// +/// The proxy runs embedded on each service instance, so every limit decision is +/// made locally with a GCRA shaper (zero blocking latency). A shared store, when +/// configured via [`sync`](ShieldConfig::sync), is reconciled asynchronously off +/// the request path to approximate a fleet-wide limit; the request path never +/// blocks on it. #[derive(Debug, Clone, Deserialize)] #[non_exhaustive] pub struct ShieldConfig { #[serde(default)] pub enabled: bool, - /// Endpoint classification (glob pattern → class → rate limit). + /// Named limit tiers referenced by rules and by tier-name resolution (JWT + /// claim / limit service). Map of profile name → `{ rate, burst }`. + #[serde(default)] + pub profiles: std::collections::HashMap, + /// Rate-limit rules, evaluated in order; the first whose pattern matches the + /// request path applies. + #[serde(default)] + pub rules: Vec, + /// Profile name applied when a matched rule resolves no other limit (JWT and + /// service resolution absent or empty, and the rule sets no explicit + /// profile). Must name an entry in `profiles`. #[serde(default)] - pub endpoint_classes: Vec, - /// Per-identifier rate limiting. + pub default_profile: Option, + /// Resolve a key's limit from claims in the validated JWT. Presence enables + /// JWT-based resolution (tier name or explicit numbers). #[serde(default)] - pub identifier_endpoints: Vec, - /// Window size in seconds (default: 60). - #[serde(default = "default_window_secs")] - pub window_secs: u64, - /// Redis URL for shared counters across replicas (e.g. "redis://127.0.0.1/"). - /// When unset, an in-process per-replica store is used. Requires the `redis` - /// build feature; otherwise the proxy logs a warning and uses the in-process - /// store. + pub jwt_limits: Option, + /// Resolve a key's limit from an external service. The lookup is cached and + /// refreshed in the background; the request path never blocks on it. #[serde(default)] - pub redis_url: Option, + pub limit_service: Option, + /// Asynchronous cross-instance reconciliation via a shared store. When unset, + /// each instance limits locally (fleet limit ≈ N × per-instance). + #[serde(default)] + pub sync: Option, /// CIDR ranges of trusted reverse proxies / load balancers (e.g. /// "10.0.0.0/8"). `X-Forwarded-For` / `X-Real-IP` are honored only when the /// direct peer falls in one of these ranges; otherwise the peer socket /// address is used as the client identity. Empty (the default) means do not - /// trust forwarding headers — set this behind a load balancer. + /// trust forwarding headers; set this behind a load balancer. #[serde(default)] pub trusted_proxies: Vec, } -fn default_window_secs() -> u64 { - 60 +/// A named limit tier: a sustained rate plus an instantaneous burst capacity. +#[derive(Debug, Clone, Deserialize)] +#[non_exhaustive] +pub struct LimitProfileConfig { + /// Sustained rate as `"/"` (e.g. `"100/min"`, units + /// `s`/`min`/`hour`) or a bare count (interpreted per minute). + pub rate: String, + /// Maximum requests admitted back-to-back before throttling to the rate. + /// Defaults to the per-window rate count (one full window of burst). + #[serde(default)] + pub burst: Option, } -/// Endpoint classification for rate limiting. +/// One rate-limit rule: a path pattern, how to key it, and an optional static +/// profile. The rule's *phase* (before or after auth) is derived automatically: +/// rules that need validated JWT claims (a `jwt_claim` key, or JWT-based limit +/// resolution) run after auth; the rest run before auth so anonymous floods are +/// shed before any signature verification. #[derive(Debug, Clone, Deserialize)] #[non_exhaustive] -pub struct EndpointClassConfig { - /// Glob pattern (e.g., "/v1/auth/**"). +pub struct RateRuleConfig { + /// Glob path pattern (`*` within a segment, `**` across segments). pub pattern: String, - /// Class name (e.g., "auth"). - pub class: String, - /// Rate limit string (e.g., "20/min"). - pub rate: String, + /// How to derive the limit key (who is limited). Defaults to client IP. + #[serde(default)] + pub key: KeySourceConfig, + /// Static profile name for this rule, used when JWT/service resolution does + /// not apply or yields nothing. Must name an entry in `profiles`. + #[serde(default)] + pub profile: Option, } -/// Per-identifier rate limiting config. +/// How a rule derives its limit key. All sources fall back to the client IP when +/// their value is absent, so a limit can't be bypassed by omitting a header or +/// authenticating anonymously. Written as a tagged map, e.g. +/// `key: { type: jwt_claim, claim: sub }`; omitting `key` defaults to `ip`. +#[derive(Debug, Clone, Deserialize, Default, PartialEq, Eq)] +#[serde(tag = "type", rename_all = "snake_case")] +#[non_exhaustive] +pub enum KeySourceConfig { + /// Client IP (trusted-proxy `X-Forwarded-For` aware). `{ type: ip }`. + #[default] + Ip, + /// Value of a named request header (e.g. an API key). + /// `{ type: header, name: x-api-key }`. + Header { + /// Header whose value identifies the client. + name: String, + }, + /// Value of a claim from the validated JWT (provider-agnostic). + /// `{ type: jwt_claim, claim: sub }`. + JwtClaim { + /// Claim whose value identifies the principal. + claim: String, + }, +} + +/// Claims that carry a key's limit inside the JWT itself. A tier-name claim maps +/// to a `profiles` entry (numbers stay tunable in config); direct numeric claims +/// set the limit explicitly. #[derive(Debug, Clone, Deserialize)] #[non_exhaustive] -pub struct IdentifierEndpointConfig { - pub path: String, - pub body_field: String, - pub rate: String, +pub struct JwtLimitConfig { + /// Claim naming a profile tier (e.g. `"premium"`). Default: `ratelimit_tier`. + #[serde(default = "default_tier_claim")] + pub tier_claim: String, + /// Claim carrying an explicit sustained rate, requests per minute. Default: + /// `ratelimit_rpm`. + #[serde(default = "default_rpm_claim")] + pub rpm_claim: String, + /// Claim carrying an explicit burst capacity. Default: `ratelimit_burst`. + #[serde(default = "default_burst_claim")] + pub burst_claim: String, +} + +fn default_tier_claim() -> String { + "ratelimit_tier".to_string() +} +fn default_rpm_claim() -> String { + "ratelimit_rpm".to_string() +} +fn default_burst_claim() -> String { + "ratelimit_burst".to_string() +} + +/// External limit-resolution service. The response names a tier or gives explicit +/// numbers; results are cached and refreshed asynchronously, never on the request +/// path. +#[derive(Debug, Clone, Deserialize)] +#[non_exhaustive] +pub struct LimitServiceConfig { + /// HTTP endpoint queried with the limit key; returns `{ tier }` or + /// `{ rate_per_min, burst }`. + pub endpoint: String, + /// How long a resolved limit is cached before a background refresh, in + /// seconds (default: 300). + #[serde(default = "default_limit_ttl_secs")] + pub ttl_secs: u64, + /// Timeout for the background fetch, in milliseconds (default: 500). + #[serde(default = "default_limit_timeout_ms")] + pub timeout_ms: u64, +} + +fn default_limit_ttl_secs() -> u64 { + 300 +} +fn default_limit_timeout_ms() -> u64 { + 500 +} + +/// Asynchronous cross-instance reconciliation via a shared store. +#[derive(Debug, Clone, Deserialize)] +#[non_exhaustive] +pub struct SyncConfig { + /// Shared-store URL (e.g. `"redis://127.0.0.1/"`). Requires the `redis` build + /// feature; without it the proxy logs a warning and stays local-only. + pub redis_url: String, + /// Background push/pull interval in milliseconds (default: 500). The + /// worst-case fleet overshoot is bounded by `(N-1) × rate × interval`. + #[serde(default = "default_sync_interval_ms")] + pub interval_ms: u64, +} + +fn default_sync_interval_ms() -> u64 { + 500 } /// OIDC discovery config. @@ -836,17 +954,20 @@ auth: shield: enabled: true - endpoint_classes: + profiles: + auth: { rate: "20/min", burst: 5 } + default: { rate: "100/min" } + premium: { rate: "1000/min", burst: 50 } + default_profile: "default" + jwt_limits: + tier_claim: "ratelimit_tier" + rules: - pattern: "/v1/auth/**" - class: "auth" - rate: "20/min" - - pattern: "/**" - class: "default" - rate: "100/min" - identifier_endpoints: - - path: "/v1/auth/opaque/login/start" - body_field: "identifier" - rate: "10/min" + key: { type: ip } + profile: "auth" + - pattern: "/v1/**" + key: { type: jwt_claim, claim: "sub" } + trusted_proxies: ["10.0.0.0/8"] oidc_discovery: enabled: true diff --git a/src/lib.rs b/src/lib.rs index a9ef65a..e30cdcf 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -524,8 +524,18 @@ impl ProxyServer { // Duplicate-route collisions (including the verify path) were already // rejected up front, before any router was built. - // Auth runs inside Shield (added first = inner): rate limiting sheds - // load before any signature verification work. + // Two-phase rate limiting around auth. The post-auth phase (rules keyed + // by a validated JWT claim) is layered first so it sits *inside* auth and + // sees the verified claims; the pre-auth phase (IP / header keys) is + // layered after auth below so it runs *first* and sheds anonymous floods + // before any signature verification. + if let Some(shield) = &shield { + router = router.layer(axum::middleware::from_fn_with_state( + shield.clone(), + shield::post_auth_middleware, + )); + } + if let Some(auth) = auth { router = router.layer(axum::middleware::from_fn_with_state(auth, auth::middleware)); } @@ -548,13 +558,15 @@ impl ProxyServer { router = router.merge(forward_auth.routes()); } - // Shield is added before maintenance so maintenance wraps it (outer - // layers run first): a request rejected by the maintenance gate must - // not be charged against its rate-limit budget. - if let Some(shield) = shield { + // Pre-auth phase, added before maintenance so maintenance wraps it (outer + // layers run first): a request rejected by the maintenance gate must not + // be charged against its rate-limit budget. Placed after the auth layer + // so it runs before auth, and after the verify route so that endpoint is + // rate-limited too (but not JWT-gated). + if let Some(shield) = &shield { router = router.layer(axum::middleware::from_fn_with_state( - shield, - shield::middleware, + shield.clone(), + shield::pre_auth_middleware, )); } diff --git a/src/shield/gcra.rs b/src/shield/gcra.rs new file mode 100644 index 0000000..3ef1081 --- /dev/null +++ b/src/shield/gcra.rs @@ -0,0 +1,242 @@ +//! GCRA (Generic Cell Rate Algorithm) rate limiter core. +//! +//! GCRA is a token-bucket equivalent that stores a single value per key: the +//! *theoretical arrival time* (TAT). Compared with a fixed window it has no +//! boundary burst (a client cannot spend two full windows across a boundary) +//! and it lets legitimate bursts through up to a configured capacity while +//! throttling sustained abuse to the steady rate. +//! +//! This module is pure arithmetic with an injected clock, so the burst +//! boundary, refill, and clock-skew behaviour are all unit-testable without a +//! store or a real clock. + +use std::time::Duration; + +/// A GCRA limiter parameterised by a steady emission interval and a burst +/// tolerance. Construct one from a [`Profile`] with [`Gcra::from_profile`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Gcra { + /// Time between two requests at the sustained rate (`window / rate`). + emission_interval: Duration, + /// Delay-variation tolerance: how far ahead of the steady schedule a burst + /// may run. `(burst - 1) * emission_interval`, so a fresh key admits exactly + /// `burst` requests instantly before throttling to the steady rate. + tau: Duration, +} + +/// A named limit tier: a sustained rate over a window plus an instantaneous +/// burst capacity. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Profile { + /// Sustained requests permitted per `window`. + pub rate: u64, + /// Length of the sustained-rate window. + pub window: Duration, + /// Maximum requests admitted back-to-back before throttling to the rate. + /// At least 1. + pub burst: u64, +} + +/// The outcome of a single GCRA check. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Verdict { + /// Whether the request is conforming (allowed). + pub allowed: bool, + /// The TAT to persist for this key (unchanged from the stored value when + /// the request is rejected). + pub new_tat: Duration, + /// Requests still admissible at this instant after this one (0 when rejected). + pub remaining: u64, + /// When rejected, how long until a retry would conform (`Retry-After`). + pub retry_after: Duration, + /// How long until the limiter drains back toward full capacity + /// (`RateLimit-Reset`). + pub reset_after: Duration, +} + +impl Gcra { + /// Build a limiter from a [`Profile`]. `burst` is clamped to at least 1 and + /// `rate` to at least 1 so the emission interval is finite. + pub fn from_profile(profile: Profile) -> Self { + let rate = profile.rate.max(1); + let burst = profile.burst.max(1); + // T = window / rate, computed in nanoseconds to avoid truncation. + let window_nanos = profile.window.as_nanos().max(1); + let emission_nanos = window_nanos / u128::from(rate); + let emission_interval = duration_from_nanos(emission_nanos.max(1)); + let tau = emission_interval * u32::try_from(burst - 1).unwrap_or(u32::MAX); + Self { + emission_interval, + tau, + } + } + + /// Evaluate a request arriving at `now` (a monotonically-increasing instant + /// expressed as a [`Duration`] since a fixed epoch), given the key's stored + /// TAT (`None` for a first-seen key). + /// + /// A `now` that moves backwards (clock skew) is tolerated: the check never + /// panics and treats the effective arrival time as `max(now, tat_floor)`. + pub fn check(&self, stored_tat: Option, now: Duration) -> Verdict { + let t = self.emission_interval; + let tau = self.tau; + // Canonical GCRA uses the stored TAT as-is (a first-seen key starts at + // `now`, an empty bucket). + let tat = stored_tat.unwrap_or(now); + + // Number of requests admissible at `now` from the stored TAT: each + // admitted request pushes TAT forward by `t`, and a request conforms + // while `now >= tat + k*t - tau`. `now + tau >= tat` means at least one + // conforms. + let admissible = if now + tau >= tat { + let slack = (now + tau) - tat; // >= 0 + 1 + div_floor(slack, t) + } else { + 0 + }; + + if admissible == 0 { + // Non-conforming: earliest conforming instant is `tat - tau`. + let retry_after = tat.saturating_sub(tau).saturating_sub(now); + Verdict { + allowed: false, + new_tat: tat, + remaining: 0, + retry_after, + reset_after: tat.saturating_sub(now), + } + } else { + let new_tat = tat.max(now) + t; + Verdict { + allowed: true, + new_tat, + remaining: admissible - 1, + retry_after: Duration::ZERO, + reset_after: new_tat.saturating_sub(now), + } + } + } +} + +/// Floor division of two `Duration`s (`a / b`), in nanoseconds. +fn div_floor(a: Duration, b: Duration) -> u64 { + let b = b.as_nanos().max(1); + u64::try_from(a.as_nanos() / b).unwrap_or(u64::MAX) +} + +/// A `Duration` from a `u128` nanosecond count, saturating at `Duration::MAX`. +fn duration_from_nanos(nanos: u128) -> Duration { + let secs = nanos / 1_000_000_000; + let sub = (nanos % 1_000_000_000) as u32; + match u64::try_from(secs) { + Ok(secs) => Duration::new(secs, sub), + Err(_) => Duration::MAX, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn profile(rate: u64, window_secs: u64, burst: u64) -> Profile { + Profile { + rate, + window: Duration::from_secs(window_secs), + burst, + } + } + + /// A fresh key admits exactly `burst` requests instantly, then rejects. + #[test] + fn burst_boundary_admits_exactly_burst() { + let g = Gcra::from_profile(profile(60, 60, 3)); // 1 req/s, burst 3 + let now = Duration::from_secs(100); + let mut tat = None; + for i in 0..3 { + let v = g.check(tat, now); + assert!(v.allowed, "request {i} should be allowed"); + assert_eq!(v.remaining, (3 - 1 - i) as u64, "remaining after {i}"); + tat = Some(v.new_tat); + } + // 4th at the same instant is rejected. + let v = g.check(tat, now); + assert!(!v.allowed); + assert_eq!(v.remaining, 0); + // Retry after one emission interval (1s). + assert_eq!(v.retry_after, Duration::from_secs(1)); + } + + /// After exhausting the burst, one request is admitted every emission + /// interval (refill at the steady rate). + #[test] + fn refills_at_steady_rate() { + let g = Gcra::from_profile(profile(60, 60, 2)); // 1 req/s, burst 2 + let start = Duration::from_secs(0); + let mut tat = None; + // Spend the burst (2) at t=0. + for _ in 0..2 { + let v = g.check(tat, start); + assert!(v.allowed); + tat = Some(v.new_tat); + } + // Immediately after: rejected. + assert!(!g.check(tat, start).allowed); + // 1 second later: exactly one slot has refilled. + let later = start + Duration::from_secs(1); + let v = g.check(tat, later); + assert!(v.allowed); + tat = Some(v.new_tat); + // A second request at the same instant is rejected (only one refilled). + assert!(!g.check(tat, later).allowed); + } + + /// A `now` that jumps backwards (clock skew) must not panic and must not + /// wrongly admit an unbounded burst. + #[test] + fn tolerates_backward_clock() { + let g = Gcra::from_profile(profile(60, 60, 1)); // 1 req/s, burst 1 + let now = Duration::from_secs(1000); + let v = g.check(None, now); + assert!(v.allowed); + let tat = Some(v.new_tat); + // Clock jumps back 500s. The next request is still governed by the + // stored TAT (which is ahead), so it is rejected, not admitted. + let back = Duration::from_secs(500); + let v = g.check(tat, back); + // The stored TAT (ahead of the rewound clock) still governs: the request + // is rejected, not wrongly admitted, and retry_after is a finite, + // conservative wait (never a panic or an overflow). + assert!(!v.allowed); + assert!(v.retry_after > Duration::ZERO); + assert!(v.retry_after <= Duration::from_secs(501)); + } + + /// `reset_after` shrinks toward zero as the bucket drains over time. + #[test] + fn reset_after_drains_over_time() { + let g = Gcra::from_profile(profile(60, 60, 5)); // 1 req/s, burst 5 + let start = Duration::from_secs(0); + let mut tat = None; + let mut last_reset = Duration::MAX; + for _ in 0..5 { + let v = g.check(tat, start); + assert!(v.allowed); + tat = Some(v.new_tat); + last_reset = v.reset_after; + } + // After the full burst, reset_after == burst * emission (5s). + assert_eq!(last_reset, Duration::from_secs(5)); + } + + /// A bare-rate profile (burst 1) admits one request per emission interval + /// with no instantaneous burst. + #[test] + fn burst_one_is_pure_rate_limit() { + let g = Gcra::from_profile(profile(2, 1, 1)); // 2 req/s, burst 1 + let now = Duration::from_secs(0); + let v = g.check(None, now); + assert!(v.allowed); + // Second request at the same instant rejected (burst 1). + assert!(!g.check(Some(v.new_tat), now).allowed); + } +} diff --git a/src/shield/global.rs b/src/shield/global.rs new file mode 100644 index 0000000..3533daa --- /dev/null +++ b/src/shield/global.rs @@ -0,0 +1,369 @@ +//! Cross-instance reconciliation of the fleet-wide limit view via a shared store. +//! +//! The request path never touches the store. It reads a cached estimate of what +//! the *rest* of the fleet has consumed for a key (updated by the background +//! task) plus this instance's own count, and gates on that. A background task +//! pushes this instance's deltas (`INCRBY`) and pulls the aggregate (`MGET`) on +//! an interval, so a store outage degrades to per-instance limiting rather than +//! failing requests. +//! +//! Counts are held in a sliding window (see [`super::window`]) so the fleet gate +//! has no boundary burst. The worst-case overshoot is bounded by one sync +//! interval of the other instances' traffic: `(N-1) × rate × interval`. + +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use super::window; + +/// Redis key namespace for a rate-limit key at a given epoch. +fn epoch_key(key: &str, epoch: u64) -> String { + format!("sp:rl:{key}:{epoch}") +} + +/// Wall-clock time since the Unix epoch (shared reference so every instance +/// agrees on window boundaries). +fn unix_now() -> Duration { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or(Duration::ZERO) +} + +/// Drop per-key state untouched for this long (bounds memory under churning +/// key cardinality). +const KEY_TTL: Duration = Duration::from_secs(600); + +/// Per-key reconciliation state. +#[derive(Debug, Clone)] +struct KeyState { + epoch: u64, + /// Window length in seconds (from the resolved profile). + window_secs: u64, + /// Requests this instance admitted in the current epoch. + local_count: u64, + /// How much of `local_count` has already been pushed to the store. + pushed: u64, + /// The rest of the fleet's sliding consumption, from the last pull. + remote_estimate: u64, + last_seen: Instant, +} + +impl KeyState { + fn new(epoch: u64, window_secs: u64) -> Self { + Self { + epoch, + window_secs, + local_count: 0, + pushed: 0, + remote_estimate: 0, + last_seen: Instant::now(), + } + } + + /// Advance to `epoch`, resetting the per-epoch counts. The remote estimate is + /// kept (the background task refreshes it) so a fresh epoch does not briefly + /// open the full budget fleet-wide. + fn roll_to(&mut self, epoch: u64) { + if self.epoch != epoch { + self.epoch = epoch; + self.local_count = 0; + self.pushed = 0; + } + } +} + +/// Cross-instance counter reconciliation over a shared Redis-protocol store. +pub struct GlobalCounters { + client: redis::Client, + conn: tokio::sync::OnceCell, + interval: Duration, + states: dashmap::DashMap, +} + +impl GlobalCounters { + /// Open the shared store and return the reconciler. + /// + /// # Errors + /// Returns the underlying error when the store URL is invalid. + pub fn build(url: &str, interval: Duration) -> Result, String> { + let client = redis::Client::open(url) + .map_err(|e| format!("invalid shared-store URL for rate limiting: {e}"))?; + Ok(Arc::new(Self { + client, + conn: tokio::sync::OnceCell::new(), + interval: interval.max(Duration::from_millis(50)), + states: dashmap::DashMap::new(), + })) + } + + /// Whether a request for `key` is within the fleet budget. Read-only: reads + /// the cached remote estimate plus this instance's count. Does not touch the + /// store and does not record the admit (call [`record`](Self::record) after + /// the local check also passes). + pub fn gate(&self, key: &str, budget: u64, window: Duration) -> bool { + let now = unix_now(); + let ep = window::epoch(now, window); + let window_secs = window.as_secs().max(1); + let mut state = self + .states + .entry(key.to_string()) + .or_insert_with(|| KeyState::new(ep, window_secs)); + state.roll_to(ep); + state.last_seen = Instant::now(); + state.remote_estimate + state.local_count < budget + } + + /// Record one admitted request for `key`, to be pushed to the store on the + /// next background tick. + pub fn record(&self, key: &str, window: Duration) { + let now = unix_now(); + let ep = window::epoch(now, window); + let window_secs = window.as_secs().max(1); + let mut state = self + .states + .entry(key.to_string()) + .or_insert_with(|| KeyState::new(ep, window_secs)); + state.roll_to(ep); + state.local_count += 1; + state.last_seen = Instant::now(); + } + + /// Spawn the background reconciliation loop. A no-op (with a warning) when + /// called outside a Tokio runtime, so the local shaper still works. + pub fn spawn(self: &Arc) { + if tokio::runtime::Handle::try_current().is_err() { + tracing::warn!("rate-limit reconciler not started: no Tokio runtime in this context"); + return; + } + let this = self.clone(); + tokio::spawn(async move { + let mut ticker = tokio::time::interval(this.interval); + loop { + ticker.tick().await; + this.reconcile().await; + } + }); + } + + async fn connection(&self) -> redis::RedisResult { + self.conn + .get_or_try_init(|| self.client.get_multiplexed_async_connection()) + .await + .cloned() + } + + /// One reconciliation pass: push this instance's deltas, pull the aggregate, + /// refresh each key's remote estimate, and evict stale keys. + async fn reconcile(&self) { + self.evict_stale(); + + // Phase 1 (locked, brief): snapshot each active key's push delta and the + // epoch keys to read. No store I/O while holding a lock. + let now = unix_now(); + let mut plans: Vec = Vec::new(); + for entry in self.states.iter() { + let key = entry.key().clone(); + let s = entry.value(); + let window = Duration::from_secs(s.window_secs); + let ep = window::epoch(now, window); + let delta = s.local_count.saturating_sub(s.pushed); + plans.push(PushPlan { + key, + epoch: ep, + window, + delta, + }); + } + if plans.is_empty() { + return; + } + + let mut conn = match self.connection().await { + Ok(c) => c, + Err(e) => { + tracing::warn!("rate-limit shared store unavailable, staying local: {e}"); + return; + } + }; + + // Phase 2 (no locks): push deltas, then read current+previous epochs. + if let Err(e) = self.push_deltas(&mut conn, &plans).await { + tracing::warn!("rate-limit delta push failed: {e}"); + return; + } + let reads = match self.read_epochs(&mut conn, &plans).await { + Ok(r) => r, + Err(e) => { + tracing::warn!("rate-limit aggregate read failed: {e}"); + return; + } + }; + + // Phase 3 (locked, brief): commit pushed deltas and the fresh estimate. + for (plan, (cur, prev)) in plans.iter().zip(reads) { + if let Some(mut s) = self.states.get_mut(&plan.key) { + if s.epoch == plan.epoch { + s.pushed += plan.delta; + } + let elapsed = window::elapsed_in_window(now, plan.window); + let est = window::sliding_estimate(cur, prev, elapsed, plan.window); + // Exclude this instance's own current-epoch contribution: the + // gate adds `local_count` separately. + let others = (est.round() as i64 - s.pushed as i64).max(0); + s.remote_estimate = others as u64; + } + } + } + + /// `INCRBY` each key with a pending delta and re-arm its TTL, in one pipeline. + async fn push_deltas( + &self, + conn: &mut redis::aio::MultiplexedConnection, + plans: &[PushPlan], + ) -> redis::RedisResult<()> { + let mut pipe = redis::pipe(); + let mut any = false; + for p in plans.iter().filter(|p| p.delta > 0) { + any = true; + let k = epoch_key(&p.key, p.epoch); + let ttl_ms = (p.window.as_millis() as u64).saturating_mul(2).max(1); + pipe.cmd("INCRBY").arg(&k).arg(p.delta).ignore(); + pipe.cmd("PEXPIRE").arg(&k).arg(ttl_ms).ignore(); + } + if !any { + return Ok(()); + } + pipe.query_async(conn).await + } + + /// `MGET` the current and previous epoch counts for every plan, returning + /// `(cur, prev)` per plan in order. + async fn read_epochs( + &self, + conn: &mut redis::aio::MultiplexedConnection, + plans: &[PushPlan], + ) -> redis::RedisResult> { + let mut keys: Vec = Vec::with_capacity(plans.len() * 2); + for p in plans { + keys.push(epoch_key(&p.key, p.epoch)); + keys.push(epoch_key(&p.key, p.epoch.saturating_sub(1))); + } + let vals: Vec> = redis::cmd("MGET").arg(&keys).query_async(conn).await?; + Ok(plans + .iter() + .enumerate() + .map(|(i, _)| { + let cur = vals.get(i * 2).copied().flatten().unwrap_or(0).max(0) as u64; + let prev = vals.get(i * 2 + 1).copied().flatten().unwrap_or(0).max(0) as u64; + (cur, prev) + }) + .collect()) + } + + /// Drop per-key state not seen within [`KEY_TTL`]. + fn evict_stale(&self) { + let now = Instant::now(); + self.states + .retain(|_, s| now.duration_since(s.last_seen) < KEY_TTL); + } +} + +/// A per-key push plan captured under lock, used without holding locks. +struct PushPlan { + key: String, + epoch: u64, + window: Duration, + delta: u64, +} + +#[cfg(test)] +mod tests { + use super::*; + + const W: Duration = Duration::from_secs(60); + + // Opening a client does not connect, so gate/record can be exercised without + // a running store. + fn counters() -> Arc { + GlobalCounters::build("redis://127.0.0.1/", Duration::from_millis(500)).unwrap() + } + + #[test] + fn gate_admits_until_local_count_reaches_budget() { + let g = counters(); + // Budget 3: three admits pass, the fourth is over budget. + for _ in 0..3 { + assert!(g.gate("k", 3, W)); + g.record("k", W); + } + assert!(!g.gate("k", 3, W)); + } + + #[test] + fn gate_accounts_for_remote_estimate() { + let g = counters(); + // Simulate the background task having observed 2 requests elsewhere. + let now = unix_now(); + let ep = window::epoch(now, W); + let mut state = KeyState::new(ep, W.as_secs()); + state.remote_estimate = 2; + g.states.insert("k".to_string(), state); + // Budget 3, remote 2 → one local admit fits, the next is over budget. + assert!(g.gate("k", 3, W)); + g.record("k", W); + assert!(!g.gate("k", 3, W)); + } + + #[test] + fn independent_keys_have_independent_budgets() { + let g = counters(); + assert!(g.gate("a", 1, W)); + g.record("a", W); + assert!(!g.gate("a", 1, W)); + // A different key is unaffected. + assert!(g.gate("b", 1, W)); + } + + /// End-to-end reconciliation against a live Redis-protocol store: one + /// instance's admits become visible to another after a reconcile pass, so + /// the two enforce one combined budget. + /// + /// Requires a reachable store; the URL comes from `SHIELD_REDIS_TEST_URL` + /// (CI sets it to its Redis service). Without it the test reports that it was + /// skipped for lack of infrastructure rather than silently passing. + #[tokio::test] + async fn reconciles_across_instances() { + let Ok(url) = std::env::var("SHIELD_REDIS_TEST_URL") else { + eprintln!("SKIP reconciles_across_instances: SHIELD_REDIS_TEST_URL not set"); + return; + }; + // Unique key per run so repeated runs / parallel suites do not collide. + let nonce = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let key = format!("it:{}:{nonce}", std::process::id()); + + let a = GlobalCounters::build(&url, Duration::from_millis(200)).unwrap(); + let b = GlobalCounters::build(&url, Duration::from_millis(200)).unwrap(); + + // Instance A admits two requests and pushes them to the store. + a.record(&key, W); + a.record(&key, W); + a.reconcile().await; + + // B's first contact with the key: it has not reconciled this key yet, so + // it admits once (the documented one-interval lag / bounded overshoot). + assert!(b.gate(&key, 3, W), "B admits its first request for the key"); + b.record(&key, W); + + // After a reconcile pass B has pushed its own admit and pulled the + // aggregate: the combined view is 3 (A's 2 + B's 1) = the budget, so B + // rejects the next request. The two instances enforce one combined limit. + b.reconcile().await; + assert!( + !b.gate(&key, 3, W), + "B must reject once the combined budget is reached" + ); + } +} diff --git a/src/shield/matcher.rs b/src/shield/matcher.rs index 734a491..0d123e5 100644 --- a/src/shield/matcher.rs +++ b/src/shield/matcher.rs @@ -1,25 +1,55 @@ -//! Compile Shield config patterns into runtime path matchers. +//! Compile Shield config into runtime matchers, limiters, and rules. + +use std::collections::HashMap; +use std::time::Duration; use globset::GlobMatcher; -use super::rate::Rate; -use crate::config::{EndpointClassConfig, IdentifierEndpointConfig}; -use std::time::Duration; +use super::gcra::{Gcra, Profile}; +use crate::config::{KeySourceConfig, LimitProfileConfig, RateRuleConfig}; -/// A compiled endpoint-class rule: requests whose path matches `matcher` are -/// limited per client at `rate`, grouped under `class`. -pub struct EndpointClass { - pub matcher: GlobMatcher, - pub class: String, - pub rate: Rate, +/// Bare rate counts (`"20"` with no unit) are interpreted per this window. +const BARE_COUNT_WINDOW: Duration = Duration::from_secs(60); + +/// A compiled limit tier: the GCRA shaper, the per-window count reported in the +/// `RateLimit-Limit` header, and the window itself (the fleet-gate epoch length). +#[derive(Debug, Clone, Copy)] +pub struct CompiledProfile { + pub gcra: Gcra, + /// Sustained request count per window (for `RateLimit-Limit`). + pub limit: u64, + /// Length of the sustained-rate window. + pub window: Duration, } -/// A compiled per-identifier rule: requests to a matching path are limited by a -/// value read from the request body field `body_field`. -pub struct IdentifierEndpoint { +/// How a rule derives its limit key. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum KeySource { + /// Client IP (trusted-proxy aware). + Ip, + /// A named request-header value (API-key style); IP fallback when absent. + Header(String), + /// A validated-JWT claim value; IP fallback for anonymous traffic. + JwtClaim(String), +} + +/// Whether a rule can be decided before auth runs, or needs validated claims. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Phase { + /// No validated claims needed: run before auth so floods are shed cheaply. + PreAuth, + /// Needs validated claims (key derived from the JWT): run after auth. + PostAuth, +} + +/// A compiled rate-limit rule. +#[derive(Debug, Clone)] +pub struct CompiledRule { pub matcher: GlobMatcher, - pub body_field: String, - pub rate: Rate, + pub key: KeySource, + /// Static profile name, if the rule pins one. + pub profile: Option, + pub phase: Phase, } /// Build a glob matcher where `*` stays within a path segment and `**` spans @@ -32,35 +62,65 @@ fn path_glob(pattern: &str) -> Result { .map_err(|e| format!("invalid glob pattern {pattern:?}: {e}")) } -/// Compile endpoint-class rules, parsing each rate against `default_window`. -pub fn compile_endpoint_classes( - configs: &[EndpointClassConfig], - default_window: Duration, -) -> Result, String> { +/// Compile a single limit profile from its config. +pub fn compile_profile(cfg: &LimitProfileConfig) -> Result { + let rate = super::rate::Rate::parse(&cfg.rate, BARE_COUNT_WINDOW)?; + // Default burst = one full window of the sustained rate. + let burst = cfg.burst.unwrap_or(rate.limit).max(1); + let gcra = Gcra::from_profile(Profile { + rate: rate.limit, + window: rate.window, + burst, + }); + Ok(CompiledProfile { + gcra, + limit: rate.limit, + window: rate.window, + }) +} + +/// Compile every named profile. +pub fn compile_profiles( + configs: &HashMap, +) -> Result, String> { configs .iter() - .map(|c| { - Ok(EndpointClass { - matcher: path_glob(&c.pattern)?, - class: c.class.clone(), - rate: Rate::parse(&c.rate, default_window)?, - }) - }) + .map(|(name, cfg)| Ok((name.clone(), compile_profile(cfg)?))) .collect() } -/// Compile per-identifier rules, parsing each rate against `default_window`. -pub fn compile_identifier_endpoints( - configs: &[IdentifierEndpointConfig], - default_window: Duration, -) -> Result, String> { +/// Compile rules, validating that any pinned `profile` names an existing tier. +/// Each rule's phase is derived from its key: a `jwt_claim` key needs validated +/// claims and runs after auth; every other key runs before auth. +pub fn compile_rules( + configs: &[RateRuleConfig], + profiles: &HashMap, +) -> Result, String> { configs .iter() .map(|c| { - Ok(IdentifierEndpoint { - matcher: path_glob(&c.path)?, - body_field: c.body_field.clone(), - rate: Rate::parse(&c.rate, default_window)?, + if let Some(name) = &c.profile { + if !profiles.contains_key(name) { + return Err(format!( + "rule {:?} references unknown profile {name:?}", + c.pattern + )); + } + } + let key = match &c.key { + KeySourceConfig::Ip => KeySource::Ip, + KeySourceConfig::Header { name } => KeySource::Header(name.clone()), + KeySourceConfig::JwtClaim { claim } => KeySource::JwtClaim(claim.clone()), + }; + let phase = match key { + KeySource::JwtClaim(_) => Phase::PostAuth, + KeySource::Ip | KeySource::Header(_) => Phase::PreAuth, + }; + Ok(CompiledRule { + matcher: path_glob(&c.pattern)?, + key, + profile: c.profile.clone(), + phase, }) }) .collect() @@ -70,43 +130,85 @@ pub fn compile_identifier_endpoints( mod tests { use super::*; - fn ec(pattern: &str, class: &str, rate: &str) -> EndpointClassConfig { - EndpointClassConfig { - pattern: pattern.to_string(), - class: class.to_string(), - rate: rate.to_string(), - } + fn profiles() -> HashMap { + let mut cfg = HashMap::new(); + cfg.insert( + "auth".to_string(), + LimitProfileConfig { + rate: "20/min".to_string(), + burst: Some(5), + }, + ); + compile_profiles(&cfg).unwrap() + } + + #[test] + fn profile_defaults_burst_to_rate_count() { + let p = compile_profile(&LimitProfileConfig { + rate: "100/min".to_string(), + burst: None, + }) + .unwrap(); + assert_eq!(p.limit, 100); } #[test] - fn endpoint_class_glob_respects_segments() { - let classes = compile_endpoint_classes( - &[ec("/api/v1/heavy-*", "heavy", "10/min")], - Duration::from_secs(60), + fn glob_respects_and_spans_segments() { + let rules = compile_rules( + &[ + RateRuleConfig { + pattern: "/api/v1/heavy-*".to_string(), + key: KeySourceConfig::Ip, + profile: Some("auth".to_string()), + }, + RateRuleConfig { + pattern: "/v1/auth/**".to_string(), + key: KeySourceConfig::Ip, + profile: None, + }, + ], + &profiles(), ) .unwrap(); - let m = &classes[0].matcher; - assert!(m.is_match("/api/v1/heavy-export")); - // `*` does not cross a path separator. - assert!(!m.is_match("/api/v1/heavy-export/sub")); - assert!(!m.is_match("/api/v1/light")); + assert!(rules[0].matcher.is_match("/api/v1/heavy-export")); + assert!(!rules[0].matcher.is_match("/api/v1/heavy-export/sub")); + assert!(rules[1].matcher.is_match("/v1/auth/opaque/start")); } #[test] - fn double_star_spans_segments() { - let classes = compile_endpoint_classes( - &[ec("/v1/auth/**", "auth", "20/min")], - Duration::from_secs(60), + fn phase_is_derived_from_key() { + let rules = compile_rules( + &[ + RateRuleConfig { + pattern: "/a".to_string(), + key: KeySourceConfig::Ip, + profile: None, + }, + RateRuleConfig { + pattern: "/b".to_string(), + key: KeySourceConfig::JwtClaim { + claim: "sub".to_string(), + }, + profile: None, + }, + ], + &profiles(), ) .unwrap(); - let m = &classes[0].matcher; - assert!(m.is_match("/v1/auth/login")); - assert!(m.is_match("/v1/auth/opaque/start")); + assert_eq!(rules[0].phase, Phase::PreAuth); + assert_eq!(rules[1].phase, Phase::PostAuth); } #[test] - fn invalid_rate_fails_compilation() { - let err = compile_endpoint_classes(&[ec("/x", "c", "nonsense")], Duration::from_secs(60)); + fn unknown_profile_reference_fails() { + let err = compile_rules( + &[RateRuleConfig { + pattern: "/x".to_string(), + key: KeySourceConfig::Ip, + profile: Some("nope".to_string()), + }], + &profiles(), + ); assert!(err.is_err()); } } diff --git a/src/shield/mod.rs b/src/shield/mod.rs index 4c04db6..785f225 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -1,37 +1,54 @@ //! Shield: request rate limiting. //! -//! Enforces the `shield` config as an axum middleware. Two rule kinds are -//! supported: endpoint *classes* (glob path → per-client limit) and per -//! *identifier* endpoints (limit by a value read from the request body). The -//! counter backend is pluggable via [`RateLimitStore`]; the default is an -//! in-process store, with an optional Redis store for multi-instance setups. +//! The proxy runs embedded on each service instance, so every decision is made +//! locally with a per-instance GCRA shaper ([`store::GcraStore`]) that adds no +//! blocking latency to the request path. When a shared store is configured, a +//! background task reconciles counters across instances asynchronously to +//! approximate a fleet-wide limit; the request path never blocks on it. +//! +//! A request is limited by the first [rule](matcher::CompiledRule) whose glob +//! matches its path. The rule's key selects *who* is limited (client IP, a +//! header value, or a validated JWT claim) and its profile selects *how much*. +pub mod gcra; +#[cfg(feature = "redis")] +pub mod global; pub mod matcher; pub mod rate; +pub mod resolve; pub mod store; +pub mod window; +use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; -use axum::body::Body; -use axum::extract::State; +use axum::extract::Request; use axum::http::{HeaderMap, StatusCode}; use axum::middleware::Next; use axum::response::{IntoResponse, Response}; use axum::Json; use crate::config::ShieldConfig; -use matcher::{EndpointClass, IdentifierEndpoint}; -use store::{Decision, MemoryStore, RateLimitStore}; - -/// Maximum request body buffered to read an identifier field (256 KiB). -const MAX_IDENTIFIER_BODY: usize = 256 * 1024; +use gcra::Verdict; +use matcher::{CompiledProfile, CompiledRule, KeySource, Phase}; +use store::GcraStore; -/// Compiled Shield rules plus the counter store. +/// Compiled Shield rules, limit tiers, and the local GCRA store. pub struct Shield { - store: Arc, - classes: Vec, - identifiers: Vec, + rules: Vec, + profiles: HashMap, + /// Applied when a matched rule resolves no other limit. + default_profile: Option, + /// Resolve a key's limit from validated JWT claims, when configured. + jwt_limits: Option, + /// Resolve a key's limit from an external service (cached, async). + limit_service: Option>, + /// Cross-instance reconciliation of the fleet-wide view (async, off the hot + /// path). Present only when a shared store is configured and compiled in. + #[cfg(feature = "redis")] + global: Option>, + store: GcraStore, /// CIDR ranges whose `X-Forwarded-For` / `X-Real-IP` headers we trust. trusted_proxies: Vec, } @@ -39,25 +56,28 @@ pub struct Shield { impl Shield { /// Build a Shield from config, or `None` when disabled / has no rules. /// - /// Uses a Redis store when `redis_url` is set and the `redis` feature is - /// compiled in; otherwise an in-process [`MemoryStore`]. - /// /// # Errors - /// Returns an error string when a glob pattern or rate fails to compile, or - /// when the configured Redis backend cannot be reached. + /// Returns an error string when a glob pattern, rate, profile reference, or + /// trusted-proxy CIDR fails to compile. pub fn build(config: &ShieldConfig) -> Result>, String> { if !config.enabled { return Ok(None); } - if config.endpoint_classes.is_empty() && config.identifier_endpoints.is_empty() { - tracing::warn!("shield enabled but no endpoint_classes or identifier_endpoints set"); + if config.rules.is_empty() { + tracing::warn!("shield enabled but no rules configured"); return Ok(None); } - let default_window = Duration::from_secs(config.window_secs.max(1)); - let classes = matcher::compile_endpoint_classes(&config.endpoint_classes, default_window)?; - let identifiers = - matcher::compile_identifier_endpoints(&config.identifier_endpoints, default_window)?; + let profiles = matcher::compile_profiles(&config.profiles)?; + let rules = matcher::compile_rules(&config.rules, &profiles)?; + + let default_profile = + match &config.default_profile { + Some(name) => Some(*profiles.get(name).ok_or_else(|| { + format!("default_profile references unknown profile {name:?}") + })?), + None => None, + }; let trusted_proxies = config .trusted_proxies @@ -65,107 +85,204 @@ impl Shield { .map(|s| parse_cidr(s)) .collect::, _>>()?; - let store = build_store(config)?; + let jwt_limits = config + .jwt_limits + .as_ref() + .map(resolve::JwtLimits::from_config); + let limit_service = match &config.limit_service { + Some(cfg) => Some(resolve::LimitService::build(cfg, profiles.clone())?), + None => None, + }; + + #[cfg(feature = "redis")] + let global = match &config.sync { + Some(sync) => { + let g = global::GlobalCounters::build( + &sync.redis_url, + Duration::from_millis(sync.interval_ms), + )?; + g.spawn(); + Some(g) + } + None => None, + }; + #[cfg(not(feature = "redis"))] + if config.sync.is_some() { + tracing::warn!( + "shield.sync is set but the `redis` feature is not compiled in; \ + staying local-only (per-instance limits)" + ); + } + Ok(Some(Arc::new(Self { - store, - classes, - identifiers, + rules, + profiles, + default_profile, + jwt_limits, + limit_service, + #[cfg(feature = "redis")] + global, + store: GcraStore::new(), trusted_proxies, }))) } - fn match_class(&self, path: &str) -> Option<&EndpointClass> { - self.classes.iter().find(|c| c.matcher.is_match(path)) + /// The first rule in `phase` whose glob matches `path`, with its index (the + /// index keys the store so distinct rules keep independent budgets). A path + /// may match one rule per phase; each phase enforces independently. + fn match_rule(&self, path: &str, phase: Phase) -> Option<(usize, &CompiledRule)> { + self.rules + .iter() + .enumerate() + .find(|(_, r)| r.phase == phase && r.matcher.is_match(path)) } - fn match_identifier(&self, path: &str) -> Option<&IdentifierEndpoint> { - self.identifiers.iter().find(|i| i.matcher.is_match(path)) + /// Resolve the limit tier for a matched rule, in priority order: the JWT + /// itself (validated claims), then the external service (cached), then the + /// rule's pinned profile, then the default profile. `None` means no limit + /// applies and the request passes unmetered. + fn resolve_limit( + &self, + rule: &CompiledRule, + claims: Option<&serde_json::Value>, + key: &str, + ) -> Option { + if let (Some(jwt), Some(claims)) = (&self.jwt_limits, claims) { + if let Some(profile) = jwt.resolve(claims, &self.profiles) { + return Some(profile); + } + } + if let Some(service) = &self.limit_service { + if let Some(profile) = service.resolve(key) { + return Some(profile); + } + } + self.static_profile(rule) } -} -/// Select the counter store from config. -fn build_store(config: &ShieldConfig) -> Result, String> { - match &config.redis_url { - Some(url) => open_redis(url), - None => Ok(Arc::new(MemoryStore::new())), + /// The rule's pinned profile, else the default profile. + fn static_profile(&self, rule: &CompiledRule) -> Option { + rule.profile + .as_ref() + .and_then(|name| self.profiles.get(name)) + .or(self.default_profile.as_ref()) + .copied() } } -#[cfg(feature = "redis")] -fn open_redis(url: &str) -> Result, String> { - store::RedisStore::open(url) - .map(|s| Arc::new(s) as Arc) - .map_err(|e| format!("invalid Redis URL for rate-limit store: {e}")) -} - -#[cfg(not(feature = "redis"))] -fn open_redis(_url: &str) -> Result, String> { - tracing::warn!( - "shield.redis_url is set but the `redis` feature is not compiled in; \ - falling back to the in-process store (per-replica limits only)" - ); - Ok(Arc::new(MemoryStore::new())) +/// Pre-auth middleware: enforces rules that need no validated claims (IP / +/// header keys). Layered outside auth so anonymous floods are shed before any +/// signature verification. +pub async fn pre_auth_middleware( + axum::extract::State(shield): axum::extract::State>, + request: Request, + next: Next, +) -> Response { + enforce(&shield, Phase::PreAuth, request, next).await } -/// Axum middleware enforcing the compiled Shield rules. -pub async fn middleware( - State(shield): State>, +/// Post-auth middleware: enforces rules keyed by a validated JWT claim. Layered +/// inside auth so the verified claims are available on the request. +pub async fn post_auth_middleware( + axum::extract::State(shield): axum::extract::State>, request: Request, next: Next, ) -> Response { - let path = request.uri().path().to_string(); + enforce(&shield, Phase::PostAuth, request, next).await +} + +/// Match a phase's rule for the request, apply its limit, and attach headers. +async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> Response { + let path = request.uri().path(); + let Some((idx, rule)) = shield.match_rule(path, phase) else { + return next.run(request).await; + }; + let peer = request .extensions() .get::>() .map(|ci| ci.0.ip()); let client = client_ip(peer, request.headers(), &shield.trusted_proxies); + let claims = request + .extensions() + .get::() + .map(|c| c.0.as_ref()); + let key = rule_key(idx, &rule.key, &client, request.headers(), claims); + + let Some(profile) = shield.resolve_limit(rule, claims, &key) else { + // No limit resolves for this rule (JWT/service/profile/default all + // absent): allow the request unmetered. + return next.run(request).await; + }; - // The decision whose budget we surface to the client via headers. - let mut report = None; - - // Endpoint-class limit (per client IP), no body needed. - if let Some(class) = shield.match_class(&path) { - let key = format!("class:{}:{client}", class.class); - let decision = shield.store.hit(&key, &class.rate).await; - if !decision.allowed { - return too_many_requests(decision); + // Fleet gate first (read-only, cached), so the local shaper is not charged + // for a request the fleet-wide budget will reject. + #[cfg(feature = "redis")] + if let Some(global) = &shield.global { + if !global.gate(&key, profile.limit, profile.window) { + return global_reject(profile.limit, profile.window); } - report = Some(decision); } - // Per-identifier limit: buffer the body, read the field, then restore it. - let request = if let Some(id_ep) = shield.match_identifier(&path) { - let (parts, body) = request.into_parts(); - let bytes = match axum::body::to_bytes(body, MAX_IDENTIFIER_BODY).await { - Ok(b) => b, - Err(_) => return payload_too_large(), - }; - // Key by the identifier value, or fall back to the client IP when the - // field is absent so the limit can't be bypassed by omitting it. - let key = match extract_body_field(&bytes, &id_ep.body_field) { - Some(ident) => format!("id:{path}:{}:{ident}", id_ep.body_field), - None => format!("id:{path}:{}:ip:{client}", id_ep.body_field), - }; - let decision = shield.store.hit(&key, &id_ep.rate).await; - if !decision.allowed { - return too_many_requests(decision); - } - // The identifier rule is more specific than a class rule, so report it. - report = Some(decision); - Request::from_parts(parts, Body::from(bytes)) - } else { - request - }; + let verdict = shield.store.check(&key, &profile.gcra); + if !verdict.allowed { + return too_many_requests(profile.limit, &verdict); + } - let mut response = next.run(request).await; - if let Some(decision) = report { - attach_rate_headers(response.headers_mut(), &decision); + // Record the admit for the next reconciliation push. + #[cfg(feature = "redis")] + if let Some(global) = &shield.global { + global.record(&key, profile.window); } + + let mut response = next.run(request).await; + attach_rate_headers(response.headers_mut(), profile.limit, &verdict); response } -/// Convenience alias for the axum request type used by the middleware. -type Request = axum::extract::Request; +/// A `429` for a request rejected by the fleet-wide gate. `Retry-After` / +/// `RateLimit-Reset` point at the end of the current window, when the sliding +/// estimate will have decayed. +#[cfg(feature = "redis")] +fn global_reject(limit: u64, window: Duration) -> Response { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or(Duration::ZERO); + let remaining = window.saturating_sub(window::elapsed_in_window(now, window)); + let verdict = Verdict { + allowed: false, + new_tat: Duration::ZERO, + remaining: 0, + retry_after: remaining, + reset_after: remaining, + }; + too_many_requests(limit, &verdict) +} + +/// Build the store key for a matched rule. The rule index namespaces the key so +/// two rules that happen to see the same client keep independent budgets. Every +/// source falls back to the client IP when its value is absent, so a limit can't +/// be dodged by omitting a header or authenticating anonymously. +fn rule_key( + idx: usize, + key: &KeySource, + client: &str, + headers: &HeaderMap, + claims: Option<&serde_json::Value>, +) -> String { + let by_ip = || format!("{idx}:ip:{client}"); + match key { + KeySource::Ip => by_ip(), + KeySource::Header(name) => match header_str(headers, name) { + Some(v) => format!("{idx}:hdr:{v}"), + None => by_ip(), + }, + KeySource::JwtClaim(claim) => match claims.and_then(|c| resolve::claim_str(c, claim)) { + Some(v) => format!("{idx}:jwt:{v}"), + None => by_ip(), + }, + } +} /// Parse a trusted-proxy entry as a CIDR range, accepting a bare IP as a /32 /// or /128 host range. @@ -252,33 +369,21 @@ fn header_str(headers: &HeaderMap, name: &str) -> Option { .map(str::to_string) } -/// Read a (possibly dotted) field from a JSON body as a string identifier. -fn extract_body_field(bytes: &[u8], field: &str) -> Option { - let value: serde_json::Value = serde_json::from_slice(bytes).ok()?; - let mut cur = &value; - for seg in field.split('.') { - cur = cur.get(seg)?; +/// Attach the draft-ietf `RateLimit-*` headers describing the remaining budget. +fn attach_rate_headers(headers: &mut HeaderMap, limit: u64, verdict: &Verdict) { + if let Ok(v) = limit.to_string().parse() { + headers.insert("ratelimit-limit", v); } - match cur { - serde_json::Value::String(s) => Some(s.clone()), - serde_json::Value::Number(n) => Some(n.to_string()), - serde_json::Value::Bool(b) => Some(b.to_string()), - _ => None, + if let Ok(v) = verdict.remaining.to_string().parse() { + headers.insert("ratelimit-remaining", v); } -} - -/// Add `X-RateLimit-*` headers describing the remaining budget. -fn attach_rate_headers(headers: &mut HeaderMap, decision: &Decision) { - if let Ok(v) = decision.limit.to_string().parse() { - headers.insert("x-ratelimit-limit", v); - } - if let Ok(v) = decision.remaining.to_string().parse() { - headers.insert("x-ratelimit-remaining", v); + if let Ok(v) = secs_ceil(verdict.reset_after).to_string().parse() { + headers.insert("ratelimit-reset", v); } } -fn too_many_requests(decision: Decision) -> Response { - let retry = decision.retry_after.unwrap_or(Duration::ZERO).as_secs(); +/// A `429` response carrying the rate-limit headers plus `Retry-After`. +fn too_many_requests(limit: u64, verdict: &Verdict) -> Response { let mut response = ( StatusCode::TOO_MANY_REQUESTS, Json(serde_json::json!({ @@ -288,251 +393,19 @@ fn too_many_requests(decision: Decision) -> Response { ) .into_response(); let headers = response.headers_mut(); - attach_rate_headers(headers, &decision); - if let Ok(v) = retry.to_string().parse() { + attach_rate_headers(headers, limit, verdict); + if let Ok(v) = secs_ceil(verdict.retry_after).to_string().parse() { headers.insert("retry-after", v); } response } -fn payload_too_large() -> Response { - ( - StatusCode::PAYLOAD_TOO_LARGE, - Json(serde_json::json!({ - "error": "INVALID_ARGUMENT", - "message": "request body too large for rate-limit identifier extraction", - })), - ) - .into_response() +/// Whole seconds, rounded up, for `Retry-After` / `RateLimit-Reset` (never report +/// `0` for a non-zero wait). +fn secs_ceil(d: Duration) -> u64 { + let millis = d.as_millis(); + u64::try_from(millis.div_ceil(1000)).unwrap_or(u64::MAX) } #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn client_ip_trusts_headers_only_from_trusted_peer() { - let mut h = HeaderMap::new(); - h.insert("x-forwarded-for", "203.0.113.7, 10.0.0.1".parse().unwrap()); - let trusted = vec![parse_cidr("10.0.0.0/8").unwrap()]; - let lb: std::net::IpAddr = "10.0.0.1".parse().unwrap(); - let direct: std::net::IpAddr = "198.51.100.9".parse().unwrap(); - - // From a trusted proxy: trust the forwarded first hop. - assert_eq!(client_ip(Some(lb), &h, &trusted), "203.0.113.7"); - // From an untrusted direct client: ignore the spoofable header, key by peer. - assert_eq!(client_ip(Some(direct), &h, &trusted), "198.51.100.9"); - } - - #[test] - fn client_ip_uses_rightmost_untrusted_hop() { - // Appending LBs (nginx proxy_add_x_forwarded_for, AWS ALB, GCP LB) add - // the connecting IP on the RIGHT, so the leftmost entry is attacker - // controlled. The real client is the rightmost hop outside trusted ranges. - let mut h = HeaderMap::new(); - h.insert("x-forwarded-for", "1.1.1.1, 203.0.113.7".parse().unwrap()); - let trusted = vec![parse_cidr("10.0.0.0/8").unwrap()]; - let lb: std::net::IpAddr = "10.0.0.5".parse().unwrap(); - assert_eq!(client_ip(Some(lb), &h, &trusted), "203.0.113.7"); - } - - #[test] - fn client_ip_without_peer_info_uses_headers_then_unknown() { - let mut h = HeaderMap::new(); - h.insert("x-real-ip", "198.51.100.2".parse().unwrap()); - assert_eq!(client_ip(None, &h, &[]), "198.51.100.2"); - assert_eq!(client_ip(None, &HeaderMap::new(), &[]), "unknown"); - } - - #[test] - fn extract_body_field_reads_string_and_dotted() { - let body = br#"{"email":"a@b.com","nested":{"id":42}}"#; - assert_eq!( - extract_body_field(body, "email"), - Some("a@b.com".to_string()) - ); - assert_eq!( - extract_body_field(body, "nested.id"), - Some("42".to_string()) - ); - assert_eq!(extract_body_field(body, "missing"), None); - assert_eq!(extract_body_field(b"not json", "email"), None); - } - - // --- middleware enforcement (real router) --- - - use crate::config::{EndpointClassConfig, IdentifierEndpointConfig, ShieldConfig}; - use axum::http::Request as HttpRequest; - use tower::ServiceExt; - - fn shield_config( - classes: Vec, - ids: Vec, - ) -> ShieldConfig { - ShieldConfig { - enabled: true, - endpoint_classes: classes, - identifier_endpoints: ids, - window_secs: 60, - redis_url: None, - trusted_proxies: Vec::new(), - } - } - - fn app(shield: Arc) -> axum::Router { - axum::Router::new() - .route("/limited", axum::routing::get(|| async { "ok" })) - .route("/login", axum::routing::post(|| async { "ok" })) - .layer(axum::middleware::from_fn_with_state(shield, middleware)) - } - - #[tokio::test] - async fn middleware_blocks_after_endpoint_class_limit() { - let shield = Shield::build(&shield_config( - vec![EndpointClassConfig { - pattern: "/limited".into(), - class: "t".into(), - rate: "2/min".into(), - }], - vec![], - )) - .unwrap() - .unwrap(); - let app = app(shield); - - let get = || HttpRequest::get("/limited").body(Body::empty()).unwrap(); - // Two allowed (no client IP header → all share the "unknown" bucket). - assert_eq!(app.clone().oneshot(get()).await.unwrap().status(), 200); - let second = app.clone().oneshot(get()).await.unwrap(); - assert_eq!(second.status(), 200); - assert_eq!(second.headers()["x-ratelimit-remaining"], "0"); - // Third over the limit → 429 with Retry-After. - let third = app.clone().oneshot(get()).await.unwrap(); - assert_eq!(third.status(), StatusCode::TOO_MANY_REQUESTS); - assert!(third.headers().contains_key("retry-after")); - } - - #[tokio::test] - async fn middleware_limits_per_identifier_value() { - let shield = Shield::build(&shield_config( - vec![], - vec![IdentifierEndpointConfig { - path: "/login".into(), - body_field: "email".into(), - rate: "1/min".into(), - }], - )) - .unwrap() - .unwrap(); - let app = app(shield); - - let post = |email: &str| { - HttpRequest::post("/login") - .header("content-type", "application/json") - .body(Body::from(format!(r#"{{"email":"{email}"}}"#))) - .unwrap() - }; - - // First request for alice is allowed, second is blocked. - assert_eq!( - app.clone().oneshot(post("alice")).await.unwrap().status(), - 200 - ); - assert_eq!( - app.clone().oneshot(post("alice")).await.unwrap().status(), - StatusCode::TOO_MANY_REQUESTS - ); - // A different identifier has its own budget. - assert_eq!( - app.clone().oneshot(post("bob")).await.unwrap().status(), - 200 - ); - } - - #[tokio::test] - async fn identifier_response_carries_ratelimit_headers() { - let shield = Shield::build(&shield_config( - vec![], - vec![IdentifierEndpointConfig { - path: "/login".into(), - body_field: "email".into(), - rate: "5/min".into(), - }], - )) - .unwrap() - .unwrap(); - let app = app(shield); - - let resp = app - .oneshot( - HttpRequest::post("/login") - .header("content-type", "application/json") - .body(Body::from(r#"{"email":"a"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(resp.status(), 200); - // Headers come from the identifier decision, not just class decisions. - assert_eq!(resp.headers()["x-ratelimit-limit"], "5"); - assert_eq!(resp.headers()["x-ratelimit-remaining"], "4"); - } - - #[tokio::test] - async fn identifier_limit_not_bypassed_when_field_absent() { - // Omitting the identifier field must not skip the limit: requests with - // no extractable identifier fall back to a client-keyed counter. - let shield = Shield::build(&shield_config( - vec![], - vec![IdentifierEndpointConfig { - path: "/login".into(), - body_field: "email".into(), - rate: "1/min".into(), - }], - )) - .unwrap() - .unwrap(); - let app = app(shield); - - let post_no_email = || { - HttpRequest::post("/login") - .header("content-type", "application/json") - .body(Body::from("{}")) - .unwrap() - }; - assert_eq!( - app.clone().oneshot(post_no_email()).await.unwrap().status(), - 200 - ); - // Second request without the field is still counted → blocked. - assert_eq!( - app.clone().oneshot(post_no_email()).await.unwrap().status(), - StatusCode::TOO_MANY_REQUESTS - ); - } - - #[tokio::test] - async fn middleware_ignores_unmatched_paths() { - let shield = Shield::build(&shield_config( - vec![EndpointClassConfig { - pattern: "/limited".into(), - class: "t".into(), - rate: "1/min".into(), - }], - vec![], - )) - .unwrap() - .unwrap(); - let app = app(shield); - - // /login is not covered by any rule → never limited. - for _ in 0..5 { - let resp = app - .clone() - .oneshot(HttpRequest::post("/login").body(Body::empty()).unwrap()) - .await - .unwrap(); - assert_eq!(resp.status(), 200); - } - } -} +mod tests; diff --git a/src/shield/resolve.rs b/src/shield/resolve.rs new file mode 100644 index 0000000..58e1c58 --- /dev/null +++ b/src/shield/resolve.rs @@ -0,0 +1,306 @@ +//! Limit resolution: derive a key's `{rate, burst}` from the JWT itself or from +//! an external service, on top of the static config profiles. +//! +//! The resolution chain is `jwt → service → rule profile → default`. JWT +//! resolution is synchronous (the validated claims ride on the request). Service +//! resolution never blocks the request path: a lookup is cached and refreshed in +//! the background (stale-while-revalidate), so a request either uses a cached +//! limit or falls through to the static profile while the fetch happens. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use dashmap::mapref::entry::Entry; +use serde::Deserialize; +use serde_json::Value; + +use super::gcra::{Gcra, Profile}; +use super::matcher::CompiledProfile; +use crate::config::{JwtLimitConfig, LimitServiceConfig}; + +/// Bare / per-minute rates use this window. +const PER_MINUTE: Duration = Duration::from_secs(60); + +/// Resolve a (possibly dotted) claim path to a scalar string value. +pub(super) fn claim_str(claims: &Value, path: &str) -> Option { + match claim_at(claims, path)? { + Value::String(s) => Some(s.clone()), + Value::Number(n) => Some(n.to_string()), + Value::Bool(b) => Some(b.to_string()), + _ => None, + } +} + +/// Resolve a (possibly dotted) claim path to an unsigned integer, accepting a +/// numeric or a numeric-string value. +fn claim_u64(claims: &Value, path: &str) -> Option { + match claim_at(claims, path)? { + Value::Number(n) => n.as_u64(), + Value::String(s) => s.trim().parse().ok(), + _ => None, + } +} + +fn claim_at<'a>(claims: &'a Value, path: &str) -> Option<&'a Value> { + let mut cur = claims; + for seg in path.split('.') { + cur = cur.get(seg)?; + } + Some(cur) +} + +/// Build a limit tier from explicit per-minute numbers. +fn profile_from_numbers(rpm: u64, burst: u64) -> CompiledProfile { + let rate = rpm.max(1); + let gcra = Gcra::from_profile(Profile { + rate, + window: PER_MINUTE, + burst: burst.max(1), + }); + CompiledProfile { + gcra, + limit: rate, + window: PER_MINUTE, + } +} + +/// Compiled JWT-based limit resolution: the claim names to read from the token. +#[derive(Debug, Clone)] +pub struct JwtLimits { + tier_claim: String, + rpm_claim: String, + burst_claim: String, +} + +impl JwtLimits { + /// Compile from config. + pub fn from_config(cfg: &JwtLimitConfig) -> Self { + Self { + tier_claim: cfg.tier_claim.clone(), + rpm_claim: cfg.rpm_claim.clone(), + burst_claim: cfg.burst_claim.clone(), + } + } + + /// Resolve a limit from the validated claims: a tier-name claim naming a + /// profile takes precedence, then explicit `rpm` (+ optional `burst`) + /// numbers. `None` when the token carries no limit hints. + pub fn resolve( + &self, + claims: &Value, + profiles: &HashMap, + ) -> Option { + if let Some(tier) = claim_str(claims, &self.tier_claim) { + if let Some(profile) = profiles.get(&tier) { + return Some(*profile); + } + } + if let Some(rpm) = claim_u64(claims, &self.rpm_claim) { + let burst = claim_u64(claims, &self.burst_claim).unwrap_or(rpm); + return Some(profile_from_numbers(rpm, burst)); + } + None + } +} + +/// External limit-service response: a tier name or explicit numbers. +#[derive(Debug, Deserialize)] +struct LimitResponse { + #[serde(default)] + tier: Option, + #[serde(default)] + rate_per_min: Option, + #[serde(default)] + burst: Option, +} + +/// A cached resolution. `profile` is `None` when the service reported no limit +/// for the key (a negative cache entry stops us re-querying every request). +#[derive(Clone, Copy)] +struct Cached { + profile: Option, + at: Instant, +} + +/// External limit-resolution service with an async, non-blocking cache. +pub struct LimitService { + endpoint: String, + ttl: Duration, + client: reqwest::Client, + /// Profiles for mapping a returned tier name to a compiled limit. + profiles: HashMap, + cache: dashmap::DashMap, + /// Keys with a background fetch already in flight (dedupes refreshes). + inflight: dashmap::DashMap, +} + +impl LimitService { + /// Build the service client from config and the compiled profiles. + /// + /// # Errors + /// Returns an error string when the HTTP client cannot be constructed. + pub fn build( + cfg: &LimitServiceConfig, + profiles: HashMap, + ) -> Result, String> { + let client = reqwest::Client::builder() + .timeout(Duration::from_millis(cfg.timeout_ms.max(1))) + .tls_backend_preconfigured(crate::auth::jwks::build_tls_config()) + .build() + .map_err(|e| format!("invalid limit_service client: {e}"))?; + Ok(Arc::new(Self { + endpoint: cfg.endpoint.clone(), + ttl: Duration::from_secs(cfg.ttl_secs.max(1)), + client, + profiles, + cache: dashmap::DashMap::new(), + inflight: dashmap::DashMap::new(), + })) + } + + /// Resolve `key`'s limit from the cache, serving a stale value while a + /// background refresh runs. Returns `None` (fall through to the static + /// profile) only when nothing is cached yet. Never blocks the request. + pub fn resolve(self: &Arc, key: &str) -> Option { + let cached = self.cache.get(key).map(|c| *c); + match cached { + Some(c) if c.at.elapsed() < self.ttl => c.profile, + Some(c) => { + // Stale: serve the last value and refresh in the background. + self.trigger_refresh(key.to_string()); + c.profile + } + None => { + self.trigger_refresh(key.to_string()); + None + } + } + } + + /// Spawn a single background fetch for `key` (deduped by `inflight`). + fn trigger_refresh(self: &Arc, key: String) { + match self.inflight.entry(key.clone()) { + Entry::Occupied(_) => return, + Entry::Vacant(v) => { + v.insert(()); + } + } + let this = self.clone(); + tokio::spawn(async move { + if let Ok(resolved) = this.fetch(&key).await { + this.cache.insert( + key.clone(), + Cached { + profile: resolved, + at: Instant::now(), + }, + ); + } + // On error, leave any stale entry in place (fail-static, not fail-open). + this.inflight.remove(&key); + }); + } + + /// Query the service for `key` and map the response to a limit. `Ok(None)` + /// means the service reported no limit; `Err` means the lookup failed and the + /// cache should be left untouched. + async fn fetch(&self, key: &str) -> Result, ()> { + let url = reqwest::Url::parse_with_params(&self.endpoint, &[("key", key)]) + .map_err(|e| tracing::warn!("invalid limit-service endpoint: {e}"))?; + let resp = self + .client + .get(url) + .send() + .await + .map_err(|e| tracing::warn!("limit-service fetch failed: {e}"))?; + if !resp.status().is_success() { + // A 404 is a definitive "no limit for this key": cache it negatively. + if resp.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(None); + } + tracing::warn!("limit-service returned {}", resp.status()); + return Err(()); + } + let body: LimitResponse = resp + .json() + .await + .map_err(|e| tracing::warn!("limit-service response parse failed: {e}"))?; + Ok(self.map_response(body)) + } + + /// Map a service response to a compiled limit: a tier name resolves against + /// the configured profiles; otherwise explicit numbers apply. + fn map_response(&self, body: LimitResponse) -> Option { + if let Some(tier) = body.tier { + return self.profiles.get(&tier).copied(); + } + body.rate_per_min + .map(|rpm| profile_from_numbers(rpm, body.burst.unwrap_or(rpm))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn profiles() -> HashMap { + let mut m = HashMap::new(); + m.insert( + "premium".to_string(), + profile_from_numbers(1000, 100), // limit 1000 + ); + m + } + + fn jwt_limits() -> JwtLimits { + JwtLimits::from_config(&JwtLimitConfig { + tier_claim: "ratelimit_tier".to_string(), + rpm_claim: "ratelimit_rpm".to_string(), + burst_claim: "ratelimit_burst".to_string(), + }) + } + + #[test] + fn jwt_tier_name_maps_to_profile() { + let claims = serde_json::json!({ "ratelimit_tier": "premium" }); + let p = jwt_limits().resolve(&claims, &profiles()).unwrap(); + assert_eq!(p.limit, 1000); + } + + #[test] + fn jwt_direct_numbers_build_a_profile() { + let claims = serde_json::json!({ "ratelimit_rpm": 300, "ratelimit_burst": 30 }); + let p = jwt_limits().resolve(&claims, &profiles()).unwrap(); + assert_eq!(p.limit, 300); + } + + #[test] + fn jwt_tier_takes_precedence_over_numbers() { + let claims = serde_json::json!({ "ratelimit_tier": "premium", "ratelimit_rpm": 5 }); + let p = jwt_limits().resolve(&claims, &profiles()).unwrap(); + assert_eq!(p.limit, 1000); + } + + #[test] + fn jwt_unknown_tier_falls_through_to_numbers_then_none() { + // Unknown tier + no numbers → no resolution. + let claims = serde_json::json!({ "ratelimit_tier": "gold" }); + assert!(jwt_limits().resolve(&claims, &profiles()).is_none()); + // Unknown tier + numbers → numbers win. + let claims = serde_json::json!({ "ratelimit_tier": "gold", "ratelimit_rpm": 42 }); + assert_eq!( + jwt_limits().resolve(&claims, &profiles()).unwrap().limit, + 42 + ); + } + + #[test] + fn numeric_string_claims_are_accepted() { + let claims = serde_json::json!({ "ratelimit_rpm": "250" }); + assert_eq!( + jwt_limits().resolve(&claims, &profiles()).unwrap().limit, + 250 + ); + } +} diff --git a/src/shield/store.rs b/src/shield/store.rs index be72866..07af58d 100644 --- a/src/shield/store.rs +++ b/src/shield/store.rs @@ -1,105 +1,93 @@ -//! Rate-limit counter storage. +//! In-process GCRA state store. //! -//! [`RateLimitStore`] abstracts where per-key counters live. [`MemoryStore`] is -//! the default and keeps counters in-process (per replica). -// The `RedisStore` reference is an intra-doc link only when the `redis` feature -// (and thus the type) is compiled in; otherwise a plain code span, so `cargo -// doc` stays warning-free in every feature combination. -#![cfg_attr( - feature = "redis", - doc = "[`RedisStore`] (behind the `redis` feature) shares counters across replicas, which is what a multi-instance deployment behind a load balancer needs for correct global limits." -)] -#![cfg_attr( - not(feature = "redis"), - doc = "`RedisStore` (behind the `redis` feature) shares counters across replicas, which is what a multi-instance deployment behind a load balancer needs for correct global limits." -)] +//! Holds one value per key: the theoretical arrival time (TAT), as milliseconds +//! since the store's base instant. The check is synchronous and lock-free across +//! keys (a `DashMap` shard lock is held only for the single key being updated), +//! so it adds no measurable latency to the request path. +//! +//! This is always the authority for the local limit decision. Cross-instance +//! reconciliation, when enabled, sits beside it and is updated asynchronously; +//! it never turns this check into a blocking operation. -use std::time::Duration; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; -use super::rate::Rate; +use dashmap::mapref::entry::Entry; -/// Outcome of recording one request against a key. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct Decision { - /// Whether the request is within the limit. - pub allowed: bool, - /// The configured limit (for the `X-RateLimit-Limit` header). - pub limit: u64, - /// Requests remaining in the current window (0 once exceeded). - pub remaining: u64, - /// How long until the window resets, when the request is rejected. - pub retry_after: Option, -} +use super::gcra::{Gcra, Verdict}; -/// A backend that records request hits and decides whether each is allowed. -#[async_trait::async_trait] -pub trait RateLimitStore: Send + Sync { - /// Record one hit for `key` and return the limiting decision for `rate`. - /// - /// A store that cannot reach its backend should fail open (allow the - /// request) rather than reject legitimate traffic. - async fn hit(&self, key: &str, rate: &Rate) -> Decision; -} +/// Run eviction of drained entries at most once per this interval. +const SWEEP_INTERVAL_MS: u64 = 60_000; -/// In-process fixed-window counter store (per replica). -/// -/// Counters are not shared between replicas, so global limits only hold for a -/// single instance. -#[cfg_attr( - feature = "redis", - doc = "Use [`RedisStore`] for multi-instance deployments." -)] -#[cfg_attr( - not(feature = "redis"), - doc = "Use `RedisStore` (feature `redis`) for multi-instance deployments." -)] +/// In-process per-instance GCRA store. +// no-std: caller-provided Clock + spin/hashbrown map. #[derive(Debug)] -pub struct MemoryStore { - // no-std: caller-provided Clock + spin/hashbrown map. - windows: dashmap::DashMap, - base: std::time::Instant, - last_sweep_ms: std::sync::atomic::AtomicU64, -} - -#[derive(Debug, Clone, Copy)] -struct Window { - expires_at: std::time::Instant, - count: u64, +pub struct GcraStore { + /// key → TAT in milliseconds since `base`. + tats: dashmap::DashMap, + base: Instant, + last_sweep_ms: AtomicU64, } -/// Run eviction of expired entries at most once per this interval. -const SWEEP_INTERVAL_MS: u64 = 60_000; - -impl Default for MemoryStore { +impl Default for GcraStore { fn default() -> Self { Self::new() } } -impl MemoryStore { +impl GcraStore { /// Create an empty store. pub fn new() -> Self { Self { - windows: dashmap::DashMap::new(), - base: std::time::Instant::now(), - last_sweep_ms: std::sync::atomic::AtomicU64::new(0), + tats: dashmap::DashMap::new(), + base: Instant::now(), + last_sweep_ms: AtomicU64::new(0), } } - /// Drop every entry whose window has elapsed. - /// - /// Entries are reset lazily on access, but keys that are never hit again - /// would otherwise linger forever; with client-controlled key cardinality - /// (IP / identifier) that is an unbounded-growth / OOM risk. - fn evict_expired(&self, now: std::time::Instant) { - self.windows.retain(|_, w| w.expires_at > now); + /// Evaluate one request for `key` against `gcra`, recording the new TAT when + /// the request is admitted. + pub fn check(&self, key: &str, gcra: &Gcra) -> Verdict { + let now = self.now(); + self.maybe_sweep(now); + + match self.tats.entry(key.to_string()) { + Entry::Occupied(mut o) => { + let stored = Some(Duration::from_millis(*o.get())); + let verdict = gcra.check(stored, now); + if verdict.allowed { + *o.get_mut() = tat_millis(verdict.new_tat); + } + verdict + } + Entry::Vacant(v) => { + let verdict = gcra.check(None, now); + if verdict.allowed { + v.insert(tat_millis(verdict.new_tat)); + } + verdict + } + } } - /// Evict expired entries at most once per [`SWEEP_INTERVAL_MS`]; the first - /// thread past the interval claims the sweep so it stays O(n) infrequently. - fn maybe_sweep(&self, now: std::time::Instant) { - use std::sync::atomic::Ordering; - let now_ms = now.duration_since(self.base).as_millis() as u64; + /// Milliseconds elapsed since the store's base instant. + fn now(&self) -> Duration { + self.base.elapsed() + } + + /// Drop entries whose TAT is in the past: a key with `TAT <= now` has fully + /// drained and is indistinguishable from a first-seen key, so re-inserting it + /// on the next hit yields the same result. Without eviction, client-controlled + /// key cardinality (IP / principal) would grow the map without bound. + fn evict_drained(&self, now: Duration) { + let now_ms = tat_millis(now); + self.tats.retain(|_, tat| *tat > now_ms); + } + + /// Evict at most once per [`SWEEP_INTERVAL_MS`]; the first caller past the + /// interval claims the sweep so it stays an infrequent O(n) pass. + fn maybe_sweep(&self, now: Duration) { + let now_ms = tat_millis(now); let last = self.last_sweep_ms.load(Ordering::Relaxed); if now_ms.saturating_sub(last) < SWEEP_INTERVAL_MS { return; @@ -109,192 +97,63 @@ impl MemoryStore { .compare_exchange(last, now_ms, Ordering::Relaxed, Ordering::Relaxed) .is_ok() { - self.evict_expired(now); + self.evict_drained(now); } } -} - -#[async_trait::async_trait] -impl RateLimitStore for MemoryStore { - async fn hit(&self, key: &str, rate: &Rate) -> Decision { - let now = std::time::Instant::now(); - self.maybe_sweep(now); - let mut entry = self.windows.entry(key.to_string()).or_insert(Window { - expires_at: now + rate.window, - count: 0, - }); - - // Reset the counter once the current window has elapsed. - if now >= entry.expires_at { - entry.expires_at = now + rate.window; - entry.count = 0; - } - entry.count += 1; - - let count = entry.count; - let expires_at = entry.expires_at; - drop(entry); - - let allowed = count <= rate.limit; - Decision { - allowed, - limit: rate.limit, - remaining: rate.limit.saturating_sub(count), - retry_after: (!allowed).then(|| expires_at.saturating_duration_since(now)), - } + /// Number of live entries (test/introspection helper). + #[cfg(test)] + fn len(&self) -> usize { + self.tats.len() } } -/// Redis-backed fixed-window counter store, shared across replicas. -/// -/// The client is opened eagerly (URL validation) but the multiplexed -/// connection is established lazily on first use, so construction stays -/// synchronous and a Redis that is briefly unavailable at startup does not -/// block the proxy from booting. -#[cfg(feature = "redis")] -pub struct RedisStore { - client: redis::Client, - conn: tokio::sync::OnceCell, -} - -#[cfg(feature = "redis")] -impl RedisStore { - /// Open a Redis client for `url` (e.g. `redis://127.0.0.1/`). - /// - /// # Errors - /// Returns the underlying Redis error when the URL is invalid. - pub fn open(url: &str) -> redis::RedisResult { - Ok(Self { - client: redis::Client::open(url)?, - conn: tokio::sync::OnceCell::new(), - }) - } - - /// Get the shared multiplexed connection, establishing it on first call. - async fn connection(&self) -> redis::RedisResult { - self.conn - .get_or_try_init(|| self.client.get_multiplexed_async_connection()) - .await - .cloned() - } - - /// Allow the request when Redis is unreachable: an outage must not take the - /// proxy down. - fn fail_open(rate: &Rate, err: redis::RedisError) -> Decision { - tracing::warn!("rate-limit store unavailable, allowing request: {err}"); - Decision { - allowed: true, - limit: rate.limit, - remaining: rate.limit, - retry_after: None, - } - } -} - -#[cfg(feature = "redis")] -#[async_trait::async_trait] -impl RateLimitStore for RedisStore { - async fn hit(&self, key: &str, rate: &Rate) -> Decision { - let window_ms = rate.window.as_millis().max(1) as u64; - let mut conn = match self.connection().await { - Ok(c) => c, - Err(e) => return Self::fail_open(rate, e), - }; - - // INCR and (re)set the TTL atomically in one round trip. Doing them in - // separate commands risks an immortal key if INCR lands but EXPIRE does - // not; the PTTL<0 guard also re-arms the TTL on any key that somehow - // lost it, so a key can never accumulate increments forever. - let script = redis::Script::new( - r"local c = redis.call('INCR', KEYS[1]) - if redis.call('PTTL', KEYS[1]) < 0 then - redis.call('PEXPIRE', KEYS[1], ARGV[1]) - end - return {c, redis.call('PTTL', KEYS[1])}", - ); - let (count, pttl): (i64, i64) = - match script.key(key).arg(window_ms).invoke_async(&mut conn).await { - Ok(v) => v, - Err(e) => return Self::fail_open(rate, e), - }; - - let count = count.max(0) as u64; - let allowed = count <= rate.limit; - Decision { - allowed, - limit: rate.limit, - remaining: rate.limit.saturating_sub(count), - retry_after: (!allowed).then(|| Duration::from_millis(pttl.max(0) as u64)), - } - } +/// A `Duration` as whole milliseconds, saturating at `u64::MAX`. +fn tat_millis(d: Duration) -> u64 { + u64::try_from(d.as_millis()).unwrap_or(u64::MAX) } #[cfg(test)] mod tests { use super::*; - #[tokio::test] - async fn memory_store_evicts_expired_entries() { - let store = MemoryStore::new(); - let rate = Rate { - limit: 5, - window: Duration::from_millis(20), - }; - store.hit("a", &rate).await; - store.hit("b", &rate).await; - assert_eq!(store.windows.len(), 2); - - tokio::time::sleep(Duration::from_millis(30)).await; - // Both windows have elapsed; eviction reclaims them so the map can't - // grow without bound from keys that are never hit again. - store.evict_expired(std::time::Instant::now()); - assert_eq!(store.windows.len(), 0); - } - - #[tokio::test] - async fn memory_store_allows_then_blocks_within_window() { - let store = MemoryStore::new(); - let rate = Rate { - limit: 2, + fn gcra(rate: u64, burst: u64) -> Gcra { + Gcra::from_profile(super::super::gcra::Profile { + rate, window: Duration::from_secs(60), - }; - - let d1 = store.hit("k", &rate).await; - assert!(d1.allowed && d1.remaining == 1); - let d2 = store.hit("k", &rate).await; - assert!(d2.allowed && d2.remaining == 0); - let d3 = store.hit("k", &rate).await; - assert!(!d3.allowed); - assert_eq!(d3.remaining, 0); - assert!(d3.retry_after.is_some()); + burst, + }) } - #[tokio::test] - async fn memory_store_resets_after_window() { - let store = MemoryStore::new(); - let rate = Rate { - limit: 1, - window: Duration::from_millis(50), - }; + #[test] + fn admits_burst_then_blocks() { + let store = GcraStore::new(); + let g = gcra(60, 2); // 1/s, burst 2 + assert!(store.check("k", &g).allowed); + assert!(store.check("k", &g).allowed); + // Burst of 2 spent within the same millisecond window. + assert!(!store.check("k", &g).allowed); + } - assert!(store.hit("k", &rate).await.allowed); - assert!(!store.hit("k", &rate).await.allowed); - tokio::time::sleep(Duration::from_millis(60)).await; - // New window: the counter has reset. - assert!(store.hit("k", &rate).await.allowed); + #[test] + fn keys_are_independent() { + let store = GcraStore::new(); + let g = gcra(60, 1); // burst 1 + assert!(store.check("a", &g).allowed); + assert!(store.check("b", &g).allowed); + assert!(!store.check("a", &g).allowed); } - #[tokio::test] - async fn memory_store_isolates_keys() { - let store = MemoryStore::new(); - let rate = Rate { - limit: 1, - window: Duration::from_secs(60), - }; - assert!(store.hit("a", &rate).await.allowed); - // A different key has its own independent counter. - assert!(store.hit("b", &rate).await.allowed); - assert!(!store.hit("a", &rate).await.allowed); + #[test] + fn eviction_reclaims_drained_keys() { + let store = GcraStore::new(); + let g = gcra(600, 1); // 10/s → 100ms emission, burst 1 + store.check("a", &g); + store.check("b", &g); + assert_eq!(store.len(), 2); + // Force the base far into the past so both TATs are drained, then sweep. + std::thread::sleep(Duration::from_millis(120)); + store.evict_drained(store.now()); + assert_eq!(store.len(), 0); } } diff --git a/src/shield/tests.rs b/src/shield/tests.rs new file mode 100644 index 0000000..4e9a868 --- /dev/null +++ b/src/shield/tests.rs @@ -0,0 +1,339 @@ +//! Middleware-level tests driving the compiled Shield through axum + tower. + +use super::*; +use crate::config::{KeySourceConfig, LimitProfileConfig, RateRuleConfig, ShieldConfig}; + +use axum::body::Body; +use axum::http::{Request, StatusCode}; +use axum::routing::get; +use axum::Router; +use std::collections::HashMap; +use tower::ServiceExt; + +/// A ShieldConfig with the given profiles + rules, everything else off. +fn config(profiles: Vec<(&str, &str, Option)>, rules: Vec) -> ShieldConfig { + let profiles = profiles + .into_iter() + .map(|(name, rate, burst)| { + ( + name.to_string(), + LimitProfileConfig { + rate: rate.to_string(), + burst, + }, + ) + }) + .collect::>(); + ShieldConfig { + enabled: true, + profiles, + rules, + default_profile: None, + jwt_limits: None, + limit_service: None, + sync: None, + trusted_proxies: Vec::new(), + } +} + +fn rule(pattern: &str, key: KeySourceConfig, profile: Option<&str>) -> RateRuleConfig { + RateRuleConfig { + pattern: pattern.to_string(), + key, + profile: profile.map(str::to_string), + } +} + +fn app(cfg: ShieldConfig) -> Router { + let shield = Shield::build(&cfg).unwrap().unwrap(); + Router::new() + .route("/api/x", get(|| async { "ok" })) + .route("/open", get(|| async { "ok" })) + .layer(axum::middleware::from_fn_with_state( + shield, + pre_auth_middleware, + )) +} + +async fn get_req(app: &Router, path: &str, xff: Option<&str>) -> Response { + let mut req = Request::builder().uri(path); + if let Some(xff) = xff { + req = req.header("x-forwarded-for", xff); + } + app.clone() + .oneshot(req.body(Body::empty()).unwrap()) + .await + .unwrap() +} + +async fn get_keyed(app: &Router, path: &str, header: (&str, &str)) -> Response { + let req = Request::builder() + .uri(path) + .header(header.0, header.1) + .body(Body::empty()) + .unwrap(); + app.clone().oneshot(req).await.unwrap() +} + +#[tokio::test] +async fn admits_burst_then_rejects() { + let app = app(config( + vec![("t", "60/min", Some(2))], // 1/s, burst 2 + vec![rule("/api/**", KeySourceConfig::Ip, Some("t"))], + )); + // Same client (same XFF) → shared budget of 2. + assert_eq!( + get_req(&app, "/api/x", Some("9.9.9.9")).await.status(), + StatusCode::OK + ); + assert_eq!( + get_req(&app, "/api/x", Some("9.9.9.9")).await.status(), + StatusCode::OK + ); + let third = get_req(&app, "/api/x", Some("9.9.9.9")).await; + assert_eq!(third.status(), StatusCode::TOO_MANY_REQUESTS); + // A rejected response carries Retry-After. + assert!(third.headers().contains_key("retry-after")); +} + +#[tokio::test] +async fn emits_ratelimit_headers_when_allowed() { + let app = app(config( + vec![("t", "100/min", Some(10))], + vec![rule("/api/**", KeySourceConfig::Ip, Some("t"))], + )); + let resp = get_req(&app, "/api/x", Some("1.2.3.4")).await; + assert_eq!(resp.status(), StatusCode::OK); + let h = resp.headers(); + assert_eq!(h.get("ratelimit-limit").unwrap(), "100"); + // After one admitted request, 9 of the burst-10 remain. + assert_eq!(h.get("ratelimit-remaining").unwrap(), "9"); + assert!(h.contains_key("ratelimit-reset")); +} + +#[tokio::test] +async fn unmatched_path_is_not_limited() { + let app = app(config( + vec![("t", "1/min", Some(1))], + vec![rule("/api/**", KeySourceConfig::Ip, Some("t"))], + )); + // `/open` matches no rule: never limited even past the /api budget. + for _ in 0..5 { + assert_eq!( + get_req(&app, "/open", Some("5.5.5.5")).await.status(), + StatusCode::OK + ); + } +} + +#[tokio::test] +async fn header_key_isolates_clients() { + let app = app(config( + vec![("t", "60/min", Some(1))], // burst 1 + vec![rule( + "/api/**", + KeySourceConfig::Header { + name: "x-api-key".to_string(), + }, + Some("t"), + )], + )); + // Key "alice": first ok, second rejected. + assert_eq!( + get_keyed(&app, "/api/x", ("x-api-key", "alice")) + .await + .status(), + StatusCode::OK + ); + assert_eq!( + get_keyed(&app, "/api/x", ("x-api-key", "alice")) + .await + .status(), + StatusCode::TOO_MANY_REQUESTS + ); + // A different key has its own budget. + assert_eq!( + get_keyed(&app, "/api/x", ("x-api-key", "bob")) + .await + .status(), + StatusCode::OK + ); +} + +#[tokio::test] +async fn rule_without_resolvable_profile_passes() { + // Rule pins no profile and there is no default_profile → cannot limit → allow. + let app = app(config( + vec![("t", "1/min", Some(1))], + vec![rule("/api/**", KeySourceConfig::Ip, None)], + )); + for _ in 0..3 { + assert_eq!( + get_req(&app, "/api/x", Some("7.7.7.7")).await.status(), + StatusCode::OK + ); + } +} + +#[test] +fn jwt_claim_key_uses_claim_then_falls_back_to_ip() { + let claims = serde_json::json!({ "sub": "alice", "org": { "id": "acme" } }); + let key = KeySource::JwtClaim("sub".to_string()); + // Present claim keys by its value. + assert_eq!( + rule_key(0, &key, "1.1.1.1", &HeaderMap::new(), Some(&claims)), + "0:jwt:alice" + ); + // Dotted path into a nested claim. + let nested = KeySource::JwtClaim("org.id".to_string()); + assert_eq!( + rule_key(0, &nested, "1.1.1.1", &HeaderMap::new(), Some(&claims)), + "0:jwt:acme" + ); + // No claims (anonymous) → IP fallback, so the limit can't be dodged. + assert_eq!( + rule_key(0, &key, "1.1.1.1", &HeaderMap::new(), None), + "0:ip:1.1.1.1" + ); + // Claim present but missing the requested field → IP fallback. + assert_eq!( + rule_key( + 0, + &KeySource::JwtClaim("missing".to_string()), + "1.1.1.1", + &HeaderMap::new(), + Some(&claims) + ), + "0:ip:1.1.1.1" + ); +} + +#[test] +fn rule_index_namespaces_identical_keys() { + // The same client under two different rules keeps independent budgets. + let key = KeySource::Ip; + let a = rule_key(0, &key, "1.1.1.1", &HeaderMap::new(), None); + let b = rule_key(1, &key, "1.1.1.1", &HeaderMap::new(), None); + assert_ne!(a, b); +} + +#[test] +fn build_rejects_unknown_default_profile() { + let mut cfg = config(vec![("t", "1/min", None)], Vec::new()); + cfg.rules = vec![rule("/x", KeySourceConfig::Ip, Some("t"))]; + cfg.default_profile = Some("missing".to_string()); + assert!(Shield::build(&cfg).is_err()); +} + +/// End-to-end two-phase test: a rule keyed by a validated JWT claim, layered in +/// the same order as the server (pre-auth → auth → post-auth), limits per +/// principal using the claims the auth middleware verifies and attaches. +mod two_phase { + use super::*; + use crate::auth::Auth; + use crate::config::{AuthConfig, JwtConfig}; + use ed25519_dalek::{Signer, SigningKey}; + + fn keypair_pem() -> (SigningKey, std::path::PathBuf) { + let sk = SigningKey::from_bytes(&[9u8; 32]); + let spki_prefix: [u8; 12] = [ + 0x30, 0x2a, 0x30, 0x05, 0x06, 0x03, 0x2b, 0x65, 0x70, 0x03, 0x21, 0x00, + ]; + let mut der = spki_prefix.to_vec(); + der.extend_from_slice(sk.verifying_key().as_bytes()); + use base64::Engine; + let b64 = base64::engine::general_purpose::STANDARD.encode(&der); + let pem = format!("-----BEGIN PUBLIC KEY-----\n{b64}\n-----END PUBLIC KEY-----\n"); + use std::sync::atomic::{AtomicU32, Ordering}; + static N: AtomicU32 = AtomicU32::new(0); + let path = std::env::temp_dir().join(format!( + "sp_shield_{}_{}.pem", + std::process::id(), + N.fetch_add(1, Ordering::Relaxed) + )); + std::fs::write(&path, pem).unwrap(); + (sk, path) + } + + fn token(sk: &SigningKey, sub: &str) -> String { + use base64::engine::general_purpose::URL_SAFE_NO_PAD; + use base64::Engine; + let header = URL_SAFE_NO_PAD.encode(br#"{"alg":"EdDSA","typ":"JWT"}"#); + // Far-future exp so default expiry validation passes. + let claims = serde_json::json!({ "sub": sub, "exp": 9_999_999_999u64 }); + let payload = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&claims).unwrap()); + let signing_input = format!("{header}.{payload}"); + let sig = sk.sign(signing_input.as_bytes()); + format!("{signing_input}.{}", URL_SAFE_NO_PAD.encode(sig.to_bytes())) + } + + fn auth(pem: std::path::PathBuf) -> std::sync::Arc { + let cfg = AuthConfig { + mode: "jwt".into(), + jwt: Some(JwtConfig { + issuer: None, + audience: None, + jwks_uri: None, + public_key_pem_file: Some(pem), + claims_headers: HashMap::new(), + roles_claim: "roles".into(), + }), + forward_auth: None, + authz: None, + }; + Auth::build(&cfg).unwrap().unwrap() + } + + fn stack(shield: std::sync::Arc, auth: std::sync::Arc) -> Router { + // Same layering order as the server: post-auth is innermost (sees the + // verified claims), auth in the middle, pre-auth outermost. + Router::new() + .route("/api/x", get(|| async { "ok" })) + .layer(axum::middleware::from_fn_with_state( + shield.clone(), + post_auth_middleware, + )) + .layer(axum::middleware::from_fn_with_state( + auth, + crate::auth::middleware, + )) + .layer(axum::middleware::from_fn_with_state( + shield, + pre_auth_middleware, + )) + } + + async fn get_with(app: &Router, bearer: &str) -> StatusCode { + let req = Request::builder() + .uri("/api/x") + .header("authorization", format!("Bearer {bearer}")) + .body(Body::empty()) + .unwrap(); + app.clone().oneshot(req).await.unwrap().status() + } + + #[tokio::test] + async fn limits_per_principal_via_validated_claim() { + let (sk, pem) = keypair_pem(); + let cfg = config( + vec![("t", "60/min", Some(1))], // burst 1 per principal + vec![rule( + "/api/**", + KeySourceConfig::JwtClaim { + claim: "sub".to_string(), + }, + Some("t"), + )], + ); + let shield = Shield::build(&cfg).unwrap().unwrap(); + let app = stack(shield, auth(pem)); + + let alice = token(&sk, "alice"); + let bob = token(&sk, "bob"); + // Alice: first request ok, second over her burst. + assert_eq!(get_with(&app, &alice).await, StatusCode::OK); + assert_eq!(get_with(&app, &alice).await, StatusCode::TOO_MANY_REQUESTS); + // Bob is a different principal with his own budget. + assert_eq!(get_with(&app, &bob).await, StatusCode::OK); + } +} diff --git a/src/shield/window.rs b/src/shield/window.rs new file mode 100644 index 0000000..afcf16d --- /dev/null +++ b/src/shield/window.rs @@ -0,0 +1,64 @@ +//! Sliding-window counter math for the global (cross-instance) limit view. +//! +//! A fixed window has a boundary burst: a client can spend a full window's quota +//! just before the boundary and another just after. The sliding-window counter +//! smooths this by blending the current epoch's count with a decaying fraction +//! of the previous epoch's, which is why the fleet-wide gate has no boundary +//! burst even though the local shaper is GCRA. + +use std::time::Duration; + +/// The epoch (window index) for `now`, aligned to wall-clock so every instance +/// agrees on window boundaries. +pub fn epoch(now: Duration, window: Duration) -> u64 { + now.as_secs() / window.as_secs().max(1) +} + +/// How far into the current window `now` sits. +pub fn elapsed_in_window(now: Duration, window: Duration) -> Duration { + let w = window.as_secs().max(1); + Duration::from_secs(now.as_secs() % w) + Duration::from_nanos(u64::from(now.subsec_nanos())) +} + +/// Blend the current-epoch count with the previous epoch's, weighted by how much +/// of the current window remains: a request early in the window still "sees" +/// most of the previous window's traffic, decaying to none by the boundary. +pub fn sliding_estimate(cur: u64, prev: u64, elapsed: Duration, window: Duration) -> f64 { + let w = window.as_secs_f64().max(f64::MIN_POSITIVE); + let frac = (elapsed.as_secs_f64() / w).clamp(0.0, 1.0); + cur as f64 + prev as f64 * (1.0 - frac) +} + +#[cfg(test)] +mod tests { + use super::*; + + const W: Duration = Duration::from_secs(60); + + #[test] + fn epoch_advances_once_per_window() { + assert_eq!(epoch(Duration::from_secs(0), W), 0); + assert_eq!(epoch(Duration::from_secs(59), W), 0); + assert_eq!(epoch(Duration::from_secs(60), W), 1); + assert_eq!(epoch(Duration::from_secs(125), W), 2); + } + + #[test] + fn estimate_at_window_start_counts_full_previous() { + // At the very start of the window, the whole previous epoch still counts. + assert_eq!(sliding_estimate(0, 100, Duration::ZERO, W), 100.0); + } + + #[test] + fn estimate_at_window_end_drops_previous() { + // At the end of the window, the previous epoch has fully decayed. + let est = sliding_estimate(10, 100, W, W); + assert!((est - 10.0).abs() < 1e-9); + } + + #[test] + fn estimate_midway_halves_previous() { + let est = sliding_estimate(10, 100, Duration::from_secs(30), W); + assert!((est - 60.0).abs() < 1e-9); // 10 + 100 * 0.5 + } +} From bedba5d8dd411fa98a3126d92909f414dace36c4 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 11:27:20 +0300 Subject: [PATCH 02/55] test(shield): add regression test for GCRA idle-key remaining cap An idle key (TAT far in the past) reports a slack-derived remaining far above the bucket capacity. The test exercises drain-then-admit and expects remaining == burst - 1; it fails on current code. --- src/shield/gcra.rs | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/src/shield/gcra.rs b/src/shield/gcra.rs index 3ef1081..5e66878 100644 --- a/src/shield/gcra.rs +++ b/src/shield/gcra.rs @@ -228,6 +228,21 @@ mod tests { assert_eq!(last_reset, Duration::from_secs(5)); } + /// After a key has fully drained (idle far past its TAT), the next request's + /// reported `remaining` must not exceed the bucket: `burst - 1`, never an + /// inflated count derived from the long idle gap. + #[test] + fn idle_key_remaining_capped_to_burst() { + let g = Gcra::from_profile(profile(60, 60, 3)); // 1/s, burst 3 + let first = g.check(None, Duration::from_secs(100)); + assert!(first.allowed); + // Idle for a long time, then one request: remaining is burst-1, not the + // huge slack-derived value. + let v = g.check(Some(first.new_tat), Duration::from_secs(1000)); + assert!(v.allowed); + assert_eq!(v.remaining, 2, "remaining must be capped to burst-1"); + } + /// A bare-rate profile (burst 1) admits one request per emission interval /// with no instantaneous burst. #[test] From aa346331e522ca17c757412f4ddf014c0b48f020 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 11:28:16 +0300 Subject: [PATCH 03/55] fix(shield): cap GCRA admissible count at bucket capacity An idle key whose TAT sits far in the past produced a slack-derived admissible count well above burst, so RateLimit-Remaining over-reported the available budget (admission itself was already correct). Cap the count at burst so remaining never exceeds a full bucket. --- src/shield/gcra.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/shield/gcra.rs b/src/shield/gcra.rs index 5e66878..5c68c69 100644 --- a/src/shield/gcra.rs +++ b/src/shield/gcra.rs @@ -22,6 +22,9 @@ pub struct Gcra { /// may run. `(burst - 1) * emission_interval`, so a fresh key admits exactly /// `burst` requests instantly before throttling to the steady rate. tau: Duration, + /// Bucket capacity: the most requests admissible at any instant. Caps the + /// reported `remaining` so an idle key cannot report above a full bucket. + burst: u64, } /// A named limit tier: a sustained rate over a window plus an instantaneous @@ -68,6 +71,7 @@ impl Gcra { Self { emission_interval, tau, + burst, } } @@ -88,9 +92,11 @@ impl Gcra { // admitted request pushes TAT forward by `t`, and a request conforms // while `now >= tat + k*t - tau`. `now + tau >= tat` means at least one // conforms. + // Cap at the bucket: an idle key (TAT far in the past) yields a large + // slack, but no more than `burst` requests can ever be admissible at once. let admissible = if now + tau >= tat { let slack = (now + tau) - tat; // >= 0 - 1 + div_floor(slack, t) + (1 + div_floor(slack, t)).min(self.burst) } else { 0 }; From 186206413b1c0109e9361aba198c020eef485903 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 11:33:15 +0300 Subject: [PATCH 04/55] fix(shield): correct cross-instance counter accounting Three reconciliation bugs in the shared-store sync: - Push each pending delta with the epoch it was accumulated in, not the current epoch. Crossing a window boundary between a record and the tick previously misattributed the delta to the new epoch counter. - Commit the pushed amount immediately after a successful INCRBY, before the aggregate read. A read failure after a successful push previously left the delta uncommitted, so the next tick pushed it again (double count). - Reset a key state when its resolved window changes, since a different window is a different accounting unit. The shared counter is always updated with an atomic INCRBY of the delta (never a SET of a computed absolute), so concurrent instances' increments accumulate instead of clobbering. The read-failure-after-push path is guarded structurally (commit precedes read); simulating a mid-reconcile store failure needs fault injection the harness lacks, so it is covered by inspection plus an idempotency test asserting repeated passes do not re-push. --- src/shield/global.rs | 154 ++++++++++++++++++++++++++++++++----------- 1 file changed, 116 insertions(+), 38 deletions(-) diff --git a/src/shield/global.rs b/src/shield/global.rs index 3533daa..a9c80ba 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -101,31 +101,39 @@ impl GlobalCounters { /// store and does not record the admit (call [`record`](Self::record) after /// the local check also passes). pub fn gate(&self, key: &str, budget: u64, window: Duration) -> bool { - let now = unix_now(); - let ep = window::epoch(now, window); - let window_secs = window.as_secs().max(1); - let mut state = self - .states - .entry(key.to_string()) - .or_insert_with(|| KeyState::new(ep, window_secs)); - state.roll_to(ep); - state.last_seen = Instant::now(); + let state = self.state_for(key, window); state.remote_estimate + state.local_count < budget } /// Record one admitted request for `key`, to be pushed to the store on the /// next background tick. pub fn record(&self, key: &str, window: Duration) { + let mut state = self.state_for(key, window); + state.local_count += 1; + } + + /// Look up (or create) the per-key state, rolled to the current epoch and + /// stamped as seen. If the resolved `window` differs from the one the entry + /// was created with, the entry is reset: a different window is a different + /// accounting unit, so mixing counts across it would corrupt the estimate. + fn state_for( + &self, + key: &str, + window: Duration, + ) -> dashmap::mapref::one::RefMut<'_, String, KeyState> { let now = unix_now(); - let ep = window::epoch(now, window); let window_secs = window.as_secs().max(1); + let ep = window::epoch(now, window); let mut state = self .states .entry(key.to_string()) .or_insert_with(|| KeyState::new(ep, window_secs)); + if state.window_secs != window_secs { + *state = KeyState::new(ep, window_secs); + } state.roll_to(ep); - state.local_count += 1; state.last_seen = Instant::now(); + state } /// Spawn the background reconciliation loop. A no-op (with a warning) when @@ -152,26 +160,31 @@ impl GlobalCounters { .cloned() } - /// One reconciliation pass: push this instance's deltas, pull the aggregate, - /// refresh each key's remote estimate, and evict stale keys. + /// One reconciliation pass at the current wall clock. async fn reconcile(&self) { + self.reconcile_at(unix_now()).await; + } + + /// One reconciliation pass at instant `now` (parameterised for tests): push + /// this instance's deltas, pull the aggregate, refresh each key's remote + /// estimate, and evict stale keys. + async fn reconcile_at(&self, now: Duration) { self.evict_stale(); - // Phase 1 (locked, brief): snapshot each active key's push delta and the - // epoch keys to read. No store I/O while holding a lock. - let now = unix_now(); + // Phase 1 (locked, brief): snapshot each key's pending delta. The delta + // is pushed to the epoch it was accumulated in (`push_epoch`), while the + // estimate reads the current epoch (`read_epoch`); the two differ when a + // window boundary is crossed between a record and this tick. let mut plans: Vec = Vec::new(); for entry in self.states.iter() { - let key = entry.key().clone(); let s = entry.value(); let window = Duration::from_secs(s.window_secs); - let ep = window::epoch(now, window); - let delta = s.local_count.saturating_sub(s.pushed); plans.push(PushPlan { - key, - epoch: ep, + key: entry.key().clone(), + push_epoch: s.epoch, + read_epoch: window::epoch(now, window), window, - delta, + delta: s.local_count.saturating_sub(s.pushed), }); } if plans.is_empty() { @@ -186,11 +199,16 @@ impl GlobalCounters { } }; - // Phase 2 (no locks): push deltas, then read current+previous epochs. + // Push each pending delta as an atomic INCRBY (never a SET, so concurrent + // instances' increments accumulate rather than clobber). On success, + // commit `pushed` immediately so a later read failure cannot cause the + // same delta to be pushed a second time on the next tick. if let Err(e) = self.push_deltas(&mut conn, &plans).await { tracing::warn!("rate-limit delta push failed: {e}"); - return; + return; // `pushed` not advanced → same delta retried next tick, no double count } + self.commit_pushed(&plans); + let reads = match self.read_epochs(&mut conn, &plans).await { Ok(r) => r, Err(e) => { @@ -198,18 +216,33 @@ impl GlobalCounters { return; } }; + self.apply_estimates(now, &plans, &reads); + } - // Phase 3 (locked, brief): commit pushed deltas and the fresh estimate. - for (plan, (cur, prev)) in plans.iter().zip(reads) { - if let Some(mut s) = self.states.get_mut(&plan.key) { - if s.epoch == plan.epoch { - s.pushed += plan.delta; + /// Advance each key's `pushed` by the delta we just pushed, but only while + /// the state still holds the epoch the delta belonged to (a concurrent roll + /// resets the counts, so there is nothing to advance). + fn commit_pushed(&self, plans: &[PushPlan]) { + for p in plans.iter().filter(|p| p.delta > 0) { + if let Some(mut s) = self.states.get_mut(&p.key) { + if s.epoch == p.push_epoch { + s.pushed += p.delta; } - let elapsed = window::elapsed_in_window(now, plan.window); - let est = window::sliding_estimate(cur, prev, elapsed, plan.window); - // Exclude this instance's own current-epoch contribution: the - // gate adds `local_count` separately. - let others = (est.round() as i64 - s.pushed as i64).max(0); + } + } + } + + /// Refresh each key's cached estimate of the rest of the fleet's consumption. + fn apply_estimates(&self, now: Duration, plans: &[PushPlan], reads: &[(u64, u64)]) { + for (p, (cur, prev)) in plans.iter().zip(reads) { + if let Some(mut s) = self.states.get_mut(&p.key) { + let elapsed = window::elapsed_in_window(now, p.window); + let est = window::sliding_estimate(*cur, *prev, elapsed, p.window); + // Subtract only our own contribution to the current epoch's count + // (the gate adds `local_count` back separately). If the key's + // pushes went to an earlier epoch, our current-epoch share is 0. + let my_current = if s.epoch == p.read_epoch { s.pushed } else { 0 }; + let others = (est.round() as i64 - my_current as i64).max(0); s.remote_estimate = others as u64; } } @@ -225,7 +258,7 @@ impl GlobalCounters { let mut any = false; for p in plans.iter().filter(|p| p.delta > 0) { any = true; - let k = epoch_key(&p.key, p.epoch); + let k = epoch_key(&p.key, p.push_epoch); let ttl_ms = (p.window.as_millis() as u64).saturating_mul(2).max(1); pipe.cmd("INCRBY").arg(&k).arg(p.delta).ignore(); pipe.cmd("PEXPIRE").arg(&k).arg(ttl_ms).ignore(); @@ -245,8 +278,8 @@ impl GlobalCounters { ) -> redis::RedisResult> { let mut keys: Vec = Vec::with_capacity(plans.len() * 2); for p in plans { - keys.push(epoch_key(&p.key, p.epoch)); - keys.push(epoch_key(&p.key, p.epoch.saturating_sub(1))); + keys.push(epoch_key(&p.key, p.read_epoch)); + keys.push(epoch_key(&p.key, p.read_epoch.saturating_sub(1))); } let vals: Vec> = redis::cmd("MGET").arg(&keys).query_async(conn).await?; Ok(plans @@ -271,7 +304,10 @@ impl GlobalCounters { /// A per-key push plan captured under lock, used without holding locks. struct PushPlan { key: String, - epoch: u64, + /// Epoch the pending delta was accumulated in (target of the `INCRBY`). + push_epoch: u64, + /// Current epoch, whose sliding window the estimate reads. + read_epoch: u64, window: Duration, delta: u64, } @@ -366,4 +402,46 @@ mod tests { "B must reject once the combined budget is reached" ); } + + /// Repeated reconcile passes with no new admits must not re-push the same + /// delta: the shared counter reflects each admit exactly once, even if a + /// pass's aggregate read had failed on an earlier tick. + #[tokio::test] + async fn repeated_reconcile_does_not_double_push() { + let Ok(url) = std::env::var("SHIELD_REDIS_TEST_URL") else { + eprintln!( + "SKIP repeated_reconcile_does_not_double_push: SHIELD_REDIS_TEST_URL not set" + ); + return; + }; + // Hour-long window so no epoch boundary is crossed mid-test. + let win = Duration::from_secs(3600); + let nonce = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let key = format!("it2:{}:{nonce}", std::process::id()); + + let a = GlobalCounters::build(&url, Duration::from_millis(200)).unwrap(); + a.record(&key, win); + a.record(&key, win); + // Three passes: only the first has a non-zero delta to push. + a.reconcile().await; + a.reconcile().await; + a.reconcile().await; + + let epoch = window::epoch(unix_now(), win); + let redis_key = epoch_key(&key, epoch); + let client = redis::Client::open(url).unwrap(); + let mut conn = client.get_multiplexed_async_connection().await.unwrap(); + let count: i64 = redis::cmd("GET") + .arg(&redis_key) + .query_async(&mut conn) + .await + .unwrap_or(0); + assert_eq!( + count, 2, + "shared counter must reflect the 2 admits exactly once" + ); + } } From ebf6e81ee9ca49984441cb182275c68ed96ab9ce Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 11:37:58 +0300 Subject: [PATCH 05/55] fix(shield): de-identify store keys and namespace by rule fingerprint Store and shared-counter keys previously embedded the raw client value (API key, principal, IP), so secrets could land in the shared store's keyspace and its logs. Hash the value into the key instead; the hash is stable across instances so reconciliation keys still agree. The limit service still receives the raw identity, since it resolves a tier per real principal. Also key by a stable per-rule fingerprint (hash of pattern + key + profile) rather than the rule's list index, so reordering rules or a rolling deploy no longer resets budgets. --- src/shield/matcher.rs | 25 ++++++++++++++ src/shield/mod.rs | 78 +++++++++++++++++++++++++++---------------- src/shield/tests.rs | 56 ++++++++++++++++++------------- 3 files changed, 107 insertions(+), 52 deletions(-) diff --git a/src/shield/matcher.rs b/src/shield/matcher.rs index 0d123e5..a9f0536 100644 --- a/src/shield/matcher.rs +++ b/src/shield/matcher.rs @@ -50,6 +50,19 @@ pub struct CompiledRule { /// Static profile name, if the rule pins one. pub profile: Option, pub phase: Phase, + /// Stable identifier for this rule, derived from its shape (pattern + key + + /// profile) rather than its position. Namespaces store keys so reordering + /// rules or a rolling deploy does not reset budgets. + pub fingerprint: String, +} + +/// A short, stable, non-reversible hash (128-bit hex) of `input`. Used both to +/// fingerprint rules and to de-identify key values before they reach the shared +/// store, and deterministic across instances so reconciliation keys agree. +pub fn short_hash(input: &str) -> String { + use sha2::{Digest, Sha256}; + let digest = Sha256::digest(input.as_bytes()); + digest[..16].iter().map(|b| format!("{b:02x}")).collect() } /// Build a glob matcher where `*` stays within a path segment and `**` spans @@ -116,11 +129,23 @@ pub fn compile_rules( KeySource::JwtClaim(_) => Phase::PostAuth, KeySource::Ip | KeySource::Header(_) => Phase::PreAuth, }; + let key_repr = match &key { + KeySource::Ip => "ip".to_string(), + KeySource::Header(name) => format!("hdr:{name}"), + KeySource::JwtClaim(claim) => format!("jwt:{claim}"), + }; + let fingerprint = short_hash(&format!( + "{}\u{1f}{}\u{1f}{}", + c.pattern, + key_repr, + c.profile.as_deref().unwrap_or("") + )); Ok(CompiledRule { matcher: path_glob(&c.pattern)?, key, profile: c.profile.clone(), phase, + fingerprint, }) }) .collect() diff --git a/src/shield/mod.rs b/src/shield/mod.rs index 785f225..21770a3 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -127,14 +127,13 @@ impl Shield { }))) } - /// The first rule in `phase` whose glob matches `path`, with its index (the - /// index keys the store so distinct rules keep independent budgets). A path - /// may match one rule per phase; each phase enforces independently. - fn match_rule(&self, path: &str, phase: Phase) -> Option<(usize, &CompiledRule)> { + /// The first rule in `phase` whose glob matches `path`. A path may match one + /// rule per phase; each phase enforces independently (two-phase by design: + /// a pre-auth IP/header rule and a post-auth claim rule can both apply). + fn match_rule(&self, path: &str, phase: Phase) -> Option<&CompiledRule> { self.rules .iter() - .enumerate() - .find(|(_, r)| r.phase == phase && r.matcher.is_match(path)) + .find(|r| r.phase == phase && r.matcher.is_match(path)) } /// Resolve the limit tier for a matched rule, in priority order: the JWT @@ -145,7 +144,7 @@ impl Shield { &self, rule: &CompiledRule, claims: Option<&serde_json::Value>, - key: &str, + identity: &str, ) -> Option { if let (Some(jwt), Some(claims)) = (&self.jwt_limits, claims) { if let Some(profile) = jwt.resolve(claims, &self.profiles) { @@ -153,7 +152,7 @@ impl Shield { } } if let Some(service) = &self.limit_service { - if let Some(profile) = service.resolve(key) { + if let Some(profile) = service.resolve(identity) { return Some(profile); } } @@ -194,7 +193,7 @@ pub async fn post_auth_middleware( /// Match a phase's rule for the request, apply its limit, and attach headers. async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> Response { let path = request.uri().path(); - let Some((idx, rule)) = shield.match_rule(path, phase) else { + let Some(rule) = shield.match_rule(path, phase) else { return next.run(request).await; }; @@ -207,9 +206,17 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> .extensions() .get::() .map(|c| c.0.as_ref()); - let key = rule_key(idx, &rule.key, &client, request.headers(), claims); - - let Some(profile) = shield.resolve_limit(rule, claims, &key) else { + let key = rule_key( + &rule.fingerprint, + &rule.key, + &client, + request.headers(), + claims, + ); + + // The limit service resolves per-principal, so it needs the real identity; + // the store / shared counter only need the de-identified `store` key. + let Some(profile) = shield.resolve_limit(rule, claims, &key.identity) else { // No limit resolves for this rule (JWT/service/profile/default all // absent): allow the request unmetered. return next.run(request).await; @@ -219,12 +226,12 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> // for a request the fleet-wide budget will reject. #[cfg(feature = "redis")] if let Some(global) = &shield.global { - if !global.gate(&key, profile.limit, profile.window) { + if !global.gate(&key.store, profile.limit, profile.window) { return global_reject(profile.limit, profile.window); } } - let verdict = shield.store.check(&key, &profile.gcra); + let verdict = shield.store.check(&key.store, &profile.gcra); if !verdict.allowed { return too_many_requests(profile.limit, &verdict); } @@ -232,7 +239,7 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> // Record the admit for the next reconciliation push. #[cfg(feature = "redis")] if let Some(global) = &shield.global { - global.record(&key, profile.window); + global.record(&key.store, profile.window); } let mut response = next.run(request).await; @@ -259,28 +266,43 @@ fn global_reject(limit: u64, window: Duration) -> Response { too_many_requests(limit, &verdict) } -/// Build the store key for a matched rule. The rule index namespaces the key so -/// two rules that happen to see the same client keep independent budgets. Every -/// source falls back to the client IP when its value is absent, so a limit can't -/// be dodged by omitting a header or authenticating anonymously. +/// The keys a matched rule derives for one request. +struct RuleKey { + /// De-identified key for the local store and shared counter: the rule's + /// stable fingerprint, the source tag, and a hash of the value. Raw client + /// values (API keys, principals, IPs) never reach the shared store or its + /// logs. Deterministic across instances so reconciliation keys agree. + store: String, + /// The raw identity for per-principal limit-service resolution (the service + /// must see the real principal to resolve its tier). Not persisted. + identity: String, +} + +/// Derive the store key and resolution identity for a matched rule. Every source +/// falls back to the client IP when its value is absent, so a limit can't be +/// dodged by omitting a header or authenticating anonymously (subject to the +/// rule's phase: a `jwt_claim` rule only runs post-auth). fn rule_key( - idx: usize, + fingerprint: &str, key: &KeySource, client: &str, headers: &HeaderMap, claims: Option<&serde_json::Value>, -) -> String { - let by_ip = || format!("{idx}:ip:{client}"); - match key { - KeySource::Ip => by_ip(), +) -> RuleKey { + let (tag, identity) = match key { + KeySource::Ip => ("ip", client.to_string()), KeySource::Header(name) => match header_str(headers, name) { - Some(v) => format!("{idx}:hdr:{v}"), - None => by_ip(), + Some(v) => ("hdr", v), + None => ("ip", client.to_string()), }, KeySource::JwtClaim(claim) => match claims.and_then(|c| resolve::claim_str(c, claim)) { - Some(v) => format!("{idx}:jwt:{v}"), - None => by_ip(), + Some(v) => ("jwt", v), + None => ("ip", client.to_string()), }, + }; + RuleKey { + store: format!("{fingerprint}:{tag}:{}", matcher::short_hash(&identity)), + identity, } } diff --git a/src/shield/tests.rs b/src/shield/tests.rs index 4e9a868..9d4c253 100644 --- a/src/shield/tests.rs +++ b/src/shield/tests.rs @@ -179,42 +179,50 @@ async fn rule_without_resolvable_profile_passes() { fn jwt_claim_key_uses_claim_then_falls_back_to_ip() { let claims = serde_json::json!({ "sub": "alice", "org": { "id": "acme" } }); let key = KeySource::JwtClaim("sub".to_string()); - // Present claim keys by its value. - assert_eq!( - rule_key(0, &key, "1.1.1.1", &HeaderMap::new(), Some(&claims)), - "0:jwt:alice" + // Present claim: identity is the raw claim value (for service resolution); + // the store key is de-identified (no raw value) and tagged `jwt`. + let k = rule_key("fp", &key, "1.1.1.1", &HeaderMap::new(), Some(&claims)); + assert_eq!(k.identity, "alice"); + assert!(k.store.starts_with("fp:jwt:")); + assert!( + !k.store.contains("alice"), + "raw value must not appear in store key" ); // Dotted path into a nested claim. let nested = KeySource::JwtClaim("org.id".to_string()); assert_eq!( - rule_key(0, &nested, "1.1.1.1", &HeaderMap::new(), Some(&claims)), - "0:jwt:acme" + rule_key("fp", &nested, "1.1.1.1", &HeaderMap::new(), Some(&claims)).identity, + "acme" ); // No claims (anonymous) → IP fallback, so the limit can't be dodged. - assert_eq!( - rule_key(0, &key, "1.1.1.1", &HeaderMap::new(), None), - "0:ip:1.1.1.1" - ); + let anon = rule_key("fp", &key, "1.1.1.1", &HeaderMap::new(), None); + assert_eq!(anon.identity, "1.1.1.1"); + assert!(anon.store.starts_with("fp:ip:")); // Claim present but missing the requested field → IP fallback. - assert_eq!( - rule_key( - 0, - &KeySource::JwtClaim("missing".to_string()), - "1.1.1.1", - &HeaderMap::new(), - Some(&claims) - ), - "0:ip:1.1.1.1" + let missing = rule_key( + "fp", + &KeySource::JwtClaim("missing".to_string()), + "1.1.1.1", + &HeaderMap::new(), + Some(&claims), ); + assert_eq!(missing.identity, "1.1.1.1"); + assert!(missing.store.starts_with("fp:ip:")); } #[test] -fn rule_index_namespaces_identical_keys() { - // The same client under two different rules keeps independent budgets. +fn store_key_namespaced_by_fingerprint_and_value() { let key = KeySource::Ip; - let a = rule_key(0, &key, "1.1.1.1", &HeaderMap::new(), None); - let b = rule_key(1, &key, "1.1.1.1", &HeaderMap::new(), None); - assert_ne!(a, b); + // Same client under two different rules (fingerprints) → independent budgets. + let a = rule_key("fpA", &key, "1.1.1.1", &HeaderMap::new(), None); + let b = rule_key("fpB", &key, "1.1.1.1", &HeaderMap::new(), None); + assert_ne!(a.store, b.store); + // Same rule + same value → identical store key (stable across instances). + let a2 = rule_key("fpA", &key, "1.1.1.1", &HeaderMap::new(), None); + assert_eq!(a.store, a2.store); + // Different value → different store key. + let c = rule_key("fpA", &key, "2.2.2.2", &HeaderMap::new(), None); + assert_ne!(a.store, c.store); } #[test] From 9a2bc8243eb97406b3ace7c7944d85c8b6c6516e Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 11:38:38 +0300 Subject: [PATCH 06/55] test(shield): add regression test for sub-second Retry-After rounding secs_ceil computed from truncated milliseconds reports 0 for a sub-millisecond wait; the test expects >= 1s and fails on current code. --- src/shield/tests.rs | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/shield/tests.rs b/src/shield/tests.rs index 9d4c253..5ebca0f 100644 --- a/src/shield/tests.rs +++ b/src/shield/tests.rs @@ -225,6 +225,18 @@ fn store_key_namespaced_by_fingerprint_and_value() { assert_ne!(a.store, c.store); } +#[test] +fn secs_ceil_rounds_sub_second_waits_up() { + use std::time::Duration; + // A sub-millisecond wait must report at least 1 second, never 0. + assert_eq!(secs_ceil(Duration::from_micros(500)), 1); + assert_eq!(secs_ceil(Duration::from_millis(1)), 1); + assert_eq!(secs_ceil(Duration::from_millis(1500)), 2); + assert_eq!(secs_ceil(Duration::from_secs(3)), 3); + // Genuinely zero stays zero. + assert_eq!(secs_ceil(Duration::ZERO), 0); +} + #[test] fn build_rejects_unknown_default_profile() { let mut cfg = config(vec![("t", "1/min", None)], Vec::new()); From e022c8f632b441db4529fe1a1ae6032944d4157a Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 11:38:57 +0300 Subject: [PATCH 07/55] fix(shield): round Retry-After / RateLimit-Reset up from nanoseconds Computing whole seconds from truncated milliseconds reported 0 for a sub-millisecond wait. Round up from the nanosecond value so a nonzero wait is never surfaced as 0. --- src/shield/mod.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/shield/mod.rs b/src/shield/mod.rs index 21770a3..66d8d5e 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -425,8 +425,10 @@ fn too_many_requests(limit: u64, verdict: &Verdict) -> Response { /// Whole seconds, rounded up, for `Retry-After` / `RateLimit-Reset` (never report /// `0` for a non-zero wait). fn secs_ceil(d: Duration) -> u64 { - let millis = d.as_millis(); - u64::try_from(millis.div_ceil(1000)).unwrap_or(u64::MAX) + // Round up from nanoseconds, not truncated millis: a sub-millisecond wait + // must still report at least one second, never zero. + let nanos = d.as_nanos(); + u64::try_from(nanos.div_ceil(1_000_000_000)).unwrap_or(u64::MAX) } #[cfg(test)] From e1037f0a4c331ce3c0b4b0c87cc3aaf62ddc0704 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 11:39:34 +0300 Subject: [PATCH 08/55] fix(shield): reject non-positive profile rate and burst A rate of 0 or an explicit burst of 0 was silently clamped to 1, hiding a misconfigured tier. Reject them at build time instead. --- src/shield/matcher.rs | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/src/shield/matcher.rs b/src/shield/matcher.rs index a9f0536..1a4a24b 100644 --- a/src/shield/matcher.rs +++ b/src/shield/matcher.rs @@ -78,8 +78,19 @@ fn path_glob(pattern: &str) -> Result { /// Compile a single limit profile from its config. pub fn compile_profile(cfg: &LimitProfileConfig) -> Result { let rate = super::rate::Rate::parse(&cfg.rate, BARE_COUNT_WINDOW)?; + // Reject non-positive limits explicitly rather than silently clamping them + // to 1, which would hide a misconfigured (effectively unlimited or dead) tier. + if rate.limit == 0 { + return Err(format!( + "profile rate must be greater than 0, got {:?}", + cfg.rate + )); + } + if cfg.burst == Some(0) { + return Err("profile burst must be greater than 0".to_string()); + } // Default burst = one full window of the sustained rate. - let burst = cfg.burst.unwrap_or(rate.limit).max(1); + let burst = cfg.burst.unwrap_or(rate.limit); let gcra = Gcra::from_profile(Profile { rate: rate.limit, window: rate.window, @@ -177,6 +188,20 @@ mod tests { assert_eq!(p.limit, 100); } + #[test] + fn zero_rate_or_burst_is_rejected() { + assert!(compile_profile(&LimitProfileConfig { + rate: "0/min".to_string(), + burst: None, + }) + .is_err()); + assert!(compile_profile(&LimitProfileConfig { + rate: "100/min".to_string(), + burst: Some(0), + }) + .is_err()); + } + #[test] fn glob_respects_and_spans_segments() { let rules = compile_rules( From 67895d5fa8522eb07297902f0bae18c6b1dcc2f8 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 11:41:09 +0300 Subject: [PATCH 09/55] fix(proxy): apply CORS outside the auth and rate-limit layers CORS was the innermost layer, so a request short-circuited by auth (401) or Shield (429) never passed back through it and returned without CORS headers, and preflight OPTIONS was subject to auth. Move CORS to the outermost layer and expose the RateLimit-* and Retry-After headers so browser clients can read the budget and back off. --- src/lib.rs | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index e30cdcf..aef3685 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -511,8 +511,10 @@ impl ProxyServer { .merge(openapi_routes) .merge(oidc_routes) .merge(embed::extra_routes_router(&self.extra_routes)) - .merge(transcode_routes) - .layer(cors); + .merge(transcode_routes); + // CORS is applied as the outermost layer below, so it wraps the auth and + // rate-limit enforcement: a short-circuited 401/429/503 still carries CORS + // headers, and preflight OPTIONS is answered before auth can reject it. // Forward-auth verification endpoint, sharing the built Auth. Mounted // after the auth layer below so the endpoint itself is not gated by the @@ -576,6 +578,9 @@ impl ProxyServer { maintenance_middleware, )) .layer(TraceLayer::new_for_http()) + // Outermost: wraps every enforcement layer so short-circuited + // responses keep CORS headers, and answers preflight before auth. + .layer(cors) .with_state(state); Ok(router) @@ -646,6 +651,11 @@ impl ProxyServer { .expose_headers([ "grpc-status".parse().unwrap(), "grpc-message".parse().unwrap(), + // Let browser clients read the rate-limit budget and back off. + "ratelimit-limit".parse().unwrap(), + "ratelimit-remaining".parse().unwrap(), + "ratelimit-reset".parse().unwrap(), + "retry-after".parse().unwrap(), ]) } } From 93ed030ff0d3f25a14532085633104ff8d13e89e Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 11:42:38 +0300 Subject: [PATCH 10/55] fix(shield): bound the limit-service resolution cache The resolution cache grew with distinct client keys and never shrank. Sweep entries not refreshed within a few refresh cycles (at least 5 minutes), at most once per minute, so client-controlled key cardinality cannot grow it without bound. --- src/shield/resolve.rs | 36 +++++++++++++++++++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/src/shield/resolve.rs b/src/shield/resolve.rs index 58e1c58..a8a09d9 100644 --- a/src/shield/resolve.rs +++ b/src/shield/resolve.rs @@ -123,16 +123,25 @@ struct Cached { at: Instant, } +/// Sweep the cache of long-idle keys at most once per this interval. +const SWEEP_INTERVAL: Duration = Duration::from_secs(60); + /// External limit-resolution service with an async, non-blocking cache. pub struct LimitService { endpoint: String, ttl: Duration, + /// Drop cache entries not refreshed within this window (a key that stopped + /// receiving requests), so client-controlled key cardinality can't grow the + /// cache without bound. + evict_after: Duration, client: reqwest::Client, /// Profiles for mapping a returned tier name to a compiled limit. profiles: HashMap, cache: dashmap::DashMap, /// Keys with a background fetch already in flight (dedupes refreshes). inflight: dashmap::DashMap, + base: Instant, + last_sweep_ms: std::sync::atomic::AtomicU64, } impl LimitService { @@ -149,20 +158,45 @@ impl LimitService { .tls_backend_preconfigured(crate::auth::jwks::build_tls_config()) .build() .map_err(|e| format!("invalid limit_service client: {e}"))?; + let ttl = Duration::from_secs(cfg.ttl_secs.max(1)); Ok(Arc::new(Self { endpoint: cfg.endpoint.clone(), - ttl: Duration::from_secs(cfg.ttl_secs.max(1)), + ttl, + // Keep an idle entry for a few refresh cycles, at least 5 minutes. + evict_after: (ttl * 4).max(Duration::from_secs(300)), client, profiles, cache: dashmap::DashMap::new(), inflight: dashmap::DashMap::new(), + base: Instant::now(), + last_sweep_ms: std::sync::atomic::AtomicU64::new(0), })) } + /// Evict cache entries not refreshed within `evict_after`, at most once per + /// [`SWEEP_INTERVAL`]; the first caller past the interval claims the sweep. + fn maybe_sweep(&self) { + use std::sync::atomic::Ordering; + let now_ms = u64::try_from(self.base.elapsed().as_millis()).unwrap_or(u64::MAX); + let last = self.last_sweep_ms.load(Ordering::Relaxed); + if now_ms.saturating_sub(last) < SWEEP_INTERVAL.as_millis() as u64 { + return; + } + if self + .last_sweep_ms + .compare_exchange(last, now_ms, Ordering::Relaxed, Ordering::Relaxed) + .is_ok() + { + let evict_after = self.evict_after; + self.cache.retain(|_, c| c.at.elapsed() < evict_after); + } + } + /// Resolve `key`'s limit from the cache, serving a stale value while a /// background refresh runs. Returns `None` (fall through to the static /// profile) only when nothing is cached yet. Never blocks the request. pub fn resolve(self: &Arc, key: &str) -> Option { + self.maybe_sweep(); let cached = self.cache.get(key).map(|c| *c); match cached { Some(c) if c.at.elapsed() < self.ttl => c.profile, From 51e3c6e10c2f144c83215c48aa03673a30bd1be1 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 11:43:58 +0300 Subject: [PATCH 11/55] fix(shield): store GCRA theoretical arrival times in nanoseconds Truncating the stored TAT to whole milliseconds lost precision for sub-millisecond emission intervals (very high per-second rates), so a fresh key could re-admit slightly early. Store nanoseconds instead; u64 nanoseconds span far beyond any process uptime. The periodic sweep still throttles on milliseconds. --- src/shield/store.rs | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/src/shield/store.rs b/src/shield/store.rs index 07af58d..33eccca 100644 --- a/src/shield/store.rs +++ b/src/shield/store.rs @@ -23,7 +23,7 @@ const SWEEP_INTERVAL_MS: u64 = 60_000; // no-std: caller-provided Clock + spin/hashbrown map. #[derive(Debug)] pub struct GcraStore { - /// key → TAT in milliseconds since `base`. + /// key → TAT in nanoseconds since `base`. tats: dashmap::DashMap, base: Instant, last_sweep_ms: AtomicU64, @@ -53,17 +53,17 @@ impl GcraStore { match self.tats.entry(key.to_string()) { Entry::Occupied(mut o) => { - let stored = Some(Duration::from_millis(*o.get())); + let stored = Some(Duration::from_nanos(*o.get())); let verdict = gcra.check(stored, now); if verdict.allowed { - *o.get_mut() = tat_millis(verdict.new_tat); + *o.get_mut() = dur_nanos(verdict.new_tat); } verdict } Entry::Vacant(v) => { let verdict = gcra.check(None, now); if verdict.allowed { - v.insert(tat_millis(verdict.new_tat)); + v.insert(dur_nanos(verdict.new_tat)); } verdict } @@ -80,14 +80,14 @@ impl GcraStore { /// on the next hit yields the same result. Without eviction, client-controlled /// key cardinality (IP / principal) would grow the map without bound. fn evict_drained(&self, now: Duration) { - let now_ms = tat_millis(now); - self.tats.retain(|_, tat| *tat > now_ms); + let now_nanos = dur_nanos(now); + self.tats.retain(|_, tat| *tat > now_nanos); } /// Evict at most once per [`SWEEP_INTERVAL_MS`]; the first caller past the /// interval claims the sweep so it stays an infrequent O(n) pass. fn maybe_sweep(&self, now: Duration) { - let now_ms = tat_millis(now); + let now_ms = u64::try_from(now.as_millis()).unwrap_or(u64::MAX); let last = self.last_sweep_ms.load(Ordering::Relaxed); if now_ms.saturating_sub(last) < SWEEP_INTERVAL_MS { return; @@ -108,9 +108,11 @@ impl GcraStore { } } -/// A `Duration` as whole milliseconds, saturating at `u64::MAX`. -fn tat_millis(d: Duration) -> u64 { - u64::try_from(d.as_millis()).unwrap_or(u64::MAX) +/// A `Duration` as whole nanoseconds, saturating at `u64::MAX`. Nanosecond TATs +/// preserve precision for sub-millisecond emission intervals (very high rates); +/// `u64` nanoseconds span ~584 years, far beyond any process uptime. +fn dur_nanos(d: Duration) -> u64 { + u64::try_from(d.as_nanos()).unwrap_or(u64::MAX) } #[cfg(test)] From 88c3bc3a1339a88f4333784ff520b87f9a4c06b9 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 11:45:13 +0300 Subject: [PATCH 12/55] docs(shield): clarify phase-local fallback, header contract, and overshoot - Document that a key's IP fallback is phase-local: an anonymous request under a jwt_claim rule is keyed by IP post-auth but not shed pre-auth, so anonymous flood protection needs a separate pre-auth IP/header rule. Note a path may match one rule per phase (defense in depth) and JWT-based limit resolution applies only to jwt_claim rules. - Document the RateLimit-* / Retry-After response contract and how clients back off, including CORS exposure. - Correct the overshoot formula to normalise the sync interval against the window so the units match, with a worked example. --- README.md | 40 +++++++++++++++++++++++++++++----------- src/shield/global.rs | 6 ++++-- 2 files changed, 33 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 266e1d9..099a835 100644 --- a/README.md +++ b/README.md @@ -187,19 +187,34 @@ lets legitimate bursts through up to a configured `burst` while throttling sustained abuse to the `rate`, with no fixed-window boundary burst. **Keying and phases.** A rule keys on the client IP, a header value (API key), -or a validated JWT claim. IP/header rules run *before* auth so anonymous floods -are shed before any signature verification; claim-keyed rules run *after* auth so -they use the verified, un-forgeable principal. The phase is derived from the key; -there is no phase setting to misconfigure. Every key falls back to the client IP -when its value is absent, so a limit can't be dodged by omitting a header or -staying anonymous. +or a validated JWT claim. The phase is derived from the key, not configured: an +IP/header rule needs no verified identity so it runs *before* auth (a fast, +purely local check that sheds anonymous floods before any signature verification, +and short-circuits so blocked clients never reach the auth layer); a `jwt_claim` +rule needs the verified principal so it runs *after* auth. A key falls back to +the client IP when its value is absent within its own phase, so a limit can't be +dodged by omitting a header. Note the fallback is phase-local: an anonymous +request under a `jwt_claim` rule keys by IP in the post-auth phase, but is *not* +shed pre-auth. For anonymous flood protection, add a separate pre-auth IP (or +header) rule covering the same paths; a path may match one rule per phase and +each is enforced independently (defense in depth). **Limit sources.** A key's `{rate, burst}` resolves in order: the JWT itself (a `ratelimit_tier` claim naming a profile, or explicit `ratelimit_rpm` / `ratelimit_burst`), then an external service (cached and refreshed in the background, never blocking), then the rule's pinned profile, then the default. -Tier-name indirection lets you retune the numbers in config without re-issuing -tokens or changing the service. +JWT-based resolution only applies to `jwt_claim` rules, since only they run with +verified claims available. Tier-name indirection lets you retune the numbers in +config without re-issuing tokens or changing the service. + +**Response headers.** Every metered response carries the +[draft-ietf-httpapi-ratelimit-headers](https://datatracker.ietf.org/doc/draft-ietf-httpapi-ratelimit-headers/) +fields: `RateLimit-Limit` (the tier's per-window quota), `RateLimit-Remaining` +(requests still admissible now), and `RateLimit-Reset` (whole seconds until the +limiter drains toward full). A rejected request returns `429` with `Retry-After` +(whole seconds until a retry would conform). Clients should back off for +`Retry-After` seconds on a `429`, and may pace themselves using `RateLimit-*` on +allowed responses. Behind a browser, these are exposed via CORS. **Deployment modes.** @@ -214,9 +229,12 @@ tokens or changing the service. unreachable, instances degrade to local limiting rather than failing requests. **Sizing the overshoot.** In reconciled mode the aggregate lags by up to one -`sync.interval_ms`, so the fleet can briefly overshoot the budget by about -`(N - 1) × rate × interval`. Shorter intervals tighten the bound at the cost of -more store traffic; the default (500 ms) suits per-minute limits. The global +`sync.interval_ms`, so within that lag each of the other instances can admit up +to a `rate` fraction of the window. With the interval expressed in the same time +unit as the window, the worst-case fleet overshoot is about +`(N - 1) × rate × (interval / window)` requests. For example, `rate = 1000/min`, +`interval = 500 ms`, `N = 4` gives `3 × 1000 × (0.5 / 60) ≈ 25` extra requests. +Shorter intervals tighten the bound at the cost of more store traffic. The global view uses a sliding-window counter, so there is no boundary burst on top of this lag. diff --git a/src/shield/global.rs b/src/shield/global.rs index a9c80ba..5c45ff5 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -8,8 +8,10 @@ //! failing requests. //! //! Counts are held in a sliding window (see [`super::window`]) so the fleet gate -//! has no boundary burst. The worst-case overshoot is bounded by one sync -//! interval of the other instances' traffic: `(N-1) × rate × interval`. +//! has no boundary burst. The worst-case fleet overshoot is bounded by one sync +//! interval of the other instances' traffic, about +//! `(N-1) * rate * (interval / window)` requests (interval and window in the +//! same unit). use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; From 7bb7c848f39bf277a3978ccb4af57320f9acafdd Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 11:50:15 +0300 Subject: [PATCH 13/55] docs(shield): explain TAT reuse across a resolved-tier change The store key omits the tier's numbers on purpose: a tier change reuses the existing absolute-time TAT, causing only a brief self-correcting transient, whereas keying by the numbers would let a client reset its budget by flipping tiers. --- src/shield/mod.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/shield/mod.rs b/src/shield/mod.rs index 66d8d5e..8d75454 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -231,6 +231,12 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> } } + // The store key intentionally excludes the profile's numbers, so if a key's + // resolved tier changes (a JWT/service tier upgrade), the existing TAT is + // reused with the new emission interval. The TAT is an absolute time, so this + // only causes a brief transient at the change and self-corrects within one + // window. Keying by the tier's numbers instead would reset the budget on + // every tier flip, which a client could exploit to shed its own limit. let verdict = shield.store.check(&key.store, &profile.gcra); if !verdict.allowed { return too_many_requests(profile.limit, &verdict); From b8ae32354162798762ffd4e2d009abbfacbfae38 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 12:01:16 +0300 Subject: [PATCH 14/55] test(shield): add regression test for evicting an active stale limit entry During a limit-service outage a key's cache entry stops refreshing, so its fetch timestamp goes stale even while the key is actively used. The test inserts a stale entry, accesses it, sweeps, and expects it to survive; it fails on current code, which evicts by fetch age rather than access age. Extracts the sweep body so the test can drive it deterministically. --- src/shield/resolve.rs | 42 ++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 40 insertions(+), 2 deletions(-) diff --git a/src/shield/resolve.rs b/src/shield/resolve.rs index a8a09d9..6e87391 100644 --- a/src/shield/resolve.rs +++ b/src/shield/resolve.rs @@ -187,11 +187,16 @@ impl LimitService { .compare_exchange(last, now_ms, Ordering::Relaxed, Ordering::Relaxed) .is_ok() { - let evict_after = self.evict_after; - self.cache.retain(|_, c| c.at.elapsed() < evict_after); + self.sweep(); } } + /// Drop entries idle longer than `evict_after`. + fn sweep(&self) { + let evict_after = self.evict_after; + self.cache.retain(|_, c| c.at.elapsed() < evict_after); + } + /// Resolve `key`'s limit from the cache, serving a stale value while a /// background refresh runs. Returns `None` (fall through to the static /// profile) only when nothing is cached yet. Never blocks the request. @@ -329,6 +334,39 @@ mod tests { ); } + #[tokio::test] + async fn actively_used_stale_entry_survives_eviction() { + use crate::config::LimitServiceConfig; + let svc = LimitService::build( + &LimitServiceConfig { + endpoint: "http://127.0.0.1:0/".to_string(), + ttl_secs: 1, + timeout_ms: 50, + }, + profiles(), + ) + .unwrap(); + // Simulate a service outage: the last successful fetch was long ago, so + // `at` is stale and well past evict_after, but the key is still in active + // use right now. + let old = Instant::now() + .checked_sub(Duration::from_secs(600)) + .expect("clock supports the offset"); + svc.cache.insert( + "k".to_string(), + Cached { + profile: Some(profile_from_numbers(10, 10)), + at: old, + }, + ); + let _ = svc.resolve("k"); + svc.sweep(); + assert!( + svc.cache.contains_key("k"), + "an actively-used stale entry must not be evicted during an outage" + ); + } + #[test] fn numeric_string_claims_are_accepted() { let claims = serde_json::json!({ "ratelimit_rpm": "250" }); From 850dea76f0fda456bfd5aa64814d2f84744679b1 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 12:02:20 +0300 Subject: [PATCH 15/55] fix(shield): evict limit cache by last access, not last fetch The cache evicted entries by fetch age, so during a limit-service outage an actively-used key (whose fetch timestamp cannot advance) was dropped after evict_after, silently falling back to the static profile. Track a separate last-access timestamp, bumped on every resolve including stale hits, and evict on that instead. --- src/shield/resolve.rs | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/src/shield/resolve.rs b/src/shield/resolve.rs index 6e87391..01d126c 100644 --- a/src/shield/resolve.rs +++ b/src/shield/resolve.rs @@ -120,7 +120,12 @@ struct LimitResponse { #[derive(Clone, Copy)] struct Cached { profile: Option, + /// When the value was last fetched (drives staleness / refresh age). at: Instant, + /// When the entry was last read (drives idle eviction). Advances on every + /// access, including stale hits, so an actively-used key is not evicted + /// during a service outage that keeps `at` from advancing. + last_access: Instant, } /// Sweep the cache of long-idle keys at most once per this interval. @@ -191,10 +196,12 @@ impl LimitService { } } - /// Drop entries idle longer than `evict_after`. + /// Drop entries not accessed within `evict_after` (idle keys), keyed on last + /// access rather than last fetch so an active-but-unrefreshable key survives. fn sweep(&self) { let evict_after = self.evict_after; - self.cache.retain(|_, c| c.at.elapsed() < evict_after); + self.cache + .retain(|_, c| c.last_access.elapsed() < evict_after); } /// Resolve `key`'s limit from the cache, serving a stale value while a @@ -202,7 +209,12 @@ impl LimitService { /// profile) only when nothing is cached yet. Never blocks the request. pub fn resolve(self: &Arc, key: &str) -> Option { self.maybe_sweep(); - let cached = self.cache.get(key).map(|c| *c); + // Bump last-access (even for a stale hit) so an actively-used key is kept + // through an outage, then read the cached value. + let cached = self.cache.get_mut(key).map(|mut c| { + c.last_access = Instant::now(); + *c + }); match cached { Some(c) if c.at.elapsed() < self.ttl => c.profile, Some(c) => { @@ -228,11 +240,13 @@ impl LimitService { let this = self.clone(); tokio::spawn(async move { if let Ok(resolved) = this.fetch(&key).await { + let now = Instant::now(); this.cache.insert( key.clone(), Cached { profile: resolved, - at: Instant::now(), + at: now, + last_access: now, }, ); } @@ -357,6 +371,7 @@ mod tests { Cached { profile: Some(profile_from_numbers(10, 10)), at: old, + last_access: old, }, ); let _ = svc.resolve("k"); From 4427cea241236b49acc853d47cd9ad9789113661 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 12:03:05 +0300 Subject: [PATCH 16/55] docs(shield): clarify JWT-limit scope and the fleet gate approximation - README: note that jwt_limits has no effect on an IP/header rule (those run pre-auth without claims); use a jwt_claim key to let the token's tier drive the limit. - global: document that the gate/record split is intentionally not an atomic reservation. The per-instance GCRA store is the hard local cap; the fleet gate is an approximate cross-instance cap whose gate/record race is bounded by the local burst plus the documented one-interval overshoot, kept lock-free on the hot path on purpose. --- README.md | 6 ++++-- src/shield/global.rs | 8 ++++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 099a835..57d799b 100644 --- a/README.md +++ b/README.md @@ -204,8 +204,10 @@ each is enforced independently (defense in depth). `ratelimit_burst`), then an external service (cached and refreshed in the background, never blocking), then the rule's pinned profile, then the default. JWT-based resolution only applies to `jwt_claim` rules, since only they run with -verified claims available. Tier-name indirection lets you retune the numbers in -config without re-issuing tokens or changing the service. +verified claims available; setting `jwt_limits` has no effect on an IP/header +rule, which runs pre-auth (use a `jwt_claim` key if you want the token's tier to +drive the limit). Tier-name indirection lets you retune the numbers in config +without re-issuing tokens or changing the service. **Response headers.** Every metered response carries the [draft-ietf-httpapi-ratelimit-headers](https://datatracker.ietf.org/doc/draft-ietf-httpapi-ratelimit-headers/) diff --git a/src/shield/global.rs b/src/shield/global.rs index 5c45ff5..1617f0b 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -102,6 +102,14 @@ impl GlobalCounters { /// the cached remote estimate plus this instance's count. Does not touch the /// store and does not record the admit (call [`record`](Self::record) after /// the local check also passes). + /// + /// This is deliberately a read then a separate record, not an atomic + /// reserve/rollback. The per-instance [`GcraStore`](super::store::GcraStore) + /// is the hard local cap and is atomic per key; this fleet gate is only an + /// approximate cross-instance cap. A race where concurrent requests all pass + /// the gate before any records is bounded by the local burst plus the + /// documented one-interval overshoot, which is the accepted trade-off for + /// keeping the hot path lock-free and non-blocking. pub fn gate(&self, key: &str, budget: u64, window: Duration) -> bool { let state = self.state_for(key, window); state.remote_estimate + state.local_count < budget From bf3cc357c314a76e4fd8a474e45290a1c03dacfa Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 12:06:26 +0300 Subject: [PATCH 17/55] fix(shield): make delta push atomic and guard stale estimate writes Two reconciliation races surfaced by fault injection: - The delta push pipeline was non-transactional, so a mid-pipeline failure could apply some INCRBYs while the caller skipped commit_pushed and re-pushed the same deltas next tick (inflating fleet usage into premature 429s). Wrap the push in MULTI/EXEC so it applies all-or-nothing. - A request could reset a key to a new window (or roll it to a newer epoch) while read_epochs awaited the store; applying the old plan's estimate then clobbered the fresh state. Skip the estimate write when the state's window or epoch has moved past the plan. --- src/shield/global.rs | 40 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/src/shield/global.rs b/src/shield/global.rs index 1617f0b..e263795 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -246,6 +246,13 @@ impl GlobalCounters { fn apply_estimates(&self, now: Duration, plans: &[PushPlan], reads: &[(u64, u64)]) { for (p, (cur, prev)) in plans.iter().zip(reads) { if let Some(mut s) = self.states.get_mut(&p.key) { + // A request may have reset the key to a new window, or rolled it + // to a newer epoch, while the read was in flight. The estimate we + // computed is for the plan's (window, epoch); applying it would + // clobber the fresh state with a stale value, so skip it. + if s.window_secs != p.window.as_secs() || s.epoch > p.read_epoch { + continue; + } let elapsed = window::elapsed_in_window(now, p.window); let est = window::sliding_estimate(*cur, *prev, elapsed, p.window); // Subtract only our own contribution to the current epoch's count @@ -258,13 +265,17 @@ impl GlobalCounters { } } - /// `INCRBY` each key with a pending delta and re-arm its TTL, in one pipeline. + /// `INCRBY` each key with a pending delta and re-arm its TTL. Wrapped in a + /// `MULTI`/`EXEC` transaction so the batch applies all-or-nothing: a + /// mid-pipeline failure can't leave some `INCRBY`s applied while the caller + /// skips `commit_pushed` and re-pushes the same deltas next tick. async fn push_deltas( &self, conn: &mut redis::aio::MultiplexedConnection, plans: &[PushPlan], ) -> redis::RedisResult<()> { let mut pipe = redis::pipe(); + pipe.atomic(); let mut any = false; for p in plans.iter().filter(|p| p.delta > 0) { any = true; @@ -360,6 +371,33 @@ mod tests { assert!(!g.gate("k", 3, W)); } + #[test] + fn stale_estimate_not_applied_after_window_reset() { + let g = counters(); + let key = "k"; + // The key currently lives on a 120s window with a fresh remote estimate. + g.record(key, Duration::from_secs(120)); + g.states.get_mut(key).unwrap().remote_estimate = 5; + + // A reconcile plan captured earlier for the OLD 60s window resumes after + // its read. Its estimate must not clobber the freshly-reset state. + let now = unix_now(); + let plan = PushPlan { + key: key.to_string(), + push_epoch: window::epoch(now, W), + read_epoch: window::epoch(now, W), + window: W, + delta: 0, + }; + g.apply_estimates(now, &[plan], &[(100, 0)]); + + assert_eq!( + g.states.get(key).unwrap().remote_estimate, + 5, + "an old-window estimate must not overwrite the reset state" + ); + } + #[test] fn independent_keys_have_independent_budgets() { let g = counters(); From 80bf2d55a81eea27067277026d369a9393250da4 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 12:12:25 +0300 Subject: [PATCH 18/55] fix(shield): preserve inner limiter's rate-limit headers With defense-in-depth (a pre-auth IP/header rule and a post-auth principal rule on the same path), a 429 from the inner principal limiter propagated out through the outer pre-auth phase, which overwrote RateLimit-Limit/Remaining/Reset with its own verdict. The client then saw the wrong rule's remaining capacity. The outer phase now leaves rate-limit headers intact when an inner limiter already set them. --- src/shield/mod.rs | 8 +++++++- src/shield/tests.rs | 48 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/src/shield/mod.rs b/src/shield/mod.rs index 8d75454..43c0865 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -249,7 +249,13 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> } let mut response = next.run(request).await; - attach_rate_headers(response.headers_mut(), profile.limit, &verdict); + // Don't overwrite rate-limit headers an inner limiter already set. With + // defense-in-depth (a pre-auth IP/header rule and a post-auth principal rule + // on the same path), the inner post-auth verdict is the more specific one and + // must reach the client intact, so the outer pre-auth phase leaves it alone. + if !response.headers().contains_key("ratelimit-limit") { + attach_rate_headers(response.headers_mut(), profile.limit, &verdict); + } response } diff --git a/src/shield/tests.rs b/src/shield/tests.rs index 5ebca0f..91005af 100644 --- a/src/shield/tests.rs +++ b/src/shield/tests.rs @@ -356,4 +356,52 @@ mod two_phase { // Bob is a different principal with his own budget. assert_eq!(get_with(&app, &bob).await, StatusCode::OK); } + + #[tokio::test] + async fn inner_principal_headers_survive_outer_limiter() { + // Defense in depth: a generous pre-auth IP rule and a tight post-auth + // principal rule both match the path. When the principal limit rejects, + // the 429's RateLimit-* must reflect that tight rule, not be overwritten + // by the outer pre-auth verdict on the way out. + let (sk, pem) = keypair_pem(); + let cfg = config( + vec![("wide", "100/min", Some(100)), ("tight", "60/min", Some(1))], + vec![ + rule("/api/**", KeySourceConfig::Ip, Some("wide")), + rule( + "/api/**", + KeySourceConfig::JwtClaim { + claim: "sub".to_string(), + }, + Some("tight"), + ), + ], + ); + let shield = Shield::build(&cfg).unwrap().unwrap(); + let app = stack(shield, auth(pem)); + let alice = token(&sk, "alice"); + + let send = |bearer: String| { + let app = app.clone(); + async move { + let req = Request::builder() + .uri("/api/x") + .header("authorization", format!("Bearer {bearer}")) + .body(Body::empty()) + .unwrap(); + app.oneshot(req).await.unwrap() + } + }; + + assert_eq!(send(alice.clone()).await.status(), StatusCode::OK); + let rejected = send(alice).await; + assert_eq!(rejected.status(), StatusCode::TOO_MANY_REQUESTS); + // The tight principal rule (limit 100/min → 100? no, "tight" is 60/min) + // owns the rejection, so its limit shows, not the wide pre-auth rule's. + assert_eq!( + rejected.headers().get("ratelimit-limit").unwrap(), + "60", + "the rejecting inner rule's RateLimit-Limit must not be overwritten" + ); + } } From 111d52ed776693f3ec3560a7a555bedaf5a5e13b Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 12:12:43 +0300 Subject: [PATCH 19/55] fix(shield): carry unpushed deltas across an epoch roll When a key accumulated admits late in a window and was seen again just after the boundary before the reconciler ran, roll_to zeroed the counts and the previous epoch's unpushed admits were dropped, so other replicas missed that window's traffic in the sliding-window prev count. Stash the unpushed remainder as a carryover owed to the epoch being left; the reconciler publishes it to that epoch's counter and clears it. Assumes the reconcile interval stays much shorter than the window (default 500ms vs seconds), so at most one unpushed epoch is outstanding between reconciles. --- src/shield/global.rs | 112 ++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 106 insertions(+), 6 deletions(-) diff --git a/src/shield/global.rs b/src/shield/global.rs index e263795..c39af28 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -45,6 +45,12 @@ struct KeyState { local_count: u64, /// How much of `local_count` has already been pushed to the store. pushed: u64, + /// Admits from the previous epoch that were never pushed before the epoch + /// rolled, still owed to `carryover_epoch`'s counter. Carried so a boundary + /// crossing between the last push and a roll does not silently drop them. + carryover: u64, + /// The epoch `carryover` is owed to. + carryover_epoch: u64, /// The rest of the fleet's sliding consumption, from the last pull. remote_estimate: u64, last_seen: Instant, @@ -57,16 +63,29 @@ impl KeyState { window_secs, local_count: 0, pushed: 0, + carryover: 0, + carryover_epoch: 0, remote_estimate: 0, last_seen: Instant::now(), } } - /// Advance to `epoch`, resetting the per-epoch counts. The remote estimate is - /// kept (the background task refreshes it) so a fresh epoch does not briefly - /// open the full budget fleet-wide. + /// Advance to `epoch`, resetting the per-epoch counts. Any admits not yet + /// pushed are stashed as `carryover` owed to the epoch being left, so the + /// reconciler still publishes them to that epoch's counter. The remote + /// estimate is kept (the background task refreshes it) so a fresh epoch does + /// not briefly open the full budget fleet-wide. + /// + /// A single carryover slot assumes the reconcile interval is much shorter + /// than the window (the default 500ms vs seconds+), so at most one unpushed + /// epoch exists between reconciles. fn roll_to(&mut self, epoch: u64) { if self.epoch != epoch { + let unpushed = self.local_count.saturating_sub(self.pushed); + if unpushed > 0 { + self.carryover += unpushed; + self.carryover_epoch = self.epoch; + } self.epoch = epoch; self.local_count = 0; self.pushed = 0; @@ -189,15 +208,28 @@ impl GlobalCounters { for entry in self.states.iter() { let s = entry.value(); let window = Duration::from_secs(s.window_secs); + let read_epoch = window::epoch(now, window); plans.push(PushPlan { key: entry.key().clone(), push_epoch: s.epoch, - read_epoch: window::epoch(now, window), + read_epoch, window, delta: s.local_count.saturating_sub(s.pushed), + is_carryover: false, }); + // Deltas from a prior epoch that rolled before they were pushed. + if s.carryover > 0 { + plans.push(PushPlan { + key: entry.key().clone(), + push_epoch: s.carryover_epoch, + read_epoch, + window, + delta: s.carryover, + is_carryover: true, + }); + } } - if plans.is_empty() { + if plans.iter().all(|p| p.delta == 0) { return; } @@ -235,7 +267,12 @@ impl GlobalCounters { fn commit_pushed(&self, plans: &[PushPlan]) { for p in plans.iter().filter(|p| p.delta > 0) { if let Some(mut s) = self.states.get_mut(&p.key) { - if s.epoch == p.push_epoch { + if p.is_carryover { + // Clear the carried-over amount now that its epoch counter has it. + if s.carryover_epoch == p.push_epoch { + s.carryover = s.carryover.saturating_sub(p.delta); + } + } else if s.epoch == p.push_epoch { s.pushed += p.delta; } } @@ -245,6 +282,11 @@ impl GlobalCounters { /// Refresh each key's cached estimate of the rest of the fleet's consumption. fn apply_estimates(&self, now: Duration, plans: &[PushPlan], reads: &[(u64, u64)]) { for (p, (cur, prev)) in plans.iter().zip(reads) { + // Carryover plans only publish a past epoch's delta; the estimate is + // driven by the current-epoch plan for the same key. + if p.is_carryover { + continue; + } if let Some(mut s) = self.states.get_mut(&p.key) { // A request may have reset the key to a new window, or rolled it // to a newer epoch, while the read was in flight. The estimate we @@ -331,6 +373,8 @@ struct PushPlan { read_epoch: u64, window: Duration, delta: u64, + /// True for a plan publishing a prior epoch's carried-over delta (no estimate). + is_carryover: bool, } #[cfg(test)] @@ -371,6 +415,22 @@ mod tests { assert!(!g.gate("k", 3, W)); } + #[test] + fn roll_preserves_unpushed_deltas_as_carryover() { + // 5 admits in an epoch, 2 already pushed, then the epoch rolls before the + // remaining 3 are pushed: they must survive as carryover owed to the old + // epoch, not be silently dropped. + let mut s = KeyState::new(10, 60); + s.local_count = 5; + s.pushed = 2; + s.roll_to(11); + assert_eq!(s.carryover, 3); + assert_eq!(s.carryover_epoch, 10); + assert_eq!(s.local_count, 0); + assert_eq!(s.pushed, 0); + assert_eq!(s.epoch, 11); + } + #[test] fn stale_estimate_not_applied_after_window_reset() { let g = counters(); @@ -388,6 +448,7 @@ mod tests { read_epoch: window::epoch(now, W), window: W, delta: 0, + is_carryover: false, }; g.apply_estimates(now, &[plan], &[(100, 0)]); @@ -451,6 +512,45 @@ mod tests { ); } + /// A carried-over delta from a rolled epoch is published to that epoch's + /// counter (not the current one) and then cleared. + #[tokio::test] + async fn carryover_is_pushed_to_its_epoch() { + let Ok(url) = std::env::var("SHIELD_REDIS_TEST_URL") else { + eprintln!("SKIP carryover_is_pushed_to_its_epoch: SHIELD_REDIS_TEST_URL not set"); + return; + }; + let win = Duration::from_secs(3600); + let nonce = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let key = format!("it3:{}:{nonce}", std::process::id()); + + let g = GlobalCounters::build(&url, Duration::from_millis(200)).unwrap(); + let cur = window::epoch(unix_now(), win); + let mut st = KeyState::new(cur, win.as_secs()); + st.carryover = 3; + st.carryover_epoch = cur - 1; + g.states.insert(key.clone(), st); + + g.reconcile().await; + + let client = redis::Client::open(url).unwrap(); + let mut conn = client.get_multiplexed_async_connection().await.unwrap(); + let prev: i64 = redis::cmd("GET") + .arg(epoch_key(&key, cur - 1)) + .query_async(&mut conn) + .await + .unwrap_or(0); + assert_eq!(prev, 3, "carryover must be published to its own epoch"); + assert_eq!( + g.states.get(&key).unwrap().carryover, + 0, + "carryover must be cleared once published" + ); + } + /// Repeated reconcile passes with no new admits must not re-push the same /// delta: the shared counter reflects each admit exactly once, even if a /// pass's aggregate read had failed on an earlier tick. From 0db9e9293160950da230bf05ab46b10970cc62d5 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 12:41:11 +0300 Subject: [PATCH 20/55] fix(shield): race-free reconcile via claim model, refresh, and eviction Rework cross-instance reconciliation to close several races surfaced by fault injection, and expose remaining fleet budget: - Claim model: a reconcile pass zeroes a key's unclaimed admits under the lock before the async push (restoring them only if the push never reached the store). A concurrent epoch roll or push failure can no longer double-publish a delta or drop it, replacing the racy pushed-cursor + carryover interaction. - Always refresh the estimate: a key rejected purely on a stale remote estimate (no local admits) is still read and decayed, so it recovers once the fleet stops spending, instead of returning 429 forever. - Eviction keeps keys with pending unpublished admits and uses an idle threshold of at least two windows, so long windows (e.g. 100/hour) and store outages don't lose fleet counts. - Expose fleet_remaining (replacing the boolean gate) so the caller can report the tighter of the local and fleet budgets. --- src/shield/global.rs | 249 ++++++++++++++++++++++++++----------------- 1 file changed, 151 insertions(+), 98 deletions(-) diff --git a/src/shield/global.rs b/src/shield/global.rs index c39af28..89942ae 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -36,22 +36,25 @@ fn unix_now() -> Duration { const KEY_TTL: Duration = Duration::from_secs(600); /// Per-key reconciliation state. +/// +/// Admit accounting uses a claim model: `local_count` holds only admits *not yet +/// claimed* by a reconcile pass. A pass claims the count (zeroes it under the +/// lock) before the async push, so a concurrent epoch roll or a push failure +/// can't double-publish or drop it. #[derive(Debug, Clone)] struct KeyState { epoch: u64, /// Window length in seconds (from the resolved profile). window_secs: u64, - /// Requests this instance admitted in the current epoch. + /// Admits in the current epoch not yet claimed for a push. local_count: u64, - /// How much of `local_count` has already been pushed to the store. - pushed: u64, - /// Admits from the previous epoch that were never pushed before the epoch - /// rolled, still owed to `carryover_epoch`'s counter. Carried so a boundary - /// crossing between the last push and a roll does not silently drop them. + /// Unclaimed admits owed to a prior epoch (`carryover_epoch`) that rolled + /// before they were claimed, so a boundary crossing doesn't drop them. carryover: u64, /// The epoch `carryover` is owed to. carryover_epoch: u64, - /// The rest of the fleet's sliding consumption, from the last pull. + /// The fleet's sliding consumption (including this instance's own claimed + /// admits, which now live in the shared counter), from the last pull. remote_estimate: u64, last_seen: Instant, } @@ -62,7 +65,6 @@ impl KeyState { epoch, window_secs, local_count: 0, - pushed: 0, carryover: 0, carryover_epoch: 0, remote_estimate: 0, @@ -70,25 +72,22 @@ impl KeyState { } } - /// Advance to `epoch`, resetting the per-epoch counts. Any admits not yet - /// pushed are stashed as `carryover` owed to the epoch being left, so the - /// reconciler still publishes them to that epoch's counter. The remote - /// estimate is kept (the background task refreshes it) so a fresh epoch does - /// not briefly open the full budget fleet-wide. + /// Advance to `epoch`, moving any unclaimed admits to `carryover` owed to the + /// epoch being left so the reconciler still publishes them to that epoch's + /// counter. The remote estimate is kept (the background task refreshes it) so + /// a fresh epoch does not briefly open the full budget fleet-wide. /// /// A single carryover slot assumes the reconcile interval is much shorter - /// than the window (the default 500ms vs seconds+), so at most one unpushed - /// epoch exists between reconciles. + /// than the window (the default 500ms vs seconds+), so at most one unclaimed + /// prior epoch exists between reconciles. fn roll_to(&mut self, epoch: u64) { if self.epoch != epoch { - let unpushed = self.local_count.saturating_sub(self.pushed); - if unpushed > 0 { - self.carryover += unpushed; + if self.local_count > 0 { + self.carryover += self.local_count; self.carryover_epoch = self.epoch; } self.epoch = epoch; self.local_count = 0; - self.pushed = 0; } } } @@ -117,10 +116,10 @@ impl GlobalCounters { })) } - /// Whether a request for `key` is within the fleet budget. Read-only: reads - /// the cached remote estimate plus this instance's count. Does not touch the - /// store and does not record the admit (call [`record`](Self::record) after - /// the local check also passes). + /// Fleet budget still available for `key` (0 = over budget). Read-only: the + /// cached remote estimate plus this instance's unclaimed count, subtracted + /// from `budget`. Does not touch the store and does not record the admit + /// (call [`record`](Self::record) after the local check also passes). /// /// This is deliberately a read then a separate record, not an atomic /// reserve/rollback. The per-instance [`GcraStore`](super::store::GcraStore) @@ -129,9 +128,9 @@ impl GlobalCounters { /// the gate before any records is bounded by the local burst plus the /// documented one-interval overshoot, which is the accepted trade-off for /// keeping the hot path lock-free and non-blocking. - pub fn gate(&self, key: &str, budget: u64, window: Duration) -> bool { + pub fn fleet_remaining(&self, key: &str, budget: u64, window: Duration) -> u64 { let state = self.state_for(key, window); - state.remote_estimate + state.local_count < budget + budget.saturating_sub(state.remote_estimate + state.local_count) } /// Record one admitted request for `key`, to be pushed to the store on the @@ -200,57 +199,68 @@ impl GlobalCounters { async fn reconcile_at(&self, now: Duration) { self.evict_stale(); - // Phase 1 (locked, brief): snapshot each key's pending delta. The delta - // is pushed to the epoch it was accumulated in (`push_epoch`), while the - // estimate reads the current epoch (`read_epoch`); the two differ when a - // window boundary is crossed between a record and this tick. + // Phase 1 (locked per key, brief): CLAIM each key's unclaimed admits by + // zeroing them now, so a concurrent epoch roll or a push failure can't + // double-publish or drop them. Emit a plan for every active key (delta + // may be 0) so its estimate is refreshed even when the key is only being + // rejected on a stale remote estimate. The delta is pushed to the epoch + // it was accumulated in (`push_epoch`); the estimate reads the current + // epoch (`read_epoch`). + let keys: Vec = self.states.iter().map(|e| e.key().clone()).collect(); + if keys.is_empty() { + return; + } let mut plans: Vec = Vec::new(); - for entry in self.states.iter() { - let s = entry.value(); - let window = Duration::from_secs(s.window_secs); - let read_epoch = window::epoch(now, window); - plans.push(PushPlan { - key: entry.key().clone(), - push_epoch: s.epoch, - read_epoch, - window, - delta: s.local_count.saturating_sub(s.pushed), - is_carryover: false, - }); - // Deltas from a prior epoch that rolled before they were pushed. - if s.carryover > 0 { + for key in keys { + if let Some(mut s) = self.states.get_mut(&key) { + let window = Duration::from_secs(s.window_secs); + let read_epoch = window::epoch(now, window); + let claim = s.local_count; + s.local_count = 0; plans.push(PushPlan { - key: entry.key().clone(), - push_epoch: s.carryover_epoch, + key: key.clone(), + push_epoch: s.epoch, read_epoch, window, - delta: s.carryover, - is_carryover: true, + delta: claim, + is_carryover: false, }); + if s.carryover > 0 { + let carry = s.carryover; + s.carryover = 0; + plans.push(PushPlan { + key: key.clone(), + push_epoch: s.carryover_epoch, + read_epoch, + window, + delta: carry, + is_carryover: true, + }); + } } } - if plans.iter().all(|p| p.delta == 0) { - return; - } let mut conn = match self.connection().await { Ok(c) => c, Err(e) => { tracing::warn!("rate-limit shared store unavailable, staying local: {e}"); + self.restore_claims(&plans); // never pushed → give the claims back return; } }; - // Push each pending delta as an atomic INCRBY (never a SET, so concurrent - // instances' increments accumulate rather than clobber). On success, - // commit `pushed` immediately so a later read failure cannot cause the - // same delta to be pushed a second time on the next tick. + // Push each claimed delta as an atomic INCRBY inside MULTI/EXEC (never a + // SET, so concurrent instances' increments accumulate). Because the push + // is all-or-nothing, a failure means nothing was applied, so restoring + // the claims for a clean retry cannot double-count. if let Err(e) = self.push_deltas(&mut conn, &plans).await { tracing::warn!("rate-limit delta push failed: {e}"); - return; // `pushed` not advanced → same delta retried next tick, no double count + self.restore_claims(&plans); + return; } - self.commit_pushed(&plans); + // Claimed admits are now in the shared counter; there is nothing to + // commit. A read failure below just skips this tick's estimate refresh. let reads = match self.read_epochs(&mut conn, &plans).await { Ok(r) => r, Err(e) => { @@ -261,25 +271,25 @@ impl GlobalCounters { self.apply_estimates(now, &plans, &reads); } - /// Advance each key's `pushed` by the delta we just pushed, but only while - /// the state still holds the epoch the delta belonged to (a concurrent roll - /// resets the counts, so there is nothing to advance). - fn commit_pushed(&self, plans: &[PushPlan]) { + /// Return claimed deltas to the state when a push never reached the store, so + /// the next tick retries them. A delta whose epoch rolled meanwhile is owed + /// to the epoch it was accumulated in, so it goes back to `carryover`. + fn restore_claims(&self, plans: &[PushPlan]) { for p in plans.iter().filter(|p| p.delta > 0) { if let Some(mut s) = self.states.get_mut(&p.key) { - if p.is_carryover { - // Clear the carried-over amount now that its epoch counter has it. - if s.carryover_epoch == p.push_epoch { - s.carryover = s.carryover.saturating_sub(p.delta); + if !p.is_carryover && s.epoch == p.push_epoch { + s.local_count += p.delta; + } else { + if s.carryover == 0 { + s.carryover_epoch = p.push_epoch; } - } else if s.epoch == p.push_epoch { - s.pushed += p.delta; + s.carryover += p.delta; } } } } - /// Refresh each key's cached estimate of the rest of the fleet's consumption. + /// Refresh each key's cached estimate of the fleet's consumption. fn apply_estimates(&self, now: Duration, plans: &[PushPlan], reads: &[(u64, u64)]) { for (p, (cur, prev)) in plans.iter().zip(reads) { // Carryover plans only publish a past epoch's delta; the estimate is @@ -296,21 +306,19 @@ impl GlobalCounters { continue; } let elapsed = window::elapsed_in_window(now, p.window); + // The counter already includes this instance's claimed admits; + // the gate adds only the still-unclaimed `local_count` on top, so + // the full sliding estimate is used with no self-subtraction. let est = window::sliding_estimate(*cur, *prev, elapsed, p.window); - // Subtract only our own contribution to the current epoch's count - // (the gate adds `local_count` back separately). If the key's - // pushes went to an earlier epoch, our current-epoch share is 0. - let my_current = if s.epoch == p.read_epoch { s.pushed } else { 0 }; - let others = (est.round() as i64 - my_current as i64).max(0); - s.remote_estimate = others as u64; + s.remote_estimate = est.round().max(0.0) as u64; } } } - /// `INCRBY` each key with a pending delta and re-arm its TTL. Wrapped in a + /// `INCRBY` each key with a claimed delta and re-arm its TTL. Wrapped in a /// `MULTI`/`EXEC` transaction so the batch applies all-or-nothing: a /// mid-pipeline failure can't leave some `INCRBY`s applied while the caller - /// skips `commit_pushed` and re-pushes the same deltas next tick. + /// restores the claims and re-pushes the same deltas next tick. async fn push_deltas( &self, conn: &mut redis::aio::MultiplexedConnection, @@ -356,11 +364,19 @@ impl GlobalCounters { .collect()) } - /// Drop per-key state not seen within [`KEY_TTL`]. + /// Drop per-key state that is idle and carries no unpublished admits. The + /// idle threshold is at least two windows, so a key on a long window (e.g. + /// `100/hour`) is not evicted mid-window; a key with pending `local_count` or + /// `carryover` is always kept so a store outage can't lose fleet counts. fn evict_stale(&self) { let now = Instant::now(); - self.states - .retain(|_, s| now.duration_since(s.last_seen) < KEY_TTL); + self.states.retain(|_, s| { + if s.local_count > 0 || s.carryover > 0 { + return true; + } + let threshold = KEY_TTL.max(Duration::from_secs(s.window_secs.saturating_mul(2))); + now.duration_since(s.last_seen) < threshold + }); } } @@ -390,18 +406,18 @@ mod tests { } #[test] - fn gate_admits_until_local_count_reaches_budget() { + fn budget_admits_until_local_count_reaches_it() { let g = counters(); - // Budget 3: three admits pass, the fourth is over budget. + // Budget 3: three admits leave room, the fourth is over budget. for _ in 0..3 { - assert!(g.gate("k", 3, W)); + assert!(g.fleet_remaining("k", 3, W) > 0); g.record("k", W); } - assert!(!g.gate("k", 3, W)); + assert_eq!(g.fleet_remaining("k", 3, W), 0); } #[test] - fn gate_accounts_for_remote_estimate() { + fn budget_accounts_for_remote_estimate() { let g = counters(); // Simulate the background task having observed 2 requests elsewhere. let now = unix_now(); @@ -410,24 +426,21 @@ mod tests { state.remote_estimate = 2; g.states.insert("k".to_string(), state); // Budget 3, remote 2 → one local admit fits, the next is over budget. - assert!(g.gate("k", 3, W)); + assert!(g.fleet_remaining("k", 3, W) > 0); g.record("k", W); - assert!(!g.gate("k", 3, W)); + assert_eq!(g.fleet_remaining("k", 3, W), 0); } #[test] - fn roll_preserves_unpushed_deltas_as_carryover() { - // 5 admits in an epoch, 2 already pushed, then the epoch rolls before the - // remaining 3 are pushed: they must survive as carryover owed to the old - // epoch, not be silently dropped. + fn roll_preserves_unclaimed_deltas_as_carryover() { + // Admits accumulated in an epoch that rolls before a reconcile claims + // them must survive as carryover owed to the old epoch, not be dropped. let mut s = KeyState::new(10, 60); s.local_count = 5; - s.pushed = 2; s.roll_to(11); - assert_eq!(s.carryover, 3); + assert_eq!(s.carryover, 5); assert_eq!(s.carryover_epoch, 10); assert_eq!(s.local_count, 0); - assert_eq!(s.pushed, 0); assert_eq!(s.epoch, 11); } @@ -462,11 +475,11 @@ mod tests { #[test] fn independent_keys_have_independent_budgets() { let g = counters(); - assert!(g.gate("a", 1, W)); + assert!(g.fleet_remaining("a", 1, W) > 0); g.record("a", W); - assert!(!g.gate("a", 1, W)); + assert_eq!(g.fleet_remaining("a", 1, W), 0); // A different key is unaffected. - assert!(g.gate("b", 1, W)); + assert!(g.fleet_remaining("b", 1, W) > 0); } /// End-to-end reconciliation against a live Redis-protocol store: one @@ -499,15 +512,19 @@ mod tests { // B's first contact with the key: it has not reconciled this key yet, so // it admits once (the documented one-interval lag / bounded overshoot). - assert!(b.gate(&key, 3, W), "B admits its first request for the key"); + assert!( + b.fleet_remaining(&key, 3, W) > 0, + "B admits its first request for the key" + ); b.record(&key, W); // After a reconcile pass B has pushed its own admit and pulled the // aggregate: the combined view is 3 (A's 2 + B's 1) = the budget, so B // rejects the next request. The two instances enforce one combined limit. b.reconcile().await; - assert!( - !b.gate(&key, 3, W), + assert_eq!( + b.fleet_remaining(&key, 3, W), + 0, "B must reject once the combined budget is reached" ); } @@ -551,6 +568,42 @@ mod tests { ); } + /// A key with no local admits (only being rejected on a stale remote + /// estimate) must still have its estimate refreshed by a reconcile pass, so + /// it can recover once the fleet stops spending the budget. + #[tokio::test] + async fn estimate_refreshes_without_local_deltas() { + let Ok(url) = std::env::var("SHIELD_REDIS_TEST_URL") else { + eprintln!( + "SKIP estimate_refreshes_without_local_deltas: SHIELD_REDIS_TEST_URL not set" + ); + return; + }; + let win = Duration::from_secs(3600); + let nonce = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let key = format!("it4:{}:{nonce}", std::process::id()); + + let g = GlobalCounters::build(&url, Duration::from_millis(200)).unwrap(); + // Seed a stale-high remote estimate with no local admits (delta 0). The + // shared counter for this fresh key is empty, so a reconcile must pull it + // and decay the estimate to 0 rather than leaving the key stuck. + let cur = window::epoch(unix_now(), win); + let mut st = KeyState::new(cur, win.as_secs()); + st.remote_estimate = 99; + g.states.insert(key.clone(), st); + + g.reconcile().await; + + assert_eq!( + g.states.get(&key).unwrap().remote_estimate, + 0, + "estimate must be refreshed even with no local deltas" + ); + } + /// Repeated reconcile passes with no new admits must not re-push the same /// delta: the shared counter reflects each admit exactly once, even if a /// pass's aggregate read had failed on an earlier tick. From 2278c1c7e55dfdc62bac86f995e3383f0c93c6a2 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 12:41:30 +0300 Subject: [PATCH 21/55] fix(shield): fail closed on empty rules and report fleet remaining - An enabled Shield with no rules now fails startup instead of silently running unmetered. An upgrade that left the old schema (endpoint_classes / identifier_endpoints / redis_url) in place deserializes to zero rules, which would otherwise disable this security control while enabled is true. - On an allowed response, report the tighter of the local GCRA remaining and the reconciled fleet remaining, so a client near the fleet cap isn't told it has ample local headroom. --- src/shield/mod.rs | 58 +++++++++++++++++++++++++++++++++-------------- 1 file changed, 41 insertions(+), 17 deletions(-) diff --git a/src/shield/mod.rs b/src/shield/mod.rs index 43c0865..9defeb2 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -64,8 +64,15 @@ impl Shield { return Ok(None); } if config.rules.is_empty() { - tracing::warn!("shield enabled but no rules configured"); - return Ok(None); + // Fail loud rather than silently running unmetered: an upgrade that + // left an old `shield` schema (endpoint_classes / identifier_endpoints + // / redis_url) in place deserializes to zero rules, which would + // otherwise disable this security control while `enabled` is true. + return Err( + "shield.enabled is true but no rules are configured (note the schema: \ + profiles + rules + sync, not the older endpoint_classes/identifier_endpoints)" + .to_string(), + ); } let profiles = matcher::compile_profiles(&config.profiles)?; @@ -225,10 +232,13 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> // Fleet gate first (read-only, cached), so the local shaper is not charged // for a request the fleet-wide budget will reject. #[cfg(feature = "redis")] - if let Some(global) = &shield.global { - if !global.gate(&key.store, profile.limit, profile.window) { - return global_reject(profile.limit, profile.window); - } + let fleet_remaining = shield + .global + .as_ref() + .map(|g| g.fleet_remaining(&key.store, profile.limit, profile.window)); + #[cfg(feature = "redis")] + if fleet_remaining == Some(0) { + return global_reject(profile.limit, profile.window); } // The store key intentionally excludes the profile's numbers, so if a key's @@ -239,14 +249,26 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> // every tier flip, which a client could exploit to shed its own limit. let verdict = shield.store.check(&key.store, &profile.gcra); if !verdict.allowed { - return too_many_requests(profile.limit, &verdict); + return too_many_requests(profile.limit, verdict.remaining, &verdict); } - // Record the admit for the next reconciliation push. + // Report the tighter of the local and (when reconciled) fleet budgets, so a + // client near the fleet cap isn't told it has ample local room. This admit + // consumes one, hence the `- 1`. + #[cfg(not(feature = "redis"))] + let reported = verdict.remaining; #[cfg(feature = "redis")] - if let Some(global) = &shield.global { - global.record(&key.store, profile.window); - } + let reported = { + let mut r = verdict.remaining; + if let Some(fr) = fleet_remaining { + r = r.min(fr.saturating_sub(1)); + // Record the admit for the next reconciliation push. + if let Some(global) = &shield.global { + global.record(&key.store, profile.window); + } + } + r + }; let mut response = next.run(request).await; // Don't overwrite rate-limit headers an inner limiter already set. With @@ -254,7 +276,7 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> // on the same path), the inner post-auth verdict is the more specific one and // must reach the client intact, so the outer pre-auth phase leaves it alone. if !response.headers().contains_key("ratelimit-limit") { - attach_rate_headers(response.headers_mut(), profile.limit, &verdict); + attach_rate_headers(response.headers_mut(), profile.limit, reported, &verdict); } response } @@ -275,7 +297,7 @@ fn global_reject(limit: u64, window: Duration) -> Response { retry_after: remaining, reset_after: remaining, }; - too_many_requests(limit, &verdict) + too_many_requests(limit, 0, &verdict) } /// The keys a matched rule derives for one request. @@ -404,11 +426,13 @@ fn header_str(headers: &HeaderMap, name: &str) -> Option { } /// Attach the draft-ietf `RateLimit-*` headers describing the remaining budget. -fn attach_rate_headers(headers: &mut HeaderMap, limit: u64, verdict: &Verdict) { +/// `remaining` is passed explicitly (rather than read from the verdict) so the +/// caller can report the tighter of the local and fleet budgets. +fn attach_rate_headers(headers: &mut HeaderMap, limit: u64, remaining: u64, verdict: &Verdict) { if let Ok(v) = limit.to_string().parse() { headers.insert("ratelimit-limit", v); } - if let Ok(v) = verdict.remaining.to_string().parse() { + if let Ok(v) = remaining.to_string().parse() { headers.insert("ratelimit-remaining", v); } if let Ok(v) = secs_ceil(verdict.reset_after).to_string().parse() { @@ -417,7 +441,7 @@ fn attach_rate_headers(headers: &mut HeaderMap, limit: u64, verdict: &Verdict) { } /// A `429` response carrying the rate-limit headers plus `Retry-After`. -fn too_many_requests(limit: u64, verdict: &Verdict) -> Response { +fn too_many_requests(limit: u64, remaining: u64, verdict: &Verdict) -> Response { let mut response = ( StatusCode::TOO_MANY_REQUESTS, Json(serde_json::json!({ @@ -427,7 +451,7 @@ fn too_many_requests(limit: u64, verdict: &Verdict) -> Response { ) .into_response(); let headers = response.headers_mut(); - attach_rate_headers(headers, limit, verdict); + attach_rate_headers(headers, limit, remaining, verdict); if let Ok(v) = secs_ceil(verdict.retry_after).to_string().parse() { headers.insert("retry-after", v); } From 4ddb41311cceacef5f45dcd801b740a3c48c1746 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 12:41:48 +0300 Subject: [PATCH 22/55] fix(shield): reject dynamic zero rates and bound fetch storms - A JWT or limit-service rate of 0 no longer clamps to 1 (which silently granted one request per minute); it yields no usable limit and falls through to the next resolver, matching how static profiles reject zero. Blocking a principal is an authorization concern, not a rate limit. - On a limit-service fetch failure for a brand-new key, negative-cache the miss so a client rotating key values during an outage can't spawn an unbounded stream of background fetches; an existing stale entry is still left in place. --- src/shield/resolve.rs | 76 ++++++++++++++++++++++++++++++------------- 1 file changed, 54 insertions(+), 22 deletions(-) diff --git a/src/shield/resolve.rs b/src/shield/resolve.rs index 01d126c..b5543c4 100644 --- a/src/shield/resolve.rs +++ b/src/shield/resolve.rs @@ -50,19 +50,25 @@ fn claim_at<'a>(claims: &'a Value, path: &str) -> Option<&'a Value> { Some(cur) } -/// Build a limit tier from explicit per-minute numbers. -fn profile_from_numbers(rpm: u64, burst: u64) -> CompiledProfile { - let rate = rpm.max(1); +/// Build a limit tier from explicit per-minute numbers. Returns `None` for a +/// zero rate: like a static profile, a dynamic `0` is not a usable limit (it +/// would otherwise clamp to 1 and silently grant a request per minute), so the +/// caller falls through to the next resolver. Blocking a principal outright is +/// an authorization concern, not a rate limit. +fn profile_from_numbers(rpm: u64, burst: u64) -> Option { + if rpm == 0 { + return None; + } let gcra = Gcra::from_profile(Profile { - rate, + rate: rpm, window: PER_MINUTE, burst: burst.max(1), }); - CompiledProfile { + Some(CompiledProfile { gcra, - limit: rate, + limit: rpm, window: PER_MINUTE, - } + }) } /// Compiled JWT-based limit resolution: the claim names to read from the token. @@ -98,7 +104,7 @@ impl JwtLimits { } if let Some(rpm) = claim_u64(claims, &self.rpm_claim) { let burst = claim_u64(claims, &self.burst_claim).unwrap_or(rpm); - return Some(profile_from_numbers(rpm, burst)); + return profile_from_numbers(rpm, burst); } None } @@ -239,18 +245,36 @@ impl LimitService { } let this = self.clone(); tokio::spawn(async move { - if let Ok(resolved) = this.fetch(&key).await { - let now = Instant::now(); - this.cache.insert( - key.clone(), - Cached { - profile: resolved, - at: now, - last_access: now, - }, - ); + match this.fetch(&key).await { + Ok(resolved) => { + let now = Instant::now(); + this.cache.insert( + key.clone(), + Cached { + profile: resolved, + at: now, + last_access: now, + }, + ); + } + Err(()) => { + // Leave an existing (stale) entry in place (fail-static). For a + // brand-new key, negative-cache the failure so a client rotating + // key values during an outage can't spawn unbounded fetch tasks; + // it is retried after the TTL like any stale entry. + if !this.cache.contains_key(&key) { + let now = Instant::now(); + this.cache.insert( + key.clone(), + Cached { + profile: None, + at: now, + last_access: now, + }, + ); + } + } } - // On error, leave any stale entry in place (fail-static, not fail-open). this.inflight.remove(&key); }); } @@ -289,7 +313,7 @@ impl LimitService { return self.profiles.get(&tier).copied(); } body.rate_per_min - .map(|rpm| profile_from_numbers(rpm, body.burst.unwrap_or(rpm))) + .and_then(|rpm| profile_from_numbers(rpm, body.burst.unwrap_or(rpm))) } } @@ -301,7 +325,7 @@ mod tests { let mut m = HashMap::new(); m.insert( "premium".to_string(), - profile_from_numbers(1000, 100), // limit 1000 + profile_from_numbers(1000, 100).unwrap(), // limit 1000 ); m } @@ -369,7 +393,7 @@ mod tests { svc.cache.insert( "k".to_string(), Cached { - profile: Some(profile_from_numbers(10, 10)), + profile: Some(profile_from_numbers(10, 10).unwrap()), at: old, last_access: old, }, @@ -382,6 +406,14 @@ mod tests { ); } + #[test] + fn jwt_zero_rate_is_not_a_usable_limit() { + // A dynamic rate of 0 must not clamp to 1; it yields no limit so the + // caller falls through to the next resolver. + let claims = serde_json::json!({ "ratelimit_rpm": 0 }); + assert!(jwt_limits().resolve(&claims, &profiles()).is_none()); + } + #[test] fn numeric_string_claims_are_accepted() { let claims = serde_json::json!({ "ratelimit_rpm": "250" }); From 065a4144a38dfd9dc5626a38bc42b90bc31cd9cf Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 13:06:25 +0300 Subject: [PATCH 23/55] fix(shield): drop unpushable deltas instead of restoring (never double) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reconciliation is now fire-and-forget: once a pass claims a key's admits (zeroing them), a push that fails or whose EXEC ack is lost simply drops them. Restoring risked republishing a committed-but-unacked batch on the next tick, inflating the fleet count into a false 429; dropping instead under-counts this instance's last interval, which is exactly the documented "store unreachable → degrade to per-instance limiting" behaviour and never rejects valid traffic. It also stops an unpushable delta from accumulating across window boundaries into a collapsed carryover. A window change still resets the entry (a different window is a different accounting unit), a bounded under-count noted at the reset. --- src/shield/global.rs | 38 ++++++++++++++------------------------ 1 file changed, 14 insertions(+), 24 deletions(-) diff --git a/src/shield/global.rs b/src/shield/global.rs index 89942ae..a435bca 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -157,6 +157,11 @@ impl GlobalCounters { .entry(key.to_string()) .or_insert_with(|| KeyState::new(ep, window_secs)); if state.window_secs != window_secs { + // A window change is a different accounting unit, so the entry resets. + // Any unpushed admits from the old window are dropped rather than + // remapped (a different window can't share a counter); this is a rare, + // bounded under-count on a tier change, consistent with the fleet + // layer's best-effort, never-double contract. *state = KeyState::new(ep, window_secs); } state.roll_to(ep); @@ -240,22 +245,25 @@ impl GlobalCounters { } } + // Claimed deltas are fire-and-forget: once claimed (zeroed), a push that + // fails or whose EXEC ack is lost simply drops them rather than restoring. + // Restoring risked publishing a committed-but-unacked batch twice (a false + // 429); dropping instead under-counts this instance's last interval, which + // is exactly the documented "store unreachable → degrade to per-instance" + // behaviour. It also means an unpushable delta never accumulates across + // window boundaries into a collapsed carryover. let mut conn = match self.connection().await { Ok(c) => c, Err(e) => { tracing::warn!("rate-limit shared store unavailable, staying local: {e}"); - self.restore_claims(&plans); // never pushed → give the claims back return; } }; // Push each claimed delta as an atomic INCRBY inside MULTI/EXEC (never a - // SET, so concurrent instances' increments accumulate). Because the push - // is all-or-nothing, a failure means nothing was applied, so restoring - // the claims for a clean retry cannot double-count. + // SET, so concurrent instances' increments accumulate). if let Err(e) = self.push_deltas(&mut conn, &plans).await { - tracing::warn!("rate-limit delta push failed: {e}"); - self.restore_claims(&plans); + tracing::warn!("rate-limit delta push failed, dropping this interval: {e}"); return; } @@ -271,24 +279,6 @@ impl GlobalCounters { self.apply_estimates(now, &plans, &reads); } - /// Return claimed deltas to the state when a push never reached the store, so - /// the next tick retries them. A delta whose epoch rolled meanwhile is owed - /// to the epoch it was accumulated in, so it goes back to `carryover`. - fn restore_claims(&self, plans: &[PushPlan]) { - for p in plans.iter().filter(|p| p.delta > 0) { - if let Some(mut s) = self.states.get_mut(&p.key) { - if !p.is_carryover && s.epoch == p.push_epoch { - s.local_count += p.delta; - } else { - if s.carryover == 0 { - s.carryover_epoch = p.push_epoch; - } - s.carryover += p.delta; - } - } - } - } - /// Refresh each key's cached estimate of the fleet's consumption. fn apply_estimates(&self, now: Duration, plans: &[PushPlan], reads: &[(u64, u64)]) { for (p, (cur, prev)) in plans.iter().zip(reads) { From 04a3ed16586ca426311aec82a1c2722cd0d34626 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 13:06:37 +0300 Subject: [PATCH 24/55] style(shield): match phase on borrowed key Match the rule phase on &key so the KeySource is unambiguously not moved before the fingerprint and CompiledRule reuse it. No behaviour change. --- src/shield/matcher.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/shield/matcher.rs b/src/shield/matcher.rs index 1a4a24b..f5e107e 100644 --- a/src/shield/matcher.rs +++ b/src/shield/matcher.rs @@ -136,7 +136,7 @@ pub fn compile_rules( KeySourceConfig::Header { name } => KeySource::Header(name.clone()), KeySourceConfig::JwtClaim { claim } => KeySource::JwtClaim(claim.clone()), }; - let phase = match key { + let phase = match &key { KeySource::JwtClaim(_) => Phase::PostAuth, KeySource::Ip | KeySource::Header(_) => Phase::PreAuth, }; From 3b0280e9979cfff7fe5d9508ece7f5b6f23d3965 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 13:06:37 +0300 Subject: [PATCH 25/55] fix(shield): validate limit-service endpoint and honor unknown-tier numbers - Parse limit_service.endpoint at build time and fail startup on a malformed URL, rather than silently disabling dynamic limits when the first background fetch fails to parse it. - When the service returns an unknown tier alongside explicit numbers (e.g. a mid rollout tier the proxy doesn't know yet), fall through to those numbers instead of treating the response as no limit, matching JWT resolution. --- src/shield/resolve.rs | 52 +++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 50 insertions(+), 2 deletions(-) diff --git a/src/shield/resolve.rs b/src/shield/resolve.rs index b5543c4..87d435c 100644 --- a/src/shield/resolve.rs +++ b/src/shield/resolve.rs @@ -164,6 +164,11 @@ impl LimitService { cfg: &LimitServiceConfig, profiles: HashMap, ) -> Result, String> { + // Validate the endpoint at build time: a malformed URL is a config error + // for a security control, so fail startup rather than silently disabling + // dynamic limits when the first background fetch fails to parse it. + reqwest::Url::parse(&cfg.endpoint) + .map_err(|e| format!("invalid limit_service.endpoint {:?}: {e}", cfg.endpoint))?; let client = reqwest::Client::builder() .timeout(Duration::from_millis(cfg.timeout_ms.max(1))) .tls_backend_preconfigured(crate::auth::jwks::build_tls_config()) @@ -309,8 +314,12 @@ impl LimitService { /// Map a service response to a compiled limit: a tier name resolves against /// the configured profiles; otherwise explicit numbers apply. fn map_response(&self, body: LimitResponse) -> Option { - if let Some(tier) = body.tier { - return self.profiles.get(&tier).copied(); + if let Some(tier) = &body.tier { + if let Some(profile) = self.profiles.get(tier) { + return Some(*profile); + } + // Unknown tier (e.g. mid rollout): fall through to explicit numbers if + // the response also carried them, matching JWT resolution. } body.rate_per_min .and_then(|rpm| profile_from_numbers(rpm, body.burst.unwrap_or(rpm))) @@ -406,6 +415,45 @@ mod tests { ); } + fn service(endpoint: &str) -> Arc { + LimitService::build( + &LimitServiceConfig { + endpoint: endpoint.to_string(), + ttl_secs: 60, + timeout_ms: 50, + }, + profiles(), + ) + .unwrap() + } + + #[test] + fn service_unknown_tier_falls_through_to_numbers() { + let svc = service("http://127.0.0.1:9/"); + // Unknown tier but explicit numbers present → numbers win (like JWT). + let p = svc + .map_response(LimitResponse { + tier: Some("gold".to_string()), + rate_per_min: Some(50), + burst: None, + }) + .unwrap(); + assert_eq!(p.limit, 50); + } + + #[test] + fn build_rejects_malformed_endpoint() { + let bad = LimitService::build( + &LimitServiceConfig { + endpoint: "not a url".to_string(), + ttl_secs: 60, + timeout_ms: 50, + }, + profiles(), + ); + assert!(bad.is_err()); + } + #[test] fn jwt_zero_rate_is_not_a_usable_limit() { // A dynamic rate of 0 must not clamp to 1; it yields no limit so the From 7412910f73f6749d7de0625ec1aa68a211103bf0 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 13:26:24 +0300 Subject: [PATCH 26/55] test(shield): add regression test for carryover epoch collapse When two window boundaries are crossed before a reconcile claims the carryover (reconcile interval misconfigured longer than the window), roll_to adds the new epoch's count onto the previous carryover under one epoch label. The test expects the latest window to be kept, not the two collapsed; it fails on current code. --- src/shield/global.rs | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/src/shield/global.rs b/src/shield/global.rs index a435bca..8378514 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -434,6 +434,25 @@ mod tests { assert_eq!(s.epoch, 11); } + #[test] + fn consecutive_rolls_do_not_collapse_epochs() { + // If reconcile does not run between two boundary crossings (interval + // misconfigured longer than the window), a second roll must not add a new + // epoch's count onto the previous carryover under one epoch label. It + // keeps the most recent window (dropping the older) rather than corrupting + // both with a collapsed count. + let mut s = KeyState::new(10, 60); + s.local_count = 3; + s.roll_to(11); // carryover = 3 owed to epoch 10 + s.local_count = 4; + s.roll_to(12); // second roll before any reconcile claimed the carryover + assert_eq!( + s.carryover, 4, + "keeps the latest window, not 3 + 4 collapsed" + ); + assert_eq!(s.carryover_epoch, 11); + } + #[test] fn stale_estimate_not_applied_after_window_reset() { let g = counters(); From bc682739b26f6f3054dacca8578b95a17cc9797a Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 13:30:29 +0300 Subject: [PATCH 27/55] fix(shield): degrade on stale estimate, gate on carryover, no epoch collapse Three fleet-gate correctness fixes: - Expire the remote estimate: once the store has been unreachable for several intervals the cached estimate is ignored and the gate falls back to per-instance counts, instead of subtracting a frozen estimate forever and rejecting a key indefinitely during an outage. - Count carryover in fleet_remaining: unpushed previous-window admits still contribute to the sliding window, so the gate must subtract them too, not just remote_estimate + local_count. - roll_to keeps only the most recent window when a second boundary is crossed before a reconcile claims the carryover (interval misconfigured longer than the window), rather than collapsing two epochs' counts under one label. --- src/shield/global.rs | 36 +++++++++++++++++++++++++++++++++--- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/src/shield/global.rs b/src/shield/global.rs index 8378514..89a9f02 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -56,11 +56,16 @@ struct KeyState { /// The fleet's sliding consumption (including this instance's own claimed /// admits, which now live in the shared counter), from the last pull. remote_estimate: u64, + /// When `remote_estimate` was last refreshed from the store. If it goes + /// stale (the store is unreachable for several intervals), the gate stops + /// trusting it and degrades to per-instance limiting. + estimate_at: Instant, last_seen: Instant, } impl KeyState { fn new(epoch: u64, window_secs: u64) -> Self { + let now = Instant::now(); Self { epoch, window_secs, @@ -68,7 +73,8 @@ impl KeyState { carryover: 0, carryover_epoch: 0, remote_estimate: 0, - last_seen: Instant::now(), + estimate_at: now, + last_seen: now, } } @@ -83,7 +89,17 @@ impl KeyState { fn roll_to(&mut self, epoch: u64) { if self.epoch != epoch { if self.local_count > 0 { - self.carryover += self.local_count; + if self.carryover > 0 && self.carryover_epoch != self.epoch { + // An unclaimed carryover from an earlier epoch still exists: + // reconcile has not run across two boundaries (interval + // misconfigured longer than the window). Keep only the most + // recent window rather than collapsing two epochs' counts + // under one label, which would corrupt both. Dropping the + // older is a bounded under-count consistent with best-effort. + self.carryover = self.local_count; + } else { + self.carryover += self.local_count; + } self.carryover_epoch = self.epoch; } self.epoch = epoch; @@ -130,7 +146,20 @@ impl GlobalCounters { /// keeping the hot path lock-free and non-blocking. pub fn fleet_remaining(&self, key: &str, budget: u64, window: Duration) -> u64 { let state = self.state_for(key, window); - budget.saturating_sub(state.remote_estimate + state.local_count) + // Ignore the remote estimate once it is stale (store unreachable for + // several intervals): keep gating on this instance's own counts only, + // which is the documented degrade-to-per-instance behaviour, instead of + // subtracting a frozen estimate forever. + let stale_after = (self.interval * 4).max(Duration::from_secs(2)); + let remote = if state.estimate_at.elapsed() < stale_after { + state.remote_estimate + } else { + 0 + }; + // Subtract carryover too: unpushed previous-window admits still count in + // the sliding window until the next tick publishes them. + let used = remote + state.local_count + state.carryover; + budget.saturating_sub(used) } /// Record one admitted request for `key`, to be pushed to the store on the @@ -301,6 +330,7 @@ impl GlobalCounters { // the full sliding estimate is used with no self-subtraction. let est = window::sliding_estimate(*cur, *prev, elapsed, p.window); s.remote_estimate = est.round().max(0.0) as u64; + s.estimate_at = Instant::now(); } } } From f0e5c3040fabcb1ae2bfb6ebd58588bf22ac975c Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 13:30:47 +0300 Subject: [PATCH 28/55] fix(shield): reject invalid header key names at startup A header-keyed rule whose configured name isn't a valid HTTP header (e.g. it contains whitespace) would never match, silently downgrading a per-API-key limit to per-IP. Parse the name as a HeaderName during compile and fail startup on an invalid value. --- src/shield/matcher.rs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/shield/matcher.rs b/src/shield/matcher.rs index f5e107e..ddbb05c 100644 --- a/src/shield/matcher.rs +++ b/src/shield/matcher.rs @@ -133,7 +133,15 @@ pub fn compile_rules( } let key = match &c.key { KeySourceConfig::Ip => KeySource::Ip, - KeySourceConfig::Header { name } => KeySource::Header(name.clone()), + KeySourceConfig::Header { name } => { + // Reject an invalid header name at build time: a typo (e.g. + // whitespace) would never match a real header, silently + // downgrading a per-API-key limit into a per-IP one. + http::HeaderName::from_bytes(name.as_bytes()).map_err(|_| { + format!("rule {:?} has invalid header name {name:?}", c.pattern) + })?; + KeySource::Header(name.clone()) + } KeySourceConfig::JwtClaim { claim } => KeySource::JwtClaim(claim.clone()), }; let phase = match &key { From dd46db939bb4d2d46243974114fc47a7f36fea3f Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 13:30:47 +0300 Subject: [PATCH 29/55] fix(shield): throttle limit-service retries for stale keys during an outage A hot key with a stale cached limit re-triggered a fetch on every request while the service was down, since only brand-new keys were negative-cached. Reset the existing entry's fetch timestamp on failure so it is served for another TTL before the next retry, capping outbound calls at one per key per TTL. --- src/shield/resolve.rs | 34 ++++++++++++++++++++-------------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/src/shield/resolve.rs b/src/shield/resolve.rs index 87d435c..95ac264 100644 --- a/src/shield/resolve.rs +++ b/src/shield/resolve.rs @@ -263,20 +263,26 @@ impl LimitService { ); } Err(()) => { - // Leave an existing (stale) entry in place (fail-static). For a - // brand-new key, negative-cache the failure so a client rotating - // key values during an outage can't spawn unbounded fetch tasks; - // it is retried after the TTL like any stale entry. - if !this.cache.contains_key(&key) { - let now = Instant::now(); - this.cache.insert( - key.clone(), - Cached { - profile: None, - at: now, - last_access: now, - }, - ); + // Throttle retries during an outage. Reset the fetch timestamp + // (`at`) so the entry is treated as fresh for another TTL and + // the next request serves it without immediately re-fetching; + // otherwise a hot stale key would spawn one outbound call per + // request. Keep the last-good profile for an existing entry + // (fail-static); negative-cache a brand-new key so a client + // rotating key values can't spawn unbounded fetch tasks. + let now = Instant::now(); + match this.cache.get_mut(&key) { + Some(mut c) => c.at = now, + None => { + this.cache.insert( + key.clone(), + Cached { + profile: None, + at: now, + last_access: now, + }, + ); + } } } } From 02aad014ad4a5968b6246061e04d26bf91a4d04c Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 13:30:47 +0300 Subject: [PATCH 30/55] fix(shield): report tightest phase budget and a decaying global Retry-After - On an allowed response where both phases matched, report the smaller remaining across phases (overwrite the inner limiter's headers only when this phase is tighter), so a client near the tighter budget backs off in time. Inner rejections (remaining 0) are still preserved. - A fleet-gate 429 now advertises a modest poll interval (a tenth of the window, min 1s) instead of the epoch boundary: the sliding-window estimate decays continuously and does not free capacity at the boundary, so a client waiting exactly to the boundary could be rejected again. --- src/shield/mod.rs | 49 +++++++++++++++++++++++++++++++---------------- 1 file changed, 33 insertions(+), 16 deletions(-) diff --git a/src/shield/mod.rs b/src/shield/mod.rs index 9defeb2..56ab405 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -271,31 +271,48 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> }; let mut response = next.run(request).await; - // Don't overwrite rate-limit headers an inner limiter already set. With - // defense-in-depth (a pre-auth IP/header rule and a post-auth principal rule - // on the same path), the inner post-auth verdict is the more specific one and - // must reach the client intact, so the outer pre-auth phase leaves it alone. - if !response.headers().contains_key("ratelimit-limit") { - attach_rate_headers(response.headers_mut(), profile.limit, reported, &verdict); - } + // Report the tightest budget across phases. With defense-in-depth (a pre-auth + // and a post-auth rule on the same path), an inner limiter may already have + // set headers on the way out; overwrite them only when this (outer) phase's + // remaining is smaller, so the client always sees the budget that will bite + // first. An inner rejection carries remaining 0, so it is never overwritten. + maybe_tighten_rate_headers(response.headers_mut(), profile.limit, reported, &verdict); response } -/// A `429` for a request rejected by the fleet-wide gate. `Retry-After` / -/// `RateLimit-Reset` point at the end of the current window, when the sliding -/// estimate will have decayed. +/// Set the `RateLimit-*` headers unless an inner limiter already advertised a +/// tighter (smaller `remaining`) budget, which must reach the client intact. +fn maybe_tighten_rate_headers( + headers: &mut HeaderMap, + limit: u64, + remaining: u64, + verdict: &Verdict, +) { + let existing = headers + .get("ratelimit-remaining") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()); + match existing { + Some(inner) if inner <= remaining => {} // inner budget is tighter (or equal): keep it + _ => attach_rate_headers(headers, limit, remaining, verdict), + } +} + +/// A `429` for a request rejected by the fleet-wide gate. The sliding-window +/// estimate decays continuously (it does not free capacity at the epoch +/// boundary), so `Retry-After` is a modest poll interval rather than the time to +/// the boundary, which a client could wait out and still be rejected. #[cfg(feature = "redis")] fn global_reject(limit: u64, window: Duration) -> Response { - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or(Duration::ZERO); - let remaining = window.saturating_sub(window::elapsed_in_window(now, window)); + // A tenth of the window (at least 1s): long enough not to hammer, short + // enough to retry as the window decays or other instances free budget. + let backoff = (window / 10).max(Duration::from_secs(1)); let verdict = Verdict { allowed: false, new_tat: Duration::ZERO, remaining: 0, - retry_after: remaining, - reset_after: remaining, + retry_after: backoff, + reset_after: backoff, }; too_many_requests(limit, 0, &verdict) } From b519fd433dd038e6eca5a1872626210b10ec0f35 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 14:29:52 +0300 Subject: [PATCH 31/55] fix(shield): auto-reconnect the shared store; note window-reset is not client-driven - Use a self-reconnecting ConnectionManager for the shared store instead of a cached MultiplexedConnection. A Redis restart or dropped TCP connection is now re-established internally, so reconciled mode recovers on its own rather than staying degraded until the proxy restarts. - Clarify that a key's window only changes when its resolved tier does (from the signed JWT, the limit service, or config, never client input), so the one-window under-count on reset is a rare, bounded, non-client-triggerable event, not an abuse vector. --- Cargo.toml | 2 +- src/shield/global.rs | 20 +++++++++++++------- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index c158c53..f454605 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -73,7 +73,7 @@ dashmap = "6" async-trait = "0.1" ipnet = "2" # Optional shared rate-limit store for multi-instance deployments. -redis = { version = "1.2", features = ["tokio-comp"], optional = true } +redis = { version = "1.2", features = ["tokio-comp", "connection-manager"], optional = true } # Auth (JWT validation, route policies). # Crypto backend is selected via this crate's `rust_crypto` / `aws_lc_rs` diff --git a/src/shield/global.rs b/src/shield/global.rs index 89a9f02..2078d48 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -111,7 +111,7 @@ impl KeyState { /// Cross-instance counter reconciliation over a shared Redis-protocol store. pub struct GlobalCounters { client: redis::Client, - conn: tokio::sync::OnceCell, + conn: tokio::sync::OnceCell, interval: Duration, states: dashmap::DashMap, } @@ -188,8 +188,11 @@ impl GlobalCounters { if state.window_secs != window_secs { // A window change is a different accounting unit, so the entry resets. // Any unpushed admits from the old window are dropped rather than - // remapped (a different window can't share a counter); this is a rare, - // bounded under-count on a tier change, consistent with the fleet + // remapped (a different window can't share a counter). The window only + // changes when the resolved tier does, and the tier comes from the + // signed JWT, the limit service, or config, never from client input, + // so this is not client-triggerable: it is a rare, bounded one-window + // under-count on a legitimate tier change, consistent with the fleet // layer's best-effort, never-double contract. *state = KeyState::new(ep, window_secs); } @@ -215,9 +218,12 @@ impl GlobalCounters { }); } - async fn connection(&self) -> redis::RedisResult { + /// A cloneable, self-reconnecting handle to the shared store. `ConnectionManager` + /// re-establishes the underlying connection internally after a drop (e.g. a + /// Redis restart), so reconciled mode recovers without a proxy restart. + async fn connection(&self) -> redis::RedisResult { self.conn - .get_or_try_init(|| self.client.get_multiplexed_async_connection()) + .get_or_try_init(|| redis::aio::ConnectionManager::new(self.client.clone())) .await .cloned() } @@ -341,7 +347,7 @@ impl GlobalCounters { /// restores the claims and re-pushes the same deltas next tick. async fn push_deltas( &self, - conn: &mut redis::aio::MultiplexedConnection, + conn: &mut redis::aio::ConnectionManager, plans: &[PushPlan], ) -> redis::RedisResult<()> { let mut pipe = redis::pipe(); @@ -364,7 +370,7 @@ impl GlobalCounters { /// `(cur, prev)` per plan in order. async fn read_epochs( &self, - conn: &mut redis::aio::MultiplexedConnection, + conn: &mut redis::aio::ConnectionManager, plans: &[PushPlan], ) -> redis::RedisResult> { let mut keys: Vec = Vec::with_capacity(plans.len() * 2); From 0697fcb068434347cd563a16bfce412172e9f755 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 14:30:06 +0300 Subject: [PATCH 32/55] fix(shield): require an http/https limit-service endpoint URL syntax validation alone let schemes like redis:// or file:// pass startup even though the background GET can never use them, silently disabling dynamic limits. Reject any non-http(s) scheme at build time. --- src/shield/resolve.rs | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/src/shield/resolve.rs b/src/shield/resolve.rs index 95ac264..6c7eda7 100644 --- a/src/shield/resolve.rs +++ b/src/shield/resolve.rs @@ -164,11 +164,17 @@ impl LimitService { cfg: &LimitServiceConfig, profiles: HashMap, ) -> Result, String> { - // Validate the endpoint at build time: a malformed URL is a config error - // for a security control, so fail startup rather than silently disabling - // dynamic limits when the first background fetch fails to parse it. - reqwest::Url::parse(&cfg.endpoint) + // Validate the endpoint at build time: a malformed or non-HTTP URL is a + // config error for a security control, so fail startup rather than + // silently disabling dynamic limits when the background fetch later fails. + let url = reqwest::Url::parse(&cfg.endpoint) .map_err(|e| format!("invalid limit_service.endpoint {:?}: {e}", cfg.endpoint))?; + if !matches!(url.scheme(), "http" | "https") { + return Err(format!( + "limit_service.endpoint must be http/https, got scheme {:?}", + url.scheme() + )); + } let client = reqwest::Client::builder() .timeout(Duration::from_millis(cfg.timeout_ms.max(1))) .tls_backend_preconfigured(crate::auth::jwks::build_tls_config()) @@ -447,6 +453,22 @@ mod tests { assert_eq!(p.limit, 50); } + #[test] + fn build_rejects_non_http_endpoint() { + // Syntactically valid URLs that reqwest GET can't use must be rejected. + for ep in ["redis://127.0.0.1/", "file:///etc/passwd", "ftp://h/x"] { + let r = LimitService::build( + &LimitServiceConfig { + endpoint: ep.to_string(), + ttl_secs: 60, + timeout_ms: 50, + }, + profiles(), + ); + assert!(r.is_err(), "expected {ep} to be rejected"); + } + } + #[test] fn build_rejects_malformed_endpoint() { let bad = LimitService::build( From c9afc81674f5ae65503f10517528054b512a654b Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 14:30:06 +0300 Subject: [PATCH 33/55] fix(shield): normalize header key names before fingerprinting A header-keyed rule configured as X-API-Key on one instance and x-api-key on another selects the same header but produced different rule fingerprints, splitting the counter namespace so one API key could consume multiple budgets. Store the parsed HeaderName's normalized (lowercased) form so the fingerprint is casing-independent. --- src/shield/matcher.rs | 51 ++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 48 insertions(+), 3 deletions(-) diff --git a/src/shield/matcher.rs b/src/shield/matcher.rs index ddbb05c..71ff787 100644 --- a/src/shield/matcher.rs +++ b/src/shield/matcher.rs @@ -136,11 +136,14 @@ pub fn compile_rules( KeySourceConfig::Header { name } => { // Reject an invalid header name at build time: a typo (e.g. // whitespace) would never match a real header, silently - // downgrading a per-API-key limit into a per-IP one. - http::HeaderName::from_bytes(name.as_bytes()).map_err(|_| { + // downgrading a per-API-key limit into a per-IP one. Store the + // normalized (lowercased) form so a casing-only config + // difference between instances (`X-API-Key` vs `x-api-key`) + // yields the same fingerprint and shares one counter namespace. + let hn = http::HeaderName::from_bytes(name.as_bytes()).map_err(|_| { format!("rule {:?} has invalid header name {name:?}", c.pattern) })?; - KeySource::Header(name.clone()) + KeySource::Header(hn.as_str().to_string()) } KeySourceConfig::JwtClaim { claim } => KeySource::JwtClaim(claim.clone()), }; @@ -196,6 +199,48 @@ mod tests { assert_eq!(p.limit, 100); } + #[test] + fn header_key_fingerprint_is_case_insensitive() { + // Casing-only differences in the configured header name must map to the + // same fingerprint so instances share one counter namespace. + let rules = compile_rules( + &[ + RateRuleConfig { + pattern: "/a".to_string(), + key: KeySourceConfig::Header { + name: "X-API-Key".to_string(), + }, + profile: Some("auth".to_string()), + }, + RateRuleConfig { + pattern: "/a".to_string(), + key: KeySourceConfig::Header { + name: "x-api-key".to_string(), + }, + profile: Some("auth".to_string()), + }, + ], + &profiles(), + ) + .unwrap(); + assert_eq!(rules[0].fingerprint, rules[1].fingerprint); + } + + #[test] + fn invalid_header_name_is_rejected() { + let err = compile_rules( + &[RateRuleConfig { + pattern: "/a".to_string(), + key: KeySourceConfig::Header { + name: "bad name".to_string(), + }, + profile: Some("auth".to_string()), + }], + &profiles(), + ); + assert!(err.is_err()); + } + #[test] fn zero_rate_or_burst_is_rejected() { assert!(compile_profile(&LimitProfileConfig { From 3488749dbc29074e337dc6552ef020993e7d7bb4 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 15:11:24 +0300 Subject: [PATCH 34/55] fix(shield): bound rate-limit state under key-cardinality floods Idle sweeps bound how long entries live, not peak cardinality. Add hard caps so client-controlled key churn (IP / principal / API key) cannot grow memory without bound: - limit-service cache: cap at 100k entries, evicting least-recently-accessed on overflow, plus a semaphore capping concurrent background fetches at 32 so a miss storm cannot fan out unbounded outbound calls or tasks - GCRA store: size-triggered drained-key eviction between timer sweeps; only drained (past-TAT) keys are dropped, never active ones (dropping a limited key would reset its budget and hand a flooder a fresh burst) - fleet counters: cap tracked keys at 500k; a new key past the cap degrades to per-instance limiting (still locally capped) rather than being tracked, matching the existing best-effort, never-double fleet contract --- src/shield/global.rs | 76 +++++++++++++++++++++++++++++++------------ src/shield/resolve.rs | 40 ++++++++++++++++++++++- src/shield/store.rs | 13 +++++++- 3 files changed, 107 insertions(+), 22 deletions(-) diff --git a/src/shield/global.rs b/src/shield/global.rs index 2078d48..297ccde 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -16,6 +16,8 @@ use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; +use dashmap::mapref::entry::Entry; + use super::window; /// Redis key namespace for a rate-limit key at a given epoch. @@ -35,6 +37,15 @@ fn unix_now() -> Duration { /// key cardinality). const KEY_TTL: Duration = Duration::from_secs(600); +/// Hard cap on tracked keys. [`evict_stale`](GlobalCounters::evict_stale) bounds +/// retention *time*, but only runs once per reconcile tick; a burst of distinct +/// keys between ticks could still grow the map. Past this cap, a *new* key is not +/// fleet-tracked and simply degrades to per-instance limiting (the local +/// [`GcraStore`](super::store::GcraStore) still caps it) rather than being fleet +/// under-counted; this is the same best-effort tolerance as a dropped push. +/// Already-tracked keys are unaffected, so a real principal's budget is kept. +const MAX_KEYS: usize = 500_000; + /// Per-key reconciliation state. /// /// Admit accounting uses a claim model: `local_count` holds only admits *not yet @@ -145,7 +156,12 @@ impl GlobalCounters { /// documented one-interval overshoot, which is the accepted trade-off for /// keeping the hot path lock-free and non-blocking. pub fn fleet_remaining(&self, key: &str, budget: u64, window: Duration) -> u64 { - let state = self.state_for(key, window); + // When the map is full and this key is new, don't fleet-gate it: report + // the full budget so the local limiter alone decides (degrade-to- + // per-instance). Tracking it would breach the memory cap under a flood. + let Some(state) = self.state_for(key, window) else { + return budget; + }; // Ignore the remote estimate once it is stale (store unreachable for // several intervals): keep gating on this instance's own counts only, // which is the documented degrade-to-per-instance behaviour, instead of @@ -165,8 +181,12 @@ impl GlobalCounters { /// Record one admitted request for `key`, to be pushed to the store on the /// next background tick. pub fn record(&self, key: &str, window: Duration) { - let mut state = self.state_for(key, window); - state.local_count += 1; + // If the map is full and this key is untracked, skip: the admit is + // enforced locally and simply isn't published to the fleet (bounded + // under-count), which is preferable to breaching the memory cap. + if let Some(mut state) = self.state_for(key, window) { + state.local_count += 1; + } } /// Look up (or create) the per-key state, rolled to the current epoch and @@ -177,14 +197,19 @@ impl GlobalCounters { &self, key: &str, window: Duration, - ) -> dashmap::mapref::one::RefMut<'_, String, KeyState> { + ) -> Option> { let now = unix_now(); let window_secs = window.as_secs().max(1); let ep = window::epoch(now, window); - let mut state = self - .states - .entry(key.to_string()) - .or_insert_with(|| KeyState::new(ep, window_secs)); + // Soft cap check before the entry: a new key past the cap is not tracked. + // The `len()` read races with concurrent inserts, but the cap is a memory + // guard, not an exact limit, so a few entries of overshoot are harmless. + let over_cap = self.states.len() >= MAX_KEYS; + let mut state = match self.states.entry(key.to_string()) { + Entry::Occupied(o) => o.into_ref(), + Entry::Vacant(_) if over_cap => return None, + Entry::Vacant(v) => v.insert(KeyState::new(ep, window_secs)), + }; if state.window_secs != window_secs { // A window change is a different accounting unit, so the entry resets. // Any unpushed admits from the old window are dropped rather than @@ -198,15 +223,18 @@ impl GlobalCounters { } state.roll_to(ep); state.last_seen = Instant::now(); - state + Some(state) } - /// Spawn the background reconciliation loop. A no-op (with a warning) when - /// called outside a Tokio runtime, so the local shaper still works. - pub fn spawn(self: &Arc) { + /// Spawn the background reconciliation loop. Returns `false` (without + /// spawning) when called outside a Tokio runtime, so the caller can drop the + /// fleet gate and fall back to per-instance limiting rather than gating on a + /// view that would never be reconciled. + #[must_use] + pub fn spawn(self: &Arc) -> bool { if tokio::runtime::Handle::try_current().is_err() { tracing::warn!("rate-limit reconciler not started: no Tokio runtime in this context"); - return; + return false; } let this = self.clone(); tokio::spawn(async move { @@ -216,6 +244,7 @@ impl GlobalCounters { this.reconcile().await; } }); + true } /// A cloneable, self-reconnecting handle to the shared store. `ConnectionManager` @@ -281,12 +310,16 @@ impl GlobalCounters { } // Claimed deltas are fire-and-forget: once claimed (zeroed), a push that - // fails or whose EXEC ack is lost simply drops them rather than restoring. - // Restoring risked publishing a committed-but-unacked batch twice (a false - // 429); dropping instead under-counts this instance's last interval, which - // is exactly the documented "store unreachable → degrade to per-instance" - // behaviour. It also means an unpushable delta never accumulates across - // window boundaries into a collapsed carryover. + // fails (connection or transaction) simply drops them rather than + // restoring. This is deliberate, not a leak: + // * Restoring risks publishing a committed-but-unacked batch twice (a + // false 429), and across a sustained outage it would accumulate + // unbounded local_count and then dump a huge spike on recovery. + // * Dropping instead under-counts only this instance's last interval, + // which is exactly the documented "store unreachable → degrade to + // per-instance limiting" behaviour, and never over-counts. + // So the failure mode is a bounded, self-correcting under-count (slight + // over-admit), never a false rejection or a recovery spike. let mut conn = match self.connection().await { Ok(c) => c, Err(e) => { @@ -335,7 +368,10 @@ impl GlobalCounters { // the gate adds only the still-unclaimed `local_count` on top, so // the full sliding estimate is used with no self-subtraction. let est = window::sliding_estimate(*cur, *prev, elapsed, p.window); - s.remote_estimate = est.round().max(0.0) as u64; + // Round the fractional sliding estimate UP: under-counting the + // fleet would let the gate admit past the budget, so bias to the + // conservative side. + s.remote_estimate = est.ceil().max(0.0) as u64; s.estimate_at = Instant::now(); } } diff --git a/src/shield/resolve.rs b/src/shield/resolve.rs index 6c7eda7..92a852a 100644 --- a/src/shield/resolve.rs +++ b/src/shield/resolve.rs @@ -137,6 +137,16 @@ struct Cached { /// Sweep the cache of long-idle keys at most once per this interval. const SWEEP_INTERVAL: Duration = Duration::from_secs(60); +/// Hard cap on cached resolutions, bounding peak memory under rapid key rotation +/// (idle eviction bounds retention time, not peak cardinality). When full, the +/// least-recently-accessed entries are dropped: an attacker's one-shot keys are +/// the oldest and get evicted, while an actively-reused key stays warm. +const MAX_CACHE_ENTRIES: usize = 100_000; + +/// Cap on concurrent background limit-service fetches, bounding outbound calls +/// and tasks when many distinct keys miss at once. +const MAX_CONCURRENT_FETCHES: usize = 32; + /// External limit-resolution service with an async, non-blocking cache. pub struct LimitService { endpoint: String, @@ -151,6 +161,8 @@ pub struct LimitService { cache: dashmap::DashMap, /// Keys with a background fetch already in flight (dedupes refreshes). inflight: dashmap::DashMap, + /// Global cap on concurrent background fetches (in addition to per-key dedup). + fetch_slots: Arc, base: Instant, last_sweep_ms: std::sync::atomic::AtomicU64, } @@ -190,6 +202,7 @@ impl LimitService { profiles, cache: dashmap::DashMap::new(), inflight: dashmap::DashMap::new(), + fetch_slots: Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_FETCHES)), base: Instant::now(), last_sweep_ms: std::sync::atomic::AtomicU64::new(0), })) @@ -215,10 +228,24 @@ impl LimitService { /// Drop entries not accessed within `evict_after` (idle keys), keyed on last /// access rather than last fetch so an active-but-unrefreshable key survives. + /// Then enforce the hard capacity, evicting the least-recently-accessed + /// entries so rapid key rotation cannot grow the cache without bound. fn sweep(&self) { let evict_after = self.evict_after; self.cache .retain(|_, c| c.last_access.elapsed() < evict_after); + if self.cache.len() > MAX_CACHE_ENTRIES { + let mut ages: Vec<(String, Instant)> = self + .cache + .iter() + .map(|e| (e.key().clone(), e.value().last_access)) + .collect(); + // Oldest first; drop everything past the cap. + ages.sort_by_key(|(_, t)| *t); + for (key, _) in ages.into_iter().take(self.cache.len() - MAX_CACHE_ENTRIES) { + self.cache.remove(&key); + } + } } /// Resolve `key`'s limit from the cache, serving a stale value while a @@ -246,7 +273,10 @@ impl LimitService { } } - /// Spawn a single background fetch for `key` (deduped by `inflight`). + /// Spawn a single background fetch for `key` (deduped by `inflight`, and + /// globally bounded by `fetch_slots`). When no fetch slot is free, skip the + /// refresh and serve the existing stale/static result rather than piling up + /// outbound calls under a burst of distinct keys. fn trigger_refresh(self: &Arc, key: String) { match self.inflight.entry(key.clone()) { Entry::Occupied(_) => return, @@ -254,8 +284,16 @@ impl LimitService { v.insert(()); } } + let permit = match self.fetch_slots.clone().try_acquire_owned() { + Ok(p) => p, + Err(_) => { + self.inflight.remove(&key); + return; + } + }; let this = self.clone(); tokio::spawn(async move { + let _permit = permit; // released when the fetch task ends match this.fetch(&key).await { Ok(resolved) => { let now = Instant::now(); diff --git a/src/shield/store.rs b/src/shield/store.rs index 33eccca..64eb464 100644 --- a/src/shield/store.rs +++ b/src/shield/store.rs @@ -19,6 +19,15 @@ use super::gcra::{Gcra, Verdict}; /// Run eviction of drained entries at most once per this interval. const SWEEP_INTERVAL_MS: u64 = 60_000; +/// When the map grows past this, run a drained-key eviction immediately instead +/// of waiting for the timer, bounding peak memory during a burst of distinct +/// keys. This only reclaims *drained* keys (TAT in the past), never active +/// (rate-limited) ones: dropping a limited key would reset its GCRA budget and +/// hand a flooding attacker a fresh burst, so the store deliberately holds +/// active limiter state. An all-active flood is shed by the fleet gate and +/// upstream, not by discarding the very state that enforces the limit. +const SWEEP_HIGH_WATER: usize = 200_000; + /// In-process per-instance GCRA store. // no-std: caller-provided Clock + spin/hashbrown map. #[derive(Debug)] @@ -89,7 +98,9 @@ impl GcraStore { fn maybe_sweep(&self, now: Duration) { let now_ms = u64::try_from(now.as_millis()).unwrap_or(u64::MAX); let last = self.last_sweep_ms.load(Ordering::Relaxed); - if now_ms.saturating_sub(last) < SWEEP_INTERVAL_MS { + // Time-based cadence, unless the map is already over the high-water mark, + // in which case sweep now regardless of when the last one ran. + if now_ms.saturating_sub(last) < SWEEP_INTERVAL_MS && self.tats.len() <= SWEEP_HIGH_WATER { return; } if self From d208a7c4cabc9b0a95194ccf18d5eb2da49ba69d Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 15:11:42 +0300 Subject: [PATCH 35/55] fix(shield): fail closed on unknown client IP and unstarted fleet gate - client_ip: when no connection peer is available, bucket as "unknown" instead of trusting client-supplied X-Forwarded-For / X-Real-IP. Without a trusted peer those headers are attacker-rotatable, so honouring them let a single client dodge a per-IP limit by varying the header. - fleet gate: keep it only when reconciliation actually started; otherwise drop to per-instance limiting rather than gating on an estimate that would never be refreshed. - document that the fleet budget is the sustained rate and per-instance burst is not pre-reserved fleet-wide, and fold burst into the README overshoot formula: (N-1) * (burst + rate * (interval/window)). --- README.md | 18 ++++++++++-------- src/shield/mod.rs | 35 ++++++++++++++--------------------- 2 files changed, 24 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index 57d799b..c7fe395 100644 --- a/README.md +++ b/README.md @@ -231,14 +231,16 @@ allowed responses. Behind a browser, these are exposed via CORS. unreachable, instances degrade to local limiting rather than failing requests. **Sizing the overshoot.** In reconciled mode the aggregate lags by up to one -`sync.interval_ms`, so within that lag each of the other instances can admit up -to a `rate` fraction of the window. With the interval expressed in the same time -unit as the window, the worst-case fleet overshoot is about -`(N - 1) × rate × (interval / window)` requests. For example, `rate = 1000/min`, -`interval = 500 ms`, `N = 4` gives `3 × 1000 × (0.5 / 60) ≈ 25` extra requests. -Shorter intervals tighten the bound at the cost of more store traffic. The global -view uses a sliding-window counter, so there is no boundary burst on top of this -lag. +`sync.interval_ms`. Within that lag each of the other instances can admit its +local `burst` plus a `rate` fraction of the window before the estimate catches +up (the fleet gate caps sustained fleet volume at `rate`, but `burst` is a +per-instance allowance the gate does not pre-reserve). With the interval in the +same time unit as the window, the worst-case fleet overshoot is about +`(N - 1) × (burst + rate × (interval / window))` requests. For example, +`burst = 100`, `rate = 1000/min`, `interval = 500 ms`, `N = 4` gives +`3 × (100 + 1000 × (0.5 / 60)) ≈ 325` extra requests. Smaller `burst` and shorter +intervals tighten the bound. The global view uses a sliding-window counter, so +there is no boundary burst on top of this lag. See the `shield:` block under [Configuration](#configuration) for the full schema. diff --git a/src/shield/mod.rs b/src/shield/mod.rs index 56ab405..d4b0e9e 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -108,8 +108,10 @@ impl Shield { &sync.redis_url, Duration::from_millis(sync.interval_ms), )?; - g.spawn(); - Some(g) + // Keep the fleet gate only if reconciliation actually started; + // otherwise fall back to per-instance limiting rather than gating + // on an estimate that would never be refreshed. + g.spawn().then_some(g) } None => None, }; @@ -230,7 +232,11 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> }; // Fleet gate first (read-only, cached), so the local shaper is not charged - // for a request the fleet-wide budget will reject. + // for a request the fleet-wide budget will reject. The fleet budget is the + // sustained rate (`profile.limit`); `burst` is deliberately a per-instance + // smoothing allowance, not a fleet-wide entitlement (honouring it fleet-wide + // would multiply the effective limit by N). A single instance's burst is + // therefore capped by the shared budget when the fleet is near it. #[cfg(feature = "redis")] let fleet_remaining = shield .global @@ -376,9 +382,10 @@ fn parse_cidr(s: &str) -> Result { /// `X-Forwarded-For` is trusted only when the direct `peer` is a configured /// trusted proxy, and even then the *rightmost* hop outside the trusted ranges /// is used: appending load balancers (nginx, ALB, GCP) add the connecting IP on -/// the right, so the leftmost entries are attacker-controlled. Without -/// connection info (e.g. a custom server that does not provide it) the headers -/// are taken as a best effort. +/// the right, so the leftmost entries are attacker-controlled. Without connection +/// info (a server not wired with `ConnectInfo`) we fail closed to a single +/// `"unknown"` bucket rather than trusting client-supplied forwarding headers, +/// which an attacker could otherwise rotate to dodge the limit. fn client_ip( peer: Option, headers: &HeaderMap, @@ -393,8 +400,7 @@ fn client_ip( } ip.to_string() } - // No connection info: fall back to the headers as a best effort. - None => best_effort_forwarded(headers).unwrap_or_else(|| "unknown".to_string()), + None => "unknown".to_string(), } } @@ -419,19 +425,6 @@ fn rightmost_untrusted(headers: &HeaderMap, trusted: &[ipnet::IpNet]) -> Option< header_str(headers, "x-real-ip") } -/// Best-effort client from forwarding headers when no peer is known: leftmost -/// `X-Forwarded-For`, then `X-Real-IP`. -fn best_effort_forwarded(headers: &HeaderMap) -> Option { - headers - .get("x-forwarded-for") - .and_then(|v| v.to_str().ok()) - .and_then(|v| v.split(',').next()) - .map(str::trim) - .filter(|s| !s.is_empty()) - .map(str::to_string) - .or_else(|| header_str(headers, "x-real-ip")) -} - /// Trimmed, non-empty value of a header. fn header_str(headers: &HeaderMap, name: &str) -> Option { headers From d2e7107633e6d92500593127583ea160c9c095ae Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 15:11:57 +0300 Subject: [PATCH 36/55] test(shield): reject unknown config fields and assert outer budget headers - config: regression test that a typo in a shield-config field is a hard deserialization error (deny_unknown_fields), so a misspelled security key cannot silently leave a limit unapplied. - two-phase: assert an allowed response advertises the tighter pre-auth remaining rather than the roomier post-auth one, so the client backs off against the binding limit. --- src/config.rs | 37 +++++++++++++++++++++++++++++++++---- src/shield/tests.rs | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 4 deletions(-) diff --git a/src/config.rs b/src/config.rs index 54287f0..37a04f3 100644 --- a/src/config.rs +++ b/src/config.rs @@ -376,6 +376,7 @@ fn default_authz_timeout_ms() -> u64 { /// the request path to approximate a fleet-wide limit; the request path never /// blocks on it. #[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] #[non_exhaustive] pub struct ShieldConfig { #[serde(default)] @@ -416,6 +417,7 @@ pub struct ShieldConfig { /// A named limit tier: a sustained rate plus an instantaneous burst capacity. #[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] #[non_exhaustive] pub struct LimitProfileConfig { /// Sustained rate as `"/"` (e.g. `"100/min"`, units @@ -428,11 +430,13 @@ pub struct LimitProfileConfig { } /// One rate-limit rule: a path pattern, how to key it, and an optional static -/// profile. The rule's *phase* (before or after auth) is derived automatically: -/// rules that need validated JWT claims (a `jwt_claim` key, or JWT-based limit -/// resolution) run after auth; the rest run before auth so anonymous floods are -/// shed before any signature verification. +/// profile. The rule's *phase* (before or after auth) is derived from its key +/// alone: a `jwt_claim` key needs validated claims so it runs after auth; `ip` +/// and `header` keys run before auth so anonymous floods are shed before any +/// signature verification. (`jwt_limits` therefore only takes effect on +/// `jwt_claim` rules, the only ones running with claims available.) #[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] #[non_exhaustive] pub struct RateRuleConfig { /// Glob path pattern (`*` within a segment, `**` across segments). @@ -475,6 +479,7 @@ pub enum KeySourceConfig { /// to a `profiles` entry (numbers stay tunable in config); direct numeric claims /// set the limit explicitly. #[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] #[non_exhaustive] pub struct JwtLimitConfig { /// Claim naming a profile tier (e.g. `"premium"`). Default: `ratelimit_tier`. @@ -503,6 +508,7 @@ fn default_burst_claim() -> String { /// numbers; results are cached and refreshed asynchronously, never on the request /// path. #[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] #[non_exhaustive] pub struct LimitServiceConfig { /// HTTP endpoint queried with the limit key; returns `{ tier }` or @@ -526,6 +532,7 @@ fn default_limit_timeout_ms() -> u64 { /// Asynchronous cross-instance reconciliation via a shared store. #[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] #[non_exhaustive] pub struct SyncConfig { /// Shared-store URL (e.g. `"redis://127.0.0.1/"`). Requires the `redis` build @@ -1074,6 +1081,28 @@ descriptors: assert_eq!(ProxyConfig::parse_rate("invalid"), None); } + #[test] + fn shield_rejects_unknown_field() { + // A typo in a shield-config field (here `profil` for `profile`) must be a + // hard error, not silently ignored: a misspelled security-control key + // would otherwise leave the intended limit unapplied. `deny_unknown_fields` + // on the shield structs turns the typo into a startup failure. + let yaml = r#" +upstream: + default: "grpc://localhost:4180" +shield: + enabled: true + profiles: + auth: { rate: "20/min", burst: 5 } + rules: + - pattern: "/v1/**" + key: { type: ip } + profil: "auth" +"#; + let err = serde_yaml::from_str::(yaml); + assert!(err.is_err(), "unknown shield field must be rejected"); + } + #[test] fn test_openapi_config_deserialize() { let yaml = r#" diff --git a/src/shield/tests.rs b/src/shield/tests.rs index 91005af..1770893 100644 --- a/src/shield/tests.rs +++ b/src/shield/tests.rs @@ -404,4 +404,37 @@ mod two_phase { "the rejecting inner rule's RateLimit-Limit must not be overwritten" ); } + + #[tokio::test] + async fn allowed_response_reports_tighter_outer_budget() { + // Pre-auth IP rule is tighter (burst 2) than the post-auth principal rule + // (burst 100). An allowed request must advertise the tighter pre-auth + // remaining, not the roomy principal one, so the client backs off in time. + let (sk, pem) = keypair_pem(); + let cfg = config( + vec![("tight", "60/min", Some(2)), ("wide", "100/min", Some(100))], + vec![ + rule("/api/**", KeySourceConfig::Ip, Some("tight")), + rule( + "/api/**", + KeySourceConfig::JwtClaim { + claim: "sub".to_string(), + }, + Some("wide"), + ), + ], + ); + let shield = Shield::build(&cfg).unwrap().unwrap(); + let app = stack(shield, auth(pem)); + + let req = Request::builder() + .uri("/api/x") + .header("authorization", format!("Bearer {}", token(&sk, "alice"))) + .body(Body::empty()) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + // tight burst 2, one admitted → 1 remaining; wide would show 99. + assert_eq!(resp.headers().get("ratelimit-remaining").unwrap(), "1"); + } } From 4c16ca65eafa45232b00e2d47dce2c1000edc515 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Fri, 17 Jul 2026 15:18:18 +0300 Subject: [PATCH 37/55] docs(shield): document burst>limit behaviour under fleet reconciliation Make explicit that a profile with burst greater than the sustained limit (e.g. rate 1/min, burst 5) is capped at the sustained limit fleet-wide, so the extra burst headroom a local-only deployment allows is not granted under reconciliation. This is the intended conservative choice: gating on burst would let the fleet sustain burst-per-window, loosening the sustained cap. --- src/shield/mod.rs | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/shield/mod.rs b/src/shield/mod.rs index d4b0e9e..fda8343 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -237,6 +237,15 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> // smoothing allowance, not a fleet-wide entitlement (honouring it fleet-wide // would multiply the effective limit by N). A single instance's burst is // therefore capped by the shared budget when the fleet is near it. + // + // Consequence for a `burst > limit` profile (e.g. `rate: "1/min", burst: 5`): + // under reconciliation the shared counter caps the key at the sustained + // `limit`, so the extra burst headroom that a local-only deployment would + // allow is not granted fleet-wide. This is intentional, not an oversight: + // gating on `burst` instead would let the fleet sustain `burst`-per-window, + // loosening the sustained cap by the burst factor. The conservative choice + // (never over-admit the fleet's sustained budget) wins; the local GCRA still + // smooths per-instance traffic. #[cfg(feature = "redis")] let fleet_remaining = shield .global From c12f87227f35b458b6466034fabbe943d993435a Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Sat, 18 Jul 2026 00:17:58 +0300 Subject: [PATCH 38/55] test(shield): add regression test for reset tie-break on equal remaining When two defense-in-depth phases both leave the same remaining, the response must advertise the longer reset (the budget that binds longest) so a paced client doesn't retry early into a still-blocked outer limit. This test drives a pre-auth 1/hour IP rule and a post-auth 1/min principal rule; it fails on current code (the inner 60s reset wins) and passes once the tie-break lands. --- src/shield/tests.rs | 48 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/src/shield/tests.rs b/src/shield/tests.rs index 1770893..51cf0e3 100644 --- a/src/shield/tests.rs +++ b/src/shield/tests.rs @@ -437,4 +437,52 @@ mod two_phase { // tight burst 2, one admitted → 1 remaining; wide would show 99. assert_eq!(resp.headers().get("ratelimit-remaining").unwrap(), "1"); } + + #[tokio::test] + async fn allowed_response_reports_longer_reset_on_remaining_tie() { + // Both phases exhaust their burst on the first admit, leaving the same + // remaining (0). The pre-auth IP rule (1/hour) binds far longer than the + // post-auth principal rule (1/min): the client must see the longer reset + // so it doesn't retry after ~60s and immediately hit the still-blocked + // hourly IP budget. On a remaining tie the longer-reset header wins. + let (sk, pem) = keypair_pem(); + let cfg = config( + vec![ + ("hourly", "1/hour", Some(1)), + ("minutely", "1/min", Some(1)), + ], + vec![ + rule("/api/**", KeySourceConfig::Ip, Some("hourly")), + rule( + "/api/**", + KeySourceConfig::JwtClaim { + claim: "sub".to_string(), + }, + Some("minutely"), + ), + ], + ); + let shield = Shield::build(&cfg).unwrap().unwrap(); + let app = stack(shield, auth(pem)); + + let req = Request::builder() + .uri("/api/x") + .header("authorization", format!("Bearer {}", token(&sk, "alice"))) + .body(Body::empty()) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(resp.headers().get("ratelimit-remaining").unwrap(), "0"); + // Longer reset (hourly IP ≈ 3600s) must win over the minutely 60s, so a + // paced client waits out the binding budget instead of retrying early. + let reset: u64 = resp + .headers() + .get("ratelimit-reset") + .unwrap() + .to_str() + .unwrap() + .parse() + .unwrap(); + assert!(reset > 120, "expected the hourly reset to win, got {reset}"); + } } From c8ae720964c4fa37d8f01c50bf950a1e56861153 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Sat, 18 Jul 2026 00:18:32 +0300 Subject: [PATCH 39/55] fix(shield): break RateLimit-Reset ties toward the longer-binding phase On a remaining tie between two defense-in-depth phases the header guard kept the inner phase's reset unconditionally, so a client pacing off an allowed response could see a short per-minute reset while an hourly IP cap still blocked it, retry early, and hit 429. The guard now keeps the inner headers on a tie only when their reset is at least as long as this phase's; otherwise the longer-binding budget's headers win. Covered by the test in c12f872. --- src/shield/mod.rs | 30 ++++++++++++++++++++++-------- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/src/shield/mod.rs b/src/shield/mod.rs index fda8343..d38b6cb 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -296,20 +296,34 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> } /// Set the `RateLimit-*` headers unless an inner limiter already advertised a -/// tighter (smaller `remaining`) budget, which must reach the client intact. +/// budget that binds at least as hard, which must reach the client intact. +/// "Binds harder" is a smaller `remaining`, and on a `remaining` tie the larger +/// `reset` wins: with both phases at `remaining=0`, a client pacing off the +/// headers must see the longest wait (e.g. an hourly IP cap over a per-minute +/// principal cap), or it retries early and immediately hits the outer limit. fn maybe_tighten_rate_headers( headers: &mut HeaderMap, limit: u64, remaining: u64, verdict: &Verdict, ) { - let existing = headers - .get("ratelimit-remaining") - .and_then(|v| v.to_str().ok()) - .and_then(|v| v.parse::().ok()); - match existing { - Some(inner) if inner <= remaining => {} // inner budget is tighter (or equal): keep it - _ => attach_rate_headers(headers, limit, remaining, verdict), + let header_u64 = |name: &str| { + headers + .get(name) + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()) + }; + let keep_inner = match header_u64("ratelimit-remaining") { + Some(inner) if inner < remaining => true, + // Tie on remaining: keep the inner headers only if their reset is at + // least as long as this phase's, so the longer-binding budget survives. + Some(inner) if inner == remaining => { + header_u64("ratelimit-reset").unwrap_or(0) >= secs_ceil(verdict.reset_after) + } + _ => false, + }; + if !keep_inner { + attach_rate_headers(headers, limit, remaining, verdict); } } From a67797e779c47f874ce86a99ecd952e7c95924e1 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Sat, 18 Jul 2026 00:19:35 +0300 Subject: [PATCH 40/55] fix(shield): throttle high-water GCRA sweeps to bound CPU The high-water bypass ran a full DashMap retain whenever the map exceeded the mark, but active TATs cannot be reclaimed, so an all-active key flood kept the map over the mark and triggered an O(n) scan on every request (a CPU DoS). High-water sweeps are now floored to one per second, bounding the scan cost while still reclaiming drained keys sooner than the normal cadence. --- src/shield/store.rs | 32 ++++++++++++++++++++++---------- 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/src/shield/store.rs b/src/shield/store.rs index 64eb464..2ff330b 100644 --- a/src/shield/store.rs +++ b/src/shield/store.rs @@ -19,15 +19,22 @@ use super::gcra::{Gcra, Verdict}; /// Run eviction of drained entries at most once per this interval. const SWEEP_INTERVAL_MS: u64 = 60_000; -/// When the map grows past this, run a drained-key eviction immediately instead -/// of waiting for the timer, bounding peak memory during a burst of distinct -/// keys. This only reclaims *drained* keys (TAT in the past), never active -/// (rate-limited) ones: dropping a limited key would reset its GCRA budget and -/// hand a flooding attacker a fresh burst, so the store deliberately holds -/// active limiter state. An all-active flood is shed by the fleet gate and -/// upstream, not by discarding the very state that enforces the limit. +/// When the map grows past this, run a drained-key eviction sooner than the +/// normal timer (throttled to [`HW_SWEEP_MIN_MS`]), bounding peak memory during +/// a burst of distinct keys. This only reclaims *drained* keys (TAT in the +/// past), never active (rate-limited) ones: dropping a limited key would reset +/// its GCRA budget and hand a flooding attacker a fresh burst, so the store +/// deliberately holds active limiter state. An all-active flood is shed by the +/// fleet gate and upstream, not by discarding the very state that enforces the +/// limit. const SWEEP_HIGH_WATER: usize = 200_000; +/// Minimum spacing between high-water sweeps. The high-water path can't reclaim +/// active TATs, so under an all-active flood the map stays over the mark and, +/// without this floor, every request would run a full O(n) `retain` (a CPU +/// DoS). Throttling to one scan per second bounds that cost. +const HW_SWEEP_MIN_MS: u64 = 1_000; + /// In-process per-instance GCRA store. // no-std: caller-provided Clock + spin/hashbrown map. #[derive(Debug)] @@ -98,9 +105,14 @@ impl GcraStore { fn maybe_sweep(&self, now: Duration) { let now_ms = u64::try_from(now.as_millis()).unwrap_or(u64::MAX); let last = self.last_sweep_ms.load(Ordering::Relaxed); - // Time-based cadence, unless the map is already over the high-water mark, - // in which case sweep now regardless of when the last one ran. - if now_ms.saturating_sub(last) < SWEEP_INTERVAL_MS && self.tats.len() <= SWEEP_HIGH_WATER { + let elapsed = now_ms.saturating_sub(last); + // Normal cadence is SWEEP_INTERVAL_MS. Over the high-water mark we sweep + // sooner, but still throttled to HW_SWEEP_MIN_MS so an all-active flood + // (whose entries can't be reclaimed) can't trigger a full O(n) retain on + // every request. + let due = elapsed >= SWEEP_INTERVAL_MS + || (elapsed >= HW_SWEEP_MIN_MS && self.tats.len() > SWEEP_HIGH_WATER); + if !due { return; } if self From b5bef303066a049c45c9fdc8ecef0ff280da8453 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Sat, 18 Jul 2026 00:21:11 +0300 Subject: [PATCH 41/55] fix(shield): enforce limit-service cache cap at insertion Capacity was only enforced from the request path before a fetch and in the 60s sweep, but completed fetches insert from a background task. A healthy service plus clients rotating distinct identities could therefore push the cache far past MAX_CACHE_ENTRIES between sweeps, and a burst that stops leaves no later request to trim it. Inserts now refuse to create a new entry once at capacity (existing keys still refresh), so peak size is bounded at write time; the sweep trim stays as a backstop for the rare concurrent-insert overshoot. --- src/shield/resolve.rs | 32 ++++++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/src/shield/resolve.rs b/src/shield/resolve.rs index 92a852a..187f717 100644 --- a/src/shield/resolve.rs +++ b/src/shield/resolve.rs @@ -228,8 +228,10 @@ impl LimitService { /// Drop entries not accessed within `evict_after` (idle keys), keyed on last /// access rather than last fetch so an active-but-unrefreshable key survives. - /// Then enforce the hard capacity, evicting the least-recently-accessed - /// entries so rapid key rotation cannot grow the cache without bound. + /// Then enforce the hard capacity as a backstop, evicting the + /// least-recently-accessed entries. Growth is primarily bounded at insertion + /// (see [`insert_capped`](Self::insert_capped)); this trim only reclaims the + /// rare overshoot from concurrent inserts racing the soft cap check. fn sweep(&self) { let evict_after = self.evict_after; self.cache @@ -277,6 +279,28 @@ impl LimitService { /// globally bounded by `fetch_slots`). When no fetch slot is free, skip the /// refresh and serve the existing stale/static result rather than piling up /// outbound calls under a burst of distinct keys. + /// Insert or refresh a cache entry, refusing to create a *new* entry once the + /// cache is at capacity. A completed fetch inserts from a background task, so + /// without this a fast service plus rotating identities could push the map + /// far past [`MAX_CACHE_ENTRIES`] between the 60s sweeps; refusing new entries + /// at the cap degrades those keys to the static profile (the same tolerance as + /// a skipped refresh). An existing key is always updated. + fn insert_capped(&self, key: String, cached: Cached) { + // Read len() before taking the entry: DashMap::len() read-locks every + // shard, and holding a shard write-lock (via entry) while doing so would + // deadlock. The cap is a soft memory guard, so the race is harmless. + let over_cap = self.cache.len() >= MAX_CACHE_ENTRIES; + match self.cache.entry(key) { + Entry::Occupied(mut o) => { + o.insert(cached); + } + Entry::Vacant(_) if over_cap => {} // at capacity: drop the new entry + Entry::Vacant(v) => { + v.insert(cached); + } + } + } + fn trigger_refresh(self: &Arc, key: String) { match self.inflight.entry(key.clone()) { Entry::Occupied(_) => return, @@ -297,7 +321,7 @@ impl LimitService { match this.fetch(&key).await { Ok(resolved) => { let now = Instant::now(); - this.cache.insert( + this.insert_capped( key.clone(), Cached { profile: resolved, @@ -318,7 +342,7 @@ impl LimitService { match this.cache.get_mut(&key) { Some(mut c) => c.at = now, None => { - this.cache.insert( + this.insert_capped( key.clone(), Cached { profile: None, From 041a87df36a239ae8463e39e2bd17d32e5084715 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 02:16:00 +0300 Subject: [PATCH 42/55] refactor(shield): extract reconciled_headers helper Pull the local-vs-fleet remaining reconciliation out of enforce() into a pure helper returning the reported remaining plus the verdict that drives the response headers. Behaviour is unchanged; this creates a unit-testable seam for the header values. --- src/shield/mod.rs | 34 ++++++++++++++++++++++++---------- 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/src/shield/mod.rs b/src/shield/mod.rs index d38b6cb..af9915f 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -268,21 +268,19 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> } // Report the tighter of the local and (when reconciled) fleet budgets, so a - // client near the fleet cap isn't told it has ample local room. This admit - // consumes one, hence the `- 1`. + // client near the fleet cap isn't told it has ample local room. #[cfg(not(feature = "redis"))] - let reported = verdict.remaining; + let (reported, header_verdict) = reconciled_headers(verdict, None); #[cfg(feature = "redis")] - let reported = { - let mut r = verdict.remaining; - if let Some(fr) = fleet_remaining { - r = r.min(fr.saturating_sub(1)); - // Record the admit for the next reconciliation push. + let (reported, header_verdict) = { + // Record the admit for the next reconciliation push (whenever the fleet + // gate is active, i.e. `fleet_remaining` was computed). + if fleet_remaining.is_some() { if let Some(global) = &shield.global { global.record(&key.store, profile.window); } } - r + reconciled_headers(verdict, fleet_remaining) }; let mut response = next.run(request).await; @@ -291,10 +289,26 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> // set headers on the way out; overwrite them only when this (outer) phase's // remaining is smaller, so the client always sees the budget that will bite // first. An inner rejection carries remaining 0, so it is never overwritten. - maybe_tighten_rate_headers(response.headers_mut(), profile.limit, reported, &verdict); + maybe_tighten_rate_headers( + response.headers_mut(), + profile.limit, + reported, + &header_verdict, + ); response } +/// Combine the local GCRA `verdict` with the optional fleet remaining into the +/// reported remaining and the verdict whose reset drives the headers. The count +/// is the tighter of local and fleet, minus this admit (`fr - 1`). +fn reconciled_headers(verdict: Verdict, fleet_remaining: Option) -> (u64, Verdict) { + let mut reported = verdict.remaining; + if let Some(fr) = fleet_remaining { + reported = reported.min(fr.saturating_sub(1)); + } + (reported, verdict) +} + /// Set the `RateLimit-*` headers unless an inner limiter already advertised a /// budget that binds at least as hard, which must reach the client intact. /// "Binds harder" is a smaller `remaining`, and on a `remaining` tie the larger From 7aefc04ab9a7afa4e86c7195a3bb794b20542698 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 02:18:24 +0300 Subject: [PATCH 43/55] test(shield): add regression test for fleet-bound RateLimit-Reset When the fleet budget binds an allowed response, the reset must reflect the fleet's sliding-window decay, not the short local GCRA reset, or a paced client retries early into the fleet gate. This threads the window into reconciled_headers (still unused) and adds a test that fails on current code (reset stays 100ms) and a companion that pins the local-binds case. The fix lands next. --- src/shield/mod.rs | 11 ++++++++--- src/shield/tests.rs | 42 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 3 deletions(-) diff --git a/src/shield/mod.rs b/src/shield/mod.rs index af9915f..ee43f76 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -270,7 +270,7 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> // Report the tighter of the local and (when reconciled) fleet budgets, so a // client near the fleet cap isn't told it has ample local room. #[cfg(not(feature = "redis"))] - let (reported, header_verdict) = reconciled_headers(verdict, None); + let (reported, header_verdict) = reconciled_headers(verdict, None, profile.window); #[cfg(feature = "redis")] let (reported, header_verdict) = { // Record the admit for the next reconciliation push (whenever the fleet @@ -280,7 +280,7 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> global.record(&key.store, profile.window); } } - reconciled_headers(verdict, fleet_remaining) + reconciled_headers(verdict, fleet_remaining, profile.window) }; let mut response = next.run(request).await; @@ -301,7 +301,12 @@ async fn enforce(shield: &Shield, phase: Phase, request: Request, next: Next) -> /// Combine the local GCRA `verdict` with the optional fleet remaining into the /// reported remaining and the verdict whose reset drives the headers. The count /// is the tighter of local and fleet, minus this admit (`fr - 1`). -fn reconciled_headers(verdict: Verdict, fleet_remaining: Option) -> (u64, Verdict) { +fn reconciled_headers( + verdict: Verdict, + fleet_remaining: Option, + window: Duration, +) -> (u64, Verdict) { + let _ = window; // used once the fleet-bound reset lands let mut reported = verdict.remaining; if let Some(fr) = fleet_remaining { reported = reported.min(fr.saturating_sub(1)); diff --git a/src/shield/tests.rs b/src/shield/tests.rs index 51cf0e3..a307e89 100644 --- a/src/shield/tests.rs +++ b/src/shield/tests.rs @@ -237,6 +237,48 @@ fn secs_ceil_rounds_sub_second_waits_up() { assert_eq!(secs_ceil(Duration::ZERO), 0); } +#[test] +fn reconciled_headers_widen_reset_when_fleet_binds() { + use crate::shield::gcra::Verdict; + use std::time::Duration; + // Fresh high-rate local key: ample local room, tiny local reset (100ms). + let verdict = Verdict { + allowed: true, + new_tat: Duration::from_millis(100), + remaining: 599, + retry_after: Duration::ZERO, + reset_after: Duration::from_millis(100), + }; + let window = Duration::from_secs(60); + // Fleet has one slot left → it binds; reported collapses to 0. The reset must + // reflect the fleet's sliding-window decay (window/10 = 6s), not the 100ms + // local reset, or a paced client retries early into the fleet gate. + let (reported, hv) = reconciled_headers(verdict, Some(1), window); + assert_eq!(reported, 0, "fleet budget must bind the reported remaining"); + assert!( + hv.reset_after >= Duration::from_secs(6), + "expected fleet-derived reset, got {:?}", + hv.reset_after + ); +} + +#[test] +fn reconciled_headers_keep_local_reset_when_local_binds() { + use crate::shield::gcra::Verdict; + use std::time::Duration; + let verdict = Verdict { + allowed: true, + new_tat: Duration::from_secs(30), + remaining: 0, + retry_after: Duration::ZERO, + reset_after: Duration::from_secs(30), + }; + // Fleet has ample room → local binds; the local 30s reset is left intact. + let (reported, hv) = reconciled_headers(verdict, Some(100), Duration::from_secs(60)); + assert_eq!(reported, 0); + assert_eq!(hv.reset_after, Duration::from_secs(30)); +} + #[test] fn build_rejects_unknown_default_profile() { let mut cfg = config(vec![("t", "1/min", None)], Vec::new()); From f6da6a5279568bb24df02961eae2708e4d1bad69 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 02:20:58 +0300 Subject: [PATCH 44/55] fix(shield): report fleet-derived reset when the fleet budget binds An allowed response whose remaining was reduced to the fleet value still carried the local GCRA reset, which for a high-rate key can be a fraction of a second even while the shared sliding-window estimate stays saturated for much longer. A client pacing off that reset retried early and immediately hit the fleet gate. reconciled_headers now widens the reset (and retry-after) to the fleet backoff whenever the fleet budget binds, never shrinking a longer local reset. The backoff is factored into fleet_backoff, shared with global_reject. Covered by the test in 7aefc04. --- src/shield/mod.rs | 30 ++++++++++++++++++++++++------ 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/src/shield/mod.rs b/src/shield/mod.rs index ee43f76..b248016 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -306,12 +306,32 @@ fn reconciled_headers( fleet_remaining: Option, window: Duration, ) -> (u64, Verdict) { - let _ = window; // used once the fleet-bound reset lands let mut reported = verdict.remaining; + let mut hv = verdict; if let Some(fr) = fleet_remaining { - reported = reported.min(fr.saturating_sub(1)); + let fleet_r = fr.saturating_sub(1); + if fleet_r <= reported { + // The fleet budget binds. Advertise the fleet-derived reset (the + // shared sliding-window estimate can stay saturated far longer than + // this instance's local GCRA), so a client pacing off the allowed + // response does not retry before the window decays and immediately + // hit the fleet gate. Widen, never shrink: keep the local reset if it + // is already the longer wait. + reported = fleet_r; + let backoff = fleet_backoff(window); + hv.reset_after = hv.reset_after.max(backoff); + hv.retry_after = hv.retry_after.max(backoff); + } } - (reported, verdict) + (reported, hv) +} + +/// Poll interval for a client blocked (or nearly blocked) by the fleet gate: a +/// tenth of the window, at least 1s. The fleet's sliding-window estimate decays +/// continuously rather than freeing at an epoch boundary, so this is a retry +/// cadence, not a wait-to-boundary. +fn fleet_backoff(window: Duration) -> Duration { + (window / 10).max(Duration::from_secs(1)) } /// Set the `RateLimit-*` headers unless an inner limiter already advertised a @@ -352,9 +372,7 @@ fn maybe_tighten_rate_headers( /// the boundary, which a client could wait out and still be rejected. #[cfg(feature = "redis")] fn global_reject(limit: u64, window: Duration) -> Response { - // A tenth of the window (at least 1s): long enough not to hammer, short - // enough to retry as the window decays or other instances free budget. - let backoff = (window / 10).max(Duration::from_secs(1)); + let backoff = fleet_backoff(window); let verdict = Verdict { allowed: false, new_tat: Duration::ZERO, From 021352f55425bff4fa787fa9ba1c29b242a55325 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 02:21:29 +0300 Subject: [PATCH 45/55] docs(shield): note reconciliation needs the redis feature and a sync block The feature summary implied enabling the `redis` feature alone turns on fleet reconciliation, but the deployment section requires a configured `sync` block too. Clarify the bullet so users don't build with the feature and expect fleet-wide limiting without configuring it. --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index c7fe395..110eda5 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ Works with **any** gRPC service via proto descriptor files. No code generation, - **Health endpoints** `/health/live`, `/health/ready` (upstream gRPC health probe), `/health/startup` - **Prometheus metrics** at `/metrics` - **CORS** with a configurable origin allow-list -- **Rate limiting (Shield)**: local GCRA shaper (no blocking latency) keyed by client IP, header, or validated JWT claim; named limit tiers as config data; optional async cross-instance reconciliation (feature `redis`) for an approximate fleet-wide limit +- **Rate limiting (Shield)**: local GCRA shaper (no blocking latency) keyed by client IP, header, or validated JWT claim; named limit tiers as config data; optional async cross-instance reconciliation for an approximate fleet-wide limit (requires both the `redis` feature and a configured `sync` block) - **JWT auth**: validate `Bearer` tokens via an Ed25519 PEM key or JWKS auto-discovery, enforce per-route `require_auth` / `required_roles`, and forward claims as headers - **OIDC discovery**: serve `/.well-known/openid-configuration` and a JWKS endpoint (Ed25519) built from config, to front an identity provider - **Forward-auth**: a verification endpoint (`/auth/verify`) for a fronting proxy (nginx `auth_request`, Traefik `forwardAuth`) to delegate auth, returning the verified identity as headers From 70d067fac716db9a892042022f5a3030803e9337 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 02:21:38 +0300 Subject: [PATCH 46/55] build(deps): bump EmbarkStudios/cargo-deny-action from 2.0.20 to 2.1.1 Update the SHA-pinned advisory-scan action to v2.1.1. --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e040f3d..9596d4f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -83,6 +83,6 @@ jobs: # Scans Cargo.lock for RustSec advisories; honors the ignore list in # deny.toml. Third-party action: pinned to a commit SHA (dependabot # keeps it fresh). - - uses: EmbarkStudios/cargo-deny-action@bb137d7af7e4fb67e5f82a49c4fce4fad40782fe # v2.0.20 + - uses: EmbarkStudios/cargo-deny-action@3c6349835b2b7b196a839186cb8b78e02f7b5f25 # v2.1.1 with: command: check advisories From 13acd1a73b59016df3a5ed0dfb257c46be7e1bae Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 10:24:16 +0300 Subject: [PATCH 47/55] test(shield): add regression test for Retry-After on tie-break overwrite When the outer phase's longer reset wins a RateLimit-Remaining tie on an error response, Retry-After must be widened to match, or a client retries after the shorter overwritten wait and immediately hits the binding outer gate. This test drives maybe_tighten_rate_headers with an inner 429's short headers and a longer outer verdict; it fails on current code (Retry-After stays 60) until the fix. --- src/shield/tests.rs | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/src/shield/tests.rs b/src/shield/tests.rs index a307e89..a8caa71 100644 --- a/src/shield/tests.rs +++ b/src/shield/tests.rs @@ -279,6 +279,31 @@ fn reconciled_headers_keep_local_reset_when_local_binds() { assert_eq!(hv.reset_after, Duration::from_secs(30)); } +#[test] +fn tighten_headers_widen_retry_after_when_outer_reset_wins() { + use crate::shield::gcra::Verdict; + use axum::http::HeaderMap; + use std::time::Duration; + // An inner 429 left exhausted headers with a short reset and Retry-After. + let mut headers = HeaderMap::new(); + headers.insert("ratelimit-remaining", "0".parse().unwrap()); + headers.insert("ratelimit-reset", "60".parse().unwrap()); + headers.insert("retry-after", "60".parse().unwrap()); + // The outer phase is also exhausted but binds far longer (1 hour). + let outer = Verdict { + allowed: false, + new_tat: Duration::ZERO, + remaining: 0, + retry_after: Duration::from_secs(3600), + reset_after: Duration::from_secs(3600), + }; + maybe_tighten_rate_headers(&mut headers, 1, 0, &outer); + // The longer outer reset wins; Retry-After must be widened to match it, or + // the client retries after 60s and immediately hits the outer gate. + assert_eq!(headers.get("ratelimit-reset").unwrap(), "3600"); + assert_eq!(headers.get("retry-after").unwrap(), "3600"); +} + #[test] fn build_rejects_unknown_default_profile() { let mut cfg = config(vec![("t", "1/min", None)], Vec::new()); From 068517bdbff2e4c2a287598571e93fd29d7dbea0 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 10:24:56 +0300 Subject: [PATCH 48/55] fix(shield): widen Retry-After when the outer budget wins a header tie maybe_tighten_rate_headers overwrote RateLimit-* with the longer-binding outer budget on a remaining tie but left the inner layer's Retry-After untouched, so a 429 could advertise the outer's hour-long reset yet tell the client to retry after the inner's minute. It now widens Retry-After to the winning reset on an error response (never shrinking, and never adding one to a 200). Covered by the test in 13acd1a. --- src/shield/mod.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/src/shield/mod.rs b/src/shield/mod.rs index b248016..3ff33bf 100644 --- a/src/shield/mod.rs +++ b/src/shield/mod.rs @@ -363,6 +363,24 @@ fn maybe_tighten_rate_headers( }; if !keep_inner { attach_rate_headers(headers, limit, remaining, verdict); + // On an error response the rejecting layer already set `Retry-After`. We + // just replaced the budget with a longer-binding one, so widen + // `Retry-After` to that reset too; otherwise the client retries after the + // overwritten (shorter) wait and immediately hits this binding budget. + // Only ever widen: a 200 has no `Retry-After` to touch, and an already + // longer wait is left intact. + if let Some(current) = headers + .get("retry-after") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()) + { + let reset = secs_ceil(verdict.reset_after); + if reset > current { + if let Ok(v) = reset.to_string().parse() { + headers.insert("retry-after", v); + } + } + } } } From 5c48ddd6f01cf6bca67df499293518282a186b2e Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 13:22:15 +0300 Subject: [PATCH 49/55] test(shield): add regression test for relative rule patterns A rule pattern without a leading slash (e.g. `api/**`) can never match the request path, which always starts with `/`, so the route runs unmetered while Shield is enabled. This test expects compile_rules to reject such a pattern; it fails on current code until the validation lands. --- src/shield/matcher.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/shield/matcher.rs b/src/shield/matcher.rs index 71ff787..46d5dd3 100644 --- a/src/shield/matcher.rs +++ b/src/shield/matcher.rs @@ -314,4 +314,20 @@ mod tests { ); assert!(err.is_err()); } + + #[test] + fn relative_pattern_is_rejected() { + // A pattern without a leading slash (a missing-slash typo like `api/**`) + // can never match `request.uri().path()`, which always starts with `/`, + // so the route would silently run unmetered. Reject it at build time. + let err = compile_rules( + &[RateRuleConfig { + pattern: "api/**".to_string(), + key: KeySourceConfig::Ip, + profile: Some("auth".to_string()), + }], + &profiles(), + ); + assert!(err.is_err()); + } } From 0652d279580923d2cf3974ab1781231965aa576e Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 13:22:41 +0300 Subject: [PATCH 50/55] fix(shield): reject rule patterns without a leading slash compile_rules now fails a rule whose pattern does not start with '/'. Patterns are matched against the request path, which always has a leading slash, so a relative pattern (a missing-slash typo) compiled a matcher that never fired and left the route unmetered. Covered by the test in 5c48ddd. --- src/shield/matcher.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/shield/matcher.rs b/src/shield/matcher.rs index 46d5dd3..d75effb 100644 --- a/src/shield/matcher.rs +++ b/src/shield/matcher.rs @@ -123,6 +123,16 @@ pub fn compile_rules( configs .iter() .map(|c| { + // Patterns are matched against the request path, which always starts + // with `/`. Reject a relative pattern (a missing-slash typo) rather + // than compiling a matcher that silently never fires, leaving the + // route unmetered. + if !c.pattern.starts_with('/') { + return Err(format!( + "rule pattern {:?} must start with '/' (matched against the request path)", + c.pattern + )); + } if let Some(name) = &c.profile { if !profiles.contains_key(name) { return Err(format!( From 50b23ab6437a90bcee5bd9f2cebd75fd5fa05c6f Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 13:23:18 +0300 Subject: [PATCH 51/55] test(shield): add regression test for unknown fields in rule keys The internally-tagged KeySourceConfig enum did not reject unknown fields, so a stray `name` on an `ip` key was silently ignored, keeping a rule IP-keyed instead of the intended per-header limit. This test expects such a config to fail deserialization; it fails on current code until the fix lands. --- src/config.rs | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/config.rs b/src/config.rs index 37a04f3..a6dbc8d 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1103,6 +1103,27 @@ shield: assert!(err.is_err(), "unknown shield field must be rejected"); } + #[test] + fn shield_rejects_unknown_field_in_rule_key() { + // A stray field inside a rule key (here `name` on an `ip` key, a copy-edit + // leftover) must be a hard error. Silently ignoring it would keep the rule + // IP-keyed instead of the intended per-header limit, weakening the control. + let yaml = r#" +upstream: + default: "grpc://localhost:4180" +shield: + enabled: true + profiles: + auth: { rate: "20/min", burst: 5 } + rules: + - pattern: "/v1/**" + key: { type: ip, name: x-api-key } + profile: "auth" +"#; + let err = serde_yaml::from_str::(yaml); + assert!(err.is_err(), "unknown field in a rule key must be rejected"); + } + #[test] fn test_openapi_config_deserialize() { let yaml = r#" From 8edfa5bbf7eaa96f25e7429d10833b2991866be6 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 13:25:34 +0300 Subject: [PATCH 52/55] fix(shield): reject unknown and mismatched fields in rule keys serde does not honour deny_unknown_fields on internally-tagged enums, so a stray field in a rule key (e.g. `name` on an `ip` key) was silently ignored, leaving the rule IP-keyed instead of the intended per-header limit. KeySourceConfig now deserializes through a flat deny_unknown_fields form plus per-variant validation that rejects a field belonging to a different key type and a genuinely unknown field alike. The wire schema is unchanged. Covered by the test in 50b23ab. --- src/config.rs | 60 +++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 58 insertions(+), 2 deletions(-) diff --git a/src/config.rs b/src/config.rs index a6dbc8d..acd3d05 100644 --- a/src/config.rs +++ b/src/config.rs @@ -454,8 +454,7 @@ pub struct RateRuleConfig { /// their value is absent, so a limit can't be bypassed by omitting a header or /// authenticating anonymously. Written as a tagged map, e.g. /// `key: { type: jwt_claim, claim: sub }`; omitting `key` defaults to `ip`. -#[derive(Debug, Clone, Deserialize, Default, PartialEq, Eq)] -#[serde(tag = "type", rename_all = "snake_case")] +#[derive(Debug, Clone, Default, PartialEq, Eq)] #[non_exhaustive] pub enum KeySourceConfig { /// Client IP (trusted-proxy `X-Forwarded-For` aware). `{ type: ip }`. @@ -475,6 +474,63 @@ pub enum KeySourceConfig { }, } +/// Flat wire form of a rule key. `deny_unknown_fields` rejects any field outside +/// this set, and the manual [`KeySourceConfig`] deserializer additionally rejects +/// a field that belongs to a *different* variant (e.g. `name` on an `ip` key), so +/// a copy-edit leftover can't silently downgrade the intended key source. serde +/// does not honour `deny_unknown_fields` on internally-tagged enums directly, +/// hence this intermediate. +#[derive(Deserialize)] +#[serde(deny_unknown_fields, rename_all = "snake_case")] +struct KeySourceRaw { + #[serde(rename = "type")] + kind: String, + #[serde(default)] + name: Option, + #[serde(default)] + claim: Option, +} + +impl<'de> Deserialize<'de> for KeySourceConfig { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + use serde::de::Error; + let raw = KeySourceRaw::deserialize(deserializer)?; + match raw.kind.as_str() { + "ip" => { + if raw.name.is_some() || raw.claim.is_some() { + return Err(D::Error::custom("key type 'ip' takes no other fields")); + } + Ok(Self::Ip) + } + "header" => { + if raw.claim.is_some() { + return Err(D::Error::custom( + "key type 'header' takes 'name', not 'claim'", + )); + } + let name = raw.name.ok_or_else(|| D::Error::missing_field("name"))?; + Ok(Self::Header { name }) + } + "jwt_claim" => { + if raw.name.is_some() { + return Err(D::Error::custom( + "key type 'jwt_claim' takes 'claim', not 'name'", + )); + } + let claim = raw.claim.ok_or_else(|| D::Error::missing_field("claim"))?; + Ok(Self::JwtClaim { claim }) + } + other => Err(D::Error::unknown_variant( + other, + &["ip", "header", "jwt_claim"], + )), + } + } +} + /// Claims that carry a key's limit inside the JWT itself. A tier-name claim maps /// to a `profiles` entry (numbers stay tunable in config); direct numeric claims /// set the limit explicitly. From f088d18d635ddc5de551af7f993bf6736f0895fc Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 13:27:34 +0300 Subject: [PATCH 53/55] refactor(shield): extract claim_plans from reconcile_at Move the phase-1 per-key claim-and-plan loop into a sync helper returning the push plans. Behaviour is unchanged (a plan per tracked key); this creates a seam to test which keys a reconcile tick refreshes. --- src/shield/global.rs | 81 +++++++++++++++++++++++--------------------- 1 file changed, 43 insertions(+), 38 deletions(-) diff --git a/src/shield/global.rs b/src/shield/global.rs index 297ccde..11ca15a 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -268,46 +268,10 @@ impl GlobalCounters { async fn reconcile_at(&self, now: Duration) { self.evict_stale(); - // Phase 1 (locked per key, brief): CLAIM each key's unclaimed admits by - // zeroing them now, so a concurrent epoch roll or a push failure can't - // double-publish or drop them. Emit a plan for every active key (delta - // may be 0) so its estimate is refreshed even when the key is only being - // rejected on a stale remote estimate. The delta is pushed to the epoch - // it was accumulated in (`push_epoch`); the estimate reads the current - // epoch (`read_epoch`). - let keys: Vec = self.states.iter().map(|e| e.key().clone()).collect(); - if keys.is_empty() { + let plans = self.claim_plans(now); + if plans.is_empty() { return; } - let mut plans: Vec = Vec::new(); - for key in keys { - if let Some(mut s) = self.states.get_mut(&key) { - let window = Duration::from_secs(s.window_secs); - let read_epoch = window::epoch(now, window); - let claim = s.local_count; - s.local_count = 0; - plans.push(PushPlan { - key: key.clone(), - push_epoch: s.epoch, - read_epoch, - window, - delta: claim, - is_carryover: false, - }); - if s.carryover > 0 { - let carry = s.carryover; - s.carryover = 0; - plans.push(PushPlan { - key: key.clone(), - push_epoch: s.carryover_epoch, - read_epoch, - window, - delta: carry, - is_carryover: true, - }); - } - } - } // Claimed deltas are fire-and-forget: once claimed (zeroed), a push that // fails (connection or transaction) simply drops them rather than @@ -347,6 +311,47 @@ impl GlobalCounters { self.apply_estimates(now, &plans, &reads); } + /// Phase 1 (locked per key, brief): CLAIM each key's unclaimed admits by + /// zeroing them now, so a concurrent epoch roll or a push failure can't + /// double-publish or drop them. Emit a plan for every tracked key (delta may + /// be 0) so its estimate is refreshed even when the key is only being rejected + /// on a stale remote estimate. The delta is pushed to the epoch it was + /// accumulated in (`push_epoch`); the estimate reads the current epoch + /// (`read_epoch`). + fn claim_plans(&self, now: Duration) -> Vec { + let keys: Vec = self.states.iter().map(|e| e.key().clone()).collect(); + let mut plans: Vec = Vec::new(); + for key in keys { + if let Some(mut s) = self.states.get_mut(&key) { + let window = Duration::from_secs(s.window_secs); + let read_epoch = window::epoch(now, window); + let claim = s.local_count; + s.local_count = 0; + plans.push(PushPlan { + key: key.clone(), + push_epoch: s.epoch, + read_epoch, + window, + delta: claim, + is_carryover: false, + }); + if s.carryover > 0 { + let carry = s.carryover; + s.carryover = 0; + plans.push(PushPlan { + key: key.clone(), + push_epoch: s.carryover_epoch, + read_epoch, + window, + delta: carry, + is_carryover: true, + }); + } + } + } + plans + } + /// Refresh each key's cached estimate of the fleet's consumption. fn apply_estimates(&self, now: Duration, plans: &[PushPlan], reads: &[(u64, u64)]) { for (p, (cur, prev)) in plans.iter().zip(reads) { From f248982d8905b09dd96f3a1f9c8dee5458127341 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 13:28:37 +0300 Subject: [PATCH 54/55] test(shield): add regression test for idle fleet-key refresh A reconcile tick planned a store read for every tracked key, so a one-shot burst of distinct keys kept issuing MGETs for idle principals until KEY_TTL expiry, turning a brief cardinality spike into sustained load. This test expects claim_plans to skip an idle zero-delta key; it fails on current code until the activity gate lands. --- src/shield/global.rs | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/src/shield/global.rs b/src/shield/global.rs index 11ca15a..1018791 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -498,6 +498,39 @@ mod tests { assert_eq!(g.fleet_remaining("k", 3, W), 0); } + #[test] + fn idle_zero_delta_keys_are_not_refreshed_each_tick() { + use std::collections::HashSet; + let g = counters(); + let now = unix_now(); + // Three keys become tracked: one with a delta, one only gated (read), one + // that goes idle after the first tick. + g.record("delta", W); + let _ = g.fleet_remaining("gated", 100, W); + g.record("idle", W); + // First tick claims everything and clears the per-tick activity. + let _ = g.claim_plans(now); + // Between ticks, only `delta` and `gated` see activity; `idle` is untouched. + g.record("delta", W); + let _ = g.fleet_remaining("gated", 100, W); + let plans = g.claim_plans(now); + let keys: HashSet<&str> = plans.iter().map(|p| p.key.as_str()).collect(); + assert!( + keys.contains("delta"), + "a key with a pending delta must be pushed" + ); + assert!( + keys.contains("gated"), + "an actively-gated key must be refreshed" + ); + // An idle key with no delta must not be MGET'd every tick, or a one-shot + // key-cardinality burst becomes sustained load for the whole KEY_TTL. + assert!( + !keys.contains("idle"), + "an idle zero-delta key must not be refreshed every tick" + ); + } + #[test] fn roll_preserves_unclaimed_deltas_as_carryover() { // Admits accumulated in an epoch that rolls before a reconcile claims From fdf7b5283b8a150def7fcccbcd789a7ad0759130 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 22 Jul 2026 13:29:35 +0300 Subject: [PATCH 55/55] fix(shield): skip idle zero-delta keys in fleet reconciliation A reconcile tick planned a store read for every tracked key, so a one-shot burst of distinct keys kept issuing per-key MGETs for idle principals until KEY_TTL (up to 10 min at the default interval), turning a brief cardinality spike into sustained CPU/store load. KeyState now carries a per-tick activity flag set when the key is recorded or gated and cleared each pass; claim_plans skips a key that is idle and has no pending delta/carryover. A later request re-marks it and the next tick refreshes its estimate. Covered by the test in f248982. --- src/shield/global.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/src/shield/global.rs b/src/shield/global.rs index 1018791..cd4252b 100644 --- a/src/shield/global.rs +++ b/src/shield/global.rs @@ -72,6 +72,12 @@ struct KeyState { /// trusting it and degrades to per-instance limiting. estimate_at: Instant, last_seen: Instant, + /// Set when the key is touched (recorded or gated) and cleared by each + /// reconcile pass. Gates whether an idle, zero-delta key still needs its + /// estimate refreshed: an untouched key's estimate isn't being read, so + /// re-reading it every tick is pure load. A later request marks it again and + /// the next tick refreshes. + seen_since_tick: bool, } impl KeyState { @@ -86,6 +92,7 @@ impl KeyState { remote_estimate: 0, estimate_at: now, last_seen: now, + seen_since_tick: true, } } @@ -223,6 +230,7 @@ impl GlobalCounters { } state.roll_to(ep); state.last_seen = Instant::now(); + state.seen_since_tick = true; Some(state) } @@ -323,6 +331,16 @@ impl GlobalCounters { let mut plans: Vec = Vec::new(); for key in keys { if let Some(mut s) = self.states.get_mut(&key) { + // Skip a key that is idle (untouched since the last tick) and has + // nothing pending: its estimate isn't being read, so refreshing it + // is wasted store load. Always clear the flag so the next tick sees + // only keys touched since. A key with a delta/carryover is active + // by definition and always processed. + let active = s.local_count > 0 || s.carryover > 0 || s.seen_since_tick; + s.seen_since_tick = false; + if !active { + continue; + } let window = Duration::from_secs(s.window_secs); let read_epoch = window::epoch(now, window); let claim = s.local_count;