diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 12a993d..9596d4f 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 @@ -69,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 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/README.md b/README.md index d0bd61e..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)**: 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 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 @@ -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,73 @@ 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. 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. +JWT-based resolution only applies to `jwt_claim` rules, since only they run with +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/) +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.** + +- **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`. 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. + ## 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..acd3d05 100644 --- a/src/config.rs +++ b/src/config.rs @@ -369,58 +369,239 @@ 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)] +#[serde(deny_unknown_fields)] #[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)] +#[serde(deny_unknown_fields)] +#[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 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 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, +} + +/// 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, Default, PartialEq, Eq)] +#[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, + }, } -/// Per-identifier rate limiting config. +/// 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. #[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] #[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)] +#[serde(deny_unknown_fields)] +#[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)] +#[serde(deny_unknown_fields)] +#[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 +1017,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 @@ -953,6 +1137,49 @@ 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 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#" diff --git a/src/lib.rs b/src/lib.rs index a9ef65a..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 @@ -524,8 +526,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 +560,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, )); } @@ -564,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) @@ -634,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(), ]) } } diff --git a/src/shield/gcra.rs b/src/shield/gcra.rs new file mode 100644 index 0000000..5c68c69 --- /dev/null +++ b/src/shield/gcra.rs @@ -0,0 +1,263 @@ +//! 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, + /// 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 +/// 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, + burst, + } + } + + /// 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. + // 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)).min(self.burst) + } 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)); + } + + /// 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] + 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..cd4252b --- /dev/null +++ b/src/shield/global.rs @@ -0,0 +1,785 @@ +//! 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 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}; + +use dashmap::mapref::entry::Entry; + +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); + +/// 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 +/// 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, + /// Admits in the current epoch not yet claimed for a push. + local_count: u64, + /// 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 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, + /// 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 { + fn new(epoch: u64, window_secs: u64) -> Self { + let now = Instant::now(); + Self { + epoch, + window_secs, + local_count: 0, + carryover: 0, + carryover_epoch: 0, + remote_estimate: 0, + estimate_at: now, + last_seen: now, + seen_since_tick: true, + } + } + + /// 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 unclaimed + /// prior epoch exists between reconciles. + fn roll_to(&mut self, epoch: u64) { + if self.epoch != epoch { + if self.local_count > 0 { + 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; + self.local_count = 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(), + })) + } + + /// 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) + /// 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 fleet_remaining(&self, key: &str, budget: u64, window: Duration) -> u64 { + // 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 + // 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 + /// next background tick. + pub fn record(&self, key: &str, window: Duration) { + // 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 + /// 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, + ) -> Option> { + let now = unix_now(); + let window_secs = window.as_secs().max(1); + let ep = window::epoch(now, window); + // 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 + // 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); + } + state.roll_to(ep); + state.last_seen = Instant::now(); + state.seen_since_tick = true; + Some(state) + } + + /// 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 false; + } + let this = self.clone(); + tokio::spawn(async move { + let mut ticker = tokio::time::interval(this.interval); + loop { + ticker.tick().await; + this.reconcile().await; + } + }); + true + } + + /// 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(|| redis::aio::ConnectionManager::new(self.client.clone())) + .await + .cloned() + } + + /// 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(); + + let plans = self.claim_plans(now); + if plans.is_empty() { + return; + } + + // Claimed deltas are fire-and-forget: once claimed (zeroed), a push that + // 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) => { + tracing::warn!("rate-limit shared store unavailable, staying local: {e}"); + return; + } + }; + + // Push each claimed delta as an atomic INCRBY inside MULTI/EXEC (never a + // SET, so concurrent instances' increments accumulate). + if let Err(e) = self.push_deltas(&mut conn, &plans).await { + tracing::warn!("rate-limit delta push failed, dropping this interval: {e}"); + return; + } + + // 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) => { + tracing::warn!("rate-limit aggregate read failed: {e}"); + return; + } + }; + 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) { + // 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; + 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) { + // 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 + // 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); + // 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); + // 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(); + } + } + } + + /// `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 + /// restores the claims and re-pushes the same deltas next tick. + async fn push_deltas( + &self, + conn: &mut redis::aio::ConnectionManager, + 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; + 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(); + } + 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::ConnectionManager, + 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.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 + .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 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| { + 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 + }); + } +} + +/// A per-key push plan captured under lock, used without holding locks. +struct PushPlan { + key: String, + /// 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, + /// True for a plan publishing a prior epoch's carried-over delta (no estimate). + is_carryover: bool, +} + +#[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 budget_admits_until_local_count_reaches_it() { + let g = counters(); + // Budget 3: three admits leave room, the fourth is over budget. + for _ in 0..3 { + assert!(g.fleet_remaining("k", 3, W) > 0); + g.record("k", W); + } + assert_eq!(g.fleet_remaining("k", 3, W), 0); + } + + #[test] + fn budget_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.fleet_remaining("k", 3, W) > 0); + g.record("k", W); + 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 + // 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.roll_to(11); + assert_eq!(s.carryover, 5); + assert_eq!(s.carryover_epoch, 10); + assert_eq!(s.local_count, 0); + 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(); + 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, + is_carryover: false, + }; + 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(); + assert!(g.fleet_remaining("a", 1, W) > 0); + g.record("a", W); + assert_eq!(g.fleet_remaining("a", 1, W), 0); + // A different key is unaffected. + assert!(g.fleet_remaining("b", 1, W) > 0); + } + + /// 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.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_eq!( + b.fleet_remaining(&key, 3, W), + 0, + "B must reject once the combined budget is reached" + ); + } + + /// 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" + ); + } + + /// 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. + #[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" + ); + } +} diff --git a/src/shield/matcher.rs b/src/shield/matcher.rs index 734a491..d75effb 100644 --- a/src/shield/matcher.rs +++ b/src/shield/matcher.rs @@ -1,25 +1,68 @@ -//! 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, + /// 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 @@ -32,35 +75,109 @@ 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)?; + // 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); + 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)?, + // 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!( + "rule {:?} references unknown profile {name:?}", + c.pattern + )); + } + } + let key = match &c.key { + KeySourceConfig::Ip => KeySource::Ip, + 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. 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(hn.as_str().to_string()) + } + KeySourceConfig::JwtClaim { claim } => KeySource::JwtClaim(claim.clone()), + }; + let phase = match &key { + 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() @@ -70,43 +187,157 @@ 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 endpoint_class_glob_respects_segments() { - let classes = compile_endpoint_classes( - &[ec("/api/v1/heavy-*", "heavy", "10/min")], - Duration::from_secs(60), + 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 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(); - 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_eq!(rules[0].fingerprint, rules[1].fingerprint); } #[test] - fn double_star_spans_segments() { - let classes = compile_endpoint_classes( - &[ec("/v1/auth/**", "auth", "20/min")], - Duration::from_secs(60), + 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 { + 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( + &[ + 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(); + 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 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 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()); } #[test] - fn invalid_rate_fails_compilation() { - let err = compile_endpoint_classes(&[ec("/x", "c", "nonsense")], Duration::from_secs(60)); + 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()); } } diff --git a/src/shield/mod.rs b/src/shield/mod.rs index 4c04db6..3ff33bf 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}; +use gcra::Verdict; +use matcher::{CompiledProfile, CompiledRule, KeySource, Phase}; +use store::GcraStore; -/// Maximum request body buffered to read an identifier field (256 KiB). -const MAX_IDENTIFIER_BODY: usize = 256 * 1024; - -/// 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,35 @@ 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"); - return Ok(None); + if config.rules.is_empty() { + // 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 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 +92,354 @@ 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), + )?; + // 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, + }; + #[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`. 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() + .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>, + identity: &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(identity) { + 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(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( + &rule.fingerprint, + &rule.key, + &client, + request.headers(), + claims, + ); - // The decision whose budget we surface to the client via headers. - let mut report = None; + // 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; + }; - // 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); - } - report = Some(decision); + // Fleet gate first (read-only, cached), so the local shaper is not charged + // 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. + // + // 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 + .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); } - // 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 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.remaining, &verdict); + } + + // 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, profile.window); + #[cfg(feature = "redis")] + 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); + } } - // 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 + reconciled_headers(verdict, fleet_remaining, profile.window) }; let mut response = next.run(request).await; - if let Some(decision) = report { - attach_rate_headers(response.headers_mut(), &decision); - } + // 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, + &header_verdict, + ); response } -/// Convenience alias for the axum request type used by the middleware. -type Request = axum::extract::Request; +/// 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, + window: Duration, +) -> (u64, Verdict) { + let mut reported = verdict.remaining; + let mut hv = verdict; + if let Some(fr) = fleet_remaining { + 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, 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 +/// 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 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); + // 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); + } + } + } + } +} + +/// 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 backoff = fleet_backoff(window); + let verdict = Verdict { + allowed: false, + new_tat: Duration::ZERO, + remaining: 0, + retry_after: backoff, + reset_after: backoff, + }; + too_many_requests(limit, 0, &verdict) +} + +/// 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( + fingerprint: &str, + key: &KeySource, + client: &str, + headers: &HeaderMap, + claims: Option<&serde_json::Value>, +) -> RuleKey { + let (tag, identity) = match key { + KeySource::Ip => ("ip", client.to_string()), + KeySource::Header(name) => match header_str(headers, name) { + Some(v) => ("hdr", v), + None => ("ip", client.to_string()), + }, + KeySource::JwtClaim(claim) => match claims.and_then(|c| resolve::claim_str(c, claim)) { + Some(v) => ("jwt", v), + None => ("ip", client.to_string()), + }, + }; + RuleKey { + store: format!("{fingerprint}:{tag}:{}", matcher::short_hash(&identity)), + identity, + } +} /// Parse a trusted-proxy entry as a CIDR range, accepting a bare IP as a /32 /// or /128 host range. @@ -186,9 +460,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, @@ -203,8 +478,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(), } } @@ -229,19 +503,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 @@ -252,33 +513,23 @@ 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. +/// `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); } - 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) = 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, remaining: u64, verdict: &Verdict) -> Response { let mut response = ( StatusCode::TOO_MANY_REQUESTS, Json(serde_json::json!({ @@ -288,251 +539,21 @@ 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, remaining, 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 { + // 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)] -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..187f717 --- /dev/null +++ b/src/shield/resolve.rs @@ -0,0 +1,563 @@ +//! 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. 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: rpm, + window: PER_MINUTE, + burst: burst.max(1), + }); + Some(CompiledProfile { + gcra, + limit: rpm, + 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 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, + /// 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. +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, + 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, + /// Global cap on concurrent background fetches (in addition to per-key dedup). + fetch_slots: Arc, + base: Instant, + last_sweep_ms: std::sync::atomic::AtomicU64, +} + +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> { + // 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()) + .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, + // 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(), + fetch_slots: Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_FETCHES)), + 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() + { + self.sweep(); + } + } + + /// 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 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 + .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 + /// 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(); + // 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) => { + // 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`, 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. + /// 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, + Entry::Vacant(v) => { + 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(); + this.insert_capped( + key.clone(), + Cached { + profile: resolved, + at: now, + last_access: now, + }, + ); + } + Err(()) => { + // 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.insert_capped( + key.clone(), + Cached { + profile: None, + at: now, + last_access: now, + }, + ); + } + } + } + } + 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 { + 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))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn profiles() -> HashMap { + let mut m = HashMap::new(); + m.insert( + "premium".to_string(), + profile_from_numbers(1000, 100).unwrap(), // 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 + ); + } + + #[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).unwrap()), + at: old, + last_access: 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" + ); + } + + 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_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( + &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 + // 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" }); + assert_eq!( + jwt_limits().resolve(&claims, &profiles()).unwrap().limit, + 250 + ); + } +} diff --git a/src/shield/store.rs b/src/shield/store.rs index be72866..2ff330b 100644 --- a/src/shield/store.rs +++ b/src/shield/store.rs @@ -1,107 +1,118 @@ -//! 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." -)] +/// 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)] -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 nanoseconds 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_nanos(*o.get())); + let verdict = gcra.check(stored, now); + if verdict.allowed { + *o.get_mut() = dur_nanos(verdict.new_tat); + } + verdict + } + Entry::Vacant(v) => { + let verdict = gcra.check(None, now); + if verdict.allowed { + v.insert(dur_nanos(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_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 = 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 { + 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 @@ -109,192 +120,65 @@ 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 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)] 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..a8caa71 --- /dev/null +++ b/src/shield/tests.rs @@ -0,0 +1,555 @@ +//! 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: 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("fp", &nested, "1.1.1.1", &HeaderMap::new(), Some(&claims)).identity, + "acme" + ); + // No claims (anonymous) → IP fallback, so the limit can't be dodged. + 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. + 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 store_key_namespaced_by_fingerprint_and_value() { + let key = KeySource::Ip; + // 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] +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 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 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()); + 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); + } + + #[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" + ); + } + + #[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"); + } + + #[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}"); + } +} 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 + } +}