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