From 97c64bdb6d1ce6ee4684f209dfa96db3545e240b Mon Sep 17 00:00:00 2001
From: Kegsay <kegan@matrix.org>
Date: Fri, 29 May 2020 15:34:15 +0100
Subject: [PATCH] p2p: de-dupe and shuffle public rooms; implement keydb
 (#1074)

- We don't want a serverKeyAPI as fetching keys doesn't need a DB.
- De-dupe rooms so we don't see them multiple times, but shuffle the
  alias we join via so we don't all flood a single server.
---
 cmd/dendritejs/keyfetcher.go             |  5 +++++
 cmd/dendritejs/main.go                   |  9 +++------
 publicroomsapi/directory/public_rooms.go | 18 ++++++++++++++++++
 3 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/cmd/dendritejs/keyfetcher.go b/cmd/dendritejs/keyfetcher.go
index ee4905d4..47a81c41 100644
--- a/cmd/dendritejs/keyfetcher.go
+++ b/cmd/dendritejs/keyfetcher.go
@@ -82,3 +82,8 @@ func (f *libp2pKeyFetcher) FetchKeys(
 func (f *libp2pKeyFetcher) FetcherName() string {
 	return "libp2pKeyFetcher"
 }
+
+// no-op function for storing keys - we don't do any work to fetch them so don't bother storing.
+func (f *libp2pKeyFetcher) StoreKeys(ctx context.Context, results map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult) error {
+	return nil
+}
diff --git a/cmd/dendritejs/main.go b/cmd/dendritejs/main.go
index b8517bf9..45f23d9a 100644
--- a/cmd/dendritejs/main.go
+++ b/cmd/dendritejs/main.go
@@ -37,7 +37,6 @@ import (
 	"github.com/matrix-org/dendrite/publicroomsapi"
 	"github.com/matrix-org/dendrite/publicroomsapi/storage"
 	"github.com/matrix-org/dendrite/roomserver"
-	"github.com/matrix-org/dendrite/serverkeyapi"
 	"github.com/matrix-org/dendrite/syncapi"
 	go_http_js_libp2p "github.com/matrix-org/go-http-js-libp2p"
 	"github.com/matrix-org/gomatrixserverlib"
@@ -197,14 +196,12 @@ func main() {
 	deviceDB := base.CreateDeviceDB()
 	federation := createFederationClient(cfg, node)
 
-	serverKeyAPI := serverkeyapi.SetupServerKeyAPIComponent(
-		base, federation,
-	)
+	fetcher := &libp2pKeyFetcher{}
 	keyRing := gomatrixserverlib.KeyRing{
 		KeyFetchers: []gomatrixserverlib.KeyFetcher{
-			&libp2pKeyFetcher{},
+			fetcher,
 		},
-		KeyDatabase: serverKeyAPI,
+		KeyDatabase: fetcher,
 	}
 	p2pPublicRoomProvider := NewLibP2PPublicRoomsProvider(node)
 
diff --git a/publicroomsapi/directory/public_rooms.go b/publicroomsapi/directory/public_rooms.go
index 7bd6740e..bc40da51 100644
--- a/publicroomsapi/directory/public_rooms.go
+++ b/publicroomsapi/directory/public_rooms.go
@@ -16,6 +16,7 @@ package directory
 
 import (
 	"context"
+	"math/rand"
 	"net/http"
 	"strconv"
 	"sync"
@@ -96,6 +97,23 @@ func GetPostPublicRoomsWithExternal(
 	// downcasting `limit` is safe as we know it isn't bigger than request.Limit which is int16
 	fedRooms := bulkFetchPublicRoomsFromServers(req.Context(), fedClient, extRoomsProvider.Homeservers(), int16(limit))
 	response.Chunk = append(response.Chunk, fedRooms...)
+
+	// de-duplicate rooms with the same room ID. We can join the room via any of these aliases as we know these servers
+	// are alive and well, so we arbitrarily pick one (purposefully shuffling them to spread the load a bit)
+	var publicRooms []gomatrixserverlib.PublicRoom
+	haveRoomIDs := make(map[string]bool)
+	rand.Shuffle(len(response.Chunk), func(i, j int) {
+		response.Chunk[i], response.Chunk[j] = response.Chunk[j], response.Chunk[i]
+	})
+	for _, r := range response.Chunk {
+		if haveRoomIDs[r.RoomID] {
+			continue
+		}
+		haveRoomIDs[r.RoomID] = true
+		publicRooms = append(publicRooms, r)
+	}
+	response.Chunk = publicRooms
+
 	return util.JSONResponse{
 		Code: http.StatusOK,
 		JSON: response,