From e3d393cb2b5f0865f21b049038d2bb2fde55385c Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Thu, 10 Nov 2022 11:07:46 +0000
Subject: [PATCH] Keys endpoint

---
 federationapi/federationapi_keys_test.go |  2 +-
 federationapi/routing/keys.go            | 22 +++++++++++++++++-----
 federationapi/routing/routing.go         |  2 +-
 3 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/federationapi/federationapi_keys_test.go b/federationapi/federationapi_keys_test.go
index 7ccc02f76..3acaa70dc 100644
--- a/federationapi/federationapi_keys_test.go
+++ b/federationapi/federationapi_keys_test.go
@@ -137,7 +137,7 @@ func (m *MockRoundTripper) RoundTrip(req *http.Request) (res *http.Response, err
 	}
 
 	// Get the keys and JSON-ify them.
-	keys := routing.LocalKeys(s.config)
+	keys := routing.LocalKeys(s.config, gomatrixserverlib.ServerName(req.Host))
 	body, err := json.MarshalIndent(keys.JSON, "", "  ")
 	if err != nil {
 		return nil, err
diff --git a/federationapi/routing/keys.go b/federationapi/routing/keys.go
index 8931830f3..070f9174a 100644
--- a/federationapi/routing/keys.go
+++ b/federationapi/routing/keys.go
@@ -16,6 +16,7 @@ package routing
 
 import (
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"time"
 
@@ -134,18 +135,21 @@ func ClaimOneTimeKeys(
 
 // LocalKeys returns the local keys for the server.
 // See https://matrix.org/docs/spec/server_server/unstable.html#publishing-keys
-func LocalKeys(cfg *config.FederationAPI) util.JSONResponse {
-	keys, err := localKeys(cfg, time.Now().Add(cfg.Matrix.KeyValidityPeriod))
+func LocalKeys(cfg *config.FederationAPI, serverName gomatrixserverlib.ServerName) util.JSONResponse {
+	keys, err := localKeys(cfg, serverName, time.Now().Add(cfg.Matrix.KeyValidityPeriod))
 	if err != nil {
 		return util.ErrorResponse(err)
 	}
 	return util.JSONResponse{Code: http.StatusOK, JSON: keys}
 }
 
-func localKeys(cfg *config.FederationAPI, validUntil time.Time) (*gomatrixserverlib.ServerKeys, error) {
+func localKeys(cfg *config.FederationAPI, serverName gomatrixserverlib.ServerName, validUntil time.Time) (*gomatrixserverlib.ServerKeys, error) {
 	var keys gomatrixserverlib.ServerKeys
+	if !cfg.Matrix.IsLocalServerName(serverName) {
+		return nil, fmt.Errorf("server name not known")
+	}
 
-	keys.ServerName = cfg.Matrix.ServerName
+	keys.ServerName = serverName
 	keys.ValidUntilTS = gomatrixserverlib.AsTimestamp(validUntil)
 
 	publicKey := cfg.Matrix.PrivateKey.Public().(ed25519.PublicKey)
@@ -186,6 +190,14 @@ func NotaryKeys(
 	fsAPI federationAPI.FederationInternalAPI,
 	req *gomatrixserverlib.PublicKeyNotaryLookupRequest,
 ) util.JSONResponse {
+	serverName := gomatrixserverlib.ServerName(httpReq.Host) // TODO: this is not ideal
+	if !cfg.Matrix.IsLocalServerName(serverName) {
+		return util.JSONResponse{
+			Code: http.StatusNotFound,
+			JSON: jsonerror.NotFound("Server name not known"),
+		}
+	}
+
 	if req == nil {
 		req = &gomatrixserverlib.PublicKeyNotaryLookupRequest{}
 		if reqErr := clienthttputil.UnmarshalJSONRequest(httpReq, &req); reqErr != nil {
@@ -201,7 +213,7 @@ func NotaryKeys(
 	for serverName, kidToCriteria := range req.ServerKeys {
 		var keyList []gomatrixserverlib.ServerKeys
 		if serverName == cfg.Matrix.ServerName {
-			if k, err := localKeys(cfg, time.Now().Add(cfg.Matrix.KeyValidityPeriod)); err == nil {
+			if k, err := localKeys(cfg, serverName, time.Now().Add(cfg.Matrix.KeyValidityPeriod)); err == nil {
 				keyList = append(keyList, *k)
 			} else {
 				return util.ErrorResponse(err)
diff --git a/federationapi/routing/routing.go b/federationapi/routing/routing.go
index 9f16e5093..0a3ab7a88 100644
--- a/federationapi/routing/routing.go
+++ b/federationapi/routing/routing.go
@@ -74,7 +74,7 @@ func Setup(
 	}
 
 	localKeys := httputil.MakeExternalAPI("localkeys", func(req *http.Request) util.JSONResponse {
-		return LocalKeys(cfg)
+		return LocalKeys(cfg, gomatrixserverlib.ServerName(req.Host))
 	})
 
 	notaryKeys := httputil.MakeExternalAPI("notarykeys", func(req *http.Request) util.JSONResponse {