Fix keyring regressions in previous P2P demo

3 files changed, 45 insertions, 41 deletions

diff --git a/federationapi/federationapi.go b/federationapi/federationapi.go
index 371b364c..02c4cfdb 100644
--- a/federationapi/federationapi.go
+++ b/federationapi/federationapi.go
@@ -73,6 +73,7 @@ func NewInternalAPI(
 	federation *gomatrixserverlib.FederationClient,
 	rsAPI roomserverAPI.RoomserverInternalAPI,
 	caches *caching.Caches,
+	keyRing *gomatrixserverlib.KeyRing,
 	resetBlacklist bool,
 ) api.FederationInternalAPI {
 	cfg := &base.Cfg.FederationAPI
@@ -125,5 +126,5 @@ func NewInternalAPI(
 		logrus.WithError(err).Panic("failed to start key server consumer")
 	}
 
-	return internal.NewFederationInternalAPI(federationDB, cfg, rsAPI, federation, stats, caches, queues)
+	return internal.NewFederationInternalAPI(federationDB, cfg, rsAPI, federation, stats, caches, queues, keyRing)
 }
diff --git a/federationapi/federationapi_keys_test.go b/federationapi/federationapi_keys_test.go
index fb4307aa..9e6c47cd 100644
--- a/federationapi/federationapi_keys_test.go
+++ b/federationapi/federationapi_keys_test.go
@@ -94,7 +94,7 @@ func TestMain(m *testing.M) {
 
 		// Finally, build the server key APIs.
 		sbase := base.NewBaseDendrite(cfg, "Monolith", base.NoCacheMetrics)
-		s.api = NewInternalAPI(sbase, s.fedclient, nil, s.cache, true)
+		s.api = NewInternalAPI(sbase, s.fedclient, nil, s.cache, nil, true)
 	}
 
 	// Now that we have built our server key APIs, start the
diff --git a/federationapi/internal/api.go b/federationapi/internal/api.go
index 73d27315..1f31b07c 100644
--- a/federationapi/internal/api.go
+++ b/federationapi/internal/api.go
@@ -39,58 +39,61 @@ func NewFederationInternalAPI(
 	statistics *statistics.Statistics,
 	caches *caching.Caches,
 	queues *queue.OutgoingQueues,
+	keyRing *gomatrixserverlib.KeyRing,
 ) *FederationInternalAPI {
 	serverKeyDB, err := cache.NewKeyDatabase(db, caches)
 	if err != nil {
 		logrus.WithError(err).Panicf("failed to set up caching wrapper for server key database")
 	}
 
-	keyRing := &gomatrixserverlib.KeyRing{
-		KeyFetchers: []gomatrixserverlib.KeyFetcher{},
-		KeyDatabase: serverKeyDB,
-	}
-
-	addDirectFetcher := func() {
-		keyRing.KeyFetchers = append(
-			keyRing.KeyFetchers,
-			&gomatrixserverlib.DirectKeyFetcher{
-				Client: federation,
-			},
-		)
-	}
+	if keyRing == nil {
+		keyRing = &gomatrixserverlib.KeyRing{
+			KeyFetchers: []gomatrixserverlib.KeyFetcher{},
+			KeyDatabase: serverKeyDB,
+		}
 
-	if cfg.PreferDirectFetch {
-		addDirectFetcher()
-	} else {
-		defer addDirectFetcher()
-	}
+		addDirectFetcher := func() {
+			keyRing.KeyFetchers = append(
+				keyRing.KeyFetchers,
+				&gomatrixserverlib.DirectKeyFetcher{
+					Client: federation,
+				},
+			)
+		}
 
-	var b64e = base64.StdEncoding.WithPadding(base64.NoPadding)
-	for _, ps := range cfg.KeyPerspectives {
-		perspective := &gomatrixserverlib.PerspectiveKeyFetcher{
-			PerspectiveServerName: ps.ServerName,
-			PerspectiveServerKeys: map[gomatrixserverlib.KeyID]ed25519.PublicKey{},
-			Client:                federation,
+		if cfg.PreferDirectFetch {
+			addDirectFetcher()
+		} else {
+			defer addDirectFetcher()
 		}
 
-		for _, key := range ps.Keys {
-			rawkey, err := b64e.DecodeString(key.PublicKey)
-			if err != nil {
-				logrus.WithError(err).WithFields(logrus.Fields{
-					"server_name": ps.ServerName,
-					"public_key":  key.PublicKey,
-				}).Warn("Couldn't parse perspective key")
-				continue
+		var b64e = base64.StdEncoding.WithPadding(base64.NoPadding)
+		for _, ps := range cfg.KeyPerspectives {
+			perspective := &gomatrixserverlib.PerspectiveKeyFetcher{
+				PerspectiveServerName: ps.ServerName,
+				PerspectiveServerKeys: map[gomatrixserverlib.KeyID]ed25519.PublicKey{},
+				Client:                federation,
 			}
-			perspective.PerspectiveServerKeys[key.KeyID] = rawkey
-		}
 
-		keyRing.KeyFetchers = append(keyRing.KeyFetchers, perspective)
+			for _, key := range ps.Keys {
+				rawkey, err := b64e.DecodeString(key.PublicKey)
+				if err != nil {
+					logrus.WithError(err).WithFields(logrus.Fields{
+						"server_name": ps.ServerName,
+						"public_key":  key.PublicKey,
+					}).Warn("Couldn't parse perspective key")
+					continue
+				}
+				perspective.PerspectiveServerKeys[key.KeyID] = rawkey
+			}
 
-		logrus.WithFields(logrus.Fields{
-			"server_name":     ps.ServerName,
-			"num_public_keys": len(ps.Keys),
-		}).Info("Enabled perspective key fetcher")
+			keyRing.KeyFetchers = append(keyRing.KeyFetchers, perspective)
+
+			logrus.WithFields(logrus.Fields{
+				"server_name":     ps.ServerName,
+				"num_public_keys": len(ps.Keys),
+			}).Info("Enabled perspective key fetcher")
+		}
 	}
 
 	return &FederationInternalAPI{