Add ratelimit callbacks to module API to allow dynamic ratelimiting (#18458)

tests/api/test_ratelimiting.py

@@ -1,6 +1,10 @@
+from typing import Optional
+
 from synapse.api.ratelimiting import LimitExceededError, Ratelimiter
 from synapse.appservice import ApplicationService
 from synapse.config.ratelimiting import RatelimitSettings
+from synapse.module_api.callbacks.ratelimit_callbacks import RatelimitModuleApiCallbacks
+from synapse.storage.databases.main.room import RatelimitOverride
 from synapse.types import create_requester
 
 from tests import unittest
@@ -440,3 +444,49 @@ class TestRatelimiter(unittest.HomeserverTestCase):
             limiter.can_do_action(requester=None, key="a", _time_now_s=20.0)
         )
         self.assertTrue(success)
+
+    def test_get_ratelimit_override_for_user_callback(self) -> None:
+        test_user_id = "@user:test"
+        test_limiter_name = "name"
+        callbacks = RatelimitModuleApiCallbacks(self.hs)
+        requester = create_requester(test_user_id)
+        limiter = Ratelimiter(
+            store=self.hs.get_datastores().main,
+            clock=self.clock,
+            cfg=RatelimitSettings(
+                test_limiter_name,
+                per_second=0.1,
+                burst_count=3,
+            ),
+            ratelimit_callbacks=callbacks,
+        )
+
+        # Observe four actions, exceeding the burst_count.
+        limiter.record_action(requester=requester, n_actions=4, _time_now_s=0.0)
+
+        # We should be prevented from taking a new action now.
+        success, _ = self.get_success_or_raise(
+            limiter.can_do_action(requester=requester, _time_now_s=0.0)
+        )
+        self.assertFalse(success)
+
+        # Now register a callback that overrides the ratelimit for this user
+        # and limiter name.
+        async def get_ratelimit_override_for_user(
+            user_id: str, limiter_name: str
+        ) -> Optional[RatelimitOverride]:
+            if user_id == test_user_id:
+                return RatelimitOverride(
+                    messages_per_second=0.1,
+                    burst_count=10,
+                )
+            return None
+
+        callbacks.register_callbacks(
+            get_ratelimit_override_for_user=get_ratelimit_override_for_user
+        )
+
+        success, _ = self.get_success_or_raise(
+            limiter.can_do_action(requester=requester, _time_now_s=0.0)
+        )
+        self.assertTrue(success)