fedsender: try to satisfy all notary key requests from the cache first (#1925)

* fedsender: try to satisfy all notary key requests from the cache first

* Linting

4 files changed, 79 insertions, 15 deletions

diff --git a/federationsender/api/api.go b/federationsender/api/api.go
index 46e0676d..82cdf9d8 100644
--- a/federationsender/api/api.go
+++ b/federationsender/api/api.go
@@ -96,8 +96,18 @@ type FederationSenderInternalAPI interface {
 }
 
 type QueryServerKeysRequest struct {
-	ServerName     gomatrixserverlib.ServerName
-	OptionalKeyIDs []gomatrixserverlib.KeyID
+	ServerName      gomatrixserverlib.ServerName
+	KeyIDToCriteria map[gomatrixserverlib.KeyID]gomatrixserverlib.PublicKeyNotaryQueryCriteria
+}
+
+func (q *QueryServerKeysRequest) KeyIDs() []gomatrixserverlib.KeyID {
+	kids := make([]gomatrixserverlib.KeyID, len(q.KeyIDToCriteria))
+	i := 0
+	for keyID := range q.KeyIDToCriteria {
+		kids[i] = keyID
+		i++
+	}
+	return kids
 }
 
 type QueryServerKeysResponse struct {
diff --git a/federationsender/internal/query.go b/federationsender/internal/query.go
index 7fd384a8..af531f7d 100644
--- a/federationsender/internal/query.go
+++ b/federationsender/internal/query.go
@@ -25,29 +25,73 @@ func (f *FederationSenderInternalAPI) QueryJoinedHostServerNamesInRoom(
 	return
 }
 
-func (a *FederationSenderInternalAPI) QueryServerKeys(
-	ctx context.Context, req *api.QueryServerKeysRequest, res *api.QueryServerKeysResponse,
-) error {
+func (a *FederationSenderInternalAPI) fetchServerKeysDirectly(ctx context.Context, serverName gomatrixserverlib.ServerName) (*gomatrixserverlib.ServerKeys, error) {
 	ctx, cancel := context.WithTimeout(ctx, time.Second*30)
 	defer cancel()
-	ires, err := a.doRequest(req.ServerName, func() (interface{}, error) {
-		return a.federation.GetServerKeys(ctx, req.ServerName)
+	ires, err := a.doRequest(serverName, func() (interface{}, error) {
+		return a.federation.GetServerKeys(ctx, serverName)
 	})
 	if err != nil {
-		// try to load from the cache
-		serverKeysResponses, dbErr := a.db.GetNotaryKeys(ctx, req.ServerName, req.OptionalKeyIDs)
+		return nil, err
+	}
+	sks := ires.(gomatrixserverlib.ServerKeys)
+	return &sks, nil
+}
+
+func (a *FederationSenderInternalAPI) fetchServerKeysFromCache(
+	ctx context.Context, req *api.QueryServerKeysRequest,
+) ([]gomatrixserverlib.ServerKeys, error) {
+	var results []gomatrixserverlib.ServerKeys
+	for keyID, criteria := range req.KeyIDToCriteria {
+		serverKeysResponses, _ := a.db.GetNotaryKeys(ctx, req.ServerName, []gomatrixserverlib.KeyID{keyID})
+		if len(serverKeysResponses) == 0 {
+			return nil, fmt.Errorf("failed to find server key response for key ID %s", keyID)
+		}
+		// we should only get 1 result as we only gave 1 key ID
+		sk := serverKeysResponses[0]
+		util.GetLogger(ctx).Infof("fetchServerKeysFromCache: minvalid:%v  keys: %+v", criteria.MinimumValidUntilTS, sk)
+		if criteria.MinimumValidUntilTS != 0 {
+			// check if it's still valid. if they have the same value that's also valid
+			if sk.ValidUntilTS < criteria.MinimumValidUntilTS {
+				return nil, fmt.Errorf(
+					"found server response for key ID %s but it is no longer valid, min: %v valid_until: %v",
+					keyID, criteria.MinimumValidUntilTS, sk.ValidUntilTS,
+				)
+			}
+		}
+		results = append(results, sk)
+	}
+	return results, nil
+}
+
+func (a *FederationSenderInternalAPI) QueryServerKeys(
+	ctx context.Context, req *api.QueryServerKeysRequest, res *api.QueryServerKeysResponse,
+) error {
+	// attempt to satisfy the entire request from the cache first
+	results, err := a.fetchServerKeysFromCache(ctx, req)
+	if err == nil {
+		// satisfied entirely from cache, return it
+		res.ServerKeys = results
+		return nil
+	}
+	util.GetLogger(ctx).WithField("server", req.ServerName).WithError(err).Warn("notary: failed to satisfy keys request entirely from cache, hitting direct")
+
+	serverKeys, err := a.fetchServerKeysDirectly(ctx, req.ServerName)
+	if err != nil {
+		// try to load as much as we can from the cache in a best effort basis
+		util.GetLogger(ctx).WithField("server", req.ServerName).WithError(err).Warn("notary: failed to ask server for keys, returning best effort keys")
+		serverKeysResponses, dbErr := a.db.GetNotaryKeys(ctx, req.ServerName, req.KeyIDs())
 		if dbErr != nil {
-			return fmt.Errorf("server returned %s, and db returned %s", err, dbErr)
+			return fmt.Errorf("notary: server returned %s, and db returned %s", err, dbErr)
 		}
 		res.ServerKeys = serverKeysResponses
 		return nil
 	}
-	serverKeys := ires.(gomatrixserverlib.ServerKeys)
 	// cache it!
-	if err = a.db.UpdateNotaryKeys(context.Background(), req.ServerName, serverKeys); err != nil {
+	if err = a.db.UpdateNotaryKeys(context.Background(), req.ServerName, *serverKeys); err != nil {
 		// non-fatal, still return the response
 		util.GetLogger(ctx).WithError(err).Warn("failed to UpdateNotaryKeys")
 	}
-	res.ServerKeys = []gomatrixserverlib.ServerKeys{serverKeys}
+	res.ServerKeys = []gomatrixserverlib.ServerKeys{*serverKeys}
 	return nil
 }
diff --git a/federationsender/storage/postgres/notary_server_keys_metadata_table.go b/federationsender/storage/postgres/notary_server_keys_metadata_table.go
index 3c3e6d14..b460dcd8 100644
--- a/federationsender/storage/postgres/notary_server_keys_metadata_table.go
+++ b/federationsender/storage/postgres/notary_server_keys_metadata_table.go
@@ -17,6 +17,7 @@ package postgres
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
 
 	"github.com/lib/pq"
 	"github.com/matrix-org/dendrite/federationsender/storage/tables"
@@ -148,7 +149,11 @@ func (s *notaryServerKeysMetadataStatements) SelectKeys(ctx context.Context, txn
 	var results []gomatrixserverlib.ServerKeys
 	for rows.Next() {
 		var sk gomatrixserverlib.ServerKeys
-		if err := rows.Scan(&sk.Raw); err != nil {
+		var raw string
+		if err = rows.Scan(&raw); err != nil {
+			return nil, err
+		}
+		if err = json.Unmarshal([]byte(raw), &sk); err != nil {
 			return nil, err
 		}
 		results = append(results, sk)
diff --git a/federationsender/storage/sqlite3/notary_server_keys_metadata_table.go b/federationsender/storage/sqlite3/notary_server_keys_metadata_table.go
index b48212bc..a2959407 100644
--- a/federationsender/storage/sqlite3/notary_server_keys_metadata_table.go
+++ b/federationsender/storage/sqlite3/notary_server_keys_metadata_table.go
@@ -17,6 +17,7 @@ package sqlite3
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
 	"fmt"
 	"strings"
 
@@ -150,7 +151,11 @@ func (s *notaryServerKeysMetadataStatements) SelectKeys(ctx context.Context, txn
 	var results []gomatrixserverlib.ServerKeys
 	for rows.Next() {
 		var sk gomatrixserverlib.ServerKeys
-		if err := rows.Scan(&sk.Raw); err != nil {
+		var raw string
+		if err = rows.Scan(&raw); err != nil {
+			return nil, err
+		}
+		if err = json.Unmarshal([]byte(raw), &sk); err != nil {
 			return nil, err
 		}
 		results = append(results, sk)