From fd52c4f3d672131fe666d1e7d5bae586c65084f9 Mon Sep 17 00:00:00 2001
From: "H. Shay" <hillerys@element.io>
Date: Mon, 16 May 2022 16:22:03 -0700
Subject: [PATCH 1/8] update get_current_state_ids and get_prev_state_ids to
 use state filter

---
 synapse/events/snapshot.py | 15 +++++++++++----
 synapse/storage/state.py   |  5 +++--
 tests/test_state.py        |  2 +-
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/synapse/events/snapshot.py b/synapse/events/snapshot.py
index 9ccd24b298bb..d507325465de 100644
--- a/synapse/events/snapshot.py
+++ b/synapse/events/snapshot.py
@@ -24,6 +24,7 @@
 if TYPE_CHECKING:
     from synapse.storage import Storage
     from synapse.storage.databases.main import DataStore
+    from synapse.storage.state import StateFilter
 
 
 @attr.s(slots=True, auto_attribs=True)
@@ -196,7 +197,7 @@ def state_group(self) -> Optional[int]:
 
         return self._state_group
 
-    async def get_current_state_ids(self) -> Optional[StateMap[str]]:
+    async def get_current_state_ids(self, state_filter: Optional["StateFilter"] = None) -> Optional[StateMap[str]]:
         """
         Gets the room state map, including this event - ie, the state in ``state_group``
 
@@ -204,6 +205,9 @@ async def get_current_state_ids(self) -> Optional[StateMap[str]]:
         not make it into the room state. This method will raise an exception if
         ``rejected`` is set.
 
+        Arg:
+           state_filter: specifies the type of state event to fetch from DB, example: EventTypes.JoinRules
+
         Returns:
             Returns None if state_group is None, which happens when the associated
             event is an outlier.
@@ -216,7 +220,7 @@ async def get_current_state_ids(self) -> Optional[StateMap[str]]:
 
         assert self._state_delta_due_to_event is not None
 
-        prev_state_ids = await self.get_prev_state_ids()
+        prev_state_ids = await self.get_prev_state_ids(state_filter)
 
         if self._state_delta_due_to_event:
             prev_state_ids = dict(prev_state_ids)
@@ -224,12 +228,15 @@ async def get_current_state_ids(self) -> Optional[StateMap[str]]:
 
         return prev_state_ids
 
-    async def get_prev_state_ids(self) -> StateMap[str]:
+    async def get_prev_state_ids(self, state_filter: Optional["StateFilter"] = None) -> StateMap[str]:
         """
         Gets the room state map, excluding this event.
 
         For a non-state event, this will be the same as get_current_state_ids().
 
+        Args:
+            state_filter: specifies the type of state event to fetch from DB, example: EventTypes.JoinRules
+
         Returns:
             Returns {} if state_group is None, which happens when the associated
             event is an outlier.
@@ -239,7 +246,7 @@ async def get_prev_state_ids(self) -> StateMap[str]:
         """
         assert self.state_group_before_event is not None
         return await self._storage.state.get_state_ids_for_group(
-            self.state_group_before_event
+            self.state_group_before_event, state_filter
         )
 
 
diff --git a/synapse/storage/state.py b/synapse/storage/state.py
index d4a1bd4f9d7d..fcd2bba6e688 100644
--- a/synapse/storage/state.py
+++ b/synapse/storage/state.py
@@ -593,16 +593,17 @@ async def get_state_groups_ids(
 
         return group_to_state
 
-    async def get_state_ids_for_group(self, state_group: int) -> StateMap[str]:
+    async def get_state_ids_for_group(self, state_group: int, state_filter: Optional[StateFilter] = None) -> StateMap[str]:
         """Get the event IDs of all the state in the given state group
 
         Args:
             state_group: A state group for which we want to get the state IDs.
+            state_filter: specifies the type of state event to fetch from DB, example: EventTypes.JoinRules
 
         Returns:
             Resolves to a map of (type, state_key) -> event_id
         """
-        group_to_state = await self._get_state_for_groups((state_group,))
+        group_to_state = await self._get_state_for_groups((state_group,), state_filter)
 
         return group_to_state[state_group]
 
diff --git a/tests/test_state.py b/tests/test_state.py
index 651ec1c7d4bd..f462b14a65e5 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -88,7 +88,7 @@ async def get_state_groups_ids(self, room_id, event_ids):
 
         return groups
 
-    async def get_state_ids_for_group(self, state_group):
+    async def get_state_ids_for_group(self, state_group, state_filter = None):
         return self._group_to_state[state_group]
 
     async def store_state_group(

From a179643b300be2cb8772f9e41fe0821d252f654c Mon Sep 17 00:00:00 2001
From: "H. Shay" <hillerys@element.io>
Date: Mon, 16 May 2022 16:34:23 -0700
Subject: [PATCH 2/8] update calls to get_prev_state_ids to use state filter

---
 synapse/handlers/federation.py  | 9 +++++++--
 synapse/handlers/message.py     | 8 ++++++--
 synapse/handlers/room_member.py | 9 +++++++--
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/synapse/handlers/federation.py b/synapse/handlers/federation.py
index be5099b507f6..0386d0a07bba 100644
--- a/synapse/handlers/federation.py
+++ b/synapse/handlers/federation.py
@@ -54,6 +54,7 @@
     ReplicationStoreRoomOnOutlierMembershipRestServlet,
 )
 from synapse.storage.databases.main.events_worker import EventRedactBehaviour
+from synapse.storage.state import StateFilter
 from synapse.types import JsonDict, StateMap, get_domain_from_id
 from synapse.util.async_helpers import Linearizer
 from synapse.util.retryutils import NotRetryingDestination
@@ -1259,7 +1260,9 @@ async def add_display_name_to_third_party_invite(
             event.content["third_party_invite"]["signed"]["token"],
         )
         original_invite = None
-        prev_state_ids = await context.get_prev_state_ids()
+        prev_state_ids = await context.get_prev_state_ids(
+            StateFilter.from_types([(EventTypes.ThirdPartyInvite, None)])
+        )
         original_invite_id = prev_state_ids.get(key)
         if original_invite_id:
             original_invite = await self.store.get_event(
@@ -1308,7 +1311,9 @@ async def _check_signature(self, event: EventBase, context: EventContext) -> Non
         signed = event.content["third_party_invite"]["signed"]
         token = signed["token"]
 
-        prev_state_ids = await context.get_prev_state_ids()
+        prev_state_ids = await context.get_prev_state_ids(
+            StateFilter.from_types([(EventTypes.ThirdPartyInvite, None)])
+        )
         invite_event_id = prev_state_ids.get((EventTypes.ThirdPartyInvite, token))
 
         invite_event = None
diff --git a/synapse/handlers/message.py b/synapse/handlers/message.py
index 0951b9c71f75..c4720ab313aa 100644
--- a/synapse/handlers/message.py
+++ b/synapse/handlers/message.py
@@ -634,7 +634,9 @@ async def create_event(
             # federation as well as those created locally. As of room v3, aliases events
             # can be created by users that are not in the room, therefore we have to
             # tolerate them in event_auth.check().
-            prev_state_ids = await context.get_prev_state_ids()
+            prev_state_ids = await context.get_prev_state_ids(
+                StateFilter.from_types([(EventTypes.Member, None)])
+            )
             prev_event_id = prev_state_ids.get((EventTypes.Member, event.sender))
             prev_event = (
                 await self.store.get_event(prev_event_id, allow_none=True)
@@ -761,7 +763,9 @@ async def deduplicate_state_event(
             # This can happen due to out of band memberships
             return None
 
-        prev_state_ids = await context.get_prev_state_ids()
+        prev_state_ids = await context.get_prev_state_ids(
+            StateFilter.from_types([(event.type, None)])
+        )
         prev_event_id = prev_state_ids.get((event.type, event.state_key))
         if not prev_event_id:
             return None
diff --git a/synapse/handlers/room_member.py b/synapse/handlers/room_member.py
index 802e57c4d0cc..ea876c168de7 100644
--- a/synapse/handlers/room_member.py
+++ b/synapse/handlers/room_member.py
@@ -38,6 +38,7 @@
 from synapse.events import EventBase
 from synapse.events.snapshot import EventContext
 from synapse.handlers.profile import MAX_AVATAR_URL_LEN, MAX_DISPLAYNAME_LEN
+from synapse.storage.state import StateFilter
 from synapse.types import (
     JsonDict,
     Requester,
@@ -362,7 +363,9 @@ async def _local_membership_update(
             historical=historical,
         )
 
-        prev_state_ids = await context.get_prev_state_ids()
+        prev_state_ids = await context.get_prev_state_ids(
+            StateFilter.from_types([(EventTypes.Member, None)])
+        )
 
         prev_member_event_id = prev_state_ids.get((EventTypes.Member, user_id), None)
 
@@ -1160,7 +1163,9 @@ async def send_membership_event(
         else:
             requester = types.create_requester(target_user)
 
-        prev_state_ids = await context.get_prev_state_ids()
+        prev_state_ids = await context.get_prev_state_ids(
+            StateFilter.from_types([(EventTypes.GuestAccess, None)])
+        )
         if event.membership == Membership.JOIN:
             if requester.is_guest:
                 guest_can_join = await self._can_guest_join(prev_state_ids)

From 545bdfaccc98c7b0e21a319a9d30929b4a2f4e97 Mon Sep 17 00:00:00 2001
From: "H. Shay" <hillerys@element.io>
Date: Mon, 16 May 2022 16:34:38 -0700
Subject: [PATCH 3/8] lint

---
 synapse/events/snapshot.py | 8 ++++++--
 synapse/storage/state.py   | 4 +++-
 tests/test_state.py        | 2 +-
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/synapse/events/snapshot.py b/synapse/events/snapshot.py
index d507325465de..7a91544119f7 100644
--- a/synapse/events/snapshot.py
+++ b/synapse/events/snapshot.py
@@ -197,7 +197,9 @@ def state_group(self) -> Optional[int]:
 
         return self._state_group
 
-    async def get_current_state_ids(self, state_filter: Optional["StateFilter"] = None) -> Optional[StateMap[str]]:
+    async def get_current_state_ids(
+        self, state_filter: Optional["StateFilter"] = None
+    ) -> Optional[StateMap[str]]:
         """
         Gets the room state map, including this event - ie, the state in ``state_group``
 
@@ -228,7 +230,9 @@ async def get_current_state_ids(self, state_filter: Optional["StateFilter"] = No
 
         return prev_state_ids
 
-    async def get_prev_state_ids(self, state_filter: Optional["StateFilter"] = None) -> StateMap[str]:
+    async def get_prev_state_ids(
+        self, state_filter: Optional["StateFilter"] = None
+    ) -> StateMap[str]:
         """
         Gets the room state map, excluding this event.
 
diff --git a/synapse/storage/state.py b/synapse/storage/state.py
index fcd2bba6e688..9e1d31b3f659 100644
--- a/synapse/storage/state.py
+++ b/synapse/storage/state.py
@@ -593,7 +593,9 @@ async def get_state_groups_ids(
 
         return group_to_state
 
-    async def get_state_ids_for_group(self, state_group: int, state_filter: Optional[StateFilter] = None) -> StateMap[str]:
+    async def get_state_ids_for_group(
+        self, state_group: int, state_filter: Optional[StateFilter] = None
+    ) -> StateMap[str]:
         """Get the event IDs of all the state in the given state group
 
         Args:
diff --git a/tests/test_state.py b/tests/test_state.py
index f462b14a65e5..8c621d216c20 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -88,7 +88,7 @@ async def get_state_groups_ids(self, room_id, event_ids):
 
         return groups
 
-    async def get_state_ids_for_group(self, state_group, state_filter = None):
+    async def get_state_ids_for_group(self, state_group, state_filter=None):
         return self._group_to_state[state_group]
 
     async def store_state_group(

From d7fc1bea65ae7c387f717211fc8309d5edfbb131 Mon Sep 17 00:00:00 2001
From: "H. Shay" <hillerys@element.io>
Date: Wed, 18 May 2022 12:52:31 -0700
Subject: [PATCH 4/8] add state filter to room.py and roommember.py

---
 synapse/handlers/room.py                     | 3 ++-
 synapse/storage/databases/main/roommember.py | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/synapse/handlers/room.py b/synapse/handlers/room.py
index a2973109adc4..f6a9b5845b1c 100644
--- a/synapse/handlers/room.py
+++ b/synapse/handlers/room.py
@@ -303,7 +303,8 @@ async def _upgrade_room(
             context=tombstone_context,
         )
 
-        old_room_state = await tombstone_context.get_current_state_ids()
+        state_filter = StateFilter.from_types([(EventTypes.CanonicalAlias, ""), (EventTypes.PowerLevels, "")])
+        old_room_state = await tombstone_context.get_current_state_ids(state_filter)
 
         # We know the tombstone event isn't an outlier so it has current state.
         assert old_room_state is not None
diff --git a/synapse/storage/databases/main/roommember.py b/synapse/storage/databases/main/roommember.py
index 48e83592e728..4e8fc880ce5b 100644
--- a/synapse/storage/databases/main/roommember.py
+++ b/synapse/storage/databases/main/roommember.py
@@ -46,6 +46,7 @@
     ProfileInfo,
     RoomsForUser,
 )
+from synapse.storage.state import StateFilter
 from synapse.types import PersistedEventPosition, get_domain_from_id
 from synapse.util.async_helpers import Linearizer
 from synapse.util.caches import intern_string

From d3ec2797cf99d2ae11524a5a4ae47317ba8c42f8 Mon Sep 17 00:00:00 2001
From: "H. Shay" <hillerys@element.io>
Date: Wed, 18 May 2022 14:06:34 -0700
Subject: [PATCH 5/8] lint

---
 synapse/handlers/room.py                     | 4 +++-
 synapse/storage/databases/main/roommember.py | 1 -
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/synapse/handlers/room.py b/synapse/handlers/room.py
index f6a9b5845b1c..1516d5326060 100644
--- a/synapse/handlers/room.py
+++ b/synapse/handlers/room.py
@@ -303,7 +303,9 @@ async def _upgrade_room(
             context=tombstone_context,
         )
 
-        state_filter = StateFilter.from_types([(EventTypes.CanonicalAlias, ""), (EventTypes.PowerLevels, "")])
+        state_filter = StateFilter.from_types(
+            [(EventTypes.CanonicalAlias, ""), (EventTypes.PowerLevels, "")]
+        )
         old_room_state = await tombstone_context.get_current_state_ids(state_filter)
 
         # We know the tombstone event isn't an outlier so it has current state.
diff --git a/synapse/storage/databases/main/roommember.py b/synapse/storage/databases/main/roommember.py
index 4e8fc880ce5b..48e83592e728 100644
--- a/synapse/storage/databases/main/roommember.py
+++ b/synapse/storage/databases/main/roommember.py
@@ -46,7 +46,6 @@
     ProfileInfo,
     RoomsForUser,
 )
-from synapse.storage.state import StateFilter
 from synapse.types import PersistedEventPosition, get_domain_from_id
 from synapse.util.async_helpers import Linearizer
 from synapse.util.caches import intern_string

From e80c6633cfd288b3087bb69670b06577ab284ce3 Mon Sep 17 00:00:00 2001
From: "H. Shay" <hillerys@element.io>
Date: Wed, 18 May 2022 14:10:37 -0700
Subject: [PATCH 6/8] newsfragment

---
 changelog.d/12791.misc | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 changelog.d/12791.misc

diff --git a/changelog.d/12791.misc b/changelog.d/12791.misc
new file mode 100644
index 000000000000..b6e92b7eafad
--- /dev/null
+++ b/changelog.d/12791.misc
@@ -0,0 +1 @@
+Update EventContext `get_current_event_ids` and `get_prev_event_ids` to accept state filters and update calls where possible.

From 2762c6656af14eb620fe267540640a36796c7333 Mon Sep 17 00:00:00 2001
From: "H. Shay" <hillerys@element.io>
Date: Wed, 18 May 2022 20:30:05 -0700
Subject: [PATCH 7/8] use state filter before compute_auth_events call

---
 synapse/handlers/federation_event.py     | 5 ++++-
 synapse/handlers/message.py              | 4 +++-
 synapse/push/bulk_push_rule_evaluator.py | 6 +++++-
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/synapse/handlers/federation_event.py b/synapse/handlers/federation_event.py
index 761caa04b726..eae47f46e1db 100644
--- a/synapse/handlers/federation_event.py
+++ b/synapse/handlers/federation_event.py
@@ -63,6 +63,7 @@
 )
 from synapse.state import StateResolutionStore
 from synapse.storage.databases.main.events_worker import EventRedactBehaviour
+from synapse.storage.state import StateFilter
 from synapse.types import (
     PersistedEventPosition,
     RoomStreamToken,
@@ -74,6 +75,7 @@
 from synapse.util.iterutils import batch_iter
 from synapse.util.retryutils import NotRetryingDestination
 from synapse.util.stringutils import shortstr
+from synapse import event_auth
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
@@ -1500,7 +1502,8 @@ async def _check_event_auth(
             return context
 
         # now check auth against what we think the auth events *should* be.
-        prev_state_ids = await context.get_prev_state_ids()
+        event_types = event_auth.auth_types_for_event(event.room_version, event)
+        prev_state_ids = await context.get_prev_state_ids(StateFilter.from_types(event_types))
         auth_events_ids = self._event_auth_handler.compute_auth_events(
             event, prev_state_ids, for_verification=True
         )
diff --git a/synapse/handlers/message.py b/synapse/handlers/message.py
index c4720ab313aa..36edc950b12e 100644
--- a/synapse/handlers/message.py
+++ b/synapse/handlers/message.py
@@ -1551,7 +1551,9 @@ async def persist_and_notify_client_event(
                         "Redacting MSC2716 events is not supported in this room version",
                     )
 
-            prev_state_ids = await context.get_prev_state_ids()
+            event_types = event_auth.auth_types_for_event(event.room_version, event)
+            prev_state_ids = await context.get_prev_state_ids(StateFilter.from_types(event_types))
+
             auth_events_ids = self._event_auth_handler.compute_auth_events(
                 event, prev_state_ids, for_verification=True
             )
diff --git a/synapse/push/bulk_push_rule_evaluator.py b/synapse/push/bulk_push_rule_evaluator.py
index 4ac2c546bf2a..a0be3fda486d 100644
--- a/synapse/push/bulk_push_rule_evaluator.py
+++ b/synapse/push/bulk_push_rule_evaluator.py
@@ -32,6 +32,8 @@
 from synapse.util.metrics import measure_func
 
 from .push_rule_evaluator import PushRuleEvaluatorForEvent
+from ..handlers import event_auth
+from ..storage.state import StateFilter
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
@@ -168,8 +170,10 @@ def _get_rules_for_room(self, room_id: str) -> "RulesForRoomData":
     async def _get_power_levels_and_sender_level(
         self, event: EventBase, context: EventContext
     ) -> Tuple[dict, int]:
-        prev_state_ids = await context.get_prev_state_ids()
+        event_types = event_auth.auth_types_for_event(event.room_version, event)
+        prev_state_ids = await context.get_prev_state_ids(StateFilter.from_types(event_types))
         pl_event_id = prev_state_ids.get(POWER_KEY)
+
         if pl_event_id:
             # fastpath: if there's a power level event, that's all we need, and
             # not having a power level event is an extreme edge case

From 79e47a30aa99705537180e0c7e66187e407c02cc Mon Sep 17 00:00:00 2001
From: "H. Shay" <hillerys@element.io>
Date: Thu, 19 May 2022 10:17:14 -0700
Subject: [PATCH 8/8] fix + lints

---
 synapse/handlers/federation_event.py     |  7 +++++--
 synapse/handlers/message.py              |  4 +++-
 synapse/push/bulk_push_rule_evaluator.py | 11 ++++++-----
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/synapse/handlers/federation_event.py b/synapse/handlers/federation_event.py
index eae47f46e1db..05c122f22491 100644
--- a/synapse/handlers/federation_event.py
+++ b/synapse/handlers/federation_event.py
@@ -30,6 +30,7 @@
 
 from prometheus_client import Counter
 
+from synapse import event_auth
 from synapse.api.constants import (
     EventContentFields,
     EventTypes,
@@ -75,7 +76,6 @@
 from synapse.util.iterutils import batch_iter
 from synapse.util.retryutils import NotRetryingDestination
 from synapse.util.stringutils import shortstr
-from synapse import event_auth
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
@@ -1503,7 +1503,10 @@ async def _check_event_auth(
 
         # now check auth against what we think the auth events *should* be.
         event_types = event_auth.auth_types_for_event(event.room_version, event)
-        prev_state_ids = await context.get_prev_state_ids(StateFilter.from_types(event_types))
+        prev_state_ids = await context.get_prev_state_ids(
+            StateFilter.from_types(event_types)
+        )
+
         auth_events_ids = self._event_auth_handler.compute_auth_events(
             event, prev_state_ids, for_verification=True
         )
diff --git a/synapse/handlers/message.py b/synapse/handlers/message.py
index 36edc950b12e..e566ff1f8ed8 100644
--- a/synapse/handlers/message.py
+++ b/synapse/handlers/message.py
@@ -1552,7 +1552,9 @@ async def persist_and_notify_client_event(
                     )
 
             event_types = event_auth.auth_types_for_event(event.room_version, event)
-            prev_state_ids = await context.get_prev_state_ids(StateFilter.from_types(event_types))
+            prev_state_ids = await context.get_prev_state_ids(
+                StateFilter.from_types(event_types)
+            )
 
             auth_events_ids = self._event_auth_handler.compute_auth_events(
                 event, prev_state_ids, for_verification=True
diff --git a/synapse/push/bulk_push_rule_evaluator.py b/synapse/push/bulk_push_rule_evaluator.py
index a0be3fda486d..4cc8a2ecca7a 100644
--- a/synapse/push/bulk_push_rule_evaluator.py
+++ b/synapse/push/bulk_push_rule_evaluator.py
@@ -20,7 +20,7 @@
 from prometheus_client import Counter
 
 from synapse.api.constants import EventTypes, Membership, RelationTypes
-from synapse.event_auth import get_user_power_level
+from synapse.event_auth import auth_types_for_event, get_user_power_level
 from synapse.events import EventBase, relation_from_event
 from synapse.events.snapshot import EventContext
 from synapse.state import POWER_KEY
@@ -31,9 +31,8 @@
 from synapse.util.caches.lrucache import LruCache
 from synapse.util.metrics import measure_func
 
-from .push_rule_evaluator import PushRuleEvaluatorForEvent
-from ..handlers import event_auth
 from ..storage.state import StateFilter
+from .push_rule_evaluator import PushRuleEvaluatorForEvent
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
@@ -170,8 +169,10 @@ def _get_rules_for_room(self, room_id: str) -> "RulesForRoomData":
     async def _get_power_levels_and_sender_level(
         self, event: EventBase, context: EventContext
     ) -> Tuple[dict, int]:
-        event_types = event_auth.auth_types_for_event(event.room_version, event)
-        prev_state_ids = await context.get_prev_state_ids(StateFilter.from_types(event_types))
+        event_types = auth_types_for_event(event.room_version, event)
+        prev_state_ids = await context.get_prev_state_ids(
+            StateFilter.from_types(event_types)
+        )
         pl_event_id = prev_state_ids.get(POWER_KEY)
 
         if pl_event_id: