Clean up old one-time prekeys.

2026-04-22 01:40:07 +01:00 · 2023-08-11 12:38:03 -04:00
parent 389b439e9a
commit d6adfea9b1
18 changed files with 572 additions and 11 deletions
--- a/app/src/main/java/org/thoughtcrime/securesms/database/KyberPreKeyTable.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/KyberPreKeyTable.kt
@@ -4,12 +4,14 @@ import android.content.Context
 import org.signal.core.util.delete
 import org.signal.core.util.exists
 import org.signal.core.util.insertInto
+import org.signal.core.util.logging.Log
 import org.signal.core.util.readToList
 import org.signal.core.util.readToSingleObject
 import org.signal.core.util.requireBoolean
 import org.signal.core.util.requireNonNullBlob
 import org.signal.core.util.select
 import org.signal.core.util.toInt
+import org.signal.core.util.update
 import org.signal.libsignal.protocol.state.KyberPreKeyRecord
 import org.whispersystems.signalservice.api.push.ServiceId

@@ -18,6 +20,8 @@ import org.whispersystems.signalservice.api.push.ServiceId
 */
 class KyberPreKeyTable(context: Context, databaseHelper: SignalDatabase) : DatabaseTable(context, databaseHelper) {
  companion object {
+    private val TAG = Log.tag(KyberPreKeyTable::class.java)
+
    const val TABLE_NAME = "kyber_prekey"
    const val ID = "_id"
    const val ACCOUNT_ID = "account_id"
@@ -25,6 +29,8 @@ class KyberPreKeyTable(context: Context, databaseHelper: SignalDatabase) : Datab
    const val TIMESTAMP = "timestamp"
    const val LAST_RESORT = "last_resort"
    const val SERIALIZED = "serialized"
+    const val STALE_TIMESTAMP = "stale_timestamp"
+
    const val CREATE_TABLE = """
      CREATE TABLE $TABLE_NAME (
        $ID INTEGER PRIMARY KEY,
@@ -33,6 +39,7 @@ class KyberPreKeyTable(context: Context, databaseHelper: SignalDatabase) : Datab
        $TIMESTAMP INTEGER NOT NULL,
        $LAST_RESORT INTEGER NOT NULL,
        $SERIALIZED BLOB NOT NULL,
+        $STALE_TIMESTAMP INTEGER NOT NULL DEFAULT 0,
        UNIQUE($ACCOUNT_ID, $KEY_ID)
    )
    """
@@ -120,6 +127,48 @@ class KyberPreKeyTable(context: Context, databaseHelper: SignalDatabase) : Datab
      .run()
  }

+  fun markAllStaleIfNecessary(serviceId: ServiceId, staleTime: Long) {
+    writableDatabase
+      .update(TABLE_NAME)
+      .values(STALE_TIMESTAMP to staleTime)
+      .where("$ACCOUNT_ID = ? AND $STALE_TIMESTAMP = 0 AND $LAST_RESORT = 0", serviceId)
+      .run()
+  }
+
+  /**
+   * Deletes all keys that have been stale since before the specified threshold.
+   * We will always keep at least [minCount] items, preferring more recent ones.
+   */
+  fun deleteAllStaleBefore(serviceId: ServiceId, threshold: Long, minCount: Int) {
+    val count = writableDatabase
+      .delete(TABLE_NAME)
+      .where(
+        """
+          $ACCOUNT_ID = ? 
+            AND $LAST_RESORT = 0
+            AND $STALE_TIMESTAMP > 0 
+            AND $STALE_TIMESTAMP < $threshold
+            AND $ID NOT IN (
+              SELECT $ID
+              FROM $TABLE_NAME
+              WHERE 
+                $ACCOUNT_ID = ?
+                AND $LAST_RESORT = 0
+              ORDER BY 
+                CASE $STALE_TIMESTAMP WHEN 0 THEN 1 ELSE 0 END DESC,
+                $STALE_TIMESTAMP DESC,
+                $ID DESC
+              LIMIT $minCount
+            )
+        """,
+        serviceId,
+        serviceId
+      )
+      .run()
+
+    Log.i(TAG, "Deleted $count stale one-time EC prekeys.")
+  }
+
  data class KyberPreKey(
    val record: KyberPreKeyRecord,
    val lastResort: Boolean
--- a/app/src/main/java/org/thoughtcrime/securesms/database/OneTimePreKeyTable.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/OneTimePreKeyTable.kt
@@ -3,8 +3,10 @@ package org.thoughtcrime.securesms.database
 import android.content.Context
 import androidx.core.content.contentValuesOf
 import org.signal.core.util.SqlUtil
+import org.signal.core.util.delete
 import org.signal.core.util.logging.Log
 import org.signal.core.util.requireNonNullString
+import org.signal.core.util.update
 import org.signal.libsignal.protocol.InvalidKeyException
 import org.signal.libsignal.protocol.ecc.Curve
 import org.signal.libsignal.protocol.ecc.ECKeyPair
@@ -23,6 +25,8 @@ class OneTimePreKeyTable(context: Context, databaseHelper: SignalDatabase) : Dat
    const val KEY_ID = "key_id"
    const val PUBLIC_KEY = "public_key"
    const val PRIVATE_KEY = "private_key"
+    const val STALE_TIMESTAMP = "stale_timestamp"
+
    const val CREATE_TABLE = """
      CREATE TABLE $TABLE_NAME (
        $ID INTEGER PRIMARY KEY,
@@ -30,6 +34,7 @@ class OneTimePreKeyTable(context: Context, databaseHelper: SignalDatabase) : Dat
        $KEY_ID INTEGER UNIQUE, 
        $PUBLIC_KEY TEXT NOT NULL, 
        $PRIVATE_KEY TEXT NOT NULL,
+        $STALE_TIMESTAMP INTEGER NOT NULL DEFAULT 0,
        UNIQUE($ACCOUNT_ID, $KEY_ID)
      )
    """
@@ -68,4 +73,43 @@ class OneTimePreKeyTable(context: Context, databaseHelper: SignalDatabase) : Dat
    val database = databaseHelper.signalWritableDatabase
    database.delete(TABLE_NAME, "$ACCOUNT_ID = ? AND $KEY_ID = ?", SqlUtil.buildArgs(serviceId, keyId))
  }
+
+  fun markAllStaleIfNecessary(serviceId: ServiceId, staleTime: Long) {
+    writableDatabase
+      .update(TABLE_NAME)
+      .values(STALE_TIMESTAMP to staleTime)
+      .where("$ACCOUNT_ID = ? AND $STALE_TIMESTAMP = 0", serviceId)
+      .run()
+  }
+
+  /**
+   * Deletes all keys that have been stale since before the specified threshold.
+   * We will always keep at least [minCount] items, preferring more recent ones.
+   */
+  fun deleteAllStaleBefore(serviceId: ServiceId, threshold: Long, minCount: Int) {
+    val count = writableDatabase
+      .delete(TABLE_NAME)
+      .where(
+        """
+          $ACCOUNT_ID = ? 
+            AND $STALE_TIMESTAMP > 0 
+            AND $STALE_TIMESTAMP < $threshold
+            AND $ID NOT IN (
+              SELECT $ID
+              FROM $TABLE_NAME
+              WHERE $ACCOUNT_ID = ?
+              ORDER BY 
+                CASE $STALE_TIMESTAMP WHEN 0 THEN 1 ELSE 0 END DESC,
+                $STALE_TIMESTAMP DESC,
+                $ID DESC
+              LIMIT $minCount
+            )
+        """,
+        serviceId,
+        serviceId
+      )
+      .run()
+
+    Log.i(TAG, "Deleted $count stale one-time EC prekeys.")
+  }
 }
--- a/app/src/main/java/org/thoughtcrime/securesms/database/helpers/SignalDatabaseMigrations.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/helpers/SignalDatabaseMigrations.kt
@@ -58,6 +58,7 @@ import org.thoughtcrime.securesms.database.helpers.migration.V199_AddThreadActiv
 import org.thoughtcrime.securesms.database.helpers.migration.V200_ResetPniColumn
 import org.thoughtcrime.securesms.database.helpers.migration.V201_RecipientTableValidations
 import org.thoughtcrime.securesms.database.helpers.migration.V202_DropMessageTableThreadDateIndex
+import org.thoughtcrime.securesms.database.helpers.migration.V203_PreKeyStaleTimestamp

 /**
 * Contains all of the database migrations for [SignalDatabase]. Broken into a separate file for cleanliness.
@@ -66,7 +67,7 @@ object SignalDatabaseMigrations {

  val TAG: String = Log.tag(SignalDatabaseMigrations.javaClass)

-  const val DATABASE_VERSION = 202
+  const val DATABASE_VERSION = 203

  @JvmStatic
  fun migrate(context: Application, db: SQLiteDatabase, oldVersion: Int, newVersion: Int) {
@@ -285,6 +286,10 @@ object SignalDatabaseMigrations {
    if (oldVersion < 202) {
      V202_DropMessageTableThreadDateIndex.migrate(context, db, oldVersion, newVersion)
    }
+
+    if (oldVersion < 203) {
+      V203_PreKeyStaleTimestamp.migrate(context, db, oldVersion, newVersion)
+    }
  }

  @JvmStatic
--- a/app/src/main/java/org/thoughtcrime/securesms/database/helpers/migration/V203_PreKeyStaleTimestamp.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/helpers/migration/V203_PreKeyStaleTimestamp.kt
@@ -0,0 +1,40 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.database.helpers.migration
+
+import android.app.Application
+import androidx.sqlite.db.SupportSQLiteDatabase
+import net.zetetic.database.sqlcipher.SQLiteDatabase
+
+/**
+ * Keep track of a "stale timestamp" for one-time prekeys so that we can know when it's safe to delete them.
+ */
+@Suppress("ClassName")
+object V203_PreKeyStaleTimestamp : SignalDatabaseMigration {
+  override fun migrate(context: Application, db: SQLiteDatabase, oldVersion: Int, newVersion: Int) {
+    // Note: Because of a sequencing issue between beta/nightly, we had two V202 migrations (of which this used to be one of them),
+    //       so we have to do some conditional migrating based on the user's current state.
+    db.execSQL("DROP INDEX IF EXISTS message_thread_date_index")
+
+    if (!columnExists(db, "one_time_prekeys", "stale_timestamp")) {
+      db.execSQL("ALTER TABLE one_time_prekeys ADD COLUMN stale_timestamp INTEGER NOT NULL DEFAULT 0")
+      db.execSQL("ALTER TABLE kyber_prekey ADD COLUMN stale_timestamp INTEGER NOT NULL DEFAULT 0")
+    }
+  }
+
+  private fun columnExists(db: SupportSQLiteDatabase, table: String, column: String): Boolean {
+    db.query("PRAGMA table_info($table)", null).use { cursor ->
+      val nameColumnIndex = cursor.getColumnIndexOrThrow("name")
+      while (cursor.moveToNext()) {
+        val name = cursor.getString(nameColumnIndex)
+        if (name == column) {
+          return true
+        }
+      }
+    }
+    return false
+  }
+}