Add a database for storing the server keys (#137)

* Add a database for storing the server keys

* Tweak wording, and comment on the resolution of the timestamp

* Update gomatrixserverlib

src/github.com/matrix-org/dendrite/cmd/dendrite-client-api-server/main.go

@@ -25,6 +25,7 @@ import (
 	"github.com/matrix-org/dendrite/clientapi/producers"
 	"github.com/matrix-org/dendrite/clientapi/routing"
 	"github.com/matrix-org/dendrite/common"
+	"github.com/matrix-org/dendrite/common/keydb"
 	"github.com/matrix-org/dendrite/roomserver/api"
 
 	"github.com/matrix-org/gomatrixserverlib"
@@ -41,6 +42,7 @@ var (
 	serverName           = gomatrixserverlib.ServerName(os.Getenv("SERVER_NAME"))
 	serverKey            = os.Getenv("SERVER_KEY")
 	accountDataSource    = os.Getenv("ACCOUNT_DATABASE")
+	keyDataSource        = os.Getenv("KEY_DATABASE")
 )
 
 func main() {
@@ -79,7 +81,7 @@ func main() {
 
 	roomserverProducer, err := producers.NewRoomserverProducer(cfg.KafkaProducerURIs, cfg.ClientAPIOutputTopic)
 	if err != nil {
-		log.Panicf("Failed to setup kafka producers(%s): %s", cfg.KafkaProducerURIs, err)
+		log.Panicf("Failed to setup kafka producers(%q): %s", cfg.KafkaProducerURIs, err)
 	}
 
 	federation := gomatrixserverlib.NewFederationClient(cfg.ServerName, cfg.KeyID, cfg.PrivateKey)
@@ -87,11 +89,15 @@ func main() {
 	queryAPI := api.NewRoomserverQueryAPIHTTP(cfg.RoomserverURL, nil)
 	accountDB, err := accounts.NewDatabase(accountDataSource, serverName)
 	if err != nil {
-		log.Panicf("Failed to setup account database(%s): %s", accountDataSource, err.Error())
+		log.Panicf("Failed to setup account database(%q): %s", accountDataSource, err.Error())
 	}
 	deviceDB, err := devices.NewDatabase(accountDataSource, serverName)
 	if err != nil {
-		log.Panicf("Failed to setup device database(%s): %s", accountDataSource, err.Error())
+		log.Panicf("Failed to setup device database(%q): %s", accountDataSource, err.Error())
+	}
+	keyDB, err := keydb.NewDatabase(keyDataSource)
+	if err != nil {
+		log.Panicf("Failed to setup key database(%q): %s", keyDataSource, err.Error())
 	}
 
 	keyRing := gomatrixserverlib.KeyRing{
@@ -99,7 +105,7 @@ func main() {
 			// TODO: Use perspective key fetchers for production.
 			&gomatrixserverlib.DirectKeyFetcher{federation.Client},
 		},
-		KeyDatabase: &dummyKeyDatabase{},
+		KeyDatabase: keyDB,
 	}
 
 	routing.Setup(
@@ -108,18 +114,3 @@ func main() {
 	)
 	log.Fatal(http.ListenAndServe(bindAddr, nil))
 }
-
-// TODO: Implement a proper key database.
-type dummyKeyDatabase struct{}
-
-func (d *dummyKeyDatabase) FetchKeys(
-	requests map[gomatrixserverlib.PublicKeyRequest]gomatrixserverlib.Timestamp,
-) (map[gomatrixserverlib.PublicKeyRequest]gomatrixserverlib.ServerKeys, error) {
-	return nil, nil
-}
-
-func (d *dummyKeyDatabase) StoreKeys(
-	map[gomatrixserverlib.PublicKeyRequest]gomatrixserverlib.ServerKeys,
-) error {
-	return nil
-}

src/github.com/matrix-org/dendrite/cmd/dendrite-federation-api-server/main.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/matrix-org/dendrite/clientapi/producers"
 	"github.com/matrix-org/dendrite/common"
+	"github.com/matrix-org/dendrite/common/keydb"
 	"github.com/matrix-org/dendrite/federationapi/config"
 	"github.com/matrix-org/dendrite/federationapi/routing"
 	"github.com/matrix-org/dendrite/roomserver/api"
@@ -47,6 +48,7 @@ var (
 	kafkaURIs            = strings.Split(os.Getenv("KAFKA_URIS"), ",")
 	roomserverURL        = os.Getenv("ROOMSERVER_URL")
 	roomserverInputTopic = os.Getenv("TOPIC_INPUT_ROOM_EVENT")
+	keyDataSource        = os.Getenv("KEY_DATABASE")
 )
 
 func main() {
@@ -95,12 +97,17 @@ func main() {
 
 	federation := gomatrixserverlib.NewFederationClient(cfg.ServerName, cfg.KeyID, cfg.PrivateKey)
 
+	keyDB, err := keydb.NewDatabase(keyDataSource)
+	if err != nil {
+		log.Panicf("Failed to setup key database(%q): %s", keyDataSource, err.Error())
+	}
+
 	keyRing := gomatrixserverlib.KeyRing{
 		KeyFetchers: []gomatrixserverlib.KeyFetcher{
 			// TODO: Use perspective key fetchers for production.
 			&gomatrixserverlib.DirectKeyFetcher{federation.Client},
 		},
-		KeyDatabase: &dummyKeyDatabase{},
+		KeyDatabase: keyDB,
 	}
 	queryAPI := api.NewRoomserverQueryAPIHTTP(roomserverURL, nil)
 
@@ -112,18 +119,3 @@ func main() {
 	routing.Setup(http.DefaultServeMux, cfg, queryAPI, roomserverProducer, keyRing, federation)
 	log.Fatal(http.ListenAndServe(bindAddr, nil))
 }
-
-// TODO: Implement a proper key database.
-type dummyKeyDatabase struct{}
-
-func (d *dummyKeyDatabase) FetchKeys(
-	requests map[gomatrixserverlib.PublicKeyRequest]gomatrixserverlib.Timestamp,
-) (map[gomatrixserverlib.PublicKeyRequest]gomatrixserverlib.ServerKeys, error) {
-	return nil, nil
-}
-
-func (d *dummyKeyDatabase) StoreKeys(
-	map[gomatrixserverlib.PublicKeyRequest]gomatrixserverlib.ServerKeys,
-) error {
-	return nil
-}

src/github.com/matrix-org/dendrite/common/keydb/keydb.go

@@ -0,0 +1,70 @@
+// 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 keydb
+
+import (
+	"database/sql"
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+// A Database implements gomatrixserverlib.KeyDatabase and is used to store
+// the public keys for other matrix servers.
+type Database struct {
+	statements serverKeyStatements
+}
+
+// NewDatabase prepares a new key database.
+// It creates the necessary tables if they don't already exist.
+// It prepares all the SQL statements that it will use.
+// Returns an error if there was a problem talking to the database.
+func NewDatabase(dataSourceName string) (*Database, error) {
+	db, err := sql.Open("postgres", dataSourceName)
+	if err != nil {
+		return nil, err
+	}
+	d := &Database{}
+	d.statements.prepare(db)
+	return d, nil
+}
+
+// FetchKeys implements gomatrixserverlib.KeyDatabase
+func (d *Database) FetchKeys(
+	requests map[gomatrixserverlib.PublicKeyRequest]gomatrixserverlib.Timestamp,
+) (map[gomatrixserverlib.PublicKeyRequest]gomatrixserverlib.ServerKeys, error) {
+	return d.statements.bulkSelectServerKeys(requests)
+}
+
+// StoreKeys implements gomatrixserverlib.KeyDatabase
+func (d *Database) StoreKeys(
+	keyMap map[gomatrixserverlib.PublicKeyRequest]gomatrixserverlib.ServerKeys,
+) error {
+	// TODO: Inserting all the keys within a single transaction may
+	// be more efficient since the transaction overhead can be quite
+	// high for a single insert statement.
+	var lastErr error
+	for request, keys := range keyMap {
+		if err := d.statements.upsertServerKeys(request, keys); err != nil {
+			// Rather than returning immediately on error we try to insert the
+			// remaining keys.
+			// Since we are inserting the keys outside of a transaction it is
+			// possible for some of the inserts to succeed even though some
+			// of the inserts have failed.
+			// Ensuring that we always insert all the keys we can means that
+			// this behaviour won't depend on the iteration order of the map.
+			lastErr = err
+		}
+	}
+	return lastErr
+}

src/github.com/matrix-org/dendrite/common/keydb/server_key_table.go

@@ -0,0 +1,125 @@
+// 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 keydb
+
+import (
+	"database/sql"
+	"encoding/json"
+	"github.com/lib/pq"
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+const serverKeysSchema = `
+-- A cache of server keys downloaded from remote servers.
+CREATE TABLE IF NOT EXISTS server_keys (
+	-- The name of the matrix server the key is for.
+	server_name TEXT NOT NULL,
+	-- The ID of the server key.
+	server_key_id TEXT NOT NULL,
+	-- Combined server name and key ID separated by the ASCII unit separator
+	-- to make it easier to run bulk queries.
+	server_name_and_key_id TEXT NOT NULL,
+	-- When the keys are valid until as a millisecond timestamp.
+	valid_until_ts BIGINT NOT NULL,
+	-- The raw JSON for the server key.
+	server_key_json TEXT NOT NULL,
+	CONSTRAINT server_keys_unique UNIQUE (server_name, server_key_id)
+);
+
+CREATE INDEX IF NOT EXISTS server_name_and_key_id ON server_keys (server_name_and_key_id);
+`
+
+const bulkSelectServerKeysSQL = "" +
+	"SELECT server_name, server_key_id, server_key_json FROM server_keys" +
+	" WHERE server_name_and_key_id = ANY($1)"
+
+const upsertServerKeysSQL = "" +
+	"INSERT INTO server_keys (server_name, server_key_id," +
+	" server_name_and_key_id, valid_until_ts, server_key_json)" +
+	" VALUES ($1, $2, $3, $4, $5)" +
+	" ON CONFLICT ON CONSTRAINT server_keys_unique" +
+	" DO UPDATE SET valid_until_ts = $4, server_key_json = $5"
+
+type serverKeyStatements struct {
+	bulkSelectServerKeysStmt *sql.Stmt
+	upsertServerKeysStmt     *sql.Stmt
+}
+
+func (s *serverKeyStatements) prepare(db *sql.DB) (err error) {
+	_, err = db.Exec(serverKeysSchema)
+	if err != nil {
+		return
+	}
+	if s.bulkSelectServerKeysStmt, err = db.Prepare(bulkSelectServerKeysSQL); err != nil {
+		return
+	}
+	if s.upsertServerKeysStmt, err = db.Prepare(upsertServerKeysSQL); err != nil {
+		return
+	}
+	return
+}
+
+func (s *serverKeyStatements) bulkSelectServerKeys(
+	requests map[gomatrixserverlib.PublicKeyRequest]gomatrixserverlib.Timestamp,
+) (map[gomatrixserverlib.PublicKeyRequest]gomatrixserverlib.ServerKeys, error) {
+	var nameAndKeyIDs []string
+	for request := range requests {
+		nameAndKeyIDs = append(nameAndKeyIDs, nameAndKeyID(request))
+	}
+	rows, err := s.bulkSelectServerKeysStmt.Query(pq.StringArray(nameAndKeyIDs))
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	results := map[gomatrixserverlib.PublicKeyRequest]gomatrixserverlib.ServerKeys{}
+	for rows.Next() {
+		var serverName string
+		var keyID string
+		var keyJSON []byte
+		if err := rows.Scan(&serverName, &keyID, &keyJSON); err != nil {
+			return nil, err
+		}
+		var serverKeys gomatrixserverlib.ServerKeys
+		if err := json.Unmarshal(keyJSON, &serverKeys); err != nil {
+			return nil, err
+		}
+		r := gomatrixserverlib.PublicKeyRequest{
+			gomatrixserverlib.ServerName(serverName), gomatrixserverlib.KeyID(keyID),
+		}
+		results[r] = serverKeys
+	}
+	return results, nil
+}
+
+func (s *serverKeyStatements) upsertServerKeys(
+	request gomatrixserverlib.PublicKeyRequest, keys gomatrixserverlib.ServerKeys,
+) error {
+	keyJSON, err := json.Marshal(keys)
+	if err != nil {
+		return err
+	}
+	_, err = s.upsertServerKeysStmt.Exec(
+		string(request.ServerName), string(request.KeyID), nameAndKeyID(request),
+		int64(keys.ValidUntilTS), keyJSON,
+	)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func nameAndKeyID(request gomatrixserverlib.PublicKeyRequest) string {
+	return string(request.ServerName) + "\x1F" + string(request.KeyID)
+}