[pseudoID] More pseudo ID fixes (#3167)

Signed-off-by: `Sam Wedgwood <sam@wedgwood.dev>`
This commit is contained in:
Sam Wedgwood 2023-08-15 12:37:04 +01:00 committed by GitHub
parent fa6c7ba456
commit 9a12420428
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
24 changed files with 472 additions and 237 deletions

View file

@ -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)

View file

@ -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
}

View file

@ -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)

View file

@ -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
}

View file

@ -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.