drempel

if Code.ensure_loaded?(Phoenix.Controller) do
  defmodule Drempel.Plugs.RateLimit do
    @moduledoc """
    Use this plug to enforce a rate limit on a Phoenix action,
    the default key is conn.remote_ip,
    the default bucket is {controller_module, action_name}.
    """

    alias Plug.Conn
    @behaviour Plug

    import Phoenix.Controller, only: [action_name: 1, controller_module: 1]

    @doc false
    def init(opts) do
      opts = Enum.into(opts, %{})

      %{
        handler: {Map.get(opts, :handler, __MODULE__), :rate_exceeded},
        key: Map.get(opts, :key, nil),
        bucket: Map.get(opts, :bucket, nil)
      }
    end

    @doc false
    def call(%Conn{} = conn, opts) do
      key = Map.get(opts, :key) || conn.remote_ip
      bucket = Map.get(opts, :bucket) || {controller_module(conn), action_name(conn)}

      case Drempel.get(bucket, key) do
        0 -> conn
        delay -> handle_limit(conn, div(delay, 1000), opts)
      end
    end

    @doc """
    Adds to the limit and assigns to :backoff if a backoff is needed
    """
    def update(%Conn{} = conn, opts) do
      opts = Enum.into(opts, %{})
      key = Map.get(opts, :key) || conn.remote_ip
      bucket = Map.get(opts, :bucket) || {controller_module(conn), action_name(conn)}

      case Drempel.update(bucket, key) do
        0 ->
          conn

        delay ->
          conn
          |> Conn.assign(:backoff, delay)
      end
    end

    @doc """
    A default rate limit exceeded handler that just sends 429 "Too Many Requests".
    """
    def rate_exceeded(conn, _params) do
      conn
      |> Conn.send_resp(429, "Too Many Requests")
    end

    defp handle_limit(conn, delay, opts) do
      conn =
        conn
        |> Conn.assign(:backoff, delay)
        |> Conn.halt()

      {mod, meth} = Map.get(opts, :handler)
      apply(mod, meth, [conn, conn.params])
    end
  end
end