//! 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 `X-Forwarded-For` only when the request's *direct
//! peer* is a configured trusted proxy, and otherwise keys on the peer address
//! and ignores the header 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.

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
/// `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 X-Forwarded-For is believed. A direct (untrusted)
        // peer keys on its own socket address, so it cannot spoof another key.
        if self.trusted_proxies.contains(&peer)
            && let Some(client) = rightmost_xff(req.headers())
        {
            return Ok(client);
        }
        Ok(peer)
    }
}

/// 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<()> {
        let mut b = Request::builder();
        if let Some(xff) = xff {
            b = b.header("x-forwarded-for", xff);
        }
        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 no_peer_errors() {
        assert!(
            extractor(&["127.0.0.1"])
                .extract(&req(None, Some("203.0.113.7")))
                .is_err()
        );
    }
}