From e5ef9a2ead670ded5fd7652b88bea6e86486c347 Mon Sep 17 00:00:00 2001 From: Till Faelligen <2353100+S7evinK@users.noreply.github.com> Date: Fri, 2 Jun 2023 14:32:13 +0200 Subject: [PATCH] Add initial support for storing user room keys --- roomserver/storage/interface.go | 4 + roomserver/storage/postgres/storage.go | 9 ++ .../storage/postgres/user_room_keys_table.go | 82 +++++++++++++++++++ roomserver/storage/shared/storage.go | 25 ++++++ roomserver/storage/shared/storage_test.go | 49 +++++++++-- roomserver/storage/sqlite3/storage.go | 8 ++ .../storage/sqlite3/user_room_keys_table.go | 82 +++++++++++++++++++ roomserver/storage/tables/interface.go | 6 ++ .../tables/user_room_keys_table_test.go | 62 ++++++++++++++ 9 files changed, 322 insertions(+), 5 deletions(-) create mode 100644 roomserver/storage/postgres/user_room_keys_table.go create mode 100644 roomserver/storage/sqlite3/user_room_keys_table.go create mode 100644 roomserver/storage/tables/user_room_keys_table_test.go diff --git a/roomserver/storage/interface.go b/roomserver/storage/interface.go index 7d22df008..ac77965fc 100644 --- a/roomserver/storage/interface.go +++ b/roomserver/storage/interface.go @@ -16,6 +16,7 @@ package storage import ( "context" + "crypto/ed25519" "github.com/matrix-org/gomatrixserverlib" "github.com/matrix-org/gomatrixserverlib/spec" @@ -188,6 +189,9 @@ type Database interface { MaybeRedactEvent( ctx context.Context, roomInfo *types.RoomInfo, eventNID types.EventNID, event gomatrixserverlib.PDU, plResolver state.PowerLevelResolver, ) (gomatrixserverlib.PDU, gomatrixserverlib.PDU, error) + + InsertUserRoomKey(ctx context.Context, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PrivateKey) error + SelectUserRoomKey(ctx context.Context, userNID types.EventStateKeyNID, roomNID types.RoomNID) (key ed25519.PrivateKey, err error) } type RoomDatabase interface { diff --git a/roomserver/storage/postgres/storage.go b/roomserver/storage/postgres/storage.go index 19cde5410..453ff45da 100644 --- a/roomserver/storage/postgres/storage.go +++ b/roomserver/storage/postgres/storage.go @@ -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 } diff --git a/roomserver/storage/postgres/user_room_keys_table.go b/roomserver/storage/postgres/user_room_keys_table.go new file mode 100644 index 000000000..1f0943c9a --- /dev/null +++ b/roomserver/storage/postgres/user_room_keys_table.go @@ -0,0 +1,82 @@ +// 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" + + "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 roomserver_user_room_keys ( + user_nid INTEGER NOT NULL, + room_nid INTEGER NOT NULL, + pseudo_id_key BYTEA 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) VALUES ($1, $2, $3)` +const selectUserRoomKeySQL = `SELECT pseudo_id_key FROM roomserver_user_room_keys WHERE user_nid = $1 AND room_nid = $2` + +type userRoomKeysStatements struct { + insertUserRoomKeyStmt *sql.Stmt + selectUserRoomKeyStmt *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.insertUserRoomKeyStmt, insertUserRoomKeySQL}, + {&s.selectUserRoomKeyStmt, selectUserRoomKeySQL}, + }.Prepare(db) +} + +func (s *userRoomKeysStatements) InsertUserRoomKey( + ctx context.Context, + txn *sql.Tx, + userNID types.EventStateKeyNID, + roomNID types.RoomNID, + key ed25519.PrivateKey, +) error { + stmt := sqlutil.TxStmtContext(ctx, txn, s.insertUserRoomKeyStmt) + defer internal.CloseAndLogIfError(ctx, stmt, "failed to close statement") + _, err := stmt.ExecContext(ctx, userNID, roomNID, key) + return err +} + +func (s *userRoomKeysStatements) SelectUserRoomKey( + ctx context.Context, + txn *sql.Tx, + userNID types.EventStateKeyNID, + roomNID types.RoomNID, +) (ed25519.PrivateKey, error) { + stmt := sqlutil.TxStmtContext(ctx, txn, s.selectUserRoomKeyStmt) + defer internal.CloseAndLogIfError(ctx, stmt, "failed to close statement") + var result ed25519.PrivateKey + err := stmt.QueryRowContext(ctx, userNID, roomNID).Scan(&result) + return result, err +} diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go index cefa58a3d..a1ea2ff59 100644 --- a/roomserver/storage/shared/storage.go +++ b/roomserver/storage/shared/storage.go @@ -2,6 +2,7 @@ package shared import ( "context" + "crypto/ed25519" "database/sql" "encoding/json" "fmt" @@ -41,6 +42,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) } @@ -1589,6 +1591,29 @@ func (d *Database) UpgradeRoom(ctx context.Context, oldRoomID, newRoomID, eventS }) } +// InsertUserRoomKey inserts a new user room key for the given user and room. +// Returns an error if a database error occurred, also if the primary constraint was violated. +func (d *Database) InsertUserRoomKey(ctx context.Context, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PrivateKey) error { + return d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error { + return d.UserRoomKeyTable.InsertUserRoomKey(ctx, txn, userNID, roomNID, key) + }) +} + +// SelectUserRoomKey queries the user room key for a given user. +// Returns the key and an error. +// TODO: should we handle absent keys (sql.ErrNoRows) as non-fatal? +func (d *Database) SelectUserRoomKey(ctx context.Context, userNID types.EventStateKeyNID, roomNID types.RoomNID) (key ed25519.PrivateKey, err error) { + err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error { + var sErr error + key, sErr = d.UserRoomKeyTable.SelectUserRoomKey(ctx, txn, userNID, roomNID) + if sErr != nil { + return sErr + } + return nil + }) + return +} + // 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! diff --git a/roomserver/storage/shared/storage_test.go b/roomserver/storage/shared/storage_test.go index 941e84802..b05fafae3 100644 --- a/roomserver/storage/shared/storage_test.go +++ b/roomserver/storage/shared/storage_test.go @@ -2,10 +2,12 @@ package shared_test import ( "context" + "crypto/ed25519" "testing" "time" "github.com/matrix-org/dendrite/internal/caching" + "github.com/matrix-org/dendrite/roomserver/types" "github.com/stretchr/testify/assert" "github.com/matrix-org/dendrite/internal/sqlutil" @@ -28,23 +30,32 @@ func mustCreateRoomserverDatabase(t *testing.T, dbType test.DBType) (*shared.Dat var membershipTable tables.Membership var stateKeyTable tables.EventStateKeys + var userRoomKeys tables.UserRoomKeys switch dbType { case test.DBTypePostgres: err = postgres.CreateEventStateKeysTable(db) assert.NoError(t, err) err = postgres.CreateMembershipTable(db) assert.NoError(t, err) + err = postgres.CreateUserRoomKeysTable(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.CreateEventStateKeysTable(db) assert.NoError(t, err) err = sqlite3.CreateMembershipTable(db) assert.NoError(t, err) + err = sqlite3.CreateUserRoomKeysTable(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) @@ -53,11 +64,12 @@ func mustCreateRoomserverDatabase(t *testing.T, dbType test.DBType) (*shared.Dat evDb := shared.EventDatabase{EventStateKeysTable: stateKeyTable, Cache: cache} return &shared.Database{ - DB: db, - EventDatabase: evDb, - MembershipTable: membershipTable, - Writer: sqlutil.NewExclusiveWriter(), - Cache: cache, + DB: db, + EventDatabase: evDb, + MembershipTable: membershipTable, + UserRoomKeyTable: userRoomKeys, + Writer: sqlutil.NewExclusiveWriter(), + Cache: cache, }, func() { clearDB() err = db.Close() @@ -97,3 +109,30 @@ func Test_GetLeftUsers(t *testing.T) { assert.ElementsMatch(t, expectedUserIDs, leftUsers) }) } + +func TestUserRoomKeys(t *testing.T) { + ctx := context.Background() + test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) { + db, close := mustCreateRoomserverDatabase(t, dbType) + defer close() + userNID := types.EventStateKeyNID(1) + roomNID := types.RoomNID(1) + _, key, err := ed25519.GenerateKey(nil) + assert.NoError(t, err) + + err = db.InsertUserRoomKey(ctx, userNID, roomNID, key) + assert.NoError(t, err) + + // again, this should result in an error now, due to the primary key on userNID/roomNID + err = db.InsertUserRoomKey(context.Background(), userNID, roomNID, key) + assert.Error(t, err) + + gotKey, err := db.SelectUserRoomKey(context.Background(), userNID, roomNID) + assert.NoError(t, err) + assert.Equal(t, key, gotKey) + + // Key doesn't exist + _, err = db.SelectUserRoomKey(context.Background(), userNID, 2) + assert.Error(t, err) + }) +} diff --git a/roomserver/storage/sqlite3/storage.go b/roomserver/storage/sqlite3/storage.go index 6ab427a84..ef51a5b08 100644 --- a/roomserver/storage/sqlite3/storage.go +++ b/roomserver/storage/sqlite3/storage.go @@ -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 } diff --git a/roomserver/storage/sqlite3/user_room_keys_table.go b/roomserver/storage/sqlite3/user_room_keys_table.go new file mode 100644 index 000000000..72452f9f3 --- /dev/null +++ b/roomserver/storage/sqlite3/user_room_keys_table.go @@ -0,0 +1,82 @@ +// 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" + + "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 roomserver_user_room_keys ( + user_nid INTEGER NOT NULL, + room_nid INTEGER NOT NULL, + pseudo_id_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) VALUES ($1, $2, $3)` +const selectUserRoomKeySQL = `SELECT pseudo_id_key FROM roomserver_user_room_keys WHERE user_nid = $1 AND room_nid = $2` + +type userRoomKeysStatements struct { + insertUserRoomKeyStmt *sql.Stmt + selectUserRoomKeyStmt *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.insertUserRoomKeyStmt, insertUserRoomKeySQL}, + {&s.selectUserRoomKeyStmt, selectUserRoomKeySQL}, + }.Prepare(db) +} + +func (s *userRoomKeysStatements) InsertUserRoomKey( + ctx context.Context, + txn *sql.Tx, + userNID types.EventStateKeyNID, + roomNID types.RoomNID, + key ed25519.PrivateKey, +) error { + stmt := sqlutil.TxStmtContext(ctx, txn, s.insertUserRoomKeyStmt) + defer internal.CloseAndLogIfError(ctx, stmt, "failed to close statement") + _, err := stmt.ExecContext(ctx, userNID, roomNID, key) + return err +} + +func (s *userRoomKeysStatements) SelectUserRoomKey( + ctx context.Context, + txn *sql.Tx, + userNID types.EventStateKeyNID, + roomNID types.RoomNID, +) (ed25519.PrivateKey, error) { + stmt := sqlutil.TxStmtContext(ctx, txn, s.selectUserRoomKeyStmt) + defer internal.CloseAndLogIfError(ctx, stmt, "failed to close statement") + var result ed25519.PrivateKey + err := stmt.QueryRowContext(ctx, userNID, roomNID).Scan(&result) + return result, err +} diff --git a/roomserver/storage/tables/interface.go b/roomserver/storage/tables/interface.go index 333483b32..2d875bffc 100644 --- a/roomserver/storage/tables/interface.go +++ b/roomserver/storage/tables/interface.go @@ -2,6 +2,7 @@ package tables import ( "context" + "crypto/ed25519" "database/sql" "errors" @@ -184,6 +185,11 @@ type Purge interface { ) error } +type UserRoomKeys interface { + InsertUserRoomKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID, key ed25519.PrivateKey) error + SelectUserRoomKey(ctx context.Context, txn *sql.Tx, userNID types.EventStateKeyNID, roomNID types.RoomNID) (ed25519.PrivateKey, error) +} + // StrippedEvent represents a stripped event for returning extracted content values. type StrippedEvent struct { RoomID string diff --git a/roomserver/storage/tables/user_room_keys_table_test.go b/roomserver/storage/tables/user_room_keys_table_test.go new file mode 100644 index 000000000..cc0550c5c --- /dev/null +++ b/roomserver/storage/tables/user_room_keys_table_test.go @@ -0,0 +1,62 @@ +package tables_test + +import ( + "context" + "crypto/ed25519" + "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" +) + +func mustCreateUserRoomKeysTable(t *testing.T, dbType test.DBType) (tab tables.UserRoomKeys, 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, close +} + +func TestUserRoomKeysTable(t *testing.T) { + test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) { + tab, close := mustCreateUserRoomKeysTable(t, dbType) + defer close() + userNID := types.EventStateKeyNID(1) + roomNID := types.RoomNID(1) + _, key, err := ed25519.GenerateKey(nil) + assert.NoError(t, err) + err = tab.InsertUserRoomKey(context.Background(), nil, userNID, roomNID, key) + assert.NoError(t, err) + // again, this should result in an error now, due to the primary key on userNID/roomNID + err = tab.InsertUserRoomKey(context.Background(), nil, userNID, roomNID, key) + assert.Error(t, err) + + gotKey, err := tab.SelectUserRoomKey(context.Background(), nil, userNID, roomNID) + assert.NoError(t, err) + assert.Equal(t, key, gotKey) + + // Key doesn't exist + _, err = tab.SelectUserRoomKey(context.Background(), nil, userNID, 2) + assert.Error(t, err) + }) +}