From 1e3c19f31da928df84f0b2d4ade3a77c495a9467 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Mon, 15 Nov 2021 13:13:33 +0000
Subject: [PATCH] Use `"event"` field from `RespSendJoin`

---
 federationapi/routing/join.go        | 26 +-----------------------
 federationsender/internal/perform.go | 30 ++++++++--------------------
 2 files changed, 9 insertions(+), 47 deletions(-)

diff --git a/federationapi/routing/join.go b/federationapi/routing/join.go
index 6b979171..c68f4d1a 100644
--- a/federationapi/routing/join.go
+++ b/federationapi/routing/join.go
@@ -571,31 +571,6 @@ func SendJoin(
 		}
 	}
 
-	// Since we performed QueryStateAndAuthChain before we sent the membership event to the
-	// roomserver, we need to ensure that the stateAndAuthChainResponse contain the new join
-	// event. It would be expensive to call QueryStateAndAuthChain again, so we'll just modify
-	// the response in situ.
-	replaced := false
-	for i := range stateAndAuthChainResponse.StateEvents {
-		stateEvent := stateAndAuthChainResponse.StateEvents[i]
-		if event.Type() != gomatrixserverlib.MRoomMember {
-			continue
-		}
-		if stateEvent.StateKeyEquals(*event.StateKey()) {
-			stateAndAuthChainResponse.StateEvents[i] = event.Headered(
-				stateAndAuthChainResponse.RoomVersion,
-			)
-			replaced = true
-			break
-		}
-	}
-	if !replaced {
-		stateAndAuthChainResponse.StateEvents = append(
-			stateAndAuthChainResponse.StateEvents,
-			event.Headered(stateAndAuthChainResponse.RoomVersion),
-		)
-	}
-
 	// sort events deterministically by depth (lower is earlier)
 	// We also do this because sytest's basic federation server isn't good at using the correct
 	// state if these lists are randomised, resulting in flakey tests. :(
@@ -606,6 +581,7 @@ func SendJoin(
 	return util.JSONResponse{
 		Code: http.StatusOK,
 		JSON: gomatrixserverlib.RespSendJoin{
+			Event:       event,
 			StateEvents: gomatrixserverlib.UnwrapEventHeaders(stateAndAuthChainResponse.StateEvents),
 			AuthEvents:  gomatrixserverlib.UnwrapEventHeaders(stateAndAuthChainResponse.AuthChainEvents),
 			Origin:      cfg.Matrix.ServerName,
diff --git a/federationsender/internal/perform.go b/federationsender/internal/perform.go
index cdc80c5c..74b97cce 100644
--- a/federationsender/internal/perform.go
+++ b/federationsender/internal/perform.go
@@ -247,6 +247,13 @@ func (r *FederationSenderInternalAPI) performJoinUsingServer(
 		defer close(waiterr)
 		defer cancel()
 
+		// If the remote server returned a signed membership event then
+		// we will use that instead. That is necessary for restricted
+		// joins to work.
+		if respSendJoin.Event != nil {
+			event = respSendJoin.Event
+		}
+
 		// TODO: Can we expand Check here to return a list of missing auth
 		// events rather than failing one at a time?
 		respState, err := respSendJoin.Check(ctx, r.keyRing, event, federatedAuthProvider(ctx, r.federation, r.keyRing, serverName))
@@ -259,27 +266,6 @@ func (r *FederationSenderInternalAPI) performJoinUsingServer(
 			return
 		}
 
-		// Find the membership event.
-		var joinEvent *gomatrixserverlib.Event
-		for _, stateEvent := range respState.StateEvents {
-			if stateEvent.Type() != gomatrixserverlib.MRoomMember {
-				continue
-			}
-			if stateEvent.StateKeyEquals(*event.StateKey()) {
-				joinEvent = event
-				break
-			}
-		}
-		if joinEvent == nil {
-			err = fmt.Errorf("The remote server did not send back our join event")
-			logrus.WithFields(logrus.Fields{
-				"room_id": roomID,
-				"user_id": userID,
-			}).WithError(err).Error("Found no membership event")
-			waiterr <- err
-			return
-		}
-
 		// If we successfully performed a send_join above then the other
 		// server now thinks we're a part of the room. Send the newly
 		// returned state to the roomserver to update our local view.
@@ -287,7 +273,7 @@ func (r *FederationSenderInternalAPI) performJoinUsingServer(
 			ctx, r.rsAPI,
 			roomserverAPI.KindNew,
 			respState,
-			joinEvent.Headered(respMakeJoin.RoomVersion),
+			event.Headered(respMakeJoin.RoomVersion),
 			nil,
 		); err != nil {
 			logrus.WithFields(logrus.Fields{