Validate registration IDs

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

@@ -8,7 +8,6 @@ import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
 
 import io.dropwizard.auth.Auth;
 import io.micrometer.core.instrument.Metrics;
-import io.micrometer.core.instrument.Tags;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.Parameter;
 import io.swagger.v3.oas.annotations.responses.ApiResponse;
@@ -228,10 +227,6 @@ public class AccountController {
     final Account account = disabledPermittedAuth.getAccount();
     final byte deviceId = disabledPermittedAuth.getAuthenticatedDevice().getId();
 
-    if (!AccountsManager.validNewAccountAttributes(attributes)) {
-      Metrics.counter(INVALID_REGISTRATION_ID, Tags.of(UserAgentTagUtil.getPlatformTag(userAgent))).increment();
-    }
-
     final Account updatedAccount = accounts.update(account, a -> {
       a.getDevice(deviceId).ifPresent(d -> {
         d.setFetchesMessages(attributes.getFetchesMessages());

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

@@ -107,10 +107,6 @@ public class RegistrationController {
     final String password = authorizationHeader.getPassword();
 
     RateLimiter.adaptLegacyException(() -> rateLimiters.getRegistrationLimiter().validate(number));
-    if (!AccountsManager.validNewAccountAttributes(registrationRequest.accountAttributes())) {
-      Metrics.counter(INVALID_ACCOUNT_ATTRS_COUNTER_NAME, Tags.of(UserAgentTagUtil.getPlatformTag(userAgent))).increment();
-      throw new WebApplicationException(Response.status(422, "account attributes invalid").build());
-    }
 
     final PhoneVerificationRequest.VerificationType verificationType = phoneVerificationTokenManager.verify(number,
         registrationRequest);

service/src/main/java/org/whispersystems/textsecuregcm/entities/AccountAttributes.java

@@ -4,12 +4,15 @@
  */
 package org.whispersystems.textsecuregcm.entities;
 
+import static org.whispersystems.textsecuregcm.util.RegistrationIdValidator.validRegistrationId;
+
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 import com.google.common.annotations.VisibleForTesting;
 import java.util.Optional;
 import java.util.OptionalInt;
 import javax.annotation.Nullable;
+import javax.validation.constraints.AssertTrue;
 import javax.validation.constraints.Size;
 import org.whispersystems.textsecuregcm.auth.UnidentifiedAccessUtil;
 import org.whispersystems.textsecuregcm.storage.Device.DeviceCapabilities;
@@ -127,4 +130,10 @@ public class AccountAttributes {
   public void setPhoneNumberIdentityRegistrationId(final Integer phoneNumberIdentityRegistrationId) {
     this.phoneNumberIdentityRegistrationId = phoneNumberIdentityRegistrationId;
   }
+
+  @AssertTrue
+  public boolean isEachRegistrationIdValid() {
+    return validRegistrationId(registrationId) &&
+        (phoneNumberIdentityRegistrationId == null || validRegistrationId(phoneNumberIdentityRegistrationId));
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/entities/ChangeNumberRequest.java

@@ -21,6 +21,7 @@ import javax.validation.constraints.NotNull;
 import org.signal.libsignal.protocol.IdentityKey;
 import org.whispersystems.textsecuregcm.util.ByteArrayAdapter;
 import org.whispersystems.textsecuregcm.util.IdentityKeyAdapter;
+import org.whispersystems.textsecuregcm.util.RegistrationIdValidator;
 
 public record ChangeNumberRequest(
     @Schema(description="""
@@ -78,4 +79,9 @@ public record ChangeNumberRequest(
     }
     return spks.isEmpty() || PreKeySignatureValidator.validatePreKeySignatures(pniIdentityKey, spks);
   }
+
+  @AssertTrue
+  public boolean isEachPniRegistrationIdValid() {
+    return pniRegistrationIds == null || pniRegistrationIds.values().stream().allMatch(RegistrationIdValidator::validRegistrationId);
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/entities/LinkDeviceRequest.java

@@ -16,6 +16,7 @@ public record LinkDeviceRequest(@Schema(requiredMode = Schema.RequiredMode.REQUI
                                 """)
                                 String verificationCode,
 
+                                @Valid
                                 AccountAttributes accountAttributes,
 
                                 @NotNull

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

@@ -33,7 +33,6 @@ import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.Optional;
-import java.util.OptionalInt;
 import java.util.Queue;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
@@ -386,21 +385,6 @@ public class AccountsManager {
     return updatedAccount.get();
   }
 
-  public static boolean validNewAccountAttributes(final AccountAttributes accountAttributes) {
-    if (!validRegistrationId(accountAttributes.getRegistrationId())) {
-      return false;
-    }
-    final OptionalInt pniRegistrationId = accountAttributes.getPhoneNumberIdentityRegistrationId();
-    if (pniRegistrationId.isPresent() && !validRegistrationId(pniRegistrationId.getAsInt())) {
-      return false;
-    }
-    return true;
-  }
-
-  private static boolean validRegistrationId(int registrationId) {
-    return registrationId > 0 && registrationId <= Device.MAX_REGISTRATION_ID;
-  }
-
   public Account updatePniKeys(final Account account,
       final IdentityKey pniIdentityKey,
       final Map<Byte, ECSignedPreKey> pniSignedPreKeys,

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

@@ -56,7 +56,6 @@ public class ChangeNumberManager {
       throw new IllegalArgumentException("PNI identity key, signed pre-keys, device messages, and registration IDs must be all null or all non-null");
     }
 
-
     if (number.equals(account.getNumber())) {
       // The client has gotten confused/desynchronized with us about their own phone number, most likely due to losing
       // our OK response to an immediately preceding change-number request, and are sending a change they don't realize

service/src/main/java/org/whispersystems/textsecuregcm/util/RegistrationIdValidator.java

@@ -0,0 +1,15 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ *
+ */
+
+package org.whispersystems.textsecuregcm.util;
+
+import org.whispersystems.textsecuregcm.storage.Device;
+
+public class RegistrationIdValidator {
+  public static boolean validRegistrationId(int registrationId) {
+    return registrationId > 0 && registrationId <= Device.MAX_REGISTRATION_ID;
+  }
+}

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

@@ -265,6 +265,36 @@ class AccountControllerV2Test {
       }
     }
 
+    @ParameterizedTest
+    @MethodSource
+    void invalidRegistrationId(final Integer pniRegistrationId, final int expectedStatusCode) {
+      when(registrationServiceClient.getSession(any(), any()))
+          .thenReturn(CompletableFuture.completedFuture(
+              Optional.of(new RegistrationServiceSession(new byte[16], NEW_NUMBER, true, null, null, null,
+                  SESSION_EXPIRATION_SECONDS))));
+      final ChangeNumberRequest changeNumberRequest = new ChangeNumberRequest(encodeSessionId("session"),
+          null, NEW_NUMBER, "123", new IdentityKey(Curve.generateKeyPair().getPublicKey()),
+          Collections.emptyList(), Collections.emptyMap(), null, Map.of((byte) 1, pniRegistrationId));
+
+      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));
+      assertEquals(expectedStatusCode, response.getStatus());
+    }
+
+    private static Stream<Arguments> invalidRegistrationId() {
+      return Stream.of(
+          Arguments.of(0x3FFF, 200),
+          Arguments.of(0, 422),
+          Arguments.of(-1, 422),
+          Arguments.of(0x3FFF + 1, 422),
+          Arguments.of(Integer.MAX_VALUE, 422)
+      );
+    }
+
     @Test
     void rateLimitedNumber() throws Exception {
       doThrow(new RateLimitExceededException(null, true))
@@ -398,13 +428,20 @@ class AccountControllerV2Test {
      * Valid request JSON with the given Recovery Password
      */
     private static String requestJsonRecoveryPassword(final byte[] recoveryPassword, final String newNumber) {
-      return requestJson("", recoveryPassword, newNumber);
+      return requestJson("", recoveryPassword, newNumber, 123);
+    }
+
+    /**
+     * Valid request JSON with the given pniRegistrationId
+     */
+    private static String requestJsonRegistrationIds(final Integer pniRegistrationId) {
+      return requestJson("", new byte[0], "+18005551234", pniRegistrationId);
     }
 
     /**
      * Valid request JSON with the give session ID and recovery password
      */
-    private static String requestJson(final String sessionId, final byte[] recoveryPassword, final String newNumber) {
+    private static String requestJson(final String sessionId, final byte[] recoveryPassword, final String newNumber, final Integer pniRegistrationId) {
       return String.format("""
           {
             "sessionId": "%s",
@@ -414,16 +451,16 @@ class AccountControllerV2Test {
             "pniIdentityKey": "%s",
             "deviceMessages": [],
             "devicePniSignedPrekeys": {},
-            "pniRegistrationIds": {}
+            "pniRegistrationIds": {"1": %d}
           }
-          """, encodeSessionId(sessionId), encodeRecoveryPassword(recoveryPassword), newNumber, Base64.getEncoder().encodeToString(IDENTITY_KEY.serialize()));
+          """, encodeSessionId(sessionId), encodeRecoveryPassword(recoveryPassword), newNumber, Base64.getEncoder().encodeToString(IDENTITY_KEY.serialize()), pniRegistrationId);
     }
 
     /**
      * Valid request JSON with the give session ID
      */
     private static String requestJson(final String sessionId, final String newNumber) {
-      return requestJson(sessionId, new byte[0], newNumber);
+      return requestJson(sessionId, new byte[0], newNumber, 123);
     }
 
     /**

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

@@ -576,6 +576,56 @@ class DeviceControllerTest {
     );
   }
 
+  @ParameterizedTest
+  @MethodSource
+  void linkDeviceRegistrationId(final int registrationId, final int expectedStatusCode) {
+    final Device existingDevice = mock(Device.class);
+    when(existingDevice.getId()).thenReturn(Device.PRIMARY_ID);
+    when(AuthHelper.VALID_ACCOUNT.getDevices()).thenReturn(List.of(existingDevice));
+
+    VerificationCode deviceCode = resources.getJerseyTest()
+        .target("/v1/devices/provisioning/code")
+        .request()
+        .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+        .get(VerificationCode.class);
+
+    final ECKeyPair aciIdentityKeyPair = Curve.generateKeyPair();
+    final ECKeyPair pniIdentityKeyPair = Curve.generateKeyPair();
+
+    final ECSignedPreKey aciSignedPreKey = KeysHelper.signedECPreKey(1, aciIdentityKeyPair);
+    final ECSignedPreKey pniSignedPreKey = KeysHelper.signedECPreKey(2, pniIdentityKeyPair);
+    final KEMSignedPreKey aciPqLastResortPreKey = KeysHelper.signedKEMPreKey(3, aciIdentityKeyPair);
+    final KEMSignedPreKey pniPqLastResortPreKey = KeysHelper.signedKEMPreKey(4, pniIdentityKeyPair);
+
+    when(account.getIdentityKey(IdentityType.ACI)).thenReturn(new IdentityKey(aciIdentityKeyPair.getPublicKey()));
+    when(account.getIdentityKey(IdentityType.PNI)).thenReturn(new IdentityKey(pniIdentityKeyPair.getPublicKey()));
+
+    when(keysManager.storeEcSignedPreKeys(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
+    when(keysManager.storePqLastResort(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
+
+    final LinkDeviceRequest request = new LinkDeviceRequest(deviceCode.verificationCode(),
+        new AccountAttributes(false, registrationId, null, null, true, null),
+        new DeviceActivationRequest(aciSignedPreKey, pniSignedPreKey, aciPqLastResortPreKey, pniPqLastResortPreKey, Optional.of(new ApnRegistrationId("apn", null)), Optional.empty()));
+
+    try (final Response response = resources.getJerseyTest()
+        .target("/v1/devices/link")
+        .request()
+        .header("Authorization", AuthHelper.getProvisioningAuthHeader(AuthHelper.VALID_NUMBER, "password1"))
+        .put(Entity.entity(request, MediaType.APPLICATION_JSON_TYPE))) {
+      assertEquals(expectedStatusCode, response.getStatus());
+    }
+  }
+
+  private static Stream<Arguments> linkDeviceRegistrationId() {
+    return Stream.of(
+        Arguments.of(0x3FFF, 200),
+        Arguments.of(0, 422),
+        Arguments.of(-1, 422),
+        Arguments.of(0x3FFF + 1, 422),
+        Arguments.of(Integer.MAX_VALUE, 422)
+    );
+  }
+
   private static ECSignedPreKey ecSignedPreKeyWithBadSignature(final ECSignedPreKey signedPreKey) {
     return new ECSignedPreKey(signedPreKey.keyId(),
         signedPreKey.publicKey(),