crank/apps/admin-api/src/rate_limit.rs

use std::net::{IpAddr, SocketAddr};

use axum::{
    extract::{ConnectInfo, Request, State},
    http::HeaderMap,
    middleware::Next,
    response::Response,
};
use crank_runtime::{RateLimitCheckError, RateLimitRejection};

use crate::{error::ApiError, state::AppState};

#[derive(Clone, Debug)]
pub struct ClientIdentityBucket(pub String);

pub async fn apply_api_rate_limit(
    State(state): State<AppState>,
    mut request: Request,
    next: Next,
) -> Result<Response, ApiError> {
    let peer_ip = request
        .extensions()
        .get::<ConnectInfo<SocketAddr>>()
        .map(|ConnectInfo(address)| address.ip());
    let key = client_rate_limit_key(
        request.headers(),
        request.uri().path(),
        peer_ip,
        &state.trusted_proxy_ips,
    );
    if let Err(error) = state.api_rate_limiter.check(&key).await {
        return match error {
            RateLimitCheckError::Rejected(rejection) => Err(ApiError::rate_limited_with_context(
                "request rate limit exceeded",
                rejection_context(rejection),
            )),
            RateLimitCheckError::StoreUnavailable => {
                Err(ApiError::internal("rate limit service unavailable"))
            }
        };
    }

    request.extensions_mut().insert(ClientIdentityBucket(key));

    Ok(next.run(request).await)
}

fn rejection_context(rejection: RateLimitRejection) -> serde_json::Value {
    serde_json::json!({
        "retry_after_ms": rejection.retry_after_ms,
    })
}

pub fn client_rate_limit_key(
    headers: &HeaderMap,
    path: &str,
    peer_ip: Option<IpAddr>,
    trusted_proxy_ips: &[IpAddr],
) -> String {
    if peer_ip.is_some_and(|peer_ip| trusted_proxy_ips.contains(&peer_ip))
        && let Some(client_ip) = forwarded_client_ip(headers)
    {
        return format!("ip:{client_ip}");
    }

    if let Some(peer_ip) = peer_ip {
        return format!("ip:{peer_ip}");
    }

    format!("anonymous:{path}")
}

/// Resolves the client IP from proxy headers after the immediate peer was
/// matched against the trusted-proxy allowlist.
///
/// `X-Real-IP` is preferred because a trusted proxy (e.g. nginx) sets it to the
/// real peer address. For `X-Forwarded-For` the proxy *appends* the observed
/// peer, so the last entry is the trustworthy hop; taking the first entry (as
/// naive implementations do) would let a client spoof its address by sending a
/// pre-populated header.
fn forwarded_client_ip(headers: &HeaderMap) -> Option<IpAddr> {
    if let Some(real_ip) = header_value(headers, "x-real-ip").and_then(parse_ip) {
        return Some(real_ip);
    }

    header_value(headers, "x-forwarded-for")?
        .split(',')
        .map(str::trim)
        .rfind(|value| !value.is_empty())
        .and_then(parse_ip)
}

fn header_value<'a>(headers: &'a HeaderMap, name: &'static str) -> Option<&'a str> {
    headers.get(name)?.to_str().ok().map(str::trim)
}

fn parse_ip(value: &str) -> Option<IpAddr> {
    value.parse().ok()
}

#[cfg(test)]
mod tests {
    use std::net::{IpAddr, Ipv4Addr};

    use axum::http::{HeaderMap, HeaderValue, header::COOKIE};

    use super::client_rate_limit_key;

    fn peer() -> Option<IpAddr> {
        Some(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 7)))
    }

    #[test]
    fn unverified_session_cookie_cannot_change_client_key() {
        let mut headers = HeaderMap::new();
        headers.insert(
            COOKIE,
            HeaderValue::from_static("theme=dark; crank_session=sess_123.secret_456"),
        );
        headers.insert("x-forwarded-for", HeaderValue::from_static("10.0.0.5"));

        assert_eq!(
            client_rate_limit_key(&headers, "/api/auth/login", peer(), &[peer().unwrap()]),
            "ip:10.0.0.5"
        );
    }

    #[test]
    fn ignores_forwarded_headers_when_untrusted() {
        let mut headers = HeaderMap::new();
        headers.insert("x-forwarded-for", HeaderValue::from_static("10.0.0.5"));
        headers.insert("x-real-ip", HeaderValue::from_static("10.0.0.9"));

        assert_eq!(
            client_rate_limit_key(&headers, "/api/auth/login", peer(), &[]),
            "ip:203.0.113.7"
        );
    }

    #[test]
    fn ignores_forwarded_headers_from_unlisted_peer() {
        let mut headers = HeaderMap::new();
        headers.insert("x-real-ip", HeaderValue::from_static("10.0.0.9"));

        assert_eq!(
            client_rate_limit_key(
                &headers,
                "/api/auth/login",
                peer(),
                &[IpAddr::V4(Ipv4Addr::new(198, 51, 100, 10))]
            ),
            "ip:203.0.113.7"
        );
    }

    #[test]
    fn prefers_real_ip_when_trusted() {
        let mut headers = HeaderMap::new();
        headers.insert(
            "x-forwarded-for",
            HeaderValue::from_static("1.2.3.4, 10.0.0.6"),
        );
        headers.insert("x-real-ip", HeaderValue::from_static("10.0.0.9"));

        assert_eq!(
            client_rate_limit_key(&headers, "/api/auth/login", peer(), &[peer().unwrap()]),
            "ip:10.0.0.9"
        );
    }

    #[test]
    fn uses_last_forwarded_hop_when_trusted() {
        // A client can prepend spoofed entries; the trusted proxy appends the
        // real peer, so the last entry is authoritative.
        let mut headers = HeaderMap::new();
        headers.insert(
            "x-forwarded-for",
            HeaderValue::from_static("1.2.3.4, 10.0.0.6"),
        );

        assert_eq!(
            client_rate_limit_key(&headers, "/api/auth/login", peer(), &[peer().unwrap()]),
            "ip:10.0.0.6"
        );
    }

    #[test]
    fn falls_back_to_peer_ip_without_headers() {
        let headers = HeaderMap::new();

        assert_eq!(
            client_rate_limit_key(&headers, "/api/auth/login", peer(), &[peer().unwrap()]),
            "ip:203.0.113.7"
        );
    }

    #[test]
    fn ignores_invalid_forwarded_ip_values() {
        let mut headers = HeaderMap::new();
        headers.insert("x-real-ip", HeaderValue::from_static("not-an-ip"));
        headers.insert(
            "x-forwarded-for",
            HeaderValue::from_static("198.51.100.8, also-not-an-ip"),
        );

        assert_eq!(
            client_rate_limit_key(&headers, "/api/auth/login", peer(), &[peer().unwrap()]),
            "ip:203.0.113.7"
        );
    }

    #[test]
    fn falls_back_to_path_without_peer() {
        let headers = HeaderMap::new();

        assert_eq!(
            client_rate_limit_key(&headers, "/api/auth/login", None, &[]),
            "anonymous:/api/auth/login"
        );
    }
}