[pseudoID] More pseudo ID fixes (#3167)

Signed-off-by: `Sam Wedgwood <sam@wedgwood.dev>`

5 files changed, 124 insertions, 16 deletions

diff --git a/roomserver/storage/interface.go b/roomserver/storage/interface.go
index e9b4609e..0638252b 100644
--- a/roomserver/storage/interface.go
+++ b/roomserver/storage/interface.go
@@ -158,7 +158,7 @@ type Database interface {
 	GetStateEvent(ctx context.Context, roomID, evType, stateKey string) (*types.HeaderedEvent, error)
 	GetStateEventsWithEventType(ctx context.Context, roomID, evType string) ([]*types.HeaderedEvent, error)
 	// GetRoomsByMembership returns a list of room IDs matching the provided membership and user ID (as state_key).
-	GetRoomsByMembership(ctx context.Context, userID, membership string) ([]string, error)
+	GetRoomsByMembership(ctx context.Context, userID spec.UserID, membership string) ([]string, error)
 	// GetBulkStateContent returns all state events which match a given room ID and a given state key tuple. Both must be satisfied for a match.
 	// If a tuple has the StateKey of '*' and allowWildcards=true then all state events with the EventType should be returned.
 	GetBulkStateContent(ctx context.Context, roomIDs []string, tuples []gomatrixserverlib.StateKeyTuple, allowWildcards bool) ([]tables.StrippedEvent, error)
diff --git a/roomserver/storage/postgres/user_room_keys_table.go b/roomserver/storage/postgres/user_room_keys_table.go
index 202b0abc..217ee957 100644
--- a/roomserver/storage/postgres/user_room_keys_table.go
+++ b/roomserver/storage/postgres/user_room_keys_table.go
@@ -56,12 +56,15 @@ 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 = ANY($1) AND pseudo_id_pub_key = ANY($2)`
 
+const selectAllUserRoomPublicKeyForUserSQL = `SELECT room_nid, pseudo_id_pub_key FROM roomserver_user_room_keys WHERE user_nid = $1`
+
 type userRoomKeysStatements struct {
-	insertUserRoomPrivateKeyStmt *sql.Stmt
-	insertUserRoomPublicKeyStmt  *sql.Stmt
-	selectUserRoomKeyStmt        *sql.Stmt
-	selectUserRoomPublicKeyStmt  *sql.Stmt
-	selectUserNIDsStmt           *sql.Stmt
+	insertUserRoomPrivateKeyStmt       *sql.Stmt
+	insertUserRoomPublicKeyStmt        *sql.Stmt
+	selectUserRoomKeyStmt              *sql.Stmt
+	selectUserRoomPublicKeyStmt        *sql.Stmt
+	selectUserNIDsStmt                 *sql.Stmt
+	selectAllUserRoomPublicKeysForUser *sql.Stmt
 }
 
 func CreateUserRoomKeysTable(db *sql.DB) error {
@@ -77,6 +80,7 @@ func PrepareUserRoomKeysTable(db *sql.DB) (tables.UserRoomKeys, error) {
 		{&s.selectUserRoomKeyStmt, selectUserRoomKeySQL},
 		{&s.selectUserRoomPublicKeyStmt, selectUserRoomPublicKeySQL},
 		{&s.selectUserNIDsStmt, selectUserNIDsSQL},
+		{&s.selectAllUserRoomPublicKeysForUser, selectAllUserRoomPublicKeyForUserSQL},
 	}.Prepare(db)
 }
 
@@ -150,3 +154,24 @@ func (s *userRoomKeysStatements) BulkSelectUserNIDs(ctx context.Context, txn *sq
 	}
 	return result, rows.Err()
 }
+
+func (s *userRoomKeysStatements) SelectAllPublicKeysForUser(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID) (map[types.RoomNID]ed25519.PublicKey, error) {
+	stmt := sqlutil.TxStmtContext(ctx, txn, s.selectAllUserRoomPublicKeysForUser)
+
+	rows, err := stmt.QueryContext(ctx, userNID)
+	if errors.Is(err, sql.ErrNoRows) {
+		return nil, nil
+	}
+
+	resultMap := make(map[types.RoomNID]ed25519.PublicKey)
+
+	var roomNID types.RoomNID
+	var pubkey ed25519.PublicKey
+	for rows.Next() {
+		if err = rows.Scan(&roomNID, &pubkey); err != nil {
+			return nil, err
+		}
+		resultMap[roomNID] = pubkey
+	}
+	return resultMap, err
+}
diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index 3c8b69c3..b09c5afb 100644
--- a/roomserver/storage/shared/storage.go
+++ b/roomserver/storage/shared/storage.go
@@ -1347,7 +1347,7 @@ func (d *Database) GetStateEventsWithEventType(ctx context.Context, roomID, evTy
 }
 
 // GetRoomsByMembership returns a list of room IDs matching the provided membership and user ID (as state_key).
-func (d *Database) GetRoomsByMembership(ctx context.Context, userID, membership string) ([]string, error) {
+func (d *Database) GetRoomsByMembership(ctx context.Context, userID spec.UserID, membership string) ([]string, error) {
 	var membershipState tables.MembershipState
 	switch membership {
 	case "join":
@@ -1361,17 +1361,73 @@ func (d *Database) GetRoomsByMembership(ctx context.Context, userID, membership
 	default:
 		return nil, fmt.Errorf("GetRoomsByMembership: invalid membership %s", membership)
 	}
-	stateKeyNID, err := d.EventStateKeysTable.SelectEventStateKeyNID(ctx, nil, userID)
+
+	// Convert provided user ID to NID
+	userNID, err := d.EventStateKeysTable.SelectEventStateKeyNID(ctx, nil, userID.String())
 	if err != nil {
 		if err == sql.ErrNoRows {
 			return nil, nil
+		} else {
+			return nil, fmt.Errorf("SelectEventStateKeyNID: cannot map user ID to state key NIDs: %w", err)
 		}
-		return nil, fmt.Errorf("GetRoomsByMembership: cannot map user ID to state key NID: %w", err)
 	}
-	roomNIDs, err := d.MembershipTable.SelectRoomsWithMembership(ctx, nil, stateKeyNID, membershipState)
+
+	// Use this NID to fetch all associated room keys (for pseudo ID rooms)
+	roomKeyMap, err := d.UserRoomKeyTable.SelectAllPublicKeysForUser(ctx, nil, userNID)
 	if err != nil {
-		return nil, fmt.Errorf("GetRoomsByMembership: failed to SelectRoomsWithMembership: %w", err)
+		if err == sql.ErrNoRows {
+			roomKeyMap = map[types.RoomNID]ed25519.PublicKey{}
+		} else {
+			return nil, fmt.Errorf("SelectAllPublicKeysForUser: could not select user room public keys for user: %w", err)
+		}
 	}
+
+	var eventStateKeyNIDs []types.EventStateKeyNID
+
+	// If there are room keys (i.e. this user is in pseudo ID rooms), then gather the appropriate NIDs
+	if len(roomKeyMap) != 0 {
+		// Convert keys to string representation
+		userRoomKeys := make([]string, len(roomKeyMap))
+		i := 0
+		for _, key := range roomKeyMap {
+			userRoomKeys[i] = spec.Base64Bytes(key).Encode()
+			i += 1
+		}
+
+		// Convert the string representation to its NID
+		pseudoIDStateKeys, sqlErr := d.EventStateKeysTable.BulkSelectEventStateKeyNID(ctx, nil, userRoomKeys)
+		if sqlErr != nil {
+			if sqlErr == sql.ErrNoRows {
+				pseudoIDStateKeys = map[string]types.EventStateKeyNID{}
+			} else {
+				return nil, fmt.Errorf("BulkSelectEventStateKeyNID: could not select state keys for public room keys: %w", err)
+			}
+		}
+
+		// Collect all NIDs together
+		eventStateKeyNIDs = make([]types.EventStateKeyNID, len(pseudoIDStateKeys)+1)
+		eventStateKeyNIDs[0] = userNID
+		i = 1
+		for _, nid := range pseudoIDStateKeys {
+			eventStateKeyNIDs[i] = nid
+			i += 1
+		}
+	} else {
+		// If there are no room keys (so no pseudo ID rooms), we only need to care about the user ID NID.
+		eventStateKeyNIDs = []types.EventStateKeyNID{userNID}
+	}
+
+	// Fetch rooms that match membership for each NID
+	roomNIDs := []types.RoomNID{}
+	for _, nid := range eventStateKeyNIDs {
+		var roomNIDsChunk []types.RoomNID
+		roomNIDsChunk, err = d.MembershipTable.SelectRoomsWithMembership(ctx, nil, nid, membershipState)
+		if err != nil {
+			return nil, fmt.Errorf("GetRoomsByMembership: failed to SelectRoomsWithMembership: %w", err)
+		}
+		roomNIDs = append(roomNIDs, roomNIDsChunk...)
+	}
+
 	roomIDs, err := d.RoomsTable.BulkSelectRoomIDs(ctx, nil, roomNIDs)
 	if err != nil {
 		return nil, fmt.Errorf("GetRoomsByMembership: failed to lookup room nids: %w", err)
diff --git a/roomserver/storage/sqlite3/user_room_keys_table.go b/roomserver/storage/sqlite3/user_room_keys_table.go
index 5d6ddc9a..434bad29 100644
--- a/roomserver/storage/sqlite3/user_room_keys_table.go
+++ b/roomserver/storage/sqlite3/user_room_keys_table.go
@@ -56,12 +56,15 @@ 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)`
 
+const selectAllUserRoomPublicKeyForUserSQL = `SELECT room_nid, pseudo_id_pub_key FROM roomserver_user_room_keys WHERE user_nid = $1`
+
 type userRoomKeysStatements struct {
-	db                           *sql.DB
-	insertUserRoomPrivateKeyStmt *sql.Stmt
-	insertUserRoomPublicKeyStmt  *sql.Stmt
-	selectUserRoomKeyStmt        *sql.Stmt
-	selectUserRoomPublicKeyStmt  *sql.Stmt
+	db                                 *sql.DB
+	insertUserRoomPrivateKeyStmt       *sql.Stmt
+	insertUserRoomPublicKeyStmt        *sql.Stmt
+	selectUserRoomKeyStmt              *sql.Stmt
+	selectUserRoomPublicKeyStmt        *sql.Stmt
+	selectAllUserRoomPublicKeysForUser *sql.Stmt
 	//selectUserNIDsStmt           *sql.Stmt //prepared at runtime
 }
 
@@ -77,6 +80,7 @@ func PrepareUserRoomKeysTable(db *sql.DB) (tables.UserRoomKeys, error) {
 		{&s.insertUserRoomPublicKeyStmt, insertUserRoomPublicKeySQL},
 		{&s.selectUserRoomKeyStmt, selectUserRoomKeySQL},
 		{&s.selectUserRoomPublicKeyStmt, selectUserRoomPublicKeySQL},
+		{&s.selectAllUserRoomPublicKeysForUser, selectAllUserRoomPublicKeyForUserSQL},
 		//{&s.selectUserNIDsStmt, selectUserNIDsSQL}, //prepared at runtime
 	}.Prepare(db)
 }
@@ -165,3 +169,24 @@ func (s *userRoomKeysStatements) BulkSelectUserNIDs(ctx context.Context, txn *sq
 	}
 	return result, rows.Err()
 }
+
+func (s *userRoomKeysStatements) SelectAllPublicKeysForUser(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID) (map[types.RoomNID]ed25519.PublicKey, error) {
+	stmt := sqlutil.TxStmtContext(ctx, txn, s.selectAllUserRoomPublicKeysForUser)
+
+	rows, err := stmt.QueryContext(ctx, userNID)
+	if errors.Is(err, sql.ErrNoRows) {
+		return nil, nil
+	}
+
+	resultMap := make(map[types.RoomNID]ed25519.PublicKey)
+
+	var roomNID types.RoomNID
+	var pubkey ed25519.PublicKey
+	for rows.Next() {
+		if err = rows.Scan(&roomNID, &pubkey); err != nil {
+			return nil, err
+		}
+		resultMap[roomNID] = pubkey
+	}
+	return resultMap, err
+}
diff --git a/roomserver/storage/tables/interface.go b/roomserver/storage/tables/interface.go
index 445c1223..0ae064e6 100644
--- a/roomserver/storage/tables/interface.go
+++ b/roomserver/storage/tables/interface.go
@@ -198,6 +198,8 @@ type UserRoomKeys interface {
 	// 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)
+	// SelectAllPublicKeysForUser returns all known public keys for a user. Returns a map from room NID -> public key
+	SelectAllPublicKeysForUser(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID) (map[types.RoomNID]ed25519.PublicKey, error)
 }
 
 // StrippedEvent represents a stripped event for returning extracted content values.