Add fetcher behaviour into server key API implementation

serverkeyapi/api/api.go

@@ -53,6 +53,11 @@ type httpServerKeyInternalAPI struct {
 }
 
 func (s *httpServerKeyInternalAPI) KeyRing() *gomatrixserverlib.KeyRing {
+	// This is a bit of a cheat - we tell gomatrixserverlib that this API is
+	// both the key database and the key fetcher. While this does have the
+	// rather unfortunate effect of preventing gomatrixserverlib from handling
+	// key fetchers directly, we can at least reimplement this behaviour on
+	// the other end of the API.
 	return &gomatrixserverlib.KeyRing{
 		KeyDatabase: s,
 		KeyFetchers: []gomatrixserverlib.KeyFetcher{s},

serverkeyapi/internal/api.go

@@ -6,14 +6,12 @@ import (
 	"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/dendrite/serverkeyapi/storage"
 	"github.com/matrix-org/gomatrixserverlib"
 )
 
 type ServerKeyAPI struct {
 	api.ServerKeyInternalAPI
 
-	DB             storage.Database
 	Cfg            *config.Dendrite
 	ImmutableCache caching.ImmutableCache
 	OurKeyRing     gomatrixserverlib.KeyRing
@@ -21,6 +19,8 @@ type ServerKeyAPI struct {
 }
 
 func (s *ServerKeyAPI) KeyRing() *gomatrixserverlib.KeyRing {
+	// Return a real keyring - one that has the real database and real
+	// fetchers.
 	return &s.OurKeyRing
 }
 
@@ -28,16 +28,46 @@ func (s *ServerKeyAPI) StoreKeys(
 	ctx context.Context,
 	results map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult,
 ) error {
-	return s.DB.StoreKeys(ctx, results)
+	// Store any keys that we were given in our database.
+	return s.OurKeyRing.KeyDatabase.StoreKeys(ctx, results)
 }
 
 func (s *ServerKeyAPI) FetchKeys(
 	ctx context.Context,
 	requests map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp,
 ) (map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult, error) {
-	return s.DB.FetchKeys(ctx, requests)
+	results := map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult{}
+	// First 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 {
+			results[req] = res
+			delete(requests, req)
+		}
+	}
+	// 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.
+	for _, fetcher := range s.OurKeyRing.KeyFetchers {
+		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 {
+				results[req] = res
+				delete(requests, req)
+			}
+		}
+	}
+	// Return the keys.
+	return results, nil
 }
 
 func (s *ServerKeyAPI) FetcherName() string {
-	return s.DB.FetcherName()
+	return s.OurKeyRing.KeyDatabase.FetcherName()
 }

serverkeyapi/internal/http.go

@@ -28,7 +28,7 @@ func (s *ServerKeyAPI) SetupHTTP(internalAPIMux *mux.Router) {
 				}
 				lookup[req] = timestamp
 			}
-			keys, err := s.DB.FetchKeys(req.Context(), lookup)
+			keys, err := s.FetchKeys(req.Context(), lookup)
 			if err != nil {
 				return util.ErrorResponse(err)
 			}
@@ -51,7 +51,7 @@ func (s *ServerKeyAPI) SetupHTTP(internalAPIMux *mux.Router) {
 				store[req] = res
 				s.ImmutableCache.StoreServerKey(req, res)
 			}
-			if err := s.DB.StoreKeys(req.Context(), store); err != nil {
+			if err := s.StoreKeys(req.Context(), store); err != nil {
 				return util.ErrorResponse(err)
 			}
 			return util.JSONResponse{Code: http.StatusOK, JSON: &response}

serverkeyapi/serverkeyapi.go

@@ -34,7 +34,6 @@ func SetupServerKeyAPIComponent(
 	}
 
 	internalAPI := internal.ServerKeyAPI{
-		DB:             serverKeyDB,
 		Cfg:            base.Cfg,
 		ImmutableCache: base.ImmutableCache,
 		FedClient:      fedClient,