//! Rate-limit key extraction that never trusts an attacker-controlled
//! `X-Forwarded-For`.
//!
//! `tower_governor`'s `SmartIpKeyExtractor` reads `X-Forwarded-For` from *every*
//! request, so any direct client can mint a fresh rate-limit bucket per request
//! by rotating the header, defeating every limiter. [`TrustedProxyKeyExtractor`]
//! closes that: it honors forwarding headers only when the request's *direct
//! peer* is a configured trusted proxy, and otherwise keys on the peer address
//! and ignores the headers entirely. There is deliberately **no `Default`**, the
//! set of trusted proxies must be supplied explicitly, so "trust this header"
//! can never be the silent default that the prior extractor made it.
//!
//! # Two hops, not one
//!
//! Rightmost-`X-Forwarded-For` identifies the real client only when exactly one
//! trusted proxy sits in front of the app. Prod has two: Cloudflare, then the
//! on-host Caddy (`forums.makenot.work` in `MNW/server/deploy/Caddyfile`).
//! Cloudflare appends the real client to `X-Forwarded-For`, then Caddy appends
//! the Cloudflare edge address, so the rightmost entry is a *Cloudflare* IP.
//! Keying on it collapses every visitor routed through the same edge into one
//! bucket: per-IP limits stop isolating abusers, and one abuser exhausts the
//! bucket shared by everyone behind that edge.
//!
//! So a trusted peer's `CF-Connecting-IP` wins over `X-Forwarded-For`.
//! Cloudflare sets that header to the true client and it survives exactly one
//! hop, and the vhost pins TLS client auth to the Cloudflare origin-pull CA
//! (`import cloudflare_tls`), so nothing but Cloudflare can reach the origin to
//! forge one. Rightmost-XFF stays as the fallback for a single-proxy
//! deployment (a tailnet-direct staging box) that has no Cloudflare in front.

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

use axum::http::{HeaderMap, Request};
use tower_governor::errors::GovernorError;
use tower_governor::key_extractor::KeyExtractor;

/// A [`KeyExtractor`] that keys on the real client IP, trusting
/// `CF-Connecting-IP` / `X-Forwarded-For` only when the direct peer is a known
/// proxy.
///
/// Construction requires the trusted-proxy set; there is no `Default`. An empty
/// set trusts no proxy, every request keys on its direct peer, which is safe
/// but collapses to a single global bucket behind a reverse proxy, so a proxied
/// deployment must list its proxy (loopback for an on-host Caddy).
#[derive(Clone, Debug)]
pub struct TrustedProxyKeyExtractor {
    trusted_proxies: Arc<[IpAddr]>,
}

impl TrustedProxyKeyExtractor {
    /// Build with the explicit set of proxy IPs whose `X-Forwarded-For` is
    /// trusted. Trusting the header is a decision the caller must make here.
    pub fn new(trusted_proxies: impl Into<Arc<[IpAddr]>>) -> Self {
        Self {
            trusted_proxies: trusted_proxies.into(),
        }
    }

    fn peer_ip<T>(req: &Request<T>) -> Option<IpAddr> {
        req.extensions()
            .get::<axum::extract::ConnectInfo<SocketAddr>>()
            .map(|ci| ci.0.ip())
    }
}

impl KeyExtractor for TrustedProxyKeyExtractor {
    type Key = IpAddr;

    // `name`/`key_name` are only trait methods under tower_governor's `tracing`
    // feature, which this crate does not enable; omitted deliberately.

    fn extract<T>(&self, req: &Request<T>) -> Result<Self::Key, GovernorError> {
        let peer = Self::peer_ip(req).ok_or(GovernorError::UnableToExtractKey)?;
        // Only a trusted proxy's forwarding headers are believed. A direct
        // (untrusted) peer keys on its own socket address, so it cannot spoof
        // another key.
        if self.trusted_proxies.contains(&peer) {
            // CF-Connecting-IP first: under Cloudflare + Caddy it names the true
            // client, while rightmost-XFF names the Cloudflare edge. See the
            // module docs.
            if let Some(client) = cf_connecting_ip(req.headers()) {
                return Ok(client);
            }
            if let Some(client) = rightmost_xff(req.headers()) {
                return Ok(client);
            }
        }
        Ok(peer)
    }
}

/// The client address Cloudflare observed, from `CF-Connecting-IP`.
///
/// Cloudflare sets exactly one address here (not a chain), replacing any value
/// the client sent, so there is no leftmost/rightmost question to get wrong.
fn cf_connecting_ip(headers: &HeaderMap) -> Option<IpAddr> {
    headers
        .get("cf-connecting-ip")
        .and_then(|hv| hv.to_str().ok())
        .and_then(|s| s.trim().parse::<IpAddr>().ok())
}

/// The rightmost parseable `X-Forwarded-For` entry, the address the trusted
/// proxy observed connecting to it. Taking the *rightmost* entry (not the
/// leftmost) means a client-injected hop to the left of the proxy's appended
/// entry is ignored, so spoofing fails even when the proxy appends rather than
/// overwrites.
fn rightmost_xff(headers: &HeaderMap) -> Option<IpAddr> {
    headers
        .get("x-forwarded-for")
        .and_then(|hv| hv.to_str().ok())
        .and_then(|s| s.rsplit(',').find_map(|p| p.trim().parse::<IpAddr>().ok()))
}

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

    fn req(peer: Option<&str>, xff: Option<&str>) -> Request<()> {
        req_cf(peer, xff, None)
    }

    fn req_cf(peer: Option<&str>, xff: Option<&str>, cf: Option<&str>) -> Request<()> {
        let mut b = Request::builder();
        if let Some(xff) = xff {
            b = b.header("x-forwarded-for", xff);
        }
        if let Some(cf) = cf {
            b = b.header("cf-connecting-ip", cf);
        }
        let mut req = b.body(()).unwrap();
        if let Some(peer) = peer {
            let addr: SocketAddr = format!("{peer}:1234").parse().unwrap();
            req.extensions_mut()
                .insert(axum::extract::ConnectInfo(addr));
        }
        req
    }

    fn extractor(trusted: &[&str]) -> TrustedProxyKeyExtractor {
        let ips: Vec<IpAddr> = trusted.iter().map(|s| s.parse().unwrap()).collect();
        TrustedProxyKeyExtractor::new(ips)
    }

    #[test]
    fn trusted_peer_uses_forwarded_client() {
        // Peer is the trusted proxy → believe XFF (rightmost = real client).
        let key = extractor(&["127.0.0.1"])
            .extract(&req(Some("127.0.0.1"), Some("203.0.113.7")))
            .unwrap();
        assert_eq!(key, "203.0.113.7".parse::<IpAddr>().unwrap());
    }

    #[test]
    fn untrusted_peer_ignores_forwarded_header() {
        // A direct client spoofing XFF gets keyed on its OWN address, not the
        // forged one, the whole point of the fix.
        let key = extractor(&["127.0.0.1"])
            .extract(&req(Some("198.51.100.9"), Some("203.0.113.7")))
            .unwrap();
        assert_eq!(key, "198.51.100.9".parse::<IpAddr>().unwrap());
    }

    #[test]
    fn trusted_peer_no_header_falls_back_to_peer() {
        let key = extractor(&["127.0.0.1"])
            .extract(&req(Some("127.0.0.1"), None))
            .unwrap();
        assert_eq!(key, "127.0.0.1".parse::<IpAddr>().unwrap());
    }

    #[test]
    fn rightmost_entry_defeats_injected_left_hop() {
        // Attacker sends "1.2.3.4", proxy appends the real client on the right.
        // Rightmost wins, so the injected left entry is ignored.
        let key = extractor(&["127.0.0.1"])
            .extract(&req(Some("127.0.0.1"), Some("1.2.3.4, 203.0.113.7")))
            .unwrap();
        assert_eq!(key, "203.0.113.7".parse::<IpAddr>().unwrap());
    }

    #[test]
    fn trusted_peer_garbage_header_falls_back_to_peer() {
        let key = extractor(&["127.0.0.1"])
            .extract(&req(Some("127.0.0.1"), Some("not-an-ip")))
            .unwrap();
        assert_eq!(key, "127.0.0.1".parse::<IpAddr>().unwrap());
    }

    #[test]
    fn empty_trusted_set_always_keys_on_peer() {
        let key = extractor(&[])
            .extract(&req(Some("203.0.113.7"), Some("10.0.0.1")))
            .unwrap();
        assert_eq!(key, "203.0.113.7".parse::<IpAddr>().unwrap());
    }

    #[test]
    fn cf_connecting_ip_beats_xff_for_a_trusted_peer() {
        // The prod shape: Cloudflare appended the real client to XFF, then Caddy
        // appended the Cloudflare edge. Rightmost-XFF would key on the edge and
        // bucket every visitor behind it together; CF-Connecting-IP is the client.
        let key = extractor(&["127.0.0.1"])
            .extract(&req_cf(
                Some("127.0.0.1"),
                Some("203.0.113.7, 172.68.1.1"),
                Some("203.0.113.7"),
            ))
            .unwrap();
        assert_eq!(key, "203.0.113.7".parse::<IpAddr>().unwrap());
    }

    #[test]
    fn untrusted_peer_ignores_cf_connecting_ip() {
        // A direct client forging the Cloudflare header still keys on itself.
        let key = extractor(&["127.0.0.1"])
            .extract(&req_cf(Some("198.51.100.9"), None, Some("203.0.113.7")))
            .unwrap();
        assert_eq!(key, "198.51.100.9".parse::<IpAddr>().unwrap());
    }

    #[test]
    fn trusted_peer_garbage_cf_header_falls_back_to_xff() {
        // An unparseable CF header must not shadow a usable XFF chain.
        let key = extractor(&["127.0.0.1"])
            .extract(&req_cf(
                Some("127.0.0.1"),
                Some("203.0.113.7"),
                Some("not-an-ip"),
            ))
            .unwrap();
        assert_eq!(key, "203.0.113.7".parse::<IpAddr>().unwrap());
    }

    #[test]
    fn trusted_peer_cf_header_alone_is_enough() {
        // Caddy can be configured to strip XFF; CF-Connecting-IP still keys.
        let key = extractor(&["127.0.0.1"])
            .extract(&req_cf(Some("127.0.0.1"), None, Some("203.0.113.7")))
            .unwrap();
        assert_eq!(key, "203.0.113.7".parse::<IpAddr>().unwrap());
    }

    #[test]
    fn no_peer_errors() {
        assert!(
            extractor(&["127.0.0.1"])
                .extract(&req(None, Some("203.0.113.7")))
                .is_err()
        );
    }
}