From 4ad297828b22083f0604a08fd54f66908095a076 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Tue, 29 Sep 2020 16:00:05 +0100
Subject: [PATCH] Add prefer_direct_fetch option

---
 dendrite-config.yaml                |  5 +++++
 internal/config/config_serverkey.go |  3 +++
 serverkeyapi/serverkeyapi.go        | 21 ++++++++++++++++-----
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/dendrite-config.yaml b/dendrite-config.yaml
index b71fb5091..74fa9b3e1 100644
--- a/dendrite-config.yaml
+++ b/dendrite-config.yaml
@@ -274,6 +274,11 @@ server_key_api:
     - key_id: ed25519:a_RXGa
       public_key: l8Hft5qXKn1vfHrg3p4+W8gELQVo8N13JkluMfmn2sQ
 
+  # This option will control whether Dendrite will prefer to look up keys directly
+  # or whether it should try perspective servers first, using direct fetches as a
+  # last resort.
+  prefer_direct_fetch: false
+
 # Configuration for the Sync API.
 sync_api:
   internal_api:
diff --git a/internal/config/config_serverkey.go b/internal/config/config_serverkey.go
index 40506d233..788a2fa05 100644
--- a/internal/config/config_serverkey.go
+++ b/internal/config/config_serverkey.go
@@ -14,6 +14,9 @@ type ServerKeyAPI struct {
 	// Perspective keyservers, to use as a backup when direct key fetch
 	// requests don't succeed
 	KeyPerspectives KeyPerspectives `yaml:"key_perspectives"`
+
+	// Should we prefer direct key fetches over perspective ones?
+	PreferDirectFetch bool `yaml:"prefer_direct_fetch"`
 }
 
 func (c *ServerKeyAPI) Defaults() {
diff --git a/serverkeyapi/serverkeyapi.go b/serverkeyapi/serverkeyapi.go
index 783402b25..92781f464 100644
--- a/serverkeyapi/serverkeyapi.go
+++ b/serverkeyapi/serverkeyapi.go
@@ -51,15 +51,26 @@ func NewInternalAPI(
 		ServerKeyValidity: cfg.Matrix.KeyValidityPeriod,
 		FedClient:         fedClient,
 		OurKeyRing: gomatrixserverlib.KeyRing{
-			KeyFetchers: []gomatrixserverlib.KeyFetcher{
-				&gomatrixserverlib.DirectKeyFetcher{
-					Client: fedClient,
-				},
-			},
+			KeyFetchers: []gomatrixserverlib.KeyFetcher{},
 			KeyDatabase: serverKeyDB,
 		},
 	}
 
+	addDirectFetcher := func() {
+		internalAPI.OurKeyRing.KeyFetchers = append(
+			internalAPI.OurKeyRing.KeyFetchers,
+			&gomatrixserverlib.DirectKeyFetcher{
+				Client: fedClient,
+			},
+		)
+	}
+
+	if cfg.PreferDirectFetch {
+		addDirectFetcher()
+	} else {
+		defer addDirectFetcher()
+	}
+
 	var b64e = base64.StdEncoding.WithPadding(base64.NoPadding)
 	for _, ps := range cfg.KeyPerspectives {
 		perspective := &gomatrixserverlib.PerspectiveKeyFetcher{