Break up FetchKeys

serverkeyapi/internal/api.go

@@ -70,7 +70,6 @@ func (s *ServerKeyAPI) StoreKeys(
 	return s.OurKeyRing.KeyDatabase.StoreKeys(ctx, results)
 }
 
-// nolint:gocyclo
 func (s *ServerKeyAPI) FetchKeys(
 	_ context.Context,
 	requests map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp,
@@ -78,15 +77,74 @@ 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()
+	now := gomatrixserverlib.AsTimestamp(time.Now())
 	results := map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult{}
 	origRequests := map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp{}
 	for k, v := range requests {
 		origRequests[k] = v
 	}
-	now := gomatrixserverlib.AsTimestamp(time.Now())
 
 	// First, check if any of these key checks are for our own keys. If
 	// they are then we will satisfy them directly.
+	if err := s.handleLocalKeys(ctx, requests, results); err != nil {
+		return nil, err
+	}
+
+	// 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 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
+	for _, fetcher := range s.OurKeyRing.KeyFetchers {
+		// If there are no more keys to look up then stop.
+		if len(requests) == 0 {
+			break
+		}
+
+		// 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
+		}
+	}
+
+	// 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
+}
+
+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(
+	ctx context.Context,
+	requests map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp,
+	results map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult,
+) error {
 	for req := range requests {
 		if req.ServerName == s.Cfg.Matrix.ServerName {
 			// We found a key request that is supposed to be for our own
@@ -98,7 +156,7 @@ func (s *ServerKeyAPI) FetchKeys(
 			request := &api.QueryLocalKeysRequest{}
 			response := &api.QueryLocalKeysResponse{}
 			if err := s.QueryLocalKeys(ctx, request, response); err != nil {
-				return nil, err
+				return err
 			}
 
 			// Depending on whether the key is expired or not, we'll need
@@ -121,10 +179,23 @@ func (s *ServerKeyAPI) FetchKeys(
 		}
 	}
 
-	// Then consult our local database and see if we have the requested
+	return nil
-	// keys. These might come from a cache, depending on the database
+}
-	// implementation used.
+
-	if dbResults, err := s.OurKeyRing.KeyDatabase.FetchKeys(ctx, requests); err == nil {
+// 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 {
 		results[req] = res
@@ -136,23 +207,32 @@ func (s *ServerKeyAPI) FetchKeys(
 			delete(requests, req)
 		}
 	}
+	return nil
 }
 
-	// For any key requests that we still have outstanding, next try to
+// handleDatabaseKeys handles cases where a fetcher can satisfy
-	// fetch them directly. We'll go through each of the key fetchers to
+// the remaining requests.
-	// ask for the remaining keys
+func (s *ServerKeyAPI) handleFetcherKeys(
-	for _, fetcher := range s.OurKeyRing.KeyFetchers {
+	ctx context.Context,
-		// If there are no more keys to look up then stop.
+	now gomatrixserverlib.Timestamp,
-		if len(requests) == 0 {
+	fetcher gomatrixserverlib.KeyFetcher,
-			break
+	requests map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp,
-		}
+	results map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult,
-		fetcherCtx, fetcherCancel := context.WithTimeout(ctx, time.Second*30)
+) error {
-		defer fetcherCancel()
 	logrus.WithFields(logrus.Fields{
 		"fetcher_name": fetcher.FetcherName(),
 	}).Infof("Fetching %d key(s)", len(requests))
 
-		if fetcherResults, err := fetcher.FetchKeys(fetcherCtx, requests); err == nil {
+	// 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.
@@ -196,7 +276,7 @@ func (s *ServerKeyAPI) FetchKeys(
 			"fetcher_name":  fetcher.FetcherName(),
 			"database_name": s.OurKeyRing.KeyDatabase.FetcherName(),
 		}).Errorf("Failed to store keys in the database")
-				return nil, fmt.Errorf("server key API failed to store retrieved keys: %w", err)
+		return fmt.Errorf("server key API failed to store retrieved keys: %w", err)
 	}
 
 	if len(storeResults) > 0 {
@@ -204,32 +284,6 @@ func (s *ServerKeyAPI) FetchKeys(
 			"fetcher_name": fetcher.FetcherName(),
 		}).Infof("Updated %d of %d key(s) in database", len(storeResults), len(results))
 	}
-		} else {
-			logrus.WithError(err).WithFields(logrus.Fields{
-				"fetcher_name": fetcher.FetcherName(),
-			}).Errorf("Failed to retrieve %d key(s)", len(requests))
-		}
-	}
 
-	// Check that we've actually satisfied all of the key requests that we
+	return nil
-	// 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
-}
-
-func (s *ServerKeyAPI) FetcherName() string {
-	return fmt.Sprintf("ServerKeyAPI (wrapping %q)", s.OurKeyRing.KeyDatabase.FetcherName())
 }