mirror of
https://github.com/1f349/dendrite.git
synced 2024-11-08 18:16:59 +00:00
Add initial support for storing user room keys (#3098)
This commit is contained in:
parent
5713c5715c
commit
832ccc32f6
@ -16,6 +16,7 @@ package storage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ed25519"
|
||||
|
||||
"github.com/matrix-org/gomatrixserverlib"
|
||||
"github.com/matrix-org/gomatrixserverlib/spec"
|
||||
@ -27,6 +28,7 @@ import (
|
||||
)
|
||||
|
||||
type Database interface {
|
||||
UserRoomKeys
|
||||
// Do we support processing input events for more than one room at a time?
|
||||
SupportsConcurrentRoomInputs() bool
|
||||
// RoomInfo returns room information for the given room ID, or nil if there is no room.
|
||||
@ -194,8 +196,22 @@ type Database interface {
|
||||
) (gomatrixserverlib.PDU, gomatrixserverlib.PDU, error)
|
||||
}
|
||||
|
||||
type UserRoomKeys interface {
|
||||
// InsertUserRoomPrivatePublicKey inserts the given private key as well as the public key for it. This should be used
|
||||
// when creating keys locally.
|
||||
InsertUserRoomPrivatePublicKey(ctx context.Context, userID spec.UserID, roomID spec.RoomID, key ed25519.PrivateKey) (result ed25519.PrivateKey, err error)
|
||||
// InsertUserRoomPublicKey inserts the given public key, this should be used for users NOT local to this server
|
||||
InsertUserRoomPublicKey(ctx context.Context, userID spec.UserID, roomID spec.RoomID, key ed25519.PublicKey) (result ed25519.PublicKey, err error)
|
||||
// SelectUserRoomPrivateKey selects the private key for the given user and room combination
|
||||
SelectUserRoomPrivateKey(ctx context.Context, userID spec.UserID, roomID spec.RoomID) (key ed25519.PrivateKey, err error)
|
||||
// SelectUserIDsForPublicKeys selects all userIDs for the requested senderKeys. Returns a map from roomID -> map from publicKey to userID.
|
||||
// If a senderKey can't be found, it is omitted in the result.
|
||||
SelectUserIDsForPublicKeys(ctx context.Context, publicKeys map[spec.RoomID][]ed25519.PublicKey) (map[spec.RoomID]map[string]string, error)
|
||||
}
|
||||
|
||||
type RoomDatabase interface {
|
||||
EventDatabase
|
||||
UserRoomKeys
|
||||
// RoomInfo returns room information for the given room ID, or nil if there is no room.
|
||||
RoomInfo(ctx context.Context, roomID string) (*types.RoomInfo, error)
|
||||
RoomInfoByNID(ctx context.Context, roomNID types.RoomNID) (*types.RoomInfo, error)
|
||||
|
@ -131,6 +131,9 @@ func (d *Database) create(db *sql.DB) error {
|
||||
if err := CreateRedactionsTable(db); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := CreateUserRoomKeysTable(db); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@ -192,6 +195,11 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
userRoomKeys, err := PrepareUserRoomKeysTable(db)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
d.Database = shared.Database{
|
||||
DB: db,
|
||||
EventDatabase: shared.EventDatabase{
|
||||
@ -215,6 +223,7 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room
|
||||
MembershipTable: membership,
|
||||
PublishedTable: published,
|
||||
Purge: purge,
|
||||
UserRoomKeyTable: userRoomKeys,
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
132
roomserver/storage/postgres/user_room_keys_table.go
Normal file
132
roomserver/storage/postgres/user_room_keys_table.go
Normal file
@ -0,0 +1,132 @@
|
||||
// Copyright 2023 The Matrix.org Foundation C.I.C.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ed25519"
|
||||
"database/sql"
|
||||
"errors"
|
||||
|
||||
"github.com/lib/pq"
|
||||
"github.com/matrix-org/dendrite/internal"
|
||||
"github.com/matrix-org/dendrite/internal/sqlutil"
|
||||
"github.com/matrix-org/dendrite/roomserver/storage/tables"
|
||||
"github.com/matrix-org/dendrite/roomserver/types"
|
||||
)
|
||||
|
||||
const userRoomKeysSchema = `
|
||||
CREATE TABLE IF NOT EXISTS roomserver_user_room_keys (
|
||||
user_nid INTEGER NOT NULL,
|
||||
room_nid INTEGER NOT NULL,
|
||||
pseudo_id_key BYTEA NULL, -- may be null for users not local to the server
|
||||
pseudo_id_pub_key BYTEA NOT NULL,
|
||||
CONSTRAINT roomserver_user_room_keys_pk PRIMARY KEY (user_nid, room_nid)
|
||||
);
|
||||
`
|
||||
|
||||
const insertUserRoomPrivateKeySQL = `
|
||||
INSERT INTO roomserver_user_room_keys (user_nid, room_nid, pseudo_id_key, pseudo_id_pub_key) VALUES ($1, $2, $3, $4)
|
||||
ON CONFLICT ON CONSTRAINT roomserver_user_room_keys_pk DO UPDATE SET pseudo_id_key = roomserver_user_room_keys.pseudo_id_key
|
||||
RETURNING (pseudo_id_key)
|
||||
`
|
||||
|
||||
const insertUserRoomPublicKeySQL = `
|
||||
INSERT INTO roomserver_user_room_keys (user_nid, room_nid, pseudo_id_pub_key) VALUES ($1, $2, $3)
|
||||
ON CONFLICT ON CONSTRAINT roomserver_user_room_keys_pk DO UPDATE SET pseudo_id_pub_key = $3
|
||||
RETURNING (pseudo_id_pub_key)
|
||||
`
|
||||
|
||||
const selectUserRoomKeySQL = `SELECT pseudo_id_key FROM roomserver_user_room_keys WHERE user_nid = $1 AND room_nid = $2`
|
||||
|
||||
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)`
|
||||
|
||||
type userRoomKeysStatements struct {
|
||||
insertUserRoomPrivateKeyStmt *sql.Stmt
|
||||
insertUserRoomPublicKeyStmt *sql.Stmt
|
||||
selectUserRoomKeyStmt *sql.Stmt
|
||||
selectUserNIDsStmt *sql.Stmt
|
||||
}
|
||||
|
||||
func CreateUserRoomKeysTable(db *sql.DB) error {
|
||||
_, err := db.Exec(userRoomKeysSchema)
|
||||
return err
|
||||
}
|
||||
|
||||
func PrepareUserRoomKeysTable(db *sql.DB) (tables.UserRoomKeys, error) {
|
||||
s := &userRoomKeysStatements{}
|
||||
return s, sqlutil.StatementList{
|
||||
{&s.insertUserRoomPrivateKeyStmt, insertUserRoomPrivateKeySQL},
|
||||
{&s.insertUserRoomPublicKeyStmt, insertUserRoomPublicKeySQL},
|
||||
{&s.selectUserRoomKeyStmt, selectUserRoomKeySQL},
|
||||
{&s.selectUserNIDsStmt, selectUserNIDsSQL},
|
||||
}.Prepare(db)
|
||||
}
|
||||
|
||||
func (s *userRoomKeysStatements) InsertUserRoomPrivatePublicKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PrivateKey) (result ed25519.PrivateKey, err error) {
|
||||
stmt := sqlutil.TxStmtContext(ctx, txn, s.insertUserRoomPrivateKeyStmt)
|
||||
err = stmt.QueryRowContext(ctx, userNID, roomNID, key, key.Public()).Scan(&result)
|
||||
return result, err
|
||||
}
|
||||
|
||||
func (s *userRoomKeysStatements) InsertUserRoomPublicKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PublicKey) (result ed25519.PublicKey, err error) {
|
||||
stmt := sqlutil.TxStmtContext(ctx, txn, s.insertUserRoomPublicKeyStmt)
|
||||
err = stmt.QueryRowContext(ctx, userNID, roomNID, key).Scan(&result)
|
||||
return result, err
|
||||
}
|
||||
|
||||
func (s *userRoomKeysStatements) SelectUserRoomPrivateKey(
|
||||
ctx context.Context,
|
||||
txn *sql.Tx,
|
||||
userNID types.EventStateKeyNID,
|
||||
roomNID types.RoomNID,
|
||||
) (ed25519.PrivateKey, error) {
|
||||
stmt := sqlutil.TxStmtContext(ctx, txn, s.selectUserRoomKeyStmt)
|
||||
var result ed25519.PrivateKey
|
||||
err := stmt.QueryRowContext(ctx, userNID, roomNID).Scan(&result)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
return result, err
|
||||
}
|
||||
|
||||
func (s *userRoomKeysStatements) BulkSelectUserNIDs(ctx context.Context, txn *sql.Tx, senderKeys map[types.RoomNID][]ed25519.PublicKey) (map[string]types.UserRoomKeyPair, error) {
|
||||
stmt := sqlutil.TxStmtContext(ctx, txn, s.selectUserNIDsStmt)
|
||||
|
||||
roomNIDs := make([]types.RoomNID, 0, len(senderKeys))
|
||||
var senders [][]byte
|
||||
for roomNID := range senderKeys {
|
||||
roomNIDs = append(roomNIDs, roomNID)
|
||||
for _, key := range senderKeys[roomNID] {
|
||||
senders = append(senders, key)
|
||||
}
|
||||
}
|
||||
rows, err := stmt.QueryContext(ctx, pq.Array(roomNIDs), pq.Array(senders))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer internal.CloseAndLogIfError(ctx, rows, "failed to close rows")
|
||||
|
||||
result := make(map[string]types.UserRoomKeyPair, len(senders)+len(roomNIDs))
|
||||
var publicKey []byte
|
||||
userRoomKeyPair := types.UserRoomKeyPair{}
|
||||
for rows.Next() {
|
||||
if err = rows.Scan(&userRoomKeyPair.EventStateKeyNID, &userRoomKeyPair.RoomNID, &publicKey); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result[string(publicKey)] = userRoomKeyPair
|
||||
}
|
||||
return result, rows.Err()
|
||||
}
|
@ -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!
|
||||
|
||||
|
@ -2,11 +2,15 @@ package shared_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ed25519"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/matrix-org/dendrite/internal/caching"
|
||||
"github.com/matrix-org/gomatrixserverlib"
|
||||
"github.com/matrix-org/gomatrixserverlib/spec"
|
||||
"github.com/stretchr/testify/assert"
|
||||
ed255192 "golang.org/x/crypto/ed25519"
|
||||
|
||||
"github.com/matrix-org/dendrite/internal/sqlutil"
|
||||
"github.com/matrix-org/dendrite/roomserver/storage/postgres"
|
||||
@ -23,41 +27,62 @@ func mustCreateRoomserverDatabase(t *testing.T, dbType test.DBType) (*shared.Dat
|
||||
connStr, clearDB := test.PrepareDBConnectionString(t, dbType)
|
||||
dbOpts := &config.DatabaseOptions{ConnectionString: config.DataSource(connStr)}
|
||||
|
||||
db, err := sqlutil.Open(dbOpts, sqlutil.NewExclusiveWriter())
|
||||
writer := sqlutil.NewExclusiveWriter()
|
||||
db, err := sqlutil.Open(dbOpts, writer)
|
||||
assert.NoError(t, err)
|
||||
|
||||
var membershipTable tables.Membership
|
||||
var stateKeyTable tables.EventStateKeys
|
||||
var userRoomKeys tables.UserRoomKeys
|
||||
var roomsTable tables.Rooms
|
||||
switch dbType {
|
||||
case test.DBTypePostgres:
|
||||
err = postgres.CreateRoomsTable(db)
|
||||
assert.NoError(t, err)
|
||||
err = postgres.CreateEventStateKeysTable(db)
|
||||
assert.NoError(t, err)
|
||||
err = postgres.CreateMembershipTable(db)
|
||||
assert.NoError(t, err)
|
||||
err = postgres.CreateUserRoomKeysTable(db)
|
||||
assert.NoError(t, err)
|
||||
roomsTable, err = postgres.PrepareRoomsTable(db)
|
||||
assert.NoError(t, err)
|
||||
membershipTable, err = postgres.PrepareMembershipTable(db)
|
||||
assert.NoError(t, err)
|
||||
stateKeyTable, err = postgres.PrepareEventStateKeysTable(db)
|
||||
assert.NoError(t, err)
|
||||
userRoomKeys, err = postgres.PrepareUserRoomKeysTable(db)
|
||||
case test.DBTypeSQLite:
|
||||
err = sqlite3.CreateRoomsTable(db)
|
||||
assert.NoError(t, err)
|
||||
err = sqlite3.CreateEventStateKeysTable(db)
|
||||
assert.NoError(t, err)
|
||||
err = sqlite3.CreateMembershipTable(db)
|
||||
assert.NoError(t, err)
|
||||
err = sqlite3.CreateUserRoomKeysTable(db)
|
||||
assert.NoError(t, err)
|
||||
roomsTable, err = sqlite3.PrepareRoomsTable(db)
|
||||
assert.NoError(t, err)
|
||||
membershipTable, err = sqlite3.PrepareMembershipTable(db)
|
||||
assert.NoError(t, err)
|
||||
stateKeyTable, err = sqlite3.PrepareEventStateKeysTable(db)
|
||||
assert.NoError(t, err)
|
||||
userRoomKeys, err = sqlite3.PrepareUserRoomKeysTable(db)
|
||||
}
|
||||
assert.NoError(t, err)
|
||||
|
||||
cache := caching.NewRistrettoCache(8*1024*1024, time.Hour, false)
|
||||
|
||||
evDb := shared.EventDatabase{EventStateKeysTable: stateKeyTable, Cache: cache}
|
||||
evDb := shared.EventDatabase{EventStateKeysTable: stateKeyTable, Cache: cache, Writer: writer}
|
||||
|
||||
return &shared.Database{
|
||||
DB: db,
|
||||
EventDatabase: evDb,
|
||||
MembershipTable: membershipTable,
|
||||
Writer: sqlutil.NewExclusiveWriter(),
|
||||
Cache: cache,
|
||||
DB: db,
|
||||
EventDatabase: evDb,
|
||||
MembershipTable: membershipTable,
|
||||
UserRoomKeyTable: userRoomKeys,
|
||||
RoomsTable: roomsTable,
|
||||
Writer: writer,
|
||||
Cache: cache,
|
||||
}, func() {
|
||||
clearDB()
|
||||
err = db.Close()
|
||||
@ -97,3 +122,80 @@ func Test_GetLeftUsers(t *testing.T) {
|
||||
assert.ElementsMatch(t, expectedUserIDs, leftUsers)
|
||||
})
|
||||
}
|
||||
|
||||
func TestUserRoomKeys(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
alice := test.NewUser(t)
|
||||
room := test.NewRoom(t, alice)
|
||||
|
||||
userID, err := spec.NewUserID(alice.ID, true)
|
||||
assert.NoError(t, err)
|
||||
roomID, err := spec.NewRoomID(room.ID)
|
||||
assert.NoError(t, err)
|
||||
|
||||
test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
|
||||
db, close := mustCreateRoomserverDatabase(t, dbType)
|
||||
defer close()
|
||||
|
||||
// create a room NID so we can query the room
|
||||
_, err = db.RoomsTable.InsertRoomNID(ctx, nil, roomID.String(), gomatrixserverlib.RoomVersionV10)
|
||||
assert.NoError(t, err)
|
||||
doesNotExist, err := spec.NewRoomID("!doesnotexist:localhost")
|
||||
assert.NoError(t, err)
|
||||
_, err = db.RoomsTable.InsertRoomNID(ctx, nil, doesNotExist.String(), gomatrixserverlib.RoomVersionV10)
|
||||
assert.NoError(t, err)
|
||||
|
||||
_, key, err := ed25519.GenerateKey(nil)
|
||||
assert.NoError(t, err)
|
||||
|
||||
gotKey, err := db.InsertUserRoomPrivatePublicKey(ctx, *userID, *roomID, key)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, gotKey, key)
|
||||
|
||||
// again, this shouldn't result in an error, but return the existing key
|
||||
_, key2, err := ed25519.GenerateKey(nil)
|
||||
assert.NoError(t, err)
|
||||
gotKey, err = db.InsertUserRoomPrivatePublicKey(ctx, *userID, *roomID, key2)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, gotKey, key)
|
||||
|
||||
gotKey, err = db.SelectUserRoomPrivateKey(context.Background(), *userID, *roomID)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, key, gotKey)
|
||||
|
||||
// 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)
|
||||
|
||||
queryUserIDs := map[spec.RoomID][]ed25519.PublicKey{
|
||||
*roomID: {key.Public().(ed25519.PublicKey)},
|
||||
}
|
||||
|
||||
userIDs, err := db.SelectUserIDsForPublicKeys(ctx, queryUserIDs)
|
||||
assert.NoError(t, err)
|
||||
wantKeys := map[spec.RoomID]map[string]string{
|
||||
*roomID: {
|
||||
string(key.Public().(ed25519.PublicKey)): userID.String(),
|
||||
},
|
||||
}
|
||||
assert.Equal(t, wantKeys, userIDs)
|
||||
|
||||
// insert key that came in over federation
|
||||
var gotPublicKey, key4 ed255192.PublicKey
|
||||
key4, _, err = ed25519.GenerateKey(nil)
|
||||
assert.NoError(t, err)
|
||||
gotPublicKey, err = db.InsertUserRoomPublicKey(context.Background(), *userID, *doesNotExist, key4)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, key4, gotPublicKey)
|
||||
|
||||
// test invalid room
|
||||
reallyDoesNotExist, err := spec.NewRoomID("!reallydoesnotexist:localhost")
|
||||
assert.NoError(t, err)
|
||||
_, err = db.InsertUserRoomPublicKey(context.Background(), *userID, *reallyDoesNotExist, key4)
|
||||
assert.Error(t, err)
|
||||
_, err = db.InsertUserRoomPrivatePublicKey(context.Background(), *userID, *reallyDoesNotExist, key)
|
||||
assert.Error(t, err)
|
||||
})
|
||||
}
|
||||
|
@ -138,6 +138,9 @@ func (d *Database) create(db *sql.DB) error {
|
||||
if err := CreateRedactionsTable(db); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := CreateUserRoomKeysTable(db); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@ -199,6 +202,10 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
userRoomKeys, err := PrepareUserRoomKeysTable(db)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
d.Database = shared.Database{
|
||||
DB: db,
|
||||
@ -224,6 +231,7 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room
|
||||
PublishedTable: published,
|
||||
GetRoomUpdaterFn: d.GetRoomUpdater,
|
||||
Purge: purge,
|
||||
UserRoomKeyTable: userRoomKeys,
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
146
roomserver/storage/sqlite3/user_room_keys_table.go
Normal file
146
roomserver/storage/sqlite3/user_room_keys_table.go
Normal file
@ -0,0 +1,146 @@
|
||||
// Copyright 2023 The Matrix.org Foundation C.I.C.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package sqlite3
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ed25519"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"strings"
|
||||
|
||||
"github.com/matrix-org/dendrite/internal"
|
||||
"github.com/matrix-org/dendrite/internal/sqlutil"
|
||||
"github.com/matrix-org/dendrite/roomserver/storage/tables"
|
||||
"github.com/matrix-org/dendrite/roomserver/types"
|
||||
)
|
||||
|
||||
const userRoomKeysSchema = `
|
||||
CREATE TABLE IF NOT EXISTS roomserver_user_room_keys (
|
||||
user_nid INTEGER NOT NULL,
|
||||
room_nid INTEGER NOT NULL,
|
||||
pseudo_id_key TEXT NULL, -- may be null for users not local to the server
|
||||
pseudo_id_pub_key TEXT NOT NULL,
|
||||
CONSTRAINT roomserver_user_room_keys_pk PRIMARY KEY (user_nid, room_nid)
|
||||
);
|
||||
`
|
||||
|
||||
const insertUserRoomKeySQL = `
|
||||
INSERT INTO roomserver_user_room_keys (user_nid, room_nid, pseudo_id_key, pseudo_id_pub_key) VALUES ($1, $2, $3, $4)
|
||||
ON CONFLICT DO UPDATE SET pseudo_id_key = roomserver_user_room_keys.pseudo_id_key
|
||||
RETURNING (pseudo_id_key)
|
||||
`
|
||||
|
||||
const insertUserRoomPublicKeySQL = `
|
||||
INSERT INTO roomserver_user_room_keys (user_nid, room_nid, pseudo_id_pub_key) VALUES ($1, $2, $3)
|
||||
ON CONFLICT DO UPDATE SET pseudo_id_pub_key = $3
|
||||
RETURNING (pseudo_id_pub_key)
|
||||
`
|
||||
|
||||
const selectUserRoomKeySQL = `SELECT pseudo_id_key FROM roomserver_user_room_keys WHERE user_nid = $1 AND room_nid = $2`
|
||||
|
||||
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 {
|
||||
insertUserRoomPrivateKeyStmt *sql.Stmt
|
||||
insertUserRoomPublicKeyStmt *sql.Stmt
|
||||
selectUserRoomKeyStmt *sql.Stmt
|
||||
//selectUserNIDsStmt *sql.Stmt //prepared at runtime
|
||||
}
|
||||
|
||||
func CreateUserRoomKeysTable(db *sql.DB) error {
|
||||
_, err := db.Exec(userRoomKeysSchema)
|
||||
return err
|
||||
}
|
||||
|
||||
func PrepareUserRoomKeysTable(db *sql.DB) (tables.UserRoomKeys, error) {
|
||||
s := &userRoomKeysStatements{}
|
||||
return s, sqlutil.StatementList{
|
||||
{&s.insertUserRoomPrivateKeyStmt, insertUserRoomKeySQL},
|
||||
{&s.insertUserRoomPublicKeyStmt, insertUserRoomPublicKeySQL},
|
||||
{&s.selectUserRoomKeyStmt, selectUserRoomKeySQL},
|
||||
//{&s.selectUserNIDsStmt, selectUserNIDsSQL}, //prepared at runtime
|
||||
}.Prepare(db)
|
||||
}
|
||||
|
||||
func (s *userRoomKeysStatements) InsertUserRoomPrivatePublicKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PrivateKey) (result ed25519.PrivateKey, err error) {
|
||||
stmt := sqlutil.TxStmtContext(ctx, txn, s.insertUserRoomPrivateKeyStmt)
|
||||
err = stmt.QueryRowContext(ctx, userNID, roomNID, key, key.Public()).Scan(&result)
|
||||
return result, err
|
||||
}
|
||||
|
||||
func (s *userRoomKeysStatements) InsertUserRoomPublicKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PublicKey) (result ed25519.PublicKey, err error) {
|
||||
stmt := sqlutil.TxStmtContext(ctx, txn, s.insertUserRoomPublicKeyStmt)
|
||||
err = stmt.QueryRowContext(ctx, userNID, roomNID, key).Scan(&result)
|
||||
return result, err
|
||||
}
|
||||
|
||||
func (s *userRoomKeysStatements) SelectUserRoomPrivateKey(
|
||||
ctx context.Context,
|
||||
txn *sql.Tx,
|
||||
userNID types.EventStateKeyNID,
|
||||
roomNID types.RoomNID,
|
||||
) (ed25519.PrivateKey, error) {
|
||||
stmt := sqlutil.TxStmtContext(ctx, txn, s.selectUserRoomKeyStmt)
|
||||
var result ed25519.PrivateKey
|
||||
err := stmt.QueryRowContext(ctx, userNID, roomNID).Scan(&result)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
return result, err
|
||||
}
|
||||
|
||||
func (s *userRoomKeysStatements) BulkSelectUserNIDs(ctx context.Context, txn *sql.Tx, senderKeys map[types.RoomNID][]ed25519.PublicKey) (map[string]types.UserRoomKeyPair, error) {
|
||||
|
||||
roomNIDs := make([]any, 0, len(senderKeys))
|
||||
var senders []any
|
||||
for roomNID := range senderKeys {
|
||||
roomNIDs = append(roomNIDs, roomNID)
|
||||
|
||||
for _, key := range senderKeys[roomNID] {
|
||||
senders = append(senders, []byte(key))
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
params := append(roomNIDs, senders...)
|
||||
|
||||
stmt := sqlutil.TxStmt(txn, selectStmt)
|
||||
defer internal.CloseAndLogIfError(ctx, stmt, "failed to close statement")
|
||||
|
||||
rows, err := stmt.QueryContext(ctx, params...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer internal.CloseAndLogIfError(ctx, rows, "failed to close rows")
|
||||
|
||||
result := make(map[string]types.UserRoomKeyPair, len(params))
|
||||
var publicKey []byte
|
||||
userRoomKeyPair := types.UserRoomKeyPair{}
|
||||
for rows.Next() {
|
||||
if err = rows.Scan(&userRoomKeyPair.EventStateKeyNID, &userRoomKeyPair.RoomNID, &publicKey); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result[string(publicKey)] = userRoomKeyPair
|
||||
}
|
||||
return result, rows.Err()
|
||||
}
|
@ -2,6 +2,7 @@ package tables
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ed25519"
|
||||
"database/sql"
|
||||
"errors"
|
||||
|
||||
@ -184,6 +185,19 @@ type Purge interface {
|
||||
) error
|
||||
}
|
||||
|
||||
type UserRoomKeys interface {
|
||||
// InsertUserRoomPrivatePublicKey inserts the given private key as well as the public key for it. This should be used
|
||||
// when creating keys locally.
|
||||
InsertUserRoomPrivatePublicKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PrivateKey) (ed25519.PrivateKey, error)
|
||||
// InsertUserRoomPublicKey inserts the given public key, this should be used for users NOT local to this server
|
||||
InsertUserRoomPublicKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PublicKey) (ed25519.PublicKey, error)
|
||||
// SelectUserRoomPrivateKey selects the private key for the given user and room combination
|
||||
SelectUserRoomPrivateKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID) (ed25519.PrivateKey, error)
|
||||
// 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)
|
||||
}
|
||||
|
||||
// StrippedEvent represents a stripped event for returning extracted content values.
|
||||
type StrippedEvent struct {
|
||||
RoomID string
|
||||
|
115
roomserver/storage/tables/user_room_keys_table_test.go
Normal file
115
roomserver/storage/tables/user_room_keys_table_test.go
Normal file
@ -0,0 +1,115 @@
|
||||
package tables_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ed25519"
|
||||
"database/sql"
|
||||
"testing"
|
||||
|
||||
"github.com/matrix-org/dendrite/internal/sqlutil"
|
||||
"github.com/matrix-org/dendrite/roomserver/storage/postgres"
|
||||
"github.com/matrix-org/dendrite/roomserver/storage/sqlite3"
|
||||
"github.com/matrix-org/dendrite/roomserver/storage/tables"
|
||||
"github.com/matrix-org/dendrite/roomserver/types"
|
||||
"github.com/matrix-org/dendrite/setup/config"
|
||||
"github.com/matrix-org/dendrite/test"
|
||||
"github.com/stretchr/testify/assert"
|
||||
ed255192 "golang.org/x/crypto/ed25519"
|
||||
)
|
||||
|
||||
func mustCreateUserRoomKeysTable(t *testing.T, dbType test.DBType) (tab tables.UserRoomKeys, db *sql.DB, close func()) {
|
||||
t.Helper()
|
||||
connStr, close := test.PrepareDBConnectionString(t, dbType)
|
||||
db, err := sqlutil.Open(&config.DatabaseOptions{
|
||||
ConnectionString: config.DataSource(connStr),
|
||||
}, sqlutil.NewExclusiveWriter())
|
||||
assert.NoError(t, err)
|
||||
switch dbType {
|
||||
case test.DBTypePostgres:
|
||||
err = postgres.CreateUserRoomKeysTable(db)
|
||||
assert.NoError(t, err)
|
||||
tab, err = postgres.PrepareUserRoomKeysTable(db)
|
||||
case test.DBTypeSQLite:
|
||||
err = sqlite3.CreateUserRoomKeysTable(db)
|
||||
assert.NoError(t, err)
|
||||
tab, err = sqlite3.PrepareUserRoomKeysTable(db)
|
||||
}
|
||||
assert.NoError(t, err)
|
||||
|
||||
return tab, db, close
|
||||
}
|
||||
|
||||
func TestUserRoomKeysTable(t *testing.T) {
|
||||
test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
|
||||
tab, db, close := mustCreateUserRoomKeysTable(t, dbType)
|
||||
defer close()
|
||||
userNID := types.EventStateKeyNID(1)
|
||||
roomNID := types.RoomNID(1)
|
||||
_, key, err := ed25519.GenerateKey(nil)
|
||||
assert.NoError(t, err)
|
||||
|
||||
err = sqlutil.WithTransaction(db, func(txn *sql.Tx) error {
|
||||
var gotKey, key2, key3 ed25519.PrivateKey
|
||||
gotKey, err = tab.InsertUserRoomPrivatePublicKey(context.Background(), txn, userNID, roomNID, key)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, gotKey, key)
|
||||
|
||||
// again, this shouldn't result in an error, but return the existing key
|
||||
_, key2, err = ed25519.GenerateKey(nil)
|
||||
assert.NoError(t, err)
|
||||
gotKey, err = tab.InsertUserRoomPrivatePublicKey(context.Background(), txn, userNID, roomNID, key2)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, gotKey, key)
|
||||
|
||||
// add another user
|
||||
_, key3, err = ed25519.GenerateKey(nil)
|
||||
assert.NoError(t, err)
|
||||
userNID2 := types.EventStateKeyNID(2)
|
||||
_, err = tab.InsertUserRoomPrivatePublicKey(context.Background(), txn, userNID2, roomNID, key3)
|
||||
assert.NoError(t, err)
|
||||
|
||||
gotKey, err = tab.SelectUserRoomPrivateKey(context.Background(), txn, userNID, roomNID)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, key, gotKey)
|
||||
|
||||
// try to update an existing key, this should only be done for users NOT on this homeserver
|
||||
var gotPubKey ed25519.PublicKey
|
||||
gotPubKey, err = tab.InsertUserRoomPublicKey(context.Background(), txn, userNID, roomNID, key2.Public().(ed25519.PublicKey))
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, key2.Public(), gotPubKey)
|
||||
|
||||
// Key doesn't exist
|
||||
gotKey, err = tab.SelectUserRoomPrivateKey(context.Background(), txn, userNID, 2)
|
||||
assert.NoError(t, err)
|
||||
assert.Nil(t, gotKey)
|
||||
|
||||
// query user NIDs for senderKeys
|
||||
var gotKeys map[string]types.UserRoomKeyPair
|
||||
query := map[types.RoomNID][]ed25519.PublicKey{
|
||||
roomNID: {key2.Public().(ed25519.PublicKey), key3.Public().(ed25519.PublicKey)},
|
||||
types.RoomNID(2): {key.Public().(ed25519.PublicKey), key3.Public().(ed25519.PublicKey)}, // doesn't exist
|
||||
}
|
||||
gotKeys, err = tab.BulkSelectUserNIDs(context.Background(), txn, query)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, gotKeys)
|
||||
|
||||
wantKeys := map[string]types.UserRoomKeyPair{
|
||||
string(key2.Public().(ed25519.PublicKey)): {RoomNID: roomNID, EventStateKeyNID: userNID},
|
||||
string(key3.Public().(ed25519.PublicKey)): {RoomNID: roomNID, EventStateKeyNID: userNID2},
|
||||
}
|
||||
assert.Equal(t, wantKeys, gotKeys)
|
||||
|
||||
// insert key that came in over federation
|
||||
var gotPublicKey, key4 ed255192.PublicKey
|
||||
key4, _, err = ed25519.GenerateKey(nil)
|
||||
assert.NoError(t, err)
|
||||
gotPublicKey, err = tab.InsertUserRoomPublicKey(context.Background(), txn, userNID, 2, key4)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, key4, gotPublicKey)
|
||||
|
||||
return nil
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
|
||||
})
|
||||
}
|
@ -44,6 +44,11 @@ type EventMetadata struct {
|
||||
RoomNID RoomNID
|
||||
}
|
||||
|
||||
type UserRoomKeyPair struct {
|
||||
RoomNID RoomNID
|
||||
EventStateKeyNID EventStateKeyNID
|
||||
}
|
||||
|
||||
// StateSnapshotNID is a numeric ID for the state at an event.
|
||||
type StateSnapshotNID int64
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user