structured_proxy/shield/

mod.rs

//! Shield: request rate limiting.
//!
//! 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::extract::Request;
use axum::http::header::RETRY_AFTER;
use axum::http::{HeaderMap, HeaderName, HeaderValue};
use axum::middleware::Next;
use axum::response::Response;

use crate::client_address::{ClientAddress, Resolution};
use crate::config::ShieldConfig;
use gcra::Verdict;
use matcher::{CompiledProfile, CompiledRule, KeySource, Phase};
use store::GcraStore;

/// Compiled Shield rules, limit tiers, and the local GCRA store.
pub struct Shield {
    rules: Vec<CompiledRule>,
    profiles: HashMap<String, CompiledProfile>,
    /// Applied when a matched rule resolves no other limit.
    default_profile: Option<CompiledProfile>,
    /// Resolve a key's limit from validated JWT claims, when configured.
    jwt_limits: Option<resolve::JwtLimits>,
    /// Resolve a key's limit from an external service (cached, async).
    limit_service: Option<Arc<resolve::LimitService>>,
    /// 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<Arc<global::GlobalCounters>>,
    store: GcraStore,
}

impl Shield {
    /// Build a Shield from config, or `None` when disabled / has no rules.
    ///
    /// # Errors
    /// Returns an error string when a glob pattern, rate or profile reference
    /// fails to compile.
    pub fn build(config: &ShieldConfig) -> Result<Option<Arc<Self>>, String> {
        if !config.enabled {
            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 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 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 {
            rules,
            profiles,
            default_profile,
            jwt_limits,
            limit_service,
            #[cfg(feature = "redis")]
            global,
            store: GcraStore::new(),
        })))
    }

    /// 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))
    }

    /// Whether a rule runs in `phase`: a phase without one needs no layer.
    pub(crate) fn enforces(&self, phase: Phase) -> bool {
        self.rules.iter().any(|rule| rule.phase == phase)
    }

    /// 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<CompiledProfile> {
        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)
    }

    /// The rule's pinned profile, else the default profile.
    fn static_profile(&self, rule: &CompiledRule) -> Option<CompiledProfile> {
        rule.profile
            .as_ref()
            .and_then(|name| self.profiles.get(name))
            .or(self.default_profile.as_ref())
            .copied()
    }
}

/// 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<Arc<Shield>>,
    request: Request,
    next: Next,
) -> Response {
    enforce(&shield, Phase::PreAuth, request, next).await
}

/// 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<Arc<Shield>>,
    request: Request,
    next: Next,
) -> Response {
    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 {
    match shield.decide(phase, &request) {
        Decision::Pass => next.run(request).await,
        Decision::Reject(response) => response,
        Decision::Report(report) => {
            let mut response = next.run(request).await;
            report.apply(&mut response);
            response
        }
    }
}

/// What a phase's limit makes of a request.
pub(crate) enum Decision {
    /// No limit applies.
    Pass,
    /// Over the limit: the answer.
    Reject(Response),
    /// Admitted: the budget its response reports.
    Report(Report),
}

/// The budget an admitted request's response reports.
pub(crate) struct Report {
    limit: u64,
    remaining: u64,
    verdict: Verdict,
}

impl Report {
    /// Report the budget on `response`: the tightest one 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; they are
    /// overwritten 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.
    pub(crate) fn apply(&self, response: &mut Response) {
        maybe_tighten_rate_headers(
            response.headers_mut(),
            self.limit,
            self.remaining,
            &self.verdict,
        );
    }
}

impl Shield {
    /// Match a phase's rule for `request` and charge its limit.
    pub(crate) fn decide<B>(&self, phase: Phase, request: &http::Request<B>) -> Decision {
        let path = request.uri().path();
        let Some(rule) = self.match_rule(path, phase) else {
            return Decision::Pass;
        };
        self.charge(rule, request)
    }

    /// Charge `request` to the limit of `rule`, which its path matched.
    fn charge<B>(&self, rule: &CompiledRule, request: &http::Request<B>) -> Decision {
        // The proxy resolved the client address before any guard. Shield used on
        // its own, in a router of the embedder's, gets it from a
        // `ClientAddressLayer` in front of it; without one it falls back to what
        // resolution gives with no trusted proxy: the peer an axum server records.
        let unresolved;
        let client = match request.extensions().get::<ClientAddress>() {
            Some(client) => client,
            None => {
                let peer = request
                    .extensions()
                    .get::<axum::extract::ConnectInfo<std::net::SocketAddr>>()
                    .map(|ci| ci.0);
                unresolved = ClientAddress::from_peer(peer);
                &unresolved
            }
        };
        let claims = request
            .extensions()
            .get::<crate::auth::ValidatedClaims>()
            .map(|c| c.0.as_ref());
        let key = rule_key(
            &rule.fingerprint,
            &rule.key,
            client,
            request.headers(),
            claims,
        );

        // The limit service resolves per-principal, so it needs the real identity;
        // the store / shared counter only need the de-identified `store` key.
        let Some(profile) = self.resolve_limit(rule, claims, key.identity.as_str()) else {
            // No limit resolves for this rule (JWT/service/profile/default all
            // absent): allow the request unmetered.
            return Decision::Pass;
        };

        // 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 = self
            .global
            .as_ref()
            .map(|g| g.fleet_remaining(key.store.as_str(), profile.limit, profile.window));
        #[cfg(feature = "redis")]
        if fleet_remaining == Some(0) {
            return Decision::Reject(global_reject(profile.limit, profile.window));
        }

        // 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 = self.store.check(key.store.as_str(), &profile.gcra);
        if !verdict.allowed {
            return Decision::Reject(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) = &self.global {
                    global.record(key.store.as_str(), profile.window);
                }
            }
            reconciled_headers(verdict, fleet_remaining, profile.window)
        };
        Decision::Report(Report {
            limit: profile.limit,
            remaining: reported,
            verdict: header_verdict,
        })
    }
}

/// Combine the local GCRA `verdict` with the optional fleet remaining into the
/// reported remaining and the verdict whose reset drives the headers. The count
/// is the tighter of local and fleet, minus this admit (`fr - 1`).
fn reconciled_headers(
    verdict: Verdict,
    fleet_remaining: Option<u64>,
    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: &HeaderName| {
        headers
            .get(name)
            .and_then(|v| v.to_str().ok())
            .and_then(|v| v.parse::<u64>().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::<u64>().ok())
        {
            let reset = secs_ceil(verdict.reset_after);
            if reset > current {
                headers.insert(RETRY_AFTER, HeaderValue::from(reset));
            }
        }
    }
}

/// 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<'r> {
    /// 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: StoreKey,
    /// The raw identity for per-principal limit-service resolution (the service
    /// must see the real principal to resolve its tier). Not persisted.
    identity: Identity<'r>,
}

/// A store key, `fingerprint:tag:hash`, held inline: the fingerprint and the
/// hash are 32 hex digits each and the tag at most three letters.
struct StoreKey {
    bytes: [u8; 72],
    len: usize,
}

impl StoreKey {
    fn new(fingerprint: &str, tag: &str, hash: &[u8; 32]) -> Self {
        let mut key = Self {
            bytes: [0; 72],
            len: 0,
        };
        for part in [fingerprint.as_bytes(), b":", tag.as_bytes(), b":", hash] {
            key.bytes[key.len..key.len + part.len()].copy_from_slice(part);
            key.len += part.len();
        }
        key
    }

    fn as_str(&self) -> &str {
        std::str::from_utf8(&self.bytes[..self.len]).expect("hex digits, a tag and colons")
    }
}

/// The value a rule keys a request's limit by, borrowed from the request where
/// it can be.
enum Identity<'r> {
    Address(AddressText),
    Borrowed(&'r str),
    Owned(String),
}

impl Identity<'_> {
    fn as_str(&self) -> &str {
        match self {
            Identity::Address(address) => address.as_str(),
            Identity::Borrowed(text) => text,
            Identity::Owned(text) => text,
        }
    }
}

/// An IP address as text, held inline: an IPv6 address is at most 45
/// characters (RFC 4291 §2.2, with an embedded IPv4 address).
struct AddressText {
    bytes: [u8; 45],
    len: usize,
}

impl AddressText {
    fn new(ip: std::net::IpAddr) -> Self {
        use std::fmt::Write;
        let mut text = Self {
            bytes: [0; 45],
            len: 0,
        };
        write!(text, "{ip}").expect("an IP address fits in 45 characters");
        text
    }

    fn as_str(&self) -> &str {
        std::str::from_utf8(&self.bytes[..self.len]).expect("an address is written as text")
    }
}

impl std::fmt::Write for AddressText {
    fn write_str(&mut self, s: &str) -> std::fmt::Result {
        let end = self.len + s.len();
        let slot = self.bytes.get_mut(self.len..end).ok_or(std::fmt::Error)?;
        slot.copy_from_slice(s.as_bytes());
        self.len = end;
        Ok(())
    }
}

/// Derive the store key and resolution identity for a matched rule. Every source
/// falls back to the client address 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<'r>(
    fingerprint: &str,
    key: &KeySource,
    client: &ClientAddress,
    headers: &'r HeaderMap,
    claims: Option<&serde_json::Value>,
) -> RuleKey<'r> {
    let (tag, identity) = match key {
        KeySource::Ip => ("ip", address_identity(client)),
        KeySource::Header(name) => match header_str(headers, name) {
            Some(v) => ("hdr", Identity::Borrowed(v)),
            None => ("ip", address_identity(client)),
        },
        KeySource::JwtClaim(claim) => match claims.and_then(|c| resolve::claim_str(c, claim)) {
            Some(v) => ("jwt", Identity::Owned(v)),
            None => ("ip", address_identity(client)),
        },
    };
    let hash = matcher::short_hash_hex(identity.as_str().as_bytes());
    RuleKey {
        store: StoreKey::new(fingerprint, tag, &hash),
        identity,
    }
}

/// The identity a client address keys a limit by: the address, or one shared
/// bucket per reason none resolved. A request that cannot name its client
/// shares its budget with every other such request rather than escaping the
/// limit, and a proxy's own address never stands in for the client's.
fn address_identity(client: &ClientAddress) -> Identity<'static> {
    match client.resolution() {
        Resolution::Peer(ip) | Resolution::Forwarded(ip) => Identity::Address(AddressText::new(ip)),
        Resolution::Invalid(_) => Identity::Borrowed("invalid"),
        Resolution::Unavailable => Identity::Borrowed("unknown"),
    }
}

/// Trimmed, non-empty value of a header.
fn header_str<'h>(headers: &'h HeaderMap, name: &str) -> Option<&'h str> {
    headers
        .get(name)
        .and_then(|v| v.to_str().ok())
        .map(str::trim)
        .filter(|s| !s.is_empty())
}

static RATELIMIT_LIMIT: HeaderName = HeaderName::from_static("ratelimit-limit");
static RATELIMIT_REMAINING: HeaderName = HeaderName::from_static("ratelimit-remaining");
static RATELIMIT_RESET: HeaderName = HeaderName::from_static("ratelimit-reset");

/// 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) {
    headers.insert(RATELIMIT_LIMIT.clone(), HeaderValue::from(limit));
    headers.insert(RATELIMIT_REMAINING.clone(), HeaderValue::from(remaining));
    headers.insert(
        RATELIMIT_RESET.clone(),
        HeaderValue::from(secs_ceil(verdict.reset_after)),
    );
}

/// 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 = crate::guard::reject(tonic::Code::ResourceExhausted, "rate limit exceeded");
    let headers = response.headers_mut();
    attach_rate_headers(headers, limit, remaining, verdict);
    headers.insert(
        RETRY_AFTER,
        HeaderValue::from(secs_ceil(verdict.retry_after)),
    );
    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;