neiam /gitgud

defmodule GitGud.RateLimit do
  @moduledoc """
  Counters that cap how often something may happen per key: failed
  sign-ins per account and per client IP, emails sent per address,
  sign-ups per IP, runner polls per runner.

  Each policy is a limit and a window. A key's window opens on its
  first hit and runs for the policy's length; once the count reaches
  the limit the key is refused until the window closes.

      case RateLimit.check(:login_account, email) do
        :ok -> ...try the password, RateLimit.record/2 on failure...
        {:error, retry_after_s} -> refuse
      end

  `check/2` only asks; `record/2` counts; `hit/2` counts and asks in
  one go, for things where every attempt costs (an email, a poll).

  Limits can be overridden per policy:

      config :git_gud, GitGud.RateLimit, login_account: {10, 900}

  Backed by a public ETS table owned by this process, swept once a
  minute. Per node: in a cluster each node keeps its own counts, so a
  client spread across N nodes gets up to N times the limit — still a
  wall against guessing, not a precise quota.
  """

  use GenServer

  @table :gitgud_rate_limit
  @sweep_interval_ms 60_000

  # policy => {limit, window in seconds}
  @policies %{
    # Failed password sign-ins (web form, git and package Basic auth).
    login_account: {10, 15 * 60},
    login_ip: {50, 15 * 60},
    # Sign-in / confirmation emails asked for.
    email_address: {5, 60 * 60},
    email_ip: {20, 60 * 60},
    # Accounts created.
    register_ip: {10, 60 * 60},
    # Runner registrations with a bad token.
    runner_register_ip: {20, 15 * 60},
    # FetchTask polls; forgejo-runner polls every 2s by default.
    runner_fetch: {120, 60},
    # A failed sign-in reaches the account's security log once per
    # address per window, so a guessing run is one entry, not hundreds.
    login_failed_audit: {1, 15 * 60},
    # What one CI run may open with its job token, keyed by run id; a
    # day outlasts any run.
    job_issues: {20, 24 * 60 * 60},
    job_pull_requests: {10, 24 * 60 * 60}
  }

  @type policy ::
          :login_account
          | :login_ip
          | :email_address
          | :email_ip
          | :register_ip
          | :runner_register_ip
          | :runner_fetch
          | :login_failed_audit
          | :job_issues
          | :job_pull_requests

  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])
    schedule_sweep()
    {:ok, %{}}
  end

  @impl true
  def handle_info(:sweep, state) do
    now = now()
    :ets.select_delete(@table, [{{:_, :_, :"$1"}, [{:<, :"$1", now}], [true]}])
    schedule_sweep()
    {:noreply, state}
  end

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

  @doc """
  `:ok` while `key` is under `policy`'s limit, otherwise
  `{:error, retry_after_seconds}`. Counts nothing.
  """
  @spec check(policy(), term()) :: :ok | {:error, pos_integer()}
  def check(policy, key) do
    {limit, _window} = limits(policy)

    case lookup({policy, key}) do
      {count, expires_at} when count >= limit -> {:error, retry_after(expires_at)}
      _ -> :ok
    end
  end

  @doc "Count one occurrence for `key` under `policy`."
  @spec record(policy(), term()) :: :ok
  def record(policy, key) do
    {_limit, window} = limits(policy)
    _ = bump({policy, key}, window)
    :ok
  end

  @doc """
  Count one occurrence and decide: `:ok` if it fits in the limit,
  `{:error, retry_after_seconds}` if it's one too many.
  """
  @spec hit(policy(), term()) :: :ok | {:error, pos_integer()}
  def hit(policy, key) do
    {limit, window} = limits(policy)

    case bump({policy, key}, window) do
      {count, _expires_at} when count <= limit -> :ok
      {_count, expires_at} -> {:error, retry_after(expires_at)}
    end
  end

  @doc "Forget `key`'s count under `policy` (e.g. after a good sign-in)."
  @spec clear(policy(), term()) :: :ok
  def clear(policy, key) do
    :ets.delete(@table, {policy, key})
    :ok
  end

  @doc "`{limit, window_seconds}` for `policy`, with any configured override."
  @spec limits(policy()) :: {pos_integer(), pos_integer()}
  def limits(policy) do
    Keyword.get(Application.get_env(:git_gud, __MODULE__, []), policy) ||
      Map.fetch!(@policies, policy)
  end

  @doc """
  The key for a client address: the address itself, or its /64 for
  IPv6, since one host is routinely handed a whole /64.
  """
  @spec ip_key(:inet.ip_address() | nil) :: String.t()
  def ip_key(nil), do: "unknown"
  def ip_key({_, _, _, _} = ip), do: ip |> :inet.ntoa() |> to_string()

  def ip_key({a, b, c, d, _, _, _, _}),
    do: ({a, b, c, d, 0, 0, 0, 0} |> :inet.ntoa() |> to_string()) <> "/64"

  @doc "Whole minutes, rounded up, for a message like \"try again in 3 minutes\"."
  @spec minutes(pos_integer()) :: pos_integer()
  def minutes(seconds), do: max(div(seconds + 59, 60), 1)

  defp lookup(key) do
    now = now()

    case :ets.lookup(@table, key) do
      [{_, count, expires_at}] when expires_at > now -> {count, expires_at}
      _ -> nil
    end
  end

  # A window that has closed starts over. Two processes racing past the
  # expiry may both reset it; that loses a count, never adds a refusal.
  defp bump(key, window) do
    now = now()

    case :ets.lookup(@table, key) do
      [{_, _count, expires_at}] when expires_at <= now ->
        :ets.insert(@table, {key, 1, now + window * 1000})
        {1, now + window * 1000}

      _ ->
        expires_at = now + window * 1000
        count = :ets.update_counter(@table, key, {2, 1}, {key, 0, expires_at})
        {count, lookup_expiry(key, expires_at)}
    end
  end

  defp lookup_expiry(key, default) do
    case :ets.lookup(@table, key) do
      [{_, _, expires_at}] -> expires_at
      _ -> default
    end
  end

  defp retry_after(expires_at), do: max(div(expires_at - now() + 999, 1000), 1)

  defp now, do: System.monotonic_time(:millisecond)
end