normogen/backend/src/middleware/client_ip.rs

//! Resolves the originating client IP for each request and exposes it to
//! handlers via the [`ClientIp`] request extension.
//!
//! Resolution order (most-trusted first):
//! 1. `X-Forwarded-For` — first hop (the original client). Normogen is deployed
//!    behind a reverse proxy on Solaria, so this is the primary source.
//! 2. `X-Real-IP` — common single-IP proxy header.
//! 3. The peer address from Axum's `ConnectInfo<SocketAddr>` (direct connection).
//! 4. `"0.0.0.0"` only if none of the above are present.

use axum::{
    extract::ConnectInfo, extract::Request, http::HeaderMap, middleware::Next, response::Response,
};
use std::net::SocketAddr;

/// Header carrying the originating client chain (proxy-populated).
const X_FORWARDED_FOR: &str = "x-forwarded-for";
/// Header some proxies set to the single originating IP.
const X_REAL_IP: &str = "x-real-ip";
/// Fallback when no IP can be determined.
const UNKNOWN_IP: &str = "0.0.0.0";

/// Extractor carrying the resolved client IP, inserted by [`client_ip_middleware`].
#[derive(Debug, Clone)]
pub struct ClientIp(pub String);

impl ClientIp {
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl std::ops::Deref for ClientIp {
    type Target = str;
    fn deref(&self) -> &str {
        &self.0
    }
}

/// Resolve the client IP from request headers, falling back to the peer socket.
///
/// Pure and unit-testable: pass the headers and an optional `ConnectInfo`.
pub fn extract_client_ip(headers: &HeaderMap, conn: Option<&ConnectInfo<SocketAddr>>) -> String {
    // X-Forwarded-For: "client, proxy1, proxy2" — the leftmost entry is the
    // original client. Only validate it looks like an IP; take the first token.
    if let Some(xff) = headers.get(X_FORWARDED_FOR).and_then(|h| h.to_str().ok()) {
        if let Some(first) = xff.split(',').map(str::trim).find(|s| !s.is_empty()) {
            return first.to_string();
        }
    }

    // X-Real-IP: single IP.
    if let Some(ip) = headers.get(X_REAL_IP).and_then(|h| h.to_str().ok()) {
        let ip = ip.trim();
        if !ip.is_empty() {
            return ip.to_string();
        }
    }

    // Direct peer address, when ConnectInfo is wired (see main.rs).
    if let Some(info) = conn {
        return info.0.ip().to_string();
    }

    UNKNOWN_IP.to_string()
}

/// Middleware that computes the client IP once and stashes it in request
/// extensions so handlers can pull it out with `Extension(ClientIp(..))`.
pub async fn client_ip_middleware(mut req: Request, next: Next) -> Response {
    let conn = req.extensions().get::<ConnectInfo<SocketAddr>>();
    let ip = extract_client_ip(req.headers(), conn);
    req.extensions_mut().insert(ClientIp(ip));
    next.run(req).await
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::extract::ConnectInfo;
    use std::net::{IpAddr, Ipv4Addr, SocketAddr};

    fn headers() -> HeaderMap {
        HeaderMap::new()
    }

    fn conn() -> ConnectInfo<SocketAddr> {
        ConnectInfo(SocketAddr::new(
            IpAddr::V4(Ipv4Addr::new(203, 0, 113, 7)),
            12345,
        ))
    }

    #[test]
    fn prefers_x_forwarded_for_first_hop() {
        let mut h = headers();
        h.insert(
            X_FORWARDED_FOR,
            "198.51.100.2, 10.0.0.1, 10.0.0.2".parse().unwrap(),
        );
        assert_eq!(extract_client_ip(&h, Some(&conn())), "198.51.100.2");
    }

    #[test]
    fn falls_back_to_x_real_ip_when_no_xff() {
        let mut h = headers();
        h.insert(X_REAL_IP, "198.51.100.9".parse().unwrap());
        assert_eq!(extract_client_ip(&h, Some(&conn())), "198.51.100.9");
    }

    #[test]
    fn xff_takes_precedence_over_x_real_ip() {
        let mut h = headers();
        h.insert(X_FORWARDED_FOR, "198.51.100.2".parse().unwrap());
        h.insert(X_REAL_IP, "10.0.0.99".parse().unwrap());
        assert_eq!(extract_client_ip(&h, Some(&conn())), "198.51.100.2");
    }

    #[test]
    fn uses_socket_when_no_proxy_headers() {
        assert_eq!(extract_client_ip(&headers(), Some(&conn())), "203.0.113.7");
    }

    #[test]
    fn unknown_when_nothing_available() {
        assert_eq!(extract_client_ip(&headers(), None), UNKNOWN_IP);
    }

    #[test]
    fn ignores_blank_xff_entries() {
        let mut h = headers();
        // Leading whitespace/empty entry should be skipped, not returned blank.
        h.insert(X_FORWARDED_FOR, "  , 10.0.0.5".parse().unwrap());
        assert_eq!(extract_client_ip(&h, None), "10.0.0.5");
    }
}