mirror of
https://github.com/hoernschen/dendrite.git
synced 2025-08-01 13:52:46 +00:00
Add initial support for storing user room keys (#3098)
This commit is contained in:
Till
GitHub
parent
5713c5715c
commit
832ccc32f6
10 changed files with 700 additions and 7 deletions
|
|
@ -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)
|
||||
})
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue