Add initial support for storing user room keys (#3098)

This commit is contained in:
Till 2023-06-12 12:45:42 +02:00 committed by GitHub
parent 5713c5715c
commit 832ccc32f6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 700 additions and 7 deletions

View File

@ -16,6 +16,7 @@ package storage
import ( import (
"context" "context"
"crypto/ed25519"
"github.com/matrix-org/gomatrixserverlib" "github.com/matrix-org/gomatrixserverlib"
"github.com/matrix-org/gomatrixserverlib/spec" "github.com/matrix-org/gomatrixserverlib/spec"
@ -27,6 +28,7 @@ import (
) )
type Database interface { type Database interface {
UserRoomKeys
// Do we support processing input events for more than one room at a time? // Do we support processing input events for more than one room at a time?
SupportsConcurrentRoomInputs() bool SupportsConcurrentRoomInputs() bool
// RoomInfo returns room information for the given room ID, or nil if there is no room. // 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) ) (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 { type RoomDatabase interface {
EventDatabase EventDatabase
UserRoomKeys
// RoomInfo returns room information for the given room ID, or nil if there is no room. // 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) RoomInfo(ctx context.Context, roomID string) (*types.RoomInfo, error)
RoomInfoByNID(ctx context.Context, roomNID types.RoomNID) (*types.RoomInfo, error) RoomInfoByNID(ctx context.Context, roomNID types.RoomNID) (*types.RoomInfo, error)

View File

@ -131,6 +131,9 @@ func (d *Database) create(db *sql.DB) error {
if err := CreateRedactionsTable(db); err != nil { if err := CreateRedactionsTable(db); err != nil {
return err return err
} }
if err := CreateUserRoomKeysTable(db); err != nil {
return err
}
return nil return nil
} }
@ -192,6 +195,11 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room
if err != nil { if err != nil {
return err return err
} }
userRoomKeys, err := PrepareUserRoomKeysTable(db)
if err != nil {
return err
}
d.Database = shared.Database{ d.Database = shared.Database{
DB: db, DB: db,
EventDatabase: shared.EventDatabase{ EventDatabase: shared.EventDatabase{
@ -215,6 +223,7 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room
MembershipTable: membership, MembershipTable: membership,
PublishedTable: published, PublishedTable: published,
Purge: purge, Purge: purge,
UserRoomKeyTable: userRoomKeys,
} }
return nil return nil
} }

View 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()
}

View File

@ -2,14 +2,18 @@ package shared
import ( import (
"context" "context"
"crypto/ed25519"
"database/sql" "database/sql"
"encoding/json" "encoding/json"
"errors"
"fmt" "fmt"
"sort" "sort"
"github.com/matrix-org/dendrite/internal/eventutil"
"github.com/matrix-org/gomatrixserverlib" "github.com/matrix-org/gomatrixserverlib"
"github.com/matrix-org/gomatrixserverlib/spec" "github.com/matrix-org/gomatrixserverlib/spec"
"github.com/matrix-org/util" "github.com/matrix-org/util"
"github.com/sirupsen/logrus"
"github.com/tidwall/gjson" "github.com/tidwall/gjson"
"github.com/matrix-org/dendrite/internal/caching" "github.com/matrix-org/dendrite/internal/caching"
@ -41,6 +45,7 @@ type Database struct {
MembershipTable tables.Membership MembershipTable tables.Membership
PublishedTable tables.Published PublishedTable tables.Published
Purge tables.Purge Purge tables.Purge
UserRoomKeyTable tables.UserRoomKeys
GetRoomUpdaterFn func(ctx context.Context, roomInfo *types.RoomInfo) (*RoomUpdater, error) 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 // 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! // it should live in this package!

View File

@ -2,11 +2,15 @@ package shared_test
import ( import (
"context" "context"
"crypto/ed25519"
"testing" "testing"
"time" "time"
"github.com/matrix-org/dendrite/internal/caching" "github.com/matrix-org/dendrite/internal/caching"
"github.com/matrix-org/gomatrixserverlib"
"github.com/matrix-org/gomatrixserverlib/spec"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
ed255192 "golang.org/x/crypto/ed25519"
"github.com/matrix-org/dendrite/internal/sqlutil" "github.com/matrix-org/dendrite/internal/sqlutil"
"github.com/matrix-org/dendrite/roomserver/storage/postgres" "github.com/matrix-org/dendrite/roomserver/storage/postgres"
@ -23,40 +27,61 @@ func mustCreateRoomserverDatabase(t *testing.T, dbType test.DBType) (*shared.Dat
connStr, clearDB := test.PrepareDBConnectionString(t, dbType) connStr, clearDB := test.PrepareDBConnectionString(t, dbType)
dbOpts := &config.DatabaseOptions{ConnectionString: config.DataSource(connStr)} 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) assert.NoError(t, err)
var membershipTable tables.Membership var membershipTable tables.Membership
var stateKeyTable tables.EventStateKeys var stateKeyTable tables.EventStateKeys
var userRoomKeys tables.UserRoomKeys
var roomsTable tables.Rooms
switch dbType { switch dbType {
case test.DBTypePostgres: case test.DBTypePostgres:
err = postgres.CreateRoomsTable(db)
assert.NoError(t, err)
err = postgres.CreateEventStateKeysTable(db) err = postgres.CreateEventStateKeysTable(db)
assert.NoError(t, err) assert.NoError(t, err)
err = postgres.CreateMembershipTable(db) err = postgres.CreateMembershipTable(db)
assert.NoError(t, err) 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) membershipTable, err = postgres.PrepareMembershipTable(db)
assert.NoError(t, err) assert.NoError(t, err)
stateKeyTable, err = postgres.PrepareEventStateKeysTable(db) stateKeyTable, err = postgres.PrepareEventStateKeysTable(db)
assert.NoError(t, err)
userRoomKeys, err = postgres.PrepareUserRoomKeysTable(db)
case test.DBTypeSQLite: case test.DBTypeSQLite:
err = sqlite3.CreateRoomsTable(db)
assert.NoError(t, err)
err = sqlite3.CreateEventStateKeysTable(db) err = sqlite3.CreateEventStateKeysTable(db)
assert.NoError(t, err) assert.NoError(t, err)
err = sqlite3.CreateMembershipTable(db) err = sqlite3.CreateMembershipTable(db)
assert.NoError(t, err) 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) membershipTable, err = sqlite3.PrepareMembershipTable(db)
assert.NoError(t, err) assert.NoError(t, err)
stateKeyTable, err = sqlite3.PrepareEventStateKeysTable(db) stateKeyTable, err = sqlite3.PrepareEventStateKeysTable(db)
assert.NoError(t, err)
userRoomKeys, err = sqlite3.PrepareUserRoomKeysTable(db)
} }
assert.NoError(t, err) assert.NoError(t, err)
cache := caching.NewRistrettoCache(8*1024*1024, time.Hour, false) 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{ return &shared.Database{
DB: db, DB: db,
EventDatabase: evDb, EventDatabase: evDb,
MembershipTable: membershipTable, MembershipTable: membershipTable,
Writer: sqlutil.NewExclusiveWriter(), UserRoomKeyTable: userRoomKeys,
RoomsTable: roomsTable,
Writer: writer,
Cache: cache, Cache: cache,
}, func() { }, func() {
clearDB() clearDB()
@ -97,3 +122,80 @@ func Test_GetLeftUsers(t *testing.T) {
assert.ElementsMatch(t, expectedUserIDs, leftUsers) 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)
})
}

View File

@ -138,6 +138,9 @@ func (d *Database) create(db *sql.DB) error {
if err := CreateRedactionsTable(db); err != nil { if err := CreateRedactionsTable(db); err != nil {
return err return err
} }
if err := CreateUserRoomKeysTable(db); err != nil {
return err
}
return nil return nil
} }
@ -199,6 +202,10 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room
if err != nil { if err != nil {
return err return err
} }
userRoomKeys, err := PrepareUserRoomKeysTable(db)
if err != nil {
return err
}
d.Database = shared.Database{ d.Database = shared.Database{
DB: db, DB: db,
@ -224,6 +231,7 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room
PublishedTable: published, PublishedTable: published,
GetRoomUpdaterFn: d.GetRoomUpdater, GetRoomUpdaterFn: d.GetRoomUpdater,
Purge: purge, Purge: purge,
UserRoomKeyTable: userRoomKeys,
} }
return nil return nil
} }

View 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()
}

View File

@ -2,6 +2,7 @@ package tables
import ( import (
"context" "context"
"crypto/ed25519"
"database/sql" "database/sql"
"errors" "errors"
@ -184,6 +185,19 @@ type Purge interface {
) error ) 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. // StrippedEvent represents a stripped event for returning extracted content values.
type StrippedEvent struct { type StrippedEvent struct {
RoomID string RoomID string

View 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)
})
}

View File

@ -44,6 +44,11 @@ type EventMetadata struct {
RoomNID RoomNID RoomNID RoomNID
} }
type UserRoomKeyPair struct {
RoomNID RoomNID
EventStateKeyNID EventStateKeyNID
}
// StateSnapshotNID is a numeric ID for the state at an event. // StateSnapshotNID is a numeric ID for the state at an event.
type StateSnapshotNID int64 type StateSnapshotNID int64