diff options
| -rw-r--r-- | synapse/api/ratelimiting.py | 67 | ||||
| -rw-r--r-- | tests/api/__init__.py | 0 | ||||
| -rw-r--r-- | tests/api/test_ratelimiting.py | 39 |
3 files changed, 87 insertions, 19 deletions
diff --git a/synapse/api/ratelimiting.py b/synapse/api/ratelimiting.py index a730abd59b..ab26c2376a 100644 --- a/synapse/api/ratelimiting.py +++ b/synapse/api/ratelimiting.py @@ -2,35 +2,64 @@ import collections class Ratelimiter(object): + """ + Ratelimit message sending by user. + """ def __init__(self): self.message_counts = collections.OrderedDict() - def prune_message_counts(self, time_now): - for user_id in self.message_counts.keys(): - message_count, time_start, msg_rate_hz = ( - self.message_counts[user_id] - ) - time_delta = time_now - time_start - if message_count - time_delta * msg_rate_hz > 0: - break - else: - del self.message_counts[user_id] - - def send_message(self, user_id, time_now, msg_rate_hz, burst_count): - self.prune_message_counts(time_now) + def send_message(self, user_id, time_now_s, msg_rate_hz, burst_count): + """Can the user send a message? + Args: + user_id: The user sending a message. + time_now_s: The time now. + msg_rate_hz: The long term number of messages a user can send in a + second. + burst_count: How many messages the user can send before being + limited. + Returns: + A pair of a bool indicating if they can send a message now and a + time in seconds of when they can next send a message. + """ + self.prune_message_counts(time_now_s) message_count, time_start, _ignored = self.message_counts.pop( - user_id, (0., time_now, None), + user_id, (0., time_now_s, None), ) - time_delta = time_now - time_start - if message_count - time_delta * msg_rate_hz < 0: - a - if message_count - (time_now - time_start) * msg_rate_hz > burst_count: + time_delta = time_now_s - time_start + sent_count = message_count - time_delta * msg_rate_hz + if sent_count < 0: + allowed = True + time_start = time_now_s + messagecount = 1. + elif sent_count > burst_count - 1.: allowed = False else: allowed = True message_count += 1 + self.message_counts[user_id] = ( message_count, time_start, msg_rate_hz ) - return allowed + + if msg_rate_hz > 0: + time_allowed = ( + time_start + (message_count - burst_count + 1) / msg_rate_hz + ) + if time_allowed < time_now_s: + time_allowed = time_now_s + else: + time_allowed = -1 + + return allowed, time_allowed + + def prune_message_counts(self, time_now_s): + for user_id in self.message_counts.keys(): + message_count, time_start, msg_rate_hz = ( + self.message_counts[user_id] + ) + time_delta = time_now_s - time_start + if message_count - time_delta * msg_rate_hz > 0: + break + else: + del self.message_counts[user_id] diff --git a/tests/api/__init__.py b/tests/api/__init__.py new file mode 100644 index 0000000000..e69de29bb2 --- /dev/null +++ b/tests/api/__init__.py diff --git a/tests/api/test_ratelimiting.py b/tests/api/test_ratelimiting.py new file mode 100644 index 0000000000..dc2f83c7eb --- /dev/null +++ b/tests/api/test_ratelimiting.py @@ -0,0 +1,39 @@ +from synapse.api.ratelimiting import Ratelimiter + +import unittest + +class TestRatelimiter(unittest.TestCase): + + def test_allowed(self): + limiter = Ratelimiter() + allowed, time_allowed = limiter.send_message( + user_id="test_id", time_now_s=0, msg_rate_hz=0.1, burst_count=1, + ) + self.assertTrue(allowed) + self.assertEquals(10., time_allowed) + + allowed, time_allowed = limiter.send_message( + user_id="test_id", time_now_s=5, msg_rate_hz=0.1, burst_count=1, + ) + self.assertFalse(allowed) + self.assertEquals(10., time_allowed) + + allowed, time_allowed = limiter.send_message( + user_id="test_id", time_now_s=10, msg_rate_hz=0.1, burst_count=1 + ) + self.assertTrue(allowed) + self.assertEquals(20., time_allowed) + + def test_pruning(self): + limiter = Ratelimiter() + allowed, time_allowed = limiter.send_message( + user_id="test_id_1", time_now_s=0, msg_rate_hz=0.1, burst_count=1, + ) + + self.assertIn("test_id_1", limiter.message_counts) + + allowed, time_allowed = limiter.send_message( + user_id="test_id_2", time_now_s=10, msg_rate_hz=0.1, burst_count=1 + ) + + self.assertNotIn("test_id_1", limiter.message_counts) |