didbot_serve/

rate_limit.rs

//! A minimal fixed-window rate limiter, shared by every anonymous route that
//! does real work before it knows who is calling.
//!
//! `bot.did.createAccount` and `POST /oauth/par` are the motivating cases:
//! the first mints a signing key and a hostname, the second fetches a
//! document from a host the caller names, and both do it for a caller this
//! server has not authenticated yet. This limiter is what stands in front of
//! that cost.
//!
//! A small fixed-window counter: one `HashMap`, one lock, no background
//! task. The per-account budget on `atproto-proxy` is a `governor` token
//! bucket instead; see `crate::proxy`.
//!
//! Fixed-window rather than a token bucket: it under-approximates a true rate
//! at the window boundary (a caller can burst up to `limit` twice, once at
//! the end of one window and once at the start of the next), which is an
//! acceptable trade for a limiter this small.
//!
//! # A caller's name
//!
//! [`caller_key`] defaults to the TCP peer address — the socket's own
//! `SocketAddr`, delivered through `axum`'s `ConnectInfo` extractor, wired at
//! [`serve_router`](crate::serve_router) and
//! [`serve_tls`](crate::serve_tls) — because that is the one address a
//! caller cannot lie about. `didbot-tls` terminates TLS itself with no
//! reverse proxy in front of it, so `X-Forwarded-For` and `X-Real-IP` are, by
//! default, ordinary caller-supplied header values: trusting them
//! unconditionally would let a caller mint a fresh budget on every request
//! just by sending a different header, which is not a degraded limiter but a
//! bypassed one. [`caller_key`] only reads those headers when
//! `trust_forwarded_headers` is `true` — an explicit, off-by-default setting
//! an operator turns on after putting a real reverse proxy in front of this
//! server. Nothing here infers that posture from a header's mere presence.
//!
//! Turned on, the *rightmost* `X-Forwarded-For` element is read: a proxy
//! appends its own observation to whatever the client sent, so the right-hand
//! end is the proxy's and the left-hand end is the client's. This holds for
//! the single trusted proxy a deployment puts in front of one server; a chain
//! of several needs a setting naming how far to walk, and there is none.
//!
//! A caller reached with no [`ConnectInfo`] at
//! all — a test harness driving a [`Router`](axum::Router) directly, or a
//! deployment that has not wired [`serve_router`](crate::serve_router) or
//! [`serve_tls`](crate::serve_tls) — shares the same `"unknown"` bucket
//! every caller got before this address was available, rather than failing
//! the request.
//!

use std::collections::HashMap;
use std::convert::Infallible;
use std::net::{IpAddr, Ipv6Addr, SocketAddr};
use std::sync::Mutex;
use std::time::{Duration, Instant};

use axum::extract::connect_info::ConnectInfo;
use axum::extract::FromRequestParts;
use axum::http::request::Parts;
use axum::http::HeaderMap;

/// A handler's own extraction of the caller's TCP peer address, for
/// [`caller_key`].
///
/// Unlike `Option<ConnectInfo<SocketAddr>>`, this never fails to extract:
/// `axum::extract::connect_info::ConnectInfo` does not implement axum's
/// `OptionalFromRequestParts`, so wrapping it in `Option` in a handler
/// signature is a compile error, not a graceful `None`. [`Peer`] reads the
/// same [`ConnectInfo`] extension directly and answers `None` when it is
/// absent — a test driving a [`Router`](axum::Router) with
/// [`tower::ServiceExt::oneshot`], or a deployment that has not wired
/// [`crate::serve_router`]/[`crate::serve_tls`] — rather than rejecting the
/// request.
pub struct Peer(pub Option<IpAddr>);

impl<S> FromRequestParts<S> for Peer
where
    S: Send + Sync,
{
    type Rejection = Infallible;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        Ok(Peer(
            parts
                .extensions
                .get::<ConnectInfo<SocketAddr>>()
                .map(|ConnectInfo(addr)| addr.ip()),
        ))
    }
}

/// The default window a caller's request count is measured over.
///
/// Not a deployment's to set yet: `plan/config.md` holds `[limits]` as
/// schema for this and `oauth/rate_limit.rs`'s own budgets, deferred until
/// the two converge on one shape rather than adding a key each now that a
/// later pass would have to reconcile.
pub const WINDOW: Duration = Duration::from_secs(60);

/// The default number of requests one key may make per [`WINDOW`], for a
/// limiter that does not pick its own. See [`WINDOW`] for why this is not a
/// key today.
pub const DEFAULT_LIMIT: u32 = 30;

/// The longest key that gets a bucket of its own.
///
/// [`caller_key`] answers an address, but a limiter keyed on something the
/// caller *sends* is handed an arbitrary string by an unauthenticated
/// caller. `MAX_DID_LENGTH` is the
/// longest an atproto identifier can be at all: a handle caps at 253 bytes
/// and a DID at this, so nothing longer than this resolves to an account on
/// any deployment. Keys past it are therefore not identifiers being checked,
/// and giving each one a bucket would mean retaining a caller-chosen string
/// per request in a map with no other bound on it.
pub const MAX_KEY_LENGTH: usize = didbot_identity::MAX_DID_LENGTH;

/// The single bucket every key over [`MAX_KEY_LENGTH`] shares.
///
/// The same answer [`caller_key`] already gives a caller it cannot tell
/// apart from another: one shared budget rather than a free one. Sharing is
/// strictly the stricter reading — the keys collapsed into it cannot name an
/// account, so nothing legitimate is being made to queue behind anything
/// else — and it is what keeps a key's *content* out of this map.
pub const OVER_LONG_KEY: &str = "over-long";

/// The key a bucket is actually filed under: `key` itself, or the shared
/// [`OVER_LONG_KEY`] when it is longer than [`MAX_KEY_LENGTH`].
fn bucket_key(key: &str) -> &str {
    if key.len() > MAX_KEY_LENGTH {
        OVER_LONG_KEY
    } else {
        key
    }
}

struct Bucket {
    window_started_at: Instant,
    count: u32,
}

/// Counts requests per key (typically a caller's address, from
/// [`caller_key`]) within a fixed window, and says whether the next one is
/// still allowed.
pub struct RateLimiter {
    limit: u32,
    window: Duration,
    state: Mutex<State>,
}

/// The buckets, and when they were last swept for elapsed windows.
struct State {
    buckets: HashMap<String, Bucket>,
    last_swept: Instant,
}

impl RateLimiter {
    /// A limiter allowing `limit` requests per `window`, per key.
    pub fn new(limit: u32, window: Duration) -> Self {
        Self {
            limit,
            window,
            state: Mutex::new(State {
                buckets: HashMap::new(),
                last_swept: Instant::now(),
            }),
        }
    }

    /// Records one request for `key` and reports whether it is within the
    /// limit. A key seen for the first time, or one whose window has
    /// elapsed, starts a fresh window.
    ///
    /// A key over [`MAX_KEY_LENGTH`] is counted against the one shared
    /// [`OVER_LONG_KEY`] bucket rather than a bucket of its own, and buckets
    /// whose windows have elapsed are dropped rather than kept — see both
    /// items' own docs for why neither changes what this answers.
    pub fn check(&self, key: &str) -> bool {
        let now = Instant::now();
        let key = bucket_key(key);
        let mut state = self.state.lock().unwrap_or_else(|p| p.into_inner());
        // An elapsed window is reset to zero the moment it is next touched,
        // so a bucket already past its window says exactly what an absent
        // one says. Dropping those is what keeps this map's size a function
        // of recent callers rather than of every key ever seen.
        //
        // Swept once per window rather than on every call, because a scan
        // per request would make one caller's cost grow with how many other
        // callers there are — an amplification of its own, and the shape of
        // the one this sweep exists to close. Amortized this way a bucket
        // outlives its window by at most one more, and the scan it costs is
        // paid once across every request in that window.
        if now.duration_since(state.last_swept) >= self.window {
            let window = self.window;
            state
                .buckets
                .retain(|_, bucket| now.duration_since(bucket.window_started_at) < window);
            state.last_swept = now;
        }
        let bucket = state.buckets.entry(key.to_owned()).or_insert(Bucket {
            window_started_at: now,
            count: 0,
        });
        if now.duration_since(bucket.window_started_at) >= self.window {
            bucket.window_started_at = now;
            bucket.count = 0;
        }
        bucket.count += 1;
        bucket.count <= self.limit
    }

    /// How many keys this limiter is currently holding a bucket for.
    ///
    /// The limiter's retained footprint, for a test that needs to assert on
    /// it and for anything that later wants to log it. A bucket whose window
    /// has elapsed but which the sweep in [`check`](Self::check) has not
    /// reached yet still counts.
    pub fn tracked_keys(&self) -> usize {
        self.state
            .lock()
            .unwrap_or_else(|p| p.into_inner())
            .buckets
            .len()
    }
}

impl Default for RateLimiter {
    fn default() -> Self {
        Self::new(DEFAULT_LIMIT, WINDOW)
    }
}

/// Names the caller a [`RateLimiter`] should key its budget on.
///
/// `peer` is the TCP socket's own address — see this module's doc for how a
/// handler gets one — and is what this returns unless `trust_forwarded_headers`
/// says to read `X-Forwarded-For`/`X-Real-IP` instead. `peer` being `None`
/// (no [`ConnectInfo`] reached the handler) and
/// `trust_forwarded_headers` being `true` with neither header present both
/// answer the same `"unknown"`, the shared bucket every caller falls back to
/// when this cannot tell them apart.
///
/// An IPv6 address is keyed on its /64, because one host routinely holds a
/// whole /64. An IPv4-mapped IPv6 address is keyed as its IPv4 address. A
/// forwarded value that is not an address at all leaves the caller on the TCP
/// peer's budget, so a caller cannot open a bucket of its own by inventing a
/// string to be keyed under.
pub fn caller_key(
    headers: &HeaderMap,
    peer: Option<IpAddr>,
    trust_forwarded_headers: bool,
) -> String {
    if trust_forwarded_headers {
        if let Some(observed) = forwarded_caller(headers) {
            return address_key(observed);
        }
    }
    peer.map(address_key)
        .unwrap_or_else(|| "unknown".to_owned())
}

/// The budget key for one address: see [`caller_key`].
fn address_key(ip: IpAddr) -> String {
    match ip.to_canonical() {
        IpAddr::V4(v4) => v4.to_string(),
        IpAddr::V6(v6) => {
            let prefix = Ipv6Addr::from_bits(v6.to_bits() & !(u128::MAX >> 64));
            format!("{prefix}/64")
        }
    }
}

/// The address the reverse proxy in front of this server saw the request
/// arrive from, for a deployment that has opted into trusting it.
///
/// The **rightmost** `X-Forwarded-For` element, because a proxy appends:
/// nginx's `proxy_add_x_forwarded_for` and an AWS ALB both put the peer they
/// saw at the right-hand end and keep whatever the client sent to its left.
/// The leftmost element is therefore the client's own text, and reading it is
/// what let a caller send a different value on every request and get a fresh
/// budget against every limiter each time.
///
/// `None` when neither header holds an address: an unparseable value is not a
/// caller's name, and the TCP peer is.
fn forwarded_caller(headers: &HeaderMap) -> Option<IpAddr> {
    let element = |name: &str, take: fn(&str) -> Option<&str>| {
        headers
            .get(name)
            .and_then(|value| value.to_str().ok())
            .and_then(take)
            .map(str::trim)
            .filter(|value| !value.is_empty())
            .and_then(|value| value.parse::<IpAddr>().ok())
    };
    element("x-forwarded-for", |value| value.rsplit(',').next())
        .or_else(|| element("x-real-ip", |value| Some(value)))
}

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

    #[test]
    fn requests_within_the_limit_are_allowed() {
        let limiter = RateLimiter::new(3, Duration::from_secs(60));
        assert!(limiter.check("1.2.3.4"));
        assert!(limiter.check("1.2.3.4"));
        assert!(limiter.check("1.2.3.4"));
    }

    #[test]
    fn a_request_past_the_limit_is_refused() {
        let limiter = RateLimiter::new(2, Duration::from_secs(60));
        assert!(limiter.check("1.2.3.4"));
        assert!(limiter.check("1.2.3.4"));
        assert!(!limiter.check("1.2.3.4"));
    }

    #[test]
    fn different_keys_have_independent_budgets() {
        let limiter = RateLimiter::new(1, Duration::from_secs(60));
        assert!(limiter.check("1.2.3.4"));
        assert!(limiter.check("5.6.7.8"));
        assert!(!limiter.check("1.2.3.4"));
    }

    #[test]
    fn a_window_reset_allows_more_requests() {
        let limiter = RateLimiter::new(1, Duration::from_millis(10));
        assert!(limiter.check("1.2.3.4"));
        assert!(!limiter.check("1.2.3.4"));
        std::thread::sleep(Duration::from_millis(20));
        assert!(limiter.check("1.2.3.4"));
    }

    #[test]
    fn buckets_for_elapsed_windows_are_not_retained() {
        let limiter = RateLimiter::new(5, Duration::from_millis(10));
        for i in 0..50 {
            assert!(limiter.check(&format!("key-{i}")));
        }
        assert_eq!(limiter.tracked_keys(), 50);
        std::thread::sleep(Duration::from_millis(20));
        // Every window above has elapsed, so every one of those buckets now
        // says exactly what a key this limiter has never seen says. Touching
        // one more key must not leave the other 50 resident.
        assert!(limiter.check("key-fresh"));
        assert_eq!(limiter.tracked_keys(), 1);
    }

    #[test]
    fn an_over_long_key_is_not_retained_verbatim() {
        let limiter = RateLimiter::new(5, Duration::from_secs(60));
        // An identifier no handle and no DID could ever be. Nothing this
        // long resolves to an account, so it must not buy the caller a
        // bucket of its own — let alone one holding a copy of it.
        let huge = "z".repeat(64 * 1024);
        assert!(limiter.check(&huge));
        let other = "y".repeat(64 * 1024);
        assert!(limiter.check(&other));
        // Both collapse into the one shared over-long bucket.
        assert_eq!(limiter.tracked_keys(), 1);
    }

    #[test]
    fn caller_key_defaults_to_the_peer_address_and_ignores_forwarded_headers() {
        let mut headers = HeaderMap::new();
        headers.insert("x-forwarded-for", "9.9.9.9, 10.0.0.1".parse().unwrap());
        headers.insert("x-real-ip", "8.8.8.8".parse().unwrap());
        let peer: IpAddr = "203.0.113.9".parse().unwrap();
        // `trust_forwarded_headers` is false: a caller-supplied header,
        // however it is shaped, must not override the real peer address.
        assert_eq!(caller_key(&headers, Some(peer), false), "203.0.113.9");
    }

    #[test]
    fn caller_key_shares_one_bucket_with_no_peer_and_no_trust() {
        let headers = HeaderMap::new();
        assert_eq!(caller_key(&headers, None, false), "unknown");
    }

    /// The rightmost element, which is the proxy's own observation: a proxy
    /// appends, so everything to its left is text the client sent.
    #[test]
    fn caller_key_reads_forwarded_headers_only_when_trusted() {
        let mut headers = HeaderMap::new();
        headers.insert("x-forwarded-for", "9.9.9.9, 10.0.0.1".parse().unwrap());
        let peer: IpAddr = "203.0.113.9".parse().unwrap();
        assert_eq!(caller_key(&headers, Some(peer), true), "10.0.0.1");
    }

    /// A caller that writes its own `X-Forwarded-For` gets one budget, not a
    /// fresh one per request: the proxy's appended element is what is read,
    /// and it is the same every time.
    #[test]
    fn a_forged_forwarded_prefix_does_not_mint_a_budget_per_request() {
        let proxy: IpAddr = "10.0.0.1".parse().unwrap();
        let forged = |claim: &str| {
            let mut headers = HeaderMap::new();
            headers.insert(
                "x-forwarded-for",
                format!("{claim}, {proxy}").parse().unwrap(),
            );
            caller_key(&headers, Some("203.0.113.9".parse().unwrap()), true)
        };
        assert_eq!(forged("9.9.9.9"), forged("198.51.100.7"));
        assert_eq!(forged("a"), forged("9.9.9.9"));
        assert_eq!(forged("9.9.9.9"), "10.0.0.1");
    }

    /// A value that is not an address is not a caller's name. Keying it
    /// verbatim gave anything a caller cared to type a bucket of its own.
    #[test]
    fn an_unparseable_forwarded_value_leaves_the_caller_on_the_peer() {
        let peer: IpAddr = "203.0.113.9".parse().unwrap();
        let mut headers = HeaderMap::new();
        headers.insert("x-forwarded-for", "not-an-address".parse().unwrap());
        assert_eq!(caller_key(&headers, Some(peer), true), "203.0.113.9");
        assert_eq!(caller_key(&headers, None, true), "unknown");
    }

    #[test]
    fn caller_key_falls_back_to_x_real_ip_when_trusted() {
        let mut headers = HeaderMap::new();
        headers.insert("x-real-ip", "8.8.8.8".parse().unwrap());
        assert_eq!(caller_key(&headers, None, true), "8.8.8.8");
    }

    #[test]
    fn caller_key_shares_one_bucket_when_trusted_with_no_headers_or_peer() {
        let headers = HeaderMap::new();
        assert_eq!(caller_key(&headers, None, true), "unknown");
    }

    /// One host can pick any address in its /64, so a per-address key would
    /// hand it an unlimited number of budgets.
    #[test]
    fn caller_key_shares_one_budget_across_an_ipv6_64() {
        let headers = HeaderMap::new();
        let key = |ip: &str| caller_key(&headers, Some(ip.parse().unwrap()), false);
        assert_eq!(
            key("2001:db8:1:2::1"),
            key("2001:db8:1:2:ffff:ffff:ffff:ffff")
        );
        assert_ne!(key("2001:db8:1:2::1"), key("2001:db8:1:3::1"));
        assert_eq!(key("::ffff:203.0.113.9"), "203.0.113.9");

        let mut forwarded = HeaderMap::new();
        forwarded.insert("x-forwarded-for", "2001:db8:1:2::9".parse().unwrap());
        assert_eq!(caller_key(&forwarded, None, true), key("2001:db8:1:2::1"));
    }
}