From 1b9695cb9807d7896dcf4582ab1da2b686ccf4ca Mon Sep 17 00:00:00 2001 From: Alex Hart Date: Wed, 1 Oct 2025 16:08:01 -0300 Subject: [PATCH] Reject last-use kyber key sets that we've seen before. --- .../database/KyberPreKeyTableTest.kt | 78 +++++++++--------- .../protocol/BufferedKyberPreKeyStoreTest.kt | 79 +++++++++++++++++++ .../securesms/util/KyberPreKeysTestUtil.kt | 71 +++++++++++++++++ .../crypto/storage/SignalKyberPreKeyStore.kt | 5 +- .../SignalServiceAccountDataStoreImpl.java | 5 +- .../securesms/database/KyberPreKeyTable.kt | 35 ++++++-- .../database/LastResortKeyTupleTable.kt | 62 +++++++++++++++ .../securesms/database/SignalDatabase.kt | 7 ++ .../helpers/SignalDatabaseMigrations.kt | 6 +- .../V293_LastResortKeyTupleTableMigration.kt | 29 +++++++ .../protocol/BufferedKyberPreKeyStore.kt | 30 +++++-- .../BufferedSignalServiceAccountDataStore.kt | 21 +++-- 12 files changed, 362 insertions(+), 66 deletions(-) create mode 100644 app/src/androidTest/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStoreTest.kt create mode 100644 app/src/androidTest/java/org/thoughtcrime/securesms/util/KyberPreKeysTestUtil.kt create mode 100644 app/src/main/java/org/thoughtcrime/securesms/database/LastResortKeyTupleTable.kt create mode 100644 app/src/main/java/org/thoughtcrime/securesms/database/helpers/migration/V293_LastResortKeyTupleTableMigration.kt diff --git a/app/src/androidTest/java/org/thoughtcrime/securesms/database/KyberPreKeyTableTest.kt b/app/src/androidTest/java/org/thoughtcrime/securesms/database/KyberPreKeyTableTest.kt index 481dd88085..99d3b1bb38 100644 --- a/app/src/androidTest/java/org/thoughtcrime/securesms/database/KyberPreKeyTableTest.kt +++ b/app/src/androidTest/java/org/thoughtcrime/securesms/database/KyberPreKeyTableTest.kt @@ -9,15 +9,10 @@ import org.junit.Assert.assertEquals import org.junit.Assert.assertNotNull import org.junit.Assert.assertNull import org.junit.Test -import org.signal.core.util.readToSingleObject -import org.signal.core.util.requireLongOrNull -import org.signal.core.util.select -import org.signal.core.util.update -import org.signal.libsignal.protocol.ecc.ECKeyPair -import org.signal.libsignal.protocol.kem.KEMKeyPair -import org.signal.libsignal.protocol.kem.KEMKeyType -import org.signal.libsignal.protocol.state.KyberPreKeyRecord -import org.whispersystems.signalservice.api.push.ServiceId +import org.signal.libsignal.protocol.ReusedBaseKeyException +import org.thoughtcrime.securesms.util.KyberPreKeysTestUtil.generateECPublicKey +import org.thoughtcrime.securesms.util.KyberPreKeysTestUtil.getStaleTime +import org.thoughtcrime.securesms.util.KyberPreKeysTestUtil.insertTestRecord import org.whispersystems.signalservice.api.push.ServiceId.ACI import org.whispersystems.signalservice.api.push.ServiceId.PNI import java.util.UUID @@ -142,42 +137,43 @@ class KyberPreKeyTableTest { assertNotNull(getStaleTime(aci, 3)) } - private fun insertTestRecord(account: ServiceId, id: Int, staleTime: Long = 0, lastResort: Boolean = false) { - val kemKeyPair = KEMKeyPair.generate(KEMKeyType.KYBER_1024) - SignalDatabase.kyberPreKeys.insert( - serviceId = account, - keyId = id, - record = KyberPreKeyRecord( - id, - System.currentTimeMillis(), - kemKeyPair, - ECKeyPair.generate().privateKey.calculateSignature(kemKeyPair.publicKey.serialize()) - ), - lastResort = lastResort + @Test(expected = ReusedBaseKeyException::class) + fun handleMarkKyberPreKeyUsed_doesNotAllowDuplicateLastResortKeyEntries() { + insertTestRecord(aci, id = 1, staleTime = 10, lastResort = true) + val publicKey = generateECPublicKey() + + SignalDatabase.kyberPreKeys.handleMarkKyberPreKeyUsed( + serviceId = aci, + kyberPreKeyId = 1, + signedPreKeyId = 1, + baseKey = publicKey ) - val count = SignalDatabase.rawDatabase - .update(KyberPreKeyTable.TABLE_NAME) - .values(KyberPreKeyTable.STALE_TIMESTAMP to staleTime) - .where("${KyberPreKeyTable.ACCOUNT_ID} = ? AND ${KyberPreKeyTable.KEY_ID} = $id", account.toAccountId()) - .run() - - assertEquals(1, count) + SignalDatabase.kyberPreKeys.handleMarkKyberPreKeyUsed( + serviceId = aci, + kyberPreKeyId = 1, + signedPreKeyId = 1, + baseKey = publicKey + ) } - private fun getStaleTime(account: ServiceId, id: Int): Long? { - return SignalDatabase.rawDatabase - .select(KyberPreKeyTable.STALE_TIMESTAMP) - .from(KyberPreKeyTable.TABLE_NAME) - .where("${KyberPreKeyTable.ACCOUNT_ID} = ? AND ${KyberPreKeyTable.KEY_ID} = $id", account.toAccountId()) - .run() - .readToSingleObject { it.requireLongOrNull(KyberPreKeyTable.STALE_TIMESTAMP) } - } + @Test + fun handleMarkKyberPreKeyUsed_allowDuplicateNonLastResortKeyEntries() { + insertTestRecord(aci, id = 1, staleTime = 10, lastResort = false) + val publicKey = generateECPublicKey() - private fun ServiceId.toAccountId(): String { - return when (this) { - is ACI -> this.toString() - is PNI -> KyberPreKeyTable.PNI_ACCOUNT_ID - } + SignalDatabase.kyberPreKeys.handleMarkKyberPreKeyUsed( + serviceId = aci, + kyberPreKeyId = 1, + signedPreKeyId = 1, + baseKey = publicKey + ) + + SignalDatabase.kyberPreKeys.handleMarkKyberPreKeyUsed( + serviceId = aci, + kyberPreKeyId = 1, + signedPreKeyId = 1, + baseKey = publicKey + ) } } diff --git a/app/src/androidTest/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStoreTest.kt b/app/src/androidTest/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStoreTest.kt new file mode 100644 index 0000000000..a230d61c10 --- /dev/null +++ b/app/src/androidTest/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStoreTest.kt @@ -0,0 +1,79 @@ +/* + * Copyright 2025 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.thoughtcrime.securesms.messages.protocol + +import org.junit.Before +import org.junit.Rule +import org.junit.Test +import org.signal.libsignal.protocol.ReusedBaseKeyException +import org.thoughtcrime.securesms.keyvalue.SignalStore +import org.thoughtcrime.securesms.testing.SignalDatabaseRule +import org.thoughtcrime.securesms.util.KyberPreKeysTestUtil +import org.whispersystems.signalservice.api.push.ServiceId + +class BufferedKyberPreKeyStoreTest { + + @get:Rule + val harness = SignalDatabaseRule() + + private lateinit var aci: ServiceId + private lateinit var testSubject: BufferedKyberPreKeyStore + private lateinit var dataStore: BufferedSignalServiceAccountDataStore + + @Before + fun setUp() { + SignalStore.account.generateAciIdentityKeyIfNecessary() + + aci = harness.localAci + testSubject = BufferedKyberPreKeyStore(aci) + dataStore = BufferedSignalServiceAccountDataStore(aci) + } + + @Test + fun givenALastResortKey_whenIMarkKyberPreKeyUsed_thenIExpectNoIssues() { + KyberPreKeysTestUtil.insertTestRecord(aci, 1, lastResort = true) + val publicKey = KyberPreKeysTestUtil.generateECPublicKey() + + testSubject.markKyberPreKeyUsed( + kyberPreKeyId = 1, + signedPreKeyId = 2, + publicKey = publicKey + ) + } + + @Test(expected = ReusedBaseKeyException::class) + fun givenALastResortKey_whenIMarkKyberPreKeyUsedTwice_thenIExpectException() { + KyberPreKeysTestUtil.insertTestRecord(aci, 1, lastResort = true) + val publicKey = KyberPreKeysTestUtil.generateECPublicKey() + + testSubject.markKyberPreKeyUsed( + kyberPreKeyId = 1, + signedPreKeyId = 2, + publicKey = publicKey + ) + + testSubject.markKyberPreKeyUsed( + kyberPreKeyId = 1, + signedPreKeyId = 2, + publicKey = publicKey + ) + } + + @Test + fun givenAMarkedLastResortKey_whenIFlushTwice_thenIExpectNoIssues() { + KyberPreKeysTestUtil.insertTestRecord(aci, 1, lastResort = true) + val publicKey = KyberPreKeysTestUtil.generateECPublicKey() + + testSubject.markKyberPreKeyUsed( + kyberPreKeyId = 1, + signedPreKeyId = 2, + publicKey = publicKey + ) + + testSubject.flushToDisk(dataStore) + testSubject.flushToDisk(dataStore) + } +} diff --git a/app/src/androidTest/java/org/thoughtcrime/securesms/util/KyberPreKeysTestUtil.kt b/app/src/androidTest/java/org/thoughtcrime/securesms/util/KyberPreKeysTestUtil.kt new file mode 100644 index 0000000000..803630235a --- /dev/null +++ b/app/src/androidTest/java/org/thoughtcrime/securesms/util/KyberPreKeysTestUtil.kt @@ -0,0 +1,71 @@ +/* + * Copyright 2025 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.thoughtcrime.securesms.util + +import org.junit.Assert.assertEquals +import org.signal.core.util.readToSingleObject +import org.signal.core.util.requireLongOrNull +import org.signal.core.util.select +import org.signal.core.util.update +import org.signal.libsignal.protocol.ecc.ECKeyPair +import org.signal.libsignal.protocol.ecc.ECPublicKey +import org.signal.libsignal.protocol.kem.KEMKeyPair +import org.signal.libsignal.protocol.kem.KEMKeyType +import org.signal.libsignal.protocol.state.KyberPreKeyRecord +import org.thoughtcrime.securesms.database.KyberPreKeyTable +import org.thoughtcrime.securesms.database.SignalDatabase +import org.whispersystems.signalservice.api.push.ServiceId +import org.whispersystems.signalservice.api.push.ServiceId.ACI +import org.whispersystems.signalservice.api.push.ServiceId.PNI +import java.security.SecureRandom + +object KyberPreKeysTestUtil { + fun insertTestRecord(account: ServiceId, id: Int, staleTime: Long = 0, lastResort: Boolean = false) { + val kemKeyPair = KEMKeyPair.generate(KEMKeyType.KYBER_1024) + SignalDatabase.kyberPreKeys.insert( + serviceId = account, + keyId = id, + record = KyberPreKeyRecord( + id, + System.currentTimeMillis(), + kemKeyPair, + ECKeyPair.generate().privateKey.calculateSignature(kemKeyPair.publicKey.serialize()) + ), + lastResort = lastResort + ) + + val count = SignalDatabase.rawDatabase + .update(KyberPreKeyTable.TABLE_NAME) + .values(KyberPreKeyTable.STALE_TIMESTAMP to staleTime) + .where("${KyberPreKeyTable.ACCOUNT_ID} = ? AND ${KyberPreKeyTable.KEY_ID} = $id", account.toAccountId()) + .run() + + assertEquals(1, count) + } + + fun getStaleTime(account: ServiceId, id: Int): Long? { + return SignalDatabase.rawDatabase + .select(KyberPreKeyTable.STALE_TIMESTAMP) + .from(KyberPreKeyTable.TABLE_NAME) + .where("${KyberPreKeyTable.ACCOUNT_ID} = ? AND ${KyberPreKeyTable.KEY_ID} = $id", account.toAccountId()) + .run() + .readToSingleObject { it.requireLongOrNull(KyberPreKeyTable.STALE_TIMESTAMP) } + } + + fun generateECPublicKey(): ECPublicKey { + val byteArray = ByteArray(ECPublicKey.KEY_SIZE - 1) + SecureRandom().nextBytes(byteArray) + + return ECPublicKey.fromPublicKeyBytes(byteArray) + } + + private fun ServiceId.toAccountId(): String { + return when (this) { + is ACI -> this.toString() + is PNI -> KyberPreKeyTable.PNI_ACCOUNT_ID + } + } +} diff --git a/app/src/main/java/org/thoughtcrime/securesms/crypto/storage/SignalKyberPreKeyStore.kt b/app/src/main/java/org/thoughtcrime/securesms/crypto/storage/SignalKyberPreKeyStore.kt index c00f6b5092..5fb8a46027 100644 --- a/app/src/main/java/org/thoughtcrime/securesms/crypto/storage/SignalKyberPreKeyStore.kt +++ b/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) } } diff --git a/app/src/main/java/org/thoughtcrime/securesms/crypto/storage/SignalServiceAccountDataStoreImpl.java b/app/src/main/java/org/thoughtcrime/securesms/crypto/storage/SignalServiceAccountDataStoreImpl.java index 0928ffe176..ac79582e7c 100644 --- a/app/src/main/java/org/thoughtcrime/securesms/crypto/storage/SignalServiceAccountDataStoreImpl.java +++ b/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 diff --git a/app/src/main/java/org/thoughtcrime/securesms/database/KyberPreKeyTable.kt b/app/src/main/java/org/thoughtcrime/securesms/database/KyberPreKeyTable.kt index 74141ae3ea..d96d8047f3 100644 --- a/app/src/main/java/org/thoughtcrime/securesms/database/KyberPreKeyTable.kt +++ b/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) { diff --git a/app/src/main/java/org/thoughtcrime/securesms/database/LastResortKeyTupleTable.kt b/app/src/main/java/org/thoughtcrime/securesms/database/LastResortKeyTupleTable.kt new file mode 100644 index 0000000000..a70c3f8467 --- /dev/null +++ b/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) + } + } +} diff --git a/app/src/main/java/org/thoughtcrime/securesms/database/SignalDatabase.kt b/app/src/main/java/org/thoughtcrime/securesms/database/SignalDatabase.kt index 2b4e98d54d..8089b55751 100644 --- a/app/src/main/java/org/thoughtcrime/securesms/database/SignalDatabase.kt +++ b/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 } } diff --git a/app/src/main/java/org/thoughtcrime/securesms/database/helpers/SignalDatabaseMigrations.kt b/app/src/main/java/org/thoughtcrime/securesms/database/helpers/SignalDatabaseMigrations.kt index e0439ebcb9..a5836fb25e 100644 --- a/app/src/main/java/org/thoughtcrime/securesms/database/helpers/SignalDatabaseMigrations.kt +++ b/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) { diff --git a/app/src/main/java/org/thoughtcrime/securesms/database/helpers/migration/V293_LastResortKeyTupleTableMigration.kt b/app/src/main/java/org/thoughtcrime/securesms/database/helpers/migration/V293_LastResortKeyTupleTableMigration.kt new file mode 100644 index 0000000000..621bc267f0 --- /dev/null +++ b/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() + ) + } +} diff --git a/app/src/main/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStore.kt b/app/src/main/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStore.kt index ba9c6ee146..80b7f31373 100644 --- a/app/src/main/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStore.kt +++ b/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 = mutableSetOf() + private val removedIfNotLastResort: MutableSet> = mutableSetOf() + + /** Tuples of last-resort key data we've already seen. */ + private val lastResortKeyTuples: MutableSet> = mutableSetOf() + + /** A separate list of tuples to flush so we don't try to flush the same one multiple times */ + private val unFlushedLastResortKeyTuples: MutableSet> = 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() } } diff --git a/app/src/main/java/org/thoughtcrime/securesms/messages/protocol/BufferedSignalServiceAccountDataStore.kt b/app/src/main/java/org/thoughtcrime/securesms/messages/protocol/BufferedSignalServiceAccountDataStore.kt index aedb9dc956..d09def8254 100644 --- a/app/src/main/java/org/thoughtcrime/securesms/messages/protocol/BufferedSignalServiceAccountDataStore.kt +++ b/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) + } } }