Fix account dropping new style registration locks

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

@@ -21,6 +21,7 @@ import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.SharedMetricRegistries;
 import com.codahale.metrics.annotation.Timed;
 import com.google.common.annotations.VisibleForTesting;
+import io.dropwizard.auth.Auth;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.AuthenticationCredentials;
@@ -83,7 +84,6 @@ import java.util.Optional;
 import java.util.UUID;
 
 import static com.codahale.metrics.MetricRegistry.name;
-import io.dropwizard.auth.Auth;
 
 @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
 @Path("/v1/accounts")
@@ -470,22 +470,13 @@ public class AccountController {
     device.setSignalingKey(attributes.getSignalingKey());
     device.setUserAgent(userAgent);
 
-    if (!Util.isEmpty(attributes.getPin())) {
+    setAccountRegistrationLockFromAttributes(account, attributes);
-      account.setPin(attributes.getPin());
-    } else if (!Util.isEmpty(attributes.getRegistrationLock())) {
-      AuthenticationCredentials credentials = new AuthenticationCredentials(attributes.getRegistrationLock());
-      account.setRegistrationLock(credentials.getHashedAuthenticationToken(), credentials.getSalt());
-    } else {
-      account.setPin(null);
-      account.setRegistrationLock(null, null);
-    }
 
     account.setUnidentifiedAccessKey(attributes.getUnidentifiedAccessKey());
     account.setUnrestrictedUnidentifiedAccess(attributes.isUnrestrictedUnidentifiedAccess());
 
     accounts.update(account);
   }
-
   @GET
   @Path("/whoami")
   @Produces(MediaType.APPLICATION_JSON)
@@ -614,7 +605,7 @@ public class AccountController {
     account.setNumber(number);
     account.setUuid(UUID.randomUUID());
     account.addDevice(device);
-    account.setPin(accountAttributes.getPin());
+    setAccountRegistrationLockFromAttributes(account, accountAttributes);
     account.setUnidentifiedAccessKey(accountAttributes.getUnidentifiedAccessKey());
     account.setUnrestrictedUnidentifiedAccess(accountAttributes.isUnrestrictedUnidentifiedAccess());
 
@@ -634,6 +625,18 @@ public class AccountController {
     return account;
   }
 
+  private void setAccountRegistrationLockFromAttributes(Account account, @Valid AccountAttributes attributes) {
+    if (!Util.isEmpty(attributes.getPin())) {
+      account.setPin(attributes.getPin());
+    } else if (!Util.isEmpty(attributes.getRegistrationLock())) {
+      AuthenticationCredentials credentials = new AuthenticationCredentials(attributes.getRegistrationLock());
+      account.setRegistrationLock(credentials.getHashedAuthenticationToken(), credentials.getSalt());
+    } else {
+      account.setPin(null);
+      account.setRegistrationLock(null, null);
+    }
+  }
+
   @VisibleForTesting protected
   VerificationCode generateVerificationCode(String number) {
     if (testDevices.containsKey(number)) {

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

@@ -1,11 +1,14 @@
 package org.whispersystems.textsecuregcm.tests.controllers;
 
 import com.google.common.collect.ImmutableSet;
+import io.dropwizard.auth.PolymorphicAuthValueFactoryProvider;
+import io.dropwizard.testing.junit.ResourceTestRule;
 import org.glassfish.jersey.test.grizzly.GrizzlyWebTestContainerFactory;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
+import org.mockito.ArgumentMatcher;
 import org.whispersystems.textsecuregcm.auth.AuthenticationCredentials;
 import org.whispersystems.textsecuregcm.auth.DisabledPermittedAccount;
 import org.whispersystems.textsecuregcm.auth.ExternalServiceCredentialGenerator;
@@ -54,11 +57,18 @@ import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 
-import io.dropwizard.auth.PolymorphicAuthValueFactoryProvider;
-import io.dropwizard.testing.junit.ResourceTestRule;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Matchers.anyString;
-import static org.mockito.Mockito.*;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.isA;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.notNull;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
 
 public class AccountControllerTest {
 
@@ -575,6 +585,36 @@ public class AccountControllerTest {
     verify(pinLimiter).validate(eq(SENDER_REG_LOCK));
   }
 
+  @Test
+  public void testVerifyRegistrationLockSetsRegistrationLockOnNewAccount() throws Exception {
+
+    AccountCreationResult result =
+        resources.getJerseyTest()
+                 .target(String.format("/v1/accounts/code/%s", "666666"))
+                 .request()
+                 .header("Authorization", AuthHelper.getAuthHeader(SENDER_REG_LOCK, "bar"))
+                 .put(Entity.entity(new AccountAttributes("keykeykeykey", false, 3333, null, null, Hex.toStringCondensed(registration_lock_key)),
+                                    MediaType.APPLICATION_JSON_TYPE), AccountCreationResult.class);
+
+    assertThat(result.getUuid()).isNotNull();
+
+    verify(pinLimiter).validate(eq(SENDER_REG_LOCK));
+
+    verify(accountsManager).create(argThat(new ArgumentMatcher<>() {
+      @Override
+      public boolean matches(final Account account) {
+        final StoredRegistrationLock regLock = account.getRegistrationLock();
+        return regLock.requiresClientRegistrationLock() && regLock.verify(Hex.toStringCondensed(registration_lock_key), null);
+      }
+
+      @Override
+      public String toString() {
+        return "Account that has registration lock set";
+      }
+    }));
+
+  }
+
   @Test
   public void testVerifyRegistrationLockOld() throws Exception {
     StoredRegistrationLock lock = senderRegLockAccount.getRegistrationLock();