diff --git a/keyserver/storage/postgres/one_time_keys_table.go b/keyserver/storage/postgres/one_time_keys_table.go index cc397ba8..0b143a1a 100644 --- a/keyserver/storage/postgres/one_time_keys_table.go +++ b/keyserver/storage/postgres/one_time_keys_table.go @@ -59,6 +59,9 @@ const deleteOneTimeKeySQL = "" + const selectKeyByAlgorithmSQL = "" + "SELECT key_id, key_json FROM keyserver_one_time_keys WHERE user_id = $1 AND device_id = $2 AND algorithm = $3 LIMIT 1" +const deleteOneTimeKeysSQL = "" + + "DELETE FROM keyserver_one_time_keys WHERE user_id = $1 AND device_id = $2" + type oneTimeKeysStatements struct { db *sql.DB upsertKeysStmt *sql.Stmt @@ -66,6 +69,7 @@ type oneTimeKeysStatements struct { selectKeysCountStmt *sql.Stmt selectKeyByAlgorithmStmt *sql.Stmt deleteOneTimeKeyStmt *sql.Stmt + deleteOneTimeKeysStmt *sql.Stmt } func NewPostgresOneTimeKeysTable(db *sql.DB) (tables.OneTimeKeys, error) { @@ -91,6 +95,9 @@ func NewPostgresOneTimeKeysTable(db *sql.DB) (tables.OneTimeKeys, error) { if s.deleteOneTimeKeyStmt, err = db.Prepare(deleteOneTimeKeySQL); err != nil { return nil, err } + if s.deleteOneTimeKeysStmt, err = db.Prepare(deleteOneTimeKeysSQL); err != nil { + return nil, err + } return s, nil } @@ -187,3 +194,8 @@ func (s *oneTimeKeysStatements) SelectAndDeleteOneTimeKey( algorithm + ":" + keyID: json.RawMessage(keyJSON), }, err } + +func (s *oneTimeKeysStatements) DeleteOneTimeKeys(ctx context.Context, txn *sql.Tx, userID, deviceID string) error { + _, err := sqlutil.TxStmt(txn, s.deleteOneTimeKeysStmt).ExecContext(ctx, userID, deviceID) + return err +} diff --git a/keyserver/storage/shared/storage.go b/keyserver/storage/shared/storage.go index deee76eb..f2790c8d 100644 --- a/keyserver/storage/shared/storage.go +++ b/keyserver/storage/shared/storage.go @@ -171,6 +171,9 @@ func (d *Database) DeleteDeviceKeys(ctx context.Context, userID string, deviceID if err := d.DeviceKeysTable.DeleteDeviceKeys(ctx, txn, userID, string(deviceID)); err != nil && err != sql.ErrNoRows { return fmt.Errorf("d.DeviceKeysTable.DeleteDeviceKeys: %w", err) } + if err := d.OneTimeKeysTable.DeleteOneTimeKeys(ctx, txn, userID, string(deviceID)); err != nil && err != sql.ErrNoRows { + return fmt.Errorf("d.OneTimeKeysTable.DeleteOneTimeKeys: %w", err) + } } return nil }) diff --git a/keyserver/storage/sqlite3/one_time_keys_table.go b/keyserver/storage/sqlite3/one_time_keys_table.go index 185b8861..897839ac 100644 --- a/keyserver/storage/sqlite3/one_time_keys_table.go +++ b/keyserver/storage/sqlite3/one_time_keys_table.go @@ -58,6 +58,9 @@ const deleteOneTimeKeySQL = "" + const selectKeyByAlgorithmSQL = "" + "SELECT key_id, key_json FROM keyserver_one_time_keys WHERE user_id = $1 AND device_id = $2 AND algorithm = $3 LIMIT 1" +const deleteOneTimeKeysSQL = "" + + "DELETE FROM keyserver_one_time_keys WHERE user_id = $1 AND device_id = $2" + type oneTimeKeysStatements struct { db *sql.DB upsertKeysStmt *sql.Stmt @@ -65,6 +68,7 @@ type oneTimeKeysStatements struct { selectKeysCountStmt *sql.Stmt selectKeyByAlgorithmStmt *sql.Stmt deleteOneTimeKeyStmt *sql.Stmt + deleteOneTimeKeysStmt *sql.Stmt } func NewSqliteOneTimeKeysTable(db *sql.DB) (tables.OneTimeKeys, error) { @@ -90,6 +94,9 @@ func NewSqliteOneTimeKeysTable(db *sql.DB) (tables.OneTimeKeys, error) { if s.deleteOneTimeKeyStmt, err = db.Prepare(deleteOneTimeKeySQL); err != nil { return nil, err } + if s.deleteOneTimeKeysStmt, err = db.Prepare(deleteOneTimeKeysSQL); err != nil { + return nil, err + } return s, nil } @@ -201,3 +208,8 @@ func (s *oneTimeKeysStatements) SelectAndDeleteOneTimeKey( algorithm + ":" + keyID: json.RawMessage(keyJSON), }, err } + +func (s *oneTimeKeysStatements) DeleteOneTimeKeys(ctx context.Context, txn *sql.Tx, userID, deviceID string) error { + _, err := sqlutil.TxStmt(txn, s.deleteOneTimeKeysStmt).ExecContext(ctx, userID, deviceID) + return err +} diff --git a/keyserver/storage/tables/interface.go b/keyserver/storage/tables/interface.go index ff70a236..cd171959 100644 --- a/keyserver/storage/tables/interface.go +++ b/keyserver/storage/tables/interface.go @@ -31,6 +31,7 @@ type OneTimeKeys interface { // SelectAndDeleteOneTimeKey selects a single one time key matching the user/device/algorithm specified and returns the algo:key_id => JSON. // Returns an empty map if the key does not exist. SelectAndDeleteOneTimeKey(ctx context.Context, txn *sql.Tx, userID, deviceID, algorithm string) (map[string]json.RawMessage, error) + DeleteOneTimeKeys(ctx context.Context, txn *sql.Tx, userID, deviceID string) error } type DeviceKeys interface {