From 2a0244647c1e3edcc7a72f95ecd288b1646b1fe4 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Thu, 23 Apr 2020 14:51:17 +0100
Subject: [PATCH] Try to build up invite state from local knowledge for local
 invites, if possible

---
 roomserver/input/events.go | 46 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/roomserver/input/events.go b/roomserver/input/events.go
index 5f67c7f49..80a959de1 100644
--- a/roomserver/input/events.go
+++ b/roomserver/input/events.go
@@ -78,6 +78,8 @@ type RoomEventDatabase interface {
 	GetRoomVersionForRoom(
 		ctx context.Context, roomID string,
 	) (gomatrixserverlib.RoomVersion, error)
+	RoomNID(ctx context.Context, roomID string) (types.RoomNID, error)
+	LatestEventIDs(ctx context.Context, roomNID types.RoomNID) ([]gomatrixserverlib.EventReference, types.StateSnapshotNID, int64, error)
 }
 
 // OutputRoomEventWriter has the APIs needed to write an event to the output logs.
@@ -248,9 +250,53 @@ func processInviteEvent(
 	event := input.Event.Unwrap()
 
 	if len(input.InviteRoomState) > 0 {
+		// If we were supplied with some invite room state then let's use that.
+		// This will ordinarily happen over federation.
 		if err = event.SetUnsignedField("invite_room_state", input.InviteRoomState); err != nil {
 			return err
 		}
+	} else {
+		// Otherwise, we should see if we know anything about the room state
+		// locally. If we have local knowledge of the room, use the locally known
+		// state to build up the invite room state.
+		if roomNID, err := db.RoomNID(ctx, roomID); err == nil && roomNID != 0 {
+			stateWanted := []gomatrixserverlib.StateKeyTuple{}
+			for _, t := range []string{
+				gomatrixserverlib.MRoomName, gomatrixserverlib.MRoomCanonicalAlias,
+				gomatrixserverlib.MRoomAliases, gomatrixserverlib.MRoomJoinRules,
+			} {
+				stateWanted = append(stateWanted, gomatrixserverlib.StateKeyTuple{
+					EventType: t,
+					StateKey:  "",
+				})
+			}
+			_, currentStateSnapshotNID, _, err := db.LatestEventIDs(ctx, roomNID)
+			if err != nil {
+				return err
+			}
+			roomState := state.NewStateResolution(db)
+			stateEntries, err := roomState.LoadStateAtSnapshotForStringTuples(
+				ctx, currentStateSnapshotNID, stateWanted,
+			)
+			if err != nil {
+				return err
+			}
+			stateNIDs := []types.EventNID{}
+			for _, stateNID := range stateEntries {
+				stateNIDs = append(stateNIDs, stateNID.EventNID)
+			}
+			stateEvents, err := db.Events(ctx, stateNIDs)
+			if err != nil {
+				return err
+			}
+			strippedState := []gomatrixserverlib.InviteV2StrippedState{}
+			for _, event := range stateEvents {
+				strippedState = append(strippedState, gomatrixserverlib.NewInviteV2StrippedState(&event.Event))
+			}
+			if err = event.SetUnsignedField("invite_room_state", strippedState); err != nil {
+				return err
+			}
+		}
 	}
 
 	outputUpdates, err := updateToInviteMembership(updater, &event, nil, input.Event.RoomVersion)