1 files changed, 9 insertions, 35 deletions
diff --git a/syncapi/routing/messages.go b/syncapi/routing/messages.go
index 4a01ec35..02d8fcc7 100644
--- a/syncapi/routing/messages.go
+++ b/syncapi/routing/messages.go
@@ -64,6 +64,7 @@ type messagesResp struct {
 // OnIncomingMessagesRequest implements the /messages endpoint from the
 // client-server API.
 // See: https://matrix.org/docs/spec/client_server/latest.html#get-matrix-client-r0-rooms-roomid-messages
+// nolint:gocyclo
 func OnIncomingMessagesRequest(
 	req *http.Request, db storage.Database, roomID string, device *userapi.Device,
 	rsAPI api.SyncRoomserverAPI,
@@ -246,7 +247,14 @@ func OnIncomingMessagesRequest(
 		Start: start.String(),
 		End:   end.String(),
 	}
-	res.applyLazyLoadMembers(req.Context(), snapshot, roomID, device, filter.LazyLoadMembers, lazyLoadCache)
+	if filter.LazyLoadMembers {
+		membershipEvents, err := applyLazyLoadMembers(req.Context(), device, snapshot, roomID, clientEvents, lazyLoadCache)
+		if err != nil {
+			util.GetLogger(req.Context()).WithError(err).Error("failed to apply lazy loading")
+			return jsonerror.InternalServerError()
+		}
+		res.State = append(res.State, gomatrixserverlib.HeaderedToClientEvents(membershipEvents, gomatrixserverlib.FormatAll)...)
+	}
 
 	// If we didn't return any events, set the end to an empty string, so it will be omitted
 	// in the response JSON.
@@ -265,40 +273,6 @@ func OnIncomingMessagesRequest(
 	}
 }
 
-// applyLazyLoadMembers loads membership events for users returned in Chunk, if the filter has
-// LazyLoadMembers enabled.
-func (m *messagesResp) applyLazyLoadMembers(
-	ctx context.Context,
-	db storage.DatabaseTransaction,
-	roomID string,
-	device *userapi.Device,
-	lazyLoad bool,
-	lazyLoadCache caching.LazyLoadCache,
-) {
-	if !lazyLoad {
-		return
-	}
-	membershipToUser := make(map[string]*gomatrixserverlib.HeaderedEvent)
-	for _, evt := range m.Chunk {
-		// Don't add membership events the client should already know about
-		if _, cached := lazyLoadCache.IsLazyLoadedUserCached(device, roomID, evt.Sender); cached {
-			continue
-		}
-		membership, err := db.GetStateEvent(ctx, roomID, gomatrixserverlib.MRoomMember, evt.Sender)
-		if err != nil {
-			util.GetLogger(ctx).WithError(err).Error("failed to get membership event for user")
-			continue
-		}
-		if membership != nil {
-			membershipToUser[evt.Sender] = membership
-			lazyLoadCache.StoreLazyLoadedUser(device, roomID, evt.Sender, membership.EventID())
-		}
-	}
-	for _, evt := range membershipToUser {
-		m.State = append(m.State, gomatrixserverlib.HeaderedToClientEvent(evt, gomatrixserverlib.FormatAll))
-	}
-}
-
 func getMembershipForUser(ctx context.Context, roomID, userID string, rsAPI api.SyncRoomserverAPI) (resp api.QueryMembershipForUserResponse, err error) {
 	req := api.QueryMembershipForUserRequest{
 		RoomID: roomID,