Validate device message sizes when distributing PNI keys

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

@@ -51,6 +51,7 @@ import org.whispersystems.textsecuregcm.entities.RegistrationLockFailure;
 import org.whispersystems.textsecuregcm.entities.StaleDevices;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
+import org.whispersystems.textsecuregcm.push.MessageTooLargeException;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.ChangeNumberManager;
@@ -94,6 +95,7 @@ public class AccountControllerV2 {
   @ApiResponse(responseCode = "403", description = "Verification failed for the provided Registration Recovery Password")
   @ApiResponse(responseCode = "409", description = "Mismatched number of devices or device ids in 'devices to notify' list", content = @Content(schema = @Schema(implementation = MismatchedDevices.class)))
   @ApiResponse(responseCode = "410", description = "Mismatched registration ids in 'devices to notify' list", content = @Content(schema = @Schema(implementation = StaleDevices.class)))
+  @ApiResponse(responseCode = "413", description = "One or more device messages was too large")
   @ApiResponse(responseCode = "422", description = "The request did not pass validation")
   @ApiResponse(responseCode = "423", content = @Content(schema = @Schema(implementation = RegistrationLockFailure.class)))
   @ApiResponse(responseCode = "429", description = "Too many attempts", headers = @Header(
@@ -143,7 +145,8 @@ public class AccountControllerV2 {
           request.devicePniSignedPrekeys(),
           request.devicePniPqLastResortPrekeys(),
           request.deviceMessages(),
-          request.pniRegistrationIds());
+          request.pniRegistrationIds(),
+          userAgentString);
 
       return AccountIdentityResponseBuilder.fromAccount(updatedAccount);
     } catch (MismatchedDevicesException e) {
@@ -159,6 +162,8 @@ public class AccountControllerV2 {
           .build());
     } catch (IllegalArgumentException e) {
       throw new BadRequestException(e);
+    } catch (MessageTooLargeException e) {
+      throw new WebApplicationException(Response.Status.REQUEST_ENTITY_TOO_LARGE);
     }
   }
 
@@ -176,6 +181,7 @@ public class AccountControllerV2 {
       content = @Content(schema = @Schema(implementation = MismatchedDevices.class)))
   @ApiResponse(responseCode = "410", description = "The registration IDs provided for some devices do not match those stored on the server.",
       content = @Content(schema = @Schema(implementation = StaleDevices.class)))
+  @ApiResponse(responseCode = "413", description = "One or more device messages was too large")
   public AccountIdentityResponse distributePhoneNumberIdentityKeys(
       @Mutable @Auth final AuthenticatedDevice authenticatedDevice,
       @HeaderParam(HttpHeaders.USER_AGENT) @Nullable final String userAgentString,
@@ -196,7 +202,8 @@ public class AccountControllerV2 {
           request.devicePniSignedPrekeys(),
           request.devicePniPqLastResortPrekeys(),
           request.deviceMessages(),
-          request.pniRegistrationIds());
+          request.pniRegistrationIds(),
+          userAgentString);
 
       return AccountIdentityResponseBuilder.fromAccount(updatedAccount);
     } catch (MismatchedDevicesException e) {
@@ -212,6 +219,8 @@ public class AccountControllerV2 {
           .build());
     } catch (IllegalArgumentException e) {
       throw new BadRequestException(e);
+    } catch (MessageTooLargeException e) {
+      throw new WebApplicationException(Response.Status.REQUEST_ENTITY_TOO_LARGE);
     }
   }
 

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

@@ -22,6 +22,7 @@ import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope;
 import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
 import org.whispersystems.textsecuregcm.push.MessageSender;
+import org.whispersystems.textsecuregcm.push.MessageTooLargeException;
 import org.whispersystems.textsecuregcm.util.DestinationDeviceValidator;
 
 public class ChangeNumberManager {
@@ -42,8 +43,9 @@ public class ChangeNumberManager {
       @Nullable final Map<Byte, ECSignedPreKey> deviceSignedPreKeys,
       @Nullable final Map<Byte, KEMSignedPreKey> devicePqLastResortPreKeys,
       @Nullable final List<IncomingMessage> deviceMessages,
-      @Nullable final Map<Byte, Integer> pniRegistrationIds)
-      throws InterruptedException, MismatchedDevicesException, StaleDevicesException {
+      @Nullable final Map<Byte, Integer> pniRegistrationIds,
+      @Nullable final String senderUserAgent)
+      throws InterruptedException, MismatchedDevicesException, StaleDevicesException, MessageTooLargeException {
 
     if (ObjectUtils.allNotNull(pniIdentityKey, deviceSignedPreKeys, deviceMessages, pniRegistrationIds)) {
       // AccountsManager validates the device set on deviceSignedPreKeys and pniRegistrationIds
@@ -63,14 +65,14 @@ public class ChangeNumberManager {
       if (pniIdentityKey == null) {
         return account;
       }
-      return updatePniKeys(account, pniIdentityKey, deviceSignedPreKeys, devicePqLastResortPreKeys, deviceMessages, pniRegistrationIds);
+      return updatePniKeys(account, pniIdentityKey, deviceSignedPreKeys, devicePqLastResortPreKeys, deviceMessages, pniRegistrationIds, senderUserAgent);
     }
 
     final Account updatedAccount = accountsManager.changeNumber(
         account, number, pniIdentityKey, deviceSignedPreKeys, devicePqLastResortPreKeys, pniRegistrationIds);
 
     if (deviceMessages != null) {
-      sendDeviceMessages(updatedAccount, deviceMessages);
+      sendDeviceMessages(updatedAccount, deviceMessages, senderUserAgent);
     }
 
     return updatedAccount;
@@ -81,7 +83,9 @@ public class ChangeNumberManager {
       final Map<Byte, ECSignedPreKey> deviceSignedPreKeys,
       @Nullable final Map<Byte, KEMSignedPreKey> devicePqLastResortPreKeys,
       final List<IncomingMessage> deviceMessages,
-      final Map<Byte, Integer> pniRegistrationIds) throws MismatchedDevicesException, StaleDevicesException {
+      final Map<Byte, Integer> pniRegistrationIds,
+      final String senderUserAgent) throws MismatchedDevicesException, StaleDevicesException, MessageTooLargeException {
+
     validateDeviceMessages(account, deviceMessages);
 
     // Don't try to be smart about ignoring unnecessary retries. If we make literally no change we will skip the ddb
@@ -89,7 +93,7 @@ public class ChangeNumberManager {
     final Account updatedAccount = accountsManager.updatePniKeys(
         account, pniIdentityKey, deviceSignedPreKeys, devicePqLastResortPreKeys, pniRegistrationIds);
 
-    sendDeviceMessages(updatedAccount, deviceMessages);
+    sendDeviceMessages(updatedAccount, deviceMessages, senderUserAgent);
     return updatedAccount;
   }
 
@@ -110,7 +114,18 @@ public class ChangeNumberManager {
         false);
   }
 
-  private void sendDeviceMessages(final Account account, final List<IncomingMessage> deviceMessages) {
+  private void sendDeviceMessages(final Account account,
+      final List<IncomingMessage> deviceMessages,
+      final String senderUserAgent) throws MessageTooLargeException {
+
+    for (final IncomingMessage message : deviceMessages) {
+      MessageSender.validateContentLength(message.content().length,
+          false,
+          true,
+          false,
+          senderUserAgent);
+    }
+
     try {
       final long serverTimestamp = System.currentTimeMillis();