//! Client address resolution behind trusted reverse proxies.

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

use axum::{
    extract::{ConnectInfo, FromRequestParts},
    http::{HeaderMap, request::Parts},
};

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

/// Parse `a.b.c.d`, `a.b.c.d/nn`, `::1` or `fd00::/8`.
pub fn parse_cidr(value: &str) -> Option<(IpAddr, u8)> {
    let value = value.trim();
    let (addr, bits) = match value.split_once('/') {
        Some((addr, bits)) => (addr, Some(bits.parse::<u8>().ok()?)),
        None => (value, None),
    };
    let ip: IpAddr = addr.parse().ok()?;
    let max = if ip.is_ipv4() { 32 } else { 128 };
    let bits = bits.unwrap_or(max);
    if bits > max {
        return None;
    }
    Some((ip, bits))
}

fn normalize(ip: IpAddr) -> IpAddr {
    match ip {
        IpAddr::V6(v6) => v6.to_ipv4_mapped().map(IpAddr::V4).unwrap_or(ip),
        other => other,
    }
}

pub fn cidr_contains(network: (IpAddr, u8), ip: IpAddr) -> bool {
    let ip = normalize(ip);
    match (normalize(network.0), ip) {
        (IpAddr::V4(net), IpAddr::V4(ip)) => {
            let bits = u32::from(network.1.min(32));
            if bits == 0 {
                return true;
            }
            let mask = u32::MAX << (32 - bits);
            (u32::from(net) & mask) == (u32::from(ip) & mask)
        }
        (IpAddr::V6(net), IpAddr::V6(ip)) => {
            let bits = u32::from(network.1.min(128));
            if bits == 0 {
                return true;
            }
            let mask = u128::MAX << (128 - bits);
            (u128::from(net) & mask) == (u128::from(ip) & mask)
        }
        _ => false,
    }
}

fn is_trusted(proxies: &[String], ip: IpAddr) -> bool {
    proxies
        .iter()
        .filter_map(|p| parse_cidr(p))
        .any(|net| cidr_contains(net, ip))
}

/// Resolve the real client address. Forwarding headers are honoured only when
/// the TCP peer is a configured trusted proxy; the right-most untrusted hop of
/// `X-Forwarded-For` wins, so a client cannot spoof its address by prepending
/// values.
pub fn resolve_client_ip(proxies: &[String], peer: IpAddr, headers: &HeaderMap) -> IpAddr {
    let peer = normalize(peer);
    if proxies.is_empty() || !is_trusted(proxies, peer) {
        return peer;
    }
    if let Some(xff) = headers.get("x-forwarded-for").and_then(|v| v.to_str().ok()) {
        let hops: Vec<IpAddr> = xff
            .split(',')
            .filter_map(|part| part.trim().parse::<IpAddr>().ok())
            .map(normalize)
            .collect();
        for hop in hops.iter().rev() {
            if !is_trusted(proxies, *hop) {
                return *hop;
            }
        }
        if let Some(first) = hops.first() {
            return *first;
        }
    }
    if let Some(real) = headers
        .get("x-real-ip")
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.trim().parse::<IpAddr>().ok())
    {
        return normalize(real);
    }
    peer
}

/// Extractor yielding the effective client IP.
pub struct ClientIp(pub IpAddr);

impl FromRequestParts<Arc<AppState>> for ClientIp {
    type Rejection = AppError;

    async fn from_request_parts(
        parts: &mut Parts,
        state: &Arc<AppState>,
    ) -> Result<Self, Self::Rejection> {
        let peer = parts
            .extensions
            .get::<ConnectInfo<SocketAddr>>()
            .map(|ConnectInfo(addr)| addr.ip())
            .unwrap_or(IpAddr::from([127, 0, 0, 1]));
        Ok(ClientIp(resolve_client_ip(
            &state.config.security.trusted_proxies,
            peer,
            &parts.headers,
        )))
    }
}

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

    #[test]
    fn cidr_matching() {
        let net = parse_cidr("10.0.0.0/8").unwrap();
        assert!(cidr_contains(net, "10.1.2.3".parse().unwrap()));
        assert!(!cidr_contains(net, "11.1.2.3".parse().unwrap()));
        let v6 = parse_cidr("fd00::/8").unwrap();
        assert!(cidr_contains(v6, "fd12::1".parse().unwrap()));
        assert!(parse_cidr("1.2.3.4/33").is_none());
        assert!(cidr_contains(
            parse_cidr("127.0.0.1").unwrap(),
            "::ffff:127.0.0.1".parse().unwrap()
        ));
    }

    #[test]
    fn forwarded_for_only_from_trusted_proxy() {
        let proxies = vec!["127.0.0.1".to_string()];
        let mut headers = HeaderMap::new();
        headers.insert("x-forwarded-for", "6.6.6.6, 203.0.113.9".parse().unwrap());
        assert_eq!(
            resolve_client_ip(&proxies, "127.0.0.1".parse().unwrap(), &headers),
            "203.0.113.9".parse::<IpAddr>().unwrap()
        );
        assert_eq!(
            resolve_client_ip(&proxies, "198.51.100.1".parse().unwrap(), &headers),
            "198.51.100.1".parse::<IpAddr>().unwrap()
        );
        assert_eq!(
            resolve_client_ip(&[], "127.0.0.1".parse().unwrap(), &headers),
            "127.0.0.1".parse::<IpAddr>().unwrap()
        );
    }
}