Merge SenderID & Per Room User Key work (#3109)

3 files changed, 51 insertions, 16 deletions

diff --git a/roomserver/storage/shared/room_updater.go b/roomserver/storage/shared/room_updater.go
index 6fb57332..70672a33 100644
--- a/roomserver/storage/shared/room_updater.go
+++ b/roomserver/storage/shared/room_updater.go
@@ -6,7 +6,6 @@ import (
 	"fmt"
 
 	"github.com/matrix-org/gomatrixserverlib"
-	"github.com/matrix-org/gomatrixserverlib/spec"
 
 	"github.com/matrix-org/dendrite/roomserver/types"
 )
@@ -251,7 +250,3 @@ func (u *RoomUpdater) MarkEventAsSent(eventNID types.EventNID) error {
 func (u *RoomUpdater) MembershipUpdater(targetUserNID types.EventStateKeyNID, targetLocal bool) (*MembershipUpdater, error) {
 	return u.d.membershipUpdaterTxn(u.ctx, u.txn, u.roomInfo.RoomNID, targetUserNID, targetLocal)
 }
-
-func (u *RoomUpdater) GetUserIDForSender(ctx context.Context, roomID string, senderID spec.SenderID) (*spec.UserID, error) {
-	return u.d.GetUserIDForSender(ctx, roomID, senderID)
-}
diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index bda51da8..61a3520a 100644
--- a/roomserver/storage/shared/storage.go
+++ b/roomserver/storage/shared/storage.go
@@ -721,6 +721,22 @@ func (d *Database) GetOrCreateRoomInfo(ctx context.Context, event gomatrixserver
 	}, err
 }
 
+func (d *Database) GetRoomVersion(ctx context.Context, roomID string) (gomatrixserverlib.RoomVersion, error) {
+	cachedRoomVersion, versionOK := d.Cache.GetRoomVersion(roomID)
+	if versionOK {
+		return cachedRoomVersion, nil
+	}
+
+	roomInfo, err := d.RoomInfo(ctx, roomID)
+	if err != nil {
+		return "", err
+	}
+	if roomInfo == nil {
+		return "", nil
+	}
+	return roomInfo.RoomVersion, nil
+}
+
 func (d *Database) GetOrCreateEventTypeNID(ctx context.Context, eventType string) (eventTypeNID types.EventTypeNID, err error) {
 	err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
 		if eventTypeNID, err = d.assignEventTypeNID(ctx, txn, eventType); err != nil {
@@ -1550,16 +1566,6 @@ func (d *Database) GetKnownUsers(ctx context.Context, userID, searchString strin
 	return d.MembershipTable.SelectKnownUsers(ctx, nil, stateKeyNID, searchString, limit)
 }
 
-func (d *Database) GetUserIDForSender(ctx context.Context, roomID string, senderID spec.SenderID) (*spec.UserID, error) {
-	// TODO: Use real logic once DB for pseudoIDs is in place
-	return spec.NewUserID(string(senderID), true)
-}
-
-func (d *Database) GetSenderIDForUser(ctx context.Context, roomID string, userID spec.UserID) (spec.SenderID, error) {
-	// TODO: Use real logic once DB for pseudoIDs is in place
-	return spec.SenderID(userID.String()), nil
-}
-
 // GetKnownRooms returns a list of all rooms we know about.
 func (d *Database) GetKnownRooms(ctx context.Context) ([]string, error) {
 	return d.RoomsTable.SelectRoomIDsWithEvents(ctx, nil)
@@ -1718,6 +1724,35 @@ func (d *Database) SelectUserRoomPrivateKey(ctx context.Context, userID spec.Use
 	return
 }
 
+// SelectUserRoomPublicKey queries the users room public key.
+// If no key exists, returns no key and no error. Otherwise returns
+// the key and a database error, if any.
+func (d *Database) SelectUserRoomPublicKey(ctx context.Context, userID spec.UserID, roomID spec.RoomID) (key 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 nil
+		}
+
+		key, sErr = d.UserRoomKeyTable.SelectUserRoomPublicKey(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))
diff --git a/roomserver/storage/shared/storage_test.go b/roomserver/storage/shared/storage_test.go
index 581d83ee..c7b915c7 100644
--- a/roomserver/storage/shared/storage_test.go
+++ b/roomserver/storage/shared/storage_test.go
@@ -163,12 +163,17 @@ func TestUserRoomKeys(t *testing.T) {
 		gotKey, err = db.SelectUserRoomPrivateKey(context.Background(), *userID, *roomID)
 		assert.NoError(t, err)
 		assert.Equal(t, key, gotKey)
+		pubKey, err := db.SelectUserRoomPublicKey(context.Background(), *userID, *roomID)
+		assert.NoError(t, err)
+		assert.Equal(t, key.Public(), pubKey)
 
 		// 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)
+		pubKey, err = db.SelectUserRoomPublicKey(context.Background(), *userID, *doesNotExist)
+		assert.NoError(t, err)
+		assert.Nil(t, pubKey)
 
 		queryUserIDs := map[spec.RoomID][]ed25519.PublicKey{
 			*roomID: {key.Public().(ed25519.PublicKey)},