matrix-synapse 1.142.0rc3__cp314-abi3-musllinux_1_2_aarch64.whl

synapse/state/v2.py

@@ -0,0 +1,985 @@
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright (C) 2023 New Vector, Ltd
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# See the GNU Affero General Public License for more details:
+# <https://www.gnu.org/licenses/agpl-3.0.html>.
+#
+# Originally licensed under the Apache License, Version 2.0:
+# <http://www.apache.org/licenses/LICENSE-2.0>.
+#
+# [This file includes modifications made by New Vector Limited]
+#
+#
+
+import heapq
+import itertools
+import logging
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Generator,
+    Iterable,
+    Literal,
+    Optional,
+    Protocol,
+    Sequence,
+    overload,
+)
+
+from synapse import event_auth
+from synapse.api.constants import CREATOR_POWER_LEVEL, EventTypes
+from synapse.api.errors import AuthError
+from synapse.api.room_versions import RoomVersion, StateResolutionVersions
+from synapse.events import EventBase, is_creator
+from synapse.storage.databases.main.event_federation import StateDifference
+from synapse.types import MutableStateMap, StateMap, StrCollection
+
+logger = logging.getLogger(__name__)
+
+
+class Clock(Protocol):
+    # This is usually synapse.util.Clock, but it's replaced with a FakeClock in tests.
+    # We only ever sleep(0) though, so that other async functions can make forward
+    # progress without waiting for stateres to complete.
+    async def sleep(self, duration_ms: float) -> None: ...
+
+
+class StateResolutionStore(Protocol):
+    # This is usually synapse.state.StateResolutionStore, but it's replaced with a
+    # TestStateResolutionStore in tests.
+    def get_events(
+        self, event_ids: StrCollection, allow_rejected: bool = False
+    ) -> Awaitable[dict[str, EventBase]]: ...
+
+    def get_auth_chain_difference(
+        self,
+        room_id: str,
+        state_sets: list[set[str]],
+        conflicted_state: Optional[set[str]],
+        additional_backwards_reachable_conflicted_events: Optional[set[str]],
+    ) -> Awaitable[StateDifference]: ...
+
+
+# We want to await to the reactor occasionally during state res when dealing
+# with large data sets, so that we don't exhaust the reactor. This is done by
+# awaiting to reactor during loops every N iterations.
+_AWAIT_AFTER_ITERATIONS = 100
+
+
+__all__ = [
+    "resolve_events_with_store",
+]
+
+
+async def resolve_events_with_store(
+    clock: Clock,
+    room_id: str,
+    room_version: RoomVersion,
+    state_sets: Sequence[StateMap[str]],
+    event_map: Optional[dict[str, EventBase]],
+    state_res_store: StateResolutionStore,
+) -> StateMap[str]:
+    """Resolves the state using the v2 state resolution algorithm
+
+    Args:
+        clock
+        room_id: the room we are working in
+        room_version: The room version
+        state_sets: List of dicts of (type, state_key) -> event_id,
+            which are the different state groups to resolve.
+        event_map:
+            a dict from event_id to event, for any events that we happen to
+            have in flight (eg, those currently being persisted). This will be
+            used as a starting point for finding the state we need; any missing
+            events will be requested via state_res_store.
+
+            If None, all events will be fetched via state_res_store.
+
+        state_res_store:
+
+    Returns:
+        A map from (type, state_key) to event_id.
+    """
+
+    logger.debug("Computing conflicted state")
+
+    # We use event_map as a cache, so if its None we need to initialize it
+    if event_map is None:
+        event_map = {}
+
+    # First split up the un/conflicted state
+    unconflicted_state, conflicted_state = _seperate(state_sets)
+
+    if not conflicted_state:
+        return unconflicted_state
+
+    logger.debug("%d conflicted state entries", len(conflicted_state))
+    logger.debug("Calculating auth chain difference")
+
+    conflicted_set: Optional[set[str]] = None
+    if room_version.state_res == StateResolutionVersions.V2_1:
+        # calculate the conflicted subgraph
+        conflicted_set = set(itertools.chain.from_iterable(conflicted_state.values()))
+    auth_diff = await _get_auth_chain_difference(
+        room_id,
+        state_sets,
+        event_map,
+        state_res_store,
+        conflicted_set,
+    )
+    full_conflicted_set = set(
+        itertools.chain(
+            itertools.chain.from_iterable(conflicted_state.values()), auth_diff
+        )
+    )
+
+    events = await state_res_store.get_events(
+        [eid for eid in full_conflicted_set if eid not in event_map],
+        allow_rejected=True,
+    )
+    event_map.update(events)
+
+    # everything in the event map should be in the right room
+    for event in event_map.values():
+        if event.room_id != room_id:
+            raise Exception(
+                "Attempting to state-resolve for room %s with event %s which is in %s"
+                % (
+                    room_id,
+                    event.event_id,
+                    event.room_id,
+                )
+            )
+
+    full_conflicted_set = {eid for eid in full_conflicted_set if eid in event_map}
+
+    logger.debug("%d full_conflicted_set entries", len(full_conflicted_set))
+
+    # Get and sort all the power events (kicks/bans/etc)
+    power_events = (
+        eid for eid in full_conflicted_set if _is_power_event(event_map[eid])
+    )
+
+    sorted_power_events = await _reverse_topological_power_sort(
+        clock, room_id, power_events, event_map, state_res_store, full_conflicted_set
+    )
+
+    logger.debug("sorted %d power events", len(sorted_power_events))
+
+    # v2.1 starts iterative auth checks from the empty set and not the unconflicted state.
+    # It relies on IAC behaviour which populates the base state with the events from auth_events
+    # if the state tuple is missing from the base state. This ensures the base state is only
+    # populated from auth_events rather than whatever the unconflicted state is (which could be
+    # completely bogus).
+    base_state = (
+        {}
+        if room_version.state_res == StateResolutionVersions.V2_1
+        else unconflicted_state
+    )
+
+    # Now sequentially auth each one
+    resolved_state = await _iterative_auth_checks(
+        clock,
+        room_id,
+        room_version,
+        event_ids=sorted_power_events,
+        base_state=base_state,
+        event_map=event_map,
+        state_res_store=state_res_store,
+    )
+
+    logger.debug("resolved power events")
+
+    # OK, so we've now resolved the power events. Now sort the remaining
+    # events using the mainline of the resolved power level.
+
+    set_power_events = set(sorted_power_events)
+    leftover_events = [
+        ev_id for ev_id in full_conflicted_set if ev_id not in set_power_events
+    ]
+
+    logger.debug("sorting %d remaining events", len(leftover_events))
+
+    pl = resolved_state.get((EventTypes.PowerLevels, ""), None)
+    leftover_events = await _mainline_sort(
+        clock, room_id, leftover_events, pl, event_map, state_res_store
+    )
+
+    logger.debug("resolving remaining events")
+
+    resolved_state = await _iterative_auth_checks(
+        clock,
+        room_id,
+        room_version,
+        leftover_events,
+        resolved_state,
+        event_map,
+        state_res_store,
+    )
+
+    logger.debug("resolved")
+
+    # We make sure that unconflicted state always still applies.
+    resolved_state.update(unconflicted_state)
+
+    logger.debug("done")
+
+    return resolved_state
+
+
+async def _get_power_level_for_sender(
+    room_id: str,
+    event_id: str,
+    event_map: dict[str, EventBase],
+    state_res_store: StateResolutionStore,
+) -> int:
+    """Return the power level of the sender of the given event according to
+    their auth events.
+
+    Args:
+        room_id
+        event_id
+        event_map
+        state_res_store
+
+    Returns:
+        The power level.
+    """
+    event = await _get_event(room_id, event_id, event_map, state_res_store)
+
+    pl = None
+    create = None
+    for aid in event.auth_event_ids():
+        aev = await _get_event(
+            room_id, aid, event_map, state_res_store, allow_none=True
+        )
+        if aev and (aev.type, aev.state_key) == (EventTypes.PowerLevels, ""):
+            pl = aev
+        if aev and (aev.type, aev.state_key) == (EventTypes.Create, ""):
+            create = aev
+
+    if event.type != EventTypes.Create:
+        # we should always have a create event
+        assert create is not None
+
+    if create and create.room_version.msc4289_creator_power_enabled:
+        if is_creator(create, event.sender):
+            return CREATOR_POWER_LEVEL
+
+    if pl is None:
+        # Couldn't find power level. Check if they're the creator of the room
+        for aid in event.auth_event_ids():
+            aev = await _get_event(
+                room_id, aid, event_map, state_res_store, allow_none=True
+            )
+            if aev and (aev.type, aev.state_key) == (EventTypes.Create, ""):
+                creator = (
+                    aev.sender
+                    if event.room_version.implicit_room_creator
+                    else aev.content.get("creator")
+                )
+                if not creator:
+                    logger.warning(
+                        "_get_power_level_for_sender: event %s has no PL in auth_events and "
+                        "creator is missing from create event %s",
+                        event_id,
+                        aev.event_id,
+                    )
+                if creator == event.sender:
+                    return 100
+                break
+        return 0
+
+    level = pl.content.get("users", {}).get(event.sender)
+    if level is None:
+        level = pl.content.get("users_default", 0)
+
+    if level is None:
+        return 0
+    else:
+        return int(level)
+
+
+async def _get_auth_chain_difference(
+    room_id: str,
+    state_sets: Sequence[StateMap[str]],
+    unpersisted_events: dict[str, EventBase],
+    state_res_store: StateResolutionStore,
+    conflicted_state: Optional[set[str]],
+) -> set[str]:
+    """Compare the auth chains of each state set and return the set of events
+    that only appear in some, but not all of the auth chains.
+
+    Args:
+        state_sets: The input state sets we are trying to resolve across.
+        unpersisted_events: A map from event ID to EventBase containing all unpersisted
+            events involved in this resolution.
+        state_res_store: A way to retrieve events and extract graph information on the auth chains.
+        conflicted_state: which event IDs are conflicted. Used in v2.1 for calculating the conflicted
+            subgraph.
+
+    Returns:
+        The auth difference of the given state sets, as a set of event IDs. Also includes the
+        conflicted subgraph if `conflicted_state` is set.
+    """
+    is_state_res_v21 = conflicted_state is not None
+    num_conflicted_state = (
+        len(conflicted_state) if conflicted_state is not None else None
+    )
+
+    # The `StateResolutionStore.get_auth_chain_difference` function assumes that
+    # all events passed to it (and their auth chains) have been persisted
+    # previously. We need to manually handle any other events that are yet to be
+    # persisted.
+    #
+    # We do this in three steps:
+    #   1. Compute the set of unpersisted events belonging to the auth difference.
+    #   2. Replacing any unpersisted events in the state_sets with their auth events,
+    #      recursively, until the state_sets contain only persisted events.
+    #      Then we call `store.get_auth_chain_difference` as normal, which computes
+    #      the set of persisted events belonging to the auth difference.
+    #   3. Adding the results of 1 and 2 together.
+
+    # Map from event ID in `unpersisted_events` to their auth event IDs, and their auth
+    # event IDs if they appear in the `unpersisted_events`. This is the intersection of
+    # the event's auth chain with the events in `unpersisted_events` *plus* their
+    # auth event IDs.
+    events_to_auth_chain: dict[str, set[str]] = {}
+    # remember the forward links when doing the graph traversal, we'll need it for v2.1 checks
+    # This is a map from an event to the set of events that contain it as an auth event.
+    event_to_next_event: dict[str, set[str]] = {}
+    for event in unpersisted_events.values():
+        chain = {event.event_id}
+        events_to_auth_chain[event.event_id] = chain
+
+        to_search = [event]
+        while to_search:
+            next_event = to_search.pop()
+            for auth_id in next_event.auth_event_ids():
+                chain.add(auth_id)
+                event_to_next_event.setdefault(auth_id, set()).add(next_event.event_id)
+                auth_event = unpersisted_events.get(auth_id)
+                if auth_event:
+                    to_search.append(auth_event)
+
+    # We now 1) calculate the auth chain difference for the unpersisted events
+    # and 2) work out the state sets to pass to the store.
+    #
+    # Note: If there are no `unpersisted_events` (which is the common case), we can do a
+    # much simpler calculation.
+    additional_backwards_reachable_conflicted_events: set[str] = set()
+    unpersisted_conflicted_events: set[str] = set()
+    if unpersisted_events:
+        # The list of state sets to pass to the store, where each state set is a set
+        # of the event ids making up the state. This is similar to `state_sets`,
+        # except that (a) we only have event ids, not the complete
+        # ((type, state_key)->event_id) mappings; and (b) we have stripped out
+        # unpersisted events and replaced them with the persisted events in
+        # their auth chain.
+        state_sets_ids: list[set[str]] = []
+
+        # For each state set, the unpersisted event IDs reachable (by their auth
+        # chain) from the events in that set.
+        unpersisted_set_ids: list[set[str]] = []
+
+        for state_set in state_sets:
+            set_ids: set[str] = set()
+            state_sets_ids.append(set_ids)
+
+            unpersisted_ids: set[str] = set()
+            unpersisted_set_ids.append(unpersisted_ids)
+
+            for event_id in state_set.values():
+                event_chain = events_to_auth_chain.get(event_id)
+                if event_chain is not None:
+                    # We have an unpersisted event. We add all the auth
+                    # events that it references which are also unpersisted.
+                    set_ids.update(
+                        e for e in event_chain if e not in unpersisted_events
+                    )
+
+                    # We also add the full chain of unpersisted event IDs
+                    # referenced by this state set, so that we can work out the
+                    # auth chain difference of the unpersisted events.
+                    unpersisted_ids.update(
+                        e for e in event_chain if e in unpersisted_events
+                    )
+                else:
+                    set_ids.add(event_id)
+        if conflicted_state:
+            for conflicted_event_id in conflicted_state:
+                # presence in this map means it is unpersisted.
+                event_chain = events_to_auth_chain.get(conflicted_event_id)
+                if event_chain is not None:
+                    unpersisted_conflicted_events.add(conflicted_event_id)
+                    # tell the DB layer that we have some unpersisted conflicted events
+                    additional_backwards_reachable_conflicted_events.update(
+                        e for e in event_chain if e not in unpersisted_events
+                    )
+        # The auth chain difference of the unpersisted events of the state sets
+        # is calculated by taking the difference between the union and
+        # intersections.
+        union = unpersisted_set_ids[0].union(*unpersisted_set_ids[1:])
+        intersection = unpersisted_set_ids[0].intersection(*unpersisted_set_ids[1:])
+
+        auth_difference_unpersisted_part: StrCollection = union - intersection
+    else:
+        auth_difference_unpersisted_part = ()
+        state_sets_ids = [set(state_set.values()) for state_set in state_sets]
+
+    if conflicted_state:
+        # to ensure that conflicted state is a subset of state set IDs, we need to remove UNPERSISTED
+        # conflicted state set ids as we removed them above.
+        conflicted_state = conflicted_state - unpersisted_conflicted_events
+
+    difference = await state_res_store.get_auth_chain_difference(
+        room_id,
+        state_sets_ids,
+        conflicted_state,
+        additional_backwards_reachable_conflicted_events,
+    )
+    difference.auth_difference.update(auth_difference_unpersisted_part)
+
+    # if we're doing v2.1 we may need to add or expand the conflicted subgraph
+    if (
+        is_state_res_v21
+        and difference.conflicted_subgraph is not None
+        and unpersisted_events
+    ):
+        # we always include the conflicted events themselves in the subgraph.
+        if conflicted_state:
+            difference.conflicted_subgraph.update(conflicted_state)
+        # we may need to expand the subgraph in the case where the subgraph starts in the DB and
+        # ends in unpersisted events. To do this, we first need to see where the subgraph got up to,
+        # which we can do by finding the intersection between the additional backwards reachable
+        # conflicted events and the conflicted subgraph. Events in both sets mean A) some unpersisted
+        # conflicted event could backwards reach it and B) some persisted conflicted event could forward
+        # reach it.
+        subgraph_frontier = difference.conflicted_subgraph.intersection(
+            additional_backwards_reachable_conflicted_events
+        )
+        # we can now combine the 2 scenarios:
+        #  - subgraph starts in DB and ends in unpersisted
+        #  - subgraph starts in unpersisted and ends in unpersisted
+        # by expanding the frontier into unpersisted events.
+        # The frontier is currently all persisted events. We want to expand this into unpersisted
+        # events. Mark every forwards reachable event from the frontier in the forwards_conflicted_set
+        # but NOT the backwards conflicted set. This mirrors what the DB layer does but in reverse:
+        # we supplied events which are backwards reachable to the DB and now the DB is providing
+        # forwards reachable events from the DB.
+        forwards_conflicted_set: set[str] = set()
+        # we include unpersisted conflicted events here to process exclusive unpersisted subgraphs
+        search_queue = subgraph_frontier.union(unpersisted_conflicted_events)
+        while search_queue:
+            frontier_event = search_queue.pop()
+            next_event_ids = event_to_next_event.get(frontier_event, set())
+            search_queue.update(next_event_ids)
+            forwards_conflicted_set.add(frontier_event)
+
+        # we've already calculated the backwards form as this is the auth chain for each
+        # unpersisted conflicted event.
+        backwards_conflicted_set: set[str] = set()
+        for uce in unpersisted_conflicted_events:
+            backwards_conflicted_set.update(events_to_auth_chain.get(uce, []))
+
+        # the unpersisted conflicted subgraph is the intersection of the backwards/forwards sets
+        conflicted_subgraph_unpersisted_part = backwards_conflicted_set.intersection(
+            forwards_conflicted_set
+        )
+        # print(f"event_to_next_event={event_to_next_event}")
+        # print(f"unpersisted_conflicted_events={unpersisted_conflicted_events}")
+        # print(f"unperssited backwards_conflicted_set={backwards_conflicted_set}")
+        # print(f"unperssited forwards_conflicted_set={forwards_conflicted_set}")
+        difference.conflicted_subgraph.update(conflicted_subgraph_unpersisted_part)
+
+    if difference.conflicted_subgraph:
+        old_events = difference.auth_difference.union(
+            conflicted_state if conflicted_state else set()
+        )
+        additional_events = difference.conflicted_subgraph.difference(old_events)
+
+        logger.debug(
+            "v2.1 %s additional events replayed=%d num_conflicts=%d conflicted_subgraph=%d auth_difference=%d",
+            room_id,
+            len(additional_events),
+            num_conflicted_state,
+            len(difference.conflicted_subgraph),
+            len(difference.auth_difference),
+        )
+        # State res v2.1 includes the conflicted subgraph in the difference
+        return difference.auth_difference.union(difference.conflicted_subgraph)
+
+    return difference.auth_difference
+
+
+def _seperate(
+    state_sets: Iterable[StateMap[str]],
+) -> tuple[StateMap[str], StateMap[set[str]]]:
+    """Return the unconflicted and conflicted state. This is different than in
+    the original algorithm, as this defines a key to be conflicted if one of
+    the state sets doesn't have that key.
+
+    Args:
+        state_sets
+
+    Returns:
+        A tuple of unconflicted and conflicted state. The conflicted state dict
+        is a map from type/state_key to set of event IDs
+    """
+    unconflicted_state = {}
+    conflicted_state = {}
+
+    for key in set(itertools.chain.from_iterable(state_sets)):
+        event_ids = {state_set.get(key) for state_set in state_sets}
+        if len(event_ids) == 1:
+            unconflicted_state[key] = event_ids.pop()
+        else:
+            event_ids.discard(None)
+            conflicted_state[key] = event_ids
+
+    # mypy doesn't understand that discarding None above means that conflicted
+    # state is StateMap[set[str]], not StateMap[set[Optional[Str]]].
+    return unconflicted_state, conflicted_state  # type: ignore[return-value]
+
+
+def _is_power_event(event: EventBase) -> bool:
+    """Return whether or not the event is a "power event", as defined by the
+    v2 state resolution algorithm
+
+    Args:
+        event
+
+    Returns:
+        True if the event is a power event.
+    """
+    if (event.type, event.state_key) in (
+        (EventTypes.PowerLevels, ""),
+        (EventTypes.JoinRules, ""),
+        (EventTypes.Create, ""),
+    ):
+        return True
+
+    if event.type == EventTypes.Member:
+        if event.membership in ("leave", "ban"):
+            return event.sender != event.state_key
+
+    return False
+
+
+async def _add_event_and_auth_chain_to_graph(
+    graph: dict[str, set[str]],
+    room_id: str,
+    event_id: str,
+    event_map: dict[str, EventBase],
+    state_res_store: StateResolutionStore,
+    full_conflicted_set: set[str],
+) -> None:
+    """Helper function for _reverse_topological_power_sort that add the event
+    and its auth chain (that is in the auth diff) to the graph
+
+    Args:
+        graph: A map from event ID to the events auth event IDs
+        room_id: the room we are working in
+        event_id: Event to add to the graph
+        event_map
+        state_res_store
+        full_conflicted_set: Set of event IDs that are in the full conflicted set.
+    """
+
+    state = [event_id]
+    while state:
+        eid = state.pop()
+        graph.setdefault(eid, set())
+
+        event = await _get_event(room_id, eid, event_map, state_res_store)
+        for aid in event.auth_event_ids():
+            if aid in full_conflicted_set:
+                if aid not in graph:
+                    state.append(aid)
+
+                graph.setdefault(eid, set()).add(aid)
+
+
+async def _reverse_topological_power_sort(
+    clock: Clock,
+    room_id: str,
+    event_ids: Iterable[str],
+    event_map: dict[str, EventBase],
+    state_res_store: StateResolutionStore,
+    full_conflicted_set: set[str],
+) -> list[str]:
+    """Returns a list of the event_ids sorted by reverse topological ordering,
+    and then by power level and origin_server_ts
+
+    Args:
+        clock
+        room_id: the room we are working in
+        event_ids: The events to sort
+        event_map
+        state_res_store
+        full_conflicted_set: Set of event IDs that are in the full conflicted set.
+
+    Returns:
+        The sorted list
+    """
+
+    graph: dict[str, set[str]] = {}
+    for idx, event_id in enumerate(event_ids, start=1):
+        await _add_event_and_auth_chain_to_graph(
+            graph, room_id, event_id, event_map, state_res_store, full_conflicted_set
+        )
+
+        # We await occasionally when we're working with large data sets to
+        # ensure that we don't block the reactor loop for too long.
+        if idx % _AWAIT_AFTER_ITERATIONS == 0:
+            await clock.sleep(0)
+
+    event_to_pl = {}
+    for idx, event_id in enumerate(graph, start=1):
+        pl = await _get_power_level_for_sender(
+            room_id, event_id, event_map, state_res_store
+        )
+        event_to_pl[event_id] = pl
+
+        # We await occasionally when we're working with large data sets to
+        # ensure that we don't block the reactor loop for too long.
+        if idx % _AWAIT_AFTER_ITERATIONS == 0:
+            await clock.sleep(0)
+
+    def _get_power_order(event_id: str) -> tuple[int, int, str]:
+        ev = event_map[event_id]
+        pl = event_to_pl[event_id]
+
+        return -pl, ev.origin_server_ts, event_id
+
+    # Note: graph is modified during the sort
+    it = lexicographical_topological_sort(graph, key=_get_power_order)
+    sorted_events = list(it)
+
+    return sorted_events
+
+
+async def _iterative_auth_checks(
+    clock: Clock,
+    room_id: str,
+    room_version: RoomVersion,
+    event_ids: list[str],
+    base_state: StateMap[str],
+    event_map: dict[str, EventBase],
+    state_res_store: StateResolutionStore,
+) -> MutableStateMap[str]:
+    """Sequentially apply auth checks to each event in given list, updating the
+    state as it goes along.
+
+    Args:
+        clock
+        room_id
+        room_version
+        event_ids: Ordered list of events to apply auth checks to
+        base_state: The set of state to start with
+        event_map
+        state_res_store
+
+    Returns:
+        Returns the final updated state
+    """
+    resolved_state = dict(base_state)
+
+    for idx, event_id in enumerate(event_ids, start=1):
+        event = event_map[event_id]
+
+        auth_events = {}
+        for aid in event.auth_event_ids():
+            ev = await _get_event(
+                room_id, aid, event_map, state_res_store, allow_none=True
+            )
+
+            if not ev:
+                logger.warning(
+                    "auth_event id %s for event %s is missing", aid, event_id
+                )
+            else:
+                if ev.rejected_reason is None:
+                    auth_events[(ev.type, ev.state_key)] = ev
+
+        for key in event_auth.auth_types_for_event(room_version, event):
+            if key in resolved_state:
+                ev_id = resolved_state[key]
+                ev = await _get_event(room_id, ev_id, event_map, state_res_store)
+
+                if ev.rejected_reason is None:
+                    auth_events[key] = event_map[ev_id]
+
+        if event.rejected_reason is not None:
+            # Do not admit previously rejected events into state.
+            # TODO: This isn't spec compliant. Events that were previously rejected due
+            #       to failing auth checks at their state, but pass auth checks during
+            #       state resolution should be accepted. Synapse does not handle the
+            #       change of rejection status well, so we preserve the previous
+            #       rejection status for now.
+            #
+            #       Note that events rejected for non-state reasons, such as having the
+            #       wrong auth events, should remain rejected.
+            #
+            #       https://spec.matrix.org/v1.2/rooms/v9/#rejected-events
+            #       https://github.com/matrix-org/synapse/issues/13797
+            continue
+
+        try:
+            event_auth.check_state_dependent_auth_rules(
+                event,
+                auth_events.values(),
+            )
+
+            resolved_state[(event.type, event.state_key)] = event_id
+        except AuthError:
+            pass
+
+        # We await occasionally when we're working with large data sets to
+        # ensure that we don't block the reactor loop for too long.
+        if idx % _AWAIT_AFTER_ITERATIONS == 0:
+            await clock.sleep(0)
+
+    return resolved_state
+
+
+async def _mainline_sort(
+    clock: Clock,
+    room_id: str,
+    event_ids: list[str],
+    resolved_power_event_id: Optional[str],
+    event_map: dict[str, EventBase],
+    state_res_store: StateResolutionStore,
+) -> list[str]:
+    """Returns a sorted list of event_ids sorted by mainline ordering based on
+    the given event resolved_power_event_id
+
+    Args:
+        clock
+        room_id: room we're working in
+        event_ids: Events to sort
+        resolved_power_event_id: The final resolved power level event ID
+        event_map
+        state_res_store
+
+    Returns:
+        The sorted list
+    """
+    if not event_ids:
+        # It's possible for there to be no event IDs here to sort, so we can
+        # skip calculating the mainline in that case.
+        return []
+
+    mainline = []
+    pl = resolved_power_event_id
+    idx = 0
+    while pl:
+        mainline.append(pl)
+        pl_ev = await _get_event(room_id, pl, event_map, state_res_store)
+        auth_events = pl_ev.auth_event_ids()
+        pl = None
+        for aid in auth_events:
+            ev = await _get_event(
+                room_id, aid, event_map, state_res_store, allow_none=True
+            )
+            if ev and (ev.type, ev.state_key) == (EventTypes.PowerLevels, ""):
+                pl = aid
+                break
+
+        # We await occasionally when we're working with large data sets to
+        # ensure that we don't block the reactor loop for too long.
+        if idx != 0 and idx % _AWAIT_AFTER_ITERATIONS == 0:
+            await clock.sleep(0)
+
+        idx += 1
+
+    mainline_map = {ev_id: i + 1 for i, ev_id in enumerate(reversed(mainline))}
+
+    event_ids = list(event_ids)
+
+    order_map = {}
+    for idx, ev_id in enumerate(event_ids, start=1):
+        depth = await _get_mainline_depth_for_event(
+            clock, event_map[ev_id], mainline_map, event_map, state_res_store
+        )
+        order_map[ev_id] = (depth, event_map[ev_id].origin_server_ts, ev_id)
+
+        # We await occasionally when we're working with large data sets to
+        # ensure that we don't block the reactor loop for too long.
+        if idx % _AWAIT_AFTER_ITERATIONS == 0:
+            await clock.sleep(0)
+
+    event_ids.sort(key=lambda ev_id: order_map[ev_id])
+
+    return event_ids
+
+
+async def _get_mainline_depth_for_event(
+    clock: Clock,
+    event: EventBase,
+    mainline_map: dict[str, int],
+    event_map: dict[str, EventBase],
+    state_res_store: StateResolutionStore,
+) -> int:
+    """Get the mainline depths for the given event based on the mainline map
+
+    Args:
+        event
+        mainline_map: Map from event_id to mainline depth for events in the mainline.
+        event_map
+        state_res_store
+
+    Returns:
+        The mainline depth
+    """
+
+    room_id = event.room_id
+    tmp_event: Optional[EventBase] = event
+
+    # We do an iterative search, replacing `event with the power level in its
+    # auth events (if any)
+    idx = 0
+    while tmp_event:
+        depth = mainline_map.get(tmp_event.event_id)
+        if depth is not None:
+            return depth
+
+        auth_events = tmp_event.auth_event_ids()
+        tmp_event = None
+
+        for aid in auth_events:
+            aev = await _get_event(
+                room_id, aid, event_map, state_res_store, allow_none=True
+            )
+            if aev and (aev.type, aev.state_key) == (EventTypes.PowerLevels, ""):
+                tmp_event = aev
+                break
+
+        idx += 1
+
+        if idx % _AWAIT_AFTER_ITERATIONS == 0:
+            await clock.sleep(0)
+
+    # Didn't find a power level auth event, so we just return 0
+    return 0
+
+
+@overload
+async def _get_event(
+    room_id: str,
+    event_id: str,
+    event_map: dict[str, EventBase],
+    state_res_store: StateResolutionStore,
+    allow_none: Literal[False] = False,
+) -> EventBase: ...
+
+
+@overload
+async def _get_event(
+    room_id: str,
+    event_id: str,
+    event_map: dict[str, EventBase],
+    state_res_store: StateResolutionStore,
+    allow_none: Literal[True],
+) -> Optional[EventBase]: ...
+
+
+async def _get_event(
+    room_id: str,
+    event_id: str,
+    event_map: dict[str, EventBase],
+    state_res_store: StateResolutionStore,
+    allow_none: bool = False,
+) -> Optional[EventBase]:
+    """Helper function to look up event in event_map, falling back to looking
+    it up in the store
+
+    Args:
+        room_id
+        event_id
+        event_map
+        state_res_store
+        allow_none: if the event is not found, return None rather than raising
+            an exception
+
+    Returns:
+        The event, or none if the event does not exist (and allow_none is True).
+    """
+    if event_id not in event_map:
+        events = await state_res_store.get_events([event_id], allow_rejected=True)
+        event_map.update(events)
+    event = event_map.get(event_id)
+
+    if event is None:
+        if allow_none:
+            return None
+        raise Exception("Unknown event %s" % (event_id,))
+
+    if event.room_id != room_id:
+        raise Exception(
+            "In state res for room %s, event %s is in %s"
+            % (room_id, event_id, event.room_id)
+        )
+    return event
+
+
+def lexicographical_topological_sort(
+    graph: dict[str, set[str]], key: Callable[[str], Any]
+) -> Generator[str, None, None]:
+    """Performs a lexicographic reverse topological sort on the graph.
+
+    This returns a reverse topological sort (i.e. if node A references B then B
+    appears before A in the sort), with ties broken lexicographically based on
+    return value of the `key` function.
+
+    NOTE: `graph` is modified during the sort.
+
+    Args:
+        graph: A representation of the graph where each node is a key in the
+            dict and its value are the nodes edges.
+        key: A function that takes a node and returns a value that is comparable
+            and used to order nodes
+
+    Yields:
+        The next node in the topological sort
+    """
+
+    # Note, this is basically Kahn's algorithm except we look at nodes with no
+    # outgoing edges, c.f.
+    # https://en.wikipedia.org/wiki/Topological_sorting#Kahn's_algorithm
+    outdegree_map = graph
+    reverse_graph: dict[str, set[str]] = {}
+
+    # Lists of nodes with zero out degree. Is actually a tuple of
+    # `(key(node), node)` so that sorting does the right thing
+    zero_outdegree = []
+
+    for node, edges in graph.items():
+        if len(edges) == 0:
+            zero_outdegree.append((key(node), node))
+
+        reverse_graph.setdefault(node, set())
+        for edge in edges:
+            reverse_graph.setdefault(edge, set()).add(node)
+
+    # heapq is a built in implementation of a sorted queue.
+    heapq.heapify(zero_outdegree)
+
+    while zero_outdegree:
+        _, node = heapq.heappop(zero_outdegree)
+
+        for parent in reverse_graph[node]:
+            out = outdegree_map[parent]
+            out.discard(node)
+            if len(out) == 0:
+                heapq.heappush(zero_outdegree, (key(parent), parent))
+
+        yield node