diff --git a/lib/rule.ex b/lib/rule.ex index 46b1097..e405636 100644 --- a/lib/rule.ex +++ b/lib/rule.ex @@ -98,7 +98,7 @@ defmodule PlugAttack.Rule do Be careful not to use the same `key` for different rules that use the same storage. - Passes `{:fail2ban, key}`, as the data to `block_action` calls when an + Passes `{:fail2ban, key, remaining_ban}`, as the data to `block_action` calls when an abusive request is detected. Each misbehaving client is blocked after each call and tracked for `:period` time. If more than `:limit` abusive requests are detected within the `:period`, the client is banned for `:ban_for`. @@ -129,15 +129,17 @@ defmodule PlugAttack.Rule do ban_for = Keyword.fetch!(opts, :ban_for) now = System.system_time(:millisecond) - if banned?(key, storage, now) do - {:block, {:fail2ban, :banned, key}} - else - track_fail2ban(key, storage, limit, period, ban_for, now) + case banned?(key, storage, now) do + {true, remaining_ban} -> {:block, {:fail2ban, :banned, key, remaining_ban}} + false -> track_fail2ban(key, storage, limit, period, ban_for, now) end end defp banned?(key, {mod, opts}, now) do - mod.read(opts, {:fail2ban_banned, key}, now) == {:ok, true} + case mod.read(opts, {:fail2ban_banned, key}, now) do + {:ok, true, remaining_ban} -> {true, remaining_ban} + _ -> false + end end defp track_fail2ban(key, {mod, opts}, limit, period, ban_for, now) do diff --git a/lib/storage/ets.ex b/lib/storage/ets.ex index fe93ebb..ea6c445 100644 --- a/lib/storage/ets.ex +++ b/lib/storage/ets.ex @@ -53,7 +53,7 @@ defmodule PlugAttack.Storage.Ets do def read(name, key, now) do case :ets.lookup(name, key) do [{^key, value, expires_at}] when expires_at > now -> - {:ok, value} + {:ok, value, expires_at - now} _ -> :error diff --git a/test/rules_test.exs b/test/rules_test.exs index 08c6521..f73df54 100644 --- a/test/rules_test.exs +++ b/test/rules_test.exs @@ -18,7 +18,7 @@ defmodule PlugAttack.RuleTest do :timer.sleep(1) assert {:allow, {:fail2ban, :counting, :key}} = fail2ban() :timer.sleep(100) - assert {:block, {:fail2ban, :banned, :key}} = fail2ban() + assert {:block, {:fail2ban, :banned, :key, 99}} = fail2ban() :timer.sleep(200) assert {:allow, {:fail2ban, :counting, :key}} = fail2ban() end diff --git a/test/storage/ets_test.exs b/test/storage/ets_test.exs index 7b46743..9e8acd1 100644 --- a/test/storage/ets_test.exs +++ b/test/storage/ets_test.exs @@ -26,7 +26,7 @@ defmodule PlugAttack.Storage.EtsTest do test "read/write" do assert :error = Ets.read(__MODULE__, :foo, now()) Ets.write(__MODULE__, :foo, true, expires_in(20)) - assert {:ok, true} == Ets.read(__MODULE__, :foo, now()) + assert {:ok, true, 20} == Ets.read(__MODULE__, :foo, now()) :timer.sleep(30) assert :error = Ets.read(__MODULE__, :foo, now()) end