dragoman 0.5.12

//! Request limits for a server that faces the internet itself: a token bucket
//! per client and tier, a cap on the requests in flight, a response deadline,
//! and — for the TLS listener — a cap on open connections per client and a
//! deadline for the TLS handshake.
//!
//! A client is its IP address, an IPv6 one cut to its /64: a host is handed a
//! whole /64, and keying on the full address would let it walk around its own
//! limit. Behind a proxy the peer is the proxy, so a trusted peer — this
//! machine, or a network given in `--trusted-proxies`, such as a CDN's edges —
//! stands for the client it names; any other peer's headers are the client's
//! own words and are ignored. See [`Proxies::client`].
//!
//! Tiers follow what a request costs, not just its path. A pid URL is a cheap
//! page for a browser but a conversion — and possibly a fetch from Crossref or
//! DataCite — for a client asking for BibTeX; `?local_only=true`, which the UI
//! sends for the records on a page, never leaves the local store and counts as
//! a page.

use std::collections::HashMap;
use std::future::Future;
use std::io;
use std::net::{IpAddr, Ipv6Addr, SocketAddr};
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::{Duration, Instant};

use axum::extract::{ConnectInfo, Request, State};
use axum::http::{HeaderMap, HeaderName, HeaderValue, StatusCode, header};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::net::TcpStream;
use tokio::sync::Semaphore;

use crate::negotiate::{self, NegotiateResult};

/// How often idle buckets are dropped. A bucket that has refilled is the same
/// as no bucket, so dropping it changes no answer.
const SWEEP_INTERVAL: Duration = Duration::from_secs(60);

/// What a request costs, and so which bucket it draws from.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Tier {
    /// Entity pages and the records the UI shows on them.
    Pages,
    /// Works lists, facets and counts.
    Lists,
    /// Search, handle lookups, exports, bibliographies and format conversion.
    Expensive,
}

impl Tier {
    fn label(self) -> &'static str {
        match self {
            Tier::Pages => "pages",
            Tier::Lists => "lists",
            Tier::Expensive => "expensive",
        }
    }
}

/// A sustained rate and the burst a client may spend at once.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Rate {
    per_second: f64,
    burst: f64,
}

/// A tier's limit as configured: a rate, or none at all.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct RateLimit(pub Option<Rate>);

/// Parse `RATE/BURST` (requests per second, bucket size), e.g. `20/60`; `0`
/// turns the tier's limit off.
pub fn parse_rate(s: &str) -> Result<RateLimit, String> {
    let s = s.trim();
    if s == "0" {
        return Ok(RateLimit(None));
    }
    let (rate, burst) = s
        .split_once('/')
        .ok_or_else(|| format!("expected RATE/BURST (e.g. 20/60) or 0, got '{s}'"))?;
    let per_second: f64 = rate
        .trim()
        .parse()
        .map_err(|_| format!("invalid rate '{rate}'"))?;
    let burst: f64 = burst
        .trim()
        .parse()
        .map_err(|_| format!("invalid burst '{burst}'"))?;
    if !(per_second > 0.0 && per_second.is_finite()) || !(burst >= 1.0 && burst.is_finite()) {
        return Err(format!(
            "rate must be positive and burst at least 1, got '{s}'"
        ));
    }
    Ok(RateLimit(Some(Rate { per_second, burst })))
}

/// The limits the request middleware applies.
#[derive(Clone, Debug)]
pub struct Limits {
    pub pages: Option<Rate>,
    pub lists: Option<Rate>,
    pub expensive: Option<Rate>,
    /// Requests served at once, across all clients; past it, 503.
    pub max_in_flight: Option<usize>,
    /// How long a response may take before it is abandoned with a 503.
    pub request_timeout: Option<Duration>,
    /// Who may name the client a request comes from.
    pub proxies: Arc<Proxies>,
}

struct Bucket {
    tokens: f64,
    updated: Instant,
}

impl Bucket {
    fn refill(&mut self, rate: Rate, now: Instant) {
        let elapsed = now.saturating_duration_since(self.updated).as_secs_f64();
        self.tokens = (self.tokens + elapsed * rate.per_second).min(rate.burst);
        self.updated = now;
    }
}

/// The token buckets and the in-flight cap, shared by every request.
pub struct Limiter {
    limits: Limits,
    buckets: Mutex<HashMap<(Tier, IpAddr), Bucket>>,
    in_flight: Option<Arc<Semaphore>>,
}

impl Limiter {
    pub fn new(limits: Limits) -> Arc<Self> {
        let in_flight = limits.max_in_flight.map(|n| Arc::new(Semaphore::new(n)));
        Arc::new(Self {
            limits,
            buckets: Mutex::new(HashMap::new()),
            in_flight,
        })
    }

    fn rate(&self, tier: Tier) -> Option<Rate> {
        match tier {
            Tier::Pages => self.limits.pages,
            Tier::Lists => self.limits.lists,
            Tier::Expensive => self.limits.expensive,
        }
    }

    /// Spend a token from `client`'s bucket for `tier`, or say how long until
    /// one is back.
    fn take(&self, tier: Tier, client: IpAddr, now: Instant) -> Result<(), Duration> {
        let Some(rate) = self.rate(tier) else {
            return Ok(());
        };
        let mut buckets = self.buckets.lock().unwrap_or_else(|e| e.into_inner());
        let bucket = buckets.entry((tier, client)).or_insert(Bucket {
            tokens: rate.burst,
            updated: now,
        });
        bucket.refill(rate, now);
        if bucket.tokens >= 1.0 {
            bucket.tokens -= 1.0;
            Ok(())
        } else {
            Err(Duration::from_secs_f64(
                (1.0 - bucket.tokens) / rate.per_second,
            ))
        }
    }

    /// Drop the buckets that have refilled, so the table holds only the clients
    /// that are spending.
    fn sweep(&self, now: Instant) {
        let mut buckets = self.buckets.lock().unwrap_or_else(|e| e.into_inner());
        buckets.retain(|(tier, _), bucket| {
            let Some(rate) = self.rate(*tier) else {
                return false;
            };
            bucket.refill(rate, now);
            bucket.tokens < rate.burst
        });
    }

    /// Sweep idle buckets for as long as the server runs.
    pub fn spawn_sweeper(self: &Arc<Self>) {
        let limiter = Arc::downgrade(self);
        tokio::spawn(async move {
            let mut tick = tokio::time::interval(SWEEP_INTERVAL);
            tick.tick().await;
            loop {
                tick.tick().await;
                let Some(limiter) = limiter.upgrade() else {
                    return;
                };
                limiter.sweep(Instant::now());
            }
        });
    }
}

/// The tier a request draws from; `None` for the page shell, its assets and
/// health checks, which cost next to nothing and every visit needs.
pub fn tier(path: &str, query: Option<&str>, headers: &HeaderMap) -> Option<Tier> {
    match path {
        "/" | "/og-image.png" | "/commonmeta_v1.0.json" | "/api/status" | "/api/metrics" => None,
        _ if path.starts_with("/assets/") => None,
        "/api/search" | "/api/export" | "/api/bibliography" => Some(Tier::Expensive),
        _ if ["/api/handles/", "/dois/", "/works/"]
            .iter()
            .any(|p| path.starts_with(p)) =>
        {
            Some(Tier::Expensive)
        }
        "/api/count"
        | "/api/random"
        | "/api/facets"
        | "/api/container/works"
        | "/api/citations"
        | "/api/citations/facets" => Some(Tier::Lists),
        // A batch of records costs what a list does while it stays in the
        // local store; one that may fetch each from Crossref or DataCite is a
        // conversion, as its single-record form would be.
        "/api/works" if local_only(query) => Some(Tier::Lists),
        "/api/works" => Some(Tier::Expensive),
        _ if path.starts_with("/api/person/") || path.starts_with("/api/org/") => Some(Tier::Lists),
        _ if converts(query, headers) => Some(Tier::Expensive),
        _ => Some(Tier::Pages),
    }
}

/// Whether a query confines the request to the local store.
fn local_only(query: Option<&str>) -> bool {
    query
        .unwrap_or("")
        .split('&')
        .any(|pair| pair == "local_only=true")
}

/// Whether a pid request asks for a format rather than the page — by
/// `?format=` or by its Accept header — without confining itself to the local
/// store.
fn converts(query: Option<&str>, headers: &HeaderMap) -> bool {
    let mut format = false;
    for pair in query.unwrap_or("").split('&') {
        match pair.split_once('=').unwrap_or((pair, "")) {
            ("local_only", "true") => return false,
            ("format", v) if !v.is_empty() => format = true,
            _ => {}
        }
    }
    format
        || headers
            .get(header::ACCEPT)
            .and_then(|v| v.to_str().ok())
            .is_some_and(|a| matches!(negotiate::negotiate(a), NegotiateResult::Format(_)))
}

/// A network, as an address and a prefix length.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Net {
    addr: IpAddr,
    prefix: u8,
}

impl Net {
    fn contains(&self, ip: IpAddr) -> bool {
        match (self.addr, unmap(ip)) {
            (IpAddr::V4(net), IpAddr::V4(ip)) => {
                let mask = u32::MAX
                    .checked_shl(32 - u32::from(self.prefix))
                    .unwrap_or(0);
                u32::from(net) & mask == u32::from(ip) & mask
            }
            (IpAddr::V6(net), IpAddr::V6(ip)) => {
                let mask = u128::MAX
                    .checked_shl(128 - u32::from(self.prefix))
                    .unwrap_or(0);
                u128::from(net) & mask == u128::from(ip) & mask
            }
            _ => false,
        }
    }
}

/// The networks given in `--trusted-proxies`.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Networks(pub Vec<Net>);

/// Parse a comma-separated list of addresses and CIDR networks, e.g.
/// `173.245.48.0/20,2400:cb00::/32,10.0.0.5`.
pub fn parse_networks(s: &str) -> Result<Networks, String> {
    s.split(',')
        .map(str::trim)
        .filter(|n| !n.is_empty())
        .map(|n| {
            let (addr, prefix) = n.split_once('/').unwrap_or((n, ""));
            let addr = addr
                .parse::<IpAddr>()
                .map_err(|_| format!("invalid address in '{n}'"))?;
            let max = if addr.is_ipv4() { 32 } else { 128 };
            let prefix = match prefix {
                "" => max,
                p => p
                    .parse::<u8>()
                    .ok()
                    .filter(|p| *p <= max)
                    .ok_or_else(|| format!("invalid prefix length in '{n}'"))?,
            };
            // An IPv4-mapped network is the IPv4 one it maps.
            Ok(match addr {
                IpAddr::V6(v6) if prefix >= 96 && v6.to_ipv4_mapped().is_some() => Net {
                    addr: unmap(addr),
                    prefix: prefix - 96,
                },
                _ => Net { addr, prefix },
            })
        })
        .collect::<Result<_, _>>()
        .map(Networks)
}

/// Parse an HTTP header name, e.g. `CF-Connecting-IP`.
pub fn parse_header_name(s: &str) -> Result<HeaderName, String> {
    s.trim()
        .parse()
        .map_err(|_| format!("invalid header name '{s}'"))
}

/// The proxies allowed to say who the client is: this machine always, and the
/// networks given in `--trusted-proxies`.
#[derive(Clone, Debug, Default)]
pub struct Proxies {
    networks: Vec<Net>,
    /// The header a trusted proxy puts the client in, e.g. `CF-Connecting-IP`.
    header: Option<HeaderName>,
}

impl Proxies {
    pub fn new(networks: Networks, header: Option<HeaderName>) -> Self {
        Self {
            networks: networks.0,
            header,
        }
    }

    /// Whether `ip` is a proxy whose word on the client is taken.
    pub fn trusts(&self, ip: IpAddr) -> bool {
        ip.is_loopback() || self.networks.iter().any(|n| n.contains(ip))
    }

    /// The client a request comes from. An untrusted peer is the client. A
    /// trusted one names it, in order:
    ///
    /// 1. the `--client-ip-header`, which a CDN sets itself, overwriting
    ///    whatever the client sent;
    /// 2. `X-Real-IP`, but only from this machine: a CDN passes on the one the
    ///    client sent;
    /// 3. `X-Forwarded-For`, read from the right, skipping the trusted proxies
    ///    that appended themselves; the entries before them are the client's.
    ///
    /// A trusted remote peer that names no client stands for itself. A local
    /// request that names none — no proxy, no peer known — is `None`: not
    /// limited.
    pub fn client(&self, peer: Option<IpAddr>, headers: &HeaderMap) -> Option<IpAddr> {
        let peer = peer.map(unmap);
        if let Some(peer) = peer.filter(|p| !self.trusts(*p)) {
            return Some(peer);
        }
        let header = |name: &HeaderName| {
            let value = headers.get(name)?.to_str().ok()?;
            value.split(',').next()?.trim().parse::<IpAddr>().ok()
        };
        let local = peer.is_none_or(|p| p.is_loopback());
        self.header
            .as_ref()
            .and_then(header)
            .or_else(|| {
                local
                    .then(|| header(&HeaderName::from_static("x-real-ip")))
                    .flatten()
            })
            .or_else(|| {
                let forwarded = headers.get("x-forwarded-for")?.to_str().ok()?;
                forwarded
                    .rsplit(',')
                    .filter_map(|e| e.trim().parse::<IpAddr>().ok())
                    .find(|ip| !self.trusts(*ip))
            })
            .map(unmap)
            .or(peer.filter(|p| !p.is_loopback()))
    }
}

/// An IPv4-mapped IPv6 address as the IPv4 address it is.
fn unmap(ip: IpAddr) -> IpAddr {
    match ip {
        IpAddr::V6(v6) => v6.to_ipv4_mapped().map_or(ip, IpAddr::V4),
        IpAddr::V4(_) => ip,
    }
}

/// The peer address the server recorded for a request, if it records one.
pub fn peer(req: &Request) -> Option<IpAddr> {
    req.extensions()
        .get::<ConnectInfo<SocketAddr>>()
        .map(|c| c.0.ip())
}

/// The key a client's limits are kept under: its IPv4 address (an IPv4-mapped
/// IPv6 one unwrapped), or its IPv6 /64.
pub fn client_key(ip: IpAddr) -> IpAddr {
    match ip {
        IpAddr::V4(_) => ip,
        IpAddr::V6(v6) => match v6.to_ipv4_mapped() {
            Some(v4) => IpAddr::V4(v4),
            None => {
                let s = v6.segments();
                IpAddr::V6(Ipv6Addr::new(s[0], s[1], s[2], s[3], 0, 0, 0, 0))
            }
        },
    }
}

/// Rate-limit, cap and time-box a request.
pub async fn middleware(State(limiter): State<Arc<Limiter>>, req: Request, next: Next) -> Response {
    if let Some(tier) = tier(req.uri().path(), req.uri().query(), req.headers())
        && let Some(client) = limiter.limits.proxies.client(peer(&req), req.headers())
        && let Err(wait) = limiter.take(tier, client_key(client), Instant::now())
    {
        metrics::counter!("http_rate_limited_total", "tier" => tier.label()).increment(1);
        tracing::debug!(tier = tier.label(), %client, "rate limited");
        return refused(StatusCode::TOO_MANY_REQUESTS, "rate limit exceeded", wait);
    }

    let _permit = match &limiter.in_flight {
        Some(slots) => match slots.clone().try_acquire_owned() {
            Ok(permit) => Some(permit),
            Err(_) => {
                metrics::counter!("http_overloaded_total").increment(1);
                return refused(
                    StatusCode::SERVICE_UNAVAILABLE,
                    "server busy",
                    Duration::from_secs(1),
                );
            }
        },
        None => None,
    };

    match limiter.limits.request_timeout {
        Some(deadline) => match tokio::time::timeout(deadline, next.run(req)).await {
            Ok(response) => response,
            Err(_) => {
                metrics::counter!("http_timed_out_total").increment(1);
                refused(
                    StatusCode::SERVICE_UNAVAILABLE,
                    "request timed out",
                    Duration::from_secs(5),
                )
            }
        },
        None => next.run(req).await,
    }
}

fn refused(status: StatusCode, message: &'static str, retry_after: Duration) -> Response {
    let secs = retry_after.as_secs_f64().ceil().max(1.0) as u64;
    let mut response = (status, message).into_response();
    let headers = response.headers_mut();
    headers.insert(header::RETRY_AFTER, HeaderValue::from(secs));
    // A refusal is about this moment and this client: a cache that kept it
    // would hand it to everyone else.
    headers.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
    response
}

// ── Connections ──────────────────────────────────────────────────────────────

/// Open connections per client, capped.
pub struct Connections {
    max: usize,
    open: Mutex<HashMap<IpAddr, usize>>,
}

impl Connections {
    pub fn new(max: usize) -> Arc<Self> {
        Arc::new(Self {
            max,
            open: Mutex::new(HashMap::new()),
        })
    }

    /// Count a new connection from `client`, or refuse it at the cap.
    fn open(self: &Arc<Self>, client: IpAddr) -> Option<ConnectionGuard> {
        let mut open = self.open.lock().unwrap_or_else(|e| e.into_inner());
        let count = open.entry(client).or_insert(0);
        if *count >= self.max {
            return None;
        }
        *count += 1;
        Some(ConnectionGuard {
            connections: self.clone(),
            client,
        })
    }
}

/// One open connection; dropping it closes the count.
pub struct ConnectionGuard {
    connections: Arc<Connections>,
    client: IpAddr,
}

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        let mut open = self
            .connections
            .open
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        if let Some(count) = open.get_mut(&self.client) {
            *count -= 1;
            if *count == 0 {
                open.remove(&self.client);
            }
        }
    }
}

/// A connection's stream, holding its client's count for as long as it lives.
/// Until the first request bytes arrive it also holds a deadline: hyper starts
/// its header deadline only once a request has begun, so a client that
/// connects — and, over TLS, shakes hands — but never sends a byte would
/// otherwise hold its connection, and its task, forever.
pub struct Guarded<S> {
    stream: S,
    first_read: Option<Pin<Box<tokio::time::Sleep>>>,
    _guard: Option<ConnectionGuard>,
}

impl<S: AsyncRead + Unpin> AsyncRead for Guarded<S> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let this = self.get_mut();
        let before = buf.filled().len();
        match Pin::new(&mut this.stream).poll_read(cx, buf) {
            Poll::Ready(result) => {
                if buf.filled().len() > before {
                    this.first_read = None;
                }
                Poll::Ready(result)
            }
            Poll::Pending => {
                let expired = this
                    .first_read
                    .as_mut()
                    .is_some_and(|deadline| deadline.as_mut().poll(cx).is_ready());
                if expired {
                    Poll::Ready(Err(io::Error::new(
                        io::ErrorKind::TimedOut,
                        "no request in time",
                    )))
                } else {
                    Poll::Pending
                }
            }
        }
    }
}

impl<S: AsyncWrite + Unpin> AsyncWrite for Guarded<S> {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.stream).poll_write(cx, buf)
    }

    fn poll_write_vectored(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &[io::IoSlice<'_>],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.stream).poll_write_vectored(cx, bufs)
    }

    fn is_write_vectored(&self) -> bool {
        self.stream.is_write_vectored()
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.stream).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.stream).poll_shutdown(cx)
    }
}

/// Wraps the TLS acceptor: refuses a connection past its client's cap, drops
/// one whose handshake does not finish in time — the ACME acceptor sets no
/// deadline of its own — and hands on the stream [`Guarded`].
#[derive(Clone)]
pub struct GuardedAcceptor<A> {
    inner: A,
    connections: Option<Arc<Connections>>,
    proxies: Arc<Proxies>,
    deadline: Option<Duration>,
}

impl<A> GuardedAcceptor<A> {
    pub fn new(
        inner: A,
        connections: Option<Arc<Connections>>,
        proxies: Arc<Proxies>,
        deadline: Option<Duration>,
    ) -> Self {
        Self {
            inner,
            connections,
            proxies,
            deadline,
        }
    }
}

impl<A> GuardedAcceptor<A> {
    /// Count a connection from `peer` against its client's cap, or name the
    /// client that is at it. A trusted proxy — this machine (cloudflared, a
    /// local proxy) or a CDN edge — carries many clients' connections, so it is
    /// held to no one client's cap.
    fn admit(&self, peer: Option<IpAddr>) -> Result<Option<ConnectionGuard>, IpAddr> {
        let (Some(connections), Some(peer)) = (&self.connections, peer) else {
            return Ok(None);
        };
        if self.proxies.trusts(peer) {
            return Ok(None);
        }
        let client = client_key(peer);
        connections.open(client).map(Some).ok_or(client)
    }
}

impl<A, S> axum_server::accept::Accept<TcpStream, S> for GuardedAcceptor<A>
where
    A: axum_server::accept::Accept<TcpStream, S>,
    A::Future: Send + 'static,
    A::Stream: Send + 'static,
    A::Service: Send + 'static,
{
    type Stream = Guarded<A::Stream>;
    type Service = A::Service;
    type Future = Pin<Box<dyn Future<Output = io::Result<(Self::Stream, A::Service)>> + Send>>;

    fn accept(&self, stream: TcpStream, service: S) -> Self::Future {
        let guard = match self.admit(stream.peer_addr().ok().map(|p| p.ip())) {
            Ok(guard) => guard,
            Err(client) => {
                metrics::counter!("http_connections_refused_total").increment(1);
                tracing::debug!(%client, "connection cap reached");
                return Box::pin(std::future::ready(Err(io::Error::other(
                    "connection cap reached",
                ))));
            }
        };
        let accept = self.inner.accept(stream, service);
        let deadline = self.deadline;
        Box::pin(async move {
            let (stream, service) = match deadline {
                Some(deadline) => tokio::time::timeout(deadline, accept)
                    .await
                    .unwrap_or_else(|_| Err(io::Error::other("TLS handshake timed out")))?,
                None => accept.await?,
            };
            Ok((
                Guarded {
                    stream,
                    first_read: deadline.map(|d| Box::pin(tokio::time::sleep(d))),
                    _guard: guard,
                },
                service,
            ))
        })
    }
}

/// The connection-level limits, which only the public TLS listener applies:
/// behind a reverse proxy every connection is the proxy's, and the proxy sets
/// its own.
pub struct ConnectionLimits {
    pub per_client: Option<Arc<Connections>>,
    /// Peers exempt from the per-client cap.
    pub proxies: Arc<Proxies>,
    /// The deadline for the TLS handshake, and for each request's headers.
    pub header_timeout: Option<Duration>,
}

impl ConnectionLimits {
    /// Put these limits on a server: its acceptor guarded, and a header
    /// deadline — which hyper keeps only when given a timer to keep it with.
    pub fn apply<A>(
        self,
        server: axum_server::Server<SocketAddr, A>,
    ) -> axum_server::Server<SocketAddr, GuardedAcceptor<A>> {
        let deadline = self.header_timeout;
        let mut server = server
            .map(|inner| GuardedAcceptor::new(inner, self.per_client, self.proxies, deadline));
        if let Some(deadline) = deadline {
            server
                .http_builder()
                .http1()
                .timer(hyper_util::rt::TokioTimer::new())
                .header_read_timeout(deadline);
        }
        server
    }
}

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

    fn headers(pairs: &[(&'static str, &str)]) -> HeaderMap {
        let mut h = HeaderMap::new();
        for (k, v) in pairs {
            h.insert(*k, v.parse().unwrap());
        }
        h
    }

    #[test]
    fn parse_rate_reads_rate_and_burst_or_off() {
        assert_eq!(
            parse_rate("20/60"),
            Ok(RateLimit(Some(Rate {
                per_second: 20.0,
                burst: 60.0
            })))
        );
        assert_eq!(
            parse_rate("0.5/5"),
            Ok(RateLimit(Some(Rate {
                per_second: 0.5,
                burst: 5.0
            })))
        );
        assert_eq!(parse_rate("0"), Ok(RateLimit(None)));
        for bad in ["20", "x/5", "5/x", "0/5", "5/0", "-1/5"] {
            assert!(parse_rate(bad).is_err(), "{bad} should not parse");
        }
    }

    #[test]
    fn tier_follows_cost() {
        let none = HeaderMap::new();
        let bibtex = headers(&[("accept", "application/x-bibtex")]);
        let html = headers(&[("accept", "text/html,application/xhtml+xml")]);
        let cases: &[(&str, Option<&str>, &HeaderMap, Option<Tier>)] = &[
            ("/", None, &none, None),
            ("/assets/index-abc.js", None, &none, None),
            ("/api/status", None, &none, None),
            ("/api/search", Some("q=x"), &none, Some(Tier::Expensive)),
            ("/api/handles/10.5555/1", None, &none, Some(Tier::Expensive)),
            ("/dois/10.5555/1", None, &none, Some(Tier::Expensive)),
            ("/api/export", None, &none, Some(Tier::Expensive)),
            (
                "/api/person/0000-0003-1419-2405/works",
                None,
                &none,
                Some(Tier::Lists),
            ),
            ("/api/facets", Some("id=x"), &none, Some(Tier::Lists)),
            (
                "/api/citations",
                Some("id=10.5555/1&limit=10"),
                &none,
                Some(Tier::Lists),
            ),
            (
                "/api/citations/facets",
                Some("id=10.5555/1"),
                &none,
                Some(Tier::Lists),
            ),
            (
                "/api/works",
                Some("local_only=true"),
                &none,
                Some(Tier::Lists),
            ),
            ("/api/works", None, &none, Some(Tier::Expensive)),
            ("/10.5555/1", None, &html, Some(Tier::Pages)),
            ("/10.5555/1", Some("noredirect"), &none, Some(Tier::Pages)),
            ("/10.5555/1", None, &bibtex, Some(Tier::Expensive)),
            (
                "/10.5555/1",
                Some("format=bibtex"),
                &none,
                Some(Tier::Expensive),
            ),
            (
                "/10.5555/1",
                Some("format=commonmeta&local_only=true"),
                &none,
                Some(Tier::Pages),
            ),
            (
                "/orcid/0000-0003-1419-2405",
                Some("format=commonmeta"),
                &none,
                Some(Tier::Expensive),
            ),
        ];
        for (path, query, h, want) in cases {
            assert_eq!(tier(path, *query, h), *want, "{path}?{query:?}");
        }
    }

    #[test]
    fn a_local_proxy_names_the_client_and_no_one_else_does() {
        let proxies = Proxies::default();
        let remote: IpAddr = "203.0.113.7".parse().unwrap();
        let local: IpAddr = "127.0.0.1".parse().unwrap();
        let spoofed = headers(&[("x-real-ip", "10.0.0.1")]);
        assert_eq!(proxies.client(Some(remote), &spoofed), Some(remote));
        assert_eq!(
            proxies.client(Some(local), &spoofed),
            Some("10.0.0.1".parse().unwrap())
        );
        // A proxy appends the peer it saw; what comes before is the client's.
        let forwarded = headers(&[("x-forwarded-for", "1.2.3.4, 198.51.100.2")]);
        assert_eq!(
            proxies.client(Some(local), &forwarded),
            Some("198.51.100.2".parse().unwrap())
        );
        assert_eq!(proxies.client(Some(local), &HeaderMap::new()), None);
    }

    #[test]
    fn a_cdn_edge_names_the_client_in_its_own_header() {
        let proxies = Proxies::new(
            parse_networks("173.245.48.0/20, 2400:cb00::/32").unwrap(),
            Some(parse_header_name("CF-Connecting-IP").unwrap()),
        );
        let edge: IpAddr = "173.245.49.1".parse().unwrap();
        let edge6: IpAddr = "2400:cb00:2049::1".parse().unwrap();
        let client: IpAddr = "198.51.100.9".parse().unwrap();
        assert!(proxies.trusts(edge) && proxies.trusts(edge6));
        assert!(!proxies.trusts(client));

        // The CDN's header wins over anything the client could have sent.
        let named = headers(&[
            ("cf-connecting-ip", "198.51.100.9"),
            ("x-real-ip", "10.6.6.6"),
            ("x-forwarded-for", "10.6.6.6"),
        ]);
        assert_eq!(proxies.client(Some(edge), &named), Some(client));
        assert_eq!(proxies.client(Some(edge6), &named), Some(client));
        // A CDN passes on X-Real-IP as the client sent it, so it is not taken.
        let real_ip = headers(&[("x-real-ip", "10.6.6.6")]);
        assert_eq!(proxies.client(Some(edge), &real_ip), Some(edge));
        // X-Forwarded-For is read from the right, past the trusted hops.
        let forwarded = headers(&[("x-forwarded-for", "10.6.6.6, 198.51.100.9, 173.245.48.7")]);
        let no_header = Proxies::new(parse_networks("173.245.48.0/20").unwrap(), None);
        assert_eq!(no_header.client(Some(edge), &forwarded), Some(client));
        // Outside the trusted networks a peer is the client, whatever it says.
        assert_eq!(proxies.client(Some(client), &named), Some(client));
    }

    #[test]
    fn parse_networks_reads_addresses_and_cidrs() {
        let nets = parse_networks("10.0.0.0/8,192.0.2.1, ::ffff:198.51.100.0/120,2001:db8::/32")
            .unwrap()
            .0;
        assert_eq!(nets.len(), 4);
        assert!(nets[0].contains("10.200.0.1".parse().unwrap()));
        assert!(!nets[0].contains("11.0.0.1".parse().unwrap()));
        assert!(nets[1].contains("192.0.2.1".parse().unwrap()));
        assert!(!nets[1].contains("192.0.2.2".parse().unwrap()));
        assert!(nets[2].contains("198.51.100.77".parse().unwrap()));
        assert!(nets[3].contains("2001:db8:ffff::1".parse().unwrap()));
        assert!(nets[0].contains("::ffff:10.1.2.3".parse().unwrap()));
        assert_eq!(parse_networks(""), Ok(Networks::default()));
        for bad in ["10.0.0.0/33", "2001:db8::/129", "nope", "10.0.0.0/x"] {
            assert!(parse_networks(bad).is_err(), "{bad} should not parse");
        }
    }

    #[test]
    fn client_key_groups_an_ipv6_slash_64() {
        let a: IpAddr = "2a01:4f8:221:2fe5::2".parse().unwrap();
        let b: IpAddr = "2a01:4f8:221:2fe5:ffff::9".parse().unwrap();
        assert_eq!(client_key(a), client_key(b));
        assert_eq!(
            client_key(a),
            "2a01:4f8:221:2fe5::".parse::<IpAddr>().unwrap()
        );
        assert_eq!(
            client_key("::ffff:192.0.2.1".parse().unwrap()),
            "192.0.2.1".parse::<IpAddr>().unwrap()
        );
    }

    fn limiter(rate: &str) -> Arc<Limiter> {
        Limiter::new(Limits {
            pages: parse_rate(rate).unwrap().0,
            lists: None,
            expensive: None,
            max_in_flight: None,
            request_timeout: None,
            proxies: Arc::default(),
        })
    }

    #[test]
    fn a_bucket_spends_its_burst_then_refills_at_the_rate() {
        let limiter = limiter("2/3");
        let client: IpAddr = "192.0.2.1".parse().unwrap();
        let other: IpAddr = "192.0.2.2".parse().unwrap();
        let t0 = Instant::now();
        for _ in 0..3 {
            assert!(limiter.take(Tier::Pages, client, t0).is_ok());
        }
        let wait = limiter.take(Tier::Pages, client, t0).unwrap_err();
        assert_eq!(wait, Duration::from_millis(500));
        // Another client, and another tier, have their own buckets.
        assert!(limiter.take(Tier::Pages, other, t0).is_ok());
        assert!(limiter.take(Tier::Lists, client, t0).is_ok());
        // Half a second at 2/s buys one token back, and only one.
        let t1 = t0 + Duration::from_millis(500);
        assert!(limiter.take(Tier::Pages, client, t1).is_ok());
        assert!(limiter.take(Tier::Pages, client, t1).is_err());
    }

    #[test]
    fn sweep_drops_only_refilled_buckets() {
        let limiter = limiter("1/2");
        let busy: IpAddr = "192.0.2.1".parse().unwrap();
        let idle: IpAddr = "192.0.2.2".parse().unwrap();
        let t0 = Instant::now();
        limiter.take(Tier::Pages, idle, t0).unwrap();
        let t1 = t0 + Duration::from_secs(5);
        limiter.take(Tier::Pages, busy, t1).unwrap();
        limiter.sweep(t1);
        let buckets = limiter.buckets.lock().unwrap();
        assert!(buckets.contains_key(&(Tier::Pages, busy)));
        assert!(!buckets.contains_key(&(Tier::Pages, idle)));
    }

    #[test]
    fn connections_cap_per_client_and_release_on_drop() {
        let connections = Connections::new(2);
        let client: IpAddr = "192.0.2.1".parse().unwrap();
        let first = connections.open(client).unwrap();
        let _second = connections.open(client).unwrap();
        assert!(connections.open(client).is_none());
        assert!(connections.open("192.0.2.2".parse().unwrap()).is_some());
        drop(first);
        assert!(connections.open(client).is_some());
    }

    /// Serve a tiny app behind the guarded acceptor, as the TLS listener does
    /// (minus TLS), and return its address.
    async fn guarded_server(
        max_conns: usize,
        header_timeout: Duration,
        proxies: Proxies,
    ) -> SocketAddr {
        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        listener.set_nonblocking(true).unwrap();
        let addr = listener.local_addr().unwrap();
        let limits = ConnectionLimits {
            per_client: Some(Connections::new(max_conns)),
            proxies: Arc::new(proxies),
            header_timeout: Some(header_timeout),
        };
        let server = limits.apply(axum_server::from_tcp(listener).unwrap());
        let app = axum::Router::new().route("/", axum::routing::get(|| async { "ok" }));
        tokio::spawn(server.serve(app.into_make_service()));
        addr
    }

    /// Whether the server closes `stream` within `within`, without answering.
    async fn closed(stream: &mut TcpStream, within: Duration) -> bool {
        use tokio::io::AsyncReadExt;
        let mut buf = [0u8; 64];
        matches!(
            tokio::time::timeout(within, stream.read(&mut buf)).await,
            Ok(Ok(0)) | Ok(Err(_))
        )
    }

    #[test]
    fn the_cap_holds_clients_but_not_trusted_proxies() {
        let acceptor = GuardedAcceptor::new(
            axum_server::accept::DefaultAcceptor,
            Some(Connections::new(2)),
            Arc::new(Proxies::new(
                parse_networks("173.245.48.0/20").unwrap(),
                None,
            )),
            None,
        );
        let client: IpAddr = "198.51.100.1".parse().unwrap();
        let first = acceptor.admit(Some(client)).unwrap();
        let _second = acceptor.admit(Some(client)).unwrap();
        assert!(first.is_some());
        assert_eq!(acceptor.admit(Some(client)).err(), Some(client));
        // Another client has its own count, and a closed connection frees a slot.
        assert!(
            acceptor
                .admit(Some("198.51.100.2".parse().unwrap()))
                .is_ok()
        );
        drop(first);
        assert!(acceptor.admit(Some(client)).unwrap().is_some());
        // This machine and a CDN edge are counted against no cap at all.
        for proxy in ["127.0.0.1", "::1", "173.245.49.1"] {
            for _ in 0..5 {
                assert!(
                    acceptor
                        .admit(Some(proxy.parse().unwrap()))
                        .unwrap()
                        .is_none()
                );
            }
        }
    }

    #[tokio::test]
    async fn the_guarded_listener_drops_idle_connections() {
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        let addr = guarded_server(2, Duration::from_millis(300), Proxies::default()).await;

        // A client that sends nothing, and one that starts its headers and
        // stalls, are both dropped at the deadline...
        let mut silent = TcpStream::connect(addr).await.unwrap();
        let mut stalled = TcpStream::connect(addr).await.unwrap();
        stalled.write_all(b"GET / HTTP/1.1\r\n").await.unwrap();
        assert!(closed(&mut silent, Duration::from_secs(2)).await);
        assert!(closed(&mut stalled, Duration::from_secs(2)).await);

        // ...and a client that asks is answered, and its connection, left idle,
        // is dropped in turn.
        let mut client = TcpStream::connect(addr).await.unwrap();
        client
            .write_all(b"GET / HTTP/1.1\r\nHost: x\r\n\r\n")
            .await
            .unwrap();
        let mut buf = [0u8; 512];
        let n = client.read(&mut buf).await.unwrap();
        assert!(buf[..n].starts_with(b"HTTP/1.1 200"));
        assert!(closed(&mut client, Duration::from_secs(2)).await);
    }
}