neiam /oh-my-ex-oh-ircd

defmodule OhMyExOhIrcd.IRC.Proxy do
  @moduledoc """
  Minimal PROXY-protocol v1 + v2 parser.

  When we sit behind a TCP proxy (Traefik, HAProxy, an ELB with PROXY
  enabled, ...) the proxy injects a small header at the start of every
  TCP connection containing the original client IP/port. Without
  parsing that header the IRC pod sees the proxy's pod IP for every
  user — useless for WHOIS, ban masks, K-lines, and rate limiting.

  Returned by `read_header/2`:

      {:ok, src_ip_string, leftover_bytes}
        -- parsed cleanly; any bytes already read past the header are
        -- handed back so the caller can replay them into the normal
        -- handler buffer.

      :no_header
        -- the bytes that arrived don't look like a PROXY header at all.
        -- Up to the caller whether to fail closed or treat as a
        -- direct connection.

      {:error, reason}
        -- the bytes look like an attempt at PROXY but parsing failed
        -- (malformed, truncated, ...). Caller should drop.
  """

  @v2_sig <<0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A>>

  @doc """
  Reads the PROXY header from `socket` synchronously (up to ~107 bytes
  for v1 / variable for v2) and returns the parsed client IP plus any
  bytes already read past the end of the header.
  """
  @spec read_header(any(), pos_integer()) ::
          {:ok, String.t(), binary()} | :no_header | {:error, term()}
  def read_header(socket, timeout_ms \\ 2_000) do
    # Peek at the first byte(s) to decide v1 vs v2. recv/3 returns whatever
    # ThousandIsland's transport has buffered; we may get more than the
    # header in one call, that's fine — we hand the leftover back.
    case ThousandIsland.Socket.recv(socket, 0, timeout_ms) do
      {:ok, data} -> parse(data)
      {:error, reason} -> {:error, reason}
    end
  end

  @doc """
  Parse a PROXY header from an in-memory binary. Public so it can be
  unit-tested without a real socket.
  """
  @spec parse(binary()) :: {:ok, String.t(), binary()} | :no_header | {:error, term()}
  def parse(<<@v2_sig, rest::binary>>), do: parse_v2(rest)
  def parse("PROXY " <> rest), do: parse_v1(rest)
  def parse(_), do: :no_header

  ## v1 (text)

  defp parse_v1(rest) do
    case :binary.split(rest, "\r\n", []) do
      [header, leftover] ->
        case String.split(header, " ") do
          ["TCP4", src_ip, _dst, _sp, _dp] -> {:ok, src_ip, leftover}
          ["TCP6", src_ip, _dst, _sp, _dp] -> {:ok, src_ip, leftover}
          ["UNKNOWN" | _] -> {:ok, "unknown", leftover}
          _ -> {:error, :malformed_v1}
        end

      _ ->
        {:error, :truncated_v1}
    end
  end

  ## v2 (binary)

  defp parse_v2(<<ver_cmd, fam_proto, addr_len::16-big, rest::binary>>) do
    version = Bitwise.bsr(ver_cmd, 4)
    command = Bitwise.band(ver_cmd, 0x0F)

    cond do
      version != 0x2 ->
        {:error, :bad_v2_version}

      byte_size(rest) < addr_len ->
        {:error, :truncated_v2}

      true ->
        <<addrs::binary-size(addr_len), leftover::binary>> = rest
        # command == 0x00 → LOCAL (proxy's own healthcheck) — no real client
        case command do
          0x0 -> {:ok, "local-healthcheck", leftover}
          0x1 -> {:ok, extract_v2_src(fam_proto, addrs), leftover}
          _ -> {:error, :unknown_v2_command}
        end
    end
  end

  defp parse_v2(_), do: {:error, :truncated_v2_header}

  # fam_proto: high nibble = family (1=IPv4, 2=IPv6, 3=UNIX), low nibble = proto.
  defp extract_v2_src(0x11, <<a, b, c, d, _da, _db, _dc, _dd, _sp::16, _dp::16, _::binary>>),
    do: "#{a}.#{b}.#{c}.#{d}"

  defp extract_v2_src(0x12, <<a, b, c, d, _::binary>>), do: "#{a}.#{b}.#{c}.#{d}"

  defp extract_v2_src(0x21, <<src::binary-size(16), _dst::binary-size(16), _sp::16, _dp::16, _::binary>>),
    do: ipv6_to_string(src)

  defp extract_v2_src(0x22, <<src::binary-size(16), _::binary>>), do: ipv6_to_string(src)

  defp extract_v2_src(_, _), do: "unknown"

  defp ipv6_to_string(<<a::16, b::16, c::16, d::16, e::16, f::16, g::16, h::16>>) do
    [a, b, c, d, e, f, g, h]
    |> Enum.map_join(":", &Integer.to_string(&1, 16))
    |> String.downcase()
  end
end