Implement pagination for MSC2666

changelog.d/19279.feature

@@ -0,0 +1 @@
+Implemented pagination for the [MSC2666](https://github.com/matrix-org/matrix-spec-proposals/pull/2666) mutual rooms endpoint. Contributed by @tulir @ Beeper.

synapse/rest/client/mutual_rooms.py

@@ -19,9 +19,12 @@
 #
 #
 import logging
+from bisect import bisect
 from http import HTTPStatus
 from typing import TYPE_CHECKING
 
+from unpaddedbase64 import decode_base64, encode_base64
+
 from synapse.api.errors import Codes, SynapseError
 from synapse.http.server import HttpServer
 from synapse.http.servlet import RestServlet, parse_strings_from_args
@@ -35,6 +38,30 @@
 
 logger = logging.getLogger(__name__)
 
+MUTUAL_ROOMS_BATCH_LIMIT = 100
+
+
+def _parse_mutual_rooms_batch_token_args(args: dict[bytes, list[bytes]]) -> str | None:
+    from_batches = parse_strings_from_args(args, "from")
+    if not from_batches:
+        return None
+    if len(from_batches) > 1:
+        raise SynapseError(
+            HTTPStatus.BAD_REQUEST,
+            "Duplicate from query parameter",
+            errcode=Codes.INVALID_PARAM,
+        )
+    if from_batches[0]:
+        try:
+            return decode_base64(from_batches[0]).decode("utf-8")
+        except Exception:
+            raise SynapseError(
+                HTTPStatus.BAD_REQUEST,
+                "Malformed from token",
+                errcode=Codes.INVALID_PARAM,
+            )
+    return None
+
 
 class UserMutualRoomsServlet(RestServlet):
     """
@@ -56,6 +83,7 @@ async def on_GET(self, request: SynapseRequest) -> tuple[int, JsonDict]:
         args: dict[bytes, list[bytes]] = request.args  # type: ignore
 
         user_ids = parse_strings_from_args(args, "user_id", required=True)
+        from_batch = _parse_mutual_rooms_batch_token_args(args)
 
         if len(user_ids) > 1:
             raise SynapseError(
@@ -64,29 +92,44 @@ async def on_GET(self, request: SynapseRequest) -> tuple[int, JsonDict]:
                 errcode=Codes.INVALID_PARAM,
             )
 
-        # We don't do batching, so a batch token is illegal by default
-        if b"batch_token" in args:
-            raise SynapseError(
-                HTTPStatus.BAD_REQUEST,
-                "Unknown batch_token",
-                errcode=Codes.INVALID_PARAM,
-            )
-
         user_id = user_ids[0]
 
         requester = await self.auth.get_user_by_req(request)
         if user_id == requester.user.to_string():
             raise SynapseError(
-                HTTPStatus.UNPROCESSABLE_ENTITY,
+                HTTPStatus.BAD_REQUEST,
                 "You cannot request a list of shared rooms with yourself",
-                errcode=Codes.INVALID_PARAM,
+                errcode=Codes.UNKNOWN,
             )
 
-        rooms = await self.store.get_mutual_rooms_between_users(
-            frozenset((requester.user.to_string(), user_id))
+        rooms = sorted(
+            await self.store.get_mutual_rooms_between_users(
+                frozenset((requester.user.to_string(), user_id))
+            )
         )
 
-        return 200, {"joined": list(rooms)}
+        if from_batch:
+            # A from_batch token was provided, so cut off any rooms where the ID is
+            # lower than or equal to the token
+            rooms = rooms[bisect(rooms, from_batch) :]
+
+        if len(rooms) < MUTUAL_ROOMS_BATCH_LIMIT:
+            # We've reached the end of the list, don't return a batch token
+            return 200, {"joined": rooms}
+
+        rooms = rooms[:MUTUAL_ROOMS_BATCH_LIMIT]
+        # We use urlsafe unpadded base64 encoding for the batch token in order to
+        # handle funny room IDs in old pre-v12 rooms properly. We also truncate it
+        # to stay within the 255-character limit of opaque tokens.
+        next_batch = encode_base64(rooms[-1].encode("utf-8"), urlsafe=True)[:255]
+        # Due to the truncation, it is technically possible to have conflicting next
+        # batches by creating hundreds of rooms with the same 191 character prefix
+        # in the room ID. In the event that some silly user does that, don't let
+        # them paginate further.
+        if next_batch == from_batch:
+            return 200, {"joined": rooms}
+
+        return 200, {"joined": list(rooms), "next_batch": next_batch}
 
 
 def register_servlets(hs: "HomeServer", http_server: HttpServer) -> None:

tests/rest/client/test_mutual_rooms.py

@@ -55,12 +55,16 @@ def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
 
     def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
         self.store = hs.get_datastores().main
+        mutual_rooms.MUTUAL_ROOMS_BATCH_LIMIT = 10
 
-    def _get_mutual_rooms(self, token: str, other_user: str) -> FakeChannel:
+    def _get_mutual_rooms(
+        self, token: str, other_user: str, since_token: str | None = None
+    ) -> FakeChannel:
         return self.make_request(
             "GET",
             "/_matrix/client/unstable/uk.half-shot.msc2666/user/mutual_rooms"
-            f"?user_id={quote(other_user)}",
+            f"?user_id={quote(other_user)}"
+            + (f"&from={quote(since_token)}" if since_token else ""),
             access_token=token,
         )
 
@@ -141,6 +145,30 @@ def _check_mutual_rooms_with(
         for room_id_id in channel.json_body["joined"]:
             self.assertIn(room_id_id, [room_id_one, room_id_two])
 
+    def test_shared_room_list_pagination(self) -> None:
+        u1 = self.register_user("user1", "pass")
+        u1_token = self.login(u1, "pass")
+        u2 = self.register_user("user2", "pass")
+        u2_token = self.login(u2, "pass")
+        self.get_success(self.hs.get_datastores().main.set_ratelimit_for_user(u1, 0, 0))
+        room_ids = []
+        for i in range(15):
+            room_id = self.helper.create_room_as(u1, is_public=i % 2 == 0, tok=u1_token)
+            self.helper.invite(room_id, src=u1, targ=u2, tok=u1_token)
+            self.helper.join(room_id, user=u2, tok=u2_token)
+            room_ids.append(room_id)
+        room_ids.sort()
+
+        channel = self._get_mutual_rooms(u1_token, u2)
+        self.assertEqual(200, channel.code, channel.result)
+        self.assertEqual(channel.json_body["joined"], room_ids[0:10])
+        self.assertIn("next_batch", channel.json_body)
+
+        channel = self._get_mutual_rooms(u1_token, u2, channel.json_body["next_batch"])
+        self.assertEqual(200, channel.code, channel.result)
+        self.assertEqual(channel.json_body["joined"], room_ids[10:20])
+        self.assertNotIn("next_batch", channel.json_body)
+
     def test_shared_room_list_after_leave(self) -> None:
         """
         A room should no longer be considered shared if the other