Introduce conditional throttle

lib/plug_attack.ex

@@ -93,8 +93,8 @@ defmodule PlugAttack do
   Defines a rule.
 
   A rule is an expression that returns either `{:allow, data}`, `{:block, data}`,
-  or `nil`. If an allow or block tuple is returned we say the rule *matched*,
-  otherwise the rule didn't match and further rules will be evaluated.
+  `nil` or updated `conn`. If an allow or block tuple is returned we say the rule
+  *matched*, otherwise the rule didn't match and further rules will be evaluated.
 
   If a rule matched the corresponding `allow_action/3` or `block_action/3`
   function on the defining module will be called passing the `conn`,
@@ -169,21 +169,30 @@ defmodule PlugAttack do
 
   @doc false
   def compile(env, rules) do
-    conn = quote(do: conn)
     opts = quote(do: opts)
-    body = Enum.reduce(rules, conn, &quote_rule(&2, &1, conn, opts, env))
-    {conn, opts, body}
+    conn = quote(do: conn)
+    chain = Enum.reduce(rules, conn, &quote_rule(&2, &1, conn, opts, env))
+    {conn, opts, quote do
+      priv = {unquote(env.module), unquote(opts)}
+      conn = Plug.Conn.put_private(unquote(conn), :plug_attack, priv)
+      unquote(chain)
+    end}
   end
 
   defp quote_rule(next, name, conn, opts, _env) do
     quote do
       case unquote(name)(unquote(conn)) do
-        {:allow, data} -> allow_action(unquote(conn), data, unquote(opts))
-        {:block, data} -> block_action(unquote(conn), data, unquote(opts))
-        nil            -> unquote(next)
+        {:allow, data} ->
+          allow_action(unquote(conn), data, unquote(opts))
+        {:block, data} ->
+          block_action(unquote(conn), data, unquote(opts))
+        %Plug.Conn{} = unquote(conn) ->
+          unquote(next)
+        nil ->
+          unquote(next)
         other ->
           raise "a PlugAttack rule should return `{:allow, data}`, " <>
-            "`{:block, data}`, or `nil`, got: #{inspect other}"
+            "`{:block, data}`, `nil`, or `conn`, got: #{inspect other}"
       end
     end
   end

lib/rule.ex

@@ -74,7 +74,7 @@ defmodule PlugAttack.Rule do
     now     = System.system_time(:milliseconds)
 
     expires_at = expires_at(now, period)
-    count      = do_throttle(storage, key, now, period, expires_at)
+    count      = increment_throttle(storage, key, now, period, expires_at)
     rem        = limit - count
     data       = [period: period, expires_at: expires_at,
                   limit: limit, remaining: max(rem, 0)]
@@ -83,11 +83,106 @@ defmodule PlugAttack.Rule do
 
   defp expires_at(now, period), do: (div(now, period) + 1) * period
 
-  defp do_throttle({mod, opts}, key, now, period, expires_at) do
+  defp increment_throttle({mod, opts}, key, now, period, expires_at) do
     full_key = {:throttle, key, div(now, period)}
     mod.increment(opts, full_key, 1, expires_at)
   end
 
+  @doc """
+  Implements a conditional request throttling algorithm.
+
+  With a request that does not use conditional headers (`If-Modified-Since`
+  or `If-None-Match` behaves exactly like `throttle/2`). For conditional
+  requests defers counting the request towards the limit to after the response
+  is computed using `Plug.Conn.register_before_send/2`. The throttle counter
+  is not incremented in case of a `304 Not Modified` response.
+
+  The `key` differentiates different throttles, you can use, for example,
+  `conn.remote_ip` for per IP throttling, or an email address for login attempts
+  limitation. If the `key` is falsey the throttling is not performed and
+  next rules are evaluated.
+
+  Be careful not to use the same `key` for different rules that use the same
+  storage.
+
+  Passes `{:throttle, data}`, as the data to both allow and block tuples, where
+  data is a keyword containing: `:period`, `:limit`, `:expires_at` - when the
+  current limit will expire as unix time in milliseconds,
+  and `:remaining` - the remaining limit. This can be useful for adding
+  "X-RateLimit-*" headers. When lazy throttling is performed the `allow_action`
+  callback is called from within the callback registered with
+  `Plug.Conn.register_before_send/2`.
+
+  ## Race conditions
+
+  Because the counter is imcremented lazily, there's a possible race condition,
+  where more requests are let-through than intended. This can happen during
+  long requests, when the counter is not incremented (yet) when new requests
+  are coming in.
+
+  ## Options
+
+    * `:storage` - required, a tuple of `PlugAttack.Storage` implementation
+      and storage options.
+    * `:limit` - required, how many requests in a period are allowed.
+    * `:period` - required, how long, in ms, is the period.
+
+  """
+  @spec conditional_throttle(Plug.Conn.t, term, Keyword.t) :: PlugAttack.rule
+  def conditional_throttle(conn, key, opts) do
+    cond do
+      key && conditional_request?(conn) ->
+        do_conditional_throttle(conn, key, opts)
+      key ->
+        do_throttle(key, opts)
+      true ->
+        nil
+    end
+  end
+
+  defp conditional_request?(conn) do
+    Plug.Conn.get_req_header(conn, "if-none-match") != []
+    or Plug.Conn.get_req_header(conn, "if-modified-since") != []
+  end
+
+  defp do_conditional_throttle(conn, key, opts) do
+    storage = Keyword.fetch!(opts, :storage)
+    limit   = Keyword.fetch!(opts, :limit)
+    period  = Keyword.fetch!(opts, :period)
+    now     = System.system_time(:milliseconds)
+
+    expires_at = expires_at(now, period)
+    count      = check_throttle(storage, key, now, period, expires_at)
+    rem        = limit - count
+    if rem >= 0 do
+      Plug.Conn.register_before_send(conn, fn conn ->
+        before_send_throttle(conn, storage, key, now, period, expires_at, limit, rem)
+      end)
+    else
+      data = [period: period, expires_at: expires_at,
+              limit: limit, remaining: max(rem, 0)]
+      {:block, {:throttle, data}}
+    end
+  end
+
+  defp before_send_throttle(conn, storage, key, now, period, expires_at, limit, rem) do
+    rem =
+      if conn.status != 304 do
+        limit - increment_throttle(storage, key, now, period, expires_at)
+      else
+        rem
+      end
+    data = [period: period, expires_at: expires_at,
+            limit: limit, remaining: max(rem, 0)]
+    {attack_module, opts} = conn.private.plug_attack
+    attack_module.allow_action(conn, {:allow, {:throttle, data}}, opts)
+  end
+
+  defp check_throttle({mod, opts}, key, now, period, expires_at) do
+    full_key = {:throttle, key, div(now, period)}
+    mod.increment(opts, full_key, 0, expires_at)
+  end
+
   @doc """
   Implements an algorithm inspired by fail2ban.
 

mix.lock

@@ -1,4 +1,4 @@
-%{"earmark": {:hex, :earmark, "1.1.1", "433136b7f2e99cde88b745b3a0cfc3fbc81fe58b918a09b40fce7f00db4d8187", [:mix], []},
-  "ex_doc": {:hex, :ex_doc, "0.15.0", "e73333785eef3488cf9144a6e847d3d647e67d02bd6fdac500687854dd5c599f", [:mix], [{:earmark, "~> 1.1", [hex: :earmark, optional: false]}]},
-  "mime": {:hex, :mime, "1.1.0", "01c1d6f4083d8aa5c7b8c246ade95139620ef8effb009edde934e0ec3b28090a", [:mix], []},
-  "plug": {:hex, :plug, "1.3.1", "aaf54675428a393370ec1f4c865e3fcf42608686960764beac93c5abeba9e655", [:mix], [{:cowboy, "~> 1.0.1 or ~> 1.1", [hex: :cowboy, optional: true]}, {:mime, "~> 1.0", [hex: :mime, optional: false]}]}}
+%{"earmark": {:hex, :earmark, "1.2.2", "f718159d6b65068e8daeef709ccddae5f7fdc770707d82e7d126f584cd925b74", [:mix], [], "hexpm"},
+  "ex_doc": {:hex, :ex_doc, "0.16.1", "b4b8a23602b4ce0e9a5a960a81260d1f7b29635b9652c67e95b0c2f7ccee5e81", [:mix], [{:earmark, "~> 1.1", [hex: :earmark, repo: "hexpm", optional: false]}], "hexpm"},
+  "mime": {:hex, :mime, "1.1.0", "01c1d6f4083d8aa5c7b8c246ade95139620ef8effb009edde934e0ec3b28090a", [:mix], [], "hexpm"},
+  "plug": {:hex, :plug, "1.3.5", "7503bfcd7091df2a9761ef8cecea666d1f2cc454cbbaf0afa0b6e259203b7031", [:mix], [{:cowboy, "~> 1.0.1 or ~> 1.1", [hex: :cowboy, repo: "hexpm", optional: true]}, {:mime, "~> 1.0", [hex: :mime, repo: "hexpm", optional: false]}], "hexpm"}}

test/plug_attack_test.exs

@@ -31,13 +31,26 @@ defmodule PlugAttackTest do
     assert function_exported?(TestPlug, :call, 2)
   end
 
+  test "stores plug & opts in private", %{conn: conn} do
+    ref = make_ref()
+    conn = TestPlug.call(conn, TestPlug.init(ref))
+    assert {TestPlug, ref} == conn.private.plug_attack
+  end
+
   test "uses the rule definition with allow", %{conn: conn} do
     Process.put(:rule, {:allow, []})
     conn = TestPlug.call(conn, TestPlug.init([]))
     refute conn.halted
     assert_received {:allow, []}
   end
 
+  test "allows returning updated conn from a rule", %{conn: conn} do
+    updated = Plug.Conn.assign(conn, :test, make_ref())
+    Process.put(:rule, updated)
+    conn = TestPlug.call(conn, TestPlug.init([]))
+    assert conn == updated
+  end
+
   test "uses the rule definition with block", %{conn: conn} do
     Process.put(:rule, {:block, []})
     conn = TestPlug.call(conn, TestPlug.init([]))

test/rules_test.exs

@@ -1,5 +1,10 @@
 defmodule PlugAttack.RuleTest do
   use ExUnit.Case, async: true
+  use Plug.Test
+
+  doctest PlugAttack.Rule
+
+  @storage {PlugAttack.Storage.Ets, __MODULE__}
 
   setup do
     {:ok, _} = PlugAttack.Storage.Ets.start_link(__MODULE__)
@@ -24,8 +29,8 @@ defmodule PlugAttack.RuleTest do
   end
 
   defp fail2ban() do
-    PlugAttack.Rule.fail2ban(:key, storage: {PlugAttack.Storage.Ets, __MODULE__},
-                                   period: 100, limit: 3, ban_for: 200)
+    PlugAttack.Rule.fail2ban(:key,
+      storage: @storage, period: 100, limit: 3, ban_for: 200)
   end
 
   test "throttle" do
@@ -57,7 +62,54 @@ defmodule PlugAttack.RuleTest do
   end
 
   defp throttle() do
-    PlugAttack.Rule.throttle(:key, storage: {PlugAttack.Storage.Ets, __MODULE__},
-                                   limit: 5, period: 100)
+    PlugAttack.Rule.throttle(:key,
+      storage: @storage, limit: 5, period: 100)
+  end
+
+  describe "conditional throttle" do
+    setup do
+      [conn: conn(:get, "/")]
+    end
+
+    test "conditional throttle with unconditional request", %{conn: conn} do
+      assert {:allow, {:throttle, _}} = conditional_throttle(conn)
+      assert {:allow, {:throttle, _}} = conditional_throttle(conn)
+      assert {:block, {:throttle, _}} = conditional_throttle(conn)
+    end
+
+    test "conditional throttle with If-None-Match request", %{conn: conn} do
+      conditional =
+        conn
+        |> put_private(:plug_attack, {__MODULE__, self()})
+        |> put_req_header("if-none-match", "x")
+
+      assert %Plug.Conn{} = req1 = conditional_throttle(conditional)
+      assert %Plug.Conn{} = req2 = conditional_throttle(conditional)
+      assert %Plug.Conn{} = req3 = conditional_throttle(conditional)
+
+      refute_received {:allow_action, _, _}
+
+      send_resp(req1, 304, "")
+      assert_received {:allow_action, _, _}
+      send_resp(req2, 200, "")
+      assert_received {:allow_action, _, _}
+      assert {:allow, _} = conditional_throttle(conn)
+      assert {:block, _} = conditional_throttle(conn)
+
+      # This is the race condition described in the docs
+      send_resp(req3, 200, "")
+      assert_received {:allow_action, _, _}
+    end
+  end
+
+  defp conditional_throttle(conn) do
+    PlugAttack.Rule.conditional_throttle(conn, :key,
+      storage: @storage, limit: 2, period: 100)
+  end
+
+  # Simulating PlugAttack for conditional throttle test
+  def allow_action(conn, data, pid) do
+    send(pid, {:allow_action, conn, data})
+    conn
   end
 end