From f60c9f2a154a871cdb16d59d2f702e5f0a2e9547 Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Sat, 7 Sep 2024 11:47:55 -0500
Subject: [PATCH] Use destination service ID from the envelope when looking up
 in shared MRM data

---
 .../textsecuregcm/storage/MessagesCache.java  | 21 +++++++++++++++----
 ...edMultiRecipientPayloadAndViewsScript.java |  5 +----
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCache.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCache.java
index d791bbea8..40e0327cc 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCache.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCache.java
@@ -41,6 +41,7 @@ import java.util.function.Predicate;
 import javax.annotation.Nullable;
 import org.reactivestreams.Publisher;
 import org.signal.libsignal.protocol.SealedSenderMultiRecipientMessage;
+import org.signal.libsignal.protocol.ServiceId;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
@@ -48,6 +49,7 @@ import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicMessagesCon
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
 import org.whispersystems.textsecuregcm.experiment.Experiment;
 import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
+import org.whispersystems.textsecuregcm.identity.ServiceIdentifier;
 import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantPubSubConnection;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
@@ -431,8 +433,9 @@ public class MessagesCache extends RedisClusterPubSubAdapter<String, String> imp
       final Experiment experiment = new Experiment(MRM_VIEWS_EXPERIMENT_NAME);
 
       final byte[] key = mrmMessage.getSharedMrmKey().toByteArray();
-      final byte[] sharedMrmViewKey = MessagesCache.getSharedMrmViewKey(new AciServiceIdentifier(destinationUuid),
-          destinationDevice);
+      final byte[] sharedMrmViewKey = MessagesCache.getSharedMrmViewKey(
+          // the message might be addressed to the account's PNI, so use the service ID from the envelope
+          ServiceIdentifier.valueOf(mrmMessage.getDestinationServiceId()), destinationDevice);
 
       final Mono<MessageProtos.Envelope> mrmMessageMono = Mono.from(redisCluster.withBinaryClusterReactive(
           conn -> conn.reactive().hmget(key, "data".getBytes(StandardCharsets.UTF_8), sharedMrmViewKey)
@@ -786,9 +789,19 @@ public class MessagesCache extends RedisClusterPubSubAdapter<String, String> imp
     return ("user_queue_persisting::{" + accountUuid + "::" + deviceId + "}").getBytes(StandardCharsets.UTF_8);
   }
 
-  static byte[] getSharedMrmViewKey(final AciServiceIdentifier serviceIdentifier, final byte deviceId) {
+  static byte[] getSharedMrmViewKey(final ServiceId serviceId, final byte deviceId) {
+    return getSharedMrmViewKey(serviceId.toServiceIdFixedWidthBinary(), deviceId);
+  }
+
+  static byte[] getSharedMrmViewKey(final ServiceIdentifier serviceIdentifier, final byte deviceId) {
+    return getSharedMrmViewKey(serviceIdentifier.toFixedWidthByteArray(), deviceId);
+  }
+
+  private static byte[] getSharedMrmViewKey(final byte[] fixedWithServiceId, final byte deviceId) {
+    assert fixedWithServiceId.length == 17;
+
     final ByteBuffer keyBb = ByteBuffer.allocate(18);
-    keyBb.put(serviceIdentifier.toFixedWidthByteArray());
+    keyBb.put(fixedWithServiceId);
     keyBb.put(deviceId);
     assert !keyBb.hasRemaining();
     return keyBb.array();
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCacheInsertSharedMultiRecipientPayloadAndViewsScript.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCacheInsertSharedMultiRecipientPayloadAndViewsScript.java
index b0f3ef5db..2d7274552 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCacheInsertSharedMultiRecipientPayloadAndViewsScript.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCacheInsertSharedMultiRecipientPayloadAndViewsScript.java
@@ -42,10 +42,7 @@ class MessagesCacheInsertSharedMultiRecipientPayloadAndViewsScript {
 
     message.getRecipients().forEach((serviceId, recipient) -> {
       for (byte device : recipient.getDevices()) {
-        final byte[] key = new byte[18];
-        System.arraycopy(serviceId.toServiceIdFixedWidthBinary(), 0, key, 0, 17);
-        key[17] = device;
-        args.add(key);
+        args.add(MessagesCache.getSharedMrmViewKey(serviceId, device));
         args.add(message.serializedRecipientView(recipient));
       }
     });