(** X-Forwarded-For parsing and trusted proxy detection.

    This module provides utilities for extracting client IP addresses from
    X-Forwarded-For headers when running behind reverse proxies, with support
    for trusted proxy validation to prevent IP spoofing.

    {2 Background}

    When a web server runs behind a reverse proxy (nginx, HAProxy, cloud load
    balancers), the direct TCP connection comes from the proxy, not the client.
    Proxies add the original client IP to the X-Forwarded-For header.

    {2 Security}

    Without trusted proxy validation, malicious clients can spoof their IP by
    sending a fake X-Forwarded-For header. This module only trusts XFF headers
    from connections originating from configured trusted proxy CIDR ranges.

    {2 References}

    - {{:https://datatracker.ietf.org/doc/html/rfc7239} RFC 7239} - Forwarded
      HTTP Extension (standardized header)
    - X-Forwarded-For is a de-facto standard documented by MDN *)

(** {1 Types} *)

type ip = Ipaddr.t
(** An IP address (IPv4 or IPv6). *)

type prefix = Ipaddr.Prefix.t
(** A CIDR prefix for matching IP ranges. *)

(** {1 CIDR Parsing} *)

let err_invalid_cidr s msg =
  Error (`Msg (Fmt.str "Invalid CIDR '%s': %s" s msg))

let parse_cidr s =
  match Ipaddr.Prefix.of_string s with
  | Ok prefix -> Ok prefix
  | Error (`Msg msg) -> err_invalid_cidr s msg

let parse_cidr_exn s =
  match parse_cidr s with
  | Ok prefix -> prefix
  | Error (`Msg msg) -> invalid_arg msg

(** {1 Trusted Proxy Detection} *)

let ip_in_prefix ip prefix = Ipaddr.Prefix.mem ip prefix

let is_trusted_proxy ip trusted_prefixes =
  List.exists (fun prefix -> ip_in_prefix ip prefix) trusted_prefixes

(** {1 X-Forwarded-For Parsing} *)

let parse_xff xff_value =
  String.split_on_char ',' xff_value
  |> List.filter_map (fun s ->
      let ip_str = String.trim s in
      if ip_str = "" then None
      else
        (* Remove port if present (e.g., "1.2.3.4:5678") *)
        let ip_str_no_port =
          match String.rindex_opt ip_str ':' with
          | Some idx ->
              let before_colon = String.sub ip_str 0 idx in
              (* Check if this looks like IPv6 (has multiple colons) *)
              if String.contains before_colon ':' then ip_str else before_colon
          | None -> ip_str
        in
        Ipaddr.of_string ip_str_no_port |> Result.to_option)

let first_xff_ip xff_value =
  match parse_xff xff_value with first :: _ -> Some first | [] -> None

(** {1 Client IP Extraction} *)

let client_ip ~socket_ip ~xff_header ~trusted_proxies =
  match (socket_ip, trusted_proxies, xff_header) with
  | Some socket_ip, Some prefixes, Some xff
    when is_trusted_proxy socket_ip prefixes -> (
      (* Connection from trusted proxy - extract from X-Forwarded-For *)
      match first_xff_ip xff with
      | Some ip -> Some ip
      | None ->
          Some socket_ip (* Failed to parse XFF - fall back to socket IP *))
  | Some socket_ip, _, _ ->
      (* No trusted proxies, no XFF header, or not from trusted proxy *)
      Some socket_ip
  | None, _, _ -> None

let client_ip_string ~socket_ip ~xff_header ~trusted_proxies =
  match client_ip ~socket_ip ~xff_header ~trusted_proxies with
  | Some ip -> Fmt.str "%a" Ipaddr.pp ip
  | None -> "unknown"

(** {1 Common Trusted Proxy Ranges} *)

let private_ranges =
  List.filter_map
    (fun s -> Result.to_option (parse_cidr s))
    [
      "10.0.0.0/8";
      (* RFC 1918 private *)
      "172.16.0.0/12";
      (* RFC 1918 private *)
      "192.168.0.0/16";
      (* RFC 1918 private *)
      "127.0.0.0/8";
      (* Loopback *)
      "::1/128";
      (* IPv6 loopback *)
      "fc00::/7";
      (* IPv6 unique local *)
      "fe80::/10" (* IPv6 link-local *);
    ]

let cloudflare_ranges =
  List.filter_map
    (fun s -> Result.to_option (parse_cidr s))
    [
      (* Cloudflare IPv4 ranges - see https://www.cloudflare.com/ips-v4 *)
      "173.245.48.0/20";
      "103.21.244.0/22";
      "103.22.200.0/22";
      "103.31.4.0/22";
      "141.101.64.0/18";
      "108.162.192.0/18";
      "190.93.240.0/20";
      "188.114.96.0/20";
      "197.234.240.0/22";
      "198.41.128.0/17";
      "162.158.0.0/15";
      "104.16.0.0/13";
      "104.24.0.0/14";
      "172.64.0.0/13";
      "131.0.72.0/22";
      (* Cloudflare IPv6 ranges - see https://www.cloudflare.com/ips-v6 *)
      "2400:cb00::/32";
      "2606:4700::/32";
      "2803:f800::/32";
      "2405:b500::/32";
      "2405:8100::/32";
      "2a06:98c0::/29";
      "2c0f:f248::/32";
    ]