fedsender: add cache tables for notary keys (#1923)

* Add notary server tables for postgres

* Add sqlite tables

* fedsender: GetServerKeys -> QueryServerKeys

As it now checks a cache and can return multiple responses

federationapi/routing/keys.go

@@ -188,40 +188,52 @@ func NotaryKeys(
 	}
 	response.ServerKeys = []json.RawMessage{}
 
-	for serverName := range req.ServerKeys {
-		var keys *gomatrixserverlib.ServerKeys
+	for serverName, kidToCriteria := range req.ServerKeys {
+		var keyList []gomatrixserverlib.ServerKeys
 		if serverName == cfg.Matrix.ServerName {
 			if k, err := localKeys(cfg, time.Now().Add(cfg.Matrix.KeyValidityPeriod)); err == nil {
-				keys = k
+				keyList = append(keyList, *k)
 			} else {
 				return util.ErrorResponse(err)
 			}
 		} else {
-			if k, err := fsAPI.GetServerKeys(httpReq.Context(), serverName); err == nil {
-				keys = &k
-			} else {
+			kids := make([]gomatrixserverlib.KeyID, len(kidToCriteria))
+			i := 0
+			for kid := range kidToCriteria {
+				kids[i] = kid
+				i++
+			}
+			var resp federationSenderAPI.QueryServerKeysResponse
+			err := fsAPI.QueryServerKeys(httpReq.Context(), &federationSenderAPI.QueryServerKeysRequest{
+				ServerName:     serverName,
+				OptionalKeyIDs: kids,
+			}, &resp)
+			if err != nil {
 				return util.ErrorResponse(err)
 			}
+			keyList = append(keyList, resp.ServerKeys...)
 		}
-		if keys == nil {
+		if len(keyList) == 0 {
 			continue
 		}
 
-		j, err := json.Marshal(keys)
-		if err != nil {
-			logrus.WithError(err).Errorf("Failed to marshal %q response", serverName)
-			return jsonerror.InternalServerError()
-		}
+		for _, keys := range keyList {
+			j, err := json.Marshal(keys)
+			if err != nil {
+				logrus.WithError(err).Errorf("Failed to marshal %q response", serverName)
+				return jsonerror.InternalServerError()
+			}
 
-		js, err := gomatrixserverlib.SignJSON(
-			string(cfg.Matrix.ServerName), cfg.Matrix.KeyID, cfg.Matrix.PrivateKey, j,
-		)
-		if err != nil {
-			logrus.WithError(err).Errorf("Failed to sign %q response", serverName)
-			return jsonerror.InternalServerError()
-		}
+			js, err := gomatrixserverlib.SignJSON(
+				string(cfg.Matrix.ServerName), cfg.Matrix.KeyID, cfg.Matrix.PrivateKey, j,
+			)
+			if err != nil {
+				logrus.WithError(err).Errorf("Failed to sign %q response", serverName)
+				return jsonerror.InternalServerError()
+			}
 
-		response.ServerKeys = append(response.ServerKeys, js)
+			response.ServerKeys = append(response.ServerKeys, js)
+		}
 	}
 
 	return util.JSONResponse{

federationsender/api/api.go

@@ -20,7 +20,6 @@ type FederationClient interface {
 	ClaimKeys(ctx context.Context, s gomatrixserverlib.ServerName, oneTimeKeys map[string]map[string]string) (res gomatrixserverlib.RespClaimKeys, err error)
 	QueryKeys(ctx context.Context, s gomatrixserverlib.ServerName, keys map[string][]string) (res gomatrixserverlib.RespQueryKeys, err error)
 	GetEvent(ctx context.Context, s gomatrixserverlib.ServerName, eventID string) (res gomatrixserverlib.Transaction, err error)
-	GetServerKeys(ctx context.Context, matrixServer gomatrixserverlib.ServerName) (gomatrixserverlib.ServerKeys, error)
 	MSC2836EventRelationships(ctx context.Context, dst gomatrixserverlib.ServerName, r gomatrixserverlib.MSC2836EventRelationshipsRequest, roomVersion gomatrixserverlib.RoomVersion) (res gomatrixserverlib.MSC2836EventRelationshipsResponse, err error)
 	MSC2946Spaces(ctx context.Context, dst gomatrixserverlib.ServerName, roomID string, r gomatrixserverlib.MSC2946SpacesRequest) (res gomatrixserverlib.MSC2946SpacesResponse, err error)
 	LookupServerKeys(ctx context.Context, s gomatrixserverlib.ServerName, keyRequests map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp) ([]gomatrixserverlib.ServerKeys, error)
@@ -41,6 +40,8 @@ func (e *FederationClientError) Error() string {
 type FederationSenderInternalAPI interface {
 	FederationClient
 
+	QueryServerKeys(ctx context.Context, request *QueryServerKeysRequest, response *QueryServerKeysResponse) error
+
 	// PerformDirectoryLookup looks up a remote room ID from a room alias.
 	PerformDirectoryLookup(
 		ctx context.Context,
@@ -94,6 +95,15 @@ type FederationSenderInternalAPI interface {
 	) error
 }
 
+type QueryServerKeysRequest struct {
+	ServerName     gomatrixserverlib.ServerName
+	OptionalKeyIDs []gomatrixserverlib.KeyID
+}
+
+type QueryServerKeysResponse struct {
+	ServerKeys []gomatrixserverlib.ServerKeys
+}
+
 type PerformDirectoryLookupRequest struct {
 	RoomAlias  string                       `json:"room_alias"`
 	ServerName gomatrixserverlib.ServerName `json:"server_name"`

federationsender/internal/api.go

@@ -202,20 +202,6 @@ func (a *FederationSenderInternalAPI) GetEvent(
 	return ires.(gomatrixserverlib.Transaction), nil
 }
 
-func (a *FederationSenderInternalAPI) GetServerKeys(
-	ctx context.Context, s gomatrixserverlib.ServerName,
-) (gomatrixserverlib.ServerKeys, error) {
-	ctx, cancel := context.WithTimeout(ctx, time.Second*30)
-	defer cancel()
-	ires, err := a.doRequest(s, func() (interface{}, error) {
-		return a.federation.GetServerKeys(ctx, s)
-	})
-	if err != nil {
-		return gomatrixserverlib.ServerKeys{}, err
-	}
-	return ires.(gomatrixserverlib.ServerKeys), nil
-}
-
 func (a *FederationSenderInternalAPI) LookupServerKeys(
 	ctx context.Context, s gomatrixserverlib.ServerName, keyRequests map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp,
 ) ([]gomatrixserverlib.ServerKeys, error) {

federationsender/internal/query.go

@@ -2,8 +2,12 @@ package internal
 
 import (
 	"context"
+	"fmt"
+	"time"
 
 	"github.com/matrix-org/dendrite/federationsender/api"
+	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/matrix-org/util"
 )
 
 // QueryJoinedHostServerNamesInRoom implements api.FederationSenderInternalAPI
@@ -20,3 +24,30 @@ func (f *FederationSenderInternalAPI) QueryJoinedHostServerNamesInRoom(
 
 	return
 }
+
+func (a *FederationSenderInternalAPI) QueryServerKeys(
+	ctx context.Context, req *api.QueryServerKeysRequest, res *api.QueryServerKeysResponse,
+) 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)
+	})
+	if err != nil {
+		// try to load from the cache
+		serverKeysResponses, dbErr := a.db.GetNotaryKeys(ctx, req.ServerName, req.OptionalKeyIDs)
+		if dbErr != nil {
+			return fmt.Errorf("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 {
+		// non-fatal, still return the response
+		util.GetLogger(ctx).WithError(err).Warn("failed to UpdateNotaryKeys")
+	}
+	res.ServerKeys = []gomatrixserverlib.ServerKeys{serverKeys}
+	return nil
+}

federationsender/inthttp/client.go

@@ -15,6 +15,7 @@ import (
 // HTTP paths for the internal HTTP API
 const (
 	FederationSenderQueryJoinedHostServerNamesInRoomPath = "/federationsender/queryJoinedHostServerNamesInRoom"
+	FederationSenderQueryServerKeysPath                  = "/federationsender/queryServerKeys"
 
 	FederationSenderPerformDirectoryLookupRequestPath = "/federationsender/performDirectoryLookup"
 	FederationSenderPerformJoinRequestPath            = "/federationsender/performJoinRequest"
@@ -31,7 +32,6 @@ const (
 	FederationSenderLookupStatePath        = "/federationsender/client/lookupState"
 	FederationSenderLookupStateIDsPath     = "/federationsender/client/lookupStateIDs"
 	FederationSenderGetEventPath           = "/federationsender/client/getEvent"
-	FederationSenderGetServerKeysPath      = "/federationsender/client/getServerKeys"
 	FederationSenderLookupServerKeysPath   = "/federationsender/client/lookupServerKeys"
 	FederationSenderEventRelationshipsPath = "/federationsender/client/msc2836eventRelationships"
 	FederationSenderSpacesSummaryPath      = "/federationsender/client/msc2946spacesSummary"
@@ -377,31 +377,14 @@ func (h *httpFederationSenderInternalAPI) GetEvent(
 	return *response.Res, nil
 }
 
-type getServerKeys struct {
-	S          gomatrixserverlib.ServerName
-	ServerKeys gomatrixserverlib.ServerKeys
-	Err        *api.FederationClientError
-}
-
-func (h *httpFederationSenderInternalAPI) GetServerKeys(
-	ctx context.Context, s gomatrixserverlib.ServerName,
-) (gomatrixserverlib.ServerKeys, error) {
-	span, ctx := opentracing.StartSpanFromContext(ctx, "GetServerKeys")
+func (h *httpFederationSenderInternalAPI) QueryServerKeys(
+	ctx context.Context, req *api.QueryServerKeysRequest, res *api.QueryServerKeysResponse,
+) error {
+	span, ctx := opentracing.StartSpanFromContext(ctx, "QueryServerKeys")
 	defer span.Finish()
 
-	request := getServerKeys{
-		S: s,
-	}
-	var response getServerKeys
-	apiURL := h.federationSenderURL + FederationSenderGetServerKeysPath
-	err := httputil.PostJSON(ctx, span, h.httpClient, apiURL, &request, &response)
-	if err != nil {
-		return gomatrixserverlib.ServerKeys{}, err
-	}
-	if response.Err != nil {
-		return gomatrixserverlib.ServerKeys{}, response.Err
-	}
-	return response.ServerKeys, nil
+	apiURL := h.federationSenderURL + FederationSenderQueryServerKeysPath
+	return httputil.PostJSON(ctx, span, h.httpClient, apiURL, req, res)
 }
 
 type lookupServerKeys struct {

federationsender/inthttp/server.go

@@ -264,25 +264,17 @@ func AddRoutes(intAPI api.FederationSenderInternalAPI, internalAPIMux *mux.Route
 		}),
 	)
 	internalAPIMux.Handle(
-		FederationSenderGetServerKeysPath,
-		httputil.MakeInternalAPI("GetServerKeys", func(req *http.Request) util.JSONResponse {
-			var request getServerKeys
+		FederationSenderQueryServerKeysPath,
+		httputil.MakeInternalAPI("QueryServerKeys", func(req *http.Request) util.JSONResponse {
+			var request api.QueryServerKeysRequest
+			var response api.QueryServerKeysResponse
 			if err := json.NewDecoder(req.Body).Decode(&request); err != nil {
 				return util.MessageResponse(http.StatusBadRequest, err.Error())
 			}
-			res, err := intAPI.GetServerKeys(req.Context(), request.S)
-			if err != nil {
-				ferr, ok := err.(*api.FederationClientError)
-				if ok {
-					request.Err = ferr
-				} else {
-					request.Err = &api.FederationClientError{
-						Err: err.Error(),
-					}
-				}
+			if err := intAPI.QueryServerKeys(req.Context(), &request, &response); err != nil {
+				return util.ErrorResponse(err)
 			}
-			request.ServerKeys = res
-			return util.JSONResponse{Code: http.StatusOK, JSON: request}
+			return util.JSONResponse{Code: http.StatusOK, JSON: &response}
 		}),
 	)
 	internalAPIMux.Handle(

federationsender/storage/interface.go

@@ -66,4 +66,10 @@ type Database interface {
 	RenewInboundPeek(ctx context.Context, serverName gomatrixserverlib.ServerName, roomID, peekID string, renewalInterval int64) error
 	GetInboundPeek(ctx context.Context, serverName gomatrixserverlib.ServerName, roomID, peekID string) (*types.InboundPeek, error)
 	GetInboundPeeks(ctx context.Context, roomID string) ([]types.InboundPeek, error)
+
+	// Update the notary with the given server keys from the given server name.
+	UpdateNotaryKeys(ctx context.Context, serverName gomatrixserverlib.ServerName, serverKeys gomatrixserverlib.ServerKeys) error
+	// Query the notary for the server keys for the given server. If `optKeyIDs` is not empty, multiple server keys may be returned (between 1 - len(optKeyIDs))
+	// such that the combination of all server keys will include all the `optKeyIDs`.
+	GetNotaryKeys(ctx context.Context, serverName gomatrixserverlib.ServerName, optKeyIDs []gomatrixserverlib.KeyID) ([]gomatrixserverlib.ServerKeys, error)
 }

federationsender/storage/postgres/notary_server_keys_json_table.go

@@ -0,0 +1,64 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// 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 postgres
+
+import (
+	"context"
+	"database/sql"
+
+	"github.com/matrix-org/dendrite/federationsender/storage/tables"
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+const notaryServerKeysJSONSchema = `
+CREATE SEQUENCE IF NOT EXISTS federationsender_notary_server_keys_json_pkey;
+CREATE TABLE IF NOT EXISTS federationsender_notary_server_keys_json (
+    notary_id BIGINT PRIMARY KEY NOT NULL DEFAULT nextval('federationsender_notary_server_keys_json_pkey'),
+	response_json TEXT NOT NULL,
+	server_name TEXT NOT NULL,
+	valid_until BIGINT NOT NULL
+);
+`
+
+const insertServerKeysJSONSQL = "" +
+	"INSERT INTO federationsender_notary_server_keys_json (response_json, server_name, valid_until) VALUES ($1, $2, $3)" +
+	" RETURNING notary_id"
+
+type notaryServerKeysStatements struct {
+	db                       *sql.DB
+	insertServerKeysJSONStmt *sql.Stmt
+}
+
+func NewPostgresNotaryServerKeysTable(db *sql.DB) (s *notaryServerKeysStatements, err error) {
+	s = &notaryServerKeysStatements{
+		db: db,
+	}
+	_, err = db.Exec(notaryServerKeysJSONSchema)
+	if err != nil {
+		return
+	}
+
+	if s.insertServerKeysJSONStmt, err = db.Prepare(insertServerKeysJSONSQL); err != nil {
+		return
+	}
+	return
+}
+
+func (s *notaryServerKeysStatements) InsertJSONResponse(
+	ctx context.Context, txn *sql.Tx, keyQueryResponseJSON gomatrixserverlib.ServerKeys, serverName gomatrixserverlib.ServerName, validUntil gomatrixserverlib.Timestamp,
+) (tables.NotaryID, error) {
+	var notaryID tables.NotaryID
+	return notaryID, txn.Stmt(s.insertServerKeysJSONStmt).QueryRowContext(ctx, string(keyQueryResponseJSON.Raw), serverName, validUntil).Scan(&notaryID)
+}

federationsender/storage/postgres/notary_server_keys_metadata_table.go

@@ -0,0 +1,162 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// 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 postgres
+
+import (
+	"context"
+	"database/sql"
+
+	"github.com/lib/pq"
+	"github.com/matrix-org/dendrite/federationsender/storage/tables"
+	"github.com/matrix-org/dendrite/internal"
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+const notaryServerKeysMetadataSchema = `
+CREATE TABLE IF NOT EXISTS federationsender_notary_server_keys_metadata (
+    notary_id BIGINT NOT NULL,
+	server_name TEXT NOT NULL,
+	key_id TEXT NOT NULL,
+	UNIQUE (server_name, key_id)
+);
+`
+
+const upsertServerKeysSQL = "" +
+	"INSERT INTO federationsender_notary_server_keys_metadata (notary_id, server_name, key_id) VALUES ($1, $2, $3)" +
+	" ON CONFLICT (server_name, key_id) DO UPDATE SET notary_id = $1"
+
+// for a given (server_name, key_id), find the existing notary ID and valid until. Used to check if we will replace it
+// JOINs with the json table
+const selectNotaryKeyMetadataSQL = `
+	SELECT federationsender_notary_server_keys_metadata.notary_id, valid_until FROM federationsender_notary_server_keys_json
+	JOIN federationsender_notary_server_keys_metadata ON
+	federationsender_notary_server_keys_metadata.notary_id = federationsender_notary_server_keys_json.notary_id
+	WHERE federationsender_notary_server_keys_metadata.server_name = $1 AND federationsender_notary_server_keys_metadata.key_id = $2
+`
+
+// select the response which has the highest valid_until value
+// JOINs with the json table
+const selectNotaryKeyResponsesSQL = `
+	SELECT response_json FROM federationsender_notary_server_keys_json
+	WHERE server_name = $1 AND valid_until = (
+		SELECT MAX(valid_until) FROM federationsender_notary_server_keys_json WHERE server_name = $1
+	)
+`
+
+// select the responses which have the given key IDs
+// JOINs with the json table
+const selectNotaryKeyResponsesWithKeyIDsSQL = `
+	SELECT response_json FROM federationsender_notary_server_keys_json
+	JOIN federationsender_notary_server_keys_metadata ON
+	federationsender_notary_server_keys_metadata.notary_id = federationsender_notary_server_keys_json.notary_id
+	WHERE federationsender_notary_server_keys_json.server_name = $1 AND federationsender_notary_server_keys_metadata.key_id = ANY ($2)
+	GROUP BY federationsender_notary_server_keys_json.notary_id
+`
+
+// JOINs with the metadata table
+const deleteUnusedServerKeysJSONSQL = `
+	DELETE FROM federationsender_notary_server_keys_json WHERE federationsender_notary_server_keys_json.notary_id NOT IN (
+		SELECT DISTINCT notary_id FROM federationsender_notary_server_keys_metadata
+	)
+`
+
+type notaryServerKeysMetadataStatements struct {
+	db                                     *sql.DB
+	upsertServerKeysStmt                   *sql.Stmt
+	selectNotaryKeyResponsesStmt           *sql.Stmt
+	selectNotaryKeyResponsesWithKeyIDsStmt *sql.Stmt
+	selectNotaryKeyMetadataStmt            *sql.Stmt
+	deleteUnusedServerKeysJSONStmt         *sql.Stmt
+}
+
+func NewPostgresNotaryServerKeysMetadataTable(db *sql.DB) (s *notaryServerKeysMetadataStatements, err error) {
+	s = &notaryServerKeysMetadataStatements{
+		db: db,
+	}
+	_, err = db.Exec(notaryServerKeysMetadataSchema)
+	if err != nil {
+		return
+	}
+
+	if s.upsertServerKeysStmt, err = db.Prepare(upsertServerKeysSQL); err != nil {
+		return
+	}
+	if s.selectNotaryKeyResponsesStmt, err = db.Prepare(selectNotaryKeyResponsesSQL); err != nil {
+		return
+	}
+	if s.selectNotaryKeyResponsesWithKeyIDsStmt, err = db.Prepare(selectNotaryKeyResponsesWithKeyIDsSQL); err != nil {
+		return
+	}
+	if s.selectNotaryKeyMetadataStmt, err = db.Prepare(selectNotaryKeyMetadataSQL); err != nil {
+		return
+	}
+	if s.deleteUnusedServerKeysJSONStmt, err = db.Prepare(deleteUnusedServerKeysJSONSQL); err != nil {
+		return
+	}
+	return
+}
+
+func (s *notaryServerKeysMetadataStatements) UpsertKey(
+	ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, keyID gomatrixserverlib.KeyID, newNotaryID tables.NotaryID, newValidUntil gomatrixserverlib.Timestamp,
+) (tables.NotaryID, error) {
+	notaryID := newNotaryID
+	// see if the existing notary ID a) exists, b) has a longer valid_until
+	var existingNotaryID tables.NotaryID
+	var existingValidUntil gomatrixserverlib.Timestamp
+	if err := txn.Stmt(s.selectNotaryKeyMetadataStmt).QueryRowContext(ctx, serverName, keyID).Scan(&existingNotaryID, &existingValidUntil); err != nil {
+		if err != sql.ErrNoRows {
+			return 0, err
+		}
+	}
+	if existingValidUntil.Time().After(newValidUntil.Time()) {
+		// the existing valid_until is valid longer, so use that.
+		return existingNotaryID, nil
+	}
+	// overwrite the notary_id for this (server_name, key_id) tuple
+	_, err := txn.Stmt(s.upsertServerKeysStmt).ExecContext(ctx, notaryID, serverName, keyID)
+	return notaryID, err
+}
+
+func (s *notaryServerKeysMetadataStatements) SelectKeys(ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, keyIDs []gomatrixserverlib.KeyID) ([]gomatrixserverlib.ServerKeys, error) {
+	var rows *sql.Rows
+	var err error
+	if len(keyIDs) == 0 {
+		rows, err = txn.Stmt(s.selectNotaryKeyResponsesStmt).QueryContext(ctx, string(serverName))
+	} else {
+		keyIDstr := make([]string, len(keyIDs))
+		for i := range keyIDs {
+			keyIDstr[i] = string(keyIDs[i])
+		}
+		rows, err = txn.Stmt(s.selectNotaryKeyResponsesWithKeyIDsStmt).QueryContext(ctx, string(serverName), pq.StringArray(keyIDstr))
+	}
+	if err != nil {
+		return nil, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "selectNotaryKeyResponsesStmt close failed")
+	var results []gomatrixserverlib.ServerKeys
+	for rows.Next() {
+		var sk gomatrixserverlib.ServerKeys
+		if err := rows.Scan(&sk.Raw); err != nil {
+			return nil, err
+		}
+		results = append(results, sk)
+	}
+	return results, nil
+}
+
+func (s *notaryServerKeysMetadataStatements) DeleteOldJSONResponses(ctx context.Context, txn *sql.Tx) error {
+	_, err := txn.Stmt(s.deleteUnusedServerKeysJSONStmt).ExecContext(ctx)
+	return err
+}

federationsender/storage/postgres/storage.go

@@ -17,6 +17,7 @@ package postgres
 
 import (
 	"database/sql"
+	"fmt"
 
 	"github.com/matrix-org/dendrite/federationsender/storage/postgres/deltas"
 	"github.com/matrix-org/dendrite/federationsender/storage/shared"
@@ -69,6 +70,14 @@ func NewDatabase(dbProperties *config.DatabaseOptions, cache caching.FederationS
 	if err != nil {
 		return nil, err
 	}
+	notaryJSON, err := NewPostgresNotaryServerKeysTable(d.db)
+	if err != nil {
+		return nil, fmt.Errorf("NewPostgresNotaryServerKeysTable: %s", err)
+	}
+	notaryMetadata, err := NewPostgresNotaryServerKeysMetadataTable(d.db)
+	if err != nil {
+		return nil, fmt.Errorf("NewPostgresNotaryServerKeysMetadataTable: %s", err)
+	}
 	m := sqlutil.NewMigrations()
 	deltas.LoadRemoveRoomsTable(m)
 	if err = m.RunDeltas(d.db, dbProperties); err != nil {
@@ -85,6 +94,8 @@ func NewDatabase(dbProperties *config.DatabaseOptions, cache caching.FederationS
 		FederationSenderBlacklist:     blacklist,
 		FederationSenderInboundPeeks:  inboundPeeks,
 		FederationSenderOutboundPeeks: outboundPeeks,
+		NotaryServerKeysJSON:          notaryJSON,
+		NotaryServerKeysMetadata:      notaryMetadata,
 	}
 	if err = d.PartitionOffsetStatements.Prepare(d.db, d.writer, "federationsender"); err != nil {
 		return nil, err

federationsender/storage/shared/storage.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"time"
 
 	"github.com/matrix-org/dendrite/federationsender/storage/tables"
 	"github.com/matrix-org/dendrite/federationsender/types"
@@ -37,6 +38,8 @@ type Database struct {
 	FederationSenderBlacklist     tables.FederationSenderBlacklist
 	FederationSenderOutboundPeeks tables.FederationSenderOutboundPeeks
 	FederationSenderInboundPeeks  tables.FederationSenderInboundPeeks
+	NotaryServerKeysJSON          tables.FederationSenderNotaryServerKeysJSON
+	NotaryServerKeysMetadata      tables.FederationSenderNotaryServerKeysMetadata
 }
 
 // An Receipt contains the NIDs of a call to GetNextTransactionPDUs/EDUs.
@@ -197,3 +200,47 @@ func (d *Database) GetInboundPeek(ctx context.Context, serverName gomatrixserver
 func (d *Database) GetInboundPeeks(ctx context.Context, roomID string) ([]types.InboundPeek, error) {
 	return d.FederationSenderInboundPeeks.SelectInboundPeeks(ctx, nil, roomID)
 }
+
+func (d *Database) UpdateNotaryKeys(ctx context.Context, serverName gomatrixserverlib.ServerName, serverKeys gomatrixserverlib.ServerKeys) error {
+	return d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
+		validUntil := serverKeys.ValidUntilTS
+		// Servers MUST use the lesser of this field and 7 days into the future when determining if a key is valid.
+		// This is to avoid a situation where an attacker publishes a key which is valid for a significant amount of
+		// time without a way for the homeserver owner to revoke it.
+		// https://spec.matrix.org/unstable/server-server-api/#querying-keys-through-another-server
+		weekIntoFuture := time.Now().Add(7 * 24 * time.Hour)
+		if weekIntoFuture.Before(validUntil.Time()) {
+			validUntil = gomatrixserverlib.AsTimestamp(weekIntoFuture)
+		}
+		notaryID, err := d.NotaryServerKeysJSON.InsertJSONResponse(ctx, txn, serverKeys, serverName, validUntil)
+		if err != nil {
+			return err
+		}
+		// update the metadata for the keys
+		for keyID := range serverKeys.OldVerifyKeys {
+			_, err = d.NotaryServerKeysMetadata.UpsertKey(ctx, txn, serverName, keyID, notaryID, validUntil)
+			if err != nil {
+				return err
+			}
+		}
+		for keyID := range serverKeys.VerifyKeys {
+			_, err = d.NotaryServerKeysMetadata.UpsertKey(ctx, txn, serverName, keyID, notaryID, validUntil)
+			if err != nil {
+				return err
+			}
+		}
+
+		// clean up old responses
+		return d.NotaryServerKeysMetadata.DeleteOldJSONResponses(ctx, txn)
+	})
+}
+
+func (d *Database) GetNotaryKeys(
+	ctx context.Context, serverName gomatrixserverlib.ServerName, optKeyIDs []gomatrixserverlib.KeyID,
+) (sks []gomatrixserverlib.ServerKeys, err error) {
+	err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
+		sks, err = d.NotaryServerKeysMetadata.SelectKeys(ctx, txn, serverName, optKeyIDs)
+		return err
+	})
+	return sks, err
+}

federationsender/storage/sqlite3/notary_server_keys_json_table.go

@@ -0,0 +1,63 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// 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 sqlite3
+
+import (
+	"context"
+	"database/sql"
+
+	"github.com/matrix-org/dendrite/federationsender/storage/tables"
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+const notaryServerKeysJSONSchema = `
+CREATE TABLE IF NOT EXISTS federationsender_notary_server_keys_json (
+    notary_id INTEGER PRIMARY KEY AUTOINCREMENT,
+	response_json TEXT NOT NULL,
+	server_name TEXT NOT NULL,
+	valid_until BIGINT NOT NULL
+);
+`
+
+const insertServerKeysJSONSQL = "" +
+	"INSERT INTO federationsender_notary_server_keys_json (response_json, server_name, valid_until) VALUES ($1, $2, $3)" +
+	" RETURNING notary_id"
+
+type notaryServerKeysStatements struct {
+	db                       *sql.DB
+	insertServerKeysJSONStmt *sql.Stmt
+}
+
+func NewSQLiteNotaryServerKeysTable(db *sql.DB) (s *notaryServerKeysStatements, err error) {
+	s = &notaryServerKeysStatements{
+		db: db,
+	}
+	_, err = db.Exec(notaryServerKeysJSONSchema)
+	if err != nil {
+		return
+	}
+
+	if s.insertServerKeysJSONStmt, err = db.Prepare(insertServerKeysJSONSQL); err != nil {
+		return
+	}
+	return
+}
+
+func (s *notaryServerKeysStatements) InsertJSONResponse(
+	ctx context.Context, txn *sql.Tx, keyQueryResponseJSON gomatrixserverlib.ServerKeys, serverName gomatrixserverlib.ServerName, validUntil gomatrixserverlib.Timestamp,
+) (tables.NotaryID, error) {
+	var notaryID tables.NotaryID
+	return notaryID, txn.Stmt(s.insertServerKeysJSONStmt).QueryRowContext(ctx, string(keyQueryResponseJSON.Raw), serverName, validUntil).Scan(&notaryID)
+}

federationsender/storage/sqlite3/notary_server_keys_metadata_table.go

@@ -0,0 +1,164 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// 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 sqlite3
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"strings"
+
+	"github.com/matrix-org/dendrite/federationsender/storage/tables"
+	"github.com/matrix-org/dendrite/internal"
+	"github.com/matrix-org/dendrite/internal/sqlutil"
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+const notaryServerKeysMetadataSchema = `
+CREATE TABLE IF NOT EXISTS federationsender_notary_server_keys_metadata (
+    notary_id BIGINT NOT NULL,
+	server_name TEXT NOT NULL,
+	key_id TEXT NOT NULL,
+	UNIQUE (server_name, key_id)
+);
+`
+
+const upsertServerKeysSQL = "" +
+	"INSERT INTO federationsender_notary_server_keys_metadata (notary_id, server_name, key_id) VALUES ($1, $2, $3)" +
+	" ON CONFLICT (server_name, key_id) DO UPDATE SET notary_id = $1"
+
+// for a given (server_name, key_id), find the existing notary ID and valid until. Used to check if we will replace it
+// JOINs with the json table
+const selectNotaryKeyMetadataSQL = `
+	SELECT federationsender_notary_server_keys_metadata.notary_id, valid_until FROM federationsender_notary_server_keys_json
+	JOIN federationsender_notary_server_keys_metadata ON
+	federationsender_notary_server_keys_metadata.notary_id = federationsender_notary_server_keys_json.notary_id
+	WHERE federationsender_notary_server_keys_metadata.server_name = $1 AND federationsender_notary_server_keys_metadata.key_id = $2
+`
+
+// select the response which has the highest valid_until value
+// JOINs with the json table
+const selectNotaryKeyResponsesSQL = `
+	SELECT response_json FROM federationsender_notary_server_keys_json
+	WHERE server_name = $1 AND valid_until = (
+		SELECT MAX(valid_until) FROM federationsender_notary_server_keys_json WHERE server_name = $1
+	)
+`
+
+// select the responses which have the given key IDs
+// JOINs with the json table
+const selectNotaryKeyResponsesWithKeyIDsSQL = `
+	SELECT response_json FROM federationsender_notary_server_keys_json
+	JOIN federationsender_notary_server_keys_metadata ON
+	federationsender_notary_server_keys_metadata.notary_id = federationsender_notary_server_keys_json.notary_id
+	WHERE federationsender_notary_server_keys_json.server_name = $1 AND federationsender_notary_server_keys_metadata.key_id IN ($2)
+	GROUP BY federationsender_notary_server_keys_json.notary_id
+`
+
+// JOINs with the metadata table
+const deleteUnusedServerKeysJSONSQL = `
+	DELETE FROM federationsender_notary_server_keys_json WHERE federationsender_notary_server_keys_json.notary_id NOT IN (
+		SELECT DISTINCT notary_id FROM federationsender_notary_server_keys_metadata
+	)
+`
+
+type notaryServerKeysMetadataStatements struct {
+	db                             *sql.DB
+	upsertServerKeysStmt           *sql.Stmt
+	selectNotaryKeyResponsesStmt   *sql.Stmt
+	selectNotaryKeyMetadataStmt    *sql.Stmt
+	deleteUnusedServerKeysJSONStmt *sql.Stmt
+}
+
+func NewSQLiteNotaryServerKeysMetadataTable(db *sql.DB) (s *notaryServerKeysMetadataStatements, err error) {
+	s = &notaryServerKeysMetadataStatements{
+		db: db,
+	}
+	_, err = db.Exec(notaryServerKeysMetadataSchema)
+	if err != nil {
+		return
+	}
+
+	if s.upsertServerKeysStmt, err = db.Prepare(upsertServerKeysSQL); err != nil {
+		return
+	}
+	if s.selectNotaryKeyResponsesStmt, err = db.Prepare(selectNotaryKeyResponsesSQL); err != nil {
+		return
+	}
+	if s.selectNotaryKeyMetadataStmt, err = db.Prepare(selectNotaryKeyMetadataSQL); err != nil {
+		return
+	}
+	if s.deleteUnusedServerKeysJSONStmt, err = db.Prepare(deleteUnusedServerKeysJSONSQL); err != nil {
+		return
+	}
+	return
+}
+
+func (s *notaryServerKeysMetadataStatements) UpsertKey(
+	ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, keyID gomatrixserverlib.KeyID, newNotaryID tables.NotaryID, newValidUntil gomatrixserverlib.Timestamp,
+) (tables.NotaryID, error) {
+	notaryID := newNotaryID
+	// see if the existing notary ID a) exists, b) has a longer valid_until
+	var existingNotaryID tables.NotaryID
+	var existingValidUntil gomatrixserverlib.Timestamp
+	if err := txn.Stmt(s.selectNotaryKeyMetadataStmt).QueryRowContext(ctx, serverName, keyID).Scan(&existingNotaryID, &existingValidUntil); err != nil {
+		if err != sql.ErrNoRows {
+			return 0, err
+		}
+	}
+	if existingValidUntil.Time().After(newValidUntil.Time()) {
+		// the existing valid_until is valid longer, so use that.
+		return existingNotaryID, nil
+	}
+	// overwrite the notary_id for this (server_name, key_id) tuple
+	_, err := txn.Stmt(s.upsertServerKeysStmt).ExecContext(ctx, notaryID, serverName, keyID)
+	return notaryID, err
+}
+
+func (s *notaryServerKeysMetadataStatements) SelectKeys(ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, keyIDs []gomatrixserverlib.KeyID) ([]gomatrixserverlib.ServerKeys, error) {
+	var rows *sql.Rows
+	var err error
+	if len(keyIDs) == 0 {
+		rows, err = txn.Stmt(s.selectNotaryKeyResponsesStmt).QueryContext(ctx, string(serverName))
+	} else {
+		iKeyIDs := make([]interface{}, len(keyIDs)+1)
+		iKeyIDs[0] = serverName
+		for i := range keyIDs {
+			iKeyIDs[i+1] = string(keyIDs[i])
+		}
+		sql := strings.Replace(selectNotaryKeyResponsesWithKeyIDsSQL, "($2)", sqlutil.QueryVariadicOffset(len(keyIDs), 1), 1)
+		fmt.Println(sql)
+		fmt.Println(iKeyIDs...)
+		rows, err = s.db.QueryContext(ctx, sql, iKeyIDs...)
+	}
+	if err != nil {
+		return nil, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "selectNotaryKeyResponsesStmt close failed")
+	var results []gomatrixserverlib.ServerKeys
+	for rows.Next() {
+		var sk gomatrixserverlib.ServerKeys
+		if err := rows.Scan(&sk.Raw); err != nil {
+			return nil, err
+		}
+		results = append(results, sk)
+	}
+	return results, nil
+}
+
+func (s *notaryServerKeysMetadataStatements) DeleteOldJSONResponses(ctx context.Context, txn *sql.Tx) error {
+	_, err := txn.Stmt(s.deleteUnusedServerKeysJSONStmt).ExecContext(ctx)
+	return err
+}

federationsender/storage/sqlite3/storage.go

@@ -71,6 +71,14 @@ func NewDatabase(dbProperties *config.DatabaseOptions, cache caching.FederationS
 	if err != nil {
 		return nil, err
 	}
+	notaryKeys, err := NewSQLiteNotaryServerKeysTable(d.db)
+	if err != nil {
+		return nil, err
+	}
+	notaryKeysMetadata, err := NewSQLiteNotaryServerKeysMetadataTable(d.db)
+	if err != nil {
+		return nil, err
+	}
 	m := sqlutil.NewMigrations()
 	deltas.LoadRemoveRoomsTable(m)
 	if err = m.RunDeltas(d.db, dbProperties); err != nil {
@@ -87,6 +95,8 @@ func NewDatabase(dbProperties *config.DatabaseOptions, cache caching.FederationS
 		FederationSenderBlacklist:     blacklist,
 		FederationSenderOutboundPeeks: outboundPeeks,
 		FederationSenderInboundPeeks:  inboundPeeks,
+		NotaryServerKeysJSON:          notaryKeys,
+		NotaryServerKeysMetadata:      notaryKeysMetadata,
 	}
 	if err = d.PartitionOffsetStatements.Prepare(d.db, d.writer, "federationsender"); err != nil {
 		return nil, err

federationsender/storage/tables/interface.go

@@ -22,6 +22,8 @@ import (
 	"github.com/matrix-org/gomatrixserverlib"
 )
 
+type NotaryID int64
+
 type FederationSenderQueuePDUs interface {
 	InsertQueuePDU(ctx context.Context, txn *sql.Tx, transactionID gomatrixserverlib.TransactionID, serverName gomatrixserverlib.ServerName, nid int64) error
 	DeleteQueuePDUs(ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, jsonNIDs []int64) error
@@ -80,3 +82,25 @@ type FederationSenderInboundPeeks interface {
 	DeleteInboundPeek(ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, roomID, peekID string) (err error)
 	DeleteInboundPeeks(ctx context.Context, txn *sql.Tx, roomID string) (err error)
 }
+
+// FederationSenderNotaryServerKeysJSON contains the byte-for-byte responses from servers which contain their keys and is signed by them.
+type FederationSenderNotaryServerKeysJSON interface {
+	// InsertJSONResponse inserts a new response JSON. Useless on its own, needs querying via FederationSenderNotaryServerKeysMetadata
+	// `validUntil` should be the value of `valid_until_ts` with the 7-day check applied from:
+	//   "Servers MUST use the lesser of this field and 7 days into the future when determining if a key is valid.
+	//    This is to avoid a situation where an attacker publishes a key which is valid for a significant amount of time
+	//    without a way for the homeserver owner to revoke it.""
+	InsertJSONResponse(ctx context.Context, txn *sql.Tx, keyQueryResponseJSON gomatrixserverlib.ServerKeys, serverName gomatrixserverlib.ServerName, validUntil gomatrixserverlib.Timestamp) (NotaryID, error)
+}
+
+// FederationSenderNotaryServerKeysMetadata persists the metadata for FederationSenderNotaryServerKeysJSON
+type FederationSenderNotaryServerKeysMetadata interface {
+	// UpsertKey updates or inserts a (server_name, key_id) tuple, pointing it via NotaryID at the the response which has the longest valid_until_ts
+	// `newNotaryID` and `newValidUntil` should be the notary ID / valid_until  which has this (server_name, key_id) tuple already, e.g one you just inserted.
+	UpsertKey(ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, keyID gomatrixserverlib.KeyID, newNotaryID NotaryID, newValidUntil gomatrixserverlib.Timestamp) (NotaryID, error)
+	// SelectKeys returns the signed JSON objects which contain the given key IDs. This will be at most the length of `keyIDs` and at least 1 (assuming
+	// the keys exist in the first place). If `keyIDs` is empty, the signed JSON object with the longest valid_until_ts will be returned.
+	SelectKeys(ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, keyIDs []gomatrixserverlib.KeyID) ([]gomatrixserverlib.ServerKeys, error)
+	// DeleteOldJSONResponses removes all responses which are not referenced in FederationSenderNotaryServerKeysMetadata
+	DeleteOldJSONResponses(ctx context.Context, txn *sql.Tx) error
+}

sytest-whitelist

@@ -533,3 +533,4 @@ Inbound federation can receive invite and reject when remote is unreachable
 Remote servers cannot set power levels in rooms without existing powerlevels
 Remote servers should reject attempts by non-creators to set the power levels
 Federation handles empty auth_events in state_ids sanely
+Key notary server should return an expired key if it can't find any others