From 501ef69f97903d5a87a442d35c5faedfbe664565 Mon Sep 17 00:00:00 2001
From: Cody Henthorne <cody@signal.org>
Date: Fri, 13 Mar 2026 15:57:51 -0400
Subject: [PATCH] Fix session establishment in message processing benchmark
 tests.

---
 .../benchmark/BenchmarkCommandReceiver.kt     | 15 ++++++--
 .../org/signal/benchmark/setup/OtherClient.kt | 34 +++++++++++++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/app/src/benchmarkShared/java/org/signal/benchmark/BenchmarkCommandReceiver.kt b/app/src/benchmarkShared/java/org/signal/benchmark/BenchmarkCommandReceiver.kt
index 90037c3d37..e18c22ba34 100644
--- a/app/src/benchmarkShared/java/org/signal/benchmark/BenchmarkCommandReceiver.kt
+++ b/app/src/benchmarkShared/java/org/signal/benchmark/BenchmarkCommandReceiver.kt
@@ -76,7 +76,7 @@ class BenchmarkCommandReceiver : BroadcastReceiver() {
 
     runBlocking {
       launch(Dispatchers.IO) {
-        Log.i(TAG, "Sending initial message form Bob to establish session.")
+        Log.i(TAG, "Sending initial message from Bob to establish session.")
         BenchmarkWebSocketConnection.addPendingMessages(listOf(encryptedEnvelope.toWebSocketPayload()))
         BenchmarkWebSocketConnection.releaseMessages()
 
@@ -85,6 +85,10 @@ class BenchmarkCommandReceiver : BroadcastReceiver() {
       }
     }
 
+    // Complete the session handshake so both sides have a proper double-ratchet session
+    Log.i(TAG, "Completing session handshake with reply.")
+    client.completeSession()
+
     // Have Bob generate N messages that will be received by Alice
     val messageCount = 500
     val envelopes = client.generateInboundEnvelopes(messageCount)
@@ -103,7 +107,7 @@ class BenchmarkCommandReceiver : BroadcastReceiver() {
 
     runBlocking {
       launch(Dispatchers.IO) {
-        Log.i(TAG, "Sending initial group messages from client to establish sessions.")
+        Log.i(TAG, "Sending initial group messages from clients to establish sessions.")
         BenchmarkWebSocketConnection.addPendingMessages(encryptedEnvelopes.map { it.toWebSocketPayload() })
         BenchmarkWebSocketConnection.releaseMessages()
 
@@ -112,6 +116,10 @@ class BenchmarkCommandReceiver : BroadcastReceiver() {
       }
     }
 
+    // Complete session handshakes so both sides have proper double-ratchet sessions
+    Log.i(TAG, "Completing session handshakes with Alice replies.")
+    clients.forEach { it.completeSession() }
+
     // Have clients generate N group messages that will be received by Alice
     val allClientMessages = clients.map { client ->
       val messageCount = 100
@@ -150,6 +158,9 @@ class BenchmarkCommandReceiver : BroadcastReceiver() {
         ThreadUtil.sleep(1000)
       }
     }
+
+    Log.i(TAG, "Completing session handshakes with Alice replies.")
+    clients.forEach { it.completeSession() }
   }
 
   private fun handleDeleteThread() {
diff --git a/app/src/benchmarkShared/java/org/signal/benchmark/setup/OtherClient.kt b/app/src/benchmarkShared/java/org/signal/benchmark/setup/OtherClient.kt
index a581181291..35038bc1c0 100644
--- a/app/src/benchmarkShared/java/org/signal/benchmark/setup/OtherClient.kt
+++ b/app/src/benchmarkShared/java/org/signal/benchmark/setup/OtherClient.kt
@@ -24,8 +24,13 @@ import org.signal.libsignal.protocol.util.KeyHelper
 import org.signal.libsignal.zkgroup.groups.GroupMasterKey
 import org.signal.libsignal.zkgroup.profiles.ProfileKey
 import org.thoughtcrime.securesms.crypto.ProfileKeyUtil
+import org.thoughtcrime.securesms.crypto.ReentrantSessionLock
+import org.thoughtcrime.securesms.crypto.SealedSenderAccessUtil
 import org.thoughtcrime.securesms.database.SignalDatabase
+import org.thoughtcrime.securesms.dependencies.AppDependencies
 import org.thoughtcrime.securesms.keyvalue.SignalStore
+import org.thoughtcrime.securesms.recipients.Recipient
+import org.thoughtcrime.securesms.recipients.RecipientId
 import org.whispersystems.signalservice.api.SignalServiceAccountDataStore
 import org.whispersystems.signalservice.api.SignalSessionLock
 import org.whispersystems.signalservice.api.crypto.EnvelopeContent
@@ -77,6 +82,35 @@ class OtherClient(val serviceId: ServiceId, val e164: String, val identityKeyPai
       .toEnvelope(timestamp, getAliceServiceId())
   }
 
+  fun decrypt(envelope: Envelope, serverDeliveredTimestamp: Long) {
+    val cipher = SignalServiceCipher(serviceAddress, 1, aciStore, sessionLock, SealedSenderAccessUtil.getCertificateValidator())
+    cipher.decrypt(envelope, serverDeliveredTimestamp)
+  }
+
+  /**
+   * Completes the Signal session handshake by having Alice (the app) encrypt a reply
+   * to this client, then decrypting it. This establishes a proper double-ratchet session
+   * on both sides.
+   *
+   * Must be called after this client's initial PreKeyMessage has been processed by Alice.
+   */
+  fun completeSession() {
+    val aliceAddress = SignalServiceAddress(Harness.SELF_ACI, Harness.SELF_E164)
+    val aliceCipher = SignalServiceCipher(aliceAddress, 1, AppDependencies.protocolStore.aci(), ReentrantSessionLock.INSTANCE, null)
+
+    val bobProtocolAddress = SignalProtocolAddress(serviceId.toString(), 1)
+    val now = System.currentTimeMillis()
+    val content = Generator.encryptedTextMessage(now)
+
+    val recipientId = RecipientId.from(SignalServiceAddress(serviceId, e164))
+    val sealedSenderAccess = SealedSenderAccessUtil.getSealedSenderAccessFor(Recipient.resolved(recipientId))
+
+    val outgoing = aliceCipher.encrypt(bobProtocolAddress, sealedSenderAccess, content)
+    val envelope = outgoing.toEnvelope(now, serviceId)
+
+    decrypt(envelope, now)
+  }
+
   fun generateInboundEnvelopes(count: Int): List<Envelope> {
     val envelopes = ArrayList<Envelope>(count)
     var now = System.currentTimeMillis()