Centralize message size validation in actual message-sending methods

service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java

@@ -322,21 +322,10 @@ public class MessageController {
 
       final Optional<byte[]> reportSpamToken = spamCheckResult.token();
 
-      int totalContentLength = 0;
-
-      for (final IncomingMessage message : messages.messages()) {
-        final int contentLength = message.content().length;
-
-        try {
-          MessageSender.validateContentLength(contentLength, false, isSyncMessage, isStory, userAgent);
-        } catch (final MessageTooLargeException e) {
-          throw new WebApplicationException(Status.REQUEST_ENTITY_TOO_LARGE);
-        }
-
-        totalContentLength += contentLength;
-      }
-
       try {
+        final int totalContentLength =
+            messages.messages().stream().mapToInt(message -> message.content().length).sum();
+
         rateLimiters.getInboundMessageBytes().validate(destinationIdentifier.uuid(), totalContentLength);
       } catch (final RateLimitExceededException e) {
         messageByteLimitEstimator.add(destinationIdentifier.uuid().toString());
@@ -409,7 +398,7 @@ public class MessageController {
           authType = AUTH_TYPE_ACCESS_KEY;
         }
 
-        messageSender.sendMessages(destination, destinationIdentifier, messagesByDeviceId, registrationIdsByDeviceId);
+        messageSender.sendMessages(destination, destinationIdentifier, messagesByDeviceId, registrationIdsByDeviceId, userAgent);
 
         Metrics.counter(SENT_MESSAGE_COUNTER_NAME, List.of(UserAgentTagUtil.getPlatformTag(userAgent),
             Tag.of(ENDPOINT_TYPE_TAG_NAME, ENDPOINT_TYPE_SINGLE),
@@ -433,6 +422,8 @@ public class MessageController {
                   e.getMismatchedDevices().extraDeviceIds()))
               .build());
         }
+      } catch (final MessageTooLargeException e) {
+        throw new WebApplicationException(Status.REQUEST_ENTITY_TOO_LARGE);
       }
     } finally {
       sample.stop(Timer.builder(SEND_MESSAGE_LATENCY_TIMER_NAME)
@@ -505,15 +496,6 @@ public class MessageController {
       throw new BadRequestException("Recipient list is empty");
     }
 
-    // Verify that the message isn't too large before performing more expensive validations
-    multiRecipientMessage.getRecipients().values().forEach(recipient -> {
-      try {
-        MessageSender.validateContentLength(multiRecipientMessage.messageSizeForRecipient(recipient), true, false, isStory, userAgent);
-      } catch (final MessageTooLargeException e) {
-        throw new WebApplicationException(Status.REQUEST_ENTITY_TOO_LARGE);
-      }
-    });
-
     // Check that the request is well-formed and doesn't contain repeated entries for the same device for the same
     // recipient
     {
@@ -616,7 +598,7 @@ public class MessageController {
 
     try {
       if (!resolvedRecipients.isEmpty()) {
-        messageSender.sendMultiRecipientMessage(multiRecipientMessage, resolvedRecipients, timestamp, isStory, online, isUrgent).get();
+        messageSender.sendMultiRecipientMessage(multiRecipientMessage, resolvedRecipients, timestamp, isStory, online, isUrgent, userAgent).get();
       }
 
       final List<ServiceIdentifier> unresolvedRecipientServiceIds;
@@ -695,6 +677,8 @@ public class MessageController {
       }
 
       throw new RuntimeException(e);
+    } catch (final MessageTooLargeException e) {
+      throw new WebApplicationException(Status.REQUEST_ENTITY_TOO_LARGE);
     }
   }
 

service/src/main/java/org/whispersystems/textsecuregcm/push/MessageSender.java

@@ -35,6 +35,7 @@ import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.util.Util;
+import javax.annotation.Nullable;
 
 /**
  * A MessageSender sends Signal messages to destination devices. Messages may be "normal" user-to-user messages,
@@ -86,11 +87,18 @@ public class MessageSender {
    * @param destinationIdentifier the service identifier to which the messages are addressed
    * @param messagesByDeviceId a map of device IDs to message payloads
    * @param registrationIdsByDeviceId a map of device IDs to device registration IDs
+   * @param userAgent the User-Agent string for the sender; may be {@code null} if not known
+   *
+   * @throws MismatchedDevicesException if the given bundle of messages did not include a message for all required
+   * devices, contained messages for devices not linked to the destination account, or devices with outdated
+   * registration IDs
+   * @throws MessageTooLargeException if the given message payload is too large
    */
   public void sendMessages(final Account destination,
       final ServiceIdentifier destinationIdentifier,
       final Map<Byte, Envelope> messagesByDeviceId,
-      final Map<Byte, Integer> registrationIdsByDeviceId) throws MismatchedDevicesException {
+      final Map<Byte, Integer> registrationIdsByDeviceId,
+      @Nullable final String userAgent) throws MismatchedDevicesException, MessageTooLargeException {
 
     if (messagesByDeviceId.isEmpty()) {
       return;
@@ -105,6 +113,10 @@ public class MessageSender {
     final boolean isSyncMessage = StringUtils.isNotBlank(firstMessage.getSourceServiceId()) &&
         destination.isIdentifiedBy(ServiceIdentifier.valueOf(firstMessage.getSourceServiceId()));
 
+    final boolean isStory = firstMessage.getStory();
+
+    validateIndividualMessageContentLength(messagesByDeviceId.values(), isSyncMessage, isStory, userAgent);
+
     final Optional<MismatchedDevices> maybeMismatchedDevices = getMismatchedDevices(destination,
         destinationIdentifier,
         registrationIdsByDeviceId,
@@ -152,15 +164,24 @@ public class MessageSender {
    * @param isEphemeral {@code true} if the message should only be delivered to devices with active connections or
    * {@code false otherwise}
    * @param isUrgent {@code true} if the message is urgent or {@code false otherwise}
+   * @param userAgent the User-Agent string for the sender; may be {@code null} if not known
    *
    * @return a future that completes when all messages have been inserted into delivery queues
+   *
+   * @throws MultiRecipientMismatchedDevicesException if the given multi-recipient message had did not have all required
+   * recipient devices for a recipient account, contained recipients for devices not linked to a destination account, or
+   * recipient devices with outdated registration IDs
+   * @throws MessageTooLargeException if the given message payload is too large
    */
   public CompletableFuture<Void> sendMultiRecipientMessage(final SealedSenderMultiRecipientMessage multiRecipientMessage,
       final Map<SealedSenderMultiRecipientMessage.Recipient, Account> resolvedRecipients,
       final long clientTimestamp,
       final boolean isStory,
       final boolean isEphemeral,
-      final boolean isUrgent) throws MultiRecipientMismatchedDevicesException {
+      final boolean isUrgent,
+      @Nullable final String userAgent) throws MultiRecipientMismatchedDevicesException, MessageTooLargeException {
+
+    validateMultiRecipientMessageContentLength(multiRecipientMessage, isStory, userAgent);
 
     final Map<ServiceIdentifier, MismatchedDevices> mismatchedDevicesByServiceIdentifier = new HashMap<>();
 
@@ -224,7 +245,8 @@ public class MessageSender {
     }
   }
 
-  public static void validateContentLength(final int contentLength,
+  @VisibleForTesting
+  static void validateContentLength(final int contentLength,
       final boolean isMultiRecipientMessage,
       final boolean isSyncMessage,
       final boolean isStory,
@@ -302,4 +324,31 @@ public class MessageSender {
         ? Optional.of(new MismatchedDevices(missingDeviceIds, extraDeviceIds, staleDeviceIds))
         : Optional.empty();
   }
+
+  private static void validateIndividualMessageContentLength(final Iterable<Envelope> messages,
+      final boolean isSyncMessage,
+      final boolean isStory,
+      @Nullable final String userAgent) throws MessageTooLargeException {
+
+    for (final Envelope message : messages) {
+      MessageSender.validateContentLength(message.getContent().size(),
+          false,
+          isSyncMessage,
+          isStory,
+          userAgent);
+    }
+  }
+
+  private static void validateMultiRecipientMessageContentLength(final SealedSenderMultiRecipientMessage multiRecipientMessage,
+      final boolean isStory,
+      @Nullable final String userAgent) throws MessageTooLargeException {
+
+    for (final SealedSenderMultiRecipientMessage.Recipient recipient : multiRecipientMessage.getRecipients().values()) {
+      MessageSender.validateContentLength(multiRecipientMessage.messageSizeForRecipient(recipient),
+          true,
+          false,
+          isStory,
+          userAgent);
+    }
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/push/ReceiptSender.java

@@ -68,7 +68,7 @@ public class ReceiptSender {
                 messageSender.sendMessages(destinationAccount,
                     destinationIdentifier,
                     messagesByDeviceId,
-                    registrationIdsByDeviceId);
+                    registrationIdsByDeviceId, null);
               } catch (final Exception e) {
                 logger.warn("Could not send delivery receipt", e);
               }

service/src/main/java/org/whispersystems/textsecuregcm/storage/ChangeNumberManager.java

@@ -19,7 +19,6 @@ import org.whispersystems.textsecuregcm.entities.IncomingMessage;
 import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope;
 import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
-import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.identity.ServiceIdentifier;
 import org.whispersystems.textsecuregcm.push.MessageSender;
 import org.whispersystems.textsecuregcm.push.MessageTooLargeException;
@@ -96,14 +95,6 @@ public class ChangeNumberManager {
       final List<IncomingMessage> deviceMessages,
       final String senderUserAgent) throws MessageTooLargeException, MismatchedDevicesException {
 
-    for (final IncomingMessage message : deviceMessages) {
-      MessageSender.validateContentLength(message.content().length,
-          false,
-          true,
-          false,
-          senderUserAgent);
-    }
-
     try {
       final long serverTimestamp = System.currentTimeMillis();
       final ServiceIdentifier serviceIdentifier = new AciServiceIdentifier(account.getUuid());
@@ -125,7 +116,7 @@ public class ChangeNumberManager {
       final Map<Byte, Integer> registrationIdsByDeviceId = account.getDevices().stream()
           .collect(Collectors.toMap(Device::getId, Device::getRegistrationId));
 
-      messageSender.sendMessages(account, serviceIdentifier, messagesByDeviceId, registrationIdsByDeviceId);
+      messageSender.sendMessages(account, serviceIdentifier, messagesByDeviceId, registrationIdsByDeviceId, senderUserAgent);
     } catch (final RuntimeException e) {
       logger.warn("Changed number but could not send all device messages on {}", account.getUuid(), e);
       throw e;