Reject last-use kyber key sets that we've seen before.

app/src/main/java/org/thoughtcrime/securesms/crypto/storage/SignalKyberPreKeyStore.kt

@@ -6,6 +6,7 @@
 package org.thoughtcrime.securesms.crypto.storage
 
 import org.signal.libsignal.protocol.InvalidKeyIdException
+import org.signal.libsignal.protocol.ecc.ECPublicKey
 import org.signal.libsignal.protocol.state.KyberPreKeyRecord
 import org.signal.libsignal.protocol.state.KyberPreKeyStore
 import org.thoughtcrime.securesms.crypto.ReentrantSessionLock
@@ -56,9 +57,9 @@ class SignalKyberPreKeyStore(private val selfServiceId: ServiceId) : SignalServi
     }
   }
 
-  override fun markKyberPreKeyUsed(kyberPreKeyId: Int) {
+  override fun markKyberPreKeyUsed(kyberPreKeyId: Int, signedPreKeyId: Int, baseKey: ECPublicKey) {
     ReentrantSessionLock.INSTANCE.acquire().use {
-      SignalDatabase.kyberPreKeys.deleteIfNotLastResort(selfServiceId, kyberPreKeyId)
+      SignalDatabase.kyberPreKeys.handleMarkKyberPreKeyUsed(selfServiceId, kyberPreKeyId, signedPreKeyId, baseKey)
     }
   }
 

app/src/main/java/org/thoughtcrime/securesms/crypto/storage/SignalServiceAccountDataStoreImpl.java

@@ -9,6 +9,7 @@ import org.signal.libsignal.protocol.IdentityKeyPair;
 import org.signal.libsignal.protocol.InvalidKeyIdException;
 import org.signal.libsignal.protocol.NoSessionException;
 import org.signal.libsignal.protocol.SignalProtocolAddress;
+import org.signal.libsignal.protocol.ecc.ECPublicKey;
 import org.signal.libsignal.protocol.groups.state.SenderKeyRecord;
 import org.signal.libsignal.protocol.state.KyberPreKeyRecord;
 import org.signal.libsignal.protocol.state.PreKeyRecord;
@@ -214,8 +215,8 @@ public class SignalServiceAccountDataStoreImpl implements SignalServiceAccountDa
   }
 
   @Override
-  public void markKyberPreKeyUsed(int kyberPreKeyId) {
-    kyberPreKeyStore.markKyberPreKeyUsed(kyberPreKeyId);
+  public void markKyberPreKeyUsed(int kyberPreKeyId, int signedKeyId, ECPublicKey publicKey) {
+    kyberPreKeyStore.markKyberPreKeyUsed(kyberPreKeyId, signedKeyId, publicKey);
   }
 
   @Override

app/src/main/java/org/thoughtcrime/securesms/database/KyberPreKeyTable.kt

@@ -7,12 +7,15 @@ 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.readToSingleInt
 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.core.util.withinTransaction
+import org.signal.libsignal.protocol.ecc.ECPublicKey
 import org.signal.libsignal.protocol.state.KyberPreKeyRecord
 import org.whispersystems.signalservice.api.push.ServiceId
 
@@ -116,11 +119,33 @@ class KyberPreKeyTable(context: Context, databaseHelper: SignalDatabase) : Datab
       .run(SQLiteDatabase.CONFLICT_REPLACE)
   }
 
-  fun deleteIfNotLastResort(serviceId: ServiceId, keyId: Int) {
-    writableDatabase
-      .delete("$TABLE_NAME INDEXED BY $INDEX_ACCOUNT_KEY")
-      .where("$ACCOUNT_ID = ? AND $KEY_ID = ? AND $LAST_RESORT = ?", serviceId.toAccountId(), keyId, 0)
-      .run()
+  /**
+   * When we mark Kyber pre-keys used, we want to keep a record of last resort tuples, which are deleted when they key
+   * itself is deleted from this table via a cascading delete.
+   *
+   * For non-last-resort keys, this method just deletes them like normal.
+   */
+  fun handleMarkKyberPreKeyUsed(serviceId: ServiceId, kyberPreKeyId: Int, signedPreKeyId: Int, baseKey: ECPublicKey) {
+    writableDatabase.withinTransaction { db ->
+      val lastResortRowId = db
+        .select(ID)
+        .from(TABLE_NAME)
+        .where("$ACCOUNT_ID = ? AND $KEY_ID = ? AND $LAST_RESORT = ?", serviceId.toAccountId(), kyberPreKeyId, 1)
+        .run()
+        .readToSingleInt(-1)
+
+      if (lastResortRowId < 0) {
+        db.delete("$TABLE_NAME INDEXED BY $INDEX_ACCOUNT_KEY")
+          .where("$ACCOUNT_ID = ? AND $KEY_ID = ? AND $LAST_RESORT = ?", serviceId.toAccountId(), kyberPreKeyId, 0)
+          .run()
+      } else {
+        SignalDatabase.lastResortKeyTuples.insert(
+          kyberPreKeyRowId = lastResortRowId,
+          signedKeyId = signedPreKeyId,
+          publicKey = baseKey
+        )
+      }
+    }
   }
 
   fun delete(serviceId: ServiceId, keyId: Int) {

app/src/main/java/org/thoughtcrime/securesms/database/LastResortKeyTupleTable.kt

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.database
+
+import android.content.Context
+import android.database.sqlite.SQLiteConstraintException
+import org.signal.core.util.insertInto
+import org.signal.core.util.logging.Log
+import org.signal.libsignal.protocol.ReusedBaseKeyException
+import org.signal.libsignal.protocol.ecc.ECPublicKey
+
+/**
+ * Stores a Tuple of (kyberPreKeyId, signedPreKeyId, baseKey) for each
+ * last-resort key when [KyberPreKeyStore#markKyberPreKeyUsed] is called.
+ *
+ * Entries in this table are unique and an error will be thrown on trying to insert a duplicate.
+ */
+class LastResortKeyTupleTable(context: Context, databaseHelper: SignalDatabase) : DatabaseTable(context, databaseHelper) {
+
+  companion object {
+    private val TAG = Log.tag(LastResortKeyTupleTable::class)
+
+    private const val TABLE_NAME = "last_resort_key_tuple"
+
+    private const val ID = "_id"
+    private const val KYBER_PREKEY = "kyber_prekey_id"
+    private const val SIGNED_KEY_ID = "signed_key_id"
+    private const val PUBLIC_KEY = "public_key"
+
+    const val CREATE_TABLE = """
+      CREATE TABLE $TABLE_NAME (
+        $ID INTEGER PRIMARY KEY,
+        $KYBER_PREKEY INTEGER NOT NULL UNIQUE REFERENCES ${KyberPreKeyTable.TABLE_NAME} (${KyberPreKeyTable.ID}) ON DELETE CASCADE,
+        $SIGNED_KEY_ID INTEGER NOT NULL,
+        $PUBLIC_KEY BLOB NOT NULL,
+        UNIQUE($KYBER_PREKEY, $SIGNED_KEY_ID, $PUBLIC_KEY)
+      )
+    """
+  }
+
+  /**
+   * Inserts the Last-resort tuple. If it already exists, we will throw a [ReusedBaseKeyException]
+   */
+  @Throws(ReusedBaseKeyException::class)
+  fun insert(kyberPreKeyRowId: Int, signedKeyId: Int, publicKey: ECPublicKey) {
+    try {
+      writableDatabase.insertInto(TABLE_NAME)
+        .values(
+          KYBER_PREKEY to kyberPreKeyRowId,
+          SIGNED_KEY_ID to signedKeyId,
+          PUBLIC_KEY to publicKey.serialize()
+        )
+        .run(conflictStrategy = SQLiteDatabase.CONFLICT_ABORT)
+    } catch (e: SQLiteConstraintException) {
+      Log.w(TAG, "Found duplicate use of last-resort kyber key set.", e)
+      throw ReusedBaseKeyException(e)
+    }
+  }
+}

app/src/main/java/org/thoughtcrime/securesms/database/SignalDatabase.kt

@@ -81,6 +81,7 @@ open class SignalDatabase(private val context: Application, databaseSecret: Data
   val chatFoldersTable: ChatFolderTables = ChatFolderTables(context, this)
   val backupMediaSnapshotTable: BackupMediaSnapshotTable = BackupMediaSnapshotTable(context, this)
   val pollTable: PollTables = PollTables(context, this)
+  val lastResortKeyTuples: LastResortKeyTupleTable = LastResortKeyTupleTable(context, this)
 
   override fun onOpen(db: net.zetetic.database.sqlcipher.SQLiteDatabase) {
     db.setForeignKeyConstraintsEnabled(true)
@@ -150,6 +151,7 @@ open class SignalDatabase(private val context: Application, databaseSecret: Data
     executeStatements(db, ChatFolderTables.CREATE_TABLE)
     executeStatements(db, PollTables.CREATE_TABLE)
     db.execSQL(BackupMediaSnapshotTable.CREATE_TABLE)
+    db.execSQL(LastResortKeyTupleTable.CREATE_TABLE)
 
     executeStatements(db, RecipientTable.CREATE_INDEXS)
     executeStatements(db, MessageTable.CREATE_INDEXS)
@@ -590,5 +592,10 @@ open class SignalDatabase(private val context: Application, databaseSecret: Data
     @get:JvmName("polls")
     val polls: PollTables
       get() = instance!!.pollTable
+
+    @get:JvmStatic
+    @get:JvmName("lastResortKeyTuples")
+    val lastResortKeyTuples: LastResortKeyTupleTable
+      get() = instance!!.lastResortKeyTuples
   }
 }

app/src/main/java/org/thoughtcrime/securesms/database/helpers/SignalDatabaseMigrations.kt

@@ -147,6 +147,7 @@ import org.thoughtcrime.securesms.database.helpers.migration.V289_AddQuoteTarget
 import org.thoughtcrime.securesms.database.helpers.migration.V290_AddArchiveThumbnailTransferStateColumn
 import org.thoughtcrime.securesms.database.helpers.migration.V291_NullOutRemoteKeyIfEmpty
 import org.thoughtcrime.securesms.database.helpers.migration.V292_AddPollTables
+import org.thoughtcrime.securesms.database.helpers.migration.V293_LastResortKeyTupleTableMigration
 import org.thoughtcrime.securesms.database.SQLiteDatabase as SignalSqliteDatabase
 
 /**
@@ -299,10 +300,11 @@ object SignalDatabaseMigrations {
     289 to V289_AddQuoteTargetContentTypeColumn,
     290 to V290_AddArchiveThumbnailTransferStateColumn,
     291 to V291_NullOutRemoteKeyIfEmpty,
-    292 to V292_AddPollTables
+    292 to V292_AddPollTables,
+    293 to V293_LastResortKeyTupleTableMigration
   )
 
-  const val DATABASE_VERSION = 292
+  const val DATABASE_VERSION = 293
 
   @JvmStatic
   fun migrate(context: Application, db: SignalSqliteDatabase, oldVersion: Int, newVersion: Int) {

app/src/main/java/org/thoughtcrime/securesms/database/helpers/migration/V293_LastResortKeyTupleTableMigration.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.database.helpers.migration
+
+import android.app.Application
+import org.thoughtcrime.securesms.database.SQLiteDatabase
+
+/**
+ * Creates the LastResortKeyTuple table.
+ */
+@Suppress("ClassName")
+object V293_LastResortKeyTupleTableMigration : SignalDatabaseMigration {
+  override fun migrate(context: Application, db: SQLiteDatabase, oldVersion: Int, newVersion: Int) {
+    db.execSQL(
+      """
+      CREATE TABLE last_resort_key_tuple (
+        _id INTEGER PRIMARY KEY,
+        kyber_prekey_id INTEGER NOT NULL UNIQUE REFERENCES kyber_prekey (_id) ON DELETE CASCADE,
+        signed_key_id INTEGER NOT NULL,
+        public_key BLOB NOT NULL,
+        UNIQUE(kyber_prekey_id, signed_key_id, public_key)
+      )
+      """.trimIndent()
+    )
+  }
+}

app/src/main/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStore.kt

@@ -6,6 +6,8 @@
 package org.thoughtcrime.securesms.messages.protocol
 
 import org.signal.libsignal.protocol.InvalidKeyIdException
+import org.signal.libsignal.protocol.ReusedBaseKeyException
+import org.signal.libsignal.protocol.ecc.ECPublicKey
 import org.signal.libsignal.protocol.state.KyberPreKeyRecord
 import org.thoughtcrime.securesms.database.KyberPreKeyTable.KyberPreKey
 import org.thoughtcrime.securesms.database.SignalDatabase
@@ -25,7 +27,13 @@ class BufferedKyberPreKeyStore(private val selfServiceId: ServiceId) : SignalSer
   private var hasLoadedAll: Boolean = false
 
   /** The kyber prekeys that have been marked as removed (if they're not last resort). */
-  private val removedIfNotLastResort: MutableSet<Int> = mutableSetOf()
+  private val removedIfNotLastResort: MutableSet<Triple<Int, Int, ECPublicKey>> = mutableSetOf()
+
+  /** Tuples of last-resort key data we've already seen. */
+  private val lastResortKeyTuples: MutableSet<Triple<Int, Int, ECPublicKey>> = mutableSetOf()
+
+  /** A separate list of tuples to flush so we don't try to flush the same one multiple times */
+  private val unFlushedLastResortKeyTuples: MutableSet<Triple<Int, Int, ECPublicKey>> = mutableSetOf()
 
   @kotlin.jvm.Throws(InvalidKeyIdException::class)
   override fun loadKyberPreKey(kyberPreKeyId: Int): KyberPreKeyRecord {
@@ -63,16 +71,21 @@ class BufferedKyberPreKeyStore(private val selfServiceId: ServiceId) : SignalSer
     return store.containsKey(kyberPreKeyId)
   }
 
-  override fun markKyberPreKeyUsed(kyberPreKeyId: Int) {
+  override fun markKyberPreKeyUsed(kyberPreKeyId: Int, signedPreKeyId: Int, publicKey: ECPublicKey) {
     loadKyberPreKey(kyberPreKeyId)
 
     store[kyberPreKeyId]?.let {
       if (!it.lastResort) {
         store.remove(kyberPreKeyId)
+        removedIfNotLastResort += Triple(kyberPreKeyId, signedPreKeyId, publicKey)
+      } else {
+        if (!lastResortKeyTuples.add(Triple(kyberPreKeyId, signedPreKeyId, publicKey))) {
+          throw ReusedBaseKeyException()
+        }
+
+        unFlushedLastResortKeyTuples += Triple(kyberPreKeyId, signedPreKeyId, publicKey)
       }
     }
-
-    removedIfNotLastResort += kyberPreKeyId
   }
 
   override fun removeKyberPreKey(kyberPreKeyId: Int) {
@@ -88,8 +101,13 @@ class BufferedKyberPreKeyStore(private val selfServiceId: ServiceId) : SignalSer
   }
 
   fun flushToDisk(persistentStore: SignalServiceAccountDataStore) {
-    for (id in removedIfNotLastResort) {
-      persistentStore.markKyberPreKeyUsed(id)
+    val tuples = removedIfNotLastResort + unFlushedLastResortKeyTuples
+    unFlushedLastResortKeyTuples.clear()
+
+    for ((key, signedKey, publicKey) in tuples) {
+      persistentStore.markKyberPreKeyUsed(key, signedKey, publicKey)
     }
+
+    unFlushedLastResortKeyTuples.clear()
   }
 }

app/src/main/java/org/thoughtcrime/securesms/messages/protocol/BufferedSignalServiceAccountDataStore.kt

@@ -1,8 +1,10 @@
 package org.thoughtcrime.securesms.messages.protocol
 
+import org.signal.core.util.withinTransaction
 import org.signal.libsignal.protocol.IdentityKey
 import org.signal.libsignal.protocol.IdentityKeyPair
 import org.signal.libsignal.protocol.SignalProtocolAddress
+import org.signal.libsignal.protocol.ecc.ECPublicKey
 import org.signal.libsignal.protocol.groups.state.SenderKeyRecord
 import org.signal.libsignal.protocol.state.IdentityKeyStore
 import org.signal.libsignal.protocol.state.IdentityKeyStore.IdentityChange
@@ -10,6 +12,7 @@ import org.signal.libsignal.protocol.state.KyberPreKeyRecord
 import org.signal.libsignal.protocol.state.PreKeyRecord
 import org.signal.libsignal.protocol.state.SessionRecord
 import org.signal.libsignal.protocol.state.SignedPreKeyRecord
+import org.thoughtcrime.securesms.database.SignalDatabase
 import org.thoughtcrime.securesms.keyvalue.SignalStore
 import org.whispersystems.signalservice.api.SignalServiceAccountDataStore
 import org.whispersystems.signalservice.api.push.DistributionId
@@ -138,8 +141,8 @@ class BufferedSignalServiceAccountDataStore(selfServiceId: ServiceId) : SignalSe
     return kyberPreKeyStore.containsKyberPreKey(kyberPreKeyId)
   }
 
-  override fun markKyberPreKeyUsed(kyberPreKeyId: Int) {
-    return kyberPreKeyStore.markKyberPreKeyUsed(kyberPreKeyId)
+  override fun markKyberPreKeyUsed(kyberPreKeyId: Int, signedPreKeyId: Int, publicKey: ECPublicKey) {
+    return kyberPreKeyStore.markKyberPreKeyUsed(kyberPreKeyId, signedPreKeyId, publicKey)
   }
 
   override fun deleteAllStaleOneTimeEcPreKeys(threshold: Long, minCount: Int) {
@@ -199,11 +202,13 @@ class BufferedSignalServiceAccountDataStore(selfServiceId: ServiceId) : SignalSe
   }
 
   fun flushToDisk(persistentStore: SignalServiceAccountDataStore) {
-    identityStore.flushToDisk(persistentStore)
-    oneTimePreKeyStore.flushToDisk(persistentStore)
-    kyberPreKeyStore.flushToDisk(persistentStore)
-    signedPreKeyStore.flushToDisk(persistentStore)
-    sessionStore.flushToDisk(persistentStore)
-    senderKeyStore.flushToDisk(persistentStore)
+    SignalDatabase.writableDatabase.withinTransaction {
+      identityStore.flushToDisk(persistentStore)
+      oneTimePreKeyStore.flushToDisk(persistentStore)
+      kyberPreKeyStore.flushToDisk(persistentStore)
+      signedPreKeyStore.flushToDisk(persistentStore)
+      sessionStore.flushToDisk(persistentStore)
+      senderKeyStore.flushToDisk(persistentStore)
+    }
   }
 }