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

roomserver/storage/interface.go

@@ -16,6 +16,7 @@ package storage
 
 import (
 	"context"
+	"crypto/ed25519"
 
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/gomatrixserverlib/spec"
@@ -27,6 +28,7 @@ import (
 )
 
 type Database interface {
+	UserRoomKeys
 	// Do we support processing input events for more than one room at a time?
 	SupportsConcurrentRoomInputs() bool
 	// RoomInfo returns room information for the given room ID, or nil if there is no room.
@@ -194,8 +196,22 @@ type Database interface {
 	) (gomatrixserverlib.PDU, gomatrixserverlib.PDU, 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, userID spec.UserID, roomID spec.RoomID, key ed25519.PrivateKey) (result ed25519.PrivateKey, err error)
+	// InsertUserRoomPublicKey inserts the given public key, this should be used for users NOT local to this server
+	InsertUserRoomPublicKey(ctx context.Context, userID spec.UserID, roomID spec.RoomID, key ed25519.PublicKey) (result ed25519.PublicKey, err error)
+	// SelectUserRoomPrivateKey selects the private key for the given user and room combination
+	SelectUserRoomPrivateKey(ctx context.Context, userID spec.UserID, roomID spec.RoomID) (key ed25519.PrivateKey, err error)
+	// SelectUserIDsForPublicKeys selects all userIDs for the requested senderKeys. Returns a map from roomID -> map from publicKey to userID.
+	// If a senderKey can't be found, it is omitted in the result.
+	SelectUserIDsForPublicKeys(ctx context.Context, publicKeys map[spec.RoomID][]ed25519.PublicKey) (map[spec.RoomID]map[string]string, error)
+}
+
 type RoomDatabase interface {
 	EventDatabase
+	UserRoomKeys
 	// RoomInfo returns room information for the given room ID, or nil if there is no room.
 	RoomInfo(ctx context.Context, roomID string) (*types.RoomInfo, error)
 	RoomInfoByNID(ctx context.Context, roomNID types.RoomNID) (*types.RoomInfo, error)

roomserver/storage/postgres/storage.go

@@ -131,6 +131,9 @@ func (d *Database) create(db *sql.DB) error {
 	if err := CreateRedactionsTable(db); err != nil {
 		return err
 	}
+	if err := CreateUserRoomKeysTable(db); err != nil {
+		return err
+	}
 
 	return nil
 }
@@ -192,6 +195,11 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room
 	if err != nil {
 		return err
 	}
+	userRoomKeys, err := PrepareUserRoomKeysTable(db)
+	if err != nil {
+		return err
+	}
+
 	d.Database = shared.Database{
 		DB: db,
 		EventDatabase: shared.EventDatabase{
@@ -215,6 +223,7 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room
 		MembershipTable:    membership,
 		PublishedTable:     published,
 		Purge:              purge,
+		UserRoomKeyTable:   userRoomKeys,
 	}
 	return nil
 }

roomserver/storage/postgres/user_room_keys_table.go

@@ -0,0 +1,132 @@
+// Copyright 2023 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"
+	"crypto/ed25519"
+	"database/sql"
+	"errors"
+
+	"github.com/lib/pq"
+	"github.com/matrix-org/dendrite/internal"
+	"github.com/matrix-org/dendrite/internal/sqlutil"
+	"github.com/matrix-org/dendrite/roomserver/storage/tables"
+	"github.com/matrix-org/dendrite/roomserver/types"
+)
+
+const userRoomKeysSchema = `
+CREATE TABLE IF NOT EXISTS roomserver_user_room_keys (     
+    user_nid    INTEGER NOT NULL,
+    room_nid    INTEGER NOT NULL,
+    pseudo_id_key BYTEA NULL, -- may be null for users not local to the server
+    pseudo_id_pub_key BYTEA NOT NULL,
+    CONSTRAINT roomserver_user_room_keys_pk PRIMARY KEY (user_nid, room_nid)
+);
+`
+
+const insertUserRoomPrivateKeySQL = `
+	INSERT INTO roomserver_user_room_keys (user_nid, room_nid, pseudo_id_key, pseudo_id_pub_key) VALUES ($1, $2, $3, $4)
+	ON CONFLICT ON CONSTRAINT roomserver_user_room_keys_pk DO UPDATE SET pseudo_id_key = roomserver_user_room_keys.pseudo_id_key
+	RETURNING (pseudo_id_key)
+`
+
+const insertUserRoomPublicKeySQL = `
+	INSERT INTO roomserver_user_room_keys (user_nid, room_nid, pseudo_id_pub_key) VALUES ($1, $2, $3)
+	ON CONFLICT ON CONSTRAINT roomserver_user_room_keys_pk DO UPDATE SET pseudo_id_pub_key = $3
+	RETURNING (pseudo_id_pub_key)
+`
+
+const selectUserRoomKeySQL = `SELECT pseudo_id_key FROM roomserver_user_room_keys WHERE user_nid = $1 AND room_nid = $2`
+
+const selectUserNIDsSQL = `SELECT user_nid, room_nid, pseudo_id_pub_key FROM roomserver_user_room_keys WHERE room_nid = ANY($1) AND pseudo_id_pub_key = ANY($2)`
+
+type userRoomKeysStatements struct {
+	insertUserRoomPrivateKeyStmt *sql.Stmt
+	insertUserRoomPublicKeyStmt  *sql.Stmt
+	selectUserRoomKeyStmt        *sql.Stmt
+	selectUserNIDsStmt           *sql.Stmt
+}
+
+func CreateUserRoomKeysTable(db *sql.DB) error {
+	_, err := db.Exec(userRoomKeysSchema)
+	return err
+}
+
+func PrepareUserRoomKeysTable(db *sql.DB) (tables.UserRoomKeys, error) {
+	s := &userRoomKeysStatements{}
+	return s, sqlutil.StatementList{
+		{&s.insertUserRoomPrivateKeyStmt, insertUserRoomPrivateKeySQL},
+		{&s.insertUserRoomPublicKeyStmt, insertUserRoomPublicKeySQL},
+		{&s.selectUserRoomKeyStmt, selectUserRoomKeySQL},
+		{&s.selectUserNIDsStmt, selectUserNIDsSQL},
+	}.Prepare(db)
+}
+
+func (s *userRoomKeysStatements) InsertUserRoomPrivatePublicKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PrivateKey) (result ed25519.PrivateKey, err error) {
+	stmt := sqlutil.TxStmtContext(ctx, txn, s.insertUserRoomPrivateKeyStmt)
+	err = stmt.QueryRowContext(ctx, userNID, roomNID, key, key.Public()).Scan(&result)
+	return result, err
+}
+
+func (s *userRoomKeysStatements) InsertUserRoomPublicKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PublicKey) (result ed25519.PublicKey, err error) {
+	stmt := sqlutil.TxStmtContext(ctx, txn, s.insertUserRoomPublicKeyStmt)
+	err = stmt.QueryRowContext(ctx, userNID, roomNID, key).Scan(&result)
+	return result, err
+}
+
+func (s *userRoomKeysStatements) SelectUserRoomPrivateKey(
+	ctx context.Context,
+	txn *sql.Tx,
+	userNID types.EventStateKeyNID,
+	roomNID types.RoomNID,
+) (ed25519.PrivateKey, error) {
+	stmt := sqlutil.TxStmtContext(ctx, txn, s.selectUserRoomKeyStmt)
+	var result ed25519.PrivateKey
+	err := stmt.QueryRowContext(ctx, userNID, roomNID).Scan(&result)
+	if errors.Is(err, sql.ErrNoRows) {
+		return nil, nil
+	}
+	return result, err
+}
+
+func (s *userRoomKeysStatements) BulkSelectUserNIDs(ctx context.Context, txn *sql.Tx, senderKeys map[types.RoomNID][]ed25519.PublicKey) (map[string]types.UserRoomKeyPair, error) {
+	stmt := sqlutil.TxStmtContext(ctx, txn, s.selectUserNIDsStmt)
+
+	roomNIDs := make([]types.RoomNID, 0, len(senderKeys))
+	var senders [][]byte
+	for roomNID := range senderKeys {
+		roomNIDs = append(roomNIDs, roomNID)
+		for _, key := range senderKeys[roomNID] {
+			senders = append(senders, key)
+		}
+	}
+	rows, err := stmt.QueryContext(ctx, pq.Array(roomNIDs), pq.Array(senders))
+	if err != nil {
+		return nil, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "failed to close rows")
+
+	result := make(map[string]types.UserRoomKeyPair, len(senders)+len(roomNIDs))
+	var publicKey []byte
+	userRoomKeyPair := types.UserRoomKeyPair{}
+	for rows.Next() {
+		if err = rows.Scan(&userRoomKeyPair.EventStateKeyNID, &userRoomKeyPair.RoomNID, &publicKey); err != nil {
+			return nil, err
+		}
+		result[string(publicKey)] = userRoomKeyPair
+	}
+	return result, rows.Err()
+}

roomserver/storage/shared/storage.go

@@ -2,14 +2,18 @@ package shared
 
 import (
 	"context"
+	"crypto/ed25519"
 	"database/sql"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"sort"
 
+	"github.com/matrix-org/dendrite/internal/eventutil"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/gomatrixserverlib/spec"
 	"github.com/matrix-org/util"
+	"github.com/sirupsen/logrus"
 	"github.com/tidwall/gjson"
 
 	"github.com/matrix-org/dendrite/internal/caching"
@@ -41,6 +45,7 @@ type Database struct {
 	MembershipTable    tables.Membership
 	PublishedTable     tables.Published
 	Purge              tables.Purge
+	UserRoomKeyTable   tables.UserRoomKeys
 	GetRoomUpdaterFn   func(ctx context.Context, roomInfo *types.RoomInfo) (*RoomUpdater, error)
 }
 
@@ -1609,6 +1614,147 @@ func (d *Database) UpgradeRoom(ctx context.Context, oldRoomID, newRoomID, eventS
 	})
 }
 
+// InsertUserRoomPrivatePublicKey inserts a new user room key for the given user and room.
+// Returns the newly inserted private key or an existing private key. If there is
+// an error talking to the database, returns that error.
+func (d *Database) InsertUserRoomPrivatePublicKey(ctx context.Context, userID spec.UserID, roomID spec.RoomID, key ed25519.PrivateKey) (result ed25519.PrivateKey, err error) {
+	uID := userID.String()
+	stateKeyNIDMap, sErr := d.eventStateKeyNIDs(ctx, nil, []string{uID})
+	if sErr != nil {
+		return nil, sErr
+	}
+	stateKeyNID := stateKeyNIDMap[uID]
+
+	err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
+		roomInfo, rErr := d.roomInfo(ctx, txn, roomID.String())
+		if rErr != nil {
+			return rErr
+		}
+		if roomInfo == nil {
+			return eventutil.ErrRoomNoExists{}
+		}
+
+		var iErr error
+		result, iErr = d.UserRoomKeyTable.InsertUserRoomPrivatePublicKey(ctx, txn, stateKeyNID, roomInfo.RoomNID, key)
+		return iErr
+	})
+	return result, err
+}
+
+// InsertUserRoomPublicKey inserts a new user room key for the given user and room.
+// Returns the newly inserted public key or an existing public key. If there is
+// an error talking to the database, returns that error.
+func (d *Database) InsertUserRoomPublicKey(ctx context.Context, userID spec.UserID, roomID spec.RoomID, key ed25519.PublicKey) (result ed25519.PublicKey, err error) {
+	uID := userID.String()
+	stateKeyNIDMap, sErr := d.eventStateKeyNIDs(ctx, nil, []string{uID})
+	if sErr != nil {
+		return nil, sErr
+	}
+	stateKeyNID := stateKeyNIDMap[uID]
+
+	err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
+		roomInfo, rErr := d.roomInfo(ctx, txn, roomID.String())
+		if rErr != nil {
+			return rErr
+		}
+		if roomInfo == nil {
+			return eventutil.ErrRoomNoExists{}
+		}
+
+		var iErr error
+		result, iErr = d.UserRoomKeyTable.InsertUserRoomPublicKey(ctx, txn, stateKeyNID, roomInfo.RoomNID, key)
+		return iErr
+	})
+	return result, err
+}
+
+// SelectUserRoomPrivateKey queries the users room private key.
+// If no key exists, returns no key and no error. Otherwise returns
+// the key and a database error, if any.
+func (d *Database) SelectUserRoomPrivateKey(ctx context.Context, userID spec.UserID, roomID spec.RoomID) (key ed25519.PrivateKey, err error) {
+	uID := userID.String()
+	stateKeyNIDMap, sErr := d.eventStateKeyNIDs(ctx, nil, []string{uID})
+	if sErr != nil {
+		return nil, sErr
+	}
+	stateKeyNID := stateKeyNIDMap[uID]
+
+	err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
+		roomInfo, rErr := d.roomInfo(ctx, txn, roomID.String())
+		if rErr != nil {
+			return rErr
+		}
+		if roomInfo == nil {
+			return nil
+		}
+
+		key, sErr = d.UserRoomKeyTable.SelectUserRoomPrivateKey(ctx, txn, stateKeyNID, roomInfo.RoomNID)
+		if !errors.Is(sErr, sql.ErrNoRows) {
+			return sErr
+		}
+		return nil
+	})
+	return
+}
+
+// SelectUserIDsForPublicKeys returns a map from roomID -> map from senderKey -> userID
+func (d *Database) SelectUserIDsForPublicKeys(ctx context.Context, publicKeys map[spec.RoomID][]ed25519.PublicKey) (result map[spec.RoomID]map[string]string, err error) {
+	result = make(map[spec.RoomID]map[string]string, len(publicKeys))
+	err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
+
+		// map all roomIDs to roomNIDs
+		query := make(map[types.RoomNID][]ed25519.PublicKey)
+		rooms := make(map[types.RoomNID]spec.RoomID)
+		for roomID, keys := range publicKeys {
+			roomNID, ok := d.Cache.GetRoomServerRoomNID(roomID.String())
+			if !ok {
+				roomInfo, rErr := d.roomInfo(ctx, txn, roomID.String())
+				if rErr != nil {
+					return rErr
+				}
+				if roomInfo == nil {
+					logrus.Warnf("missing room info for %s, there will be missing users in the response", roomID.String())
+					continue
+				}
+				roomNID = roomInfo.RoomNID
+			}
+
+			query[roomNID] = keys
+			rooms[roomNID] = roomID
+		}
+
+		// get the user room key pars
+		userRoomKeyPairMap, sErr := d.UserRoomKeyTable.BulkSelectUserNIDs(ctx, txn, query)
+		if sErr != nil {
+			return sErr
+		}
+		nids := make([]types.EventStateKeyNID, 0, len(userRoomKeyPairMap))
+		for _, nid := range userRoomKeyPairMap {
+			nids = append(nids, nid.EventStateKeyNID)
+		}
+		// get the userIDs
+		nidMap, seErr := d.EventStateKeys(ctx, nids)
+		if seErr != nil {
+			return seErr
+		}
+
+		// build the result map (roomID -> map publicKey -> userID)
+		for publicKey, userRoomKeyPair := range userRoomKeyPairMap {
+			userID := nidMap[userRoomKeyPair.EventStateKeyNID]
+			roomID := rooms[userRoomKeyPair.RoomNID]
+			resMap, exists := result[roomID]
+			if !exists {
+				resMap = map[string]string{}
+			}
+			resMap[publicKey] = userID
+			result[roomID] = resMap
+		}
+
+		return nil
+	})
+	return result, err
+}
+
 // FIXME TODO: Remove all this - horrible dupe with roomserver/state. Can't use the original impl because of circular loops
 // it should live in this package!
 

roomserver/storage/shared/storage_test.go

@@ -2,11 +2,15 @@ package shared_test
 
 import (
 	"context"
+	"crypto/ed25519"
 	"testing"
 	"time"
 
 	"github.com/matrix-org/dendrite/internal/caching"
+	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/matrix-org/gomatrixserverlib/spec"
 	"github.com/stretchr/testify/assert"
+	ed255192 "golang.org/x/crypto/ed25519"
 
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/roomserver/storage/postgres"
@@ -23,40 +27,61 @@ func mustCreateRoomserverDatabase(t *testing.T, dbType test.DBType) (*shared.Dat
 	connStr, clearDB := test.PrepareDBConnectionString(t, dbType)
 	dbOpts := &config.DatabaseOptions{ConnectionString: config.DataSource(connStr)}
 
-	db, err := sqlutil.Open(dbOpts, sqlutil.NewExclusiveWriter())
+	writer := sqlutil.NewExclusiveWriter()
+	db, err := sqlutil.Open(dbOpts, writer)
 	assert.NoError(t, err)
 
 	var membershipTable tables.Membership
 	var stateKeyTable tables.EventStateKeys
+	var userRoomKeys tables.UserRoomKeys
+	var roomsTable tables.Rooms
 	switch dbType {
 	case test.DBTypePostgres:
+		err = postgres.CreateRoomsTable(db)
+		assert.NoError(t, err)
 		err = postgres.CreateEventStateKeysTable(db)
 		assert.NoError(t, err)
 		err = postgres.CreateMembershipTable(db)
 		assert.NoError(t, err)
+		err = postgres.CreateUserRoomKeysTable(db)
+		assert.NoError(t, err)
+		roomsTable, err = postgres.PrepareRoomsTable(db)
+		assert.NoError(t, err)
 		membershipTable, err = postgres.PrepareMembershipTable(db)
 		assert.NoError(t, err)
 		stateKeyTable, err = postgres.PrepareEventStateKeysTable(db)
+		assert.NoError(t, err)
+		userRoomKeys, err = postgres.PrepareUserRoomKeysTable(db)
 	case test.DBTypeSQLite:
+		err = sqlite3.CreateRoomsTable(db)
+		assert.NoError(t, err)
 		err = sqlite3.CreateEventStateKeysTable(db)
 		assert.NoError(t, err)
 		err = sqlite3.CreateMembershipTable(db)
 		assert.NoError(t, err)
+		err = sqlite3.CreateUserRoomKeysTable(db)
+		assert.NoError(t, err)
+		roomsTable, err = sqlite3.PrepareRoomsTable(db)
+		assert.NoError(t, err)
 		membershipTable, err = sqlite3.PrepareMembershipTable(db)
 		assert.NoError(t, err)
 		stateKeyTable, err = sqlite3.PrepareEventStateKeysTable(db)
+		assert.NoError(t, err)
+		userRoomKeys, err = sqlite3.PrepareUserRoomKeysTable(db)
 	}
 	assert.NoError(t, err)
 
 	cache := caching.NewRistrettoCache(8*1024*1024, time.Hour, false)
 
-	evDb := shared.EventDatabase{EventStateKeysTable: stateKeyTable, Cache: cache}
+	evDb := shared.EventDatabase{EventStateKeysTable: stateKeyTable, Cache: cache, Writer: writer}
 
 	return &shared.Database{
 			DB:               db,
 			EventDatabase:    evDb,
 			MembershipTable:  membershipTable,
-			Writer:          sqlutil.NewExclusiveWriter(),
+			UserRoomKeyTable: userRoomKeys,
+			RoomsTable:       roomsTable,
+			Writer:           writer,
 			Cache:            cache,
 		}, func() {
 			clearDB()
@@ -97,3 +122,80 @@ func Test_GetLeftUsers(t *testing.T) {
 		assert.ElementsMatch(t, expectedUserIDs, leftUsers)
 	})
 }
+
+func TestUserRoomKeys(t *testing.T) {
+	ctx := context.Background()
+	alice := test.NewUser(t)
+	room := test.NewRoom(t, alice)
+
+	userID, err := spec.NewUserID(alice.ID, true)
+	assert.NoError(t, err)
+	roomID, err := spec.NewRoomID(room.ID)
+	assert.NoError(t, err)
+
+	test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
+		db, close := mustCreateRoomserverDatabase(t, dbType)
+		defer close()
+
+		// create a room NID so we can query the room
+		_, err = db.RoomsTable.InsertRoomNID(ctx, nil, roomID.String(), gomatrixserverlib.RoomVersionV10)
+		assert.NoError(t, err)
+		doesNotExist, err := spec.NewRoomID("!doesnotexist:localhost")
+		assert.NoError(t, err)
+		_, err = db.RoomsTable.InsertRoomNID(ctx, nil, doesNotExist.String(), gomatrixserverlib.RoomVersionV10)
+		assert.NoError(t, err)
+
+		_, key, err := ed25519.GenerateKey(nil)
+		assert.NoError(t, err)
+
+		gotKey, err := db.InsertUserRoomPrivatePublicKey(ctx, *userID, *roomID, 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 = db.InsertUserRoomPrivatePublicKey(ctx, *userID, *roomID, key2)
+		assert.NoError(t, err)
+		assert.Equal(t, gotKey, key)
+
+		gotKey, err = db.SelectUserRoomPrivateKey(context.Background(), *userID, *roomID)
+		assert.NoError(t, err)
+		assert.Equal(t, key, gotKey)
+
+		// Key doesn't exist, we shouldn't get anything back
+		assert.NoError(t, err)
+		gotKey, err = db.SelectUserRoomPrivateKey(context.Background(), *userID, *doesNotExist)
+		assert.NoError(t, err)
+		assert.Nil(t, gotKey)
+
+		queryUserIDs := map[spec.RoomID][]ed25519.PublicKey{
+			*roomID: {key.Public().(ed25519.PublicKey)},
+		}
+
+		userIDs, err := db.SelectUserIDsForPublicKeys(ctx, queryUserIDs)
+		assert.NoError(t, err)
+		wantKeys := map[spec.RoomID]map[string]string{
+			*roomID: {
+				string(key.Public().(ed25519.PublicKey)): userID.String(),
+			},
+		}
+		assert.Equal(t, wantKeys, userIDs)
+
+		// insert key that came in over federation
+		var gotPublicKey, key4 ed255192.PublicKey
+		key4, _, err = ed25519.GenerateKey(nil)
+		assert.NoError(t, err)
+		gotPublicKey, err = db.InsertUserRoomPublicKey(context.Background(), *userID, *doesNotExist, key4)
+		assert.NoError(t, err)
+		assert.Equal(t, key4, gotPublicKey)
+
+		// test invalid room
+		reallyDoesNotExist, err := spec.NewRoomID("!reallydoesnotexist:localhost")
+		assert.NoError(t, err)
+		_, err = db.InsertUserRoomPublicKey(context.Background(), *userID, *reallyDoesNotExist, key4)
+		assert.Error(t, err)
+		_, err = db.InsertUserRoomPrivatePublicKey(context.Background(), *userID, *reallyDoesNotExist, key)
+		assert.Error(t, err)
+	})
+}

roomserver/storage/sqlite3/storage.go

@@ -138,6 +138,9 @@ func (d *Database) create(db *sql.DB) error {
 	if err := CreateRedactionsTable(db); err != nil {
 		return err
 	}
+	if err := CreateUserRoomKeysTable(db); err != nil {
+		return err
+	}
 
 	return nil
 }
@@ -199,6 +202,10 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room
 	if err != nil {
 		return err
 	}
+	userRoomKeys, err := PrepareUserRoomKeysTable(db)
+	if err != nil {
+		return err
+	}
 
 	d.Database = shared.Database{
 		DB: db,
@@ -224,6 +231,7 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room
 		PublishedTable:     published,
 		GetRoomUpdaterFn:   d.GetRoomUpdater,
 		Purge:              purge,
+		UserRoomKeyTable:   userRoomKeys,
 	}
 	return nil
 }

roomserver/storage/sqlite3/user_room_keys_table.go

@@ -0,0 +1,146 @@
+// Copyright 2023 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"
+	"crypto/ed25519"
+	"database/sql"
+	"errors"
+	"strings"
+
+	"github.com/matrix-org/dendrite/internal"
+	"github.com/matrix-org/dendrite/internal/sqlutil"
+	"github.com/matrix-org/dendrite/roomserver/storage/tables"
+	"github.com/matrix-org/dendrite/roomserver/types"
+)
+
+const userRoomKeysSchema = `
+CREATE TABLE IF NOT EXISTS roomserver_user_room_keys (     
+    user_nid    INTEGER NOT NULL,
+    room_nid    INTEGER NOT NULL,
+    pseudo_id_key TEXT NULL, -- may be null for users not local to the server
+    pseudo_id_pub_key TEXT NOT NULL,
+    CONSTRAINT roomserver_user_room_keys_pk PRIMARY KEY (user_nid, room_nid)
+);
+`
+
+const insertUserRoomKeySQL = `
+	INSERT INTO roomserver_user_room_keys (user_nid, room_nid, pseudo_id_key, pseudo_id_pub_key) VALUES ($1, $2, $3, $4)
+	ON CONFLICT DO UPDATE SET pseudo_id_key = roomserver_user_room_keys.pseudo_id_key
+	RETURNING (pseudo_id_key)
+`
+
+const insertUserRoomPublicKeySQL = `
+	INSERT INTO roomserver_user_room_keys (user_nid, room_nid, pseudo_id_pub_key) VALUES ($1, $2, $3)
+	ON CONFLICT DO UPDATE SET pseudo_id_pub_key = $3
+	RETURNING (pseudo_id_pub_key)
+`
+
+const selectUserRoomKeySQL = `SELECT pseudo_id_key FROM roomserver_user_room_keys WHERE user_nid = $1 AND room_nid = $2`
+
+const selectUserNIDsSQL = `SELECT user_nid, room_nid, pseudo_id_pub_key FROM roomserver_user_room_keys WHERE room_nid IN ($1) AND pseudo_id_pub_key IN ($2)`
+
+type userRoomKeysStatements struct {
+	insertUserRoomPrivateKeyStmt *sql.Stmt
+	insertUserRoomPublicKeyStmt  *sql.Stmt
+	selectUserRoomKeyStmt        *sql.Stmt
+	//selectUserNIDsStmt           *sql.Stmt //prepared at runtime
+}
+
+func CreateUserRoomKeysTable(db *sql.DB) error {
+	_, err := db.Exec(userRoomKeysSchema)
+	return err
+}
+
+func PrepareUserRoomKeysTable(db *sql.DB) (tables.UserRoomKeys, error) {
+	s := &userRoomKeysStatements{}
+	return s, sqlutil.StatementList{
+		{&s.insertUserRoomPrivateKeyStmt, insertUserRoomKeySQL},
+		{&s.insertUserRoomPublicKeyStmt, insertUserRoomPublicKeySQL},
+		{&s.selectUserRoomKeyStmt, selectUserRoomKeySQL},
+		//{&s.selectUserNIDsStmt, selectUserNIDsSQL}, //prepared at runtime
+	}.Prepare(db)
+}
+
+func (s *userRoomKeysStatements) InsertUserRoomPrivatePublicKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PrivateKey) (result ed25519.PrivateKey, err error) {
+	stmt := sqlutil.TxStmtContext(ctx, txn, s.insertUserRoomPrivateKeyStmt)
+	err = stmt.QueryRowContext(ctx, userNID, roomNID, key, key.Public()).Scan(&result)
+	return result, err
+}
+
+func (s *userRoomKeysStatements) InsertUserRoomPublicKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PublicKey) (result ed25519.PublicKey, err error) {
+	stmt := sqlutil.TxStmtContext(ctx, txn, s.insertUserRoomPublicKeyStmt)
+	err = stmt.QueryRowContext(ctx, userNID, roomNID, key).Scan(&result)
+	return result, err
+}
+
+func (s *userRoomKeysStatements) SelectUserRoomPrivateKey(
+	ctx context.Context,
+	txn *sql.Tx,
+	userNID types.EventStateKeyNID,
+	roomNID types.RoomNID,
+) (ed25519.PrivateKey, error) {
+	stmt := sqlutil.TxStmtContext(ctx, txn, s.selectUserRoomKeyStmt)
+	var result ed25519.PrivateKey
+	err := stmt.QueryRowContext(ctx, userNID, roomNID).Scan(&result)
+	if errors.Is(err, sql.ErrNoRows) {
+		return nil, nil
+	}
+	return result, err
+}
+
+func (s *userRoomKeysStatements) BulkSelectUserNIDs(ctx context.Context, txn *sql.Tx, senderKeys map[types.RoomNID][]ed25519.PublicKey) (map[string]types.UserRoomKeyPair, error) {
+
+	roomNIDs := make([]any, 0, len(senderKeys))
+	var senders []any
+	for roomNID := range senderKeys {
+		roomNIDs = append(roomNIDs, roomNID)
+
+		for _, key := range senderKeys[roomNID] {
+			senders = append(senders, []byte(key))
+		}
+	}
+
+	selectSQL := strings.Replace(selectUserNIDsSQL, "($2)", sqlutil.QueryVariadicOffset(len(senders), len(senderKeys)), 1)
+	selectSQL = strings.Replace(selectSQL, "($1)", sqlutil.QueryVariadic(len(senderKeys)), 1) // replace $1 with the roomNIDs
+
+	selectStmt, err := txn.Prepare(selectSQL)
+	if err != nil {
+		return nil, err
+	}
+
+	params := append(roomNIDs, senders...)
+
+	stmt := sqlutil.TxStmt(txn, selectStmt)
+	defer internal.CloseAndLogIfError(ctx, stmt, "failed to close statement")
+
+	rows, err := stmt.QueryContext(ctx, params...)
+	if err != nil {
+		return nil, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "failed to close rows")
+
+	result := make(map[string]types.UserRoomKeyPair, len(params))
+	var publicKey []byte
+	userRoomKeyPair := types.UserRoomKeyPair{}
+	for rows.Next() {
+		if err = rows.Scan(&userRoomKeyPair.EventStateKeyNID, &userRoomKeyPair.RoomNID, &publicKey); err != nil {
+			return nil, err
+		}
+		result[string(publicKey)] = userRoomKeyPair
+	}
+	return result, rows.Err()
+}

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)
+
+	})
+}

roomserver/types/types.go

@@ -44,6 +44,11 @@ type EventMetadata struct {
 	RoomNID  RoomNID
 }
 
+type UserRoomKeyPair struct {
+	RoomNID          RoomNID
+	EventStateKeyNID EventStateKeyNID
+}
+
 // StateSnapshotNID is a numeric ID for the state at an event.
 type StateSnapshotNID int64