Add an admin API to manage ratelimit for a specific user

changelog.d/9648.feature

@@ -0,0 +1 @@
+Add an admin API to manage ratelimit for a specific user.

docs/admin_api/user_admin_api.rst

@@ -823,3 +823,114 @@ The following parameters should be set in the URL:
 
 - ``user_id`` - The fully qualified MXID: for example, ``@user:server.com``. The user must
   be local.
+
+Override ratelimiting for users
+===============================
+
+This API allows to override or disable rate limiting for a specific user.
+There are specific APIs to set, get and delete a ratelimit.
+
+Get status of ratelimit
+-----------------------
+
+The API is::
+
+  GET /_synapse/admin/v1/users/<user_id>/override_ratelimit
+
+To use it, you will need to authenticate by providing an ``access_token`` for a
+server admin: see `README.rst <README.rst>`_.
+
+A response body like the following is returned:
+
+.. code:: json
+
+    {
+      "messages_per_second": 0,
+      "burst_count": 0
+    }
+
+**Parameters**
+
+The following parameters should be set in the URL:
+
+- ``user_id`` - The fully qualified MXID: for example, ``@user:server.com``. The user must
+  be local.
+
+**Response**
+
+The following fields are returned in the JSON response body:
+
+- ``messages_per_second`` - integer - The number of actions that can
+  be performed in a second.
+- ``burst_count`` - integer - How many actions that can be performed
+  before being limited.
+
+If **no** custom ratelimit is set, the values are ``null``.
+
+Set ratelimit
+-------------
+
+The API is::
+
+  POST /_synapse/admin/v1/users/<user_id>/override_ratelimit
+
+To use it, you will need to authenticate by providing an ``access_token`` for a
+server admin: see `README.rst <README.rst>`_.
+
+A response body like the following is returned:
+
+.. code:: json
+
+    {
+      "messages_per_second": 0,
+      "burst_count": 0
+    }
+
+**Parameters**
+
+The following parameters should be set in the URL:
+
+- ``user_id`` - The fully qualified MXID: for example, ``@user:server.com``. The user must
+  be local.
+
+Body parameters:
+
+- ``messages_per_second`` - positive integer, optional. The number of actions that can
+  be performed in a second. Defaults to ``0``.
+- ``burst_count`` - positive integer, optional. How many actions that can be performed
+  before being limited. Defaults to ``0``.
+
+To disable users' ratelimit set both values to ``0``.
+
+**Response**
+
+The following fields are returned in the JSON response body:
+
+- ``messages_per_second`` - integer - The number of actions that can
+  be performed in a second.
+- ``burst_count`` - integer - How many actions that can be performed
+  before being limited.
+
+Delete ratelimit
+----------------
+
+The API is::
+
+  DELETE /_synapse/admin/v1/users/<user_id>/override_ratelimit
+
+To use it, you will need to authenticate by providing an ``access_token`` for a
+server admin: see `README.rst <README.rst>`_.
+
+An empty JSON dict is returned.
+
+.. code:: json
+
+    {}
+
+**Parameters**
+
+The following parameters should be set in the URL:
+
+- ``user_id`` - The fully qualified MXID: for example, ``@user:server.com``. The user must
+  be local.
+

synapse/rest/admin/__init__.py

@@ -54,6 +54,7 @@
     AccountValidityRenewServlet,
     DeactivateAccountRestServlet,
     PushersRestServlet,
+    RateLimitRestServlet,
     ResetPasswordRestServlet,
     SearchUsersRestServlet,
     ShadowBanRestServlet,
@@ -240,6 +241,7 @@ def register_servlets(hs, http_server):
     ShadowBanRestServlet(hs).register(http_server)
     ForwardExtremitiesRestServlet(hs).register(http_server)
     RoomEventContextServlet(hs).register(http_server)
+    RateLimitRestServlet(hs).register(http_server)
 
 
 def register_servlets_for_client_rest_resource(hs, http_server):

synapse/rest/admin/users.py

@@ -982,3 +982,112 @@ async def on_POST(
         await self.store.set_shadow_banned(UserID.from_string(user_id), True)
 
         return 200, {}
+
+
+class RateLimitRestServlet(RestServlet):
+    """An admin API to override ratelimiting for an user.
+
+    Example:
+        POST /_synapse/admin/v1/users/@test:example.com/override_ratelimit
+        {
+          "messages_per_second": 0,
+          "burst_count": 0
+        }
+        200 OK
+        {
+          "messages_per_second": 0,
+          "burst_count": 0
+        }
+    """
+
+    PATTERNS = admin_patterns("/users/(?P<user_id>[^/]*)/override_ratelimit")
+
+    def __init__(self, hs: "HomeServer"):
+        self.hs = hs
+        self.store = hs.get_datastore()
+        self.auth = hs.get_auth()
+
+    async def on_GET(
+        self, request: SynapseRequest, user_id: str
+    ) -> Tuple[int, JsonDict]:
+        await assert_requester_is_admin(self.auth, request)
+
+        if not self.hs.is_mine_id(user_id):
+            raise SynapseError(400, "Can only lookup local users")
+
+        if not await self.store.get_user_by_id(user_id):
+            raise NotFoundError("User not found")
+
+        ratelimit = await self.store.get_ratelimit_for_user(user_id)
+
+        if ratelimit:
+            messages_per_second = ratelimit.messages_per_second
+            burst_count = ratelimit.burst_count
+        else:
+            messages_per_second = None
+            burst_count = None
+
+        ret = {
+            "messages_per_second": messages_per_second,
+            "burst_count": burst_count,
+        }
+
+        return 200, ret
+
+    async def on_POST(
+        self, request: SynapseRequest, user_id: str
+    ) -> Tuple[int, JsonDict]:
+        await assert_requester_is_admin(self.auth, request)
+
+        if not self.hs.is_mine_id(user_id):
+            raise SynapseError(400, "Only local users can be ratelimited")
+
+        if not await self.store.get_user_by_id(user_id):
+            raise NotFoundError("User not found")
+
+        body = parse_json_object_from_request(request, allow_empty_body=True)
+
+        messages_per_second = body.get("messages_per_second", 0)
+        burst_count = body.get("burst_count", 0)
+
+        if not isinstance(messages_per_second, int) or messages_per_second < 0:
+            raise SynapseError(
+                400,
+                "%r parameter must be a positive int" % (messages_per_second,),
+                errcode=Codes.INVALID_PARAM,
+            )
+
+        if not isinstance(burst_count, int) or burst_count < 0:
+            raise SynapseError(
+                400,
+                "%r parameter must be a positive int" % (burst_count,),
+                errcode=Codes.INVALID_PARAM,
+            )
+
+        await self.store.set_ratelimit_for_user(
+            user_id, messages_per_second, burst_count
+        )
+        ratelimit = await self.store.get_ratelimit_for_user(user_id)
+        assert ratelimit is not None
+
+        ret = {
+            "messages_per_second": ratelimit.messages_per_second,
+            "burst_count": ratelimit.burst_count,
+        }
+
+        return 200, ret
+
+    async def on_DELETE(
+        self, request: SynapseRequest, user_id: str
+    ) -> Tuple[int, JsonDict]:
+        await assert_requester_is_admin(self.auth, request)
+
+        if not self.hs.is_mine_id(user_id):
+            raise SynapseError(400, "Only local users can be ratelimited")
+
+        if not await self.store.get_user_by_id(user_id):
+            raise NotFoundError("User not found")
+
+        await self.store.delete_ratelimit_for_user(user_id)
+
+        return 200, {}

synapse/storage/databases/main/room.py

@@ -521,13 +521,11 @@ def _get_rooms_paginate_txn(txn):
         )
 
     @cached(max_entries=10000)
-    async def get_ratelimit_for_user(self, user_id):
-        """Check if there are any overrides for ratelimiting for the given
-        user
+    async def get_ratelimit_for_user(self, user_id) -> Optional[RatelimitOverride]:
+        """Check if there are any overrides for ratelimiting for the given user
 
         Args:
-            user_id (str)
-
+            user_id
         Returns:
             RatelimitOverride if there is an override, else None. If the contents
             of RatelimitOverride are None or 0 then ratelimitng has been
@@ -549,6 +547,63 @@ async def get_ratelimit_for_user(self, user_id):
         else:
             return None
 
+    async def set_ratelimit_for_user(
+        self, user_id: str, messages_per_second: int, burst_count: int
+    ) -> None:
+        """Sets whether a user shadow-banned.
+        Args:
+            user: user ID of the user to test
+            messages_per_second:
+            burst_count:
+        """
+
+        def set_ratelimit_txn(txn):
+            self.db_pool.simple_upsert_txn(
+                txn,
+                table="ratelimit_override",
+                keyvalues={"user_id": user_id},
+                values={
+                    "messages_per_second": messages_per_second,
+                    "burst_count": burst_count,
+                },
+            )
+
+            self._invalidate_cache_and_stream(
+                txn, self.get_ratelimit_for_user, (user_id,)
+            )
+
+        await self.db_pool.runInteraction("set_ratelimit", set_ratelimit_txn)
+
+    async def delete_ratelimit_for_user(self, user_id: str) -> None:
+        """Sets whether a user shadow-banned.
+        Args:
+            user: user ID of the user to test
+            shadow_banned: true iff the user is to be shadow-banned, false otherwise.
+        """
+
+        def delete_ratelimit_txn(txn):
+            row = self.db_pool.simple_select_one_txn(
+                txn,
+                table="ratelimit_override",
+                keyvalues={"user_id": user_id},
+                retcols=["user_id"],
+                allow_none=True,
+            )
+
+            if not row:
+                return
+
+            # They are there, delete them.
+            self.db_pool.simple_delete_one_txn(
+                txn, "ratelimit_override", keyvalues={"user_id": user_id}
+            )
+
+            self._invalidate_cache_and_stream(
+                txn, self.get_ratelimit_for_user, (user_id,)
+            )
+
+        await self.db_pool.runInteraction("delete_ratelimit", delete_ratelimit_txn)
+
     @cached()
     async def get_retention_policy_for_room(self, room_id):
         """Get the retention policy for a given room.