perfect-postcode/server-rs/src/ratelimit.rs

//! Defense-in-depth guards for unlicensed ("demo") API traffic: a per-key rate
//! limit and a server-side filter-count cap. These complement the client-side UX
//! limits so a scripted client can't trivially bypass them. Licensed users and
//! admins are exempt.
//!
//! This is a soft anti-abuse measure, not a hard security boundary: the rate-limit
//! key falls back to the (spoofable) client IP for anonymous users, and the demo
//! gate is best-effort by design (see demo_zone.rs). Logged-in free accounts are keyed
//! by their stable account id, which spoofing the IP header does not evade.

use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::time::{Duration, Instant};

use axum::http::HeaderMap;

use axum::extract::{ConnectInfo, Request};
use axum::http::StatusCode;
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use parking_lot::Mutex;
use rustc_hash::FxHashMap;
use serde_json::json;
use url::form_urlencoded;

use crate::auth::OptionalUser;
use crate::consts::{
    DEMO_MAX_FILTERS, DEMO_RATE_LIMIT_MAX, DEMO_RATE_LIMIT_MAX_KEYS, DEMO_RATE_LIMIT_WINDOW_SECS,
};
use crate::state::AppState;

/// Extract the client IP, preferring reverse-proxy/CDN headers (we sit behind one
/// in production) and falling back to the socket peer. Client-supplied and
/// spoofable — only used as a coarse rate-limit key for anonymous visitors.
fn client_ip(headers: &HeaderMap, peer: Option<IpAddr>) -> Option<IpAddr> {
    let from_header = |name: &str, take_first: bool| -> Option<IpAddr> {
        let raw = headers.get(name)?.to_str().ok()?;
        let candidate = if take_first {
            raw.split(',').next().unwrap_or(raw)
        } else {
            raw
        };
        candidate.trim().parse().ok()
    };

    from_header("cf-connecting-ip", false)
        .or_else(|| from_header("x-forwarded-for", true))
        .or_else(|| from_header("x-real-ip", false))
        .or(peer)
}

/// Sliding-window request counter keyed by account id or client IP.
pub struct DemoRateLimiter {
    windows: Mutex<FxHashMap<String, Vec<Instant>>>,
}

impl DemoRateLimiter {
    pub fn new() -> Self {
        Self {
            windows: Mutex::new(FxHashMap::default()),
        }
    }

    /// Record a hit for `key`; returns `false` when the per-window limit is
    /// exceeded (caller should reject with 429).
    pub fn check(&self, key: &str) -> bool {
        let now = Instant::now();
        let window = Duration::from_secs(DEMO_RATE_LIMIT_WINDOW_SECS);
        let mut map = self.windows.lock();

        // Bound memory: when the table grows large, drop keys with no recent hits.
        // If even that can't get us back under the cap (e.g. a spoofed-IP flood
        // minting a fresh key per request), clear the table outright. That resets
        // everyone's window — acceptable under attack — and keeps both the size and
        // the cost of this scan bounded (it then won't re-run for ~MAX_KEYS inserts,
        // instead of scanning O(n) on every request once full).
        if map.len() > DEMO_RATE_LIMIT_MAX_KEYS {
            map.retain(|_, hits| hits.iter().any(|t| now.duration_since(*t) < window));
            if map.len() > DEMO_RATE_LIMIT_MAX_KEYS {
                map.clear();
            }
        }

        let hits = map.entry(key.to_string()).or_default();
        hits.retain(|t| now.duration_since(*t) < window);
        if hits.len() >= DEMO_RATE_LIMIT_MAX {
            return false;
        }
        hits.push(now);
        true
    }
}

impl Default for DemoRateLimiter {
    fn default() -> Self {
        Self::new()
    }
}

/// Whether a request is internal/direct rather than from a public visitor via the
/// edge proxy. Internal callers (the screenshot/OG service, health checks, local
/// dev) reach the server container directly: a private/loopback socket peer and no
/// edge-proxy forwarding headers. We exempt them so demo guards never throttle our
/// own services. (Behind the public proxy, real visitors carry forwarding headers.)
fn is_internal_request(headers: &HeaderMap, peer: Option<IpAddr>) -> bool {
    let has_forwarding = headers.contains_key("cf-connecting-ip")
        || headers.contains_key("x-forwarded-for")
        || headers.contains_key("x-real-ip");
    if has_forwarding {
        return false;
    }
    match peer {
        Some(IpAddr::V4(v4)) => {
            v4.is_private() || v4.is_loopback() || v4.is_link_local() || v4.is_unspecified()
        }
        Some(IpAddr::V6(v6)) => {
            v6.is_loopback()
                || v6.is_unspecified()
                || (v6.segments()[0] & 0xfe00) == 0xfc00 // fc00::/7 unique-local
                || (v6.segments()[0] & 0xffc0) == 0xfe80 // fe80::/10 link-local
        }
        None => true,
    }
}

/// Count non-empty `;;`-separated entries in the `filters` query parameter.
fn filter_count(query: &str) -> usize {
    for (key, value) in form_urlencoded::parse(query.as_bytes()) {
        if key == "filters" {
            return value.split(";;").filter(|entry| !entry.trim().is_empty()).count();
        }
    }
    0
}

/// Whether the request carries a non-empty `share` code. Such requests view a shared
/// dashboard whose saved filter set may legitimately exceed the demo cap, so they're
/// exempt from the filter cap — the actual share-bounds authorization still runs in
/// the handler.
fn has_share_code(query: &str) -> bool {
    form_urlencoded::parse(query.as_bytes())
        .any(|(key, value)| key == "share" && !value.trim().is_empty())
}

/// Middleware applying the demo filter cap and rate limit to unlicensed `/api/`
/// traffic. Must run after `auth_middleware` (for `OptionalUser`) and the state
/// injection layer (for `AppState`).
pub async fn demo_guard_middleware(req: Request, next: Next) -> Response {
    let path = req.uri().path();
    // Map/overlay tiles aren't property-data pulls, and the basemap visibly breaks if
    // they're throttled, so the demo guards only cover the data endpoints.
    if !path.starts_with("/api/")
        || path.starts_with("/api/tiles/")
        || path.starts_with("/api/overlays/")
    {
        return next.run(req).await;
    }

    // Licensed users and admins bypass demo guards entirely.
    let user = req.extensions().get::<OptionalUser>().and_then(|u| u.0.clone());
    if let Some(u) = &user {
        if u.is_admin || u.subscription == "licensed" {
            return next.run(req).await;
        }
    }

    let peer = req
        .extensions()
        .get::<ConnectInfo<SocketAddr>>()
        .map(|info| info.0.ip());

    // Exempt our own internal services (screenshots/OG, health checks, dev).
    if is_internal_request(req.headers(), peer) {
        return next.run(req).await;
    }

    // Server-side filter cap (mirrors the client demo limit; blunts oversized
    // aggregation requests). Share recipients are exempt — their saved view may
    // legitimately carry >5 filters and the share grant is checked in the handler.
    if let Some(query) = req.uri().query() {
        if !has_share_code(query) && filter_count(query) > DEMO_MAX_FILTERS {
            return (
                StatusCode::BAD_REQUEST,
                axum::Json(json!({
                    "error": "filter_limit",
                    "message": format!("The demo is limited to {DEMO_MAX_FILTERS} filters"),
                })),
            )
                .into_response();
        }
    }

    // Rate limit per account (stable, not IP-spoofable) or, for anonymous
    // visitors, per client IP.
    let key = match &user {
        Some(u) => format!("acct:{}", u.id),
        None => match client_ip(req.headers(), peer) {
            Some(ip) => format!("ip:{ip}"),
            None => "anon".to_string(),
        },
    };

    if let Some(state) = req.extensions().get::<Arc<AppState>>() {
        if !state.demo_rate_limiter.check(&key) {
            return (
                StatusCode::TOO_MANY_REQUESTS,
                axum::Json(json!({
                    "error": "rate_limited",
                    "message": "Too many requests — slow down, or sign in for full access",
                })),
            )
                .into_response();
        }
    }

    next.run(req).await
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn counts_filters_in_query() {
        assert_eq!(filter_count(""), 0);
        assert_eq!(filter_count("resolution=9"), 0);
        assert_eq!(filter_count("filters=price%3A1%3A2"), 1);
        assert_eq!(
            filter_count("filters=price%3A1%3A2%3B%3Bbeds%3A2%3A4&resolution=9"),
            2
        );
        // Trailing/empty entries are ignored.
        assert_eq!(filter_count("filters=price%3A1%3A2%3B%3B"), 1);
    }

    #[test]
    fn detects_share_code() {
        assert!(!has_share_code(""));
        assert!(!has_share_code("filters=a%3B%3Bb"));
        assert!(!has_share_code("share=")); // empty value doesn't count
        assert!(has_share_code("share=abc123"));
        assert!(has_share_code("bounds=1,2,3,4&share=xyz&filters=a"));
    }

    #[test]
    fn rate_limiter_blocks_past_the_window_limit() {
        let limiter = DemoRateLimiter::new();
        // The first DEMO_RATE_LIMIT_MAX hits for a key pass; the next is blocked.
        for _ in 0..DEMO_RATE_LIMIT_MAX {
            assert!(limiter.check("acct:abc"));
        }
        assert!(!limiter.check("acct:abc"));
        // A different key has its own independent budget.
        assert!(limiter.check("acct:xyz"));
    }
}