From 6fa6c3c81c70589c91b1b60290b0e440b81cfb04 Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Tue, 23 Jul 2024 15:36:26 -0500
Subject: [PATCH] Ensure multi-recipient messages are addressed to ACI service
 IDs

---
 .../controllers/MessageController.java        |  7 ++++
 .../controllers/MessageControllerTest.java    | 34 ++++++-------------
 2 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
index 51738dc7a..306efc6d0 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
@@ -717,6 +717,13 @@ public class MessageController {
       final @NotNull CombinedUnidentifiedSenderAccessKeys accessKeys,
       final Collection<MultiRecipientDeliveryData> destinations) {
     final int keyLength = UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH;
+
+    if (destinations.stream()
+        .anyMatch(destination -> IdentityType.PNI.equals(destination.serviceIdentifier.identityType()))) {
+      throw new WebApplicationException("Multi-recipient messages must be addressed to ACI service IDs",
+          Status.UNAUTHORIZED);
+    }
+
     final byte[] combinedUnidentifiedAccessKeys = destinations.stream()
         .map(MultiRecipientDeliveryData::account)
         .filter(Predicate.not(Account::isUnrestrictedUnidentifiedAccess))
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/MessageControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/MessageControllerTest.java
index e8699ea5b..471143e65 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/MessageControllerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/MessageControllerTest.java
@@ -1343,10 +1343,10 @@ class MessageControllerTest {
             new MultiRecipientMessageTestCase(bothAccountsMixed, auth, story, 200, 4),
             new MultiRecipientMessageTestCase(realAndFakeMixed, auth, story, 200, 4),
 
-            new MultiRecipientMessageTestCase(singleDevicePni, auth, notStory, 200, 1),
-            new MultiRecipientMessageTestCase(singleDeviceAciAndPni, unauth, story, 200, 2),
-            new MultiRecipientMessageTestCase(multiDevicePni, auth, notStory, 200, 3),
-            new MultiRecipientMessageTestCase(bothAccountsMixed, auth, notStory, 200, 4),
+            new MultiRecipientMessageTestCase(singleDevicePni, auth, notStory, 401, 0),
+            new MultiRecipientMessageTestCase(singleDeviceAciAndPni, auth, notStory, 401, 0),
+            new MultiRecipientMessageTestCase(multiDevicePni, auth, notStory, 401, 0),
+            new MultiRecipientMessageTestCase(bothAccountsMixed, auth, notStory, 401, 0),
             new MultiRecipientMessageTestCase(realAndFakeMixed, auth, notStory, 404, 0))
         .argumentsForNextParameter(false, true); // urgent
   }
@@ -1580,10 +1580,10 @@ class MessageControllerTest {
     );
   }
 
-  @ParameterizedTest
-  @MethodSource
-  void sendMultiRecipientMessageMismatchedDevices(final ServiceIdentifier serviceIdentifier)
-      throws JsonProcessingException {
+  @Test
+  void sendMultiRecipientMessageMismatchedDevices() throws JsonProcessingException {
+
+    final ServiceIdentifier serviceIdentifier = MULTI_DEVICE_ACI_ID;
 
     final byte extraDeviceId = MULTI_DEVICE_ID3 + 1;
 
@@ -1627,15 +1627,9 @@ class MessageControllerTest {
     }
   }
 
-  private static Stream<Arguments> sendMultiRecipientMessageMismatchedDevices() {
-    return Stream.of(
-        Arguments.of(MULTI_DEVICE_ACI_ID),
-        Arguments.of(MULTI_DEVICE_PNI_ID));
-  }
-
-  @ParameterizedTest
-  @MethodSource
-  void sendMultiRecipientMessageStaleDevices(final ServiceIdentifier serviceIdentifier) throws JsonProcessingException {
+  @Test
+  void sendMultiRecipientMessageStaleDevices() throws JsonProcessingException {
+    final ServiceIdentifier serviceIdentifier = MULTI_DEVICE_ACI_ID;
     final List<Recipient> recipients = List.of(
         new Recipient(serviceIdentifier, MULTI_DEVICE_ID1, MULTI_DEVICE_REG_ID1 + 1, new byte[48]),
         new Recipient(serviceIdentifier, MULTI_DEVICE_ID2, MULTI_DEVICE_REG_ID2 + 1, new byte[48]),
@@ -1677,12 +1671,6 @@ class MessageControllerTest {
     }
   }
 
-  private static Stream<Arguments> sendMultiRecipientMessageStaleDevices() {
-    return Stream.of(
-        Arguments.of(MULTI_DEVICE_ACI_ID),
-        Arguments.of(MULTI_DEVICE_PNI_ID));
-  }
-
   @Test
   void sendMultiRecipientMessageStoryRateLimited() {
     final List<Recipient> recipients = List.of(new Recipient(SINGLE_DEVICE_ACI_ID, SINGLE_DEVICE_ID1, SINGLE_DEVICE_REG_ID1, new byte[48]));