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))