X-Git-Url: http://git.squeep.com/?a=blobdiff_plain;f=lib%2Fpleroma%2Fplugs%2Frate_limiter.ex;h=31388f574c1bfa70c2b19171b50a763e672217fb;hb=17c237ba808d4356bb1e202e459680563b79ef99;hp=9ba5875fa31aa21434f6a129d99da2e9b1b64634;hpb=a0c65bbd6c708b555f457bf24ec07d2d41c3fe4a;p=akkoma

diff --git a/lib/pleroma/plugs/rate_limiter.ex b/lib/pleroma/plugs/rate_limiter.ex
index 9ba5875fa..31388f574 100644
--- a/lib/pleroma/plugs/rate_limiter.ex
+++ b/lib/pleroma/plugs/rate_limiter.ex
@@ -31,12 +31,28 @@ defmodule Pleroma.Plugs.RateLimiter do
 
   ## Usage
 
+  AllowedSyntax:
+
+      plug(Pleroma.Plugs.RateLimiter, :limiter_name)
+      plug(Pleroma.Plugs.RateLimiter, {:limiter_name, options})
+
+  Allowed options:
+
+      * `bucket_name` overrides bucket name (e.g. to have a separate limit for a set of actions)
+      * `params` appends values of specified request params (e.g. ["id"]) to bucket name
+
   Inside a controller:
 
       plug(Pleroma.Plugs.RateLimiter, :one when action == :one)
       plug(Pleroma.Plugs.RateLimiter, :two when action in [:two, :three])
 
-  or inside a router pipiline:
+      plug(
+        Pleroma.Plugs.RateLimiter,
+        {:status_id_action, bucket_name: "status_id_action:fav_unfav", params: ["id"]}
+        when action in ~w(fav_status unfav_status)a
+      )
+
+  or inside a router pipeline:
 
       pipeline :api do
         ...
@@ -44,39 +60,61 @@ defmodule Pleroma.Plugs.RateLimiter do
         ...
       end
   """
-
-  import Phoenix.Controller, only: [json: 2]
+  import Pleroma.Web.TranslationHelpers
   import Plug.Conn
 
   alias Pleroma.User
 
-  def init(limiter_name) do
+  def init(limiter_name) when is_atom(limiter_name) do
+    init({limiter_name, []})
+  end
+
+  def init({limiter_name, opts}) do
     case Pleroma.Config.get([:rate_limit, limiter_name]) do
       nil -> nil
-      config -> {limiter_name, config}
+      config -> {limiter_name, config, opts}
     end
   end
 
-  # do not limit if there is no limiter configuration
+  # Do not limit if there is no limiter configuration
   def call(conn, nil), do: conn
 
-  def call(conn, opts) do
-    case check_rate(conn, opts) do
-      {:ok, _count} -> conn
-      {:error, _count} -> render_error(conn)
+  def call(conn, settings) do
+    case check_rate(conn, settings) do
+      {:ok, _count} ->
+        conn
+
+      {:error, _count} ->
+        render_throttled_error(conn)
+    end
+  end
+
+  defp bucket_name(conn, limiter_name, opts) do
+    bucket_name = opts[:bucket_name] || limiter_name
+
+    if params_names = opts[:params] do
+      params_values = for p <- Enum.sort(params_names), do: conn.params[p]
+      Enum.join([bucket_name] ++ params_values, ":")
+    else
+      bucket_name
     end
   end
 
-  defp check_rate(%{assigns: %{user: %User{id: user_id}}}, {limiter_name, [_, {scale, limit}]}) do
-    ExRated.check_rate("#{limiter_name}:#{user_id}", scale, limit)
+  defp check_rate(
+         %{assigns: %{user: %User{id: user_id}}} = conn,
+         {limiter_name, [_, {scale, limit}], opts}
+       ) do
+    bucket_name = bucket_name(conn, limiter_name, opts)
+    ExRated.check_rate("#{bucket_name}:#{user_id}", scale, limit)
   end
 
-  defp check_rate(conn, {limiter_name, [{scale, limit} | _]}) do
-    ExRated.check_rate("#{limiter_name}:#{ip(conn)}", scale, limit)
+  defp check_rate(conn, {limiter_name, [{scale, limit} | _], opts}) do
+    bucket_name = bucket_name(conn, limiter_name, opts)
+    ExRated.check_rate("#{bucket_name}:#{ip(conn)}", scale, limit)
   end
 
-  defp check_rate(conn, {limiter_name, {scale, limit}}) do
-    check_rate(conn, {limiter_name, [{scale, limit}]})
+  defp check_rate(conn, {limiter_name, {scale, limit}, opts}) do
+    check_rate(conn, {limiter_name, [{scale, limit}, {scale, limit}], opts})
   end
 
   def ip(%{remote_ip: remote_ip}) do
@@ -85,10 +123,9 @@ defmodule Pleroma.Plugs.RateLimiter do
     |> Enum.join(".")
   end
 
-  defp render_error(conn) do
+  defp render_throttled_error(conn) do
     conn
-    |> put_status(:too_many_requests)
-    |> json(%{error: "Throttled"})
+    |> render_error(:too_many_requests, "Throttled")
     |> halt()
   end
 end