Add `MXIDMapping` for pseudoID rooms (#3112)

Add `MXIDMapping` on membership events when
creating/joining rooms.

3 files changed, 58 insertions, 54 deletions

diff --git a/roomserver/storage/interface.go b/roomserver/storage/interface.go
index 7156c11c..e9b4609e 100644
--- a/roomserver/storage/interface.go
+++ b/roomserver/storage/interface.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"crypto/ed25519"
 
+	"github.com/matrix-org/dendrite/roomserver/api"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/gomatrixserverlib/spec"
 
@@ -190,7 +191,7 @@ type Database interface {
 	GetOrCreateEventTypeNID(ctx context.Context, eventType string) (eventTypeNID types.EventTypeNID, err error)
 	GetOrCreateEventStateKeyNID(ctx context.Context, eventStateKey *string) (types.EventStateKeyNID, error)
 	MaybeRedactEvent(
-		ctx context.Context, roomInfo *types.RoomInfo, eventNID types.EventNID, event gomatrixserverlib.PDU, plResolver state.PowerLevelResolver,
+		ctx context.Context, roomInfo *types.RoomInfo, eventNID types.EventNID, event gomatrixserverlib.PDU, plResolver state.PowerLevelResolver, querier api.QuerySenderIDAPI,
 	) (gomatrixserverlib.PDU, gomatrixserverlib.PDU, error)
 }
 
@@ -251,7 +252,7 @@ type EventDatabase interface {
 	// MaybeRedactEvent returns the redaction event and the redacted event if this call resulted in a redaction, else an error
 	// (nil if there was nothing to do)
 	MaybeRedactEvent(
-		ctx context.Context, roomInfo *types.RoomInfo, eventNID types.EventNID, event gomatrixserverlib.PDU, plResolver state.PowerLevelResolver,
+		ctx context.Context, roomInfo *types.RoomInfo, eventNID types.EventNID, event gomatrixserverlib.PDU, plResolver state.PowerLevelResolver, querier api.QuerySenderIDAPI,
 	) (gomatrixserverlib.PDU, gomatrixserverlib.PDU, error)
 	StoreEvent(ctx context.Context, event gomatrixserverlib.PDU, roomInfo *types.RoomInfo, eventTypeNID types.EventTypeNID, eventStateKeyNID types.EventStateKeyNID, authEventNIDs []types.EventNID, isRejected bool) (types.EventNID, types.StateAtEvent, error)
 }
diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index 61a3520a..fc3ace6a 100644
--- a/roomserver/storage/shared/storage.go
+++ b/roomserver/storage/shared/storage.go
@@ -10,6 +10,7 @@ import (
 	"sort"
 
 	"github.com/matrix-org/dendrite/internal/eventutil"
+	"github.com/matrix-org/dendrite/roomserver/api"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/gomatrixserverlib/spec"
 	"github.com/matrix-org/util"
@@ -991,6 +992,7 @@ func extractRoomVersionFromCreateEvent(event gomatrixserverlib.PDU) (
 // Returns the redaction event and the redacted event if this call resulted in a redaction.
 func (d *EventDatabase) MaybeRedactEvent(
 	ctx context.Context, roomInfo *types.RoomInfo, eventNID types.EventNID, event gomatrixserverlib.PDU, plResolver state.PowerLevelResolver,
+	querier api.QuerySenderIDAPI,
 ) (gomatrixserverlib.PDU, gomatrixserverlib.PDU, error) {
 	var (
 		redactionEvent, redactedEvent *types.Event
@@ -1030,15 +1032,18 @@ func (d *EventDatabase) MaybeRedactEvent(
 			return nil
 		}
 
-		// TODO: Don't hack senderID into userID here (pseudoIDs)
+		var validRoomID *spec.RoomID
+		validRoomID, err = spec.NewRoomID(redactedEvent.RoomID())
+		if err != nil {
+			return err
+		}
 		sender1Domain := ""
-		sender1, err1 := spec.NewUserID(string(redactedEvent.SenderID()), true)
+		sender1, err1 := querier.QueryUserIDForSender(ctx, *validRoomID, redactedEvent.SenderID())
 		if err1 == nil {
 			sender1Domain = string(sender1.Domain())
 		}
-		// TODO: Don't hack senderID into userID here (pseudoIDs)
 		sender2Domain := ""
-		sender2, err2 := spec.NewUserID(string(redactionEvent.SenderID()), true)
+		sender2, err2 := querier.QueryUserIDForSender(ctx, *validRoomID, redactionEvent.SenderID())
 		if err2 == nil {
 			sender2Domain = string(sender2.Domain())
 		}
@@ -1698,6 +1703,7 @@ func (d *Database) InsertUserRoomPublicKey(ctx context.Context, userID spec.User
 // 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.
+// TODO: Cache this?
 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})
@@ -1756,58 +1762,54 @@ func (d *Database) SelectUserRoomPublicKey(ctx context.Context, userID spec.User
 // 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
+	// 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, nil, roomID.String())
+			if rErr != nil {
+				return nil, rErr
 			}
-
-			query[roomNID] = keys
-			rooms[roomNID] = roomID
+			if roomInfo == nil {
+				logrus.Warnf("missing room info for %s, there will be missing users in the response", roomID.String())
+				continue
+			}
+			roomNID = roomInfo.RoomNID
 		}
 
-		// 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
-		}
+		query[roomNID] = keys
+		rooms[roomNID] = roomID
+	}
 
-		// 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
-		}
+	// get the user room key pars
+	userRoomKeyPairMap, sErr := d.UserRoomKeyTable.BulkSelectUserNIDs(ctx, nil, query)
+	if sErr != nil {
+		return nil, 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 nil, seErr
+	}
 
-		return nil
-	})
+	// 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 result, err
 }
 
diff --git a/roomserver/storage/sqlite3/user_room_keys_table.go b/roomserver/storage/sqlite3/user_room_keys_table.go
index d58b8ac3..5d6ddc9a 100644
--- a/roomserver/storage/sqlite3/user_room_keys_table.go
+++ b/roomserver/storage/sqlite3/user_room_keys_table.go
@@ -57,6 +57,7 @@ const selectUserRoomPublicKeySQL = `SELECT pseudo_id_pub_key FROM roomserver_use
 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 {
+	db                           *sql.DB
 	insertUserRoomPrivateKeyStmt *sql.Stmt
 	insertUserRoomPublicKeyStmt  *sql.Stmt
 	selectUserRoomKeyStmt        *sql.Stmt
@@ -70,7 +71,7 @@ func CreateUserRoomKeysTable(db *sql.DB) error {
 }
 
 func PrepareUserRoomKeysTable(db *sql.DB) (tables.UserRoomKeys, error) {
-	s := &userRoomKeysStatements{}
+	s := &userRoomKeysStatements{db: db}
 	return s, sqlutil.StatementList{
 		{&s.insertUserRoomPrivateKeyStmt, insertUserRoomKeySQL},
 		{&s.insertUserRoomPublicKeyStmt, insertUserRoomPublicKeySQL},
@@ -137,7 +138,7 @@ func (s *userRoomKeysStatements) BulkSelectUserNIDs(ctx context.Context, txn *sq
 	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)
+	selectStmt, err := s.db.Prepare(selectSQL)
 	if err != nil {
 		return nil, err
 	}