summary refs log tree commit diff
path: root/tests/api/test_ratelimiting.py
blob: 1e1f30d790e0117f137a8492d59fc9d863de2d07 (plain) (blame)
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
from synapse.api.ratelimiting import LimitExceededError, Ratelimiter
from synapse.appservice import ApplicationService
from synapse.types import create_requester

from tests import unittest


class TestRatelimiter(unittest.TestCase):
    def test_allowed_via_can_do_action(self):
        limiter = Ratelimiter(clock=None, rate_hz=0.1, burst_count=1)
        allowed, time_allowed = limiter.can_do_action(key="test_id", _time_now_s=0)
        self.assertTrue(allowed)
        self.assertEquals(10.0, time_allowed)

        allowed, time_allowed = limiter.can_do_action(key="test_id", _time_now_s=5)
        self.assertFalse(allowed)
        self.assertEquals(10.0, time_allowed)

        allowed, time_allowed = limiter.can_do_action(key="test_id", _time_now_s=10)
        self.assertTrue(allowed)
        self.assertEquals(20.0, time_allowed)

    def test_allowed_user_via_can_requester_do_action(self):
        user_requester = create_requester("@user:example.com")
        limiter = Ratelimiter(clock=None, rate_hz=0.1, burst_count=1)
        allowed, time_allowed = limiter.can_requester_do_action(
            user_requester, _time_now_s=0
        )
        self.assertTrue(allowed)
        self.assertEquals(10.0, time_allowed)

        allowed, time_allowed = limiter.can_requester_do_action(
            user_requester, _time_now_s=5
        )
        self.assertFalse(allowed)
        self.assertEquals(10.0, time_allowed)

        allowed, time_allowed = limiter.can_requester_do_action(
            user_requester, _time_now_s=10
        )
        self.assertTrue(allowed)
        self.assertEquals(20.0, time_allowed)

    def test_allowed_appservice_ratelimited_via_can_requester_do_action(self):
        appservice = ApplicationService(
            None, "example.com", id="foo", rate_limited=True,
        )
        as_requester = create_requester("@user:example.com", app_service=appservice)

        limiter = Ratelimiter(clock=None, rate_hz=0.1, burst_count=1)
        allowed, time_allowed = limiter.can_requester_do_action(
            as_requester, _time_now_s=0
        )
        self.assertTrue(allowed)
        self.assertEquals(10.0, time_allowed)

        allowed, time_allowed = limiter.can_requester_do_action(
            as_requester, _time_now_s=5
        )
        self.assertFalse(allowed)
        self.assertEquals(10.0, time_allowed)

        allowed, time_allowed = limiter.can_requester_do_action(
            as_requester, _time_now_s=10
        )
        self.assertTrue(allowed)
        self.assertEquals(20.0, time_allowed)

    def test_allowed_appservice_via_can_requester_do_action(self):
        appservice = ApplicationService(
            None, "example.com", id="foo", rate_limited=False,
        )
        as_requester = create_requester("@user:example.com", app_service=appservice)

        limiter = Ratelimiter(clock=None, rate_hz=0.1, burst_count=1)
        allowed, time_allowed = limiter.can_requester_do_action(
            as_requester, _time_now_s=0
        )
        self.assertTrue(allowed)
        self.assertEquals(-1, time_allowed)

        allowed, time_allowed = limiter.can_requester_do_action(
            as_requester, _time_now_s=5
        )
        self.assertTrue(allowed)
        self.assertEquals(-1, time_allowed)

        allowed, time_allowed = limiter.can_requester_do_action(
            as_requester, _time_now_s=10
        )
        self.assertTrue(allowed)
        self.assertEquals(-1, time_allowed)

    def test_allowed_via_ratelimit(self):
        limiter = Ratelimiter(clock=None, rate_hz=0.1, burst_count=1)

        # Shouldn't raise
        limiter.ratelimit(key="test_id", _time_now_s=0)

        # Should raise
        with self.assertRaises(LimitExceededError) as context:
            limiter.ratelimit(key="test_id", _time_now_s=5)
        self.assertEqual(context.exception.retry_after_ms, 5000)

        # Shouldn't raise
        limiter.ratelimit(key="test_id", _time_now_s=10)

    def test_allowed_via_can_do_action_and_overriding_parameters(self):
        """Test that we can override options of can_do_action that would otherwise fail
        an action
        """
        # Create a Ratelimiter with a very low allowed rate_hz and burst_count
        limiter = Ratelimiter(clock=None, rate_hz=0.1, burst_count=1)

        # First attempt should be allowed
        allowed, time_allowed = limiter.can_do_action(("test_id",), _time_now_s=0,)
        self.assertTrue(allowed)
        self.assertEqual(10.0, time_allowed)

        # Second attempt, 1s later, will fail
        allowed, time_allowed = limiter.can_do_action(("test_id",), _time_now_s=1,)
        self.assertFalse(allowed)
        self.assertEqual(10.0, time_allowed)

        # But, if we allow 10 actions/sec for this request, we should be allowed
        # to continue.
        allowed, time_allowed = limiter.can_do_action(
            ("test_id",), _time_now_s=1, rate_hz=10.0
        )
        self.assertTrue(allowed)
        self.assertEqual(1.1, time_allowed)

        # Similarly if we allow a burst of 10 actions
        allowed, time_allowed = limiter.can_do_action(
            ("test_id",), _time_now_s=1, burst_count=10
        )
        self.assertTrue(allowed)
        self.assertEqual(1.0, time_allowed)

    def test_allowed_via_ratelimit_and_overriding_parameters(self):
        """Test that we can override options of the ratelimit method that would otherwise
        fail an action
        """
        # Create a Ratelimiter with a very low allowed rate_hz and burst_count
        limiter = Ratelimiter(clock=None, rate_hz=0.1, burst_count=1)

        # First attempt should be allowed
        limiter.ratelimit(key=("test_id",), _time_now_s=0)

        # Second attempt, 1s later, will fail
        with self.assertRaises(LimitExceededError) as context:
            limiter.ratelimit(key=("test_id",), _time_now_s=1)
        self.assertEqual(context.exception.retry_after_ms, 9000)

        # But, if we allow 10 actions/sec for this request, we should be allowed
        # to continue.
        limiter.ratelimit(key=("test_id",), _time_now_s=1, rate_hz=10.0)

        # Similarly if we allow a burst of 10 actions
        limiter.ratelimit(key=("test_id",), _time_now_s=1, burst_count=10)

    def test_pruning(self):
        limiter = Ratelimiter(clock=None, rate_hz=0.1, burst_count=1)
        limiter.can_do_action(key="test_id_1", _time_now_s=0)

        self.assertIn("test_id_1", limiter.actions)

        limiter.can_do_action(key="test_id_2", _time_now_s=10)

        self.assertNotIn("test_id_1", limiter.actions)