From 8c9b2e402a1ae93092de26fc26fe5810bbbe8f05 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Tue, 25 Jan 2022 11:37:56 +0000
Subject: [PATCH] Don't limit servers, randomise order using map

---
 roomserver/internal/input/input_events.go  | 10 +++++-----
 roomserver/internal/input/input_missing.go |  6 +++---
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/roomserver/internal/input/input_events.go b/roomserver/internal/input/input_events.go
index 918a435b1..7fcc9cd59 100644
--- a/roomserver/internal/input/input_events.go
+++ b/roomserver/internal/input/input_events.go
@@ -133,10 +133,7 @@ func (r *Inputer) processRoomEvent(
 		}
 	}
 	if input.Origin != "" {
-		serverRes.ServerNames = append([]gomatrixserverlib.ServerName{input.Origin}, serverRes.ServerNames...)
-	}
-	if len(serverRes.ServerNames) > 5 {
-		serverRes.ServerNames = serverRes.ServerNames[:5]
+		serverRes.ServerNames = append(serverRes.ServerNames, input.Origin)
 	}
 
 	// First of all, check that the auth events of the event are known.
@@ -200,10 +197,13 @@ func (r *Inputer) processRoomEvent(
 				federation: r.FSAPI,
 				keys:       r.KeyRing,
 				roomsMu:    internal.NewMutexByRoom(),
-				servers:    serverRes.ServerNames,
+				servers:    map[gomatrixserverlib.ServerName]struct{}{},
 				hadEvents:  map[string]bool{},
 				haveEvents: map[string]*gomatrixserverlib.HeaderedEvent{},
 			}
+			for _, serverName := range serverRes.ServerNames {
+				missingState.servers[serverName] = struct{}{}
+			}
 			if err = missingState.processEventWithMissingState(ctx, event, headered.RoomVersion); err != nil {
 				isRejected = true
 				rejectionErr = fmt.Errorf("missingState.processEventWithMissingState: %w", err)
diff --git a/roomserver/internal/input/input_missing.go b/roomserver/internal/input/input_missing.go
index 41fda126b..aff11e3d3 100644
--- a/roomserver/internal/input/input_missing.go
+++ b/roomserver/internal/input/input_missing.go
@@ -25,7 +25,7 @@ type missingStateReq struct {
 	keys            gomatrixserverlib.JSONVerifier
 	federation      fedapi.FederationInternalAPI
 	roomsMu         *internal.MutexByRoom
-	servers         []gomatrixserverlib.ServerName
+	servers         map[gomatrixserverlib.ServerName]struct{}
 	hadEvents       map[string]bool
 	hadEventsMutex  sync.Mutex
 	haveEvents      map[string]*gomatrixserverlib.HeaderedEvent
@@ -371,7 +371,7 @@ func (t *missingStateReq) getMissingEvents(ctx context.Context, e *gomatrixserve
 	}
 
 	var missingResp *gomatrixserverlib.RespMissingEvents
-	for _, server := range t.servers {
+	for server := range t.servers {
 		var m gomatrixserverlib.RespMissingEvents
 		rctx, cancel := context.WithTimeout(ctx, time.Second*30)
 		defer cancel()
@@ -637,7 +637,7 @@ func (t *missingStateReq) lookupEvent(ctx context.Context, roomVersion gomatrixs
 	}
 	var event *gomatrixserverlib.Event
 	found := false
-	for _, serverName := range t.servers {
+	for serverName := range t.servers {
 		txn, err := t.federation.GetEvent(ctx, serverName, missingEventID)
 		if err != nil || len(txn.PDUs) == 0 {
 			util.GetLogger(ctx).WithError(err).WithField("event_id", missingEventID).Warn("Failed to get missing /event for event ID")