Better check whether we need to exclude partially stated rooms

2 files changed, 183 insertions, 40 deletions

diff --git a/synapse/handlers/sliding_sync.py b/synapse/handlers/sliding_sync.py
index 45739ccd35..dbbe705dce 100644
--- a/synapse/handlers/sliding_sync.py
+++ b/synapse/handlers/sliding_sync.py
@@ -24,6 +24,7 @@ from itertools import chain
 from typing import (
     TYPE_CHECKING,
     Any,
+    Callable,
     Dict,
     Final,
     List,
@@ -363,6 +364,73 @@ class RoomSyncConfig:
                     else:
                         self.required_state_map[state_type].add(state_key)
 
+    def must_await_full_state(
+        self,
+        *,
+        is_mine_id: Callable[[str], bool],
+        is_mine_server_name: Callable[[str], bool],
+    ) -> bool:
+        """
+        Check if we have a we're only requesting `required_state` which is completely
+        satisfied even with partial state, then we don't need to `await_full_state` before
+        we can return it.
+
+        Also see `StateFilter.must_await_full_state(...)` for comparison
+
+        Args:
+            is_mine_id: a callable which confirms if a given state_key matches a mxid
+               of a local user
+            is_mine_server_name: a callable which confirms if a given server name
+                matches the local server name
+        """
+        # Partially-stated rooms should have all state events except for remote
+        # membership events so if we require a remote membership event anywhere, then we
+        # need to return `False`.
+
+        wildcard_state_keys = self.required_state_map.get(StateValues.WILDCARD)
+        if (
+            wildcard_state_keys is not None
+            and StateValues.WILDCARD in wildcard_state_keys
+        ):
+            return True
+
+        if wildcard_state_keys is not None:
+            for possible_user_id in wildcard_state_keys:
+                if not possible_user_id[0].startswith(UserID.SIGIL):
+                    # Not a user ID
+                    continue
+
+                localpart_hostname = possible_user_id.split(":", 1)
+                if len(localpart_hostname) < 2:
+                    # Not a user ID
+                    continue
+
+                if not is_mine_server_name(localpart_hostname[1]):
+                    return True
+
+        membership_state_keys = self.required_state_map.get(EventTypes.Member)
+        # We aren't requesting any membership events at all
+        if membership_state_keys is None:
+            return False
+
+        # If we're requesting entirely local users, TODO
+        for user_id in membership_state_keys:
+            if user_id == StateValues.ME:
+                continue
+            # We're lazy-loading membership so we can just return the state we have.
+            # Lazy-loading means we include membership for any event `sender` in the
+            # timeline but since we had to auth those timeline events, we will have the
+            # membership state for them (including from remote senders).
+            elif user_id == StateValues.LAZY:
+                continue
+            elif user_id == StateValues.WILDCARD:
+                return False
+            elif not is_mine_id(user_id):
+                return True
+
+        # local users only
+        return False
+
 
 class StateValues:
     """
@@ -392,6 +460,8 @@ class SlidingSyncHandler:
         self.device_handler = hs.get_device_handler()
         self.push_rules_handler = hs.get_push_rules_handler()
         self.rooms_to_exclude_globally = hs.config.server.rooms_to_exclude_from_sync
+        self.is_mine_id = hs.is_mine_id
+        self.is_mine_server_name = hs.is_mine_server_name
 
         self.connection_store = SlidingSyncConnectionStore()
 
@@ -572,16 +642,11 @@ class SlidingSyncHandler:
                     # Since creating the `RoomSyncConfig` takes some work, let's just do it
                     # once and make a copy whenever we need it.
                     room_sync_config = RoomSyncConfig.from_room_config(list_config)
-                    membership_state_keys = room_sync_config.required_state_map.get(
-                        EventTypes.Member
-                    )
-                    # Also see `StateFilter.must_await_full_state(...)` for comparison
-                    lazy_loading = (
-                        membership_state_keys is not None
-                        and StateValues.LAZY in membership_state_keys
-                    )
 
-                    if not lazy_loading:
+                    if room_sync_config.must_await_full_state(
+                        is_mine_id=self.is_mine_id,
+                        is_mine_server_name=self.is_mine_server_name,
+                    ):
                         # Exclude partially-stated rooms unless the `required_state`
                         # only has `["m.room.member", "$LAZY"]` for membership
                         # (lazy-loading room members).
diff --git a/tests/rest/client/sliding_sync/test_rooms_required_state.py b/tests/rest/client/sliding_sync/test_rooms_required_state.py
index a13cad223f..40f7c9d3d8 100644
--- a/tests/rest/client/sliding_sync/test_rooms_required_state.py
+++ b/tests/rest/client/sliding_sync/test_rooms_required_state.py
@@ -649,43 +649,123 @@ class SlidingSyncRoomsRequiredStateTestCase(SlidingSyncBase):
             mark_event_as_partial_state(self.hs, join_response2["event_id"], room_id2)
         )
 
-        # Make the Sliding Sync request (NOT lazy-loading room members)
+        # Make the Sliding Sync request with examples where `must_await_full_state()` is
+        # `False`
         sync_body = {
             "lists": {
-                "foo-list": {
+                "no-state-list": {
+                    "ranges": [[0, 1]],
+                    "required_state": [],
+                    "timeline_limit": 0,
+                },
+                "other-state-list": {
                     "ranges": [[0, 1]],
                     "required_state": [
                         [EventTypes.Create, ""],
                     ],
                     "timeline_limit": 0,
                 },
+                "lazy-load-list": {
+                    "ranges": [[0, 1]],
+                    "required_state": [
+                        [EventTypes.Create, ""],
+                        # Lazy-load room members
+                        [EventTypes.Member, StateValues.LAZY],
+                        # Local member
+                        [EventTypes.Member, user2_id],
+                    ],
+                    "timeline_limit": 0,
+                },
+                "local-members-only-list": {
+                    "ranges": [[0, 1]],
+                    "required_state": [
+                        # Own user ID
+                        [EventTypes.Member, user1_id],
+                        # Local member
+                        [EventTypes.Member, user2_id],
+                    ],
+                    "timeline_limit": 0,
+                },
+                "me-list": {
+                    "ranges": [[0, 1]],
+                    "required_state": [
+                        # Own user ID
+                        [EventTypes.Member, StateValues.ME],
+                        # Local member
+                        [EventTypes.Member, user2_id],
+                    ],
+                    "timeline_limit": 0,
+                },
+                "wildcard-type-local-state-key-list": {
+                    "ranges": [[0, 1]],
+                    "required_state": [
+                        ["*", user1_id],
+                        # Not a user ID
+                        ["*", "foobarbaz"],
+                        # Not a user ID
+                        ["*", "foo.bar.baz"],
+                        # Not a user ID
+                        ["*", "@foo"],
+                    ],
+                    "timeline_limit": 0,
+                },
             }
         }
         response_body, _ = self.do_sync(sync_body, tok=user1_tok)
 
-        # Make sure the list includes room1 but room2 is excluded because it's still
-        # partially-stated
-        self.assertListEqual(
-            list(response_body["lists"]["foo-list"]["ops"]),
-            [
-                {
-                    "op": "SYNC",
-                    "range": [0, 1],
-                    "room_ids": [room_id1],
-                }
-            ],
-            response_body["lists"]["foo-list"],
-        )
-
-        # Make the Sliding Sync request (with lazy-loading room members)
+        # The list should include both rooms now because we don't need full state
+        for list_key in response_body["lists"].keys():
+            self.assertIncludes(
+                set(response_body["lists"][list_key]["ops"][0]["room_ids"]),
+                {room_id2, room_id1},
+                exact=True,
+                message=f"Expected all rooms to show up for list_key={list_key}. Response "
+                + str(response_body["lists"][list_key]),
+            )
+
+        # Make the Sliding Sync request with examples where `must_await_full_state()` is
+        # `True`
         sync_body = {
             "lists": {
-                "foo-list": {
+                "wildcard-list": {
+                    "ranges": [[0, 1]],
+                    "required_state": [
+                        ["*", "*"],
+                    ],
+                    "timeline_limit": 0,
+                },
+                "wildcard-type-remote-state-key-list": {
+                    "ranges": [[0, 1]],
+                    "required_state": [
+                        ["*", "@some:remote"],
+                        # Not a user ID
+                        ["*", "foobarbaz"],
+                        # Not a user ID
+                        ["*", "foo.bar.baz"],
+                        # Not a user ID
+                        ["*", "@foo"],
+                    ],
+                    "timeline_limit": 0,
+                },
+                "remote-member-list": {
+                    "ranges": [[0, 1]],
+                    "required_state": [
+                        # Own user ID
+                        [EventTypes.Member, user1_id],
+                        # Remote member
+                        [EventTypes.Member, "@some:remote"],
+                        # Local member
+                        [EventTypes.Member, user2_id],
+                    ],
+                    "timeline_limit": 0,
+                },
+                "lazy-but-remote-member-list": {
                     "ranges": [[0, 1]],
                     "required_state": [
-                        [EventTypes.Create, ""],
                         # Lazy-load room members
                         [EventTypes.Member, StateValues.LAZY],
+                        # Remote member
+                        [EventTypes.Member, "@some:remote"],
                     ],
                     "timeline_limit": 0,
                 },
@@ -693,15 +773,13 @@ class SlidingSyncRoomsRequiredStateTestCase(SlidingSyncBase):
         }
         response_body, _ = self.do_sync(sync_body, tok=user1_tok)
 
-        # The list should include both rooms now because we're lazy-loading room members
-        self.assertListEqual(
-            list(response_body["lists"]["foo-list"]["ops"]),
-            [
-                {
-                    "op": "SYNC",
-                    "range": [0, 1],
-                    "room_ids": [room_id2, room_id1],
-                }
-            ],
-            response_body["lists"]["foo-list"],
-        )
+        # Make sure the list includes room1 but room2 is excluded because it's still
+        # partially-stated
+        for list_key in response_body["lists"].keys():
+            self.assertIncludes(
+                set(response_body["lists"][list_key]["ops"][0]["room_ids"]),
+                {room_id1},
+                exact=True,
+                message=f"Expected only fully-stated rooms to show up for list_key={list_key}. Response "
+                + str(response_body["lists"][list_key]),
+            )