//! Client IP for throttling.
//!
//! The TCP peer is the client unless it is listed in `config.trusted_proxies`. A trusted peer's
//! `X-Forwarded-For` is read right to left, skipping trusted proxies, and the first untrusted
//! address is the client. Any unusable header falls back to the peer. IPv6 is bucketed by /64, so
//! one client holding a /64 cannot multiply its attempt budget.

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

use ipnet::IpNet;
use topcoat::context::{Cx, app_context};
use topcoat::router::request::{headers, remote_addr};

use crate::state::AppState;

const FORWARDED_FOR: &str = "x-forwarded-for";

/// The client IP for this request, bucketed for throttling. Unix-socket requests have no peer and
/// map to the unspecified address.
pub fn client_ip(cx: &Cx) -> IpAddr {
    let trusted = &app_context::<AppState>(cx).config.trusted_proxies;
    let peer = remote_addr(cx).map(|addr| addr.ip());
    let forwarded = headers(cx)
        .get_all(FORWARDED_FOR)
        .iter()
        .filter_map(|value| value.to_str().ok())
        .collect::<Vec<_>>()
        .join(",");
    resolve(peer, &forwarded, trusted)
}

/// Pure form of [`client_ip`]: `forwarded` is the combined `X-Forwarded-For` value, or empty.
pub fn resolve(peer: Option<IpAddr>, forwarded: &str, trusted: &[IpNet]) -> IpAddr {
    let Some(peer) = peer.map(normalize) else {
        return bucket(IpAddr::V4(Ipv4Addr::UNSPECIFIED));
    };
    if !is_trusted(peer, trusted) {
        return bucket(peer);
    }
    for entry in forwarded.rsplit(',') {
        let Ok(ip) = entry.trim().parse::<IpAddr>() else {
            return bucket(peer);
        };
        let ip = normalize(ip);
        if !is_trusted(ip, trusted) {
            return bucket(ip);
        }
    }
    bucket(peer)
}

fn is_trusted(ip: IpAddr, trusted: &[IpNet]) -> bool {
    trusted.iter().any(|net| net.contains(&ip))
}

/// IPv4-mapped IPv6 addresses are reported as IPv4.
fn normalize(ip: IpAddr) -> IpAddr {
    match ip {
        IpAddr::V6(v6) => v6.to_ipv4_mapped().map_or(IpAddr::V6(v6), IpAddr::V4),
        IpAddr::V4(_) => ip,
    }
}

/// Clears the host part of IPv6 addresses so every address in a /64 shares one bucket.
fn bucket(ip: IpAddr) -> IpAddr {
    match ip {
        IpAddr::V4(_) => ip,
        IpAddr::V6(v6) => {
            let mut octets = v6.octets();
            octets[8..].fill(0);
            IpAddr::V6(Ipv6Addr::from(octets))
        }
    }
}

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

    fn ip(text: &str) -> IpAddr {
        text.parse().unwrap()
    }

    fn net(text: &str) -> IpNet {
        text.parse().unwrap()
    }

    #[test]
    fn forwarded_for_is_ignored_from_untrusted_peer() {
        let trusted = [net("10.0.0.0/8")];
        let got = resolve(Some(ip("203.0.113.9")), "198.51.100.7", &trusted);
        assert_eq!(got, ip("203.0.113.9"));
    }

    #[test]
    fn forwarded_for_is_honored_from_trusted_peer() {
        let trusted = [net("127.0.0.1/32")];
        let got = resolve(Some(ip("127.0.0.1")), "198.51.100.7", &trusted);
        assert_eq!(got, ip("198.51.100.7"));
    }

    #[test]
    fn right_most_untrusted_entry_wins() {
        let trusted = [net("127.0.0.1/32"), net("10.0.0.0/8")];
        // The client may prepend anything; the proxy appends the address it saw.
        let got = resolve(
            Some(ip("127.0.0.1")),
            "1.2.3.4, 198.51.100.7, 10.0.0.5",
            &trusted,
        );
        assert_eq!(got, ip("198.51.100.7"));
    }

    #[test]
    fn unusable_header_falls_back_to_peer() {
        let trusted = [net("127.0.0.1/32")];
        assert_eq!(
            resolve(Some(ip("127.0.0.1")), "", &trusted),
            ip("127.0.0.1")
        );
        assert_eq!(
            resolve(Some(ip("127.0.0.1")), "garbage", &trusted),
            ip("127.0.0.1")
        );
        assert_eq!(
            resolve(Some(ip("127.0.0.1")), "127.0.0.1", &trusted),
            ip("127.0.0.1")
        );
    }

    #[test]
    fn ipv6_is_bucketed_by_slash_64() {
        let trusted: [IpNet; 0] = [];
        let a = resolve(Some(ip("2001:db8:1:2:aaaa::1")), "", &trusted);
        let b = resolve(Some(ip("2001:db8:1:2:bbbb::9")), "", &trusted);
        let c = resolve(Some(ip("2001:db8:1:3::1")), "", &trusted);
        assert_eq!(a, b);
        assert_ne!(a, c);
        assert_eq!(a, ip("2001:db8:1:2::"));
    }

    #[test]
    fn ipv4_mapped_peer_is_reported_as_ipv4() {
        let trusted: [IpNet; 0] = [];
        assert_eq!(
            resolve(Some(ip("::ffff:192.0.2.5")), "", &trusted),
            ip("192.0.2.5")
        );
    }

    #[test]
    fn unknown_peer_maps_to_unspecified() {
        let trusted: [IpNet; 0] = [];
        assert_eq!(resolve(None, "1.2.3.4", &trusted), ip("0.0.0.0"));
    }
}