neiam /gitgud

defmodule GitGud.Federation.RateLimiter do
  @moduledoc """
  Fixed-window counters for inbox deliveries, keyed by source host.

  Backed by an ETS table populated with atomic `:ets.update_counter/3`.
  A `Task` running on a 60s tick clears the table — the window is
  intentionally coarse; a rogue host gets at most 2× the threshold in
  any aligned 60s slice, which is plenty cheap protection for the
  threat model and avoids any per-key timer state.

  Configure:

      config :git_gud, GitGud.Federation.RateLimiter,
        # per source host
        inbox_per_minute: 60,
        # special-case for known-good hosts; nil to disable
        allowlist_multiplier: 5

  Use `check/1` from the inbox path:

      case RateLimiter.check(host) do
        :ok -> ...
        {:error, :rate_limited, retry_after_s} -> 429
      end
  """

  use GenServer
  import Ecto.Query, warn: false
  require Logger

  @table :gitgud_fed_rate_limits
  @sweep_interval_ms 60_000

  # ── child spec / start ───────────────────────────────────────────────

  def start_link(opts), do: GenServer.start_link(__MODULE__, opts, name: __MODULE__)

  @impl true
  def init(_opts) do
    :ets.new(@table, [
      :public,
      :named_table,
      :set,
      write_concurrency: true,
      read_concurrency: true
    ])

    schedule_sweep()
    {:ok, %{}}
  end

  @impl true
  def handle_info(:sweep, state) do
    :ets.delete_all_objects(@table)
    schedule_sweep()
    {:noreply, state}
  end

  defp schedule_sweep do
    Process.send_after(self(), :sweep, @sweep_interval_ms)
  end

  # ── public API ───────────────────────────────────────────────────────

  @doc """
  Increment the counter for `host` and decide.

  Returns `:ok` or `{:error, :rate_limited, retry_after_seconds}`. The
  retry hint is a worst-case wait, not the actual remaining window —
  good enough for HTTP `Retry-After`.
  """
  @spec check(String.t() | nil) :: :ok | {:error, :rate_limited, integer()}
  def check(nil), do: :ok

  def check(host) when is_binary(host) do
    threshold = effective_threshold(host)
    count = :ets.update_counter(@table, host, {2, 1}, {host, 0})

    if count > threshold do
      {:error, :rate_limited, 60}
    else
      :ok
    end
  end

  defp effective_threshold(host) do
    cfg = Application.get_env(:git_gud, __MODULE__, [])
    base = Keyword.get(cfg, :inbox_per_minute, 60)
    mult = Keyword.get(cfg, :allowlist_multiplier, 5)

    if mult && allowlisted?(host) do
      base * mult
    else
      base
    end
  end

  defp allowlisted?(host) do
    # Allowlisted hosts get a higher cap (they're already trusted, so a
    # spike is more likely legitimate). Avoid coupling RateLimiter to
    # the Repo at call time when we're hot-pathing 60+ rps — cache the
    # answer for 60s in ETS itself.
    case :ets.lookup(@table, {:allowlist, host}) do
      [{_, true}] ->
        true

      [{_, false}] ->
        false

      [] ->
        verdict =
          GitGud.Repo.exists?(
            from a in GitGud.Federation.InstanceAllowlist, where: a.host == ^host
          )

        :ets.insert(@table, {{:allowlist, host}, verdict})
        verdict
    end
  end

  @doc "Reset counters for tests."
  def reset do
    :ets.delete_all_objects(@table)
    :ok
  end

  @doc "Snapshot current counters for admin/observability."
  def snapshot do
    @table
    |> :ets.tab2list()
    |> Enum.flat_map(fn
      {host, count} when is_binary(host) -> [{host, count}]
      _ -> []
    end)
    |> Enum.sort_by(fn {_, c} -> -c end)
  end
end