diff options
Diffstat (limited to 'tests')
| -rw-r--r-- | tests/api/test_ratelimiting.py | 51 |
1 files changed, 40 insertions, 11 deletions
diff --git a/tests/api/test_ratelimiting.py b/tests/api/test_ratelimiting.py index f661a9ff8e..18649c2c05 100644 --- a/tests/api/test_ratelimiting.py +++ b/tests/api/test_ratelimiting.py @@ -246,7 +246,7 @@ class TestRatelimiter(unittest.HomeserverTestCase): self.assertTrue(allowed) self.assertEqual(10.0, time_allowed) - # Test that, after doing these 3 actions, we can't do any more action without + # Test that, after doing these 3 actions, we can't do any more actions without # waiting. allowed, time_allowed = self.get_success_or_raise( limiter.can_do_action(None, key="test_id", n_actions=1, _time_now_s=0) @@ -254,7 +254,8 @@ class TestRatelimiter(unittest.HomeserverTestCase): self.assertFalse(allowed) self.assertEqual(10.0, time_allowed) - # Test that after waiting we can do only 1 action. + # Test that after waiting we would be able to do only 1 action. + # Note that we don't actually do it (update=False) here. allowed, time_allowed = self.get_success_or_raise( limiter.can_do_action( None, @@ -265,23 +266,51 @@ class TestRatelimiter(unittest.HomeserverTestCase): ) ) self.assertTrue(allowed) - # The time allowed is the current time because we could still repeat the action - # once. - self.assertEqual(10.0, time_allowed) + # We would be able to do the 5th action at t=20. + self.assertEqual(20.0, time_allowed) + # Attempt (but fail) to perform TWO actions at t=10. + # Those would be the 4th and 5th actions. allowed, time_allowed = self.get_success_or_raise( limiter.can_do_action(None, key="test_id", n_actions=2, _time_now_s=10) ) self.assertFalse(allowed) - # The time allowed doesn't change despite allowed being False because, while we - # don't allow 2 actions, we could still do 1. + # The returned time allowed for the next action is now even though we weren't + # allowed to perform the action because whilst we don't allow 2 actions, + # we could still do 1. self.assertEqual(10.0, time_allowed) - # Test that after waiting a bit more we can do 2 actions. + # Test that after waiting until t=20, we can do perform 2 actions. + # These are the 4th and 5th actions. allowed, time_allowed = self.get_success_or_raise( limiter.can_do_action(None, key="test_id", n_actions=2, _time_now_s=20) ) self.assertTrue(allowed) - # The time allowed is the current time because we could still repeat the action - # once. - self.assertEqual(20.0, time_allowed) + # We would be able to do the 6th action at t=30. + self.assertEqual(30.0, time_allowed) + + def test_rate_limit_burst_only_given_once(self) -> None: + """ + Regression test against a bug that meant that you could build up + extra tokens by timing requests. + """ + limiter = Ratelimiter( + store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=3 + ) + + def consume_at(time: float) -> bool: + success, _ = self.get_success_or_raise( + limiter.can_do_action(requester=None, key="a", _time_now_s=time) + ) + return success + + # Use all our 3 burst tokens + self.assertTrue(consume_at(0.0)) + self.assertTrue(consume_at(0.1)) + self.assertTrue(consume_at(0.2)) + + # Wait to recover 1 token (10 seconds at 0.1 Hz). + self.assertTrue(consume_at(10.1)) + + # Check that we get rate limited after using that token. + self.assertFalse(consume_at(11.1)) |