plug_attack/test/rules_test.exs at 015ff07e87f62376e8e78badef4b78476326888d · michalmuskala/plug_attack · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
defmodule PlugAttack.RuleTest do
  use ExUnit.Case, async: true

  setup do
    {:ok, _} = PlugAttack.Storage.Ets.start_link(__MODULE__)
    :ok
  end

  test "fail2ban" do
    assert {:allow, {:fail2ban, :counting, :key}} = fail2ban()
    :timer.sleep(1)
    assert {:allow, {:fail2ban, :counting, :key}} = fail2ban()
    :timer.sleep(150)

    assert {:allow, {:fail2ban, :counting, :key}} = fail2ban()
    :timer.sleep(1)
    assert {:allow, {:fail2ban, :counting, :key}} = fail2ban()
    :timer.sleep(1)
    assert {:allow, {:fail2ban, :counting, :key}} = fail2ban()
    :timer.sleep(100)
    assert {:block, {:fail2ban, :banned, :key, 99}} = fail2ban()
    :timer.sleep(200)
    assert {:allow, {:fail2ban, :counting, :key}} = fail2ban()
  end

  defp fail2ban() do
    PlugAttack.Rule.fail2ban(:key,
      storage: {PlugAttack.Storage.Ets, __MODULE__},
      period: 100,
      limit: 3,
      ban_for: 200
    )
  end

  test "throttle" do
    assert {:allow, {:throttle, data}} = throttle()

    expires = (div(System.system_time(:millisecond), 100) + 1) * 100
    assert data[:period] == 100
    assert data[:limit] == 5
    assert data[:remaining] == 4
    assert data[:expires_at] == expires

    assert {:allow, _} = throttle()
    assert {:allow, _} = throttle()
    assert {:allow, _} = throttle()
    assert {:allow, _} = throttle()

    assert {:block, {:throttle, data}} = throttle()
    assert data[:period] == 100
    assert data[:limit] == 5
    assert data[:remaining] == 0
    assert data[:expires_at] == expires

    :timer.sleep(90)
    assert {:allow, {:throttle, data}} = throttle()
    assert data[:period] == 100
    assert data[:limit] == 5
    assert data[:remaining] == 4
    assert data[:expires_at] == expires + 100
  end

  defp throttle() do
    PlugAttack.Rule.throttle(:key,
      storage: {PlugAttack.Storage.Ets, __MODULE__},
      limit: 5,
      period: 100
    )
  end
end