[auth] timeout based throttling at 10 requests per second

Author: patatoid

apps/boruta_web/lib/boruta_web/plugs/rate_limit.ex

@@ -2,10 +2,13 @@ defmodule BorutaWeb.Plugs.RateLimit do
   @moduledoc false
 
   defmodule Counter do
+    @moduledoc false
     use Agent
 
     @base_unit :millisecond
 
+    @memory_length 50
+
     @time_unit_stamps [
       millisecond: 1,
       second: 1_000,
@@ -19,19 +22,51 @@ defmodule BorutaWeb.Plugs.RateLimit do
     def get(ip, time_unit) do
       Agent.get(__MODULE__, fn counter ->
         Map.get(counter, ip, [])
+      end)
+      |> Enum.count(fn timestamp ->
+        timestamp > :os.system_time(@base_unit) - @time_unit_stamps[time_unit]
+      end)
+    end
+
+    def throttling_timeout(ip, count, time_unit, penality) do
+      now = :os.system_time(@base_unit)
+
+      request_rates = Agent.get(__MODULE__, fn counter ->
+        Map.get(counter, ip, [])
         |> Enum.filter(fn timestamp ->
-          timestamp > :os.system_time(@base_unit) - @time_unit_stamps[time_unit]
+          timestamp > now - @memory_length * @time_unit_stamps[time_unit]
         end)
-        |> Enum.count()
       end)
+      |> Enum.group_by(fn timestamp ->
+        div(timestamp, @time_unit_stamps[time_unit])
+      end)
+
+      timeout = Enum.map(0..@memory_length - 1, fn i ->
+        current = floor(now - (i * @time_unit_stamps[time_unit]))
+
+        Map.get(request_rates, div(current, @time_unit_stamps[time_unit]), [])
+      end)
+      |> Enum.reverse()
+      |> Enum.map(fn
+        [] -> count / @time_unit_stamps[time_unit]
+        timestamps -> Enum.count(timestamps) / @time_unit_stamps[time_unit]
+      end)
+      |> Enum.reduce(1, fn factor, acc ->
+        acc * factor * (@time_unit_stamps[time_unit] / count)
+      end)
+
+      case timeout <= 1 do
+        true -> 0
+        false -> floor(timeout * penality)
+      end
     end
 
     def increment(ip, time_unit) do
       Agent.update(__MODULE__, fn counter ->
         timestamps =
           Map.get(counter, ip, [])
           |> Enum.filter(fn timestamp ->
-            timestamp > :os.system_time(@base_unit) - @time_unit_stamps[time_unit]
+            timestamp > :os.system_time(@base_unit) - @memory_length * @time_unit_stamps[time_unit]
           end)
 
         Map.put(
@@ -52,11 +87,18 @@ defmodule BorutaWeb.Plugs.RateLimit do
 
     Counter.increment(remote_ip, options[:time_unit])
 
-    case Counter.get(remote_ip, options[:time_unit]) > options[:count] do
-      false ->
-        conn
+    max_timeout = options[:timeout]
 
-      true ->
+    case Counter.throttling_timeout(
+      remote_ip,
+      options[:count],
+      options[:time_unit],
+      options[:penality]
+    ) do
+      timeout when timeout < max_timeout ->
+        :timer.sleep(timeout)
+        conn
+      _ ->
         send_resp(conn, 429, "")
     end
   end

apps/boruta_web/lib/boruta_web/router.ex

@@ -2,6 +2,8 @@ defmodule BorutaWeb.Router do
   use BorutaWeb, :router
   use Plug.ErrorHandler
 
+  alias BorutaWeb.Plugs.RateLimit
+
   import BorutaIdentityWeb.Sessions,
     only: [
       fetch_current_user: 2
@@ -25,6 +27,7 @@ defmodule BorutaWeb.Router do
 
   pipeline :api do
     plug(:accepts, ["json", "jwt"])
+    plug RateLimit, count: 10, time_unit: :second, penality: 500, timeout: 5_000
   end
 
   scope "/", BorutaWeb do

apps/boruta_web/test/boruta_web/plugs/rate_limit_test.exs

@@ -17,19 +17,9 @@ defmodule BorutaWeb.Plugs.RateLimitTest do
     test "request is throttled", %{conn: conn} do
       :timer.sleep(1000)
       b_conn = %{conn | remote_ip: {127, 0, 0, 2}}
-      options = [time_unit: :second, count: 5]
+      options = [time_unit: :second, count: 5, penality: 500]
       assert RateLimit.call(conn, options) == conn
       assert RateLimit.call(b_conn, options) == b_conn
-      assert RateLimit.call(conn, options) == conn
-      assert RateLimit.call(b_conn, options) == b_conn
-      assert RateLimit.call(conn, options) == conn
-      assert RateLimit.call(b_conn, options) == b_conn
-      assert RateLimit.call(conn, options) == conn
-      assert RateLimit.call(b_conn, options) == b_conn
-      assert RateLimit.call(conn, options) == conn
-      assert RateLimit.call(b_conn, options) == b_conn
-      assert RateLimit.call(conn, options).status == 429
-      assert RateLimit.call(b_conn, options).status == 429
     end
   end
 
@@ -57,32 +47,88 @@ defmodule BorutaWeb.Plugs.RateLimitTest do
     end
   end
 
+  describe "Counter.throttling_timeout" do
+    test "gives the timeout within the time unit range" do
+      :timer.sleep(1000)
+      ip = :ip
+      time_unit = :second
+      penality = 100
+      count = 1
+
+      Agent.update(RateLimit.Counter, fn _counter -> %{} end)
+      assert RateLimit.Counter.throttling_timeout(ip, count, time_unit, penality) == 0
+
+      Agent.update(RateLimit.Counter, fn _counter -> %{ip => [:os.system_time(:millisecond)]} end)
+      assert RateLimit.Counter.throttling_timeout(ip, count, time_unit, penality) == 0
+
+      Agent.update(RateLimit.Counter, fn _counter ->
+        %{
+          ip => [
+            :os.system_time(:millisecond),
+            :os.system_time(:millisecond),
+            :os.system_time(:millisecond),
+            :os.system_time(:millisecond),
+            :os.system_time(:millisecond),
+            :os.system_time(:millisecond),
+            :os.system_time(:millisecond)
+          ]
+        }
+      end)
+
+      assert RateLimit.Counter.throttling_timeout(ip, count, time_unit, penality) == 700
+
+      Agent.update(RateLimit.Counter, fn _counter ->
+        %{
+          ip => [
+            :os.system_time(:millisecond) - 1000,
+            :os.system_time(:millisecond) - 800,
+            :os.system_time(:millisecond)
+          ]
+        }
+      end)
+
+      assert RateLimit.Counter.throttling_timeout(ip, count, time_unit, penality) == 200
+
+      Agent.update(RateLimit.Counter, fn _counter ->
+        %{ip => [:os.system_time(:millisecond), :os.system_time(:millisecond) - 1000]}
+      end)
+
+      assert RateLimit.Counter.throttling_timeout(ip, count, time_unit, penality) == 0
+    end
+  end
+
   describe "Counter.increment" do
     test "updates the counter" do
       :timer.sleep(1000)
       ip = :ip
       time_unit = :second
 
       assert Agent.get(RateLimit.Counter, fn counter ->
-               Map.get(counter, ip, []) |> Enum.count()
+               Map.get(counter, ip, [])
+               |> Enum.count(fn timestamp -> timestamp > :os.system_time(:millisecond) - 1000 end)
              end) == 0
 
       RateLimit.Counter.increment(ip, time_unit)
 
       assert Agent.get(RateLimit.Counter, fn %{^ip => timestamps} ->
-               Enum.count(timestamps)
+               timestamps
+               |> Enum.count(fn timestamp -> timestamp > :os.system_time(:millisecond) - 1000 end)
              end) == 1
 
       RateLimit.Counter.increment(ip, time_unit)
 
       assert Agent.get(RateLimit.Counter, fn %{^ip => timestamps} ->
-               Enum.count(timestamps)
+               timestamps
+               |> Enum.count(fn timestamp -> timestamp > :os.system_time(:millisecond) - 1000 end)
              end) == 2
 
       :timer.sleep(1000)
       RateLimit.Counter.increment(ip, time_unit)
 
-      assert Agent.get(RateLimit.Counter, fn %{^ip => timestamps} -> Enum.count(timestamps) end) ==
+      assert Agent.get(RateLimit.Counter, fn %{^ip => timestamps} ->
+               timestamps
+               |> Enum.count(fn timestamp -> timestamp > :os.system_time(:millisecond) - 1000 end)
+             end) ==
                1
     end
   end