Validate registration ids for new accounts

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

@@ -131,6 +131,7 @@ public class AccountController {
   private static final String ACCOUNT_VERIFY_COUNTER_NAME = name(AccountController.class, "verify");
   private static final String CAPTCHA_ATTEMPT_COUNTER_NAME = name(AccountController.class, "captcha");
   private static final String CHALLENGE_ISSUED_COUNTER_NAME = name(AccountController.class, "challengeIssued");
+  private static final String INVALID_ACCOUNT_ATTRS_COUNTER_NAME = name(AccountController.class, "invalidAccountAttrs");
 
   private static final DistributionSummary REREGISTRATION_IDLE_DAYS_DISTRIBUTION = DistributionSummary
       .builder(name(AccountController.class, "reregistrationIdleDays"))
@@ -391,6 +392,11 @@ public AccountIdentityResponse verifyAccount(@PathParam("verification_code") Str
 
     rateLimiters.getVerifyLimiter().validate(number);
 
+    if (!AccountsManager.validNewAccountAttributes(accountAttributes)) {
+      Metrics.counter(INVALID_ACCOUNT_ATTRS_COUNTER_NAME, Tags.of(UserAgentTagUtil.getPlatformTag(userAgent))).increment();
+      throw new WebApplicationException(Response.status(422, "account attributes invalid").build());
+    }
+
     // Note that successful verification depends on being able to find a stored verification code for the given number.
     // We check that numbers are normalized before we store verification codes, and so don't need to re-assert
     // normalization here.

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

@@ -69,6 +69,7 @@ public class RegistrationController {
   private static final String REGION_CODE_TAG_NAME = "regionCode";
   private static final String VERIFICATION_TYPE_TAG_NAME = "verification";
   private static final String ACCOUNT_ACTIVATED_TAG_NAME = "accountActivated";
+  private static final String INVALID_ACCOUNT_ATTRS_COUNTER_NAME = name(RegistrationController.class, "invalidAccountAttrs");
 
   private final AccountsManager accounts;
   private final PhoneVerificationTokenManager phoneVerificationTokenManager;
@@ -118,6 +119,10 @@ public AccountIdentityResponse register(
     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/storage/AccountsManager.java

@@ -28,6 +28,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.OptionalInt;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicReference;
@@ -330,6 +331,21 @@ public Account changeNumber(final Account account,
     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 byte[] pniIdentityKey,
       final Map<Long, SignedPreKey> pniSignedPreKeys,

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

@@ -20,6 +20,7 @@ public class Device {
 
   public static final long MASTER_ID = 1;
   public static final int MAXIMUM_DEVICE_ID = 256;
+  public static final int MAX_REGISTRATION_ID = 0x3FFF;
   public static final List<Long> ALL_POSSIBLE_DEVICE_IDS = LongStream.range(1, MAXIMUM_DEVICE_ID).boxed().collect(Collectors.toList());
 
   @JsonProperty

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

@@ -1081,11 +1081,13 @@ void testVerifyCode() throws Exception {
     when(registrationServiceClient.checkVerificationCode(sessionId, "1234", AccountController.REGISTRATION_RPC_TIMEOUT))
         .thenReturn(CompletableFuture.completedFuture(true));
 
+    final AccountAttributes attrs = new AccountAttributes(true, 1, "test", "", true, new Device.DeviceCapabilities());
+
     resources.getJerseyTest()
         .target("/v1/accounts/code/1234")
         .request()
         .header(HttpHeaders.AUTHORIZATION, AuthHelper.getProvisioningAuthHeader(SENDER, "bar"))
-        .put(Entity.entity(new AccountAttributes(), MediaType.APPLICATION_JSON_TYPE), AccountIdentityResponse.class);
+        .put(Entity.entity(attrs, MediaType.APPLICATION_JSON_TYPE), AccountIdentityResponse.class);
 
     verify(accountsManager).create(eq(SENDER), eq("bar"), any(), any(), anyList());
 

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

@@ -16,12 +16,14 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.google.i18n.phonenumbers.PhoneNumberUtil;
 import io.dropwizard.testing.junit5.DropwizardExtensionsSupport;
 import io.dropwizard.testing.junit5.ResourceExtension;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.Base64;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
@@ -129,6 +131,51 @@ void unprocessableRequestJson() {
     }
   }
 
+  static Stream<Arguments> invalidRegistrationId() {
+    return Stream.of(
+        Arguments.of(Optional.of(1), Optional.of(1), 200),
+        Arguments.of(Optional.of(1), Optional.empty(), 200),
+        Arguments.of(Optional.of(0x3FFF), Optional.empty(), 200),
+        Arguments.of(Optional.empty(), Optional.of(1), 422),
+        Arguments.of(Optional.of(Integer.MAX_VALUE), Optional.empty(), 422),
+        Arguments.of(Optional.of(0x3FFF + 1), Optional.empty(), 422),
+        Arguments.of(Optional.of(1), Optional.of(0x3FFF + 1), 422)
+    );
+  }
+
+  @ParameterizedTest
+  @MethodSource()
+  void invalidRegistrationId(Optional<Integer> registrationId, Optional<Integer> pniRegistrationId, int statusCode) throws InterruptedException, JsonProcessingException {
+    final Invocation.Builder request = resources.getJerseyTest()
+        .target("/v1/registration")
+        .request()
+        .header(HttpHeaders.AUTHORIZATION, AuthHelper.getProvisioningAuthHeader(NUMBER, PASSWORD));
+    when(registrationServiceClient.getSession(any(), any()))
+        .thenReturn(
+            CompletableFuture.completedFuture(
+                Optional.of(new RegistrationServiceSession(new byte[16], NUMBER, true, null, null, null,
+                    SESSION_EXPIRATION_SECONDS))));
+    when(accountsManager.create(any(), any(), any(), any(), any()))
+        .thenReturn(mock(Account.class));
+
+    final String recoveryPassword = encodeRecoveryPassword(new byte[0]);
+
+    final Map<String, Object> accountAttrs = new HashMap<>();
+    accountAttrs.put("recoveryPassword", recoveryPassword);
+    registrationId.ifPresent(id -> accountAttrs.put("registrationId", id));
+    pniRegistrationId.ifPresent(id -> accountAttrs.put("pniRegistrationId", id));
+    final String json = SystemMapper.jsonMapper().writeValueAsString(Map.of(
+        "sessionId", encodeSessionId("sessionId"),
+        "recoveryPassword", recoveryPassword,
+        "accountAttributes", accountAttrs,
+        "skipDeviceTransfer", true
+    ));
+
+    try (Response response = request.post(Entity.json(json))) {
+      assertEquals(statusCode, response.getStatus());
+    }
+  }
+
   @Test
   void missingBasicAuthorization() {
     final Invocation.Builder request = resources.getJerseyTest()
@@ -745,7 +792,8 @@ private static String requestJson(final String sessionId, final byte[] recoveryP
           "sessionId": "%s",
           "recoveryPassword": "%s",
           "accountAttributes": {
-            "recoveryPassword": "%s"
+            "recoveryPassword": "%s",
+            "registrationId": 1
           },
           "skipDeviceTransfer": %s
         }