From 016736c45509d6e1959b9ffe964131a10de30dd7 Mon Sep 17 00:00:00 2001 From: Nicholas Date: Mon, 20 Nov 2023 10:26:20 -0500 Subject: [PATCH] Encrypting for multiple senders benchmark. --- .../microbenchmark/ProtocolBenchmarks.kt | 51 ++++++++++++++--- .../InMemorySignalServiceAccountDataStore.kt | 4 +- .../java/org/signal/util/SignalClient.kt | 57 ++++++++++++++----- 3 files changed, 88 insertions(+), 24 deletions(-) diff --git a/microbenchmark/src/androidTest/java/org/signal/microbenchmark/ProtocolBenchmarks.kt b/microbenchmark/src/androidTest/java/org/signal/microbenchmark/ProtocolBenchmarks.kt index 3b96c52189..d22e4529bd 100644 --- a/microbenchmark/src/androidTest/java/org/signal/microbenchmark/ProtocolBenchmarks.kt +++ b/microbenchmark/src/androidTest/java/org/signal/microbenchmark/ProtocolBenchmarks.kt @@ -11,6 +11,8 @@ import org.junit.runner.RunWith import org.signal.libsignal.protocol.logging.SignalProtocolLogger import org.signal.libsignal.protocol.logging.SignalProtocolLoggerProvider import org.signal.util.SignalClient +import org.whispersystems.signalservice.api.push.DistributionId +import java.util.Optional /** * Benchmarks for decrypting messages. @@ -74,16 +76,49 @@ class ProtocolBenchmarks { } } + @Test + fun multi_encrypt_sealedSender() { + val recipientCount = 10 + val clients = buildAndInitializeClients(recipientCount) + val alice = clients.first() + val others = clients.filterNot { it == alice } + val distributionId = DistributionId.create() + + clients.forEach { + it.initializedGroupSession(distributionId) + } + + benchmarkRule.measureRepeated { + alice.multiEncryptSealedSender(distributionId, others, Optional.empty()) + } + } + private fun buildAndInitializeClients(): Pair { - val alice = SignalClient() - val bob = SignalClient() + val clients = buildAndInitializeClients(2) + return clients[0] to clients[1] + } - // Do initial prekey dance - alice.initializeSession(bob) - bob.initializeSession(alice) - alice.decryptMessage(bob.encryptUnsealedSender(alice)) - bob.decryptMessage(alice.encryptUnsealedSender(bob)) + private fun buildAndInitializeClients(recipientCount: Int): List { + val clients = ArrayList(recipientCount) + for (n in 1..recipientCount) { + clients.add(SignalClient()) + } - return alice to bob + clients.forEach { alice -> + clients.filterNot { it == alice }.forEach { bob -> + alice.initializeSession(bob) + bob.initializeSession(alice) + + alice.decryptMessage(bob.encryptUnsealedSender(alice)) + + bob.decryptMessage(alice.encryptUnsealedSender(bob)) + + alice.decryptMessage(bob.encryptSealedSender(alice)) + + bob.decryptMessage(alice.encryptSealedSender(bob)) + } + } + + return clients } } diff --git a/microbenchmark/src/androidTest/java/org/signal/util/InMemorySignalServiceAccountDataStore.kt b/microbenchmark/src/androidTest/java/org/signal/util/InMemorySignalServiceAccountDataStore.kt index b2166ba898..28d7037b82 100644 --- a/microbenchmark/src/androidTest/java/org/signal/util/InMemorySignalServiceAccountDataStore.kt +++ b/microbenchmark/src/androidTest/java/org/signal/util/InMemorySignalServiceAccountDataStore.kt @@ -118,8 +118,8 @@ class InMemorySignalServiceAccountDataStore : SignalServiceAccountDataStore { senderKeys[SenderKeyLocator(sender, distributionId)] = record } - override fun loadSenderKey(sender: SignalProtocolAddress, distributionId: UUID): SenderKeyRecord { - return senderKeys[SenderKeyLocator(sender, distributionId)]!! + override fun loadSenderKey(sender: SignalProtocolAddress, distributionId: UUID): SenderKeyRecord? { + return senderKeys[SenderKeyLocator(sender, distributionId)] } override fun loadKyberPreKey(kyberPreKeyId: Int): KyberPreKeyRecord { diff --git a/microbenchmark/src/androidTest/java/org/signal/util/SignalClient.kt b/microbenchmark/src/androidTest/java/org/signal/util/SignalClient.kt index e299844222..96e3ec1a51 100644 --- a/microbenchmark/src/androidTest/java/org/signal/util/SignalClient.kt +++ b/microbenchmark/src/androidTest/java/org/signal/util/SignalClient.kt @@ -12,6 +12,8 @@ import org.signal.libsignal.protocol.SignalProtocolAddress import org.signal.libsignal.protocol.ecc.Curve import org.signal.libsignal.protocol.ecc.ECKeyPair import org.signal.libsignal.protocol.ecc.ECPublicKey +import org.signal.libsignal.protocol.groups.GroupSessionBuilder +import org.signal.libsignal.protocol.message.SenderKeyDistributionMessage import org.signal.libsignal.protocol.state.PreKeyBundle import org.signal.libsignal.protocol.state.PreKeyRecord import org.signal.libsignal.protocol.state.SignedPreKeyRecord @@ -19,8 +21,10 @@ import org.whispersystems.signalservice.api.SignalServiceAccountDataStore import org.whispersystems.signalservice.api.SignalSessionLock import org.whispersystems.signalservice.api.crypto.ContentHint import org.whispersystems.signalservice.api.crypto.EnvelopeContent +import org.whispersystems.signalservice.api.crypto.SignalGroupSessionBuilder import org.whispersystems.signalservice.api.crypto.SignalServiceCipher import org.whispersystems.signalservice.api.crypto.UnidentifiedAccess +import org.whispersystems.signalservice.api.push.DistributionId import org.whispersystems.signalservice.api.push.ServiceId.ACI import org.whispersystems.signalservice.api.push.SignalServiceAddress import org.whispersystems.signalservice.internal.push.Content @@ -43,20 +47,13 @@ class SignalClient { private val trustRoot: ECKeyPair = Curve.generateKeyPair() } + private val lock = TestSessionLock() + private val aci: ACI = ACI.from(UUID.randomUUID()) private val store: SignalServiceAccountDataStore = InMemorySignalServiceAccountDataStore() - private val preKeyBundle: PreKeyBundle = let { - val preKeyRecord = PreKeyRecord(1, Curve.generateKeyPair()) - val signedPreKeyPair = Curve.generateKeyPair() - val signedPreKeySignature = Curve.calculateSignature(store.identityKeyPair.privateKey, signedPreKeyPair.publicKey.serialize()) - - store.storePreKey(1, preKeyRecord) - store.storeSignedPreKey(1, SignedPreKeyRecord(1, System.currentTimeMillis(), signedPreKeyPair, signedPreKeySignature)) - - PreKeyBundle(1, 1, 1, preKeyRecord.keyPair.publicKey, 1, signedPreKeyPair.publicKey, signedPreKeySignature, store.identityKeyPair.publicKey) - } + private var prekeyIndex = 0 private val unidentifiedAccessKey: ByteArray = Util.getSecretBytes(32) @@ -69,15 +66,19 @@ class SignalClient { expires = Long.MAX_VALUE ) - private val cipher = SignalServiceCipher(SignalServiceAddress(aci), 1, store, TestSessionLock(), CertificateValidator(trustRoot.publicKey)) + private val cipher = SignalServiceCipher(SignalServiceAddress(aci), 1, store, lock, CertificateValidator(trustRoot.publicKey)) /** - * Sets up sessions using the [to] client's [preKeyBundle]. Note that you can only initialize a client once - * since we currently only make a single prekey bundle. + * Sets up sessions using the [to] client's [preKeyBundles]. Note that you can only initialize a client up to 1,000 times because that's how many prekeys we have. */ fun initializeSession(to: SignalClient) { val address = SignalProtocolAddress(to.aci.toString(), 1) - SessionBuilder(store, address).process(to.preKeyBundle) + SessionBuilder(store, address).process(to.createPreKeyBundle()) + } + + fun initializedGroupSession(distributionId: DistributionId): SenderKeyDistributionMessage { + val self = SignalProtocolAddress(aci.toString(), 1) + return SignalGroupSessionBuilder(lock, GroupSessionBuilder(store)).create(self, distributionId.asUuid()) } fun encryptUnsealedSender(to: SignalClient): Envelope { @@ -142,9 +143,37 @@ class SignalClient { ) } + fun multiEncryptSealedSender(distributionId: DistributionId, others: List, groupId: Optional): ByteArray { + val sentTimestamp = System.currentTimeMillis() + + val content = Content( + dataMessage = DataMessage( + body = "Test Message", + timestamp = sentTimestamp + ) + ) + val destinations = others.map { bob -> + SignalProtocolAddress(bob.aci.toString(), 1) + } + + return cipher.encryptForGroup(distributionId, destinations, senderCertificate, content.encode(), ContentHint.DEFAULT, groupId) + } + fun decryptMessage(envelope: Envelope) { cipher.decrypt(envelope, System.currentTimeMillis()) } + + private fun createPreKeyBundle(): PreKeyBundle { + val prekeyId = prekeyIndex++ + val preKeyRecord = PreKeyRecord(prekeyId, Curve.generateKeyPair()) + val signedPreKeyPair = Curve.generateKeyPair() + val signedPreKeySignature = Curve.calculateSignature(store.identityKeyPair.privateKey, signedPreKeyPair.publicKey.serialize()) + + store.storePreKey(prekeyId, preKeyRecord) + store.storeSignedPreKey(prekeyId, SignedPreKeyRecord(prekeyId, System.currentTimeMillis(), signedPreKeyPair, signedPreKeySignature)) + + return PreKeyBundle(prekeyId, prekeyId, prekeyId, preKeyRecord.keyPair.publicKey, prekeyId, signedPreKeyPair.publicKey, signedPreKeySignature, store.identityKeyPair.publicKey) + } } private fun createCertificateFor(trustRoot: ECKeyPair, uuid: UUID, e164: String, deviceId: Int, identityKey: ECPublicKey, expires: Long): SenderCertificate {