From b86529f3f87c0f64427aa3026aa6439f4935c724 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Mon, 24 Oct 2022 09:32:28 +0100
Subject: [PATCH] Cap timeout, move cross-signing section

---
 clientapi/routing/keys.go      | 6 +++++-
 keyserver/internal/internal.go | 8 ++++----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/clientapi/routing/keys.go b/clientapi/routing/keys.go
index 5c3681382..0c12b1117 100644
--- a/clientapi/routing/keys.go
+++ b/clientapi/routing/keys.go
@@ -99,7 +99,11 @@ func (r *queryKeysRequest) GetTimeout() time.Duration {
 	if r.Timeout == 0 {
 		return 10 * time.Second
 	}
-	return time.Duration(r.Timeout) * time.Millisecond
+	timeout := time.Duration(r.Timeout) * time.Millisecond
+	if timeout > time.Second*20 {
+		timeout = time.Second * 20
+	}
+	return timeout
 }
 
 func QueryKeys(req *http.Request, keyAPI api.ClientKeyAPI, device *userapi.Device) util.JSONResponse {
diff --git a/keyserver/internal/internal.go b/keyserver/internal/internal.go
index 82c14b3bf..26adaffbd 100644
--- a/keyserver/internal/internal.go
+++ b/keyserver/internal/internal.go
@@ -257,10 +257,7 @@ func (a *KeyInternalAPI) QueryKeys(ctx context.Context, req *api.QueryKeysReques
 	res.UserSigningKeys = make(map[string]gomatrixserverlib.CrossSigningKey)
 	res.Failures = make(map[string]interface{})
 
-	logrus.Print("QueryKeys:", req.UserID, req.UserToDevices)
-
-	// get cross-signing keys from the database
-	a.crossSigningKeysFromDatabase(ctx, req, res)
+	logrus.Print("QueryKeys:", req.UserID, req.Timeout, req.UserToDevices)
 
 	// make a map from domain to device keys
 	domainToDeviceKeys := make(map[string]map[string][]string)
@@ -338,6 +335,9 @@ func (a *KeyInternalAPI) QueryKeys(ctx context.Context, req *api.QueryKeysReques
 		a.queryRemoteKeys(ctx, req.Timeout, res, domainToDeviceKeys, domainToCrossSigningKeys)
 	}
 
+	// get cross-signing keys from the database
+	a.crossSigningKeysFromDatabase(ctx, req, res)
+
 	// Finally, append signatures that we know about
 	// TODO: This is horrible because we need to round-trip the signature from
 	// JSON, add the signatures and marshal it again, for some reason?