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)
+      @Nullable final Map<Byte, Integer> pniRegistrationIds,
-      throws InterruptedException, MismatchedDevicesException, StaleDevicesException {
+      @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();
 

service/src/test/java/org/whispersystems/textsecuregcm/controllers/AccountControllerV2Test.java

@@ -16,6 +16,7 @@ import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -79,6 +80,7 @@ import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.mappers.ImpossiblePhoneNumberExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.NonNormalizedPhoneNumberExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper;
+import org.whispersystems.textsecuregcm.push.MessageTooLargeException;
 import org.whispersystems.textsecuregcm.registration.RegistrationServiceClient;
 import org.whispersystems.textsecuregcm.spam.RegistrationRecoveryChecker;
 import org.whispersystems.textsecuregcm.storage.Account;
@@ -139,7 +141,7 @@ class AccountControllerV2Test {
     void setUp() throws Exception {
       when(rateLimiters.getRegistrationLimiter()).thenReturn(registrationLimiter);
 
-      when(changeNumberManager.changeNumber(any(), any(), any(), any(), any(), any(), any())).thenAnswer(
+      when(changeNumberManager.changeNumber(any(), any(), any(), any(), any(), any(), any(), any())).thenAnswer(
           (Answer<Account>) invocation -> {
             final Account account = invocation.getArgument(0);
             final String number = invocation.getArgument(1);
@@ -189,7 +191,7 @@ class AccountControllerV2Test {
                   MediaType.APPLICATION_JSON_TYPE), AccountIdentityResponse.class);
 
       verify(changeNumberManager).changeNumber(eq(AuthHelper.VALID_ACCOUNT), eq(NEW_NUMBER), any(), any(), any(),
-          any(), any());
+          any(), any(), any());
 
       assertEquals(AuthHelper.VALID_UUID, accountIdentityResponse.uuid());
       assertEquals(NEW_NUMBER, accountIdentityResponse.number());
@@ -212,7 +214,7 @@ class AccountControllerV2Test {
                   MediaType.APPLICATION_JSON_TYPE), AccountIdentityResponse.class);
 
       verify(changeNumberManager).changeNumber(eq(AuthHelper.VALID_ACCOUNT), eq(AuthHelper.VALID_NUMBER), any(), any(), any(),
-          any(), any());
+          any(), any(), any());
 
       assertEquals(AuthHelper.VALID_UUID, accountIdentityResponse.uuid());
       assertEquals(AuthHelper.VALID_NUMBER, accountIdentityResponse.number());
@@ -220,7 +222,7 @@ class AccountControllerV2Test {
     }
 
     @Test
-    void changeNumberNonNormalizedNumber() throws Exception {
+    void changeNumberNonNormalizedNumber() {
       try (Response response = resources.getJerseyTest()
           .target("/v2/accounts/number")
           .request()
@@ -293,14 +295,16 @@ class AccountControllerV2Test {
           null, NEW_NUMBER, "123", new IdentityKey(Curve.generateKeyPair().getPublicKey()),
           Collections.emptyList(), Collections.emptyMap(), null, Map.of((byte) 1, pniRegistrationId));
 
-      final Response response = resources.getJerseyTest()
+      try (final Response response = resources.getJerseyTest()
           .target("/v2/accounts/number")
           .request()
           .header(HttpHeaders.AUTHORIZATION,
               AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
-          .put(Entity.entity(changeNumberRequest, MediaType.APPLICATION_JSON_TYPE));
+          .put(Entity.entity(changeNumberRequest, MediaType.APPLICATION_JSON_TYPE))) {
+
         assertEquals(expectedStatusCode, response.getStatus());
       }
+    }
 
     private static Stream<Arguments> invalidRegistrationId() {
       return Stream.of(
@@ -422,7 +426,7 @@ class AccountControllerV2Test {
         final AccountIdentityResponse accountIdentityResponse = response.readEntity(AccountIdentityResponse.class);
 
         verify(changeNumberManager).changeNumber(eq(AuthHelper.VALID_ACCOUNT), eq(NEW_NUMBER), any(), any(), any(),
-            any(), any());
+            any(), any(), any());
 
         assertEquals(AuthHelper.VALID_UUID, accountIdentityResponse.uuid());
         assertEquals(NEW_NUMBER, accountIdentityResponse.number());
@@ -481,6 +485,33 @@ class AccountControllerV2Test {
       }
     }
 
+    @Test
+    void deviceMessageTooLarge() throws Exception {
+
+      when(registrationServiceClient.getSession(any(), any()))
+          .thenReturn(CompletableFuture.completedFuture(
+              Optional.of(new RegistrationServiceSession(new byte[16], NEW_NUMBER, true, null, null, null,
+                  SESSION_EXPIRATION_SECONDS))));
+
+      reset(changeNumberManager);
+      when(changeNumberManager.changeNumber(any(), any(), any(), any(), any(), any(), any(), any()))
+          .thenThrow(MessageTooLargeException.class);
+
+      try (final Response response = resources.getJerseyTest()
+              .target("/v2/accounts/number")
+              .request()
+              .header(HttpHeaders.AUTHORIZATION,
+                  AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+              .put(Entity.entity(
+                  new ChangeNumberRequest(encodeSessionId("session"), null, NEW_NUMBER, "123", new IdentityKey(Curve.generateKeyPair().getPublicKey()),
+                      Collections.emptyList(),
+                      Collections.emptyMap(), null, Collections.emptyMap()),
+                  MediaType.APPLICATION_JSON_TYPE))) {
+
+        assertEquals(413, response.getStatus());
+      }
+    }
+
     /**
      * Valid request JSON with the given Recovery Password
      */
@@ -558,7 +589,7 @@ class AccountControllerV2Test {
 
     @BeforeEach
     void setUp() throws Exception {
-      when(changeNumberManager.updatePniKeys(any(), any(), any(), any(), any(), any())).thenAnswer(
+      when(changeNumberManager.updatePniKeys(any(), any(), any(), any(), any(), any(), any())).thenAnswer(
           (Answer<Account>) invocation -> {
             final Account account = invocation.getArgument(0);
             final IdentityKey pniIdentityKey = invocation.getArgument(1);
@@ -594,7 +625,7 @@ class AccountControllerV2Test {
               AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
           .put(Entity.json(requestJson()), AccountIdentityResponse.class);
 
-      verify(changeNumberManager).updatePniKeys(eq(AuthHelper.VALID_ACCOUNT), eq(IDENTITY_KEY), any(), any(), any(), any());
+      verify(changeNumberManager).updatePniKeys(eq(AuthHelper.VALID_ACCOUNT), eq(IDENTITY_KEY), any(), any(), any(), any(), any());
 
       assertEquals(AuthHelper.VALID_UUID, accountIdentityResponse.uuid());
       assertEquals(AuthHelper.VALID_NUMBER, accountIdentityResponse.number());
@@ -646,6 +677,23 @@ class AccountControllerV2Test {
       }
     }
 
+    @Test
+    void deviceMessageTooLarge() throws Exception {
+      reset(changeNumberManager);
+      when(changeNumberManager.updatePniKeys(any(), any(), any(), any(), any(), any(), any()))
+          .thenThrow(MessageTooLargeException.class);
+
+      try (final Response response = resources.getJerseyTest()
+              .target("/v2/accounts/phone_number_identity_key_distribution")
+              .request()
+              .header(HttpHeaders.AUTHORIZATION,
+                  AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+              .put(Entity.json(requestJson()))) {
+
+        assertEquals(413, response.getStatus());
+      }
+    }
+
     /**
      * Valid request JSON for a {@link org.whispersystems.textsecuregcm.entities.PhoneNumberIdentityKeyDistributionRequest}
      */

service/src/test/java/org/whispersystems/textsecuregcm/storage/ChangeNumberManagerTest.java

@@ -102,7 +102,7 @@ public class ChangeNumberManagerTest {
   void changeNumberNoMessages() throws Exception {
     Account account = mock(Account.class);
     when(account.getNumber()).thenReturn("+18005551234");
-    changeNumberManager.changeNumber(account, "+18025551234", null, null, null, null, null);
+    changeNumberManager.changeNumber(account, "+18025551234", null, null, null, null, null, null);
     verify(accountsManager).changeNumber(account, "+18025551234", null, null, null, null);
     verify(accountsManager, never()).updateDevice(any(), anyByte(), any());
     verify(messageSender, never()).sendMessages(eq(account), any());
@@ -117,7 +117,7 @@ public class ChangeNumberManagerTest {
     final Map<Byte, ECSignedPreKey> prekeys = Map.of(Device.PRIMARY_ID,
         KeysHelper.signedECPreKey(1, pniIdentityKeyPair));
 
-    changeNumberManager.changeNumber(account, "+18025551234", pniIdentityKey, prekeys, null, Collections.emptyList(), Collections.emptyMap());
+    changeNumberManager.changeNumber(account, "+18025551234", pniIdentityKey, prekeys, null, Collections.emptyList(), Collections.emptyMap(), null);
     verify(accountsManager).changeNumber(account, "+18025551234", pniIdentityKey, prekeys, null, Collections.emptyMap());
     verify(messageSender, never()).sendMessages(eq(account), any());
   }
@@ -152,7 +152,7 @@ public class ChangeNumberManagerTest {
     when(msg.destinationDeviceId()).thenReturn(deviceId2);
     when(msg.content()).thenReturn(new byte[]{1});
 
-    changeNumberManager.changeNumber(account, changedE164, pniIdentityKey, prekeys, null, List.of(msg), registrationIds);
+    changeNumberManager.changeNumber(account, changedE164, pniIdentityKey, prekeys, null, List.of(msg), registrationIds, null);
 
     verify(accountsManager).changeNumber(account, changedE164, pniIdentityKey, prekeys, null, registrationIds);
 
@@ -205,7 +205,7 @@ public class ChangeNumberManagerTest {
     when(msg.destinationDeviceId()).thenReturn(deviceId2);
     when(msg.content()).thenReturn(new byte[]{1});
 
-    changeNumberManager.changeNumber(account, changedE164, pniIdentityKey, prekeys, pqPrekeys, List.of(msg), registrationIds);
+    changeNumberManager.changeNumber(account, changedE164, pniIdentityKey, prekeys, pqPrekeys, List.of(msg), registrationIds, null);
 
     verify(accountsManager).changeNumber(account, changedE164, pniIdentityKey, prekeys, pqPrekeys, registrationIds);
 
@@ -256,7 +256,7 @@ public class ChangeNumberManagerTest {
     when(msg.destinationDeviceId()).thenReturn(deviceId2);
     when(msg.content()).thenReturn(new byte[]{1});
 
-    changeNumberManager.changeNumber(account, originalE164, pniIdentityKey, prekeys, pqPrekeys, List.of(msg), registrationIds);
+    changeNumberManager.changeNumber(account, originalE164, pniIdentityKey, prekeys, pqPrekeys, List.of(msg), registrationIds, null);
 
     verify(accountsManager).updatePniKeys(account, pniIdentityKey, prekeys, pqPrekeys, registrationIds);
 
@@ -303,7 +303,7 @@ public class ChangeNumberManagerTest {
     when(msg.destinationDeviceId()).thenReturn(deviceId2);
     when(msg.content()).thenReturn(new byte[]{1});
 
-    changeNumberManager.updatePniKeys(account, pniIdentityKey, prekeys, null, List.of(msg), registrationIds);
+    changeNumberManager.updatePniKeys(account, pniIdentityKey, prekeys, null, List.of(msg), registrationIds, null);
 
     verify(accountsManager).updatePniKeys(account, pniIdentityKey, prekeys, null, registrationIds);
 
@@ -352,7 +352,7 @@ public class ChangeNumberManagerTest {
     when(msg.destinationDeviceId()).thenReturn(deviceId2);
     when(msg.content()).thenReturn(new byte[]{1});
 
-    changeNumberManager.updatePniKeys(account, pniIdentityKey, prekeys, pqPrekeys, List.of(msg), registrationIds);
+    changeNumberManager.updatePniKeys(account, pniIdentityKey, prekeys, pqPrekeys, List.of(msg), registrationIds, null);
 
     verify(accountsManager).updatePniKeys(account, pniIdentityKey, prekeys, pqPrekeys, registrationIds);
 
@@ -407,7 +407,7 @@ public class ChangeNumberManagerTest {
         destinationDeviceId3, 89);
 
     assertThrows(StaleDevicesException.class,
-        () -> changeNumberManager.changeNumber(account, "+18005559876", new IdentityKey(Curve.generateKeyPair().getPublicKey()), preKeys, null, messages, registrationIds));
+        () -> changeNumberManager.changeNumber(account, "+18005559876", new IdentityKey(Curve.generateKeyPair().getPublicKey()), preKeys, null, messages, registrationIds, null));
   }
 
   @Test
@@ -445,7 +445,7 @@ public class ChangeNumberManagerTest {
         destinationDeviceId3, 89);
 
     assertThrows(StaleDevicesException.class,
-        () -> changeNumberManager.updatePniKeys(account, new IdentityKey(Curve.generateKeyPair().getPublicKey()), preKeys, null, messages, registrationIds));
+        () -> changeNumberManager.updatePniKeys(account, new IdentityKey(Curve.generateKeyPair().getPublicKey()), preKeys, null, messages, registrationIds, null));
   }
 
   @Test
@@ -475,6 +475,6 @@ public class ChangeNumberManagerTest {
     final Map<Byte, Integer> registrationIds = Map.of((byte) 1, 17, destinationDeviceId2, 47, destinationDeviceId3, 89);
 
     assertThrows(IllegalArgumentException.class,
-        () -> changeNumberManager.changeNumber(account, "+18005559876", new IdentityKey(Curve.generateKeyPair().getPublicKey()), null, null, messages, registrationIds));
+        () -> changeNumberManager.changeNumber(account, "+18005559876", new IdentityKey(Curve.generateKeyPair().getPublicKey()), null, null, messages, registrationIds, null));
   }
 }