Delete one-time keys when deleting a device (#2208)

1 files changed, 12 insertions, 0 deletions

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
+}