//! axum extractor and middleware to retrieve the client IP.

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

use axum::{
    Extension, RequestPartsExt as _,
    extract::{ConnectInfo, FromRequestParts, Request},
    http::request::Parts,
    middleware::Next,
    response::Response,
};
use tracing::{Span, instrument};

use crate::{accept::proxy_protocol::ProxyProtocolState, extract::trusted_proxy::TrustedProxy};

/// Client IP.
///
/// The [`client_ip_middleware`] must be added to the router before using this extractor,
/// otherwise this will result in requests erroring.
#[derive(Clone, Copy, Debug)]
pub struct ClientIp(pub IpAddr);

impl<S> FromRequestParts<S> for ClientIp
where
    S: Send + Sync,
{
    type Rejection = <Extension<Self> as FromRequestParts<S>>::Rejection;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        Extension::<Self>::from_request_parts(parts, state)
            .await
            .map(|Extension(client_ip)| client_ip)
    }
}

/// Get the client IP from the request.
#[instrument(skip_all)]
async fn extract_client_ip(parts: &mut Parts) -> IpAddr {
    let is_trusted = parts
        .extract::<TrustedProxy>()
        .await
        .unwrap_or(TrustedProxy(false))
        .0;

    if is_trusted {
        if let Ok(ip) = client_ip::rightmost_x_forwarded_for(&parts.headers) {
            return ip;
        }

        if let Ok(ip) = client_ip::x_real_ip(&parts.headers) {
            return ip;
        }

        if let Ok(ip) = client_ip::rightmost_forwarded(&parts.headers) {
            return ip;
        }

        if let Ok(Extension(proxy_protocol_state)) =
            parts.extract::<Extension<ProxyProtocolState>>().await
            && let Some(header) = &proxy_protocol_state.header
            && let Some(addr) = header.proxied_address()
        {
            return addr.source.ip();
        }
    }

    if let Ok(ConnectInfo(addr)) = parts.extract::<ConnectInfo<SocketAddr>>().await {
        addr.ip()
    } else {
        // No connect info means we received a request via a Unix socket, hence localhost
        // as default.
        Ipv6Addr::LOCALHOST.into()
    }
}

/// Middleware required by the [`ClientIp`] extractor.
///
/// Use with [`axum::middleware::from_fn`].
pub async fn client_ip_middleware(request: Request, next: Next) -> Response {
    let (mut parts, body) = request.into_parts();

    let client_ip = extract_client_ip(&mut parts).await;
    Span::current().record("remote", client_ip.to_string());
    parts.extensions.insert::<ClientIp>(ClientIp(client_ip));

    let request = Request::from_parts(parts, body);

    next.run(request).await
}

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

    use axum::{body::Body, http::Request};

    use super::*;

    #[tokio::test]
    async fn x_forwarded_for_trusted() {
        let (mut parts, _) = Request::builder()
            .uri("http://example.com/path")
            .header("x-forwarded-for", "192.0.2.51, 192.0.2.42")
            .extension(TrustedProxy(true))
            .body(Body::empty())
            .expect("Failed to create request")
            .into_parts();

        let client_ip = extract_client_ip(&mut parts).await;

        assert_eq!(client_ip, Ipv4Addr::new(192, 0, 2, 42),);
    }

    #[tokio::test]
    async fn x_real_ip_trusted() {
        let (mut parts, _) = Request::builder()
            .uri("http://example.com/path")
            .header("x-real-ip", "192.0.2.42")
            .extension(TrustedProxy(true))
            .body(Body::empty())
            .expect("Failed to create request")
            .into_parts();

        let client_ip = extract_client_ip(&mut parts).await;

        assert_eq!(client_ip, Ipv4Addr::new(192, 0, 2, 42),);
    }

    #[tokio::test]
    async fn forwarded_header_trusted() {
        let (mut parts, _) = Request::builder()
            .uri("http://example.com/path")
            .header("forwarded", "for=192.0.2.42")
            .extension(TrustedProxy(true))
            .body(Body::empty())
            .expect("Failed to create request")
            .into_parts();

        let client_ip = extract_client_ip(&mut parts).await;

        assert_eq!(client_ip, Ipv4Addr::new(192, 0, 2, 42),);
    }

    #[tokio::test]
    async fn from_connect_info() {
        let connect_addr: SocketAddr = "192.0.2.42:34932"
            .parse()
            .expect("Failed to parse socket address");
        let (mut parts, _) = Request::builder()
            .uri("http://example.com/path")
            .extension(ConnectInfo(connect_addr))
            .extension(TrustedProxy(false))
            .body(Body::empty())
            .expect("Failed to create request")
            .into_parts();

        let client_ip = extract_client_ip(&mut parts).await;

        assert_eq!(client_ip, Ipv4Addr::new(192, 0, 2, 42),);
    }

    #[tokio::test]
    async fn headers_untrusted() {
        let (mut parts, _) = Request::builder()
            .uri("http://example.com/path")
            .header("x-forwarded-for", "192.0.2.42")
            .extension(TrustedProxy(false))
            .body(Body::empty())
            .expect("Failed to create request")
            .into_parts();

        let client_ip = extract_client_ip(&mut parts).await;

        assert_eq!(client_ip, Ipv6Addr::LOCALHOST);
    }

    #[tokio::test]
    async fn priority_order() {
        // Test that X-Forwarded-For takes priority over other headers when trusted
        let (mut parts, _) = Request::builder()
            .uri("http://example.com/path")
            .header("x-forwarded-for", "192.0.2.1")
            .header("x-real-ip", "192.0.2.2")
            .header("forwarded", "for=192.0.2.3")
            .extension(TrustedProxy(true))
            .body(Body::empty())
            .expect("Failed to create request")
            .into_parts();

        let client_ip = extract_client_ip(&mut parts).await;

        assert_eq!(client_ip, Ipv4Addr::new(192, 0, 2, 1),);
    }

    #[tokio::test]
    async fn no_ip_found() {
        let (mut parts, _) = Request::builder()
            .uri("http://example.com/path")
            .body(Body::empty())
            .expect("Failed to create request")
            .into_parts();

        let client_ip = extract_client_ip(&mut parts).await;

        assert_eq!(client_ip, Ipv6Addr::LOCALHOST);
    }

    #[tokio::test]
    async fn ipv6() {
        let (mut parts, _) = Request::builder()
            .uri("http://example.com/path")
            .header("x-forwarded-for", "2001:db8::42")
            .extension(TrustedProxy(true))
            .body(Body::empty())
            .expect("Failed to create request")
            .into_parts();

        let client_ip = extract_client_ip(&mut parts).await;

        assert_eq!(client_ip, Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0x42),);
    }

    #[tokio::test]
    async fn multiple_x_forwarded_for() {
        let (mut parts, _) = Request::builder()
            .uri("http://example.com/path")
            .header("x-forwarded-for", "192.0.2.1, 192.0.2.2, 192.0.2.3")
            .extension(TrustedProxy(true))
            .body(Body::empty())
            .expect("Failed to create request")
            .into_parts();

        let client_ip = extract_client_ip(&mut parts).await;

        assert_eq!(client_ip, Ipv4Addr::new(192, 0, 2, 3),);
    }
}
