From c1d43b31e63c536cafe8bbad5abdd76ef61b55ce Mon Sep 17 00:00:00 2001 From: Greyson Parrelli Date: Thu, 24 Sep 2026 10:58:22 -0400 Subject: [PATCH] Move BufferedKyberPreKeyStoreTest to a unit test. --- .../protocol/BufferedKyberPreKeyStoreTest.kt | 79 ------------- .../securesms/util/KyberPreKeysTestUtil.kt | 71 ------------ .../protocol/BufferedKyberPreKeyStoreTest.kt | 108 ++++++++++++++++++ 3 files changed, 108 insertions(+), 150 deletions(-) delete mode 100644 app/src/androidTest/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStoreTest.kt delete mode 100644 app/src/androidTest/java/org/thoughtcrime/securesms/util/KyberPreKeysTestUtil.kt create mode 100644 app/src/test/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStoreTest.kt 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 deleted file mode 100644 index ed63acd9f1..0000000000 --- a/app/src/androidTest/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStoreTest.kt +++ /dev/null @@ -1,79 +0,0 @@ -/* - * 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.core.models.ServiceId -import org.signal.libsignal.protocol.ReusedBaseKeyException -import org.thoughtcrime.securesms.keyvalue.SignalStore -import org.thoughtcrime.securesms.testing.SignalDatabaseRule -import org.thoughtcrime.securesms.util.KyberPreKeysTestUtil - -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 deleted file mode 100644 index 059e8b92db..0000000000 --- a/app/src/androidTest/java/org/thoughtcrime/securesms/util/KyberPreKeysTestUtil.kt +++ /dev/null @@ -1,71 +0,0 @@ -/* - * 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.models.ServiceId -import org.signal.core.models.ServiceId.ACI -import org.signal.core.models.ServiceId.PNI -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 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/test/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStoreTest.kt b/app/src/test/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStoreTest.kt new file mode 100644 index 0000000000..2d5029e78c --- /dev/null +++ b/app/src/test/java/org/thoughtcrime/securesms/messages/protocol/BufferedKyberPreKeyStoreTest.kt @@ -0,0 +1,108 @@ +/* + * Copyright 2025 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.thoughtcrime.securesms.messages.protocol + +import android.app.Application +import io.mockk.mockk +import io.mockk.verify +import org.junit.Rule +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.RobolectricTestRunner +import org.robolectric.annotation.Config +import org.signal.core.models.ServiceId.ACI +import org.signal.libsignal.protocol.ReusedBaseKeyException +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.SignalDatabase +import org.thoughtcrime.securesms.testutil.MockAppDependenciesRule +import org.thoughtcrime.securesms.testutil.SignalDatabaseRule +import org.whispersystems.signalservice.api.SignalServiceAccountDataStore +import java.util.UUID + +@RunWith(RobolectricTestRunner::class) +@Config(manifest = Config.NONE, application = Application::class) +class BufferedKyberPreKeyStoreTest { + + @get:Rule + val appDependencies = MockAppDependenciesRule() + + @get:Rule + val signalDatabaseRule = SignalDatabaseRule() + + private val aci: ACI = ACI.from(UUID.randomUUID()) + private val testSubject = BufferedKyberPreKeyStore(aci) + private val dataStore: SignalServiceAccountDataStore = mockk(relaxed = true) + + @Test + fun givenALastResortKey_whenIMarkKyberPreKeyUsed_thenIExpectNoIssues() { + insertLastResortKey(id = 1) + val publicKey = generateECPublicKey() + + testSubject.markKyberPreKeyUsed( + kyberPreKeyId = 1, + signedPreKeyId = 2, + publicKey = publicKey + ) + } + + @Test(expected = ReusedBaseKeyException::class) + fun givenALastResortKey_whenIMarkKyberPreKeyUsedTwice_thenIExpectException() { + insertLastResortKey(id = 1) + val publicKey = generateECPublicKey() + + testSubject.markKyberPreKeyUsed( + kyberPreKeyId = 1, + signedPreKeyId = 2, + publicKey = publicKey + ) + + testSubject.markKyberPreKeyUsed( + kyberPreKeyId = 1, + signedPreKeyId = 2, + publicKey = publicKey + ) + } + + @Test + fun givenAMarkedLastResortKey_whenIFlushTwice_thenIExpectOnlyOneWrite() { + insertLastResortKey(id = 1) + val publicKey = generateECPublicKey() + + testSubject.markKyberPreKeyUsed( + kyberPreKeyId = 1, + signedPreKeyId = 2, + publicKey = publicKey + ) + + testSubject.flushToDisk(dataStore) + testSubject.flushToDisk(dataStore) + + verify(exactly = 1) { dataStore.markKyberPreKeyUsed(1, 2, publicKey) } + } + + private fun insertLastResortKey(id: Int) { + val kemKeyPair = KEMKeyPair.generate(KEMKeyType.KYBER_1024) + SignalDatabase.kyberPreKeys.insert( + serviceId = aci, + keyId = id, + record = KyberPreKeyRecord( + id, + System.currentTimeMillis(), + kemKeyPair, + ECKeyPair.generate().privateKey.calculateSignature(kemKeyPair.publicKey.serialize()) + ), + lastResort = true + ) + } + + private fun generateECPublicKey(): ECPublicKey { + return ECKeyPair.generate().publicKey + } +}