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

4 files changed, 28 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
+}
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 {