More server key updates, tests (#1129)

* More key tweaks

* Start testing stuff

* Move responsibility for generating local keys into server key API, don't register prom in caches unless needed, start tests

* Don't store our own keys in the database

* Don't store our own keys in the database

* Don't run tests for now

* Tweak caching behaviour, update tests

* Update comments, add fixes from forward-merge

* Debug logging

* Debug logging

* Perform final comparison against original set of requests

* oops

* Fetcher timeouts

* Fetcher timeouts

* missing func

* Tweaks

* Update gomatrixserverlib

* Fix Federation API test

* Break up FetchKeys

* Add comments to caching

* Add URL check in test

* Partially revert "Move responsibility for generating local keys into server key API, don't register prom in caches unless needed, start tests"

This reverts commit d7eb54c5b30b2f6a9d6514b643e32e6ad2b602f3.

* Fix federation API test

* Fix internal cache stuff again

* Fix server key API test

* Update comments

* Update comments from review

* Fix lint

13 files changed, 538 insertions, 111 deletions

diff --git a/cmd/roomserver-integration-tests/main.go b/cmd/roomserver-integration-tests/main.go
index 43aca078..3860ca1f 100644
--- a/cmd/roomserver-integration-tests/main.go
+++ b/cmd/roomserver-integration-tests/main.go
@@ -255,7 +255,7 @@ func testRoomserver(input []string, wantOutput []string, checkQueries func(api.R
 		panic(err)
 	}
 
-	cache, err := caching.NewInMemoryLRUCache()
+	cache, err := caching.NewInMemoryLRUCache(false)
 	if err != nil {
 		panic(err)
 	}
diff --git a/go.mod b/go.mod
index 2f787408..b2451d85 100644
--- a/go.mod
+++ b/go.mod
@@ -20,7 +20,7 @@ require (
 	github.com/matrix-org/go-http-js-libp2p v0.0.0-20200518170932-783164aeeda4
 	github.com/matrix-org/go-sqlite3-js v0.0.0-20200522092705-bc8506ccbcf3
 	github.com/matrix-org/gomatrix v0.0.0-20190528120928-7df988a63f26
-	github.com/matrix-org/gomatrixserverlib v0.0.0-20200608125510-defe251235b1
+	github.com/matrix-org/gomatrixserverlib v0.0.0-20200615161710-f69539c86ea5
 	github.com/matrix-org/naffka v0.0.0-20200422140631-181f1ee7401f
 	github.com/matrix-org/util v0.0.0-20190711121626-527ce5ddefc7
 	github.com/mattn/go-sqlite3 v2.0.2+incompatible
@@ -38,7 +38,6 @@ require (
 	github.com/yggdrasil-network/yggdrasil-go v0.3.15-0.20200530233943-aec82d7a391b
 	go.uber.org/atomic v1.4.0
 	golang.org/x/crypto v0.0.0-20200221231518-2aa609cf4a9d
-	golang.org/x/tools v0.0.0-20200612022331-742c5eb664c2 // indirect
 	gopkg.in/h2non/bimg.v1 v1.0.18
 	gopkg.in/yaml.v2 v2.2.8
 )
diff --git a/go.sum b/go.sum
index 301066b9..2578e175 100644
--- a/go.sum
+++ b/go.sum
@@ -131,7 +131,6 @@ github.com/hashicorp/golang-lru v0.5.3/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uG
 github.com/hashicorp/golang-lru v0.5.4 h1:YDjusn29QI/Das2iO9M0BHnIbxPeyuCHsjMW+lJfyTc=
 github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
-github.com/hjson/hjson-go v3.0.1-0.20190209023717-9147687966d9+incompatible/go.mod h1:qsetwF8NlsTsOTwZTApNlTCerV+b2GjYRRcIk4JMFio=
 github.com/hjson/hjson-go v3.0.2-0.20200316202735-d5d0e8b0617d+incompatible/go.mod h1:qsetwF8NlsTsOTwZTApNlTCerV+b2GjYRRcIk4JMFio=
 github.com/hpcloud/tail v1.0.0 h1:nfCOvKYfkgYP8hkirhJocXT2+zOD8yUNjXaWfTlyFKI=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
@@ -372,8 +371,8 @@ github.com/matrix-org/go-sqlite3-js v0.0.0-20200522092705-bc8506ccbcf3 h1:Yb+Wlf
 github.com/matrix-org/go-sqlite3-js v0.0.0-20200522092705-bc8506ccbcf3/go.mod h1:e+cg2q7C7yE5QnAXgzo512tgFh1RbQLC0+jozuegKgo=
 github.com/matrix-org/gomatrix v0.0.0-20190528120928-7df988a63f26 h1:Hr3zjRsq2bhrnp3Ky1qgx/fzCtCALOoGYylh2tpS9K4=
 github.com/matrix-org/gomatrix v0.0.0-20190528120928-7df988a63f26/go.mod h1:3fxX6gUjWyI/2Bt7J1OLhpCzOfO/bB3AiX0cJtEKud0=
-github.com/matrix-org/gomatrixserverlib v0.0.0-20200608125510-defe251235b1 h1:BfrvDrbjoPBvYua/3F/FmrqiZTRGrvtoMRgCVnrufMI=
-github.com/matrix-org/gomatrixserverlib v0.0.0-20200608125510-defe251235b1/go.mod h1:JsAzE1Ll3+gDWS9JSUHPJiiyAksvOOnGWF2nXdg4ZzU=
+github.com/matrix-org/gomatrixserverlib v0.0.0-20200615161710-f69539c86ea5 h1:VN7DoSFVkQF9Bv+TWuBWHLgAz9Nw9UiahFfe2oE6uiQ=
+github.com/matrix-org/gomatrixserverlib v0.0.0-20200615161710-f69539c86ea5/go.mod h1:JsAzE1Ll3+gDWS9JSUHPJiiyAksvOOnGWF2nXdg4ZzU=
 github.com/matrix-org/naffka v0.0.0-20200422140631-181f1ee7401f h1:pRz4VTiRCO4zPlEMc3ESdUOcW4PXHH4Kj+YDz1XyE+Y=
 github.com/matrix-org/naffka v0.0.0-20200422140631-181f1ee7401f/go.mod h1:y0oDTjZDv5SM9a2rp3bl+CU+bvTRINQsdb7YlDql5Go=
 github.com/matrix-org/util v0.0.0-20190711121626-527ce5ddefc7 h1:ntrLa/8xVzeSs8vHFHK25k0C+NV74sYMJnNSg5NoSRo=
@@ -566,11 +565,8 @@ github.com/xdg/scram v0.0.0-20180814205039-7eeb5667e42c/go.mod h1:lB8K/P019DLNhe
 github.com/xdg/stringprep v1.0.0/go.mod h1:Jhud4/sHMO4oL310DaZAKk9ZaJ08SJfe+sJh0HrGL1Y=
 github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q=
 github.com/yggdrasil-network/yggdrasil-extras v0.0.0-20200525205615-6c8a4a2e8855/go.mod h1:xQdsh08Io6nV4WRnOVTe6gI8/2iTvfLDQ0CYa5aMt+I=
-github.com/yggdrasil-network/yggdrasil-go v0.3.14 h1:vWzYzCQxOruS+J5FkLfXOS0JhCJx1yI9Erj/h2wfZ/E=
-github.com/yggdrasil-network/yggdrasil-go v0.3.14/go.mod h1:rkQzLzVHlFdzsEMG+bDdTI+KeWPCZq1HpXRFzwinf6M=
 github.com/yggdrasil-network/yggdrasil-go v0.3.15-0.20200530233943-aec82d7a391b h1:ELOisSxFXCcptRs4LFub+Hz5fYUvV12wZrTps99Eb3E=
 github.com/yggdrasil-network/yggdrasil-go v0.3.15-0.20200530233943-aec82d7a391b/go.mod h1:d+Nz6SPeG6kmeSPFL0cvfWfgwEql75fUnZiAONgvyBE=
-github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 go.opencensus.io v0.21.0/go.mod h1:mSImk1erAIZhrmZN+AvHh14ztQfjbGwt4TtuofqLduU=
 go.opencensus.io v0.22.1/go.mod h1:Ap50jQcDJrx6rB6VgeeFPtuPIf3wMRvRfrfYDO6+BmA=
 go.opencensus.io v0.22.2/go.mod h1:yxeiOL68Rb0Xd1ddK5vPZ/oVn4vY4Ynel7k9FzqtOIw=
@@ -604,8 +600,6 @@ golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL
 golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
 golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
 golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
-golang.org/x/mod v0.2.0 h1:KU7oHjnv3XNWfa5COkzUifxZmxp1TyI7ImMXqFxLwvQ=
-golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
@@ -675,9 +669,6 @@ golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGm
 golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=
 golang.org/x/tools v0.0.0-20190311212946-11955173bddd h1:/e+gpKk9r3dJobndpTytxS2gOy6m5uvpg+ISQoEcusQ=
 golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
-golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
-golang.org/x/tools v0.0.0-20200612022331-742c5eb664c2 h1:DVqHa33CzfnTKwUV6be+I4hp31W6iXn3ZiEcdKGzLyI=
-golang.org/x/tools v0.0.0-20200612022331-742c5eb664c2/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
diff --git a/internal/caching/cache_serverkeys.go b/internal/caching/cache_serverkeys.go
index b5e31575..4697fb4d 100644
--- a/internal/caching/cache_serverkeys.go
+++ b/internal/caching/cache_serverkeys.go
@@ -2,7 +2,6 @@ package caching
 
 import (
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/gomatrixserverlib"
 )
@@ -16,22 +15,29 @@ const (
 // ServerKeyCache contains the subset of functions needed for
 // a server key cache.
 type ServerKeyCache interface {
-	GetServerKey(request gomatrixserverlib.PublicKeyLookupRequest) (response gomatrixserverlib.PublicKeyLookupResult, ok bool)
+	// request -> timestamp is emulating gomatrixserverlib.FetchKeys:
+	// https://github.com/matrix-org/gomatrixserverlib/blob/f69539c86ea55d1e2cc76fd8e944e2d82d30397c/keyring.go#L95
+	// The timestamp should be the timestamp of the event that is being
+	// verified. We will not return keys from the cache that are not valid
+	// at this timestamp.
+	GetServerKey(request gomatrixserverlib.PublicKeyLookupRequest, timestamp gomatrixserverlib.Timestamp) (response gomatrixserverlib.PublicKeyLookupResult, ok bool)
+
+	// request -> result is emulating gomatrixserverlib.StoreKeys:
+	// https://github.com/matrix-org/gomatrixserverlib/blob/f69539c86ea55d1e2cc76fd8e944e2d82d30397c/keyring.go#L112
 	StoreServerKey(request gomatrixserverlib.PublicKeyLookupRequest, response gomatrixserverlib.PublicKeyLookupResult)
 }
 
 func (c Caches) GetServerKey(
 	request gomatrixserverlib.PublicKeyLookupRequest,
+	timestamp gomatrixserverlib.Timestamp,
 ) (gomatrixserverlib.PublicKeyLookupResult, bool) {
 	key := fmt.Sprintf("%s/%s", request.ServerName, request.KeyID)
-	now := gomatrixserverlib.AsTimestamp(time.Now())
 	val, found := c.ServerKeys.Get(key)
 	if found && val != nil {
 		if keyLookupResult, ok := val.(gomatrixserverlib.PublicKeyLookupResult); ok {
-			if !keyLookupResult.WasValidAt(now, true) {
-				// We appear to be past the key validity so don't return this
-				// with the results. This ensures that the cache doesn't return
-				// values that are not useful to us.
+			if !keyLookupResult.WasValidAt(timestamp, true) {
+				// The key wasn't valid at the requested timestamp so don't
+				// return it. The caller will have to work out what to do.
 				c.ServerKeys.Unset(key)
 				return gomatrixserverlib.PublicKeyLookupResult{}, false
 			}
diff --git a/internal/caching/impl_inmemorylru.go b/internal/caching/impl_inmemorylru.go
index 158deca4..7bb791dd 100644
--- a/internal/caching/impl_inmemorylru.go
+++ b/internal/caching/impl_inmemorylru.go
@@ -8,11 +8,12 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promauto"
 )
 
-func NewInMemoryLRUCache() (*Caches, error) {
+func NewInMemoryLRUCache(enablePrometheus bool) (*Caches, error) {
 	roomVersions, err := NewInMemoryLRUCachePartition(
 		RoomVersionCacheName,
 		RoomVersionCacheMutable,
 		RoomVersionCacheMaxEntries,
+		enablePrometheus,
 	)
 	if err != nil {
 		return nil, err
@@ -21,6 +22,7 @@ func NewInMemoryLRUCache() (*Caches, error) {
 		ServerKeyCacheName,
 		ServerKeyCacheMutable,
 		ServerKeyCacheMaxEntries,
+		enablePrometheus,
 	)
 	if err != nil {
 		return nil, err
@@ -38,7 +40,7 @@ type InMemoryLRUCachePartition struct {
 	lru        *lru.Cache
 }
 
-func NewInMemoryLRUCachePartition(name string, mutable bool, maxEntries int) (*InMemoryLRUCachePartition, error) {
+func NewInMemoryLRUCachePartition(name string, mutable bool, maxEntries int, enablePrometheus bool) (*InMemoryLRUCachePartition, error) {
 	var err error
 	cache := InMemoryLRUCachePartition{
 		name:       name,
@@ -49,13 +51,15 @@ func NewInMemoryLRUCachePartition(name string, mutable bool, maxEntries int) (*I
 	if err != nil {
 		return nil, err
 	}
-	promauto.NewGaugeFunc(prometheus.GaugeOpts{
-		Namespace: "dendrite",
-		Subsystem: "caching_in_memory_lru",
-		Name:      name,
-	}, func() float64 {
-		return float64(cache.lru.Len())
-	})
+	if enablePrometheus {
+		promauto.NewGaugeFunc(prometheus.GaugeOpts{
+			Namespace: "dendrite",
+			Subsystem: "caching_in_memory_lru",
+			Name:      name,
+		}, func() float64 {
+			return float64(cache.lru.Len())
+		})
+	}
 	return &cache, nil
 }
 
diff --git a/internal/setup/base.go b/internal/setup/base.go
index 414b7964..59bdfd2e 100644
--- a/internal/setup/base.go
+++ b/internal/setup/base.go
@@ -96,7 +96,7 @@ func NewBaseDendrite(cfg *config.Dendrite, componentName string, useHTTPAPIs boo
 		kafkaConsumer, kafkaProducer = setupKafka(cfg)
 	}
 
-	cache, err := caching.NewInMemoryLRUCache()
+	cache, err := caching.NewInMemoryLRUCache(true)
 	if err != nil {
 		logrus.WithError(err).Warnf("Failed to create cache")
 	}
diff --git a/serverkeyapi/internal/api.go b/serverkeyapi/internal/api.go
index 7a35aa8e..02028c60 100644
--- a/serverkeyapi/internal/api.go
+++ b/serverkeyapi/internal/api.go
@@ -2,16 +2,23 @@ package internal
 
 import (
 	"context"
+	"crypto/ed25519"
 	"fmt"
 	"time"
 
 	"github.com/matrix-org/dendrite/serverkeyapi/api"
 	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/sirupsen/logrus"
 )
 
 type ServerKeyAPI struct {
 	api.ServerKeyInternalAPI
 
+	ServerName        gomatrixserverlib.ServerName
+	ServerPublicKey   ed25519.PublicKey
+	ServerKeyID       gomatrixserverlib.KeyID
+	ServerKeyValidity time.Duration
+
 	OurKeyRing gomatrixserverlib.KeyRing
 	FedClient  *gomatrixserverlib.FederationClient
 }
@@ -33,6 +40,7 @@ func (s *ServerKeyAPI) StoreKeys(
 	// Run in a background context - we don't want to stop this work just
 	// because the caller gives up waiting.
 	ctx := context.Background()
+
 	// Store any keys that we were given in our database.
 	return s.OurKeyRing.KeyDatabase.StoreKeys(ctx, results)
 }
@@ -44,52 +52,57 @@ func (s *ServerKeyAPI) FetchKeys(
 	// Run in a background context - we don't want to stop this work just
 	// because the caller gives up waiting.
 	ctx := context.Background()
-	results := map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult{}
 	now := gomatrixserverlib.AsTimestamp(time.Now())
-	// First consult our local database and see if we have the requested
+	results := map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult{}
+	origRequests := map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp{}
+	for k, v := range requests {
+		origRequests[k] = v
+	}
+
+	// First, check if any of these key checks are for our own keys. If
+	// they are then we will satisfy them directly.
+	s.handleLocalKeys(ctx, requests, results)
+
+	// Then consult our local database and see if we have the requested
 	// keys. These might come from a cache, depending on the database
 	// implementation used.
-	if dbResults, err := s.OurKeyRing.KeyDatabase.FetchKeys(ctx, requests); err == nil {
-		// We successfully got some keys. Add them to the results and
-		// remove them from the request list.
-		for req, res := range dbResults {
-			if !res.WasValidAt(now, true) {
-				// We appear to be past the key validity. Don't return this
-				// key with the results.
-				continue
-			}
-			results[req] = res
-			delete(requests, req)
-		}
+	if err := s.handleDatabaseKeys(ctx, now, requests, results); err != nil {
+		return nil, err
 	}
+
 	// For any key requests that we still have outstanding, next try to
 	// fetch them directly. We'll go through each of the key fetchers to
-	// ask for the remaining keys.
+	// ask for the remaining keys
 	for _, fetcher := range s.OurKeyRing.KeyFetchers {
+		// If there are no more keys to look up then stop.
 		if len(requests) == 0 {
 			break
 		}
-		if fetcherResults, err := fetcher.FetchKeys(ctx, requests); err == nil {
-			// We successfully got some keys. Add them to the results and
-			// remove them from the request list.
-			for req, res := range fetcherResults {
-				if !res.WasValidAt(now, true) {
-					// We appear to be past the key validity. Don't return this
-					// key with the results.
-					continue
-				}
-				results[req] = res
-				delete(requests, req)
-			}
-			if err = s.OurKeyRing.KeyDatabase.StoreKeys(ctx, fetcherResults); err != nil {
-				return nil, fmt.Errorf("server key API failed to store retrieved keys: %w", err)
-			}
+
+		// Ask the fetcher to look up our keys.
+		if err := s.handleFetcherKeys(ctx, now, fetcher, requests, results); err != nil {
+			logrus.WithError(err).WithFields(logrus.Fields{
+				"fetcher_name": fetcher.FetcherName(),
+			}).Errorf("Failed to retrieve %d key(s)", len(requests))
+			continue
 		}
 	}
-	// If we failed to fetch any keys then we should report an error.
-	if len(requests) > 0 {
-		return results, fmt.Errorf("server key API failed to fetch %d keys", len(requests))
+
+	// Check that we've actually satisfied all of the key requests that we
+	// were given. We should report an error if we didn't.
+	for req := range origRequests {
+		if _, ok := results[req]; !ok {
+			// The results don't contain anything for this specific request, so
+			// we've failed to satisfy it from local keys, database keys or from
+			// all of the fetchers. Report an error.
+			logrus.Warnf("Failed to retrieve key %q for server %q", req.KeyID, req.ServerName)
+			return results, fmt.Errorf(
+				"server key API failed to satisfy key request for server %q key ID %q",
+				req.ServerName, req.KeyID,
+			)
+		}
 	}
+
 	// Return the keys.
 	return results, nil
 }
@@ -97,3 +110,141 @@ func (s *ServerKeyAPI) FetchKeys(
 func (s *ServerKeyAPI) FetcherName() string {
 	return fmt.Sprintf("ServerKeyAPI (wrapping %q)", s.OurKeyRing.KeyDatabase.FetcherName())
 }
+
+// handleLocalKeys handles cases where the key request contains
+// a request for our own server keys.
+func (s *ServerKeyAPI) handleLocalKeys(
+	_ context.Context,
+	requests map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp,
+	results map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult,
+) {
+	for req := range requests {
+		if req.ServerName == s.ServerName {
+			// We found a key request that is supposed to be for our own
+			// keys. Remove it from the request list so we don't hit the
+			// database or the fetchers for it.
+			delete(requests, req)
+
+			// Insert our own key into the response.
+			results[req] = gomatrixserverlib.PublicKeyLookupResult{
+				VerifyKey: gomatrixserverlib.VerifyKey{
+					Key: gomatrixserverlib.Base64Bytes(s.ServerPublicKey),
+				},
+				ExpiredTS:    gomatrixserverlib.PublicKeyNotExpired,
+				ValidUntilTS: gomatrixserverlib.AsTimestamp(time.Now().Add(s.ServerKeyValidity)),
+			}
+		}
+	}
+}
+
+// handleDatabaseKeys handles cases where the key requests can be
+// satisfied from our local database/cache.
+func (s *ServerKeyAPI) handleDatabaseKeys(
+	ctx context.Context,
+	now gomatrixserverlib.Timestamp,
+	requests map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp,
+	results map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult,
+) error {
+	// Ask the database/cache for the keys.
+	dbResults, err := s.OurKeyRing.KeyDatabase.FetchKeys(ctx, requests)
+	if err != nil {
+		return err
+	}
+
+	// We successfully got some keys. Add them to the results.
+	for req, res := range dbResults {
+		// The key we've retrieved from the database/cache might
+		// have passed its validity period, but right now, it's
+		// the best thing we've got, and it might be sufficient to
+		// verify a past event.
+		results[req] = res
+
+		// If the key is valid right now then we can also remove it
+		// from the request list as we don't need to fetch it again
+		// in that case. If the key isn't valid right now, then by
+		// leaving it in the 'requests' map, we'll try to update the
+		// key using the fetchers in handleFetcherKeys.
+		if res.WasValidAt(now, true) {
+			delete(requests, req)
+		}
+	}
+	return nil
+}
+
+// handleFetcherKeys handles cases where a fetcher can satisfy
+// the remaining requests.
+func (s *ServerKeyAPI) handleFetcherKeys(
+	ctx context.Context,
+	now gomatrixserverlib.Timestamp,
+	fetcher gomatrixserverlib.KeyFetcher,
+	requests map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp,
+	results map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult,
+) error {
+	logrus.WithFields(logrus.Fields{
+		"fetcher_name": fetcher.FetcherName(),
+	}).Infof("Fetching %d key(s)", len(requests))
+
+	// Create a context that limits our requests to 30 seconds.
+	fetcherCtx, fetcherCancel := context.WithTimeout(ctx, time.Second*30)
+	defer fetcherCancel()
+
+	// Try to fetch the keys.
+	fetcherResults, err := fetcher.FetchKeys(fetcherCtx, requests)
+	if err != nil {
+		return err
+	}
+
+	// Build a map of the results that we want to commit to the
+	// database. We do this in a separate map because otherwise we
+	// might end up trying to rewrite database entries.
+	storeResults := map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult{}
+
+	// Now let's look at the results that we got from this fetcher.
+	for req, res := range fetcherResults {
+		if prev, ok := results[req]; ok {
+			// We've already got a previous entry for this request
+			// so let's see if the newly retrieved one contains a more
+			// up-to-date validity period.
+			if res.ValidUntilTS > prev.ValidUntilTS {
+				// This key is newer than the one we had so let's store
+				// it in the database.
+				if req.ServerName != s.ServerName {
+					storeResults[req] = res
+				}
+			}
+		} else {
+			// We didn't already have a previous entry for this request
+			// so store it in the database anyway for now.
+			if req.ServerName != s.ServerName {
+				storeResults[req] = res
+			}
+		}
+
+		// Update the results map with this new result. If nothing
+		// else, we can try verifying against this key.
+		results[req] = res
+
+		// If the key is valid right now then we can remove it from the
+		// request list as we won't need to re-fetch it.
+		if res.WasValidAt(now, true) {
+			delete(requests, req)
+		}
+	}
+
+	// Store the keys from our store map.
+	if err = s.OurKeyRing.KeyDatabase.StoreKeys(ctx, storeResults); err != nil {
+		logrus.WithError(err).WithFields(logrus.Fields{
+			"fetcher_name":  fetcher.FetcherName(),
+			"database_name": s.OurKeyRing.KeyDatabase.FetcherName(),
+		}).Errorf("Failed to store keys in the database")
+		return fmt.Errorf("server key API failed to store retrieved keys: %w", err)
+	}
+
+	if len(storeResults) > 0 {
+		logrus.WithFields(logrus.Fields{
+			"fetcher_name": fetcher.FetcherName(),
+		}).Infof("Updated %d of %d key(s) in database", len(storeResults), len(results))
+	}
+
+	return nil
+}
diff --git a/serverkeyapi/inthttp/client.go b/serverkeyapi/inthttp/client.go
index e84cf47f..39ab8c6c 100644
--- a/serverkeyapi/inthttp/client.go
+++ b/serverkeyapi/inthttp/client.go
@@ -90,7 +90,7 @@ func (s *httpServerKeyInternalAPI) FetchKeys(
 		Results: make(map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult),
 	}
 	for req, ts := range requests {
-		if res, ok := s.cache.GetServerKey(req); ok {
+		if res, ok := s.cache.GetServerKey(req, ts); ok {
 			result[req] = res
 			continue
 		}
diff --git a/serverkeyapi/serverkeyapi.go b/serverkeyapi/serverkeyapi.go
index 58ca00b7..cddd392e 100644
--- a/serverkeyapi/serverkeyapi.go
+++ b/serverkeyapi/serverkeyapi.go
@@ -46,7 +46,11 @@ func NewInternalAPI(
 	}
 
 	internalAPI := internal.ServerKeyAPI{
-		FedClient: fedClient,
+		ServerName:        cfg.Matrix.ServerName,
+		ServerPublicKey:   cfg.Matrix.PrivateKey.Public().(ed25519.PublicKey),
+		ServerKeyID:       cfg.Matrix.KeyID,
+		ServerKeyValidity: cfg.Matrix.KeyValidityPeriod,
+		FedClient:         fedClient,
 		OurKeyRing: gomatrixserverlib.KeyRing{
 			KeyFetchers: []gomatrixserverlib.KeyFetcher{
 				&gomatrixserverlib.DirectKeyFetcher{
diff --git a/serverkeyapi/serverkeyapi_test.go b/serverkeyapi/serverkeyapi_test.go
new file mode 100644
index 00000000..3368f5b2
--- /dev/null
+++ b/serverkeyapi/serverkeyapi_test.go
@@ -0,0 +1,315 @@
+package serverkeyapi
+
+import (
+	"bytes"
+	"context"
+	"crypto/ed25519"
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"os"
+	"reflect"
+	"testing"
+	"time"
+
+	"github.com/matrix-org/dendrite/federationapi/routing"
+	"github.com/matrix-org/dendrite/internal/caching"
+	"github.com/matrix-org/dendrite/internal/config"
+	"github.com/matrix-org/dendrite/serverkeyapi/api"
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+type server struct {
+	name      gomatrixserverlib.ServerName        // server name
+	validity  time.Duration                       // key validity duration from now
+	config    *config.Dendrite                    // skeleton config, from TestMain
+	fedclient *gomatrixserverlib.FederationClient // uses MockRoundTripper
+	cache     *caching.Caches                     // server-specific cache
+	api       api.ServerKeyInternalAPI            // server-specific server key API
+}
+
+func (s *server) renew() {
+	// This updates the validity period to be an hour in the
+	// future, which is particularly useful in server A and
+	// server C's cases which have validity either as now or
+	// in the past.
+	s.validity = time.Hour
+	s.config.Matrix.KeyValidityPeriod = s.validity
+}
+
+var (
+	serverKeyID = gomatrixserverlib.KeyID("ed25519:auto")
+	serverA     = &server{name: "a.com", validity: time.Duration(0)} // expires now
+	serverB     = &server{name: "b.com", validity: time.Hour}        // expires in an hour
+	serverC     = &server{name: "c.com", validity: -time.Hour}       // expired an hour ago
+)
+
+var servers = map[string]*server{
+	"a.com": serverA,
+	"b.com": serverB,
+	"c.com": serverC,
+}
+
+func TestMain(m *testing.M) {
+	// Set up the server key API for each "server" that we
+	// will use in our tests.
+	for _, s := range servers {
+		// Generate a new key.
+		_, testPriv, err := ed25519.GenerateKey(nil)
+		if err != nil {
+			panic("can't generate identity key: " + err.Error())
+		}
+
+		// Create a new cache but don't enable prometheus!
+		s.cache, err = caching.NewInMemoryLRUCache(false)
+		if err != nil {
+			panic("can't create cache: " + err.Error())
+		}
+
+		// Draw up just enough Dendrite config for the server key
+		// API to work.
+		s.config = &config.Dendrite{}
+		s.config.SetDefaults()
+		s.config.Matrix.ServerName = gomatrixserverlib.ServerName(s.name)
+		s.config.Matrix.PrivateKey = testPriv
+		s.config.Matrix.KeyID = serverKeyID
+		s.config.Matrix.KeyValidityPeriod = s.validity
+		s.config.Database.ServerKey = config.DataSource("file::memory:")
+
+		// Create a transport which redirects federation requests to
+		// the mock round tripper. Since we're not *really* listening for
+		// federation requests then this will return the key instead.
+		transport := &http.Transport{}
+		transport.RegisterProtocol("matrix", &MockRoundTripper{})
+
+		// Create the federation client.
+		s.fedclient = gomatrixserverlib.NewFederationClientWithTransport(
+			s.config.Matrix.ServerName, serverKeyID, testPriv, transport,
+		)
+
+		// Finally, build the server key APIs.
+		s.api = NewInternalAPI(s.config, s.fedclient, s.cache)
+	}
+
+	// Now that we have built our server key APIs, start the
+	// rest of the tests.
+	os.Exit(m.Run())
+}
+
+type MockRoundTripper struct{}
+
+func (m *MockRoundTripper) RoundTrip(req *http.Request) (res *http.Response, err error) {
+	// Check if the request is looking for keys from a server that
+	// we know about in the test. The only reason this should go wrong
+	// is if the test is broken.
+	s, ok := servers[req.Host]
+	if !ok {
+		return nil, fmt.Errorf("server not known: %s", req.Host)
+	}
+
+	// We're intercepting /matrix/key/v2/server requests here, so check
+	// that the URL supplied in the request is for that.
+	if req.URL.Path != "/_matrix/key/v2/server" {
+		return nil, fmt.Errorf("unexpected request path: %s", req.URL.Path)
+	}
+
+	// Get the keys and JSON-ify them.
+	keys := routing.LocalKeys(s.config)
+	body, err := json.MarshalIndent(keys.JSON, "", "  ")
+	if err != nil {
+		return nil, err
+	}
+
+	// And respond.
+	res = &http.Response{
+		StatusCode: 200,
+		Body:       ioutil.NopCloser(bytes.NewReader(body)),
+	}
+	return
+}
+
+func TestServersRequestOwnKeys(t *testing.T) {
+	// Each server will request its own keys. There's no reason
+	// for this to fail as each server should know its own keys.
+
+	for name, s := range servers {
+		req := gomatrixserverlib.PublicKeyLookupRequest{
+			ServerName: s.name,
+			KeyID:      serverKeyID,
+		}
+		res, err := s.api.FetchKeys(
+			context.Background(),
+			map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp{
+				req: gomatrixserverlib.AsTimestamp(time.Now()),
+			},
+		)
+		if err != nil {
+			t.Fatalf("server could not fetch own key: %s", err)
+		}
+		if _, ok := res[req]; !ok {
+			t.Fatalf("server didn't return its own key in the results")
+		}
+		t.Logf("%s's key expires at %s\n", name, res[req].ValidUntilTS.Time())
+	}
+}
+
+func TestCachingBehaviour(t *testing.T) {
+	// Server A will request Server B's key, which has a validity
+	// period of an hour from now. We should retrieve the key and
+	// it should make it into the cache automatically.
+
+	req := gomatrixserverlib.PublicKeyLookupRequest{
+		ServerName: serverB.name,
+		KeyID:      serverKeyID,
+	}
+	ts := gomatrixserverlib.AsTimestamp(time.Now())
+
+	res, err := serverA.api.FetchKeys(
+		context.Background(),
+		map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp{
+			req: ts,
+		},
+	)
+	if err != nil {
+		t.Fatalf("server A failed to retrieve server B key: %s", err)
+	}
+	if len(res) != 1 {
+		t.Fatalf("server B should have returned one key but instead returned %d keys", len(res))
+	}
+	if _, ok := res[req]; !ok {
+		t.Fatalf("server B isn't included in the key fetch response")
+	}
+
+	// At this point, if the previous key request was a success,
+	// then the cache should now contain the key. Check if that's
+	// the case - if it isn't then there's something wrong with
+	// the cache implementation or we failed to get the key.
+
+	cres, ok := serverA.cache.GetServerKey(req, ts)
+	if !ok {
+		t.Fatalf("server B key should be in cache but isn't")
+	}
+	if !reflect.DeepEqual(cres, res[req]) {
+		t.Fatalf("the cached result from server B wasn't what server B gave us")
+	}
+
+	// If we ask the cache for the same key but this time for an event
+	// that happened in +30 minutes. Since the validity period is for
+	// another hour, then we should get a response back from the cache.
+
+	_, ok = serverA.cache.GetServerKey(
+		req,
+		gomatrixserverlib.AsTimestamp(time.Now().Add(time.Minute*30)),
+	)
+	if !ok {
+		t.Fatalf("server B key isn't in cache when it should be (+30 minutes)")
+	}
+
+	// If we ask the cache for the same key but this time for an event
+	// that happened in +90 minutes then we should expect to get no
+	// cache result. This is because the cache shouldn't return a result
+	// that is obviously past the validity of the event.
+
+	_, ok = serverA.cache.GetServerKey(
+		req,
+		gomatrixserverlib.AsTimestamp(time.Now().Add(time.Minute*90)),
+	)
+	if ok {
+		t.Fatalf("server B key is in cache when it shouldn't be (+90 minutes)")
+	}
+}
+
+func TestRenewalBehaviour(t *testing.T) {
+	// Server A will request Server C's key but their validity period
+	// is an hour in the past. We'll retrieve the key as, even though it's
+	// past its validity, it will be able to verify past events.
+
+	req := gomatrixserverlib.PublicKeyLookupRequest{
+		ServerName: serverC.name,
+		KeyID:      serverKeyID,
+	}
+
+	res, err := serverA.api.FetchKeys(
+		context.Background(),
+		map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp{
+			req: gomatrixserverlib.AsTimestamp(time.Now()),
+		},
+	)
+	if err != nil {
+		t.Fatalf("server A failed to retrieve server C key: %s", err)
+	}
+	if len(res) != 1 {
+		t.Fatalf("server C should have returned one key but instead returned %d keys", len(res))
+	}
+	if _, ok := res[req]; !ok {
+		t.Fatalf("server C isn't included in the key fetch response")
+	}
+
+	// If we ask the cache for the server key for an event that happened
+	// 90 minutes ago then we should get a cache result, as the key hadn't
+	// passed its validity by that point. The fact that the key is now in
+	// the cache is, in itself, proof that we successfully retrieved the
+	// key before.
+
+	oldcached, ok := serverA.cache.GetServerKey(
+		req,
+		gomatrixserverlib.AsTimestamp(time.Now().Add(-time.Minute*90)),
+	)
+	if !ok {
+		t.Fatalf("server C key isn't in cache when it should be (-90 minutes)")
+	}
+
+	// If we now ask the cache for the same key but this time for an event
+	// that only happened 30 minutes ago then we shouldn't get a cached
+	// result, as the event happened after the key validity expired. This
+	// is really just for sanity checking.
+
+	_, ok = serverA.cache.GetServerKey(
+		req,
+		gomatrixserverlib.AsTimestamp(time.Now().Add(-time.Minute*30)),
+	)
+	if ok {
+		t.Fatalf("server B key is in cache when it shouldn't be (-30 minutes)")
+	}
+
+	// We're now going to kick server C into renewing its key. Since we're
+	// happy at this point that the key that we already have is from the past
+	// then repeating a key fetch should cause us to try and renew the key.
+	// If so, then the new key will end up in our cache.
+
+	serverC.renew()
+
+	res, err = serverA.api.FetchKeys(
+		context.Background(),
+		map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp{
+			req: gomatrixserverlib.AsTimestamp(time.Now()),
+		},
+	)
+	if err != nil {
+		t.Fatalf("server A failed to retrieve server C key: %s", err)
+	}
+	if len(res) != 1 {
+		t.Fatalf("server C should have returned one key but instead returned %d keys", len(res))
+	}
+	if _, ok = res[req]; !ok {
+		t.Fatalf("server C isn't included in the key fetch response")
+	}
+
+	// We're now going to ask the cache what the new key validity is. If
+	// it is still the same as the previous validity then we've failed to
+	// retrieve the renewed key. If it's newer then we've successfully got
+	// the renewed key.
+
+	newcached, ok := serverA.cache.GetServerKey(
+		req,
+		gomatrixserverlib.AsTimestamp(time.Now().Add(-time.Minute*30)),
+	)
+	if !ok {
+		t.Fatalf("server B key isn't in cache when it shouldn't be (post-renewal)")
+	}
+	if oldcached.ValidUntilTS >= newcached.ValidUntilTS {
+		t.Fatalf("the server B key should have been renewed but wasn't")
+	}
+	t.Log(res)
+}
diff --git a/serverkeyapi/storage/cache/keydb.go b/serverkeyapi/storage/cache/keydb.go
index b662e4fd..2063dfc5 100644
--- a/serverkeyapi/storage/cache/keydb.go
+++ b/serverkeyapi/storage/cache/keydb.go
@@ -39,8 +39,8 @@ func (d *KeyDatabase) FetchKeys(
 	requests map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp,
 ) (map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult, error) {
 	results := make(map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult)
-	for req := range requests {
-		if res, cached := d.cache.GetServerKey(req); cached {
+	for req, ts := range requests {
+		if res, cached := d.cache.GetServerKey(req, ts); cached {
 			results[req] = res
 			delete(requests, req)
 		}
diff --git a/serverkeyapi/storage/postgres/keydb.go b/serverkeyapi/storage/postgres/keydb.go
index 32cdf951..aaa4409b 100644
--- a/serverkeyapi/storage/postgres/keydb.go
+++ b/serverkeyapi/storage/postgres/keydb.go
@@ -17,7 +17,6 @@ package postgres
 
 import (
 	"context"
-	"time"
 
 	"golang.org/x/crypto/ed25519"
 
@@ -51,28 +50,6 @@ func NewDatabase(
 	if err != nil {
 		return nil, err
 	}
-	// Store our own keys so that we don't end up making HTTP requests to find our
-	// own keys
-	index := gomatrixserverlib.PublicKeyLookupRequest{
-		ServerName: serverName,
-		KeyID:      serverKeyID,
-	}
-	value := gomatrixserverlib.PublicKeyLookupResult{
-		VerifyKey: gomatrixserverlib.VerifyKey{
-			Key: gomatrixserverlib.Base64Bytes(serverKey),
-		},
-		ValidUntilTS: gomatrixserverlib.AsTimestamp(time.Now().Add(100 * 365 * 24 * time.Hour)),
-		ExpiredTS:    gomatrixserverlib.PublicKeyNotExpired,
-	}
-	err = d.StoreKeys(
-		context.Background(),
-		map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult{
-			index: value,
-		},
-	)
-	if err != nil {
-		return nil, err
-	}
 	return d, nil
 }
 
diff --git a/serverkeyapi/storage/sqlite3/keydb.go b/serverkeyapi/storage/sqlite3/keydb.go
index 268c7542..dc72b79e 100644
--- a/serverkeyapi/storage/sqlite3/keydb.go
+++ b/serverkeyapi/storage/sqlite3/keydb.go
@@ -17,7 +17,6 @@ package sqlite3
 
 import (
 	"context"
-	"time"
 
 	"golang.org/x/crypto/ed25519"
 
@@ -56,25 +55,6 @@ func NewDatabase(
 	if err != nil {
 		return nil, err
 	}
-	// Store our own keys so that we don't end up making HTTP requests to find our
-	// own keys
-	index := gomatrixserverlib.PublicKeyLookupRequest{
-		ServerName: serverName,
-		KeyID:      serverKeyID,
-	}
-	value := gomatrixserverlib.PublicKeyLookupResult{
-		VerifyKey: gomatrixserverlib.VerifyKey{
-			Key: gomatrixserverlib.Base64Bytes(serverKey),
-		},
-		ValidUntilTS: gomatrixserverlib.AsTimestamp(time.Now().Add(100 * 365 * 24 * time.Hour)),
-		ExpiredTS:    gomatrixserverlib.PublicKeyNotExpired,
-	}
-	err = d.StoreKeys(
-		context.Background(),
-		map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult{
-			index: value,
-		},
-	)
 	if err != nil {
 		return nil, err
 	}