Add initial support for storing user room keys (#3098)

roomserver/storage/tables/interface.go

@@ -2,6 +2,7 @@ package tables
 
 import (
 	"context"
+	"crypto/ed25519"
 	"database/sql"
 	"errors"
 
@@ -184,6 +185,19 @@ type Purge interface {
 	) error
 }
 
+type UserRoomKeys interface {
+	// InsertUserRoomPrivatePublicKey inserts the given private key as well as the public key for it. This should be used
+	// when creating keys locally.
+	InsertUserRoomPrivatePublicKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PrivateKey) (ed25519.PrivateKey, error)
+	// InsertUserRoomPublicKey inserts the given public key, this should be used for users NOT local to this server
+	InsertUserRoomPublicKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PublicKey) (ed25519.PublicKey, error)
+	// SelectUserRoomPrivateKey selects the private key for the given user and room combination
+	SelectUserRoomPrivateKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID) (ed25519.PrivateKey, error)
+	// BulkSelectUserNIDs selects all userIDs for the requested senderKeys. Returns a map from publicKey -> types.UserRoomKeyPair.
+	// If a senderKey can't be found, it is omitted in the result.
+	BulkSelectUserNIDs(ctx context.Context, txn *sql.Tx, senderKeys map[types.RoomNID][]ed25519.PublicKey) (map[string]types.UserRoomKeyPair, error)
+}
+
 // StrippedEvent represents a stripped event for returning extracted content values.
 type StrippedEvent struct {
 	RoomID       string

roomserver/storage/tables/user_room_keys_table_test.go

@@ -0,0 +1,115 @@
+package tables_test
+
+import (
+	"context"
+	"crypto/ed25519"
+	"database/sql"
+	"testing"
+
+	"github.com/matrix-org/dendrite/internal/sqlutil"
+	"github.com/matrix-org/dendrite/roomserver/storage/postgres"
+	"github.com/matrix-org/dendrite/roomserver/storage/sqlite3"
+	"github.com/matrix-org/dendrite/roomserver/storage/tables"
+	"github.com/matrix-org/dendrite/roomserver/types"
+	"github.com/matrix-org/dendrite/setup/config"
+	"github.com/matrix-org/dendrite/test"
+	"github.com/stretchr/testify/assert"
+	ed255192 "golang.org/x/crypto/ed25519"
+)
+
+func mustCreateUserRoomKeysTable(t *testing.T, dbType test.DBType) (tab tables.UserRoomKeys, db *sql.DB, close func()) {
+	t.Helper()
+	connStr, close := test.PrepareDBConnectionString(t, dbType)
+	db, err := sqlutil.Open(&config.DatabaseOptions{
+		ConnectionString: config.DataSource(connStr),
+	}, sqlutil.NewExclusiveWriter())
+	assert.NoError(t, err)
+	switch dbType {
+	case test.DBTypePostgres:
+		err = postgres.CreateUserRoomKeysTable(db)
+		assert.NoError(t, err)
+		tab, err = postgres.PrepareUserRoomKeysTable(db)
+	case test.DBTypeSQLite:
+		err = sqlite3.CreateUserRoomKeysTable(db)
+		assert.NoError(t, err)
+		tab, err = sqlite3.PrepareUserRoomKeysTable(db)
+	}
+	assert.NoError(t, err)
+
+	return tab, db, close
+}
+
+func TestUserRoomKeysTable(t *testing.T) {
+	test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
+		tab, db, close := mustCreateUserRoomKeysTable(t, dbType)
+		defer close()
+		userNID := types.EventStateKeyNID(1)
+		roomNID := types.RoomNID(1)
+		_, key, err := ed25519.GenerateKey(nil)
+		assert.NoError(t, err)
+
+		err = sqlutil.WithTransaction(db, func(txn *sql.Tx) error {
+			var gotKey, key2, key3 ed25519.PrivateKey
+			gotKey, err = tab.InsertUserRoomPrivatePublicKey(context.Background(), txn, userNID, roomNID, key)
+			assert.NoError(t, err)
+			assert.Equal(t, gotKey, key)
+
+			// again, this shouldn't result in an error, but return the existing key
+			_, key2, err = ed25519.GenerateKey(nil)
+			assert.NoError(t, err)
+			gotKey, err = tab.InsertUserRoomPrivatePublicKey(context.Background(), txn, userNID, roomNID, key2)
+			assert.NoError(t, err)
+			assert.Equal(t, gotKey, key)
+
+			// add another user
+			_, key3, err = ed25519.GenerateKey(nil)
+			assert.NoError(t, err)
+			userNID2 := types.EventStateKeyNID(2)
+			_, err = tab.InsertUserRoomPrivatePublicKey(context.Background(), txn, userNID2, roomNID, key3)
+			assert.NoError(t, err)
+
+			gotKey, err = tab.SelectUserRoomPrivateKey(context.Background(), txn, userNID, roomNID)
+			assert.NoError(t, err)
+			assert.Equal(t, key, gotKey)
+
+			// try to update an existing key, this should only be done for users NOT on this homeserver
+			var gotPubKey ed25519.PublicKey
+			gotPubKey, err = tab.InsertUserRoomPublicKey(context.Background(), txn, userNID, roomNID, key2.Public().(ed25519.PublicKey))
+			assert.NoError(t, err)
+			assert.Equal(t, key2.Public(), gotPubKey)
+
+			// Key doesn't exist
+			gotKey, err = tab.SelectUserRoomPrivateKey(context.Background(), txn, userNID, 2)
+			assert.NoError(t, err)
+			assert.Nil(t, gotKey)
+
+			// query user NIDs for senderKeys
+			var gotKeys map[string]types.UserRoomKeyPair
+			query := map[types.RoomNID][]ed25519.PublicKey{
+				roomNID:          {key2.Public().(ed25519.PublicKey), key3.Public().(ed25519.PublicKey)},
+				types.RoomNID(2): {key.Public().(ed25519.PublicKey), key3.Public().(ed25519.PublicKey)}, // doesn't exist
+			}
+			gotKeys, err = tab.BulkSelectUserNIDs(context.Background(), txn, query)
+			assert.NoError(t, err)
+			assert.NotNil(t, gotKeys)
+
+			wantKeys := map[string]types.UserRoomKeyPair{
+				string(key2.Public().(ed25519.PublicKey)): {RoomNID: roomNID, EventStateKeyNID: userNID},
+				string(key3.Public().(ed25519.PublicKey)): {RoomNID: roomNID, EventStateKeyNID: userNID2},
+			}
+			assert.Equal(t, wantKeys, gotKeys)
+
+			// insert key that came in over federation
+			var gotPublicKey, key4 ed255192.PublicKey
+			key4, _, err = ed25519.GenerateKey(nil)
+			assert.NoError(t, err)
+			gotPublicKey, err = tab.InsertUserRoomPublicKey(context.Background(), txn, userNID, 2, key4)
+			assert.NoError(t, err)
+			assert.Equal(t, key4, gotPublicKey)
+
+			return nil
+		})
+		assert.NoError(t, err)
+
+	})
+}