From 607819f42507d6a3b18ef7c44f98ed8f862a7f78 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Thu, 17 Nov 2022 09:26:56 +0000
Subject: [PATCH] Fix `/key/v2/server`, add HTTP `Host` matching

---
 federationapi/routing/keys.go | 31 ++++++++++++++-----------------
 setup/config/config_global.go |  5 +++++
 2 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/federationapi/routing/keys.go b/federationapi/routing/keys.go
index 8194c9905..b2ef1dba4 100644
--- a/federationapi/routing/keys.go
+++ b/federationapi/routing/keys.go
@@ -16,7 +16,6 @@ package routing
 
 import (
 	"encoding/json"
-	"net"
 	"net/http"
 	"time"
 
@@ -146,14 +145,26 @@ func LocalKeys(cfg *config.FederationAPI, serverName gomatrixserverlib.ServerNam
 func localKeys(cfg *config.FederationAPI, serverName gomatrixserverlib.ServerName) (*gomatrixserverlib.ServerKeys, error) {
 	var keys gomatrixserverlib.ServerKeys
 	var virtualHost *config.VirtualHost
+loop:
 	for _, v := range cfg.Matrix.VirtualHosts {
 		if v.ServerName == serverName {
 			virtualHost = v
-			break
+			break loop
+		}
+		for _, httpHost := range v.MatchHTTPHosts {
+			if httpHost == serverName {
+				virtualHost = v
+				break loop
+			}
 		}
 	}
 
-	if virtualHost == nil {
+	identity, err := cfg.Matrix.SigningIdentityFor(serverName)
+	if err != nil {
+		identity, _ = cfg.Matrix.SigningIdentityFor(cfg.Matrix.ServerName)
+	}
+
+	if identity.ServerName == serverName {
 		publicKey := cfg.Matrix.PrivateKey.Public().(ed25519.PublicKey)
 		keys.ServerName = cfg.Matrix.ServerName
 		keys.ValidUntilTS = gomatrixserverlib.AsTimestamp(time.Now().Add(cfg.Matrix.KeyValidityPeriod))
@@ -189,20 +200,6 @@ func localKeys(cfg *config.FederationAPI, serverName gomatrixserverlib.ServerNam
 		return nil, err
 	}
 
-	identity, err := cfg.Matrix.SigningIdentityFor(serverName)
-	if err != nil {
-		// TODO: This is a bit of a hack because the Host header can contain a port
-		// number if it's specified in the well-known file. Try getting a signing
-		// identity without it to see if that helps.
-		var h string
-		if h, _, err = net.SplitHostPort(string(serverName)); err == nil {
-			identity, err = cfg.Matrix.SigningIdentityFor(gomatrixserverlib.ServerName(h))
-		}
-		if err != nil {
-			return nil, err
-		}
-	}
-
 	keys.Raw, err = gomatrixserverlib.SignJSON(
 		string(identity.ServerName), identity.KeyID, identity.PrivateKey, toSign,
 	)
diff --git a/setup/config/config_global.go b/setup/config/config_global.go
index f2fdd021e..722230d9a 100644
--- a/setup/config/config_global.go
+++ b/setup/config/config_global.go
@@ -195,6 +195,11 @@ type VirtualHost struct {
 	// Defaults to 24 hours.
 	KeyValidityPeriod time.Duration `yaml:"key_validity_period"`
 
+	// Match these HTTP Host headers on the `/key/v2/server` endpoint, this needs
+	// to match all delegated names, likely including the port number too if
+	// the well-known delegation includes that also.
+	MatchHTTPHosts []gomatrixserverlib.ServerName `yaml:"match_http_hosts"`
+
 	// Is registration enabled on this virtual host?
 	AllowRegistration bool `json:"allow_registration"`
 }