Implement pagination for MSC2666 (#19279)

Co-authored-by: Andrew Morgan <1342360+anoadragon453@users.noreply.github.com>

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,10 +38,34 @@ if TYPE_CHECKING:
 
 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):
     """
-    GET /uk.half-shot.msc2666/user/mutual_rooms?user_id={user_id} HTTP/1.1
+    GET /uk.half-shot.msc2666/user/mutual_rooms?user_id={user_id}&from={token} HTTP/1.1
     """
 
     PATTERNS = client_patterns(
@@ -56,6 +83,7 @@ class UserMutualRoomsServlet(RestServlet):
         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,52 @@ class UserMutualRoomsServlet(RestServlet):
                 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(
+        # Sort here instead of the database function, so that we don't expose
+        # clients to any unrelated changes to the sorting algorithm.
+        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. This method doesn't care whether the
+            # provided token room still exists, nor whether it's even a real room ID.
+            #
+            # However, if rooms with a lower ID are added after the token was issued,
+            # they will not be included until the client makes a new request without a
+            # from token. This is considered acceptable, as clients generally won't
+            # persist these results for long periods.
+            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 @@ class UserMutualRoomsTest(unittest.HomeserverTestCase):
 
     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,52 @@ class UserMutualRoomsTest(unittest.HomeserverTestCase):
         for room_id_id in channel.json_body["joined"]:
             self.assertIn(room_id_id, [room_id_one, room_id_two])
 
+    def _create_rooms_for_pagination_test(
+        self, count: int
+    ) -> tuple[str, str, list[str]]:
+        u1 = self.register_user("user1", "pass")
+        u1_token = self.login(u1, "pass")
+        u2 = self.register_user("user2", "pass")
+        u2_token = self.login(u2, "pass")
+        room_ids = []
+        for i in range(count):
+            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()
+        return u1_token, u2, room_ids
+
+    def test_shared_room_list_pagination_two_pages(self) -> None:
+        u1_token, u2, room_ids = self._create_rooms_for_pagination_test(15)
+
+        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_pagination_one_page(self) -> None:
+        u1_token, u2, room_ids = self._create_rooms_for_pagination_test(10)
+
+        channel = self._get_mutual_rooms(u1_token, u2)
+        self.assertEqual(200, channel.code, channel.result)
+        self.assertEqual(channel.json_body["joined"], room_ids)
+        self.assertNotIn("next_batch", channel.json_body)
+
+    def test_shared_room_list_pagination_invalid_token(self) -> None:
+        u1_token, u2, room_ids = self._create_rooms_for_pagination_test(10)
+
+        channel = self._get_mutual_rooms(u1_token, u2, "!<>##faketoken")
+        self.assertEqual(400, channel.code, channel.result)
+        self.assertEqual(
+            "M_INVALID_PARAM", channel.json_body["errcode"], channel.result
+        )
+
     def test_shared_room_list_after_leave(self) -> None:
         """
         A room should no longer be considered shared if the other
@@ -172,3 +222,14 @@ class UserMutualRoomsTest(unittest.HomeserverTestCase):
         channel = self._get_mutual_rooms(u2_token, u1)
         self.assertEqual(200, channel.code, channel.result)
         self.assertEqual(len(channel.json_body["joined"]), 0)
+
+    def test_shared_room_list_nonexistent_user(self) -> None:
+        u1 = self.register_user("user1", "pass")
+        u1_token = self.login(u1, "pass")
+
+        # Check shared rooms from user1's perspective.
+        # We should see the one room in common
+        channel = self._get_mutual_rooms(u1_token, "@meow:example.com")
+        self.assertEqual(200, channel.code, channel.result)
+        self.assertEqual(len(channel.json_body["joined"]), 0)
+        self.assertNotIn("next_batch", channel.json_body)

tests/utils.py

@@ -198,7 +198,9 @@ def default_config(
         "rc_invites": {
             "per_room": {"per_second": 10000, "burst_count": 10000},
             "per_user": {"per_second": 10000, "burst_count": 10000},
+            "per_issuer": {"per_second": 10000, "burst_count": 10000},
         },
+        "rc_room_creation": {"per_second": 10000, "burst_count": 10000},
         "rc_3pid_validation": {"per_second": 10000, "burst_count": 10000},
         "rc_presence": {"per_user": {"per_second": 10000, "burst_count": 10000}},
         "saml2_enabled": False,