From 769692fe79fc4d5d9d785a1c1a3c26a4d68b4474 Mon Sep 17 00:00:00 2001 From: Michal Muskala Date: Wed, 7 Jun 2017 12:13:37 +0200 Subject: [PATCH 1/2] Update deps --- mix.lock | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mix.lock b/mix.lock index 370a7c9..7c88fa9 100644 --- a/mix.lock +++ b/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"}} From 878e0e63cd82fb7c7ba6dddcad8ca98fe7843e83 Mon Sep 17 00:00:00 2001 From: Michal Muskala Date: Wed, 7 Jun 2017 18:31:22 +0200 Subject: [PATCH 2/2] Add conditional_throttle --- lib/plug_attack.ex | 27 +++++++---- lib/rule.ex | 99 ++++++++++++++++++++++++++++++++++++++- test/plug_attack_test.exs | 13 +++++ test/rules_test.exs | 60 ++++++++++++++++++++++-- 4 files changed, 184 insertions(+), 15 deletions(-) diff --git a/lib/plug_attack.ex b/lib/plug_attack.ex index 042904f..9598ec7 100644 --- a/lib/plug_attack.ex +++ b/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, "e_rule(&2, &1, conn, opts, env)) - {conn, opts, body} + conn = quote(do: conn) + chain = Enum.reduce(rules, conn, "e_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 diff --git a/lib/rule.ex b/lib/rule.ex index 784c5bb..aa2ae63 100644 --- a/lib/rule.ex +++ b/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. diff --git a/test/plug_attack_test.exs b/test/plug_attack_test.exs index d610966..e915fb0 100644 --- a/test/plug_attack_test.exs +++ b/test/plug_attack_test.exs @@ -31,6 +31,12 @@ 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([])) @@ -38,6 +44,13 @@ defmodule PlugAttackTest do 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([])) diff --git a/test/rules_test.exs b/test/rules_test.exs index 97cc92d..a04ee18 100644 --- a/test/rules_test.exs +++ b/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