From a386fbed2c3696cd28307e7cfe02822dff76e4f9 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Mon, 21 Feb 2022 12:30:43 +0000
Subject: [PATCH] Delete one-time keys when deleting a device (#2208)

---
 keyserver/storage/postgres/one_time_keys_table.go | 12 ++++++++++++
 keyserver/storage/shared/storage.go               |  3 +++
 keyserver/storage/sqlite3/one_time_keys_table.go  | 12 ++++++++++++
 keyserver/storage/tables/interface.go             |  1 +
 4 files changed, 28 insertions(+)

diff --git a/keyserver/storage/postgres/one_time_keys_table.go b/keyserver/storage/postgres/one_time_keys_table.go
index cc397ba84..0b143a1aa 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 deee76eb4..f2790c8df 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 185b88612..897839aca 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 ff70a2366..cd1719598 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 {