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

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!