From 0013459677a7f489f0cf226c04419858b5a5fe5f Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Thu, 10 Nov 2022 11:42:20 +0000
Subject: [PATCH] Fix federation API

---
 federationapi/federationapi.go          | 18 ++++++++++++-----
 federationapi/queue/destinationqueue.go |  2 +-
 federationapi/queue/queue.go            | 26 +++++++++++++------------
 federationapi/queue/queue_test.go       | 10 ++++++----
 4 files changed, 34 insertions(+), 22 deletions(-)

diff --git a/federationapi/federationapi.go b/federationapi/federationapi.go
index 202da6c51..19e9e8640 100644
--- a/federationapi/federationapi.go
+++ b/federationapi/federationapi.go
@@ -120,15 +120,23 @@ func NewInternalAPI(
 
 	js, _ := base.NATS.Prepare(base.ProcessContext, &cfg.Matrix.JetStream)
 
+	signingInfo := map[gomatrixserverlib.ServerName]*queue.SigningInfo{}
+	for _, serverName := range append(
+		[]gomatrixserverlib.ServerName{base.Cfg.Global.ServerName},
+		base.Cfg.Global.SecondaryServerNames...,
+	) {
+		signingInfo[serverName] = &queue.SigningInfo{
+			KeyID:      cfg.Matrix.KeyID,
+			PrivateKey: cfg.Matrix.PrivateKey,
+			ServerName: serverName,
+		}
+	}
+
 	queues := queue.NewOutgoingQueues(
 		federationDB, base.ProcessContext,
 		cfg.Matrix.DisableFederation,
 		cfg.Matrix.ServerName, federation, rsAPI, &stats,
-		&queue.SigningInfo{
-			KeyID:      cfg.Matrix.KeyID,
-			PrivateKey: cfg.Matrix.PrivateKey,
-			ServerName: cfg.Matrix.ServerName,
-		},
+		signingInfo,
 	)
 
 	rsConsumer := consumers.NewOutputRoomEventConsumer(
diff --git a/federationapi/queue/destinationqueue.go b/federationapi/queue/destinationqueue.go
index a638a5742..bf04ee99a 100644
--- a/federationapi/queue/destinationqueue.go
+++ b/federationapi/queue/destinationqueue.go
@@ -50,7 +50,7 @@ type destinationQueue struct {
 	queues             *OutgoingQueues
 	db                 storage.Database
 	process            *process.ProcessContext
-	signing            *SigningInfo
+	signing            map[gomatrixserverlib.ServerName]*SigningInfo
 	rsAPI              api.FederationRoomserverAPI
 	client             fedapi.FederationClient         // federation client
 	origin             gomatrixserverlib.ServerName    // origin of requests
diff --git a/federationapi/queue/queue.go b/federationapi/queue/queue.go
index b5d0552c6..68f354993 100644
--- a/federationapi/queue/queue.go
+++ b/federationapi/queue/queue.go
@@ -46,7 +46,7 @@ type OutgoingQueues struct {
 	origin      gomatrixserverlib.ServerName
 	client      fedapi.FederationClient
 	statistics  *statistics.Statistics
-	signing     *SigningInfo
+	signing     map[gomatrixserverlib.ServerName]*SigningInfo
 	queuesMutex sync.Mutex // protects the below
 	queues      map[gomatrixserverlib.ServerName]*destinationQueue
 }
@@ -91,7 +91,7 @@ func NewOutgoingQueues(
 	client fedapi.FederationClient,
 	rsAPI api.FederationRoomserverAPI,
 	statistics *statistics.Statistics,
-	signing *SigningInfo,
+	signing map[gomatrixserverlib.ServerName]*SigningInfo,
 ) *OutgoingQueues {
 	queues := &OutgoingQueues{
 		disabled:   disabled,
@@ -199,11 +199,10 @@ func (oqs *OutgoingQueues) SendEvent(
 		log.Trace("Federation is disabled, not sending event")
 		return nil
 	}
-	if origin != oqs.origin {
-		// TODO: Support virtual hosting; gh issue #577.
+	if _, ok := oqs.signing[origin]; !ok {
 		return fmt.Errorf(
-			"sendevent: unexpected server to send as: got %q expected %q",
-			origin, oqs.origin,
+			"sendevent: unexpected server to send as %q",
+			origin,
 		)
 	}
 
@@ -214,7 +213,9 @@ func (oqs *OutgoingQueues) SendEvent(
 		destmap[d] = struct{}{}
 	}
 	delete(destmap, oqs.origin)
-	delete(destmap, oqs.signing.ServerName)
+	for local := range oqs.signing {
+		delete(destmap, local)
+	}
 
 	// Check if any of the destinations are prohibited by server ACLs.
 	for destination := range destmap {
@@ -288,11 +289,10 @@ func (oqs *OutgoingQueues) SendEDU(
 		log.Trace("Federation is disabled, not sending EDU")
 		return nil
 	}
-	if origin != oqs.origin {
-		// TODO: Support virtual hosting; gh issue #577.
+	if _, ok := oqs.signing[origin]; !ok {
 		return fmt.Errorf(
-			"sendevent: unexpected server to send as: got %q expected %q",
-			origin, oqs.origin,
+			"sendevent: unexpected server to send as %q",
+			origin,
 		)
 	}
 
@@ -303,7 +303,9 @@ func (oqs *OutgoingQueues) SendEDU(
 		destmap[d] = struct{}{}
 	}
 	delete(destmap, oqs.origin)
-	delete(destmap, oqs.signing.ServerName)
+	for local := range oqs.signing {
+		delete(destmap, local)
+	}
 
 	// There is absolutely no guarantee that the EDU will have a room_id
 	// field, as it is not required by the spec. However, if it *does*
diff --git a/federationapi/queue/queue_test.go b/federationapi/queue/queue_test.go
index 7ef4646f7..58745c607 100644
--- a/federationapi/queue/queue_test.go
+++ b/federationapi/queue/queue_test.go
@@ -350,10 +350,12 @@ func testSetup(failuresUntilBlacklist uint32, shouldTxSucceed bool, t *testing.T
 	}
 	rs := &stubFederationRoomServerAPI{}
 	stats := statistics.NewStatistics(db, failuresUntilBlacklist)
-	signingInfo := &SigningInfo{
-		KeyID:      "ed21019:auto",
-		PrivateKey: test.PrivateKeyA,
-		ServerName: "localhost",
+	signingInfo := map[gomatrixserverlib.ServerName]*SigningInfo{
+		"localhost": {
+			KeyID:      "ed21019:auto",
+			PrivateKey: test.PrivateKeyA,
+			ServerName: "localhost",
+		},
 	}
 	queues := NewOutgoingQueues(db, processContext, false, "localhost", fc, rs, &stats, signingInfo)