3 files changed, 129 insertions, 12 deletions
diff --git a/federationapi/storage/postgres/server_key_table.go b/federationapi/storage/postgres/server_key_table.go
index c62446da..fa58f1ea 100644
--- a/federationapi/storage/postgres/server_key_table.go
+++ b/federationapi/storage/postgres/server_key_table.go
@@ -94,12 +94,14 @@ func (s *serverSigningKeyStatements) BulkSelectServerKeys(
 	}
 	defer internal.CloseAndLogIfError(ctx, rows, "bulkSelectServerKeys: rows.close() failed")
 	results := map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult{}
+
+	var serverName string
+	var keyID string
+	var key string
+	var validUntilTS int64
+	var expiredTS int64
+	var vk gomatrixserverlib.VerifyKey
 	for rows.Next() {
-		var serverName string
-		var keyID string
-		var key string
-		var validUntilTS int64
-		var expiredTS int64
 		if err = rows.Scan(&serverName, &keyID, &validUntilTS, &expiredTS, &key); err != nil {
 			return nil, err
 		}
@@ -107,7 +109,6 @@ func (s *serverSigningKeyStatements) BulkSelectServerKeys(
 			ServerName: spec.ServerName(serverName),
 			KeyID:      gomatrixserverlib.KeyID(keyID),
 		}
-		vk := gomatrixserverlib.VerifyKey{}
 		err = vk.Key.Decode(key)
 		if err != nil {
 			return nil, err
diff --git a/federationapi/storage/sqlite3/server_key_table.go b/federationapi/storage/sqlite3/server_key_table.go
index f28b8994..65a854ce 100644
--- a/federationapi/storage/sqlite3/server_key_table.go
+++ b/federationapi/storage/sqlite3/server_key_table.go
@@ -98,12 +98,13 @@ func (s *serverSigningKeyStatements) BulkSelectServerKeys(
 	err := sqlutil.RunLimitedVariablesQuery(
 		ctx, bulkSelectServerSigningKeysSQL, s.db, iKeyIDs, sqlutil.SQLite3MaxVariables,
 		func(rows *sql.Rows) error {
+			var serverName string
+			var keyID string
+			var key string
+			var validUntilTS int64
+			var expiredTS int64
+			var vk gomatrixserverlib.VerifyKey
 			for rows.Next() {
-				var serverName string
-				var keyID string
-				var key string
-				var validUntilTS int64
-				var expiredTS int64
 				if err := rows.Scan(&serverName, &keyID, &validUntilTS, &expiredTS, &key); err != nil {
 					return fmt.Errorf("bulkSelectServerKeys: %v", err)
 				}
@@ -111,7 +112,6 @@ func (s *serverSigningKeyStatements) BulkSelectServerKeys(
 					ServerName: spec.ServerName(serverName),
 					KeyID:      gomatrixserverlib.KeyID(keyID),
 				}
-				vk := gomatrixserverlib.VerifyKey{}
 				err := vk.Key.Decode(key)
 				if err != nil {
 					return fmt.Errorf("bulkSelectServerKeys: %v", err)
diff --git a/federationapi/storage/tables/server_key_table_test.go b/federationapi/storage/tables/server_key_table_test.go
new file mode 100644
index 00000000..e79a086b
--- /dev/null
+++ b/federationapi/storage/tables/server_key_table_test.go
@@ -0,0 +1,116 @@
+package tables_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/matrix-org/dendrite/federationapi/storage/postgres"
+	"github.com/matrix-org/dendrite/federationapi/storage/sqlite3"
+	"github.com/matrix-org/dendrite/federationapi/storage/tables"
+	"github.com/matrix-org/dendrite/internal/sqlutil"
+	"github.com/matrix-org/dendrite/setup/config"
+	"github.com/matrix-org/dendrite/test"
+	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/matrix-org/gomatrixserverlib/spec"
+	"github.com/stretchr/testify/assert"
+)
+
+func mustCreateServerKeyDB(t *testing.T, dbType test.DBType) (tables.FederationServerSigningKeys, func()) {
+	connStr, close := test.PrepareDBConnectionString(t, dbType)
+	db, err := sqlutil.Open(&config.DatabaseOptions{
+		ConnectionString: config.DataSource(connStr),
+	}, sqlutil.NewExclusiveWriter())
+	if err != nil {
+		t.Fatalf("failed to open database: %s", err)
+	}
+	var tab tables.FederationServerSigningKeys
+	switch dbType {
+	case test.DBTypePostgres:
+		tab, err = postgres.NewPostgresServerSigningKeysTable(db)
+	case test.DBTypeSQLite:
+		tab, err = sqlite3.NewSQLiteServerSigningKeysTable(db)
+	}
+	if err != nil {
+		t.Fatalf("failed to create table: %s", err)
+	}
+	return tab, close
+}
+
+func TestServerKeysTable(t *testing.T) {
+	test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
+		ctx, cancel := context.WithCancel(context.Background())
+		tab, close := mustCreateServerKeyDB(t, dbType)
+		t.Cleanup(func() {
+			close()
+			cancel()
+		})
+
+		req := gomatrixserverlib.PublicKeyLookupRequest{
+			ServerName: "localhost",
+			KeyID:      "ed25519:test",
+		}
+		expectedTimestamp := spec.AsTimestamp(time.Now().Add(time.Hour))
+		res := gomatrixserverlib.PublicKeyLookupResult{
+			VerifyKey:    gomatrixserverlib.VerifyKey{Key: make(spec.Base64Bytes, 0)},
+			ExpiredTS:    0,
+			ValidUntilTS: expectedTimestamp,
+		}
+
+		// Insert the key
+		err := tab.UpsertServerKeys(ctx, nil, req, res)
+		assert.NoError(t, err)
+
+		selectKeys := map[gomatrixserverlib.PublicKeyLookupRequest]spec.Timestamp{
+			req: spec.AsTimestamp(time.Now()),
+		}
+		gotKeys, err := tab.BulkSelectServerKeys(ctx, nil, selectKeys)
+		assert.NoError(t, err)
+
+		// Now we should have a key for the req above
+		assert.NotNil(t, gotKeys[req])
+		assert.Equal(t, res, gotKeys[req])
+
+		// "Expire" the key by setting ExpireTS to a non-zero value and ValidUntilTS to 0
+		expectedTimestamp = spec.AsTimestamp(time.Now())
+		res.ExpiredTS = expectedTimestamp
+		res.ValidUntilTS = 0
+
+		// Update the key
+		err = tab.UpsertServerKeys(ctx, nil, req, res)
+		assert.NoError(t, err)
+
+		gotKeys, err = tab.BulkSelectServerKeys(ctx, nil, selectKeys)
+		assert.NoError(t, err)
+
+		// The key should be expired
+		assert.NotNil(t, gotKeys[req])
+		assert.Equal(t, res, gotKeys[req])
+
+		// Upsert a different key to validate querying multiple keys
+		req2 := gomatrixserverlib.PublicKeyLookupRequest{
+			ServerName: "notlocalhost",
+			KeyID:      "ed25519:test2",
+		}
+		expectedTimestamp2 := spec.AsTimestamp(time.Now().Add(time.Hour))
+		res2 := gomatrixserverlib.PublicKeyLookupResult{
+			VerifyKey:    gomatrixserverlib.VerifyKey{Key: make(spec.Base64Bytes, 0)},
+			ExpiredTS:    0,
+			ValidUntilTS: expectedTimestamp2,
+		}
+
+		err = tab.UpsertServerKeys(ctx, nil, req2, res2)
+		assert.NoError(t, err)
+
+		// Select multiple keys
+		selectKeys[req2] = spec.AsTimestamp(time.Now())
+
+		gotKeys, err = tab.BulkSelectServerKeys(ctx, nil, selectKeys)
+		assert.NoError(t, err)
+
+		// We now should receive two keys, one of which is expired
+		assert.Equal(t, 2, len(gotKeys))
+		assert.Equal(t, res2, gotKeys[req2])
+		assert.Equal(t, res, gotKeys[req])
+	})
+}