MSC4242: State DAGs (Federation)

Builds off #19424

Adds federation compatibility for state DAG rooms. Overview:
 - Adds extra HTTP API fields as per the MSC.
 - Adds methods for walking and extracting the state DAG for a room (for `/get_missing_events` and `/send_join` respectively).
 - Adds impl for processing the federation requests, as well as `/send`.

Author: kegsay

synapse/federation/federation_client.py

@@ -42,7 +42,12 @@
 import attr
 from prometheus_client import Counter
 
-from synapse.api.constants import Direction, EventContentFields, EventTypes, Membership
+from synapse.api.constants import (
+    Direction,
+    EventContentFields,
+    EventTypes,
+    Membership,
+)
 from synapse.api.errors import (
     CodeMessageException,
     Codes,
@@ -119,6 +124,8 @@ class SendJoinResult:
     origin: str
     state: list[EventBase]
     auth_chain: list[EventBase]
+    # Only valid for state DAG rooms (MSC4242)
+    state_dag: list[EventBase] | None
 
     # True if 'state' elides non-critical membership events
     partial_state: bool
@@ -658,7 +665,10 @@ async def get_pdu(
     @trace
     @tag_args
     async def get_room_state_ids(
-        self, destination: str, room_id: str, event_id: str
+        self,
+        destination: str,
+        room_id: str,
+        event_id: str,
     ) -> tuple[list[str], list[str]]:
         """Calls the /state_ids endpoint to fetch the state at a particular point
         in the room, and the auth events for the given event
@@ -670,7 +680,9 @@ async def get_room_state_ids(
             InvalidResponseError: if fields in the response have the wrong type.
         """
         result = await self.transport_layer.get_room_state_ids(
-            destination, room_id, event_id=event_id
+            destination,
+            room_id,
+            event_id=event_id,
         )
 
         state_event_ids = result["pdu_ids"]
@@ -1178,7 +1190,6 @@ async def send_request(destination: str) -> SendJoinResult:
             response = await self._do_send_join(
                 room_version, destination, pdu, omit_members=partial_state
             )
-
             # If an event was returned (and expected to be returned):
             #
             # * Ensure it has the same event ID (note that the event ID is a hash
@@ -1201,14 +1212,22 @@ async def send_request(destination: str) -> SendJoinResult:
                 event = pdu
 
             state = response.state
-            auth_chain = response.auth_events
-
+            auth_events = response.auth_events
             create_event = None
             for e in state:
                 if (e.type, e.state_key) == (EventTypes.Create, ""):
                     create_event = e
                     break
 
+            if room_version.msc4242_state_dags and response.state_dag:
+                # assign to auth_events to reuse the below code which ultimately just does
+                # sig/hash checks. We'll set the right field in SendJoinResult later.
+                auth_events = response.state_dag
+                for e in response.state_dag:
+                    if (e.type, e.state_key) == (EventTypes.Create, ""):
+                        create_event = e
+                        break
+
             if create_event is None:
                 # If the state doesn't have a create event then the room is
                 # invalid, and it would fail auth checks anyway.
@@ -1227,7 +1246,7 @@ async def send_request(destination: str) -> SendJoinResult:
                 )
 
             logger.info(
-                "Processing from send_join %d events", len(state) + len(auth_chain)
+                "Processing from send_join %d events", len(state) + len(auth_events)
             )
 
             # We now go and check the signatures and hashes for the event. Note
@@ -1246,7 +1265,7 @@ async def _execute(pdu: EventBase) -> None:
                     valid_pdus_map[valid_pdu.event_id] = valid_pdu
 
             await concurrently_execute(
-                _execute, itertools.chain(state, auth_chain), 10000
+                _execute, itertools.chain(state, auth_events), 10000
             )
 
             # NB: We *need* to copy to ensure that we don't have multiple
@@ -1259,27 +1278,28 @@ async def _execute(pdu: EventBase) -> None:
 
             signed_auth = [
                 valid_pdus_map[p.event_id]
-                for p in auth_chain
+                for p in auth_events
                 if p.event_id in valid_pdus_map
             ]
 
             # NB: We *need* to copy to ensure that we don't have multiple
             # references being passed on, as that causes... issues.
+            # TODO(kegan): It's unclear why we only need to do this for state and not auth_events
             for s in signed_state:
                 s.internal_metadata = s.internal_metadata.copy()
 
-            # double-check that the auth chain doesn't include a different create event
-            auth_chain_create_events = [
+            # double-check that the auth events doesn't include a different create event
+            auth_events_create_events = [
                 e.event_id
                 for e in signed_auth
                 if (e.type, e.state_key) == (EventTypes.Create, "")
             ]
-            if auth_chain_create_events and auth_chain_create_events != [
+            if auth_events_create_events and auth_events_create_events != [
                 create_event.event_id
             ]:
                 raise InvalidResponseError(
                     "Unexpected create event(s) in auth chain: %s"
-                    % (auth_chain_create_events,)
+                    % (auth_events_create_events,)
                 )
 
             servers_in_room = None
@@ -1301,10 +1321,21 @@ async def _execute(pdu: EventBase) -> None:
                 # Fix things up in case the remote homeserver is badly behaved.
                 servers_in_room.add(destination)
 
+            signed_auth_events = signed_auth
+            signed_state_dag = None
+            if room_version.msc4242_state_dags:
+                # We previously set the state dag to auth_events so re-assign it correctly
+                signed_state_dag = signed_auth
+                # Ensure the caller cannot accidentally use these values even if the server
+                # returned them.
+                signed_state = []
+                signed_auth_events = []
+
             return SendJoinResult(
                 event=event,
                 state=signed_state,
-                auth_chain=signed_auth,
+                auth_chain=signed_auth_events,
+                state_dag=signed_state_dag,
                 origin=destination,
                 partial_state=response.members_omitted,
                 servers_in_room=servers_in_room or frozenset(),
@@ -1608,6 +1639,7 @@ async def get_missing_events(
         limit: int,
         min_depth: int,
         timeout: int,
+        state_dag: bool = False,
     ) -> list[EventBase]:
         """Tries to fetch events we are missing. This is called when we receive
         an event without having received all of its ancestors.
@@ -1623,6 +1655,7 @@ async def get_missing_events(
             limit: Maximum number of events to return.
             min_depth: Minimum depth of events to return.
             timeout: Max time to wait in ms
+            state_dag: True to walk the state DAG (MSC4242 rooms)
         """
         try:
             content = await self.transport_layer.get_missing_events(
@@ -1633,6 +1666,7 @@ async def get_missing_events(
                 limit=limit,
                 min_depth=min_depth,
                 timeout=timeout,
+                state_dag=state_dag,
             )
 
             room_version = await self.store.get_room_version(room_id)

synapse/federation/federation_server.py

@@ -28,6 +28,8 @@
     Callable,
     Collection,
     Mapping,
+    Sequence,
+    cast,
 )
 
 from prometheus_client import Counter, Gauge, Histogram
@@ -53,7 +55,7 @@
 )
 from synapse.api.room_versions import KNOWN_ROOM_VERSIONS, RoomVersion
 from synapse.crypto.event_signing import compute_event_signature
-from synapse.events import EventBase
+from synapse.events import EventBase, FrozenEventVMSC4242
 from synapse.events.snapshot import EventPersistencePair
 from synapse.federation.federation_base import (
     FederationBase,
@@ -596,7 +598,10 @@ async def _process_edu(edu_dict: JsonDict) -> None:
         )
 
     async def on_room_state_request(
-        self, origin: str, room_id: str, event_id: str
+        self,
+        origin: str,
+        room_id: str,
+        event_id: str,
     ) -> tuple[int, JsonDict]:
         await self._event_auth_handler.assert_host_in_room(room_id, origin)
         origin_host, _ = parse_server_name(origin)
@@ -620,7 +625,10 @@ async def on_room_state_request(
     @trace
     @tag_args
     async def on_state_ids_request(
-        self, origin: str, room_id: str, event_id: str
+        self,
+        origin: str,
+        room_id: str,
+        event_id: str,
     ) -> tuple[int, JsonDict]:
         if not event_id:
             raise NotImplementedError("Specify an event")
@@ -641,17 +649,27 @@ async def on_state_ids_request(
     @trace
     @tag_args
     async def _on_state_ids_request_compute(
-        self, room_id: str, event_id: str
+        self,
+        room_id: str,
+        event_id: str,
     ) -> JsonDict:
-        state_ids = await self.handler.get_state_ids_for_pdu(room_id, event_id)
+        state_ids = await self.handler.get_state_ids_for_pdu(
+            room_id,
+            event_id,
+        )
         auth_chain_ids = await self.store.get_auth_chain_ids(room_id, state_ids)
         return {"pdu_ids": state_ids, "auth_chain_ids": list(auth_chain_ids)}
 
     async def _on_context_state_request_compute(
-        self, room_id: str, event_id: str
+        self,
+        room_id: str,
+        event_id: str,
     ) -> dict[str, list]:
         pdus: Collection[EventBase]
-        event_ids = await self.handler.get_state_ids_for_pdu(room_id, event_id)
+        event_ids = await self.handler.get_state_ids_for_pdu(
+            room_id,
+            event_id,
+        )
         pdus = await self.store.get_events_as_list(event_ids)
 
         auth_chain = await self.store.get_auth_chain(
@@ -759,6 +777,10 @@ async def on_send_join_request(
             origin, content, Membership.JOIN, room_id
         )
 
+        if event.room_version.msc4242_state_dags and caller_supports_partial_state:
+            # TODO(kegan): for now, MSC4242 won't support partial state for ease of prototyping.
+            caller_supports_partial_state = False
+
         prev_state_ids = await context.get_prev_state_ids()
 
         state_event_ids: Collection[str]
@@ -769,15 +791,31 @@ async def on_send_join_request(
                 event, prev_state_ids, summary
             )
             servers_in_room = await self.state.get_hosts_in_room_at_events(
-                room_id, event_ids=event.prev_event_ids()
+                room_id,
+                event_ids=event.prev_state_events
+                if isinstance(event, FrozenEventVMSC4242)
+                else event.prev_event_ids(),
             )
         else:
             state_event_ids = prev_state_ids.values()
             servers_in_room = None
 
-        auth_chain_event_ids = await self.store.get_auth_chain_ids(
-            room_id, state_event_ids
-        )
+        state_dag = None
+        auth_chain_event_ids = set()
+        if event.room_version.msc4242_state_dags:
+            assert isinstance(event, FrozenEventVMSC4242)
+            # NOTE: we don't return the state dag for forward extremities that aren't part of this
+            # join event to make it easier for the receiving server to set their own forward
+            # extremities (they are equal to the join event's prev_state_events). This means we may
+            # fail to sync concurrent forks not on the path to the join event, but this is an
+            # outstanding problem in general.
+            state_dag = await self.store.get_state_dag(
+                room_id, set(event.prev_state_events)
+            )
+        else:
+            auth_chain_event_ids = await self.store.get_auth_chain_ids(
+                room_id, state_event_ids
+            )
 
         # if the caller has opted in, we can omit any auth_chain events which are
         # already in state_event_ids
@@ -794,13 +832,18 @@ async def on_send_join_request(
         resp = {
             "event": event_json,
             "state": serialize_and_filter_pdus(state_events, time_now),
-            "auth_chain": serialize_and_filter_pdus(auth_chain_events, time_now),
             "members_omitted": caller_supports_partial_state,
         }
+        if state_dag is None:
+            resp["auth_chain"] = serialize_and_filter_pdus(auth_chain_events, time_now)
+        else:
+            resp["state_dag"] = serialize_and_filter_pdus(
+                cast(Sequence[EventBase], state_dag.values()), time_now
+            )
+            del resp["state"]
 
         if servers_in_room is not None:
             resp["servers_in_room"] = list(servers_in_room)
-
         return resp
 
     async def on_make_leave_request(
@@ -1097,6 +1140,7 @@ async def on_get_missing_events(
         earliest_events: list[str],
         latest_events: list[str],
         limit: int,
+        walk_state_dag: bool = False,
     ) -> dict[str, list]:
         async with self._server_linearizer.queue((origin, room_id)):
             origin_host, _ = parse_server_name(origin)
@@ -1111,7 +1155,7 @@ async def on_get_missing_events(
             )
 
             missing_events = await self.handler.on_get_missing_events(
-                origin, room_id, earliest_events, latest_events, limit
+                origin, room_id, earliest_events, latest_events, limit, walk_state_dag
             )
 
             if len(missing_events) < 5:

synapse/federation/sender/__init__.py

@@ -148,7 +148,7 @@
 import synapse.metrics
 from synapse.api.constants import EventTypes, Membership
 from synapse.api.presence import UserPresenceState
-from synapse.events import EventBase
+from synapse.events import EventBase, FrozenEventVMSC4242
 from synapse.federation.sender.per_destination_queue import (
     CATCHUP_RETRY_INTERVAL,
     PerDestinationQueue,
@@ -660,7 +660,10 @@ async def handle_event(event: EventBase) -> None:
                             # banned then it won't receive the event because it won't
                             # be in the room after the ban.
                             destinations = await self.state.get_hosts_in_room_at_events(
-                                event.room_id, event_ids=event.prev_event_ids()
+                                event.room_id,
+                                event_ids=event.prev_state_events
+                                if isinstance(event, FrozenEventVMSC4242)
+                                else event.prev_event_ids(),
                             )
                         except Exception:
                             logger.exception(