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

cmd/dendrite-federation-api-server/main.go

@@ -27,15 +27,14 @@ func main() {
 	accountDB := base.CreateAccountsDB()
 	deviceDB := base.CreateDeviceDB()
 	federation := base.CreateFederationClient()
-	serverKeyAPI := base.ServerKeyAPIClient()
-	keyRing := serverKeyAPI.KeyRing()
 	fsAPI := base.FederationSenderHTTPClient()
 	rsAPI := base.RoomserverHTTPClient()
 	asAPI := base.AppserviceHTTPClient()
+	skAPI := base.ServerKeyAPIClient()
 
 	federationapi.AddPublicRoutes(
-		base.PublicAPIMux, base.Cfg, accountDB, deviceDB, federation, keyRing,
-		rsAPI, asAPI, fsAPI, base.EDUServerClient(),
+		base.PublicAPIMux, base.Cfg, accountDB, deviceDB, federation,
+		skAPI, rsAPI, asAPI, fsAPI, base.EDUServerClient(),
 	)
 
 	base.SetupAndServeHTTP(string(base.Cfg.Bind.FederationAPI), string(base.Cfg.Listen.FederationAPI))

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)
 	}

federationapi/federationapi.go

@@ -23,6 +23,7 @@ import (
 	federationSenderAPI "github.com/matrix-org/dendrite/federationsender/api"
 	"github.com/matrix-org/dendrite/internal/config"
 	roomserverAPI "github.com/matrix-org/dendrite/roomserver/api"
+	serverkeyAPI "github.com/matrix-org/dendrite/serverkeyapi/api"
 
 	"github.com/matrix-org/dendrite/federationapi/routing"
 	"github.com/matrix-org/gomatrixserverlib"
@@ -35,16 +36,15 @@ func AddPublicRoutes(
 	accountsDB accounts.Database,
 	deviceDB devices.Database,
 	federation *gomatrixserverlib.FederationClient,
-	keyRing *gomatrixserverlib.KeyRing,
+	skAPI serverkeyAPI.ServerKeyInternalAPI,
 	rsAPI roomserverAPI.RoomserverInternalAPI,
 	asAPI appserviceAPI.AppServiceQueryAPI,
 	federationSenderAPI federationSenderAPI.FederationSenderInternalAPI,
 	eduAPI eduserverAPI.EDUServerInputAPI,
 ) {
-
 	routing.Setup(
 		router, cfg, rsAPI, asAPI,
-		eduAPI, federationSenderAPI, *keyRing,
+		eduAPI, federationSenderAPI, skAPI,
 		federation, accountsDB, deviceDB,
 	)
 }

federationapi/routing/keys.go

@@ -1,68 +0,0 @@
-// Copyright 2017 Vector Creations Ltd
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package routing
-
-import (
-	"encoding/json"
-	"net/http"
-	"time"
-
-	"github.com/matrix-org/dendrite/internal/config"
-	"github.com/matrix-org/gomatrixserverlib"
-	"github.com/matrix-org/util"
-	"golang.org/x/crypto/ed25519"
-)
-
-// LocalKeys returns the local keys for the server.
-// See https://matrix.org/docs/spec/server_server/unstable.html#publishing-keys
-func LocalKeys(cfg *config.Dendrite) util.JSONResponse {
-	keys, err := localKeys(cfg, 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.Dendrite, validUntil time.Time) (*gomatrixserverlib.ServerKeys, error) {
-	var keys gomatrixserverlib.ServerKeys
-
-	keys.ServerName = cfg.Matrix.ServerName
-
-	publicKey := cfg.Matrix.PrivateKey.Public().(ed25519.PublicKey)
-
-	keys.VerifyKeys = map[gomatrixserverlib.KeyID]gomatrixserverlib.VerifyKey{
-		cfg.Matrix.KeyID: {
-			Key: gomatrixserverlib.Base64Bytes(publicKey),
-		},
-	}
-
-	keys.TLSFingerprints = cfg.Matrix.TLSFingerPrints
-	keys.OldVerifyKeys = map[gomatrixserverlib.KeyID]gomatrixserverlib.OldVerifyKey{}
-	keys.ValidUntilTS = gomatrixserverlib.AsTimestamp(validUntil)
-
-	toSign, err := json.Marshal(keys.ServerKeyFields)
-	if err != nil {
-		return nil, err
-	}
-
-	keys.Raw, err = gomatrixserverlib.SignJSON(
-		string(cfg.Matrix.ServerName), cfg.Matrix.KeyID, cfg.Matrix.PrivateKey, toSign,
-	)
-	if err != nil {
-		return nil, err
-	}
-
-	return &keys, nil
-}

federationapi/routing/routing.go

@@ -26,6 +26,7 @@ import (
 	"github.com/matrix-org/dendrite/internal"
 	"github.com/matrix-org/dendrite/internal/config"
 	roomserverAPI "github.com/matrix-org/dendrite/roomserver/api"
+	serverkeyAPI "github.com/matrix-org/dendrite/serverkeyapi/api"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/util"
 )
@@ -51,11 +52,12 @@ func Setup(
 	asAPI appserviceAPI.AppServiceQueryAPI,
 	eduAPI eduserverAPI.EDUServerInputAPI,
 	fsAPI federationSenderAPI.FederationSenderInternalAPI,
-	keys gomatrixserverlib.KeyRing,
+	skAPI serverkeyAPI.ServerKeyInternalAPI,
 	federation *gomatrixserverlib.FederationClient,
 	accountDB accounts.Database,
 	deviceDB devices.Database,
 ) {
+	keys := *skAPI.KeyRing()
 	v2keysmux := publicAPIMux.PathPrefix(pathPrefixV2Keys).Subrouter()
 	v1fedmux := publicAPIMux.PathPrefix(pathPrefixV1Federation).Subrouter()
 	v2fedmux := publicAPIMux.PathPrefix(pathPrefixV2Federation).Subrouter()
@@ -65,7 +67,15 @@ func Setup(
 	}
 
 	localKeys := internal.MakeExternalAPI("localkeys", func(req *http.Request) util.JSONResponse {
-		return LocalKeys(cfg)
+		request := &serverkeyAPI.QueryLocalKeysRequest{}
+		response := &serverkeyAPI.QueryLocalKeysResponse{}
+		if err := skAPI.QueryLocalKeys(req.Context(), request, response); err != nil {
+			return util.ErrorResponse(err)
+		}
+		return util.JSONResponse{
+			Code: http.StatusOK,
+			JSON: response.ServerKeys,
+		}
 	})
 
 	// Ignore the {keyID} argument as we only have a single server key so we always

internal/basecomponent/base.go

@@ -95,7 +95,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")
 	}

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
 }
 

internal/setup/monolith.go

@@ -61,7 +61,7 @@ func (m *Monolith) AddAllPublicRoutes(publicMux *mux.Router) {
 	keyserver.AddPublicRoutes(publicMux, m.Config, m.DeviceDB, m.AccountDB)
 	federationapi.AddPublicRoutes(
 		publicMux, m.Config, m.AccountDB, m.DeviceDB, m.FedClient,
-		m.KeyRing, m.RoomserverAPI, m.AppserviceAPI, m.FederationSenderAPI,
+		m.ServerKeyAPI, m.RoomserverAPI, m.AppserviceAPI, m.FederationSenderAPI,
 		m.EDUInternalAPI,
 	)
 	mediaapi.AddPublicRoutes(publicMux, m.Config, m.DeviceDB)

serverkeyapi/api/api.go

@@ -22,6 +22,12 @@ type ServerKeyInternalAPI interface {
 		request *QueryPublicKeysRequest,
 		response *QueryPublicKeysResponse,
 	) error
+
+	QueryLocalKeys(
+		ctx context.Context,
+		request *QueryLocalKeysRequest,
+		response *QueryLocalKeysResponse,
+	) error
 }
 
 type QueryPublicKeysRequest struct {
@@ -38,3 +44,10 @@ type InputPublicKeysRequest struct {
 
 type InputPublicKeysResponse struct {
 }
+
+type QueryLocalKeysRequest struct {
+}
+
+type QueryLocalKeysResponse struct {
+	ServerKeys gomatrixserverlib.ServerKeys `json:"server_keys"`
+}

serverkeyapi/internal/api.go

@@ -2,9 +2,12 @@ package internal
 
 import (
 	"context"
+	"crypto/ed25519"
+	"encoding/json"
 	"fmt"
 	"time"
 
+	"github.com/matrix-org/dendrite/internal/config"
 	"github.com/matrix-org/dendrite/serverkeyapi/api"
 	"github.com/matrix-org/gomatrixserverlib"
 )
@@ -12,10 +15,35 @@ import (
 type ServerKeyAPI struct {
 	api.ServerKeyInternalAPI
 
+	Cfg        *config.Dendrite
 	OurKeyRing gomatrixserverlib.KeyRing
 	FedClient  *gomatrixserverlib.FederationClient
 }
 
+func (s *ServerKeyAPI) QueryLocalKeys(ctx context.Context, request *api.QueryLocalKeysRequest, response *api.QueryLocalKeysResponse) error {
+	response.ServerKeys.ServerName = s.Cfg.Matrix.ServerName
+
+	publicKey := s.Cfg.Matrix.PrivateKey.Public().(ed25519.PublicKey)
+	response.ServerKeys.VerifyKeys = map[gomatrixserverlib.KeyID]gomatrixserverlib.VerifyKey{
+		s.Cfg.Matrix.KeyID: {
+			Key: gomatrixserverlib.Base64Bytes(publicKey),
+		},
+	}
+	response.ServerKeys.TLSFingerprints = s.Cfg.Matrix.TLSFingerPrints
+	response.ServerKeys.OldVerifyKeys = map[gomatrixserverlib.KeyID]gomatrixserverlib.OldVerifyKey{}
+	response.ServerKeys.ValidUntilTS = gomatrixserverlib.AsTimestamp(time.Now().Add(s.Cfg.Matrix.KeyValidityPeriod))
+
+	toSign, err := json.Marshal(response.ServerKeys.ServerKeyFields)
+	if err != nil {
+		return err
+	}
+
+	response.ServerKeys.Raw, err = gomatrixserverlib.SignJSON(
+		string(s.Cfg.Matrix.ServerName), s.Cfg.Matrix.KeyID, s.Cfg.Matrix.PrivateKey, toSign,
+	)
+	return err
+}
+
 func (s *ServerKeyAPI) KeyRing() *gomatrixserverlib.KeyRing {
 	// Return a keyring that forces requests to be proxied through the
 	// below functions. That way we can enforce things like validity
@@ -50,16 +78,15 @@ func (s *ServerKeyAPI) FetchKeys(
 	// 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.
+		// We successfully got some keys. Add them to the results.
 		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 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 res.WasValidAt(now, true) {
+				delete(requests, req)
+			}
 		}
 	}
 	// For any key requests that we still have outstanding, next try to
@@ -70,18 +97,37 @@ func (s *ServerKeyAPI) FetchKeys(
 			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.
+			// 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 !res.WasValidAt(now, true) {
-					// We appear to be past the key validity. Don't return this
-					// key with the results.
-					continue
+				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.
+						storeResults[req] = res
+					}
+				} else {
+					// We didn't already have a previous entry for this request
+					// so store it in the database anyway for now.
+					storeResults[req] = res
 				}
+				// Update the results map with this new result. If nothing
+				// else, we can try verifying against this key.
 				results[req] = res
-				delete(requests, req)
+				// 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)
+				}
 			}
-			if err = s.OurKeyRing.KeyDatabase.StoreKeys(ctx, fetcherResults); err != nil {
+			// Store the keys from our store map.
+			if err = s.OurKeyRing.KeyDatabase.StoreKeys(ctx, storeResults); err != nil {
 				return nil, fmt.Errorf("server key API failed to store retrieved keys: %w", err)
 			}
 		}

serverkeyapi/inthttp/client.go

@@ -16,6 +16,7 @@ import (
 const (
 	ServerKeyInputPublicKeyPath = "/serverkeyapi/inputPublicKey"
 	ServerKeyQueryPublicKeyPath = "/serverkeyapi/queryPublicKey"
+	ServerKeyQueryLocalKeysPath = "/serverkeyapi/queryLocalKeys"
 )
 
 // NewServerKeyClient creates a ServerKeyInternalAPI implemented by talking to a HTTP POST API.
@@ -130,3 +131,15 @@ func (h *httpServerKeyInternalAPI) QueryPublicKeys(
 	apiURL := h.serverKeyAPIURL + ServerKeyQueryPublicKeyPath
 	return internalHTTP.PostJSON(ctx, span, h.httpClient, apiURL, request, response)
 }
+
+func (h *httpServerKeyInternalAPI) QueryLocalKeys(
+	ctx context.Context,
+	request *api.QueryLocalKeysRequest,
+	response *api.QueryLocalKeysResponse,
+) error {
+	span, ctx := opentracing.StartSpanFromContext(ctx, "QueryLocalKeys")
+	defer span.Finish()
+
+	apiURL := h.serverKeyAPIURL + ServerKeyQueryLocalKeysPath
+	return internalHTTP.PostJSON(ctx, span, h.httpClient, apiURL, request, response)
+}

serverkeyapi/inthttp/server.go

@@ -12,6 +12,19 @@ import (
 )
 
 func AddRoutes(s api.ServerKeyInternalAPI, internalAPIMux *mux.Router, cache caching.ServerKeyCache) {
+	internalAPIMux.Handle(ServerKeyQueryLocalKeysPath,
+		internal.MakeInternalAPI("queryLocalKeys", func(req *http.Request) util.JSONResponse {
+			request := api.QueryLocalKeysRequest{}
+			response := api.QueryLocalKeysResponse{}
+			if err := json.NewDecoder(req.Body).Decode(&request); err != nil {
+				return util.MessageResponse(http.StatusBadRequest, err.Error())
+			}
+			if err := s.QueryLocalKeys(req.Context(), &request, &response); err != nil {
+				return util.ErrorResponse(err)
+			}
+			return util.JSONResponse{Code: http.StatusOK, JSON: &response}
+		}),
+	)
 	internalAPIMux.Handle(ServerKeyQueryPublicKeyPath,
 		internal.MakeInternalAPI("queryPublicKeys", func(req *http.Request) util.JSONResponse {
 			request := api.QueryPublicKeysRequest{}

serverkeyapi/serverkeyapi.go

@@ -46,6 +46,7 @@ func NewInternalAPI(
 	}
 
 	internalAPI := internal.ServerKeyAPI{
+		Cfg:       cfg,
 		FedClient: fedClient,
 		OurKeyRing: gomatrixserverlib.KeyRing{
 			KeyFetchers: []gomatrixserverlib.KeyFetcher{

serverkeyapi/serverkeyapi_test.go

@@ -8,10 +8,11 @@ import (
 	"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"
@@ -20,15 +21,19 @@ import (
 
 type server struct {
 	name      gomatrixserverlib.ServerName
+	validity  time.Duration
 	config    *config.Dendrite
 	fedclient *gomatrixserverlib.FederationClient
 	cache     *caching.Caches
 	api       api.ServerKeyInternalAPI
 }
 
-var serverA = &server{name: "a.com"}
-var serverB = &server{name: "b.com"}
-var serverC = &server{name: "c.com"}
+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,
@@ -37,49 +42,59 @@ var servers = map[string]*server{
 }
 
 func TestMain(m *testing.M) {
-	caches, err := caching.NewInMemoryLRUCache()
-	if err != nil {
-		panic("can't create cache: " + err.Error())
-	}
-
-	for _, s := range []*server{serverA, serverB, serverC} {
+	for _, s := range servers {
 		_, testPriv, err := ed25519.GenerateKey(nil)
 		if err != nil {
 			panic("can't generate identity key: " + err.Error())
 		}
 
+		s.cache, err = caching.NewInMemoryLRUCache(false)
+		if err != nil {
+			panic("can't create cache: " + err.Error())
+		}
+
 		s.config = &config.Dendrite{}
 		s.config.SetDefaults()
+		s.config.Matrix.KeyValidityPeriod = s.validity
 		s.config.Matrix.ServerName = gomatrixserverlib.ServerName(s.name)
 		s.config.Matrix.PrivateKey = testPriv
-		s.config.Matrix.KeyID = "ed25519:test"
+		s.config.Matrix.KeyID = serverKeyID
 		s.config.Database.ServerKey = config.DataSource("file::memory:")
 
 		transport := &http.Transport{}
 		transport.RegisterProtocol("matrix", &MockRoundTripper{})
 
 		s.fedclient = gomatrixserverlib.NewFederationClientWithTransport(
-			s.config.Matrix.ServerName, "ed25519:test", testPriv, transport,
+			s.config.Matrix.ServerName, serverKeyID, testPriv, transport,
 		)
 
-		s.cache = caches
 		s.api = NewInternalAPI(s.config, s.fedclient, s.cache)
 	}
+
+	os.Exit(m.Run())
 }
 
 type MockRoundTripper struct{}
 
 func (m *MockRoundTripper) RoundTrip(req *http.Request) (res *http.Response, err error) {
-	serv, ok := servers[req.Host]
+	s, ok := servers[req.Host]
 	if !ok {
 		return nil, fmt.Errorf("server not known: %s", req.Host)
 	}
 
-	keys := routing.LocalKeys(serv.config).JSON
-	body, err := json.MarshalIndent(keys, "", "  ")
+	request := &api.QueryLocalKeysRequest{}
+	response := &api.QueryLocalKeysResponse{}
+	if err = s.api.QueryLocalKeys(context.Background(), request, response); err != nil {
+		return nil, err
+	}
+
+	body, err := json.MarshalIndent(response.ServerKeys, "", "  ")
 	if err != nil {
 		return nil, err
 	}
+
+	fmt.Println("Round-tripper says:", string(body))
+
 	res = &http.Response{
 		StatusCode: 200,
 		Body:       ioutil.NopCloser(bytes.NewReader(body)),
@@ -87,18 +102,136 @@ func (m *MockRoundTripper) RoundTrip(req *http.Request) (res *http.Response, err
 	return
 }
 
-func TestServerKeyAPIDirect(t *testing.T) {
+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")
+		}
+		fmt.Printf("%s's key expires at %d\n", name, res[req].ValidUntilTS)
+	}
+}
+
+func TestServerARequestsServerBKey(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,
+	}
+
 	res, err := serverA.api.FetchKeys(
 		context.Background(),
 		map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp{
-			{
-				ServerName: serverA.name,
-				KeyID:      "ed25519:test",
-			}: gomatrixserverlib.AsTimestamp(time.Now()),
+			req: gomatrixserverlib.AsTimestamp(time.Now()),
+		},
+	)
+	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.
+	*/
+
+	cres, ok := serverA.cache.GetServerKey(req)
+	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")
+	}
+
+	/*
+		Server A will then request Server B's key for an event that
+		happened two hours ago, which *should* pass since the key was
+		valid then too and it's already in the cache.
+	*/
+
+	_, err = serverA.api.FetchKeys(
+		context.Background(),
+		map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp{
+			req: gomatrixserverlib.AsTimestamp(time.Now().Add(-time.Hour * 2)),
+		},
+	)
+	if err != nil {
+		t.Fatalf("server A failed to retrieve server B key: %s", err)
+	}
+}
+
+func TestServerARequestsServerCKey(t *testing.T) {
+	/*
+		Server A will request Server C's key for an event that came
+		in just now, but their validity period is an hour in the
+		past. This *should* fail since the key isn't valid now.
+	*/
+
+	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")
+	}
+
+	t.Log("server C's key expires at", res[req].ValidUntilTS.Time())
+
+	/*
+		Server A will then request Server C's key for an event that
+		happened two hours ago, which *should* pass since the key was
+		valid then.
+	*/
+
+	_, err = serverA.api.FetchKeys(
+		context.Background(),
+		map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp{
+			req: gomatrixserverlib.AsTimestamp(time.Now().Add(-time.Hour)),
 		},
 	)
 	if err != nil {
 		t.Fatalf("serverKeyAPI.FetchKeys: %s", err)
 	}
-	t.Log(res)
 }

serverkeyapi/storage/postgres/keydb.go

@@ -17,7 +17,6 @@ package postgres
 
 import (
 	"context"
-	"time"
 
 	"golang.org/x/crypto/ed25519"
 
@@ -54,26 +53,28 @@ func NewDatabase(
 	}
 	// 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
-	}
+	/*
+		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
 }
 

serverkeyapi/storage/sqlite3/keydb.go

@@ -17,7 +17,6 @@ package sqlite3
 
 import (
 	"context"
-	"time"
 
 	"golang.org/x/crypto/ed25519"
 
@@ -59,23 +58,25 @@ func NewDatabase(
 	}
 	// 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,
-		},
-	)
+	/*
+		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
 	}