Fix a long-standing bug which meant that rate limiting was not restrictive enough in some cases. (#13018)

This commit is contained in:
reivilibre 2022-06-15 15:11:55 +01:00 committed by GitHub
parent 417f4cf40b
commit 0dbdc39940
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 45 additions and 12 deletions

1
changelog.d/13018.bugfix Normal file
View File

@ -0,0 +1 @@
Fix a long-standing bug which meant that rate limiting was not restrictive enough in some cases.

View File

@ -128,6 +128,9 @@ class Ratelimiter:
performed_count = action_count - time_delta * rate_hz performed_count = action_count - time_delta * rate_hz
if performed_count < 0: if performed_count < 0:
performed_count = 0 performed_count = 0
# Reset the start time and forgive all actions
action_count = 0
time_start = time_now_s time_start = time_now_s
# This check would be easier read as performed_count + n_actions > burst_count, # This check would be easier read as performed_count + n_actions > burst_count,
@ -140,7 +143,7 @@ class Ratelimiter:
else: else:
# We haven't reached our limit yet # We haven't reached our limit yet
allowed = True allowed = True
action_count = performed_count + n_actions action_count = action_count + n_actions
if update: if update:
self.actions[key] = (action_count, time_start, rate_hz) self.actions[key] = (action_count, time_start, rate_hz)

View File

@ -246,7 +246,7 @@ class TestRatelimiter(unittest.HomeserverTestCase):
self.assertTrue(allowed) self.assertTrue(allowed)
self.assertEqual(10.0, time_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. # waiting.
allowed, time_allowed = self.get_success_or_raise( allowed, time_allowed = self.get_success_or_raise(
limiter.can_do_action(None, key="test_id", n_actions=1, _time_now_s=0) 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.assertFalse(allowed)
self.assertEqual(10.0, time_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( allowed, time_allowed = self.get_success_or_raise(
limiter.can_do_action( limiter.can_do_action(
None, None,
@ -265,23 +266,51 @@ class TestRatelimiter(unittest.HomeserverTestCase):
) )
) )
self.assertTrue(allowed) self.assertTrue(allowed)
# The time allowed is the current time because we could still repeat the action # We would be able to do the 5th action at t=20.
# once. self.assertEqual(20.0, time_allowed)
self.assertEqual(10.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( allowed, time_allowed = self.get_success_or_raise(
limiter.can_do_action(None, key="test_id", n_actions=2, _time_now_s=10) limiter.can_do_action(None, key="test_id", n_actions=2, _time_now_s=10)
) )
self.assertFalse(allowed) self.assertFalse(allowed)
# The time allowed doesn't change despite allowed being False because, while we # The returned time allowed for the next action is now even though we weren't
# don't allow 2 actions, we could still do 1. # allowed to perform the action because whilst we don't allow 2 actions,
# we could still do 1.
self.assertEqual(10.0, time_allowed) 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( allowed, time_allowed = self.get_success_or_raise(
limiter.can_do_action(None, key="test_id", n_actions=2, _time_now_s=20) limiter.can_do_action(None, key="test_id", n_actions=2, _time_now_s=20)
) )
self.assertTrue(allowed) self.assertTrue(allowed)
# The time allowed is the current time because we could still repeat the action # We would be able to do the 6th action at t=30.
# once. self.assertEqual(30.0, time_allowed)
self.assertEqual(20.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))