Switch to a disallowed prefix model instead of a disallowed country code model

pom.xml

@@ -337,6 +337,7 @@
         <version>0.6.1</version>
         <configuration>
           <protocArtifact>com.google.protobuf:protoc:3.18.0:exe:${os.detected.classifier}</protocArtifact>
+          <checkStaleness>true</checkStaleness>
         </configuration>
         <executions>
           <execution>

service/config/sample.yml

@@ -83,6 +83,10 @@ directory:
       replicationPassword:           # CDS replication endpoint password
       replicationCaCertificate:      # CDS replication endpoint TLS certificate trust root
 
+directoryV2:
+  client: # Configuration for interfacing with Contact Discovery Service v2 cluster
+    userAuthenticationTokenSharedSecret: # base64-encoded secret shared with CDS to generate auth tokens for Signal users
+
 messageCache: # Redis server configuration for message store cache
   persistDelayMinutes:
 

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java

@@ -27,6 +27,7 @@ import org.whispersystems.textsecuregcm.configuration.DatabaseConfiguration;
 import org.whispersystems.textsecuregcm.configuration.DatadogConfiguration;
 import org.whispersystems.textsecuregcm.configuration.DeletedAccountsDynamoDbConfiguration;
 import org.whispersystems.textsecuregcm.configuration.DirectoryConfiguration;
+import org.whispersystems.textsecuregcm.configuration.DirectoryV2Configuration;
 import org.whispersystems.textsecuregcm.configuration.DonationConfiguration;
 import org.whispersystems.textsecuregcm.configuration.DynamoDbClientConfiguration;
 import org.whispersystems.textsecuregcm.configuration.DynamoDbConfiguration;
@@ -125,6 +126,11 @@ public class WhisperServerConfiguration extends Configuration {
   @JsonProperty
   private DirectoryConfiguration directory;
 
+  @NotNull
+  @Valid
+  @JsonProperty
+  private DirectoryV2Configuration directoryV2;
+
   @NotNull
   @Valid
   @JsonProperty
@@ -165,6 +171,11 @@ public class WhisperServerConfiguration extends Configuration {
   @JsonProperty
   private AccountsDynamoDbConfiguration accountsDynamoDb;
 
+  @Valid
+  @NotNull
+  @JsonProperty
+  private DynamoDbConfiguration phoneNumberIdentifiersDynamoDb;
+
   @Valid
   @NotNull
   @JsonProperty
@@ -390,6 +401,10 @@ public class WhisperServerConfiguration extends Configuration {
     return directory;
   }
 
+  public DirectoryV2Configuration getDirectoryV2Configuration() {
+    return directoryV2;
+  }
+
   public SecureStorageServiceConfiguration getSecureStorageServiceConfiguration() {
     return storageService;
   }
@@ -426,6 +441,10 @@ public class WhisperServerConfiguration extends Configuration {
     return accountsDynamoDb;
   }
 
+  public DynamoDbConfiguration getPhoneNumberIdentifiersDynamoDbConfiguration() {
+    return phoneNumberIdentifiersDynamoDb;
+  }
+
   public DeletedAccountsDynamoDbConfiguration getDeletedAccountsDynamoDbConfiguration() {
     return deletedAccountsDynamoDb;
   }

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -53,7 +53,6 @@ import java.util.concurrent.TimeUnit;
 import javax.servlet.DispatcherType;
 import javax.servlet.FilterRegistration;
 import javax.servlet.ServletRegistration;
-import javax.ws.rs.container.DynamicFeature;
 import org.eclipse.jetty.servlets.CrossOriginFilter;
 import org.glassfish.jersey.server.ServerProperties;
 import org.jdbi.v3.core.Jdbi;
@@ -88,6 +87,7 @@ import org.whispersystems.textsecuregcm.controllers.CertificateController;
 import org.whispersystems.textsecuregcm.controllers.ChallengeController;
 import org.whispersystems.textsecuregcm.controllers.DeviceController;
 import org.whispersystems.textsecuregcm.controllers.DirectoryController;
+import org.whispersystems.textsecuregcm.controllers.DirectoryV2Controller;
 import org.whispersystems.textsecuregcm.controllers.DonationController;
 import org.whispersystems.textsecuregcm.controllers.KeepAliveController;
 import org.whispersystems.textsecuregcm.controllers.KeysController;
@@ -115,6 +115,7 @@ import org.whispersystems.textsecuregcm.limits.RateLimitResetMetricsManager;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.limits.UnsealedSenderRateLimiter;
 import org.whispersystems.textsecuregcm.liquibase.NameableMigrationsBundle;
+import org.whispersystems.textsecuregcm.mappers.CompletionExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.DeviceLimitExceededExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.IOExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.ImpossiblePhoneNumberExceptionMapper;
@@ -169,6 +170,7 @@ import org.whispersystems.textsecuregcm.storage.AccountDatabaseCrawlerCache;
 import org.whispersystems.textsecuregcm.storage.AccountDatabaseCrawlerListener;
 import org.whispersystems.textsecuregcm.storage.Accounts;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.AssignPhoneNumberIdentifierCrawlerListener;
 import org.whispersystems.textsecuregcm.storage.ContactDiscoveryWriter;
 import org.whispersystems.textsecuregcm.storage.DeletedAccounts;
 import org.whispersystems.textsecuregcm.storage.DeletedAccountsDirectoryReconciler;
@@ -185,6 +187,7 @@ import org.whispersystems.textsecuregcm.storage.MessagesCache;
 import org.whispersystems.textsecuregcm.storage.MessagesDynamoDb;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.NonNormalizedAccountCrawlerListener;
+import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
 import org.whispersystems.textsecuregcm.storage.Profiles;
 import org.whispersystems.textsecuregcm.storage.ProfilesManager;
 import org.whispersystems.textsecuregcm.storage.PubSubManager;
@@ -329,6 +332,10 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     DynamoDbClient accountsDynamoDbClient = DynamoDbFromConfig.client(config.getAccountsDynamoDbConfiguration(),
         software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
 
+    DynamoDbClient phoneNumberIdentifiersDynamoDbClient =
+        DynamoDbFromConfig.client(config.getPhoneNumberIdentifiersDynamoDbConfiguration(),
+            software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
+
     DynamoDbClient deletedAccountsDynamoDbClient = DynamoDbFromConfig.client(config.getDeletedAccountsDynamoDbConfiguration(),
         software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
 
@@ -364,7 +371,10 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     Accounts accounts = new Accounts(accountsDynamoDbClient,
         config.getAccountsDynamoDbConfiguration().getTableName(),
         config.getAccountsDynamoDbConfiguration().getPhoneNumberTableName(),
+        config.getAccountsDynamoDbConfiguration().getPhoneNumberIdentifierTableName(),
         config.getAccountsDynamoDbConfiguration().getScanPageSize());
+    PhoneNumberIdentifiers phoneNumberIdentifiers = new PhoneNumberIdentifiers(phoneNumberIdentifiersDynamoDbClient,
+        config.getPhoneNumberIdentifiersDynamoDbConfiguration().getTableName());
     Usernames usernames = new Usernames(accountDatabase);
     ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase);
     Profiles profiles = new Profiles(accountDatabase);
@@ -423,9 +433,12 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     StripeManager stripeManager = new StripeManager(config.getStripe().getApiKey(), stripeExecutor,
         config.getStripe().getIdempotencyKeyGenerator());
 
-    ExternalServiceCredentialGenerator directoryCredentialsGenerator = new ExternalServiceCredentialGenerator(config.getDirectoryConfiguration().getDirectoryClientConfiguration().getUserAuthenticationTokenSharedSecret(),
-            config.getDirectoryConfiguration().getDirectoryClientConfiguration().getUserAuthenticationTokenUserIdSecret(),
-            true);
+    ExternalServiceCredentialGenerator directoryCredentialsGenerator = new ExternalServiceCredentialGenerator(
+        config.getDirectoryConfiguration().getDirectoryClientConfiguration().getUserAuthenticationTokenSharedSecret(),
+        config.getDirectoryConfiguration().getDirectoryClientConfiguration().getUserAuthenticationTokenUserIdSecret());
+    ExternalServiceCredentialGenerator directoryV2CredentialsGenerator = new ExternalServiceCredentialGenerator(
+        config.getDirectoryV2Configuration().getDirectoryV2ClientConfiguration()
+            .getUserAuthenticationTokenSharedSecret(), false);
 
     DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager =
         new DynamicConfigurationManager<>(config.getAppConfig().getApplication(),
@@ -440,9 +453,12 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     TwilioVerifyExperimentEnrollmentManager verifyExperimentEnrollmentManager = new TwilioVerifyExperimentEnrollmentManager(
         config.getVoiceVerificationConfiguration(), experimentEnrollmentManager);
 
-    ExternalServiceCredentialGenerator storageCredentialsGenerator   = new ExternalServiceCredentialGenerator(config.getSecureStorageServiceConfiguration().getUserAuthenticationTokenSharedSecret(), new byte[0], false);
-    ExternalServiceCredentialGenerator backupCredentialsGenerator    = new ExternalServiceCredentialGenerator(config.getSecureBackupServiceConfiguration().getUserAuthenticationTokenSharedSecret(), new byte[0], false);
-    ExternalServiceCredentialGenerator paymentsCredentialsGenerator  = new ExternalServiceCredentialGenerator(config.getPaymentsServiceConfiguration().getUserAuthenticationTokenSharedSecret(), new byte[0], false);
+    ExternalServiceCredentialGenerator storageCredentialsGenerator = new ExternalServiceCredentialGenerator(
+        config.getSecureStorageServiceConfiguration().getUserAuthenticationTokenSharedSecret(), true);
+    ExternalServiceCredentialGenerator backupCredentialsGenerator = new ExternalServiceCredentialGenerator(
+        config.getSecureBackupServiceConfiguration().getUserAuthenticationTokenSharedSecret(), true);
+    ExternalServiceCredentialGenerator paymentsCredentialsGenerator = new ExternalServiceCredentialGenerator(
+        config.getPaymentsServiceConfiguration().getUserAuthenticationTokenSharedSecret(), true);
 
     SecureBackupClient         secureBackupClient         = new SecureBackupClient(backupCredentialsGenerator, backupServiceExecutor, config.getSecureBackupServiceConfiguration());
     SecureStorageClient        secureStorageClient        = new SecureStorageClient(storageCredentialsGenerator, storageServiceExecutor, config.getSecureStorageServiceConfiguration());
@@ -453,14 +469,14 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     UsernamesManager           usernamesManager           = new UsernamesManager(usernames, reservedUsernames, cacheCluster);
     ProfilesManager            profilesManager            = new ProfilesManager(profiles, cacheCluster);
     MessagesCache              messagesCache              = new MessagesCache(messagesCluster, messagesCluster, keyspaceNotificationDispatchExecutor);
-    PushLatencyManager         pushLatencyManager         = new PushLatencyManager(metricsCluster);
+    PushLatencyManager         pushLatencyManager         = new PushLatencyManager(metricsCluster, dynamicConfigurationManager);
     ReportMessageManager       reportMessageManager       = new ReportMessageManager(reportMessageDynamoDb, rateLimitersCluster, Metrics.globalRegistry, config.getReportMessageConfiguration().getCounterTtl());
     MessagesManager            messagesManager            = new MessagesManager(messagesDynamoDb, messagesCache, pushLatencyManager, reportMessageManager);
     DeletedAccountsManager deletedAccountsManager = new DeletedAccountsManager(deletedAccounts,
         deletedAccountsLockDynamoDbClient, config.getDeletedAccountsLockDynamoDbConfiguration().getTableName());
-    AccountsManager accountsManager = new AccountsManager(accounts, cacheCluster,
+    AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
         deletedAccountsManager, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager,
-        pendingAccountsManager, secureStorageClient, secureBackupClient, clientPresenceManager);
+        pendingAccountsManager, secureStorageClient, secureBackupClient, clientPresenceManager, clock);
     RemoteConfigsManager remoteConfigsManager = new RemoteConfigsManager(remoteConfigs);
     DeadLetterHandler          deadLetterHandler          = new DeadLetterHandler(accountsManager, messagesManager);
     DispatchManager            dispatchManager            = new DispatchManager(pubSubClientFactory, Optional.of(deadLetterHandler));
@@ -510,9 +526,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
 
     MessagePersister messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager, dynamicConfigurationManager, Duration.ofMinutes(config.getMessageCacheConfiguration().getPersistDelayMinutes()));
 
-    // TODO listeners must be ordered so that ones that directly update accounts come last, so that read-only ones are not working with stale data
-    final List<AccountDatabaseCrawlerListener> accountDatabaseCrawlerListeners = new ArrayList<>();
-
+    final List<AccountDatabaseCrawlerListener> directoryReconciliationAccountDatabaseCrawlerListeners = new ArrayList<>();
     final List<DeletedAccountsDirectoryReconciler> deletedAccountsDirectoryReconcilers = new ArrayList<>();
     for (DirectoryServerConfiguration directoryServerConfiguration : config.getDirectoryConfiguration()
         .getDirectoryServerConfiguration()) {
@@ -522,25 +536,34 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
           directoryServerConfiguration.getReplicationName(), directoryReconciliationClient,
           dynamicConfigurationManager);
       // reconcilers are read-only
-      accountDatabaseCrawlerListeners.add(directoryReconciler);
+      directoryReconciliationAccountDatabaseCrawlerListeners.add(directoryReconciler);
 
       final DeletedAccountsDirectoryReconciler deletedAccountsDirectoryReconciler = new DeletedAccountsDirectoryReconciler(
           directoryServerConfiguration.getReplicationName(), directoryReconciliationClient);
       deletedAccountsDirectoryReconcilers.add(deletedAccountsDirectoryReconciler);
     }
-    accountDatabaseCrawlerListeners.add(new NonNormalizedAccountCrawlerListener(accountsManager, metricsCluster));
-    accountDatabaseCrawlerListeners.add(new ContactDiscoveryWriter(accountsManager));
-    // PushFeedbackProcessor may update device properties
-    accountDatabaseCrawlerListeners.add(new PushFeedbackProcessor(accountsManager));
-    // delete accounts last
-    accountDatabaseCrawlerListeners.add(new AccountCleaner(accountsManager));
 
-    HttpClient                currencyClient  = HttpClient.newBuilder().version(HttpClient.Version.HTTP_2).connectTimeout(Duration.ofSeconds(10)).build();
-    FixerClient               fixerClient     = new FixerClient(currencyClient, config.getPaymentsServiceConfiguration().getFixerApiKey());
-    FtxClient                 ftxClient       = new FtxClient(currencyClient);
-    CurrencyConversionManager currencyManager = new CurrencyConversionManager(fixerClient, ftxClient, config.getPaymentsServiceConfiguration().getPaymentCurrencies());
+    AccountDatabaseCrawlerCache directoryReconciliationAccountDatabaseCrawlerCache = new AccountDatabaseCrawlerCache(
+        cacheCluster, AccountDatabaseCrawlerCache.DIRECTORY_RECONCILER_PREFIX);
+    AccountDatabaseCrawler directoryReconciliationAccountDatabaseCrawler = new AccountDatabaseCrawler(accountsManager,
+        directoryReconciliationAccountDatabaseCrawlerCache, directoryReconciliationAccountDatabaseCrawlerListeners,
+        config.getAccountDatabaseCrawlerConfiguration().getChunkSize(),
+        config.getAccountDatabaseCrawlerConfiguration().getChunkIntervalMs()
+    );
 
-    AccountDatabaseCrawlerCache accountDatabaseCrawlerCache = new AccountDatabaseCrawlerCache(cacheCluster);
+    // TODO listeners must be ordered so that ones that directly update accounts come last, so that read-only ones are not working with stale data
+    final List<AccountDatabaseCrawlerListener> accountDatabaseCrawlerListeners = List.of(
+        new NonNormalizedAccountCrawlerListener(accountsManager, metricsCluster),
+        new ContactDiscoveryWriter(accountsManager),
+        new AssignPhoneNumberIdentifierCrawlerListener(accountsManager, phoneNumberIdentifiers),
+        // PushFeedbackProcessor may update device properties
+        new PushFeedbackProcessor(accountsManager),
+        // delete accounts last
+        new AccountCleaner(accountsManager)
+    );
+
+    AccountDatabaseCrawlerCache accountDatabaseCrawlerCache = new AccountDatabaseCrawlerCache(cacheCluster,
+        AccountDatabaseCrawlerCache.GENERAL_PURPOSE_PREFIX);
     AccountDatabaseCrawler accountDatabaseCrawler = new AccountDatabaseCrawler(accountsManager,
         accountDatabaseCrawlerCache, accountDatabaseCrawlerListeners,
         config.getAccountDatabaseCrawlerConfiguration().getChunkSize(),
@@ -549,12 +572,18 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
 
     DeletedAccountsTableCrawler deletedAccountsTableCrawler = new DeletedAccountsTableCrawler(deletedAccountsManager, deletedAccountsDirectoryReconcilers, cacheCluster, recurringJobExecutor);
 
+    HttpClient                currencyClient  = HttpClient.newBuilder().version(HttpClient.Version.HTTP_2).connectTimeout(Duration.ofSeconds(10)).build();
+    FixerClient               fixerClient     = new FixerClient(currencyClient, config.getPaymentsServiceConfiguration().getFixerApiKey());
+    FtxClient                 ftxClient       = new FtxClient(currencyClient);
+    CurrencyConversionManager currencyManager = new CurrencyConversionManager(fixerClient, ftxClient, config.getPaymentsServiceConfiguration().getPaymentCurrencies());
+
     apnSender.setApnFallbackManager(apnFallbackManager);
     environment.lifecycle().manage(new ApplicationShutdownMonitor());
     environment.lifecycle().manage(apnFallbackManager);
     environment.lifecycle().manage(pubSubManager);
     environment.lifecycle().manage(messageSender);
     environment.lifecycle().manage(accountDatabaseCrawler);
+    environment.lifecycle().manage(directoryReconciliationAccountDatabaseCrawler);
     environment.lifecycle().manage(deletedAccountsTableCrawler);
     environment.lifecycle().manage(remoteConfigsManager);
     environment.lifecycle().manage(messagesCache);
@@ -632,6 +661,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         new ChallengeController(rateLimitChallengeManager),
         new DeviceController(pendingDevicesManager, accountsManager, messagesManager, keysDynamoDb, rateLimiters, config.getMaxDevices()),
         new DirectoryController(directoryCredentialsGenerator),
+        new DirectoryV2Controller(directoryV2CredentialsGenerator),
         new DonationController(clock, zkReceiptOperations, redeemedReceiptsManager, accountsManager, config.getBadges(),
             ReceiptCredentialPresentation::new, stripeExecutor, config.getDonationConfiguration(), config.getStripe()),
         new MessageController(rateLimiters, messageSender, receiptSender, accountsManager, messagesManager, unsealedSenderRateLimiter, apnFallbackManager,
@@ -744,6 +774,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
 
     List.of(
         new LoggingUnhandledExceptionMapper(),
+        new CompletionExceptionMapper(),
         new IOExceptionMapper(),
         new RateLimitExceededExceptionMapper(),
         new InvalidWebsocketAddressExceptionMapper(),

service/src/main/java/org/whispersystems/textsecuregcm/auth/AuthEnablementRefreshRequirementProvider.java

@@ -76,7 +76,7 @@ public class AuthEnablementRefreshRequirementProvider implements WebsocketRefres
       @SuppressWarnings("unchecked") final Map<Long, Boolean> initialDevicesEnabled =
           (Map<Long, Boolean>) requestEvent.getContainerRequest().getProperty(DEVICES_ENABLED);
 
-      return accountsManager.get((UUID) requestEvent.getContainerRequest().getProperty(ACCOUNT_UUID)).map(account -> {
+      return accountsManager.getByAccountIdentifier((UUID) requestEvent.getContainerRequest().getProperty(ACCOUNT_UUID)).map(account -> {
         final Set<Long> deviceIdsToDisplace;
         final Map<Long, Boolean> currentDevicesEnabled = buildDevicesEnabledMap(account);
 

service/src/main/java/org/whispersystems/textsecuregcm/auth/BaseAccountAuthenticator.java

@@ -78,7 +78,7 @@ public class BaseAccountAuthenticator {
         deviceId = identifierAndDeviceId.second();
       }
 
-      Optional<Account> account = accountsManager.get(accountUuid);
+      Optional<Account> account = accountsManager.getByAccountIdentifier(accountUuid);
 
       if (account.isEmpty()) {
         failureReason = "noSuchAccount";

service/src/main/java/org/whispersystems/textsecuregcm/auth/ExternalServiceCredentialGenerator.java

@@ -1,97 +1,71 @@
 /*
- * Copyright 2013-2020 Signal Messenger, LLC
+ * Copyright 2013-2021 Signal Messenger, LLC
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 
 package org.whispersystems.textsecuregcm.auth;
 
-import org.apache.commons.codec.DecoderException;
-import org.apache.commons.codec.binary.Hex;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import org.whispersystems.textsecuregcm.util.Util;
-
+import com.google.common.annotations.VisibleForTesting;
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
+import java.time.Clock;
 import javax.crypto.Mac;
 import javax.crypto.spec.SecretKeySpec;
-import java.security.InvalidKeyException;
-import java.security.MessageDigest;
-import java.security.NoSuchAlgorithmException;
-import java.util.concurrent.TimeUnit;
+import org.apache.commons.codec.binary.Hex;
+import org.whispersystems.textsecuregcm.util.Util;
 
 public class ExternalServiceCredentialGenerator {
 
-  private final Logger logger = LoggerFactory.getLogger(ExternalServiceCredentialGenerator.class);
-
   private final byte[] key;
   private final byte[] userIdKey;
   private final boolean usernameDerivation;
+  private final boolean prependUsername;
+  private final Clock clock;
 
-  public ExternalServiceCredentialGenerator(byte[] key, byte[] userIdKey, boolean usernameDerivation) {
-    this.key                = key;
-    this.userIdKey          = userIdKey;
-    this.usernameDerivation = usernameDerivation;
+  public ExternalServiceCredentialGenerator(byte[] key, byte[] userIdKey) {
+    this(key, userIdKey, true, true);
   }
 
-  public ExternalServiceCredentials generateFor(String number) {
-    Mac    mac                = getMacInstance();
-    String username           = getUserId(number, mac, usernameDerivation);
-    long   currentTimeSeconds = System.currentTimeMillis() / 1000;
-    String prefix             = username + ":"  + currentTimeSeconds;
-    String output             = Hex.encodeHexString(Util.truncate(getHmac(key, prefix.getBytes(), mac), 10));
-    String token              = prefix + ":" + output;
+  public ExternalServiceCredentialGenerator(byte[] key, boolean prependUsername) {
+    this(key, new byte[0], false, prependUsername);
+  }
+
+  @VisibleForTesting
+  public ExternalServiceCredentialGenerator(byte[] key, byte[] userIdKey, boolean usernameDerivation) {
+    this(key, userIdKey, usernameDerivation, true);
+  }
+
+  private ExternalServiceCredentialGenerator(byte[] key, byte[] userIdKey, boolean usernameDerivation,
+      boolean prependUsername) {
+    this(key, userIdKey, usernameDerivation, prependUsername, Clock.systemUTC());
+  }
+
+  @VisibleForTesting
+  public ExternalServiceCredentialGenerator(byte[] key, byte[] userIdKey, boolean usernameDerivation,
+      boolean prependUsername, Clock clock) {
+    this.key = key;
+    this.userIdKey = userIdKey;
+    this.usernameDerivation = usernameDerivation;
+    this.prependUsername = prependUsername;
+    this.clock = clock;
+  }
+
+  public ExternalServiceCredentials generateFor(String identity) {
+    Mac mac = getMacInstance();
+    String username = getUserId(identity, mac, usernameDerivation);
+    long currentTimeSeconds = clock.millis() / 1000;
+    String prefix = username + ":" + currentTimeSeconds;
+    String output = Hex.encodeHexString(Util.truncate(getHmac(key, prefix.getBytes(), mac), 10));
+    String token = (prependUsername ? prefix : currentTimeSeconds) + ":" + output;
 
     return new ExternalServiceCredentials(username, token);
   }
 
-
-  public boolean isValid(String token, String number, long currentTimeMillis) {
-    String[] parts = token.split(":");
-    Mac      mac   = getMacInstance();
-
-    if (parts.length != 3) {
-      return false;
-    }
-
-    if (!getUserId(number, mac, usernameDerivation).equals(parts[0])) {
-      return false;
-    }
-
-    if (!isValidTime(parts[1], currentTimeMillis)) {
-      return false;
-    }
-
-    return isValidSignature(parts[0] + ":" + parts[1], parts[2], mac);
-  }
-
   private String getUserId(String number, Mac mac, boolean usernameDerivation) {
     if (usernameDerivation) return Hex.encodeHexString(Util.truncate(getHmac(userIdKey, number.getBytes(), mac), 10));
     else                    return number;
   }
 
-  private boolean isValidTime(String timeString, long currentTimeMillis) {
-    try {
-      long tokenTime = Long.parseLong(timeString);
-      long ourTime   = TimeUnit.MILLISECONDS.toSeconds(currentTimeMillis);
-
-      return TimeUnit.SECONDS.toHours(Math.abs(ourTime - tokenTime)) < 24;
-    } catch (NumberFormatException e) {
-      logger.warn("Number Format", e);
-      return false;
-    }
-  }
-
-  private boolean isValidSignature(String prefix, String suffix, Mac mac) {
-    try {
-      byte[] ourSuffix   = Util.truncate(getHmac(key, prefix.getBytes(), mac), 10);
-      byte[] theirSuffix = Hex.decodeHex(suffix.toCharArray());
-
-      return MessageDigest.isEqual(ourSuffix, theirSuffix);
-    } catch (DecoderException e) {
-      logger.warn("DirectoryCredentials", e);
-      return false;
-    }
-  }
-
   private Mac getMacInstance() {
     try {
       return Mac.getInstance("HmacSHA256");

service/src/main/java/org/whispersystems/textsecuregcm/configuration/AccountsDynamoDbConfiguration.java

@@ -8,6 +8,9 @@ public class AccountsDynamoDbConfiguration extends DynamoDbConfiguration {
   @NotNull
   private String phoneNumberTableName;
 
+  @NotNull
+  private String phoneNumberIdentifierTableName;
+
   private int scanPageSize = 100;
 
   @JsonProperty
@@ -15,6 +18,11 @@ public class AccountsDynamoDbConfiguration extends DynamoDbConfiguration {
     return phoneNumberTableName;
   }
 
+  @JsonProperty
+  public String getPhoneNumberIdentifierTableName() {
+    return phoneNumberIdentifierTableName;
+  }
+
   @JsonProperty
   public int getScanPageSize() {
     return scanPageSize;

service/src/main/java/org/whispersystems/textsecuregcm/configuration/DirectoryV2ClientConfiguration.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2013-2020 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+package org.whispersystems.textsecuregcm.configuration;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import org.whispersystems.textsecuregcm.util.ExactlySize;
+
+public class DirectoryV2ClientConfiguration {
+
+  private final byte[] userAuthenticationTokenSharedSecret;
+
+  @JsonCreator
+  public DirectoryV2ClientConfiguration(
+      @JsonProperty("userAuthenticationTokenSharedSecret") final byte[] userAuthenticationTokenSharedSecret) {
+    this.userAuthenticationTokenSharedSecret = userAuthenticationTokenSharedSecret;
+  }
+
+  @ExactlySize({32})
+  public byte[] getUserAuthenticationTokenSharedSecret() {
+    return userAuthenticationTokenSharedSecret;
+  }
+
+}

service/src/main/java/org/whispersystems/textsecuregcm/configuration/DirectoryV2Configuration.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2013-2020 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+package org.whispersystems.textsecuregcm.configuration;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import javax.validation.Valid;
+
+public class DirectoryV2Configuration {
+
+  private final DirectoryV2ClientConfiguration clientConfiguration;
+
+  @JsonCreator
+  public DirectoryV2Configuration(@JsonProperty("client") DirectoryV2ClientConfiguration clientConfiguration) {
+    this.clientConfiguration = clientConfiguration;
+  }
+
+  @Valid
+  public DirectoryV2ClientConfiguration getDirectoryV2ClientConfiguration() {
+    return clientConfiguration;
+  }
+}

service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfiguration.java

@@ -51,6 +51,10 @@ public class DynamicConfiguration {
   @JsonProperty
   private DynamicDirectoryReconcilerConfiguration directoryReconciler = new DynamicDirectoryReconcilerConfiguration();
 
+  @JsonProperty
+  @Valid
+  private DynamicPushLatencyConfiguration pushLatency = new DynamicPushLatencyConfiguration(Collections.emptyMap());
+
   public Optional<DynamicExperimentEnrollmentConfiguration> getExperimentEnrollmentConfiguration(
       final String experimentName) {
     return Optional.ofNullable(experiments.get(experimentName));
@@ -101,4 +105,8 @@ public class DynamicConfiguration {
   public DynamicDirectoryReconcilerConfiguration getDirectoryReconcilerConfiguration() {
     return directoryReconciler;
   }
+
+  public DynamicPushLatencyConfiguration getPushLatencyConfiguration() {
+    return pushLatency;
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicPaymentsConfiguration.java

@@ -7,14 +7,14 @@ package org.whispersystems.textsecuregcm.configuration.dynamic;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 import java.util.Collections;
-import java.util.Set;
+import java.util.List;
 
 public class DynamicPaymentsConfiguration {
 
   @JsonProperty
-  private Set<String> allowedCountryCodes = Collections.emptySet();
+  private List<String> disallowedPrefixes = Collections.emptyList();
 
-  public Set<String> getAllowedCountryCodes() {
-    return allowedCountryCodes;
+  public List<String> getDisallowedPrefixes() {
+    return disallowedPrefixes;
   }
 }

service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicPushLatencyConfiguration.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2013-2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.configuration.dynamic;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.vdurmont.semver4j.Semver;
+import org.whispersystems.textsecuregcm.util.ua.ClientPlatform;
+import java.util.Map;
+import java.util.Set;
+
+public class DynamicPushLatencyConfiguration {
+
+  private final Map<ClientPlatform, Set<Semver>> instrumentedVersions;
+
+  @JsonCreator
+  public DynamicPushLatencyConfiguration(@JsonProperty("instrumentedVersions") final Map<ClientPlatform, Set<Semver>> instrumentedVersions) {
+    this.instrumentedVersions = instrumentedVersions;
+  }
+
+  public Map<ClientPlatform, Set<Semver>> getInstrumentedVersions() {
+    return instrumentedVersions;
+  }
+}

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

@@ -353,11 +353,11 @@ public class AccountController {
     storedVerificationCode.flatMap(StoredVerificationCode::getTwilioVerificationSid)
         .ifPresent(smsSender::reportVerificationSucceeded);
 
-      Optional<Account> existingAccount = accounts.get(number);
+      Optional<Account> existingAccount = accounts.getByE164(number);
 
-      if (existingAccount.isPresent()) {
-        verifyRegistrationLock(existingAccount.get(), accountAttributes.getRegistrationLock());
-      }
+    if (existingAccount.isPresent()) {
+      verifyRegistrationLock(existingAccount.get(), accountAttributes.getRegistrationLock());
+    }
 
     if (availableForTransfer.orElse(false) && existingAccount.map(Account::isTransferSupported).orElse(false)) {
       throw new WebApplicationException(Response.status(409).build());
@@ -365,7 +365,8 @@ public class AccountController {
 
     rateLimiters.getVerifyLimiter().clear(number);
 
-    Account account = accounts.create(number, password, signalAgent, accountAttributes);
+    Account account = accounts.create(number, password, signalAgent, accountAttributes,
+        existingAccount.map(Account::getBadges).orElseGet(ArrayList::new));
 
     {
       metricRegistry.meter(name(AccountController.class, "verify", Util.getCountryCode(number))).mark();
@@ -411,7 +412,7 @@ public class AccountController {
     storedVerificationCode.flatMap(StoredVerificationCode::getTwilioVerificationSid)
         .ifPresent(smsSender::reportVerificationSucceeded);
 
-    final Optional<Account> existingAccount = accounts.get(request.getNumber());
+    final Optional<Account> existingAccount = accounts.getByE164(request.getNumber());
 
     if (existingAccount.isPresent()) {
       verifyRegistrationLock(existingAccount.get(), request.getRegistrationLock());

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

@@ -168,7 +168,7 @@ public class DeviceController {
       throw new WebApplicationException(Response.status(403).build());
     }
 
-    Optional<Account> account = accounts.get(number);
+    Optional<Account> account = accounts.getByE164(number);
 
     if (!account.isPresent()) {
       throw new WebApplicationException(Response.status(403).build());

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

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2013-2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+package org.whispersystems.textsecuregcm.controllers;
+
+import com.codahale.metrics.annotation.Timed;
+import io.dropwizard.auth.Auth;
+import java.util.Base64;
+import java.util.UUID;
+import javax.ws.rs.GET;
+import javax.ws.rs.Path;
+import javax.ws.rs.Produces;
+import javax.ws.rs.core.MediaType;
+import javax.ws.rs.core.Response;
+import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
+import org.whispersystems.textsecuregcm.auth.ExternalServiceCredentialGenerator;
+import org.whispersystems.textsecuregcm.auth.ExternalServiceCredentials;
+import org.whispersystems.textsecuregcm.util.ByteUtil;
+import org.whispersystems.textsecuregcm.util.UUIDUtil;
+import org.whispersystems.textsecuregcm.util.Util;
+
+@Path("/v2/directory")
+public class DirectoryV2Controller {
+
+  private final ExternalServiceCredentialGenerator directoryServiceTokenGenerator;
+
+  public DirectoryV2Controller(ExternalServiceCredentialGenerator userTokenGenerator) {
+    this.directoryServiceTokenGenerator = userTokenGenerator;
+  }
+
+  @Timed
+  @GET
+  @Path("/auth")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response getAuthToken(@Auth AuthenticatedAccount auth) {
+
+    final UUID uuid = auth.getAccount().getUuid();
+    final String e164 = auth.getAccount().getNumber();
+    final long e164AsLong = Long.parseLong(e164, e164.indexOf('+'), e164.length() - 1, 10);
+
+    final byte[] uuidAndNumber = ByteUtil.combine(UUIDUtil.toBytes(uuid), Util.longToByteArray(e164AsLong));
+    final String username = Base64.getEncoder().encodeToString(uuidAndNumber);
+
+    final ExternalServiceCredentials credentials = directoryServiceTokenGenerator.generateFor(username);
+
+    return Response.ok().entity(credentials).build();
+  }
+}

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

@@ -159,7 +159,7 @@ public class DonationController {
 
             @Override
             public boolean block() {
-              final Optional<Account> optionalAccount = accountsManager.get(auth.getAccount().getUuid());
+              final Optional<Account> optionalAccount = accountsManager.getByAccountIdentifier(auth.getAccount().getUuid());
               optionalAccount.ifPresent(account -> {
                 accountsManager.update(account, a -> {
                   a.addBadge(clock, new AccountBadge(badgeId, receiptExpiration, request.isVisible()));

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

@@ -129,7 +129,7 @@ public class KeysController {
 
     final Optional<Account> account = auth.map(AuthenticatedAccount::getAccount);
 
-    Optional<Account> target = accounts.get(targetUuid);
+    Optional<Account> target = accounts.getByAccountIdentifier(targetUuid);
     OptionalAccess.verify(account, accessKey, target, deviceId);
 
     assert (target.isPresent());

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

@@ -215,7 +215,7 @@ public class MessageController {
       Optional<Account> destination;
 
       if (!isSyncMessage) {
-        destination = accountsManager.get(destinationUuid);
+        destination = accountsManager.getByAccountIdentifier(destinationUuid);
       } else {
         destination = source.map(AuthenticatedAccount::getAccount);
       }
@@ -311,7 +311,7 @@ public class MessageController {
         .map(Recipient::getUuid)
         .distinct()
         .collect(Collectors.toUnmodifiableMap(Function.identity(), uuid -> {
-          Optional<Account> account = accountsManager.get(uuid);
+          Optional<Account> account = accountsManager.getByAccountIdentifier(uuid);
           if (account.isEmpty()) {
             throw new WebApplicationException(Status.NOT_FOUND);
           }

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

@@ -19,7 +19,6 @@ import java.util.Locale;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
-import java.util.Set;
 import java.util.UUID;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -76,7 +75,6 @@ import org.whispersystems.textsecuregcm.storage.UsernamesManager;
 import org.whispersystems.textsecuregcm.storage.VersionedProfile;
 import org.whispersystems.textsecuregcm.util.ExactlySize;
 import org.whispersystems.textsecuregcm.util.Pair;
-import org.whispersystems.textsecuregcm.util.Util;
 import software.amazon.awssdk.services.s3.S3Client;
 import software.amazon.awssdk.services.s3.model.DeleteObjectRequest;
 
@@ -137,13 +135,14 @@ public class ProfileController {
   @Produces(MediaType.APPLICATION_JSON)
   @Consumes(MediaType.APPLICATION_JSON)
   public Response setProfile(@Auth AuthenticatedAccount auth, @Valid CreateProfileRequest request) {
-    final Set<String> allowedPaymentsCountryCodes =
-        dynamicConfigurationManager.getConfiguration().getPaymentsConfiguration().getAllowedCountryCodes();
+    if (StringUtils.isNotBlank(request.getPaymentAddress())) {
+      final boolean hasDisallowedPrefix =
+          dynamicConfigurationManager.getConfiguration().getPaymentsConfiguration().getDisallowedPrefixes().stream()
+              .anyMatch(prefix -> auth.getAccount().getNumber().startsWith(prefix));
 
-    if (StringUtils.isNotBlank(request.getPaymentAddress()) &&
-        !allowedPaymentsCountryCodes.contains(Util.getCountryCode(auth.getAccount().getNumber()))) {
-
-      return Response.status(Status.FORBIDDEN).build();
+      if (hasDisallowedPrefix) {
+        return Response.status(Status.FORBIDDEN).build();
+      }
     }
 
     Optional<VersionedProfile> currentProfile = profilesManager.get(auth.getAccount().getUuid(), request.getVersion());
@@ -251,7 +250,7 @@ public class ProfileController {
         isSelf = uuid.equals(authedUuid);
       }
 
-      Optional<Account> accountProfile = accountsManager.get(uuid);
+      Optional<Account> accountProfile = accountsManager.getByAccountIdentifier(uuid);
       OptionalAccess.verify(requestAccount, accessKey, accountProfile);
 
       assert(accountProfile.isPresent());
@@ -316,7 +315,7 @@ public class ProfileController {
 
     final boolean isSelf = auth.getAccount().getUuid().equals(uuid.get());
 
-    Optional<Account> accountProfile = accountsManager.get(uuid.get());
+    Optional<Account> accountProfile = accountsManager.getByAccountIdentifier(uuid.get());
 
     if (accountProfile.isEmpty()) {
       throw new WebApplicationException(Response.status(Response.Status.NOT_FOUND).build());
@@ -398,7 +397,7 @@ public class ProfileController {
       isSelf = authedUuid.equals(identifier);
     }
 
-    Optional<Account> accountProfile = accountsManager.get(identifier);
+    Optional<Account> accountProfile = accountsManager.getByAccountIdentifier(identifier);
     OptionalAccess.verify(auth.map(AuthenticatedAccount::getAccount), accessKey, accountProfile);
 
     Optional<String> username = usernamesManager.get(accountProfile.get().getUuid());

service/src/main/java/org/whispersystems/textsecuregcm/mappers/CompletionExceptionMapper.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2020 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.mappers;
+
+import java.util.concurrent.CompletionException;
+import javax.ws.rs.core.Context;
+import javax.ws.rs.core.Response;
+import javax.ws.rs.ext.ExceptionMapper;
+import javax.ws.rs.ext.Provider;
+import javax.ws.rs.ext.Providers;
+
+@Provider
+public class CompletionExceptionMapper implements ExceptionMapper<CompletionException> {
+
+  @Context
+  private Providers providers;
+
+  @Override
+  public Response toResponse(final CompletionException exception) {
+    final Throwable cause = exception.getCause();
+
+    if (cause != null) {
+      final Class type = cause.getClass();
+      return providers.getExceptionMapper(type).toResponse(cause);
+    }
+
+    return Response.serverError().build();
+  }
+}

service/src/main/java/org/whispersystems/textsecuregcm/metrics/PushLatencyManager.java

@@ -6,17 +6,34 @@
 package org.whispersystems.textsecuregcm.metrics;
 
 import com.codahale.metrics.MetricRegistry;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.google.common.annotations.VisibleForTesting;
+import com.vdurmont.semver4j.Semver;
 import io.lettuce.core.SetArgs;
-import io.lettuce.core.cluster.api.async.RedisAdvancedClusterAsyncCommands;
 import io.micrometer.core.instrument.Metrics;
-import io.micrometer.core.instrument.Tags;
-import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
-
+import io.micrometer.core.instrument.Tag;
+import java.time.Clock;
 import java.time.Duration;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.commons.lang3.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
+import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
+import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
+import org.whispersystems.textsecuregcm.util.SystemMapper;
+import org.whispersystems.textsecuregcm.util.ua.UnrecognizedUserAgentException;
+import org.whispersystems.textsecuregcm.util.ua.UserAgent;
+import org.whispersystems.textsecuregcm.util.ua.UserAgentUtil;
 
 /**
  * Measures and records the latency between sending a push notification to a device and that device draining its queue
@@ -28,48 +45,139 @@ import java.util.concurrent.TimeUnit;
  * device and records the time elapsed since the push notification timestamp as a latency observation.
  */
 public class PushLatencyManager {
-    private static final String TIMER_NAME = MetricRegistry.name(PushLatencyManager.class, "latency");
-    private static final int    TTL        = (int)Duration.ofDays(1).toSeconds();
 
-    private final FaultTolerantRedisCluster redisCluster;
+  private final FaultTolerantRedisCluster redisCluster;
+  private final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager;
 
-    public PushLatencyManager(final FaultTolerantRedisCluster redisCluster) {
-        this.redisCluster = redisCluster;
+  private final Clock clock;
+
+  private static final String TIMER_NAME = MetricRegistry.name(PushLatencyManager.class, "latency");
+  private static final int TTL = (int) Duration.ofDays(1).toSeconds();
+
+  private static final Logger log = LoggerFactory.getLogger(PushLatencyManager.class);
+
+  @VisibleForTesting
+  enum PushType {
+    STANDARD,
+    VOIP
+  }
+
+  @VisibleForTesting
+  static class PushRecord {
+    private final Instant timestamp;
+
+    @Nullable
+    private final PushType pushType;
+
+    @JsonCreator
+    PushRecord(@JsonProperty("timestamp") final Instant timestamp,
+        @JsonProperty("pushType") @Nullable final PushType pushType) {
+
+      this.timestamp = timestamp;
+      this.pushType = pushType;
     }
 
-    public void recordPushSent(final UUID accountUuid, final long deviceId) {
-        recordPushSent(accountUuid, deviceId, System.currentTimeMillis());
+    public Instant getTimestamp() {
+      return timestamp;
     }
 
-    @VisibleForTesting
-    void recordPushSent(final UUID accountUuid, final long deviceId, final long currentTime) {
-        redisCluster.useCluster(connection ->
-                connection.async().set(getFirstUnacknowledgedPushKey(accountUuid, deviceId), String.valueOf(currentTime), SetArgs.Builder.nx().ex(TTL)));
+    @Nullable
+    public PushType getPushType() {
+      return pushType;
     }
+  }
 
-    public void recordQueueRead(final UUID accountUuid, final long deviceId, final String userAgent) {
-        getLatencyAndClearTimestamp(accountUuid, deviceId, System.currentTimeMillis()).thenAccept(latency -> {
-            if (latency != null) {
-                Metrics.timer(TIMER_NAME, Tags.of(UserAgentTagUtil.getPlatformTag(userAgent))).record(latency, TimeUnit.MILLISECONDS);
+  public PushLatencyManager(final FaultTolerantRedisCluster redisCluster,
+      final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager) {
+
+    this(redisCluster, dynamicConfigurationManager, Clock.systemUTC());
+  }
+
+  @VisibleForTesting
+  PushLatencyManager(final FaultTolerantRedisCluster redisCluster,
+      final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager,
+      final Clock clock) {
+
+    this.redisCluster = redisCluster;
+    this.dynamicConfigurationManager = dynamicConfigurationManager;
+    this.clock = clock;
+  }
+
+  public void recordPushSent(final UUID accountUuid, final long deviceId, final boolean isVoip) {
+    try {
+      final String recordJson = SystemMapper.getMapper().writeValueAsString(
+          new PushRecord(Instant.now(clock), isVoip ? PushType.VOIP : PushType.STANDARD));
+
+      redisCluster.useCluster(connection ->
+          connection.async().set(getFirstUnacknowledgedPushKey(accountUuid, deviceId),
+              recordJson,
+              SetArgs.Builder.nx().ex(TTL)));
+    } catch (final JsonProcessingException e) {
+      // This should never happen
+      log.error("Failed to write push latency record JSON", e);
+    }
+  }
+
+  public void recordQueueRead(final UUID accountUuid, final long deviceId, final String userAgentString) {
+    takePushRecord(accountUuid, deviceId).thenAccept(pushRecord -> {
+      if (pushRecord != null) {
+        final Duration latency = Duration.between(pushRecord.getTimestamp(), Instant.now());
+
+        final List<Tag> tags = new ArrayList<>(2);
+
+        tags.add(UserAgentTagUtil.getPlatformTag(userAgentString));
+
+        if (pushRecord.getPushType() != null) {
+          tags.add(Tag.of("pushType", pushRecord.getPushType().name().toLowerCase()));
+        }
+
+        try {
+          final UserAgent userAgent = UserAgentUtil.parseUserAgentString(userAgentString);
+
+          final Set<Semver> instrumentedVersions =
+              dynamicConfigurationManager.getConfiguration().getPushLatencyConfiguration().getInstrumentedVersions()
+                  .getOrDefault(userAgent.getPlatform(), Collections.emptySet());
+
+          if (instrumentedVersions.contains(userAgent.getVersion())) {
+            tags.add(Tag.of("clientVersion", userAgent.getVersion().toString()));
+          }
+        } catch (UnrecognizedUserAgentException ignored) {
+        }
+
+        Metrics.timer(TIMER_NAME, tags).record(latency);
+      }
+    });
+  }
+
+  @VisibleForTesting
+  CompletableFuture<PushRecord> takePushRecord(final UUID accountUuid, final long deviceId) {
+    final String key = getFirstUnacknowledgedPushKey(accountUuid, deviceId);
+
+    return redisCluster.withCluster(connection -> {
+      final CompletableFuture<PushRecord> getFuture = connection.async().get(key).toCompletableFuture()
+          .thenApply(recordJson -> {
+            if (StringUtils.isNotEmpty(recordJson)) {
+              try {
+                return SystemMapper.getMapper().readValue(recordJson, PushRecord.class);
+              } catch (JsonProcessingException e) {
+                return null;
+              }
+            } else {
+              return null;
             }
-        });
-    }
+          });
 
-    @VisibleForTesting
-    CompletableFuture<Long> getLatencyAndClearTimestamp(final UUID accountUuid, final long deviceId, final long currentTimeMillis) {
-        final String key = getFirstUnacknowledgedPushKey(accountUuid, deviceId);
+      getFuture.whenComplete((record, cause) -> {
+        if (cause == null) {
+          connection.async().del(key);
+        }
+      });
 
-        return redisCluster.withCluster(connection -> {
-            final RedisAdvancedClusterAsyncCommands<String, String> commands = connection.async();
+      return getFuture;
+    });
+  }
 
-            final CompletableFuture<String> getFuture = commands.get(key).toCompletableFuture();
-            commands.del(key);
-
-            return getFuture.thenApply(timestampString -> timestampString != null ? currentTimeMillis - Long.parseLong(timestampString, 10) : null);
-        });
-    }
-
-    private static String getFirstUnacknowledgedPushKey(final UUID accountUuid, final long deviceId) {
-        return "push_latency::" + accountUuid.toString() + "::" + deviceId;
-    }
+  private static String getFirstUnacknowledgedPushKey(final UUID accountUuid, final long deviceId) {
+    return "push_latency::v2::" + accountUuid.toString() + "::" + deviceId;
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/push/APNSender.java

@@ -125,7 +125,7 @@ public class APNSender implements Managed {
   private void handleUnregisteredUser(String registrationId, UUID uuid, long deviceId) {
 //    logger.info("Got APN Unregistered: " + number + "," + deviceId);
 
-    Optional<Account> account = accountsManager.get(uuid);
+    Optional<Account> account = accountsManager.getByAccountIdentifier(uuid);
 
     if (account.isEmpty()) {
       logger.info("No account found: {}", uuid);

service/src/main/java/org/whispersystems/textsecuregcm/push/ApnFallbackManager.java

@@ -99,7 +99,7 @@ public class ApnFallbackManager implements Managed {
 
           final Optional<Account> maybeAccount = separated.map(Pair::first)
                                                           .map(UUID::fromString)
-                                                          .flatMap(accountsManager::get);
+                                                          .flatMap(accountsManager::getByAccountIdentifier);
 
           final Optional<Device> maybeDevice = separated.map(Pair::second)
                                                         .flatMap(deviceId -> maybeAccount.flatMap(account -> account.getDevice(deviceId)));

service/src/main/java/org/whispersystems/textsecuregcm/push/GCMSender.java

@@ -139,7 +139,7 @@ public class GCMSender {
   }
 
   private Optional<Account> getAccountForEvent(GcmMessage message) {
-    Optional<Account> account = message.getUuid().flatMap(accountsManager::get);
+    Optional<Account> account = message.getUuid().flatMap(accountsManager::getByAccountIdentifier);
 
     if (account.isPresent()) {
       Optional<Device> device = account.get().getDevice(message.getDeviceId());

service/src/main/java/org/whispersystems/textsecuregcm/push/MessageSender.java

@@ -124,22 +124,24 @@ public class MessageSender implements Managed {
 
     gcmSender.sendMessage(gcmMessage);
 
-    RedisOperation.unchecked(() -> pushLatencyManager.recordPushSent(account.getUuid(), device.getId()));
+    RedisOperation.unchecked(() -> pushLatencyManager.recordPushSent(account.getUuid(), device.getId(), false));
   }
 
   private void sendApnNotification(Account account, Device device) {
     ApnMessage apnMessage;
 
-    if (!Util.isEmpty(device.getVoipApnId())) {
-      apnMessage = new ApnMessage(device.getVoipApnId(), account.getUuid(), device.getId(), true, Type.NOTIFICATION, Optional.empty());
+    final boolean useVoip = !Util.isEmpty(device.getVoipApnId());
+
+    if (useVoip) {
+      apnMessage = new ApnMessage(device.getVoipApnId(), account.getUuid(), device.getId(), useVoip, Type.NOTIFICATION, Optional.empty());
       RedisOperation.unchecked(() -> apnFallbackManager.schedule(account, device));
     } else {
-      apnMessage = new ApnMessage(device.getApnId(), account.getUuid(), device.getId(), false, Type.NOTIFICATION, Optional.empty());
+      apnMessage = new ApnMessage(device.getApnId(), account.getUuid(), device.getId(), useVoip, Type.NOTIFICATION, Optional.empty());
     }
 
     apnSender.sendMessage(apnMessage);
 
-    RedisOperation.unchecked(() -> pushLatencyManager.recordPushSent(account.getUuid(), device.getId()));
+    RedisOperation.unchecked(() -> pushLatencyManager.recordPushSent(account.getUuid(), device.getId(), useVoip));
   }
 
   @Override

service/src/main/java/org/whispersystems/textsecuregcm/push/ReceiptSender.java

@@ -33,7 +33,7 @@ public class ReceiptSender {
       return;
     }
 
-    final Account destinationAccount = accountManager.get(destinationUuid)
+    final Account destinationAccount = accountManager.getByAccountIdentifier(destinationUuid)
         .orElseThrow(() -> new NoSuchUserException(destinationUuid));
 
     final Envelope.Builder message = Envelope.newBuilder()

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

@@ -23,6 +23,7 @@ import org.whispersystems.textsecuregcm.auth.AuthenticationCredentials;
 import org.whispersystems.textsecuregcm.auth.StoredRegistrationLock;
 import org.whispersystems.textsecuregcm.entities.AccountAttributes;
 import org.whispersystems.textsecuregcm.util.Util;
+import javax.annotation.Nullable;
 
 public class Account {
 
@@ -32,6 +33,11 @@ public class Account {
   @JsonIgnore
   private UUID uuid;
 
+  // Nullable only until initial migration is complete
+  @JsonProperty("pni")
+  @Nullable
+  private UUID phoneNumberIdentifier;
+
   @JsonProperty
   private String number;
 
@@ -80,9 +86,10 @@ public class Account {
   public Account() {}
 
   @VisibleForTesting
-  public Account(String number, UUID uuid, Set<Device> devices, byte[] unidentifiedAccessKey) {
+  public Account(String number, UUID uuid, final UUID phoneNumberIdentifier, Set<Device> devices, byte[] unidentifiedAccessKey) {
     this.number                = number;
     this.uuid                  = uuid;
+    this.phoneNumberIdentifier = phoneNumberIdentifier;
     this.devices               = devices;
     this.unidentifiedAccessKey = unidentifiedAccessKey;
   }
@@ -98,16 +105,11 @@ public class Account {
     this.uuid = uuid;
   }
 
-  public void setNumber(String number) {
+  // Optional only until initial migration is complete
+  public Optional<UUID> getPhoneNumberIdentifier() {
     requireNotStale();
 
-    this.number = number;
-  }
-
-  public void setCanonicallyDiscoverable(boolean canonicallyDiscoverable) {
-    requireNotStale();
-
-    this.canonicallyDiscoverable = canonicallyDiscoverable;
+    return Optional.ofNullable(phoneNumberIdentifier);
   }
 
   public String getNumber() {
@@ -116,6 +118,13 @@ public class Account {
     return number;
   }
 
+  public void setNumber(String number, UUID phoneNumberIdentifier) {
+    requireNotStale();
+
+    this.number = number;
+    this.phoneNumberIdentifier = phoneNumberIdentifier;
+  }
+
   public void addDevice(Device device) {
     requireNotStale();
 
@@ -247,6 +256,12 @@ public class Account {
     return canonicallyDiscoverable;
   }
 
+  public void setCanonicallyDiscoverable(boolean canonicallyDiscoverable) {
+    requireNotStale();
+
+    this.canonicallyDiscoverable = canonicallyDiscoverable;
+  }
+
   public Optional<String> getRelay() {
     requireNotStale();
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2020 Signal Messenger, LLC
+ * Copyright 2013-2021 Signal Messenger, LLC
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 package org.whispersystems.textsecuregcm.storage;
@@ -16,6 +16,9 @@ import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
 @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
 public class AccountDatabaseCrawlerCache {
 
+  public static final String GENERAL_PURPOSE_PREFIX = "";
+  public static final String DIRECTORY_RECONCILER_PREFIX = "directory-reconciler";
+
   private static final String ACTIVE_WORKER_KEY = "account_database_crawler_cache_active_worker";
   private static final String LAST_UUID_KEY = "account_database_crawler_cache_last_uuid";
   private static final String ACCELERATE_KEY = "account_database_crawler_cache_accelerate";
@@ -27,17 +30,21 @@ public class AccountDatabaseCrawlerCache {
   private final FaultTolerantRedisCluster cacheCluster;
   private final ClusterLuaScript unlockClusterScript;
 
-  public AccountDatabaseCrawlerCache(FaultTolerantRedisCluster cacheCluster) throws IOException {
+  private final String prefix;
+
+  public AccountDatabaseCrawlerCache(FaultTolerantRedisCluster cacheCluster, String prefix) throws IOException {
     this.cacheCluster = cacheCluster;
     this.unlockClusterScript = ClusterLuaScript.fromResource(cacheCluster, "lua/account_database_crawler/unlock.lua",
         ScriptOutputType.INTEGER);
+
+    this.prefix = prefix + "::";
   }
 
   public void setAccelerated(final boolean accelerated) {
     if (accelerated) {
-      cacheCluster.useCluster(connection -> connection.sync().set(ACCELERATE_KEY, "1"));
+      cacheCluster.useCluster(connection -> connection.sync().set(getPrefixedKey(ACCELERATE_KEY), "1"));
     } else {
-      cacheCluster.useCluster(connection -> connection.sync().del(ACCELERATE_KEY));
+      cacheCluster.useCluster(connection -> connection.sync().del(getPrefixedKey(ACCELERATE_KEY)));
     }
   }
 
@@ -47,16 +54,16 @@ public class AccountDatabaseCrawlerCache {
 
   public boolean claimActiveWork(String workerId, long ttlMs) {
     return "OK".equals(cacheCluster.withCluster(connection -> connection.sync()
-        .set(ACTIVE_WORKER_KEY, workerId, SetArgs.Builder.nx().px(ttlMs))));
+        .set(getPrefixedKey(ACTIVE_WORKER_KEY), workerId, SetArgs.Builder.nx().px(ttlMs))));
   }
 
   public void releaseActiveWork(String workerId) {
-    unlockClusterScript.execute(List.of(ACTIVE_WORKER_KEY), List.of(workerId));
+    unlockClusterScript.execute(List.of(getPrefixedKey(ACTIVE_WORKER_KEY)), List.of(workerId));
   }
 
   public Optional<UUID> getLastUuid() {
     final String lastUuidString = cacheCluster.withCluster(
-        connection -> connection.sync().get(LAST_UUID_KEY));
+        connection -> connection.sync().get(getPrefixedKey(LAST_UUID_KEY)));
 
     if (lastUuidString == null) {
       return Optional.empty();
@@ -68,15 +75,15 @@ public class AccountDatabaseCrawlerCache {
   public void setLastUuid(Optional<UUID> lastUuid) {
     if (lastUuid.isPresent()) {
       cacheCluster.useCluster(connection -> connection.sync()
-          .psetex(LAST_UUID_KEY, LAST_NUMBER_TTL_MS, lastUuid.get().toString()));
+          .psetex(getPrefixedKey(LAST_UUID_KEY), LAST_NUMBER_TTL_MS, lastUuid.get().toString()));
     } else {
-      cacheCluster.useCluster(connection -> connection.sync().del(LAST_UUID_KEY));
+      cacheCluster.useCluster(connection -> connection.sync().del(getPrefixedKey(LAST_UUID_KEY)));
     }
   }
 
   public Optional<UUID> getLastUuidDynamo() {
     final String lastUuidString = cacheCluster.withCluster(
-        connection -> connection.sync().get(LAST_UUID_DYNAMO_KEY));
+        connection -> connection.sync().get(getPrefixedKey(LAST_UUID_DYNAMO_KEY)));
 
     if (lastUuidString == null) {
       return Optional.empty();
@@ -89,9 +96,14 @@ public class AccountDatabaseCrawlerCache {
     if (lastUuid.isPresent()) {
       cacheCluster.useCluster(
           connection -> connection.sync()
-              .psetex(LAST_UUID_DYNAMO_KEY, LAST_NUMBER_TTL_MS, lastUuid.get().toString()));
+              .psetex(getPrefixedKey(LAST_UUID_DYNAMO_KEY), LAST_NUMBER_TTL_MS, lastUuid.get().toString()));
     } else {
-      cacheCluster.useCluster(connection -> connection.sync().del(LAST_UUID_DYNAMO_KEY));
+      cacheCluster.useCluster(connection -> connection.sync().del(getPrefixedKey(LAST_UUID_DYNAMO_KEY)));
     }
   }
+
+  private String getPrefixedKey(final String key) {
+    return prefix + key;
+  }
+
 }

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

@@ -13,23 +13,24 @@ import io.micrometer.core.instrument.Timer;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.stream.Collectors;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
 import org.whispersystems.textsecuregcm.util.SystemMapper;
 import org.whispersystems.textsecuregcm.util.UUIDUtil;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.services.dynamodb.model.CancellationReason;
-import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
 import software.amazon.awssdk.services.dynamodb.model.Delete;
 import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
 import software.amazon.awssdk.services.dynamodb.model.Put;
-import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
 import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure;
 import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
@@ -37,13 +38,13 @@ import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
 import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledException;
 import software.amazon.awssdk.services.dynamodb.model.TransactionConflictException;
 import software.amazon.awssdk.services.dynamodb.model.Update;
-import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
-import software.amazon.awssdk.services.dynamodb.model.UpdateItemResponse;
 
 public class Accounts extends AbstractDynamoDbStore {
 
   // uuid, primary key
   static final String KEY_ACCOUNT_UUID = "U";
+  // uuid, attribute on account table, primary key for PNI table
+  static final String ATTR_PNI_UUID = "PNI";
   // phone number
   static final String ATTR_ACCOUNT_E164 = "P";
   // account, serialized to JSON
@@ -55,7 +56,8 @@ public class Accounts extends AbstractDynamoDbStore {
 
   private final DynamoDbClient client;
 
-  private final String phoneNumbersTableName;
+  private final String phoneNumberConstraintTableName;
+  private final String phoneNumberIdentifierConstraintTableName;
   private final String accountsTableName;
 
   private final int scanPageSize;
@@ -64,19 +66,22 @@ public class Accounts extends AbstractDynamoDbStore {
   private static final Timer CHANGE_NUMBER_TIMER = Metrics.timer(name(Accounts.class, "changeNumber"));
   private static final Timer UPDATE_TIMER = Metrics.timer(name(Accounts.class, "update"));
   private static final Timer GET_BY_NUMBER_TIMER = Metrics.timer(name(Accounts.class, "getByNumber"));
+  private static final Timer GET_BY_PNI_TIMER = Metrics.timer(name(Accounts.class, "getByPni"));
   private static final Timer GET_BY_UUID_TIMER = Metrics.timer(name(Accounts.class, "getByUuid"));
   private static final Timer GET_ALL_FROM_START_TIMER = Metrics.timer(name(Accounts.class, "getAllFrom"));
   private static final Timer GET_ALL_FROM_OFFSET_TIMER = Metrics.timer(name(Accounts.class, "getAllFromOffset"));
   private static final Timer DELETE_TIMER = Metrics.timer(name(Accounts.class, "delete"));
 
+  private static final Logger log = LoggerFactory.getLogger(Accounts.class);
 
-  public Accounts(DynamoDbClient client, String accountsTableName, String phoneNumbersTableName,
-      final int scanPageSize) {
+  public Accounts(DynamoDbClient client, String accountsTableName, String phoneNumberConstraintTableName,
+      String phoneNumberIdentifierConstraintTableName, final int scanPageSize) {
 
     super(client);
 
     this.client = client;
-    this.phoneNumbersTableName = phoneNumbersTableName;
+    this.phoneNumberConstraintTableName = phoneNumberConstraintTableName;
+    this.phoneNumberIdentifierConstraintTableName = phoneNumberIdentifierConstraintTableName;
     this.accountsTableName = accountsTableName;
     this.scanPageSize = scanPageSize;
   }
@@ -85,35 +90,98 @@ public class Accounts extends AbstractDynamoDbStore {
     return CREATE_TIMER.record(() -> {
 
       try {
-        TransactWriteItem phoneNumberConstraintPut = buildPutWriteItemForPhoneNumberConstraint(account, account.getUuid());
-        TransactWriteItem accountPut = buildPutWriteItemForAccount(account, account.getUuid(), Put.builder()
-            .conditionExpression("attribute_not_exists(#number) OR #number = :number")
-            .expressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164))
-            .expressionAttributeValues(Map.of(":number", AttributeValues.fromString(account.getNumber()))));
+        TransactWriteItem phoneNumberConstraintPut = TransactWriteItem.builder()
+            .put(
+                Put.builder()
+                    .tableName(phoneNumberConstraintTableName)
+                    .item(Map.of(
+                        ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
+                        KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
+                    .conditionExpression(
+                        "attribute_not_exists(#number) OR (attribute_exists(#number) AND #uuid = :uuid)")
+                    .expressionAttributeNames(
+                        Map.of("#uuid", KEY_ACCOUNT_UUID,
+                            "#number", ATTR_ACCOUNT_E164))
+                    .expressionAttributeValues(
+                        Map.of(":uuid", AttributeValues.fromUUID(account.getUuid())))
+                    .returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
+                    .build())
+            .build();
+
+        assert account.getPhoneNumberIdentifier().isPresent();
+
+        if (account.getPhoneNumberIdentifier().isEmpty()) {
+          log.error("Account {} is missing a phone number identifier", account.getUuid());
+        }
+
+        TransactWriteItem phoneNumberIdentifierConstraintPut = TransactWriteItem.builder()
+            .put(
+                Put.builder()
+                    .tableName(phoneNumberIdentifierConstraintTableName)
+                    .item(Map.of(
+                        ATTR_PNI_UUID, AttributeValues.fromUUID(account.getPhoneNumberIdentifier().get()),
+                        KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
+                    .conditionExpression(
+                        "attribute_not_exists(#pni) OR (attribute_exists(#pni) AND #uuid = :uuid)")
+                    .expressionAttributeNames(
+                        Map.of("#uuid", KEY_ACCOUNT_UUID,
+                            "#pni", ATTR_PNI_UUID))
+                    .expressionAttributeValues(
+                        Map.of(":uuid", AttributeValues.fromUUID(account.getUuid())))
+                    .returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
+                    .build())
+            .build();
+
+        final Map<String, AttributeValue> item = new HashMap<>(Map.of(
+            KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid()),
+            ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
+            ATTR_ACCOUNT_DATA, AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
+            ATTR_VERSION, AttributeValues.fromInt(account.getVersion()),
+            ATTR_CANONICALLY_DISCOVERABLE, AttributeValues.fromBool(account.shouldBeVisibleInDirectory())));
+
+        account.getPhoneNumberIdentifier().ifPresent(pni -> item.put(ATTR_PNI_UUID, AttributeValues.fromUUID(pni)));
+
+        TransactWriteItem accountPut = TransactWriteItem.builder()
+            .put(Put.builder()
+                .conditionExpression("attribute_not_exists(#number) OR #number = :number")
+                .expressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164))
+                .expressionAttributeValues(Map.of(":number", AttributeValues.fromString(account.getNumber())))
+                .tableName(accountsTableName)
+                .item(item)
+                .build())
+            .build();
 
         final TransactWriteItemsRequest request = TransactWriteItemsRequest.builder()
-            .transactItems(phoneNumberConstraintPut, accountPut)
+            .transactItems(phoneNumberConstraintPut, phoneNumberIdentifierConstraintPut, accountPut)
             .build();
 
         try {
           client.transactWriteItems(request);
         } catch (TransactionCanceledException e) {
 
-          final CancellationReason accountCancellationReason = e.cancellationReasons().get(1);
+          final CancellationReason accountCancellationReason = e.cancellationReasons().get(2);
 
           if ("ConditionalCheckFailed".equals(accountCancellationReason.code())) {
-            throw new IllegalArgumentException("uuid present with different phone number");
+            throw new IllegalArgumentException("account identifier present with different phone number");
           }
 
           final CancellationReason phoneNumberConstraintCancellationReason = e.cancellationReasons().get(0);
+          final CancellationReason phoneNumberIdentifierConstraintCancellationReason = e.cancellationReasons().get(1);
 
-          if ("ConditionalCheckFailed".equals(phoneNumberConstraintCancellationReason.code())) {
+          if ("ConditionalCheckFailed".equals(phoneNumberConstraintCancellationReason.code()) ||
+              "ConditionalCheckFailed".equals(phoneNumberIdentifierConstraintCancellationReason.code())) {
 
-            ByteBuffer actualAccountUuid = phoneNumberConstraintCancellationReason.item().get(KEY_ACCOUNT_UUID).b().asByteBuffer();
+            // In theory, both reasons should trip in tandem and either should give us the information we need. Even so,
+            // we'll be cautious here and make sure we're choosing a condition check that really failed.
+            final CancellationReason reason = "ConditionalCheckFailed".equals(phoneNumberConstraintCancellationReason.code()) ?
+                phoneNumberConstraintCancellationReason : phoneNumberIdentifierConstraintCancellationReason;
+
+            ByteBuffer actualAccountUuid = reason.item().get(KEY_ACCOUNT_UUID).b().asByteBuffer();
             account.setUuid(UUIDUtil.fromByteBuffer(actualAccountUuid));
 
-            final int version = get(account.getUuid()).get().getVersion();
-            account.setVersion(version);
+            final Account existingAccount = getByAccountIdentifier(account.getUuid()).orElseThrow();
+            account.setNumber(existingAccount.getNumber(), existingAccount.getPhoneNumberIdentifier().orElse(account.getPhoneNumberIdentifier().orElseThrow()));
+            account.setVersion(existingAccount.getVersion());
 
             update(account);
 
@@ -125,7 +193,7 @@ public class Accounts extends AbstractDynamoDbStore {
             throw new ContestedOptimisticLockException();
           }
 
-          // this shouldn’t happen
+          // this shouldn't happen
           throw new RuntimeException("could not create account: " + extractCancellationReasonCodes(e));
         }
       } catch (JsonProcessingException e) {
@@ -136,44 +204,11 @@ public class Accounts extends AbstractDynamoDbStore {
     });
   }
 
-  private TransactWriteItem buildPutWriteItemForAccount(Account account, UUID uuid, Put.Builder putBuilder) throws JsonProcessingException {
-    return TransactWriteItem.builder()
-        .put(putBuilder
-            .tableName(accountsTableName)
-            .item(Map.of(
-                KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid),
-                ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
-                ATTR_ACCOUNT_DATA, AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
-                ATTR_VERSION, AttributeValues.fromInt(account.getVersion()),
-                ATTR_CANONICALLY_DISCOVERABLE, AttributeValues.fromBool(account.shouldBeVisibleInDirectory())))
-            .build())
-        .build();
-  }
-
-  private TransactWriteItem buildPutWriteItemForPhoneNumberConstraint(Account account, UUID uuid) {
-    return TransactWriteItem.builder()
-        .put(
-            Put.builder()
-                .tableName(phoneNumbersTableName)
-                .item(Map.of(
-                    ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
-                    KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid)))
-                .conditionExpression(
-                    "attribute_not_exists(#number) OR (attribute_exists(#number) AND #uuid = :uuid)")
-                .expressionAttributeNames(
-                    Map.of("#uuid", KEY_ACCOUNT_UUID,
-                        "#number", ATTR_ACCOUNT_E164))
-                .expressionAttributeValues(
-                    Map.of(":uuid", AttributeValues.fromUUID(uuid)))
-                .returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
-                .build())
-        .build();
-  }
-
   /**
    * Changes the phone number for the given account. The given account's number should be its current, pre-change
-   * number. If this method succeeds, the account's number will be changed to the new number. If the update fails for
-   * any reason, the account's number will be unchanged.
+   * number. If this method succeeds, the account's number will be changed to the new number and its phone number
+   * identifier will be changed to the given phone number identifier. If the update fails for any reason, the account's
+   * number and PNI will be unchanged.
    * <p/>
    * This method expects that any accounts with conflicting numbers will have been removed by the time this method is
    * called. This method may fail with an unspecified {@link RuntimeException} if another account with the same number
@@ -182,26 +217,28 @@ public class Accounts extends AbstractDynamoDbStore {
    * @param account the account for which to change the phone number
    * @param number the new phone number
    */
-  public void changeNumber(final Account account, final String number) {
+  public void changeNumber(final Account account, final String number, final UUID phoneNumberIdentifier) {
     CHANGE_NUMBER_TIMER.record(() -> {
       final String originalNumber = account.getNumber();
+      final Optional<UUID> originalPni = account.getPhoneNumberIdentifier();
+
       boolean succeeded = false;
 
-      account.setNumber(number);
+      account.setNumber(number, phoneNumberIdentifier);
 
       try {
         final List<TransactWriteItem> writeItems = new ArrayList<>();
 
         writeItems.add(TransactWriteItem.builder()
             .delete(Delete.builder()
-                .tableName(phoneNumbersTableName)
+                .tableName(phoneNumberConstraintTableName)
                 .key(Map.of(ATTR_ACCOUNT_E164, AttributeValues.fromString(originalNumber)))
                 .build())
             .build());
 
         writeItems.add(TransactWriteItem.builder()
             .put(Put.builder()
-                .tableName(phoneNumbersTableName)
+                .tableName(phoneNumberConstraintTableName)
                 .item(Map.of(
                     KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid()),
                     ATTR_ACCOUNT_E164, AttributeValues.fromString(number)))
@@ -211,20 +248,41 @@ public class Accounts extends AbstractDynamoDbStore {
                 .build())
             .build());
 
+        originalPni.ifPresent(pni -> writeItems.add(TransactWriteItem.builder()
+            .delete(Delete.builder()
+                .tableName(phoneNumberIdentifierConstraintTableName)
+                .key(Map.of(ATTR_PNI_UUID, AttributeValues.fromUUID(pni)))
+                .build())
+            .build()));
+
+        writeItems.add(TransactWriteItem.builder()
+            .put(Put.builder()
+                .tableName(phoneNumberIdentifierConstraintTableName)
+                .item(Map.of(
+                    KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid()),
+                    ATTR_PNI_UUID, AttributeValues.fromUUID(phoneNumberIdentifier)))
+                .conditionExpression("attribute_not_exists(#pni)")
+                .expressionAttributeNames(Map.of("#pni", ATTR_PNI_UUID))
+                .returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
+                .build())
+            .build());
+
         writeItems.add(
             TransactWriteItem.builder()
                 .update(Update.builder()
                     .tableName(accountsTableName)
                     .key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
-                    .updateExpression("SET #data = :data, #number = :number, #cds = :cds ADD #version :version_increment")
+                    .updateExpression("SET #data = :data, #number = :number, #pni = :pni, #cds = :cds ADD #version :version_increment")
                     .conditionExpression("attribute_exists(#number) AND #version = :version")
                     .expressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164,
                         "#data", ATTR_ACCOUNT_DATA,
                         "#cds", ATTR_CANONICALLY_DISCOVERABLE,
+                        "#pni", ATTR_PNI_UUID,
                         "#version", ATTR_VERSION))
                     .expressionAttributeValues(Map.of(
                         ":data", AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
                         ":number", AttributeValues.fromString(number),
+                        ":pni", AttributeValues.fromUUID(phoneNumberIdentifier),
                         ":cds", AttributeValues.fromBool(account.shouldBeVisibleInDirectory()),
                         ":version", AttributeValues.fromInt(account.getVersion()),
                         ":version_increment", AttributeValues.fromInt(1)))
@@ -243,7 +301,7 @@ public class Accounts extends AbstractDynamoDbStore {
         throw new IllegalArgumentException(e);
       } finally {
         if (!succeeded) {
-          account.setNumber(originalNumber);
+          account.setNumber(originalNumber, originalPni.orElse(null));
         }
       }
     });
@@ -251,57 +309,121 @@ public class Accounts extends AbstractDynamoDbStore {
 
   public void update(Account account) throws ContestedOptimisticLockException {
     UPDATE_TIMER.record(() -> {
-      UpdateItemRequest updateItemRequest;
-      try {
-        updateItemRequest = UpdateItemRequest.builder()
-            .tableName(accountsTableName)
-            .key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
-            .updateExpression("SET #data = :data, #cds = :cds ADD #version :version_increment")
-            .conditionExpression("attribute_exists(#number) AND #version = :version")
-            .expressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164,
-                "#data", ATTR_ACCOUNT_DATA,
-                "#cds", ATTR_CANONICALLY_DISCOVERABLE,
-                "#version", ATTR_VERSION))
-            .expressionAttributeValues(Map.of(
-                ":data", AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
-                ":cds", AttributeValues.fromBool(account.shouldBeVisibleInDirectory()),
-                ":version", AttributeValues.fromInt(account.getVersion()),
-                ":version_increment", AttributeValues.fromInt(1)))
-            .returnValues(ReturnValue.UPDATED_NEW)
-            .build();
+      final List<TransactWriteItem> transactWriteItems = new ArrayList<>(2);
 
+      try {
+        final TransactWriteItem updateAccountWriteItem;
+
+        if (account.getPhoneNumberIdentifier().isPresent()) {
+          updateAccountWriteItem = TransactWriteItem.builder()
+              .update(Update.builder()
+                  .tableName(accountsTableName)
+                  .key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
+                  .updateExpression("SET #data = :data, #cds = :cds, #pni = :pni ADD #version :version_increment")
+                  .conditionExpression("attribute_exists(#number) AND #version = :version")
+                  .expressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164,
+                      "#data", ATTR_ACCOUNT_DATA,
+                      "#cds", ATTR_CANONICALLY_DISCOVERABLE,
+                      "#version", ATTR_VERSION,
+                      "#pni", ATTR_PNI_UUID))
+                  .expressionAttributeValues(Map.of(
+                      ":data", AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
+                      ":cds", AttributeValues.fromBool(account.shouldBeVisibleInDirectory()),
+                      ":version", AttributeValues.fromInt(account.getVersion()),
+                      ":version_increment", AttributeValues.fromInt(1),
+                      ":pni", AttributeValues.fromUUID(account.getPhoneNumberIdentifier().get())))
+                  .build())
+              .build();
+        } else {
+          updateAccountWriteItem = TransactWriteItem.builder()
+              .update(Update.builder()
+                  .tableName(accountsTableName)
+                  .key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
+                  .updateExpression("SET #data = :data, #cds = :cds ADD #version :version_increment")
+                  .conditionExpression("attribute_exists(#number) AND #version = :version")
+                  .expressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164,
+                      "#data", ATTR_ACCOUNT_DATA,
+                      "#cds", ATTR_CANONICALLY_DISCOVERABLE,
+                      "#version", ATTR_VERSION))
+                  .expressionAttributeValues(Map.of(
+                      ":data", AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
+                      ":cds", AttributeValues.fromBool(account.shouldBeVisibleInDirectory()),
+                      ":version", AttributeValues.fromInt(account.getVersion()),
+                      ":version_increment", AttributeValues.fromInt(1)))
+                  .build())
+              .build();
+        }
+
+        transactWriteItems.add(updateAccountWriteItem);
+
+        // TODO Remove after initial migration to phone number identifiers
+        account.getPhoneNumberIdentifier().ifPresent(phoneNumberIdentifier -> transactWriteItems.add(
+            TransactWriteItem.builder()
+                .put(Put.builder()
+                    .tableName(phoneNumberIdentifierConstraintTableName)
+                    .item(Map.of(
+                        ATTR_PNI_UUID, AttributeValues.fromUUID(account.getPhoneNumberIdentifier().get()),
+                        KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
+                    .conditionExpression("attribute_not_exists(#pni) OR (attribute_exists(#pni) AND #uuid = :uuid)")
+                    .expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#pni", ATTR_PNI_UUID))
+                    .expressionAttributeValues(
+                        Map.of(":uuid", AttributeValues.fromUUID(account.getUuid())))
+                    .returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
+                    .build())
+                .build()));
       } catch (JsonProcessingException e) {
         throw new IllegalArgumentException(e);
       }
 
       try {
-        UpdateItemResponse response = client.updateItem(updateItemRequest);
+        client.transactWriteItems(TransactWriteItemsRequest.builder()
+            .transactItems(transactWriteItems)
+            .build());
 
-        account.setVersion(AttributeValues.getInt(response.attributes(), "V", account.getVersion() + 1));
+        account.setVersion(account.getVersion() + 1);
       } catch (final TransactionConflictException e) {
 
         throw new ContestedOptimisticLockException();
 
-      } catch (final ConditionalCheckFailedException e) {
+      } catch (final TransactionCanceledException e) {
 
-        // the exception doesn’t give details about which condition failed,
-        // but we can infer it was an optimistic locking failure if the UUID is known
-        throw get(account.getUuid()).isPresent() ? new ContestedOptimisticLockException() : e;
+        if (e.cancellationReasons().size() > 1 && "ConditionalCheckFailed".equals(e.cancellationReasons().get(1).code())) {
+          log.error("Conflicting phone number mapping exists for account {}, PNI {}", account.getUuid(), account.getPhoneNumberIdentifier());
+          throw e;
+        }
+
+        // We can infer an optimistic locking failure if the UUID is known
+        throw getByAccountIdentifier(account.getUuid()).isPresent() ? new ContestedOptimisticLockException() : e;
       }
     });
   }
 
-  public Optional<Account> get(String number) {
+  public Optional<Account> getByE164(String number) {
     return GET_BY_NUMBER_TIMER.record(() -> {
 
       final GetItemResponse response = client.getItem(GetItemRequest.builder()
-          .tableName(phoneNumbersTableName)
+          .tableName(phoneNumberConstraintTableName)
           .key(Map.of(ATTR_ACCOUNT_E164, AttributeValues.fromString(number)))
           .build());
 
       return Optional.ofNullable(response.item())
           .map(item -> item.get(KEY_ACCOUNT_UUID))
-          .map(uuid -> accountByUuid(uuid))
+          .map(this::accountByUuid)
+          .map(Accounts::fromItem);
+    });
+  }
+
+  public Optional<Account> getByPhoneNumberIdentifier(final UUID phoneNumberIdentifier) {
+    return GET_BY_PNI_TIMER.record(() -> {
+
+      final GetItemResponse response = client.getItem(GetItemRequest.builder()
+          .tableName(phoneNumberIdentifierConstraintTableName)
+          .key(Map.of(ATTR_PNI_UUID, AttributeValues.fromUUID(phoneNumberIdentifier)))
+          .build());
+
+      return Optional.ofNullable(response.item())
+          .map(item -> item.get(KEY_ACCOUNT_UUID))
+          .map(this::accountByUuid)
           .map(Accounts::fromItem);
     });
   }
@@ -315,7 +437,7 @@ public class Accounts extends AbstractDynamoDbStore {
     return r.item().isEmpty() ? null : r.item();
   }
 
-  public Optional<Account> get(UUID uuid) {
+  public Optional<Account> getByAccountIdentifier(UUID uuid) {
     return GET_BY_UUID_TIMER.record(() ->
         Optional.ofNullable(accountByUuid(AttributeValues.fromUUID(uuid)))
             .map(Accounts::fromItem));
@@ -324,13 +446,11 @@ public class Accounts extends AbstractDynamoDbStore {
   public void delete(UUID uuid) {
     DELETE_TIMER.record(() -> {
 
-      Optional<Account> maybeAccount = get(uuid);
-
-      maybeAccount.ifPresent(account -> {
+      getByAccountIdentifier(uuid).ifPresent(account -> {
 
         TransactWriteItem phoneNumberDelete = TransactWriteItem.builder()
             .delete(Delete.builder()
-                .tableName(phoneNumbersTableName)
+                .tableName(phoneNumberConstraintTableName)
                 .key(Map.of(ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber())))
                 .build())
             .build();
@@ -342,8 +462,17 @@ public class Accounts extends AbstractDynamoDbStore {
                 .build())
             .build();
 
+        final List<TransactWriteItem> transactWriteItems = new ArrayList<>(List.of(phoneNumberDelete, accountDelete));
+
+        account.getPhoneNumberIdentifier().ifPresent(pni -> transactWriteItems.add(TransactWriteItem.builder()
+            .delete(Delete.builder()
+                .tableName(phoneNumberIdentifierConstraintTableName)
+                .key(Map.of(ATTR_PNI_UUID, AttributeValues.fromUUID(pni)))
+                .build())
+            .build()));
+
         TransactWriteItemsRequest request = TransactWriteItemsRequest.builder()
-            .transactItems(phoneNumberDelete, accountDelete).build();
+            .transactItems(transactWriteItems).build();
 
         client.transactWriteItems(request);
       });
@@ -393,7 +522,7 @@ public class Accounts extends AbstractDynamoDbStore {
     }
     try {
       Account account = SystemMapper.getMapper().readValue(item.get(ATTR_ACCOUNT_DATA).b().asByteArray(), Account.class);
-      account.setNumber(item.get(ATTR_ACCOUNT_E164).s());
+      account.setNumber(item.get(ATTR_ACCOUNT_E164).s(), AttributeValues.getUUID(item, ATTR_PNI_UUID, null));
       account.setUuid(UUIDUtil.fromByteBuffer(item.get(KEY_ACCOUNT_UUID).b().asByteBuffer()));
       account.setVersion(Integer.parseInt(item.get(ATTR_VERSION).n()));
       account.setCanonicallyDiscoverable(Optional.ofNullable(item.get(ATTR_CANONICALLY_DISCOVERABLE)).map(av -> av.bool()).orElse(false));

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

@@ -17,6 +17,9 @@ import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
 import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Tags;
 import java.io.IOException;
+import java.time.Clock;
+import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
@@ -49,6 +52,7 @@ public class AccountsManager {
 
   private static final Timer redisSetTimer       = metricRegistry.timer(name(AccountsManager.class, "redisSet"      ));
   private static final Timer redisNumberGetTimer = metricRegistry.timer(name(AccountsManager.class, "redisNumberGet"));
+  private static final Timer redisPniGetTimer    = metricRegistry.timer(name(AccountsManager.class, "redisPniGet"));
   private static final Timer redisUuidGetTimer   = metricRegistry.timer(name(AccountsManager.class, "redisUuidGet"  ));
   private static final Timer redisDeleteTimer    = metricRegistry.timer(name(AccountsManager.class, "redisDelete"   ));
 
@@ -60,18 +64,20 @@ public class AccountsManager {
   private final Logger logger = LoggerFactory.getLogger(AccountsManager.class);
 
   private final Accounts accounts;
+  private final PhoneNumberIdentifiers phoneNumberIdentifiers;
   private final FaultTolerantRedisCluster cacheCluster;
   private final DeletedAccountsManager deletedAccountsManager;
-  private final DirectoryQueue            directoryQueue;
-  private final KeysDynamoDb              keysDynamoDb;
+  private final DirectoryQueue directoryQueue;
+  private final KeysDynamoDb keysDynamoDb;
   private final MessagesManager messagesManager;
   private final UsernamesManager usernamesManager;
-  private final ProfilesManager           profilesManager;
+  private final ProfilesManager profilesManager;
   private final StoredVerificationCodeManager pendingAccounts;
-  private final SecureStorageClient       secureStorageClient;
-  private final SecureBackupClient        secureBackupClient;
+  private final SecureStorageClient secureStorageClient;
+  private final SecureBackupClient secureBackupClient;
   private final ClientPresenceManager clientPresenceManager;
-  private final ObjectMapper              mapper;
+  private final ObjectMapper mapper;
+  private final Clock clock;
 
   public enum DeletionReason {
     ADMIN_DELETED("admin"),
@@ -85,17 +91,22 @@ public class AccountsManager {
     }
   }
 
-  public AccountsManager(Accounts accounts, FaultTolerantRedisCluster cacheCluster,
+  public AccountsManager(final Accounts accounts,
+      final PhoneNumberIdentifiers phoneNumberIdentifiers,
+      final FaultTolerantRedisCluster cacheCluster,
       final DeletedAccountsManager deletedAccountsManager,
       final DirectoryQueue directoryQueue,
-      final KeysDynamoDb keysDynamoDb, final MessagesManager messagesManager,
+      final KeysDynamoDb keysDynamoDb,
+      final MessagesManager messagesManager,
       final UsernamesManager usernamesManager,
       final ProfilesManager profilesManager,
       final StoredVerificationCodeManager pendingAccounts,
       final SecureStorageClient secureStorageClient,
       final SecureBackupClient secureBackupClient,
-      final ClientPresenceManager clientPresenceManager) {
+      final ClientPresenceManager clientPresenceManager,
+      final Clock clock) {
     this.accounts = accounts;
+    this.phoneNumberIdentifiers = phoneNumberIdentifiers;
     this.cacheCluster = cacheCluster;
     this.deletedAccountsManager = deletedAccountsManager;
     this.directoryQueue = directoryQueue;
@@ -108,12 +119,14 @@ public class AccountsManager {
     this.secureBackupClient  = secureBackupClient;
     this.clientPresenceManager = clientPresenceManager;
     this.mapper              = SystemMapper.getMapper();
+    this.clock = Objects.requireNonNull(clock);
   }
 
   public Account create(final String number,
       final String password,
       final String signalAgent,
-      final AccountAttributes accountAttributes) throws InterruptedException {
+      final AccountAttributes accountAttributes,
+      final List<AccountBadge> accountBadges) throws InterruptedException {
 
     try (Timer.Context ignored = createTimer.time()) {
       final Account account = new Account();
@@ -130,17 +143,18 @@ public class AccountsManager {
         device.setLastSeen(Util.todayInMillis());
         device.setUserAgent(signalAgent);
 
-        account.setNumber(number);
+        account.setNumber(number, phoneNumberIdentifiers.getPhoneNumberIdentifier(number));
         account.setUuid(maybeRecentlyDeletedUuid.orElseGet(UUID::randomUUID));
         account.addDevice(device);
         account.setRegistrationLockFromAttributes(accountAttributes);
         account.setUnidentifiedAccessKey(accountAttributes.getUnidentifiedAccessKey());
         account.setUnrestrictedUnidentifiedAccess(accountAttributes.isUnrestrictedUnidentifiedAccess());
         account.setDiscoverableByPhoneNumber(accountAttributes.isDiscoverableByPhoneNumber());
+        account.setBadges(clock, accountBadges);
 
         final UUID originalUuid = account.getUuid();
 
-        boolean freshUser = dynamoCreate(account);
+        boolean freshUser = accounts.create(account);
 
         // create() sometimes updates the UUID, if there was a number conflict.
         // for metrics, we want secondary to run with the same original UUID
@@ -202,7 +216,7 @@ public class AccountsManager {
     deletedAccountsManager.lockAndPut(account.getNumber(), number, () -> {
       redisDelete(account);
 
-      final Optional<Account> maybeExistingAccount = get(number);
+      final Optional<Account> maybeExistingAccount = getByE164(number);
       final Optional<UUID> displacedUuid;
 
       if (maybeExistingAccount.isPresent()) {
@@ -213,12 +227,13 @@ public class AccountsManager {
       }
 
       final UUID uuid = account.getUuid();
+      final UUID phoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(number);
 
       final Account numberChangedAccount = updateWithRetries(
           account,
           a -> true,
-          a -> dynamoChangeNumber(a, number),
-          () -> dynamoGet(uuid).orElseThrow());
+          a -> accounts.changeNumber(a, number, phoneNumberIdentifier),
+          () -> accounts.getByAccountIdentifier(uuid).orElseThrow());
 
       updatedAccount.set(numberChangedAccount);
       directoryQueue.changePhoneNumber(numberChangedAccount, originalNumber, number);
@@ -278,7 +293,10 @@ public class AccountsManager {
       final UUID uuid = account.getUuid();
       final String originalNumber = account.getNumber();
 
-      updatedAccount = updateWithRetries(account, updater, this::dynamoUpdate, () -> dynamoGet(uuid).get());
+      updatedAccount = updateWithRetries(account,
+          updater,
+          accounts::update,
+          () -> accounts.getByAccountIdentifier(uuid).orElseThrow());
 
       assert updatedAccount.getNumber().equals(originalNumber);
 
@@ -347,12 +365,12 @@ public class AccountsManager {
     });
   }
 
-  public Optional<Account> get(String number) {
+  public Optional<Account> getByE164(String number) {
     try (Timer.Context ignored = getByNumberTimer.time()) {
-      Optional<Account> account = redisGet(number);
+      Optional<Account> account = redisGetByE164(number);
 
       if (account.isEmpty()) {
-        account = dynamoGet(number);
+        account = accounts.getByE164(number);
         account.ifPresent(this::redisSet);
       }
 
@@ -360,12 +378,25 @@ public class AccountsManager {
     }
   }
 
-  public Optional<Account> get(UUID uuid) {
-    try (Timer.Context ignored = getByUuidTimer.time()) {
-      Optional<Account> account = redisGet(uuid);
+  public Optional<Account> getByPhoneNumberIdentifier(UUID pni) {
+    try (Timer.Context ignored = getByNumberTimer.time()) {
+      Optional<Account> account = redisGetByPhoneNumberIdentifier(pni);
 
       if (account.isEmpty()) {
-        account = dynamoGet(uuid);
+        account = accounts.getByPhoneNumberIdentifier(pni);
+        account.ifPresent(this::redisSet);
+      }
+
+      return account;
+    }
+  }
+
+  public Optional<Account> getByAccountIdentifier(UUID uuid) {
+    try (Timer.Context ignored = getByUuidTimer.time()) {
+      Optional<Account> account = redisGetByAccountIdentifier(uuid);
+
+      if (account.isEmpty()) {
+        account = accounts.getByAccountIdentifier(uuid);
         account.ifPresent(this::redisSet);
       }
 
@@ -409,19 +440,24 @@ public class AccountsManager {
     keysDynamoDb.delete(account.getUuid());
     messagesManager.clear(account.getUuid());
 
+    account.getPhoneNumberIdentifier().ifPresent(pni -> {
+      keysDynamoDb.delete(pni);
+      messagesManager.clear(pni);
+    });
+
     deleteStorageServiceDataFuture.join();
     deleteBackupServiceDataFuture.join();
 
     redisDelete(account);
-    dynamoDelete(account);
+    accounts.delete(account.getUuid());
 
     RedisOperation.unchecked(() ->
         account.getDevices().forEach(device ->
             clientPresenceManager.displacePresence(account.getUuid(), device.getId())));
   }
 
-  private String getAccountMapKey(String number) {
-    return "AccountMap::" + number;
+  private String getAccountMapKey(String key) {
+    return "AccountMap::" + key;
   }
 
   private String getAccountEntityKey(UUID uuid) {
@@ -435,6 +471,9 @@ public class AccountsManager {
       cacheCluster.useCluster(connection -> {
         final RedisAdvancedClusterCommands<String, String> commands = connection.sync();
 
+        account.getPhoneNumberIdentifier().ifPresent(pni ->
+            commands.set(getAccountMapKey(pni.toString()), account.getUuid().toString()));
+
         commands.set(getAccountMapKey(account.getNumber()), account.getUuid().toString());
         commands.set(getAccountEntityKey(account.getUuid()), accountJson);
       });
@@ -443,11 +482,19 @@ public class AccountsManager {
     }
   }
 
-  private Optional<Account> redisGet(String number) {
-    try (Timer.Context ignored = redisNumberGetTimer.time()) {
-      final String uuid = cacheCluster.withCluster(connection -> connection.sync().get(getAccountMapKey(number)));
+  private Optional<Account> redisGetByPhoneNumberIdentifier(UUID uuid) {
+    return redisGetBySecondaryKey(uuid.toString(), redisPniGetTimer);
+  }
 
-      if (uuid != null) return redisGet(UUID.fromString(uuid));
+  private Optional<Account> redisGetByE164(String e164) {
+    return redisGetBySecondaryKey(e164, redisNumberGetTimer);
+  }
+
+  private Optional<Account> redisGetBySecondaryKey(String secondaryKey, Timer timer) {
+    try (Timer.Context ignored = timer.time()) {
+      final String uuid = cacheCluster.withCluster(connection -> connection.sync().get(getAccountMapKey(secondaryKey)));
+
+      if (uuid != null) return redisGetByAccountIdentifier(UUID.fromString(uuid));
       else              return Optional.empty();
     } catch (IllegalArgumentException e) {
       logger.warn("Deserialization error", e);
@@ -458,7 +505,7 @@ public class AccountsManager {
     }
   }
 
-  private Optional<Account> redisGet(UUID uuid) {
+  private Optional<Account> redisGetByAccountIdentifier(UUID uuid) {
     try (Timer.Context ignored = redisUuidGetTimer.time()) {
       final String json = cacheCluster.withCluster(connection -> connection.sync().get(getAccountEntityKey(uuid)));
 
@@ -481,32 +528,11 @@ public class AccountsManager {
 
   private void redisDelete(final Account account) {
     try (final Timer.Context ignored = redisDeleteTimer.time()) {
-      cacheCluster.useCluster(connection -> connection.sync()
-          .del(getAccountMapKey(account.getNumber()), getAccountEntityKey(account.getUuid())));
+      cacheCluster.useCluster(connection -> {
+        connection.sync().del(getAccountMapKey(account.getNumber()), getAccountEntityKey(account.getUuid()));
+
+        account.getPhoneNumberIdentifier().ifPresent(pni -> connection.sync().del(getAccountMapKey(pni.toString())));
+      });
     }
   }
-
-  private Optional<Account> dynamoGet(String number) {
-    return accounts.get(number);
-  }
-
-  private Optional<Account> dynamoGet(UUID uuid) {
-    return accounts.get(uuid);
-  }
-
-  private boolean dynamoCreate(Account account) {
-    return accounts.create(account);
-  }
-
-  private void dynamoUpdate(Account account) {
-    accounts.update(account);
-  }
-
-  private void dynamoDelete(final Account account) {
-    accounts.delete(account.getUuid());
-  }
-
-  private void dynamoChangeNumber(final Account account, final String number) {
-    accounts.changeNumber(account, number);
-  }
 }

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

@@ -0,0 +1,72 @@
+/*
+ * Copyright 2013-2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import io.micrometer.core.instrument.Counter;
+import io.micrometer.core.instrument.Metrics;
+import java.util.List;
+import java.util.Optional;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.stream.Collectors;
+
+import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
+
+// TODO Remove this crawler when PNIs have been assigned to all existing accounts
+public class AssignPhoneNumberIdentifierCrawlerListener extends AccountDatabaseCrawlerListener {
+
+  private final AccountsManager accountsManager;
+  private final PhoneNumberIdentifiers phoneNumberIdentifiers;
+
+  private final ExecutorService updateExecutor = Executors.newFixedThreadPool(16);
+
+  private static final Counter ASSIGNED_PNI_COUNTER =
+      Metrics.counter(name(AssignPhoneNumberIdentifierCrawlerListener.class, "assignPni"));
+
+  public AssignPhoneNumberIdentifierCrawlerListener(final AccountsManager accountsManager,
+      final PhoneNumberIdentifiers phoneNumberIdentifiers) {
+
+    this.accountsManager = accountsManager;
+    this.phoneNumberIdentifiers = phoneNumberIdentifiers;
+  }
+
+  @Override
+  public void onCrawlStart() {
+  }
+
+  @Override
+  public void onCrawlEnd(final Optional<UUID> fromUuid) {
+  }
+
+  @Override
+  protected void onCrawlChunk(final Optional<UUID> fromUuid, final List<Account> chunkAccounts) {
+    // There are exactly two ways an account can get a phone number identifier (PNI):
+    //
+    // 1. It's assigned at construction time for accounts created after the introduction of PNIs
+    // 2. It's assigned by this crawler
+    //
+    // That means that we don't need to worry about accidentally overwriting a PNI assigned by another source; if an
+    // account doesn't have a PNI when it winds up in a crawled chunk, there's no danger that it will have one after a
+    // refresh, and so we can blindly assign a random PNI.
+    CompletableFuture.allOf(chunkAccounts.stream()
+        .filter(account -> account.getPhoneNumberIdentifier().isEmpty())
+        .map(Account::getUuid)
+        .map(accountsManager::getByAccountIdentifier)
+        .filter(Optional::isPresent)
+        .map(Optional::get)
+        .map(accountWithoutPni -> CompletableFuture.runAsync(() -> {
+          final String number = accountWithoutPni.getNumber();
+          final UUID phoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(number);
+
+          accountsManager.update(accountWithoutPni, a -> a.setNumber(number, phoneNumberIdentifier));
+
+          ASSIGNED_PNI_COUNTER.increment();
+        }, updateExecutor))
+        .toArray(CompletableFuture[]::new)).join();
+  }
+}

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

@@ -43,7 +43,7 @@ public class ContactDiscoveryWriter extends AccountDatabaseCrawlerListener {
         // It’s less than ideal, but crawler listeners currently must not call update()
         // with the accounts from the chunk, because updates cause the account instance to become stale. Instead, they
         // must get a new copy, which they are free to update.
-        accounts.get(account.getUuid()).ifPresent(a -> accounts.update(a, NOOP_UPDATER));
+        accounts.getByAccountIdentifier(account.getUuid()).ifPresent(a -> accounts.update(a, NOOP_UPDATER));
       }
     }
   }

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

@@ -146,7 +146,7 @@ public class MessagePersister implements Managed {
 
     @VisibleForTesting
     void persistQueue(final UUID accountUuid, final long deviceId) {
-        final Optional<Account> maybeAccount = accountsManager.get(accountUuid);
+        final Optional<Account> maybeAccount = accountsManager.getByAccountIdentifier(accountUuid);
 
         if (maybeAccount.isEmpty()) {
             logger.error("No account record found for account {}", accountUuid);

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

@@ -62,7 +62,7 @@ public class NonNormalizedAccountCrawlerListener extends AccountDatabaseCrawlerL
           workingNonNormalizedNumbers++;
 
           try {
-            final Optional<Account> maybeConflictingAccount = accountsManager.get(getNormalizedNumber(account));
+            final Optional<Account> maybeConflictingAccount = accountsManager.getByE164(getNormalizedNumber(account));
 
             if (maybeConflictingAccount.isPresent()) {
               workingConflictingNumbers++;

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

@@ -0,0 +1,85 @@
+/*
+ * Copyright 2013-2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import com.google.common.annotations.VisibleForTesting;
+import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.Timer;
+import org.whispersystems.textsecuregcm.util.AttributeValues;
+import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
+import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
+import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
+import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
+import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
+import software.amazon.awssdk.services.dynamodb.model.UpdateItemResponse;
+import java.util.Map;
+import java.util.UUID;
+
+import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
+
+/**
+ * Manages a global, persistent mapping of phone numbers to phone number identifiers regardless of whether those
+ * numbers/identifiers are actually associated with an account.
+ */
+public class PhoneNumberIdentifiers {
+
+  private final DynamoDbClient dynamoDbClient;
+  private final String tableName;
+
+  @VisibleForTesting
+  static final String KEY_E164 = "P";
+  private static final String ATTR_PHONE_NUMBER_IDENTIFIER = "PNI";
+
+  private static final Timer GET_PNI_TIMER = Metrics.timer(name(PhoneNumberIdentifiers.class, "get"));
+  private static final Timer SET_PNI_TIMER = Metrics.timer(name(PhoneNumberIdentifiers.class, "set"));
+
+  public PhoneNumberIdentifiers(final DynamoDbClient dynamoDbClient, final String tableName) {
+    this.dynamoDbClient = dynamoDbClient;
+    this.tableName = tableName;
+  }
+
+  /**
+   * Returns the phone number identifier (PNI) associated with the given phone number.
+   *
+   * @param phoneNumber the phone number for which to retrieve a phone number identifier
+   * @return the phone number identifier associated with the given phone number
+   */
+  public UUID getPhoneNumberIdentifier(final String phoneNumber) {
+    final GetItemResponse response = GET_PNI_TIMER.record(() -> dynamoDbClient.getItem(GetItemRequest.builder()
+        .tableName(tableName)
+        .key(Map.of(KEY_E164, AttributeValues.fromString(phoneNumber)))
+        .projectionExpression(ATTR_PHONE_NUMBER_IDENTIFIER)
+        .build()));
+
+    final UUID phoneNumberIdentifier;
+
+    if (response.hasItem()) {
+      phoneNumberIdentifier = AttributeValues.getUUID(response.item(), ATTR_PHONE_NUMBER_IDENTIFIER, null);
+    } else {
+      phoneNumberIdentifier = generatePhoneNumberIdentifierIfNotExists(phoneNumber);
+    }
+
+    if (phoneNumberIdentifier == null) {
+      throw new RuntimeException("Could not retrieve phone number identifier from stored item");
+    }
+
+    return phoneNumberIdentifier;
+  }
+
+  @VisibleForTesting
+  UUID generatePhoneNumberIdentifierIfNotExists(final String phoneNumber) {
+    final UpdateItemResponse response = SET_PNI_TIMER.record(() -> dynamoDbClient.updateItem(UpdateItemRequest.builder()
+        .tableName(tableName)
+        .key(Map.of(KEY_E164, AttributeValues.fromString(phoneNumber)))
+        .updateExpression("SET #pni = if_not_exists(#pni, :pni)")
+        .expressionAttributeNames(Map.of("#pni", ATTR_PHONE_NUMBER_IDENTIFIER))
+        .expressionAttributeValues(Map.of(":pni", AttributeValues.fromUUID(UUID.randomUUID())))
+        .returnValues(ReturnValue.ALL_NEW)
+        .build()));
+
+    return AttributeValues.getUUID(response.attributes(), ATTR_PHONE_NUMBER_IDENTIFIER, null);
+  }
+}

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

@@ -58,7 +58,7 @@ public class PushFeedbackProcessor extends AccountDatabaseCrawlerListener {
 
       if (update) {
         // fetch a new version, since the chunk is shared and implicitly read-only
-        accountsManager.get(account.getUuid()).ifPresent(accountToUpdate -> {
+        accountsManager.getByAccountIdentifier(account.getUuid()).ifPresent(accountToUpdate -> {
           accountsManager.update(accountToUpdate, a -> {
             for (Device device : a.getDevices()) {
               if (deviceNeedsUpdate(device)) {

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

@@ -24,7 +24,7 @@ public class RefreshingAccountAndDeviceSupplier implements Supplier<Pair<Account
   @Override
   public Pair<Account, Device> get() {
     if (account.isStale()) {
-      account = accountsManager.get(account.getUuid())
+      account = accountsManager.getByAccountIdentifier(account.getUuid())
           .orElseThrow(() -> new RuntimeException("Could not find account"));
       device = account.getDevice(device.getId())
           .orElseThrow(() -> new RefreshingAccountAndDeviceNotFoundException("Could not find device"));

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

@@ -119,6 +119,6 @@ public class AttributeValues {
   }
 
   public static UUID getUUID(Map<String, AttributeValue> item, String key, UUID defaultValue) {
-    return AttributeValues.get(item, key).map(AttributeValues::toUUID).orElse(defaultValue);
+    return AttributeValues.get(item, key).filter(av -> av.b() != null).map(AttributeValues::toUUID).orElse(defaultValue);
   }
 }

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

@@ -10,6 +10,7 @@ import com.google.i18n.phonenumbers.PhoneNumberUtil.PhoneNumberFormat;
 import com.google.i18n.phonenumbers.Phonenumber.PhoneNumber;
 import java.io.UnsupportedEncodingException;
 import java.net.URLEncoder;
+import java.nio.ByteBuffer;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.security.SecureRandom;
@@ -107,16 +108,6 @@ public class Util {
     return param == null || param.length() == 0;
   }
 
-  public static byte[] combine(byte[] one, byte[] two, byte[] three, byte[] four) {
-    byte[] combined = new byte[one.length + two.length + three.length + four.length];
-    System.arraycopy(one, 0, combined, 0, one.length);
-    System.arraycopy(two, 0, combined, one.length, two.length);
-    System.arraycopy(three, 0, combined, one.length + two.length, three.length);
-    System.arraycopy(four, 0, combined, one.length + two.length + three.length, four.length);
-
-    return combined;
-  }
-
   public static byte[] truncate(byte[] element, int length) {
     byte[] result = new byte[length];
     System.arraycopy(element, 0, result, 0, result.length);
@@ -124,7 +115,6 @@ public class Util {
     return result;
   }
 
-
   public static byte[][] split(byte[] input, int firstLength, int secondLength) {
     byte[][] parts = new byte[2][];
 
@@ -161,11 +151,19 @@ public class Util {
     return data;
   }
 
+  public static byte[] longToByteArray(long value) {
+    final ByteBuffer longBuffer = ByteBuffer.allocate(Long.BYTES);
+
+    longBuffer.putLong(value);
+
+    return longBuffer.array();
+  }
+
   public static int toIntExact(long value) {
-    if ((int)value != value) {
+    if ((int) value != value) {
       throw new ArithmeticException("integer overflow");
     }
-    return (int)value;
+    return (int) value;
   }
 
   public static int currentDaysSinceEpoch() {

service/src/main/java/org/whispersystems/textsecuregcm/websocket/DeadLetterHandler.java

@@ -47,7 +47,7 @@ public class DeadLetterHandler implements DispatchChannel {
       switch (pubSubMessage.getType().getNumber()) {
         case PubSubMessage.Type.DELIVER_VALUE:
           Envelope          message      = Envelope.parseFrom(pubSubMessage.getContent());
-          Optional<Account> maybeAccount = accountsManager.get(address.getNumber());
+          Optional<Account> maybeAccount = accountsManager.getByE164(address.getNumber());
 
           if (maybeAccount.isPresent()) {
             messagesManager.insert(maybeAccount.get().getUuid(), address.getDeviceId(), message);

service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java

@@ -18,6 +18,7 @@ import io.dropwizard.jdbi3.JdbiFactory;
 import io.dropwizard.setup.Environment;
 import io.lettuce.core.resource.ClientResources;
 import io.micrometer.core.instrument.Metrics;
+import java.time.Clock;
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -47,6 +48,7 @@ import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
 import org.whispersystems.textsecuregcm.storage.MessagesCache;
 import org.whispersystems.textsecuregcm.storage.MessagesDynamoDb;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
+import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
 import org.whispersystems.textsecuregcm.storage.Profiles;
 import org.whispersystems.textsecuregcm.storage.ProfilesManager;
 import org.whispersystems.textsecuregcm.storage.ReportMessageDynamoDb;
@@ -88,6 +90,7 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
       throws Exception
   {
     try {
+      Clock clock = Clock.systemUTC();
       String[] users = namespace.getString("user").split(",");
 
       environment.getObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
@@ -110,6 +113,9 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
       DynamoDbClient deletedAccountsDynamoDbClient = DynamoDbFromConfig.client(
           configuration.getDeletedAccountsDynamoDbConfiguration(),
           software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
+      DynamoDbClient phoneNumberIdentifiersDynamoDbClient =
+          DynamoDbFromConfig.client(configuration.getPhoneNumberIdentifiersDynamoDbConfiguration(),
+              software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
 
       FaultTolerantRedisCluster cacheCluster = new FaultTolerantRedisCluster("main_cache_cluster",
           configuration.getCacheClusterConfiguration(), redisClusterClientResources);
@@ -122,11 +128,9 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
           .executorService(name(getClass(), "storageService-%d")).maxThreads(8).minThreads(1).build();
 
       ExternalServiceCredentialGenerator backupCredentialsGenerator = new ExternalServiceCredentialGenerator(
-          configuration.getSecureBackupServiceConfiguration().getUserAuthenticationTokenSharedSecret(), new byte[0],
-          false);
+          configuration.getSecureBackupServiceConfiguration().getUserAuthenticationTokenSharedSecret(), true);
       ExternalServiceCredentialGenerator storageCredentialsGenerator = new ExternalServiceCredentialGenerator(
-          configuration.getSecureStorageServiceConfiguration().getUserAuthenticationTokenSharedSecret(), new byte[0],
-          false);
+          configuration.getSecureStorageServiceConfiguration().getUserAuthenticationTokenSharedSecret(), true);
 
       DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager = new DynamicConfigurationManager<>(
           configuration.getAppConfig().getApplication(), configuration.getAppConfig().getEnvironment(),
@@ -156,7 +160,10 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
       Accounts accounts = new Accounts(accountsDynamoDbClient,
           configuration.getAccountsDynamoDbConfiguration().getTableName(),
           configuration.getAccountsDynamoDbConfiguration().getPhoneNumberTableName(),
+          configuration.getAccountsDynamoDbConfiguration().getPhoneNumberIdentifierTableName(),
           configuration.getAccountsDynamoDbConfiguration().getScanPageSize());
+      PhoneNumberIdentifiers phoneNumberIdentifiers = new PhoneNumberIdentifiers(phoneNumberIdentifiersDynamoDbClient,
+          configuration.getPhoneNumberIdentifiersDynamoDbConfiguration().getTableName());
       Usernames usernames = new Usernames(accountDatabase);
       Profiles profiles = new Profiles(accountDatabase);
       ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase);
@@ -183,7 +190,7 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
           Executors.newSingleThreadScheduledExecutor(), keyspaceNotificationDispatchExecutor);
       MessagesCache messagesCache = new MessagesCache(messageInsertCacheCluster, messageReadDeleteCluster,
           keyspaceNotificationDispatchExecutor);
-      PushLatencyManager pushLatencyManager = new PushLatencyManager(metricsCluster);
+      PushLatencyManager pushLatencyManager = new PushLatencyManager(metricsCluster, dynamicConfigurationManager);
       DirectoryQueue directoryQueue = new DirectoryQueue(
           configuration.getDirectoryConfiguration().getSqsConfiguration());
       UsernamesManager usernamesManager = new UsernamesManager(usernames, reservedUsernames, cacheCluster);
@@ -199,12 +206,12 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
           deletedAccountsLockDynamoDbClient,
           configuration.getDeletedAccountsLockDynamoDbConfiguration().getTableName());
       StoredVerificationCodeManager pendingAccountsManager = new StoredVerificationCodeManager(pendingAccounts);
-      AccountsManager accountsManager = new AccountsManager(accounts, cacheCluster,
+      AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
           deletedAccountsManager, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager,
-          pendingAccountsManager, secureStorageClient, secureBackupClient, clientPresenceManager);
+          pendingAccountsManager, secureStorageClient, secureBackupClient, clientPresenceManager, clock);
 
       for (String user : users) {
-        Optional<Account> account = accountsManager.get(user);
+        Optional<Account> account = accountsManager.getByE164(user);
 
         if (account.isPresent()) {
           accountsManager.delete(account.get(), DeletionReason.ADMIN_DELETED);

service/src/main/java/org/whispersystems/textsecuregcm/workers/SetUserDiscoverabilityCommand.java

@@ -18,6 +18,7 @@ import io.dropwizard.jdbi3.JdbiFactory;
 import io.dropwizard.setup.Environment;
 import io.lettuce.core.resource.ClientResources;
 import io.micrometer.core.instrument.Metrics;
+import java.time.Clock;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.ExecutorService;
@@ -45,6 +46,7 @@ import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
 import org.whispersystems.textsecuregcm.storage.MessagesCache;
 import org.whispersystems.textsecuregcm.storage.MessagesDynamoDb;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
+import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
 import org.whispersystems.textsecuregcm.storage.Profiles;
 import org.whispersystems.textsecuregcm.storage.ProfilesManager;
 import org.whispersystems.textsecuregcm.storage.ReportMessageDynamoDb;
@@ -91,6 +93,7 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
       final WhisperServerConfiguration configuration) throws Exception {
 
     try {
+      Clock clock = Clock.systemUTC();
       environment.getObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
 
       JdbiFactory jdbiFactory = new JdbiFactory();
@@ -112,6 +115,9 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
       DynamoDbClient deletedAccountsDynamoDbClient = DynamoDbFromConfig
           .client(configuration.getDeletedAccountsDynamoDbConfiguration(),
               software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
+      DynamoDbClient phoneNumberIdentifiersDynamoDbClient =
+          DynamoDbFromConfig.client(configuration.getPhoneNumberIdentifiersDynamoDbConfiguration(),
+              software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
 
       FaultTolerantRedisCluster cacheCluster = new FaultTolerantRedisCluster("main_cache_cluster",
           configuration.getCacheClusterConfiguration(), redisClusterClientResources);
@@ -126,11 +132,9 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
           .executorService(name(getClass(), "storageService-%d")).maxThreads(8).minThreads(1).build();
 
       ExternalServiceCredentialGenerator backupCredentialsGenerator = new ExternalServiceCredentialGenerator(
-          configuration.getSecureBackupServiceConfiguration().getUserAuthenticationTokenSharedSecret(), new byte[0],
-          false);
+          configuration.getSecureBackupServiceConfiguration().getUserAuthenticationTokenSharedSecret(), true);
       ExternalServiceCredentialGenerator storageCredentialsGenerator = new ExternalServiceCredentialGenerator(
-          configuration.getSecureStorageServiceConfiguration().getUserAuthenticationTokenSharedSecret(), new byte[0],
-          false);
+          configuration.getSecureStorageServiceConfiguration().getUserAuthenticationTokenSharedSecret(), true);
 
       DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager = new DynamicConfigurationManager<>(
           configuration.getAppConfig().getApplication(), configuration.getAppConfig().getEnvironment(),
@@ -161,7 +165,10 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
       Accounts accounts = new Accounts(accountsDynamoDbClient,
           configuration.getAccountsDynamoDbConfiguration().getTableName(),
           configuration.getAccountsDynamoDbConfiguration().getPhoneNumberTableName(),
+          configuration.getAccountsDynamoDbConfiguration().getPhoneNumberIdentifierTableName(),
           configuration.getAccountsDynamoDbConfiguration().getScanPageSize());
+      PhoneNumberIdentifiers phoneNumberIdentifiers = new PhoneNumberIdentifiers(phoneNumberIdentifiersDynamoDbClient,
+          configuration.getPhoneNumberIdentifiersDynamoDbConfiguration().getTableName());
       Usernames usernames = new Usernames(accountDatabase);
       Profiles profiles = new Profiles(accountDatabase);
       ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase);
@@ -186,7 +193,7 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
           Executors.newSingleThreadScheduledExecutor(), keyspaceNotificationDispatchExecutor);
       MessagesCache messagesCache = new MessagesCache(messageInsertCacheCluster, messageReadDeleteCluster,
           keyspaceNotificationDispatchExecutor);
-      PushLatencyManager pushLatencyManager = new PushLatencyManager(metricsCluster);
+      PushLatencyManager pushLatencyManager = new PushLatencyManager(metricsCluster, dynamicConfigurationManager);
       DirectoryQueue directoryQueue = new DirectoryQueue(
           configuration.getDirectoryConfiguration().getSqsConfiguration());
       UsernamesManager usernamesManager = new UsernamesManager(usernames, reservedUsernames, cacheCluster);
@@ -202,16 +209,16 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
           deletedAccountsLockDynamoDbClient,
           configuration.getDeletedAccountsLockDynamoDbConfiguration().getTableName());
       StoredVerificationCodeManager pendingAccountsManager = new StoredVerificationCodeManager(pendingAccounts);
-      AccountsManager accountsManager = new AccountsManager(accounts, cacheCluster,
+      AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
           deletedAccountsManager, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager,
-          pendingAccountsManager, secureStorageClient, secureBackupClient, clientPresenceManager);
+          pendingAccountsManager, secureStorageClient, secureBackupClient, clientPresenceManager, clock);
 
       Optional<Account> maybeAccount;
 
       try {
-        maybeAccount = accountsManager.get(UUID.fromString(namespace.getString("user")));
+        maybeAccount = accountsManager.getByAccountIdentifier(UUID.fromString(namespace.getString("user")));
       } catch (final IllegalArgumentException e) {
-        maybeAccount = accountsManager.get(namespace.getString("user"));
+        maybeAccount = accountsManager.getByE164(namespace.getString("user"));
       }
 
       maybeAccount.ifPresentOrElse(account -> {

service/src/main/resources/org/signal/badges/Badges.properties

@@ -16,13 +16,13 @@ TEST3_name = Test Badge Gamma
 TEST3_description = {short_name} is testing the gamma test badge.
 
 R_LOW_name = Signal Star
-R_LOW_description = {short_name} supports Signal by making a monthly donation.
+R_LOW_description = {short_name} is a Sustainer. Signal is a nonprofit with no advertisers or investors, supported only by people like you.
 
 R_MED_name = Signal Planet
-R_MED_description = {short_name} supports Signal by making a monthly donation.
+R_MED_description = {short_name} is a Sustainer. Signal is a nonprofit with no advertisers or investors, supported only by people like you.
 
 R_HIGH_name = Signal Sun
-R_HIGH_description = {short_name} supports Signal by making a monthly donation.
+R_HIGH_description = {short_name} is a Sustainer. Signal is a nonprofit with no advertisers or investors, supported only by people like you.
 
 BOOST_name = Signal Boost
-BOOST_description = {short_name} supported Signal by making a donation.
+BOOST_description = {short_name} supported Signal with a boost. Signal is a nonprofit with no advertisers or investors, supported only by people like you.

service/src/test/java/org/whispersystems/textsecuregcm/auth/AuthEnablementRefreshRequirementProviderTest.java

@@ -128,7 +128,7 @@ class AuthEnablementRefreshRequirementProviderTest {
     account.addDevice(authenticatedDevice);
     LongStream.range(2, 4).forEach(deviceId -> account.addDevice(DevicesHelper.createDevice(deviceId)));
 
-    when(accountsManager.get(uuid)).thenReturn(Optional.of(account));
+    when(accountsManager.getByAccountIdentifier(uuid)).thenReturn(Optional.of(account));
 
     account.getDevices()
         .forEach(device -> when(clientPresenceManager.isPresent(uuid, device.getId())).thenReturn(true));
@@ -310,7 +310,7 @@ class AuthEnablementRefreshRequirementProviderTest {
 
     assertEquals(200, response.getStatus());
 
-    verify(accountsManager, never()).get(any(UUID.class));
+    verify(accountsManager, never()).getByAccountIdentifier(any(UUID.class));
   }
 
   @Nested

service/src/test/java/org/whispersystems/textsecuregcm/auth/BaseAccountAuthenticatorTest.java

@@ -31,9 +31,6 @@ import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
-import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
-import org.whispersystems.textsecuregcm.auth.AuthenticationCredentials;
-import org.whispersystems.textsecuregcm.auth.BaseAccountAuthenticator;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
@@ -62,11 +59,14 @@ class BaseAccountAuthenticatorTest {
     clock = mock(Clock.class);
     baseAccountAuthenticator = new BaseAccountAuthenticator(accountsManager, clock);
 
-    acct1 = new Account("+14088675309", AuthHelper.getRandomUUID(random), Set.of(new Device(1, null, null, null,
+    acct1 = new Account("+14088675309", AuthHelper.getRandomUUID(random), UUID.randomUUID(),
+        Set.of(new Device(1, null, null, null,
         null, null, null, false, 0, null, yesterday, 0, null, 0, null)), null);
-    acct2 = new Account("+14098675309", AuthHelper.getRandomUUID(random), Set.of(new Device(1, null, null, null,
+    acct2 = new Account("+14098675309", AuthHelper.getRandomUUID(random), UUID.randomUUID(),
+        Set.of(new Device(1, null, null, null,
         null, null, null, false, 0, null, yesterday, 0, null, 0, null)), null);
-    oldAccount = new Account("+14108675309", AuthHelper.getRandomUUID(random), Set.of(new Device(1, null, null, null,
+    oldAccount = new Account("+14108675309", AuthHelper.getRandomUUID(random), UUID.randomUUID(),
+        Set.of(new Device(1, null, null, null,
         null, null, null, false, 0, null, oldTime, 0, null, 0, null)), null);
 
     AccountsHelper.setupMockUpdate(accountsManager);
@@ -156,7 +156,7 @@ class BaseAccountAuthenticatorTest {
     final AuthenticationCredentials credentials = mock(AuthenticationCredentials.class);
 
     when(clock.instant()).thenReturn(Instant.now());
-    when(accountsManager.get(uuid)).thenReturn(Optional.of(account));
+    when(accountsManager.getByAccountIdentifier(uuid)).thenReturn(Optional.of(account));
     when(account.getUuid()).thenReturn(uuid);
     when(account.getDevice(deviceId)).thenReturn(Optional.of(device));
     when(account.isEnabled()).thenReturn(true);
@@ -184,7 +184,7 @@ class BaseAccountAuthenticatorTest {
     final AuthenticationCredentials credentials = mock(AuthenticationCredentials.class);
 
     when(clock.instant()).thenReturn(Instant.now());
-    when(accountsManager.get(uuid)).thenReturn(Optional.of(account));
+    when(accountsManager.getByAccountIdentifier(uuid)).thenReturn(Optional.of(account));
     when(account.getUuid()).thenReturn(uuid);
     when(account.getDevice(deviceId)).thenReturn(Optional.of(device));
     when(account.isEnabled()).thenReturn(true);
@@ -213,7 +213,7 @@ class BaseAccountAuthenticatorTest {
     final AuthenticationCredentials credentials = mock(AuthenticationCredentials.class);
 
     when(clock.instant()).thenReturn(Instant.now());
-    when(accountsManager.get(uuid)).thenReturn(Optional.of(account));
+    when(accountsManager.getByAccountIdentifier(uuid)).thenReturn(Optional.of(account));
     when(account.getUuid()).thenReturn(uuid);
     when(account.getDevice(deviceId)).thenReturn(Optional.of(device));
     when(account.isEnabled()).thenReturn(false);
@@ -251,7 +251,7 @@ class BaseAccountAuthenticatorTest {
     final AuthenticationCredentials credentials = mock(AuthenticationCredentials.class);
 
     when(clock.instant()).thenReturn(Instant.now());
-    when(accountsManager.get(uuid)).thenReturn(Optional.of(account));
+    when(accountsManager.getByAccountIdentifier(uuid)).thenReturn(Optional.of(account));
     when(account.getUuid()).thenReturn(uuid);
     when(account.getDevice(deviceId)).thenReturn(Optional.of(device));
     when(account.isEnabled()).thenReturn(true);
@@ -278,7 +278,7 @@ class BaseAccountAuthenticatorTest {
     final AuthenticationCredentials credentials = mock(AuthenticationCredentials.class);
 
     when(clock.instant()).thenReturn(Instant.now());
-    when(accountsManager.get(uuid)).thenReturn(Optional.of(account));
+    when(accountsManager.getByAccountIdentifier(uuid)).thenReturn(Optional.of(account));
     when(account.getUuid()).thenReturn(uuid);
     when(account.getDevice(deviceId)).thenReturn(Optional.of(device));
     when(account.isEnabled()).thenReturn(true);
@@ -303,7 +303,7 @@ class BaseAccountAuthenticatorTest {
         () -> baseAccountAuthenticator.authenticate(new BasicCredentials(username, "password"), true));
 
     assertThat(maybeAuthenticatedAccount).isEmpty();
-    verify(accountsManager, never()).get(any(UUID.class));
+    verify(accountsManager, never()).getByAccountIdentifier(any(UUID.class));
   }
 
   private static Stream<String> testAuthenticateMalformedCredentials() {

service/src/test/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfigurationTest.java

@@ -272,20 +272,20 @@ class DynamicConfigurationTest {
       final DynamicConfiguration emptyConfig =
           DynamicConfigurationManager.parseConfiguration(emptyConfigYaml, DynamicConfiguration.class).orElseThrow();
 
-      assertTrue(emptyConfig.getPaymentsConfiguration().getAllowedCountryCodes().isEmpty());
+      assertTrue(emptyConfig.getPaymentsConfiguration().getDisallowedPrefixes().isEmpty());
     }
 
     {
       final String paymentsConfigYaml =
           "payments:\n"
-              + "  allowedCountryCodes:\n"
-              + "    - 44";
+              + "  disallowedPrefixes:\n"
+              + "    - +44";
 
       final DynamicPaymentsConfiguration config =
           DynamicConfigurationManager.parseConfiguration(paymentsConfigYaml, DynamicConfiguration.class).orElseThrow()
               .getPaymentsConfiguration();
 
-      assertEquals(Set.of("44"), config.getAllowedCountryCodes());
+      assertEquals(List.of("+44"), config.getDisallowedPrefixes());
     }
   }
 

service/src/test/java/org/whispersystems/textsecuregcm/metrics/PushLatencyManagerTest.java

@@ -6,38 +6,68 @@
 package org.whispersystems.textsecuregcm.metrics;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
+import java.time.Clock;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.util.Collections;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
-import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
+import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicPushLatencyConfiguration;
+import org.whispersystems.textsecuregcm.metrics.PushLatencyManager.PushRecord;
+import org.whispersystems.textsecuregcm.metrics.PushLatencyManager.PushType;
 import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
+import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 
 class PushLatencyManagerTest {
 
   @RegisterExtension
   static final RedisClusterExtension REDIS_CLUSTER_EXTENSION = RedisClusterExtension.builder().build();
 
-  @Test
-  void testGetLatency() throws ExecutionException, InterruptedException {
+  private DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager;
 
-    final PushLatencyManager pushLatencyManager = new PushLatencyManager(REDIS_CLUSTER_EXTENSION.getRedisCluster());
+  @BeforeEach
+  void setUp() {
+    //noinspection unchecked
+    dynamicConfigurationManager = mock(DynamicConfigurationManager.class);
+    final DynamicConfiguration dynamicConfiguration = mock(DynamicConfiguration.class);
+    final DynamicPushLatencyConfiguration dynamicPushLatencyConfiguration = mock(DynamicPushLatencyConfiguration.class);
+
+    when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration);
+    when(dynamicConfiguration.getPushLatencyConfiguration()).thenReturn(dynamicPushLatencyConfiguration);
+    when(dynamicPushLatencyConfiguration.getInstrumentedVersions()).thenReturn(Collections.emptyMap());
+  }
+
+  @ParameterizedTest
+  @ValueSource(booleans = {true, false})
+  void testTakeRecord(final boolean isVoip) throws ExecutionException, InterruptedException {
     final UUID accountUuid = UUID.randomUUID();
     final long deviceId = 1;
-    final long expectedLatency = 1234;
-    final long pushSentTimestamp = System.currentTimeMillis();
-    final long clearQueueTimestamp = pushSentTimestamp + expectedLatency;
 
-    assertNull(pushLatencyManager.getLatencyAndClearTimestamp(accountUuid, deviceId, System.currentTimeMillis()).get());
+    final Instant pushTimestamp = Instant.now();
 
-    {
-      pushLatencyManager.recordPushSent(accountUuid, deviceId, pushSentTimestamp);
+    final PushLatencyManager pushLatencyManager = new PushLatencyManager(REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+        dynamicConfigurationManager, Clock.fixed(pushTimestamp, ZoneId.systemDefault()));
 
-      assertEquals(expectedLatency,
-          (long) pushLatencyManager.getLatencyAndClearTimestamp(accountUuid, deviceId, clearQueueTimestamp).get());
-      assertNull(
-          pushLatencyManager.getLatencyAndClearTimestamp(accountUuid, deviceId, System.currentTimeMillis()).get());
-    }
+    assertNull(pushLatencyManager.takePushRecord(accountUuid, deviceId).get());
+
+    pushLatencyManager.recordPushSent(accountUuid, deviceId, isVoip);
+
+    final PushRecord pushRecord = pushLatencyManager.takePushRecord(accountUuid, deviceId).get();
+
+    assertNotNull(pushRecord);
+    assertEquals(pushTimestamp, pushRecord.getTimestamp());
+    assertEquals(isVoip ? PushType.VOIP : PushType.STANDARD, pushRecord.getPushType());
+
+    assertNull(pushLatencyManager.takePushRecord(accountUuid, deviceId).get());
   }
 }

service/src/test/java/org/whispersystems/textsecuregcm/push/ApnFallbackManagerTest.java

@@ -57,8 +57,8 @@ class ApnFallbackManagerTest {
     when(account.getDevice(DEVICE_ID)).thenReturn(Optional.of(device));
 
     final AccountsManager accountsManager = mock(AccountsManager.class);
-    when(accountsManager.get(ACCOUNT_NUMBER)).thenReturn(Optional.of(account));
-    when(accountsManager.get(ACCOUNT_UUID)).thenReturn(Optional.of(account));
+    when(accountsManager.getByE164(ACCOUNT_NUMBER)).thenReturn(Optional.of(account));
+    when(accountsManager.getByAccountIdentifier(ACCOUNT_UUID)).thenReturn(Optional.of(account));
 
     apnSender = mock(APNSender.class);
 

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

@@ -57,7 +57,7 @@ public class AccountDatabaseCrawlerIntegrationTest extends AbstractRedisClusterT
         .thenReturn(new AccountCrawlChunk(List.of(secondAccount), SECOND_UUID))
         .thenReturn(new AccountCrawlChunk(Collections.emptyList(), null));
 
-    final AccountDatabaseCrawlerCache crawlerCache = new AccountDatabaseCrawlerCache(getRedisCluster());
+    final AccountDatabaseCrawlerCache crawlerCache = new AccountDatabaseCrawlerCache(getRedisCluster(), "test");
     accountDatabaseCrawler = new AccountDatabaseCrawler(accountsManager, crawlerCache, List.of(listener), CHUNK_SIZE,
         CHUNK_INTERVAL_MS);
   }

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

@@ -6,6 +6,7 @@
 package org.whispersystems.textsecuregcm.storage;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
@@ -15,6 +16,8 @@ import static org.mockito.Mockito.when;
 import com.opentable.db.postgres.embedded.LiquibasePreparer;
 import com.opentable.db.postgres.junit5.EmbeddedPostgresExtension;
 import com.opentable.db.postgres.junit5.PreparedDbExtension;
+import java.time.Clock;
+import java.util.ArrayList;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
@@ -47,6 +50,8 @@ class AccountsManagerChangeNumberIntegrationTest {
 
   private static final String ACCOUNTS_TABLE_NAME = "accounts_test";
   private static final String NUMBERS_TABLE_NAME = "numbers_test";
+  private static final String PNI_ASSIGNMENT_TABLE_NAME = "pni_assignment_test";
+  private static final String PNI_TABLE_NAME = "pni_test";
   private static final String NEEDS_RECONCILIATION_INDEX_NAME = "needs_reconciliation_test";
   private static final String DELETED_ACCOUNTS_LOCK_TABLE_NAME = "deleted_accounts_lock_test";
   private static final int SCAN_PAGE_SIZE = 1;
@@ -90,6 +95,16 @@ class AccountsManagerChangeNumberIntegrationTest {
           .attributeType(ScalarAttributeType.S).build())
       .build();
 
+  @RegisterExtension
+  static DynamoDbExtension PNI_DYNAMO_EXTENSION = DynamoDbExtension.builder()
+      .tableName(PNI_TABLE_NAME)
+      .hashKey(PhoneNumberIdentifiers.KEY_E164)
+      .attributeDefinition(AttributeDefinition.builder()
+          .attributeName(PhoneNumberIdentifiers.KEY_E164)
+          .attributeType(ScalarAttributeType.S)
+          .build())
+      .build();
+
   @RegisterExtension
   static RedisClusterExtension CACHE_CLUSTER_EXTENSION = RedisClusterExtension.builder().build();
 
@@ -118,14 +133,33 @@ class AccountsManagerChangeNumberIntegrationTest {
       ACCOUNTS_DYNAMO_EXTENSION.getDynamoDbClient().createTable(createNumbersTableRequest);
     }
 
+    {
+      CreateTableRequest createPhoneNumberIdentifierTableRequest = CreateTableRequest.builder()
+          .tableName(PNI_ASSIGNMENT_TABLE_NAME)
+          .keySchema(KeySchemaElement.builder()
+              .attributeName(Accounts.ATTR_PNI_UUID)
+              .keyType(KeyType.HASH)
+              .build())
+          .attributeDefinitions(AttributeDefinition.builder()
+              .attributeName(Accounts.ATTR_PNI_UUID)
+              .attributeType(ScalarAttributeType.B)
+              .build())
+          .provisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT)
+          .build();
+
+      ACCOUNTS_DYNAMO_EXTENSION.getDynamoDbClient().createTable(createPhoneNumberIdentifierTableRequest);
+    }
+
     final Accounts accounts = new Accounts(
         ACCOUNTS_DYNAMO_EXTENSION.getDynamoDbClient(),
         ACCOUNTS_DYNAMO_EXTENSION.getTableName(),
         NUMBERS_TABLE_NAME,
+        PNI_ASSIGNMENT_TABLE_NAME,
         SCAN_PAGE_SIZE);
 
     {
-      final DynamicConfigurationManager dynamicConfigurationManager = mock(DynamicConfigurationManager.class);
+      @SuppressWarnings("unchecked") final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager =
+          mock(DynamicConfigurationManager.class);
 
       DynamicConfiguration dynamicConfiguration = new DynamicConfiguration();
       when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration);
@@ -146,8 +180,12 @@ class AccountsManagerChangeNumberIntegrationTest {
 
       clientPresenceManager = mock(ClientPresenceManager.class);
 
+      final PhoneNumberIdentifiers phoneNumberIdentifiers =
+          new PhoneNumberIdentifiers(PNI_DYNAMO_EXTENSION.getDynamoDbClient(), PNI_TABLE_NAME);
+
       accountsManager = new AccountsManager(
           accounts,
+          phoneNumberIdentifiers,
           CACHE_CLUSTER_EXTENSION.getRedisCluster(),
           deletedAccountsManager,
           mock(DirectoryQueue.class),
@@ -158,7 +196,8 @@ class AccountsManagerChangeNumberIntegrationTest {
           mock(StoredVerificationCodeManager.class),
           secureStorageClient,
           secureBackupClient,
-          clientPresenceManager);
+          clientPresenceManager,
+          mock(Clock.class));
     }
   }
 
@@ -167,17 +206,19 @@ class AccountsManagerChangeNumberIntegrationTest {
     final String originalNumber = "+18005551111";
     final String secondNumber = "+18005552222";
 
-    final Account account = accountsManager.create(originalNumber, "password", null, new AccountAttributes());
+    final Account account = accountsManager.create(originalNumber, "password", null, new AccountAttributes(), new ArrayList<>());
     final UUID originalUuid = account.getUuid();
+    final UUID originalPni = account.getPhoneNumberIdentifier().orElseThrow();
 
     accountsManager.changeNumber(account, secondNumber);
 
-    assertTrue(accountsManager.get(originalNumber).isEmpty());
+    assertTrue(accountsManager.getByE164(originalNumber).isEmpty());
 
-    assertTrue(accountsManager.get(secondNumber).isPresent());
-    assertEquals(Optional.of(originalUuid), accountsManager.get(secondNumber).map(Account::getUuid));
+    assertTrue(accountsManager.getByE164(secondNumber).isPresent());
+    assertEquals(originalUuid, accountsManager.getByE164(secondNumber).map(Account::getUuid).orElseThrow());
+    assertNotEquals(originalPni, accountsManager.getByE164(secondNumber).flatMap(Account::getPhoneNumberIdentifier).orElseThrow());
 
-    assertEquals(secondNumber, accountsManager.get(originalUuid).map(Account::getNumber).orElseThrow());
+    assertEquals(secondNumber, accountsManager.getByAccountIdentifier(originalUuid).map(Account::getNumber).orElseThrow());
 
     assertEquals(Optional.empty(), deletedAccounts.findUuid(originalNumber));
     assertEquals(Optional.empty(), deletedAccounts.findUuid(secondNumber));
@@ -188,18 +229,20 @@ class AccountsManagerChangeNumberIntegrationTest {
     final String originalNumber = "+18005551111";
     final String secondNumber = "+18005552222";
 
-    Account account = accountsManager.create(originalNumber, "password", null, new AccountAttributes());
+    Account account = accountsManager.create(originalNumber, "password", null, new AccountAttributes(), new ArrayList<>());
     final UUID originalUuid = account.getUuid();
+    final UUID originalPni = account.getPhoneNumberIdentifier().orElseThrow();
 
     account = accountsManager.changeNumber(account, secondNumber);
     accountsManager.changeNumber(account, originalNumber);
 
-    assertTrue(accountsManager.get(originalNumber).isPresent());
-    assertEquals(Optional.of(originalUuid), accountsManager.get(originalNumber).map(Account::getUuid));
+    assertTrue(accountsManager.getByE164(originalNumber).isPresent());
+    assertEquals(originalUuid, accountsManager.getByE164(originalNumber).map(Account::getUuid).orElseThrow());
+    assertEquals(originalPni, accountsManager.getByE164(originalNumber).flatMap(Account::getPhoneNumberIdentifier).orElseThrow());
 
-    assertTrue(accountsManager.get(secondNumber).isEmpty());
+    assertTrue(accountsManager.getByE164(secondNumber).isEmpty());
 
-    assertEquals(originalNumber, accountsManager.get(originalUuid).map(Account::getNumber).orElseThrow());
+    assertEquals(originalNumber, accountsManager.getByAccountIdentifier(originalUuid).map(Account::getNumber).orElseThrow());
 
     assertEquals(Optional.empty(), deletedAccounts.findUuid(originalNumber));
     assertEquals(Optional.empty(), deletedAccounts.findUuid(secondNumber));
@@ -210,22 +253,22 @@ class AccountsManagerChangeNumberIntegrationTest {
     final String originalNumber = "+18005551111";
     final String secondNumber = "+18005552222";
 
-    final Account account = accountsManager.create(originalNumber, "password", null, new AccountAttributes());
+    final Account account = accountsManager.create(originalNumber, "password", null, new AccountAttributes(), new ArrayList<>());
     final UUID originalUuid = account.getUuid();
 
-    final Account existingAccount = accountsManager.create(secondNumber, "password", null, new AccountAttributes());
+    final Account existingAccount = accountsManager.create(secondNumber, "password", null, new AccountAttributes(), new ArrayList<>());
     final UUID existingAccountUuid = existingAccount.getUuid();
 
     accountsManager.update(existingAccount, a -> a.addDevice(new Device(Device.MASTER_ID, "test", "token", "salt", null, null, null, true, 1, null, 0, 0, null, 0, new DeviceCapabilities())));
 
     accountsManager.changeNumber(account, secondNumber);
 
-    assertTrue(accountsManager.get(originalNumber).isEmpty());
+    assertTrue(accountsManager.getByE164(originalNumber).isEmpty());
 
-    assertTrue(accountsManager.get(secondNumber).isPresent());
-    assertEquals(Optional.of(originalUuid), accountsManager.get(secondNumber).map(Account::getUuid));
+    assertTrue(accountsManager.getByE164(secondNumber).isPresent());
+    assertEquals(Optional.of(originalUuid), accountsManager.getByE164(secondNumber).map(Account::getUuid));
 
-    assertEquals(secondNumber, accountsManager.get(originalUuid).map(Account::getNumber).orElseThrow());
+    assertEquals(secondNumber, accountsManager.getByAccountIdentifier(originalUuid).map(Account::getNumber).orElseThrow());
 
     verify(clientPresenceManager).displacePresence(existingAccountUuid, Device.MASTER_ID);
 
@@ -238,24 +281,28 @@ class AccountsManagerChangeNumberIntegrationTest {
     final String originalNumber = "+18005551111";
     final String secondNumber = "+18005552222";
 
-    final Account account = accountsManager.create(originalNumber, "password", null, new AccountAttributes());
+    final Account account = accountsManager.create(originalNumber, "password", null, new AccountAttributes(), new ArrayList<>());
     final UUID originalUuid = account.getUuid();
+    final UUID originalPni = account.getPhoneNumberIdentifier().orElseThrow();
 
-    final Account existingAccount = accountsManager.create(secondNumber, "password", null, new AccountAttributes());
+    final Account existingAccount = accountsManager.create(secondNumber, "password", null, new AccountAttributes(), new ArrayList<>());
     final UUID existingAccountUuid = existingAccount.getUuid();
 
-    accountsManager.changeNumber(account, secondNumber);
+    final Account changedNumberAccount = accountsManager.changeNumber(account, secondNumber);
+    final UUID secondPni = changedNumberAccount.getPhoneNumberIdentifier().orElseThrow();
 
-    final Account reRegisteredAccount = accountsManager.create(originalNumber, "password", null, new AccountAttributes());
+    final Account reRegisteredAccount = accountsManager.create(originalNumber, "password", null, new AccountAttributes(), new ArrayList<>());
 
     assertEquals(existingAccountUuid, reRegisteredAccount.getUuid());
+    assertEquals(originalPni, reRegisteredAccount.getPhoneNumberIdentifier().orElseThrow());
 
     assertEquals(Optional.empty(), deletedAccounts.findUuid(originalNumber));
     assertEquals(Optional.empty(), deletedAccounts.findUuid(secondNumber));
 
-    accountsManager.changeNumber(reRegisteredAccount, secondNumber);
+    final Account changedNumberReRegisteredAccount = accountsManager.changeNumber(reRegisteredAccount, secondNumber);
 
     assertEquals(Optional.of(originalUuid), deletedAccounts.findUuid(originalNumber));
     assertEquals(Optional.empty(), deletedAccounts.findUuid(secondNumber));
+    assertEquals(secondPni, changedNumberReRegisteredAccount.getPhoneNumberIdentifier().orElseThrow());
   }
 }

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

@@ -15,10 +15,13 @@ import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
 import java.io.IOException;
+import java.time.Clock;
 import java.time.Instant;
+import java.util.ArrayList;
 import java.util.Optional;
 import java.util.Random;
 import java.util.UUID;
@@ -33,6 +36,8 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.mockito.ArgumentCaptor;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
 import org.whispersystems.textsecuregcm.auth.AuthenticationCredentials;
 import org.whispersystems.textsecuregcm.entities.AccountAttributes;
 import org.whispersystems.textsecuregcm.entities.SignedPreKey;
@@ -53,6 +58,7 @@ class AccountsManagerConcurrentModificationIntegrationTest {
 
   private static final String ACCOUNTS_TABLE_NAME = "accounts_test";
   private static final String NUMBERS_TABLE_NAME = "numbers_test";
+  private static final String PNI_TABLE_NAME = "pni_test";
 
   private static final int SCAN_PAGE_SIZE = 1;
 
@@ -94,10 +100,28 @@ class AccountsManagerConcurrentModificationIntegrationTest {
       dynamoDbExtension.getDynamoDbClient().createTable(createNumbersTableRequest);
     }
 
+    {
+      CreateTableRequest createPhoneNumberIdentifierTableRequest = CreateTableRequest.builder()
+          .tableName(PNI_TABLE_NAME)
+          .keySchema(KeySchemaElement.builder()
+              .attributeName(Accounts.ATTR_PNI_UUID)
+              .keyType(KeyType.HASH)
+              .build())
+          .attributeDefinitions(AttributeDefinition.builder()
+              .attributeName(Accounts.ATTR_PNI_UUID)
+              .attributeType(ScalarAttributeType.B)
+              .build())
+          .provisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT)
+          .build();
+
+      dynamoDbExtension.getDynamoDbClient().createTable(createPhoneNumberIdentifierTableRequest);
+    }
+
     accounts = new Accounts(
         dynamoDbExtension.getDynamoDbClient(),
         dynamoDbExtension.getTableName(),
         NUMBERS_TABLE_NAME,
+        PNI_TABLE_NAME,
         SCAN_PAGE_SIZE);
 
     {
@@ -112,8 +136,13 @@ class AccountsManagerConcurrentModificationIntegrationTest {
         return null;
       }).when(deletedAccountsManager).lockAndTake(anyString(), any());
 
+      final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
+      when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString()))
+          .thenAnswer((Answer<UUID>) invocation -> UUID.randomUUID());
+
       accountsManager = new AccountsManager(
           accounts,
+          phoneNumberIdentifiers,
           RedisClusterHelper.buildMockRedisCluster(commands),
           deletedAccountsManager,
           mock(DirectoryQueue.class),
@@ -124,7 +153,8 @@ class AccountsManagerConcurrentModificationIntegrationTest {
           mock(StoredVerificationCodeManager.class),
           mock(SecureStorageClient.class),
           mock(SecureBackupClient.class),
-          mock(ClientPresenceManager.class)
+          mock(ClientPresenceManager.class),
+          mock(Clock.class)
       );
     }
   }
@@ -135,7 +165,7 @@ class AccountsManagerConcurrentModificationIntegrationTest {
     final UUID uuid;
     {
       final Account account = accountsManager.update(
-          accountsManager.create("+14155551212", "password", null, new AccountAttributes()),
+          accountsManager.create("+14155551212", "password", null, new AccountAttributes(), new ArrayList<>()),
           a -> {
             a.setUnidentifiedAccessKey(new byte[16]);
 
@@ -182,8 +212,8 @@ class AccountsManagerConcurrentModificationIntegrationTest {
         modifyDevice(uuid, Device.MASTER_ID, device -> device.setName("deviceName"))
     ).join();
 
-    final Account managerAccount = accountsManager.get(uuid).orElseThrow();
-    final Account dynamoAccount = accounts.get(uuid).orElseThrow();
+    final Account managerAccount = accountsManager.getByAccountIdentifier(uuid).orElseThrow();
+    final Account dynamoAccount = accounts.getByAccountIdentifier(uuid).orElseThrow();
 
     final Account redisAccount = getLastAccountFromRedisMock(commands);
 
@@ -222,7 +252,7 @@ class AccountsManagerConcurrentModificationIntegrationTest {
   private CompletableFuture<?> modifyAccount(final UUID uuid, final Consumer<Account> accountMutation) {
 
     return CompletableFuture.runAsync(() -> {
-      final Account account = accountsManager.get(uuid).orElseThrow();
+      final Account account = accountsManager.getByAccountIdentifier(uuid).orElseThrow();
       accountsManager.update(account, accountMutation);
     }, mutationExecutor);
   }
@@ -230,7 +260,7 @@ class AccountsManagerConcurrentModificationIntegrationTest {
   private CompletableFuture<?> modifyDevice(final UUID uuid, final long deviceId, final Consumer<Device> deviceMutation) {
 
     return CompletableFuture.runAsync(() -> {
-      final Account account = accountsManager.get(uuid).orElseThrow();
+      final Account account = accountsManager.getByAccountIdentifier(uuid).orElseThrow();
       accountsManager.updateDevice(account, deviceId, deviceMutation);
     }, mutationExecutor);
   }

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

@@ -12,12 +12,14 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.uuid.UUIDComparator;
 import io.github.resilience4j.circuitbreaker.CallNotPermittedException;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -25,11 +27,6 @@ import java.util.Optional;
 import java.util.Random;
 import java.util.Set;
 import java.util.UUID;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.LinkedBlockingDeque;
-import java.util.concurrent.ThreadPoolExecutor;
-import java.util.concurrent.TimeUnit;
-import org.assertj.core.api.AssertionsForClassTypes;
 import org.jdbi.v3.core.transaction.TransactionException;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Disabled;
@@ -38,15 +35,21 @@ import org.junit.jupiter.api.extension.RegisterExtension;
 import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
 import org.whispersystems.textsecuregcm.entities.SignedPreKey;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
+import org.whispersystems.textsecuregcm.util.SystemMapper;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
-import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
+import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+import software.amazon.awssdk.services.dynamodb.model.CancellationReason;
 import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
 import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement;
 import software.amazon.awssdk.services.dynamodb.model.KeyType;
+import software.amazon.awssdk.services.dynamodb.model.Put;
+import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
+import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure;
 import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
+import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
 import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledException;
 import software.amazon.awssdk.services.dynamodb.model.TransactionConflictException;
@@ -55,7 +58,8 @@ import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
 class AccountsTest {
 
   private static final String ACCOUNTS_TABLE_NAME = "accounts_test";
-  private static final String NUMBERS_TABLE_NAME = "numbers_test";
+  private static final String NUMBER_CONSTRAINT_TABLE_NAME = "numbers_test";
+  private static final String PNI_CONSTRAINT_TABLE_NAME = "pni_test";
 
   private static final int SCAN_PAGE_SIZE = 1;
 
@@ -74,7 +78,7 @@ class AccountsTest {
   @BeforeEach
   void setupAccountsDao() {
     CreateTableRequest createNumbersTableRequest = CreateTableRequest.builder()
-        .tableName(NUMBERS_TABLE_NAME)
+        .tableName(NUMBER_CONSTRAINT_TABLE_NAME)
         .keySchema(KeySchemaElement.builder()
             .attributeName(Accounts.ATTR_ACCOUNT_E164)
             .keyType(KeyType.HASH)
@@ -88,27 +92,48 @@ class AccountsTest {
 
     dynamoDbExtension.getDynamoDbClient().createTable(createNumbersTableRequest);
 
+    CreateTableRequest createPhoneNumberIdentifierTableRequest = CreateTableRequest.builder()
+        .tableName(PNI_CONSTRAINT_TABLE_NAME)
+        .keySchema(KeySchemaElement.builder()
+            .attributeName(Accounts.ATTR_PNI_UUID)
+            .keyType(KeyType.HASH)
+            .build())
+        .attributeDefinitions(AttributeDefinition.builder()
+            .attributeName(Accounts.ATTR_PNI_UUID)
+            .attributeType(ScalarAttributeType.B)
+            .build())
+        .provisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT)
+        .build();
+
+    dynamoDbExtension.getDynamoDbClient().createTable(createPhoneNumberIdentifierTableRequest);
+
     this.accounts = new Accounts(
         dynamoDbExtension.getDynamoDbClient(),
         dynamoDbExtension.getTableName(),
-        NUMBERS_TABLE_NAME,
+        NUMBER_CONSTRAINT_TABLE_NAME,
+        PNI_CONSTRAINT_TABLE_NAME,
         SCAN_PAGE_SIZE);
   }
 
   @Test
   void testStore() {
-    Device  device  = generateDevice (1                                            );
-    Account account = generateAccount("+14151112222", UUID.randomUUID(), Collections.singleton(device));
+    Device device = generateDevice(1);
+    Account account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), Collections.singleton(device));
 
     boolean freshUser = accounts.create(account);
 
     assertThat(freshUser).isTrue();
-    verifyStoredState("+14151112222", account.getUuid(), account, true);
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier().orElseThrow(), account, true);
+
+    assertPhoneNumberConstraintExists("+14151112222", account.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(account.getPhoneNumberIdentifier().orElseThrow(), account.getUuid());
 
     freshUser = accounts.create(account);
     assertThat(freshUser).isTrue();
-    verifyStoredState("+14151112222", account.getUuid(), account, true);
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier().orElseThrow(), account, true);
 
+    assertPhoneNumberConstraintExists("+14151112222", account.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(account.getPhoneNumberIdentifier().orElseThrow(), account.getUuid());
   }
 
   @Test
@@ -117,11 +142,14 @@ class AccountsTest {
     devices.add(generateDevice(1));
     devices.add(generateDevice(2));
 
-    Account account = generateAccount("+14151112222", UUID.randomUUID(), devices);
+    Account account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), devices);
 
     accounts.create(account);
 
-    verifyStoredState("+14151112222", account.getUuid(), account, true);
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier().orElseThrow(), account, true);
+
+    assertPhoneNumberConstraintExists("+14151112222", account.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(account.getPhoneNumberIdentifier().orElseThrow(), account.getUuid());
   }
 
   @Test
@@ -131,46 +159,123 @@ class AccountsTest {
     devicesFirst.add(generateDevice(2));
 
     UUID uuidFirst = UUID.randomUUID();
-    Account accountFirst = generateAccount("+14151112222", uuidFirst, devicesFirst);
+    UUID pniFirst = UUID.randomUUID();
+    Account accountFirst = generateAccount("+14151112222", uuidFirst, pniFirst, devicesFirst);
 
     Set<Device> devicesSecond = new HashSet<>();
     devicesSecond.add(generateDevice(1));
     devicesSecond.add(generateDevice(2));
 
     UUID uuidSecond = UUID.randomUUID();
-    Account accountSecond = generateAccount("+14152221111", uuidSecond, devicesSecond);
+    UUID pniSecond = UUID.randomUUID();
+    Account accountSecond = generateAccount("+14152221111", uuidSecond, pniSecond, devicesSecond);
 
     accounts.create(accountFirst);
     accounts.create(accountSecond);
 
-    Optional<Account> retrievedFirst = accounts.get("+14151112222");
-    Optional<Account> retrievedSecond = accounts.get("+14152221111");
+    Optional<Account> retrievedFirst = accounts.getByE164("+14151112222");
+    Optional<Account> retrievedSecond = accounts.getByE164("+14152221111");
 
     assertThat(retrievedFirst.isPresent()).isTrue();
     assertThat(retrievedSecond.isPresent()).isTrue();
 
-    verifyStoredState("+14151112222", uuidFirst, retrievedFirst.get(), accountFirst);
-    verifyStoredState("+14152221111", uuidSecond, retrievedSecond.get(), accountSecond);
+    verifyStoredState("+14151112222", uuidFirst, pniFirst, retrievedFirst.get(), accountFirst);
+    verifyStoredState("+14152221111", uuidSecond, pniSecond, retrievedSecond.get(), accountSecond);
 
-    retrievedFirst = accounts.get(uuidFirst);
-    retrievedSecond = accounts.get(uuidSecond);
+    retrievedFirst = accounts.getByAccountIdentifier(uuidFirst);
+    retrievedSecond = accounts.getByAccountIdentifier(uuidSecond);
 
     assertThat(retrievedFirst.isPresent()).isTrue();
     assertThat(retrievedSecond.isPresent()).isTrue();
 
-    verifyStoredState("+14151112222", uuidFirst, retrievedFirst.get(), accountFirst);
-    verifyStoredState("+14152221111", uuidSecond, retrievedSecond.get(), accountSecond);
+    verifyStoredState("+14151112222", uuidFirst, pniFirst, retrievedFirst.get(), accountFirst);
+    verifyStoredState("+14152221111", uuidSecond, pniSecond, retrievedSecond.get(), accountSecond);
+
+    retrievedFirst = accounts.getByPhoneNumberIdentifier(pniFirst);
+    retrievedSecond = accounts.getByPhoneNumberIdentifier(pniSecond);
+
+    assertThat(retrievedFirst.isPresent()).isTrue();
+    assertThat(retrievedSecond.isPresent()).isTrue();
+
+    verifyStoredState("+14151112222", uuidFirst, pniFirst, retrievedFirst.get(), accountFirst);
+    verifyStoredState("+14152221111", uuidSecond, pniSecond, retrievedSecond.get(), accountSecond);
+  }
+
+  @Test
+  void testRetrieveNoPni() throws JsonProcessingException {
+    final Set<Device> devices = new HashSet<>();
+    devices.add(generateDevice(1));
+    devices.add(generateDevice(2));
+
+    final UUID uuid = UUID.randomUUID();
+
+    final Account account = generateAccount("+14151112222", uuid, null, devices);
+
+    // Accounts#create enforces that newly-created accounts have a PNI, so we need to make a bit of an end-run around it
+    // to simulate an existing account with no PNI.
+    {
+      final TransactWriteItem phoneNumberConstraintPut = TransactWriteItem.builder()
+          .put(
+              Put.builder()
+                  .tableName(NUMBER_CONSTRAINT_TABLE_NAME)
+                  .item(Map.of(
+                      Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
+                      Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
+                  .conditionExpression(
+                      "attribute_not_exists(#number) OR (attribute_exists(#number) AND #uuid = :uuid)")
+                  .expressionAttributeNames(
+                      Map.of("#uuid", Accounts.KEY_ACCOUNT_UUID,
+                          "#number", Accounts.ATTR_ACCOUNT_E164))
+                  .expressionAttributeValues(
+                      Map.of(":uuid", AttributeValues.fromUUID(account.getUuid())))
+                  .returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
+                  .build())
+          .build();
+
+      final TransactWriteItem accountPut = TransactWriteItem.builder()
+          .put(Put.builder()
+              .tableName(ACCOUNTS_TABLE_NAME)
+              .conditionExpression("attribute_not_exists(#number) OR #number = :number")
+              .expressionAttributeNames(Map.of("#number", Accounts.ATTR_ACCOUNT_E164))
+              .expressionAttributeValues(Map.of(":number", AttributeValues.fromString(account.getNumber())))
+              .item(Map.of(
+                  Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid),
+                  Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
+                  Accounts.ATTR_ACCOUNT_DATA, AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
+                  Accounts.ATTR_VERSION, AttributeValues.fromInt(account.getVersion()),
+                  Accounts.ATTR_CANONICALLY_DISCOVERABLE, AttributeValues.fromBool(account.shouldBeVisibleInDirectory())))
+              .build())
+          .build();
+
+      dynamoDbExtension.getDynamoDbClient().transactWriteItems(TransactWriteItemsRequest.builder()
+          .transactItems(phoneNumberConstraintPut, accountPut)
+          .build());
+    }
+
+    Optional<Account> retrieved = accounts.getByE164("+14151112222");
+
+    assertThat(retrieved.isPresent()).isTrue();
+    verifyStoredState("+14151112222", uuid, null, retrieved.get(), account);
+
+    retrieved = accounts.getByAccountIdentifier(uuid);
+
+    assertThat(retrieved.isPresent()).isTrue();
+    verifyStoredState("+14151112222", uuid, null, retrieved.get(), account);
   }
 
   @Test
   void testOverwrite() {
-    Device  device  = generateDevice (1                                            );
-    UUID    firstUuid = UUID.randomUUID();
-    Account account   = generateAccount("+14151112222", firstUuid, Collections.singleton(device));
+    Device device = generateDevice(1);
+    UUID firstUuid = UUID.randomUUID();
+    UUID firstPni = UUID.randomUUID();
+    Account account = generateAccount("+14151112222", firstUuid, firstPni, Collections.singleton(device));
 
     accounts.create(account);
 
-    verifyStoredState("+14151112222", account.getUuid(), account, true);
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier().orElseThrow(), account, true);
+
+    assertPhoneNumberConstraintExists("+14151112222", firstUuid);
+    assertPhoneNumberIdentifierConstraintExists(firstPni, firstUuid);
 
     account.setProfileName("name");
 
@@ -179,14 +284,17 @@ class AccountsTest {
     UUID secondUuid = UUID.randomUUID();
 
     device = generateDevice(1);
-    account = generateAccount("+14151112222", secondUuid, Collections.singleton(device));
+    account = generateAccount("+14151112222", secondUuid, UUID.randomUUID(), Collections.singleton(device));
 
     final boolean freshUser = accounts.create(account);
     assertThat(freshUser).isFalse();
-    verifyStoredState("+14151112222", firstUuid, account, true);
+    verifyStoredState("+14151112222", firstUuid, firstPni, account, true);
+
+    assertPhoneNumberConstraintExists("+14151112222", firstUuid);
+    assertPhoneNumberIdentifierConstraintExists(firstPni, firstUuid);
 
     device = generateDevice(1);
-    Account invalidAccount = generateAccount("+14151113333", firstUuid, Collections.singleton(device));
+    Account invalidAccount = generateAccount("+14151113333", firstUuid, UUID.randomUUID(), Collections.singleton(device));
 
     assertThatThrownBy(() -> accounts.create(invalidAccount));
   }
@@ -194,28 +302,34 @@ class AccountsTest {
   @Test
   void testUpdate() {
     Device  device  = generateDevice (1                                            );
-    Account account = generateAccount("+14151112222", UUID.randomUUID(), Collections.singleton(device));
+    Account account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), Collections.singleton(device));
 
     accounts.create(account);
 
+    assertPhoneNumberConstraintExists("+14151112222", account.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(account.getPhoneNumberIdentifier().orElseThrow(), account.getUuid());
+
     device.setName("foobar");
 
     accounts.update(account);
 
-    Optional<Account> retrieved = accounts.get("+14151112222");
+    assertPhoneNumberConstraintExists("+14151112222", account.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(account.getPhoneNumberIdentifier().orElseThrow(), account.getUuid());
+
+    Optional<Account> retrieved = accounts.getByE164("+14151112222");
 
     assertThat(retrieved.isPresent()).isTrue();
-    verifyStoredState("+14151112222", account.getUuid(), retrieved.get(), account);
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier().orElseThrow(), retrieved.get(), account);
 
-    retrieved = accounts.get(account.getUuid());
+    retrieved = accounts.getByAccountIdentifier(account.getUuid());
 
     assertThat(retrieved.isPresent()).isTrue();
-    verifyStoredState("+14151112222", account.getUuid(), account, true);
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier().orElseThrow(), account, true);
 
     device = generateDevice(1);
-    Account unknownAccount = generateAccount("+14151113333", UUID.randomUUID(), Collections.singleton(device));
+    Account unknownAccount = generateAccount("+14151113333", UUID.randomUUID(), UUID.randomUUID(), Collections.singleton(device));
 
-    assertThatThrownBy(() -> accounts.update(unknownAccount)).isInstanceOfAny(ConditionalCheckFailedException.class);
+    assertThatThrownBy(() -> accounts.update(unknownAccount)).isInstanceOfAny(TransactionCanceledException.class);
 
     account.setProfileName("name");
 
@@ -223,7 +337,7 @@ class AccountsTest {
 
     assertThat(account.getVersion()).isEqualTo(2);
 
-    verifyStoredState("+14151112222", account.getUuid(), account, true);
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier().orElseThrow(), account, true);
 
     account.setVersion(1);
 
@@ -234,7 +348,7 @@ class AccountsTest {
 
     accounts.update(account);
 
-    verifyStoredState("+14151112222", account.getUuid(), account, true);
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier().orElseThrow(), account, true);
   }
 
   @Test
@@ -242,23 +356,46 @@ class AccountsTest {
 
     final DynamoDbClient dynamoDbClient = mock(DynamoDbClient.class);
     accounts = new Accounts(dynamoDbClient,
-        dynamoDbExtension.getTableName(), NUMBERS_TABLE_NAME, SCAN_PAGE_SIZE);
+        dynamoDbExtension.getTableName(), NUMBER_CONSTRAINT_TABLE_NAME, PNI_CONSTRAINT_TABLE_NAME, SCAN_PAGE_SIZE);
 
-    when(dynamoDbClient.updateItem(any(UpdateItemRequest.class)))
+    when(dynamoDbClient.transactWriteItems(any(TransactWriteItemsRequest.class)))
         .thenThrow(TransactionConflictException.class);
 
     Device device = generateDevice(1);
-    Account account = generateAccount("+14151112222", UUID.randomUUID(), Collections.singleton(device));
+    Account account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), Collections.singleton(device));
 
     assertThatThrownBy(() -> accounts.update(account)).isInstanceOfAny(ContestedOptimisticLockException.class);
   }
 
+  @Test
+  // TODO Remove after initial migration is complete
+  void testUpdateWithMockTransactionCancellationException() {
+
+    final DynamoDbClient dynamoDbClient = mock(DynamoDbClient.class);
+    accounts = new Accounts(dynamoDbClient,
+        dynamoDbExtension.getTableName(), NUMBER_CONSTRAINT_TABLE_NAME, PNI_CONSTRAINT_TABLE_NAME, SCAN_PAGE_SIZE);
+
+    when(dynamoDbClient.transactWriteItems(any(TransactWriteItemsRequest.class)))
+        .thenThrow(TransactionCanceledException.builder()
+            .cancellationReasons(CancellationReason.builder()
+                .code("Test")
+                .build())
+            .build());
+
+    when(dynamoDbClient.getItem(any(GetItemRequest.class))).thenReturn(GetItemResponse.builder().build());
+
+    Device device = generateDevice(1);
+    Account account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), Collections.singleton(device));
+
+    assertThatThrownBy(() -> accounts.update(account)).isInstanceOfAny(TransactionCanceledException.class);
+  }
+
   @Test
   void testRetrieveFrom() {
     List<Account> users = new ArrayList<>();
 
     for (int i = 1; i <= 100; i++) {
-      Account account = generateAccount("+1" + String.format("%03d", i), UUID.randomUUID());
+      Account account = generateAccount("+1" + String.format("%03d", i), UUID.randomUUID(), UUID.randomUUID());
       users.add(account);
       accounts.create(account);
     }
@@ -276,7 +413,7 @@ class AccountsTest {
           .findAny()
           .orElseThrow();
 
-      verifyStoredState(expectedAccount.getNumber(), expectedAccount.getUuid(), retrievedAccount, expectedAccount);
+      verifyStoredState(expectedAccount.getNumber(), expectedAccount.getUuid(), expectedAccount.getPhoneNumberIdentifier().orElseThrow(), retrievedAccount, expectedAccount);
 
       users.remove(expectedAccount);
     }
@@ -293,7 +430,7 @@ class AccountsTest {
             .findAny()
             .orElseThrow();
 
-        verifyStoredState(expectedAccount.getNumber(), expectedAccount.getUuid(), retrievedAccount, expectedAccount);
+        verifyStoredState(expectedAccount.getNumber(), expectedAccount.getUuid(), expectedAccount.getPhoneNumberIdentifier().orElseThrow(), retrievedAccount, expectedAccount);
 
         users.remove(expectedAccount);
       }
@@ -306,48 +443,59 @@ class AccountsTest {
   void testDelete() {
     final Device deletedDevice = generateDevice(1);
     final Account deletedAccount = generateAccount("+14151112222", UUID.randomUUID(),
-        Collections.singleton(deletedDevice));
+        UUID.randomUUID(), Collections.singleton(deletedDevice));
     final Device retainedDevice = generateDevice(1);
     final Account retainedAccount = generateAccount("+14151112345", UUID.randomUUID(),
-        Collections.singleton(retainedDevice));
+        UUID.randomUUID(), Collections.singleton(retainedDevice));
 
     accounts.create(deletedAccount);
     accounts.create(retainedAccount);
 
-    assertThat(accounts.get(deletedAccount.getUuid())).isPresent();
-    assertThat(accounts.get(retainedAccount.getUuid())).isPresent();
+    assertPhoneNumberConstraintExists("+14151112222", deletedAccount.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(deletedAccount.getPhoneNumberIdentifier().orElseThrow(), deletedAccount.getUuid());
+    assertPhoneNumberConstraintExists("+14151112345", retainedAccount.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(retainedAccount.getPhoneNumberIdentifier().orElseThrow(), retainedAccount.getUuid());
+
+    assertThat(accounts.getByAccountIdentifier(deletedAccount.getUuid())).isPresent();
+    assertThat(accounts.getByAccountIdentifier(retainedAccount.getUuid())).isPresent();
 
     accounts.delete(deletedAccount.getUuid());
 
-    assertThat(accounts.get(deletedAccount.getUuid())).isNotPresent();
+    assertThat(accounts.getByAccountIdentifier(deletedAccount.getUuid())).isNotPresent();
 
-    verifyStoredState(retainedAccount.getNumber(), retainedAccount.getUuid(),
-        accounts.get(retainedAccount.getUuid()).get(), retainedAccount);
+    assertPhoneNumberConstraintDoesNotExist(deletedAccount.getNumber());
+    assertPhoneNumberIdentifierConstraintDoesNotExist(deletedAccount.getPhoneNumberIdentifier().orElseThrow());
+
+    verifyStoredState(retainedAccount.getNumber(), retainedAccount.getUuid(), retainedAccount.getPhoneNumberIdentifier().orElseThrow(),
+        accounts.getByAccountIdentifier(retainedAccount.getUuid()).get(), retainedAccount);
 
     {
       final Account recreatedAccount = generateAccount(deletedAccount.getNumber(), UUID.randomUUID(),
-          Collections.singleton(generateDevice(1)));
+          UUID.randomUUID(), Collections.singleton(generateDevice(1)));
 
       final boolean freshUser = accounts.create(recreatedAccount);
 
       assertThat(freshUser).isTrue();
-      assertThat(accounts.get(recreatedAccount.getUuid())).isPresent();
-      verifyStoredState(recreatedAccount.getNumber(), recreatedAccount.getUuid(),
-          accounts.get(recreatedAccount.getUuid()).get(), recreatedAccount);
+      assertThat(accounts.getByAccountIdentifier(recreatedAccount.getUuid())).isPresent();
+      verifyStoredState(recreatedAccount.getNumber(), recreatedAccount.getUuid(), recreatedAccount.getPhoneNumberIdentifier().orElseThrow(),
+          accounts.getByAccountIdentifier(recreatedAccount.getUuid()).get(), recreatedAccount);
+
+      assertPhoneNumberConstraintExists(recreatedAccount.getNumber(), recreatedAccount.getUuid());
+      assertPhoneNumberIdentifierConstraintExists(recreatedAccount.getPhoneNumberIdentifier().orElseThrow(), recreatedAccount.getUuid());
     }
   }
 
   @Test
   void testMissing() {
     Device  device  = generateDevice (1                                            );
-    Account account = generateAccount("+14151112222", UUID.randomUUID(), Collections.singleton(device));
+    Account account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), Collections.singleton(device));
 
     accounts.create(account);
 
-    Optional<Account> retrieved = accounts.get("+11111111");
+    Optional<Account> retrieved = accounts.getByE164("+11111111");
     assertThat(retrieved.isPresent()).isFalse();
 
-    retrieved = accounts.get(UUID.randomUUID());
+    retrieved = accounts.getByAccountIdentifier(UUID.randomUUID());
     assertThat(retrieved.isPresent()).isFalse();
   }
 
@@ -369,8 +517,9 @@ class AccountsTest {
     when(client.updateItem(any(UpdateItemRequest.class)))
         .thenThrow(RuntimeException.class);
 
-    Accounts accounts = new Accounts(client, ACCOUNTS_TABLE_NAME, NUMBERS_TABLE_NAME, SCAN_PAGE_SIZE);
-    Account account = generateAccount("+14151112222", UUID.randomUUID());
+    Accounts accounts = new Accounts(client, ACCOUNTS_TABLE_NAME, NUMBER_CONSTRAINT_TABLE_NAME,
+        PNI_CONSTRAINT_TABLE_NAME, SCAN_PAGE_SIZE);
+    Account account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID());
 
     try {
       accounts.update(account);
@@ -406,17 +555,16 @@ class AccountsTest {
   @Test
   void testCanonicallyDiscoverableSet() {
     Device device = generateDevice(1);
-    UUID uuid = UUID.randomUUID();
-    Account account = generateAccount("+14151112222", uuid, Collections.singleton(device));
+    Account account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), Collections.singleton(device));
     account.setDiscoverableByPhoneNumber(false);
     accounts.create(account);
-    verifyStoredState("+14151112222", account.getUuid(), account, false);
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier().orElseThrow(), account, false);
     account.setDiscoverableByPhoneNumber(true);
     accounts.update(account);
-    verifyStoredState("+14151112222", account.getUuid(), account, true);
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier().orElseThrow(), account, true);
     account.setDiscoverableByPhoneNumber(false);
     accounts.update(account);
-    verifyStoredState("+14151112222", account.getUuid(), account, false);
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier().orElseThrow(), account, false);
   }
 
   @Test
@@ -424,27 +572,45 @@ class AccountsTest {
     final String originalNumber = "+14151112222";
     final String targetNumber = "+14151113333";
 
+    final UUID originalPni = UUID.randomUUID();
+    final UUID targetPni = UUID.randomUUID();
+
     final Device device = generateDevice(1);
-    final Account account = generateAccount(originalNumber, UUID.randomUUID(), Collections.singleton(device));
+    final Account account = generateAccount(originalNumber, UUID.randomUUID(), originalPni, Collections.singleton(device));
 
     accounts.create(account);
 
+    assertThat(accounts.getByPhoneNumberIdentifier(originalPni)).isPresent();
+
+    assertPhoneNumberConstraintExists(originalNumber, account.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(originalPni, account.getUuid());
+
     {
-      final Optional<Account> retrieved = accounts.get(originalNumber);
+      final Optional<Account> retrieved = accounts.getByE164(originalNumber);
       assertThat(retrieved).isPresent();
 
-      verifyStoredState(originalNumber, account.getUuid(), retrieved.get(), account);
+      verifyStoredState(originalNumber, account.getUuid(), account.getPhoneNumberIdentifier().orElseThrow(), retrieved.get(), account);
     }
 
-    accounts.changeNumber(account, targetNumber);
+    accounts.changeNumber(account, targetNumber, targetPni);
 
-    assertThat(accounts.get(originalNumber)).isEmpty();
+    assertThat(accounts.getByE164(originalNumber)).isEmpty();
+    assertThat(accounts.getByAccountIdentifier(originalPni)).isEmpty();
+
+    assertPhoneNumberConstraintDoesNotExist(originalNumber);
+    assertPhoneNumberIdentifierConstraintDoesNotExist(originalPni);
+    assertPhoneNumberConstraintExists(targetNumber, account.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(targetPni, account.getUuid());
 
     {
-      final Optional<Account> retrieved = accounts.get(targetNumber);
+      final Optional<Account> retrieved = accounts.getByE164(targetNumber);
       assertThat(retrieved).isPresent();
 
-      verifyStoredState(targetNumber, account.getUuid(), retrieved.get(), account);
+      verifyStoredState(targetNumber, account.getUuid(), account.getPhoneNumberIdentifier().orElseThrow(), retrieved.get(), account);
+
+      assertThat(retrieved.get().getPhoneNumberIdentifier()).isPresent();
+      assertThat(retrieved.get().getPhoneNumberIdentifier().get()).isEqualTo(targetPni);
+      assertThat(accounts.getByPhoneNumberIdentifier(targetPni)).isPresent();
     }
   }
 
@@ -453,16 +619,267 @@ class AccountsTest {
     final String originalNumber = "+14151112222";
     final String targetNumber = "+14151113333";
 
+    final UUID originalPni = UUID.randomUUID();
+    final UUID targetPni = UUID.randomUUID();
+
     final Device existingDevice = generateDevice(1);
-    final Account existingAccount = generateAccount(targetNumber, UUID.randomUUID(), Collections.singleton(existingDevice));
+    final Account existingAccount = generateAccount(targetNumber, UUID.randomUUID(), targetPni, Collections.singleton(existingDevice));
 
     final Device device = generateDevice(1);
-    final Account account = generateAccount(originalNumber, UUID.randomUUID(), Collections.singleton(device));
+    final Account account = generateAccount(originalNumber, UUID.randomUUID(), originalPni, Collections.singleton(device));
 
     accounts.create(account);
     accounts.create(existingAccount);
 
-    assertThrows(TransactionCanceledException.class, () -> accounts.changeNumber(account, targetNumber));
+    assertThrows(TransactionCanceledException.class, () -> accounts.changeNumber(account, targetNumber, targetPni));
+
+    assertPhoneNumberConstraintExists(originalNumber, account.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(originalPni, account.getUuid());
+    assertPhoneNumberConstraintExists(targetNumber, existingAccount.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(targetPni, existingAccount.getUuid());
+  }
+
+  @Test
+  public void testChangeNumberPhoneNumberIdentifierConflict() {
+    final String originalNumber = "+14151112222";
+    final String targetNumber = "+14151113333";
+
+    final Device device = generateDevice(1);
+    final Account account = generateAccount(originalNumber, UUID.randomUUID(), UUID.randomUUID(), Collections.singleton(device));
+
+    accounts.create(account);
+
+    final UUID existingAccountIdentifier = UUID.randomUUID();
+    final UUID existingPhoneNumberIdentifier = UUID.randomUUID();
+
+    // Artificially inject a conflicting PNI entry
+    dynamoDbExtension.getDynamoDbClient().putItem(PutItemRequest.builder()
+        .tableName(PNI_CONSTRAINT_TABLE_NAME)
+        .item(Map.of(
+            Accounts.ATTR_PNI_UUID, AttributeValues.fromUUID(existingPhoneNumberIdentifier),
+            Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(existingAccountIdentifier)))
+        .conditionExpression(
+            "attribute_not_exists(#pni) OR (attribute_exists(#pni) AND #uuid = :uuid)")
+        .expressionAttributeNames(
+            Map.of("#uuid", Accounts.KEY_ACCOUNT_UUID,
+                "#pni", Accounts.ATTR_PNI_UUID))
+        .expressionAttributeValues(
+            Map.of(":uuid", AttributeValues.fromUUID(existingAccountIdentifier)))
+        .build());
+
+    assertThrows(TransactionCanceledException.class, () -> accounts.changeNumber(account, targetNumber, existingPhoneNumberIdentifier));
+  }
+
+  @Test
+  // TODO Remove or adapt after initial PNI migration
+  void testReregistrationFromAccountWithoutPhoneNumberIdentifier() throws JsonProcessingException {
+    final String number = "+18005551234";
+    final UUID originalUuid = UUID.randomUUID();
+
+    // Artificially inject Dynamo items for a legacy account without an assigned PNI
+    {
+      final Account account = generateAccount(number, originalUuid, null);
+
+      final TransactWriteItem phoneNumberConstraintPut = TransactWriteItem.builder()
+          .put(
+              Put.builder()
+                  .tableName(NUMBER_CONSTRAINT_TABLE_NAME)
+                  .item(Map.of(
+                      Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
+                      Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
+                  .conditionExpression(
+                      "attribute_not_exists(#number) OR (attribute_exists(#number) AND #uuid = :uuid)")
+                  .expressionAttributeNames(
+                      Map.of("#uuid", Accounts.KEY_ACCOUNT_UUID,
+                          "#number", Accounts.ATTR_ACCOUNT_E164))
+                  .expressionAttributeValues(
+                      Map.of(":uuid", AttributeValues.fromUUID(account.getUuid())))
+                  .returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
+                  .build())
+          .build();
+
+      final Map<String, AttributeValue> item = new HashMap<>(Map.of(
+          Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid()),
+          Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
+          Accounts.ATTR_ACCOUNT_DATA,
+          AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
+          Accounts.ATTR_VERSION, AttributeValues.fromInt(account.getVersion()),
+          Accounts.ATTR_CANONICALLY_DISCOVERABLE, AttributeValues.fromBool(account.shouldBeVisibleInDirectory())));
+
+      final TransactWriteItem accountPut = TransactWriteItem.builder()
+          .put(Put.builder()
+              .conditionExpression("attribute_not_exists(#number) OR #number = :number")
+              .expressionAttributeNames(Map.of("#number", Accounts.ATTR_ACCOUNT_E164))
+              .expressionAttributeValues(Map.of(":number", AttributeValues.fromString(account.getNumber())))
+              .tableName(ACCOUNTS_TABLE_NAME)
+              .item(item)
+              .build())
+          .build();
+
+      dynamoDbExtension.getDynamoDbClient().transactWriteItems(TransactWriteItemsRequest.builder()
+          .transactItems(phoneNumberConstraintPut, accountPut)
+          .build());
+    }
+
+    final Account reregisteredAccount = generateAccount(number, UUID.randomUUID(), UUID.randomUUID());
+    accounts.create(reregisteredAccount);
+
+    assertPhoneNumberConstraintExists(number, originalUuid);
+    assertPhoneNumberIdentifierConstraintExists(reregisteredAccount.getPhoneNumberIdentifier().orElseThrow(), originalUuid);
+  }
+
+  @Test
+  // TODO Remove or adapt after initial PNI migration
+  void testUpdateAccountAddingPniWithoutPhoneNumberIdentifier() throws JsonProcessingException {
+    final String number = "+18005551234";
+    final UUID uuid = UUID.randomUUID();
+
+    // Artificially inject Dynamo items for a legacy account without an assigned PNI
+    {
+      final Account account = generateAccount(number, uuid, null);
+
+      final TransactWriteItem phoneNumberConstraintPut = TransactWriteItem.builder()
+          .put(
+              Put.builder()
+                  .tableName(NUMBER_CONSTRAINT_TABLE_NAME)
+                  .item(Map.of(
+                      Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
+                      Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
+                  .conditionExpression(
+                      "attribute_not_exists(#number) OR (attribute_exists(#number) AND #uuid = :uuid)")
+                  .expressionAttributeNames(
+                      Map.of("#uuid", Accounts.KEY_ACCOUNT_UUID,
+                          "#number", Accounts.ATTR_ACCOUNT_E164))
+                  .expressionAttributeValues(
+                      Map.of(":uuid", AttributeValues.fromUUID(account.getUuid())))
+                  .returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
+                  .build())
+          .build();
+
+      final Map<String, AttributeValue> item = new HashMap<>(Map.of(
+          Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid()),
+          Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
+          Accounts.ATTR_ACCOUNT_DATA,
+          AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
+          Accounts.ATTR_VERSION, AttributeValues.fromInt(account.getVersion()),
+          Accounts.ATTR_CANONICALLY_DISCOVERABLE, AttributeValues.fromBool(account.shouldBeVisibleInDirectory())));
+
+      final TransactWriteItem accountPut = TransactWriteItem.builder()
+          .put(Put.builder()
+              .conditionExpression("attribute_not_exists(#number) OR #number = :number")
+              .expressionAttributeNames(Map.of("#number", Accounts.ATTR_ACCOUNT_E164))
+              .expressionAttributeValues(Map.of(":number", AttributeValues.fromString(account.getNumber())))
+              .tableName(ACCOUNTS_TABLE_NAME)
+              .item(item)
+              .build())
+          .build();
+
+      dynamoDbExtension.getDynamoDbClient().transactWriteItems(TransactWriteItemsRequest.builder()
+          .transactItems(phoneNumberConstraintPut, accountPut)
+          .build());
+    }
+
+    assertThat(accounts.getByAccountIdentifier(uuid)).hasValueSatisfying(account -> {
+      assertThat(account.getUuid()).isEqualTo(uuid);
+      assertThat(account.getNumber()).isEqualTo(number);
+      assertThat(account.getPhoneNumberIdentifier()).isEmpty();
+    });
+
+    final UUID phoneNumberIdentifier = UUID.randomUUID();
+
+    {
+      final Account accountToUpdate = accounts.getByAccountIdentifier(uuid).orElseThrow();
+      accountToUpdate.setNumber(number, phoneNumberIdentifier);
+
+      assertThat(accountToUpdate.getPhoneNumberIdentifier()).hasValueSatisfying(pni ->
+          assertThat(pni).isEqualTo(phoneNumberIdentifier));
+
+      accounts.update(accountToUpdate);
+
+      assertThat(accountToUpdate.getPhoneNumberIdentifier()).hasValueSatisfying(pni ->
+          assertThat(pni).isEqualTo(phoneNumberIdentifier));
+    }
+
+    assertThat(accounts.getByAccountIdentifier(uuid)).hasValueSatisfying(account -> {
+      assertThat(account.getUuid()).isEqualTo(uuid);
+      assertThat(account.getNumber()).isEqualTo(number);
+      assertThat(account.getPhoneNumberIdentifier()).hasValueSatisfying(pni ->
+          assertThat(pni).isEqualTo(phoneNumberIdentifier));
+    });
+  }
+
+  @Test
+    // TODO Remove or adapt after initial PNI migration
+  void testUpdateAccountWithoutPhoneNumberIdentifier() throws JsonProcessingException {
+    final String number = "+18005551234";
+    final UUID uuid = UUID.randomUUID();
+
+    // Artificially inject Dynamo items for a legacy account without an assigned PNI
+    {
+      final Account account = generateAccount(number, uuid, null);
+
+      final TransactWriteItem phoneNumberConstraintPut = TransactWriteItem.builder()
+          .put(
+              Put.builder()
+                  .tableName(NUMBER_CONSTRAINT_TABLE_NAME)
+                  .item(Map.of(
+                      Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
+                      Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
+                  .conditionExpression(
+                      "attribute_not_exists(#number) OR (attribute_exists(#number) AND #uuid = :uuid)")
+                  .expressionAttributeNames(
+                      Map.of("#uuid", Accounts.KEY_ACCOUNT_UUID,
+                          "#number", Accounts.ATTR_ACCOUNT_E164))
+                  .expressionAttributeValues(
+                      Map.of(":uuid", AttributeValues.fromUUID(account.getUuid())))
+                  .returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
+                  .build())
+          .build();
+
+      final Map<String, AttributeValue> item = new HashMap<>(Map.of(
+          Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid()),
+          Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
+          Accounts.ATTR_ACCOUNT_DATA,
+          AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
+          Accounts.ATTR_VERSION, AttributeValues.fromInt(account.getVersion()),
+          Accounts.ATTR_CANONICALLY_DISCOVERABLE, AttributeValues.fromBool(account.shouldBeVisibleInDirectory())));
+
+      final TransactWriteItem accountPut = TransactWriteItem.builder()
+          .put(Put.builder()
+              .conditionExpression("attribute_not_exists(#number) OR #number = :number")
+              .expressionAttributeNames(Map.of("#number", Accounts.ATTR_ACCOUNT_E164))
+              .expressionAttributeValues(Map.of(":number", AttributeValues.fromString(account.getNumber())))
+              .tableName(ACCOUNTS_TABLE_NAME)
+              .item(item)
+              .build())
+          .build();
+
+      dynamoDbExtension.getDynamoDbClient().transactWriteItems(TransactWriteItemsRequest.builder()
+          .transactItems(phoneNumberConstraintPut, accountPut)
+          .build());
+    }
+
+    assertThat(accounts.getByAccountIdentifier(uuid)).hasValueSatisfying(account -> {
+      assertThat(account.getUuid()).isEqualTo(uuid);
+      assertThat(account.getNumber()).isEqualTo(number);
+      assertThat(account.getPhoneNumberIdentifier()).isEmpty();
+    });
+
+    final String updatedName = "An updated name!";
+
+    {
+      final Account accountToUpdate = accounts.getByAccountIdentifier(uuid).orElseThrow();
+      accountToUpdate.setProfileName(updatedName);
+
+      accounts.update(accountToUpdate);
+    }
+
+    assertThat(accounts.getByAccountIdentifier(uuid)).hasValueSatisfying(account -> {
+      assertThat(account.getUuid()).isEqualTo(uuid);
+      assertThat(account.getNumber()).isEqualTo(number);
+      assertThat(account.getPhoneNumberIdentifier()).isEmpty();
+      assertThat(account.getProfileName()).isEqualTo(updatedName);
+    });
   }
 
   private Device generateDevice(long id) {
@@ -473,20 +890,62 @@ class AccountsTest {
         random.nextBoolean(), random.nextBoolean(), random.nextBoolean()));
   }
 
-  private Account generateAccount(String number, UUID uuid) {
+  private Account generateAccount(String number, UUID uuid, final UUID pni) {
     Device device = generateDevice(1);
-    return generateAccount(number, uuid, Collections.singleton(device));
+    return generateAccount(number, uuid, pni, Collections.singleton(device));
   }
 
-  private Account generateAccount(String number, UUID uuid, Set<Device> devices) {
+  private Account generateAccount(String number, UUID uuid, final UUID pni, Set<Device> devices) {
     byte[]       unidentifiedAccessKey = new byte[16];
     Random random = new Random(System.currentTimeMillis());
     Arrays.fill(unidentifiedAccessKey, (byte)random.nextInt(255));
 
-    return new Account(number, uuid, devices, unidentifiedAccessKey);
+    return new Account(number, uuid, pni, devices, unidentifiedAccessKey);
   }
 
-  private void verifyStoredState(String number, UUID uuid, Account expecting, boolean canonicallyDiscoverable) {
+  private void assertPhoneNumberConstraintExists(final String number, final UUID uuid) {
+    final GetItemResponse numberConstraintResponse = dynamoDbExtension.getDynamoDbClient().getItem(
+        GetItemRequest.builder()
+            .tableName(NUMBER_CONSTRAINT_TABLE_NAME)
+            .key(Map.of(Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(number)))
+            .build());
+
+    assertThat(numberConstraintResponse.hasItem()).isTrue();
+    assertThat(AttributeValues.getUUID(numberConstraintResponse.item(), Accounts.KEY_ACCOUNT_UUID, null)).isEqualTo(uuid);
+  }
+
+  private void assertPhoneNumberConstraintDoesNotExist(final String number) {
+    final GetItemResponse numberConstraintResponse = dynamoDbExtension.getDynamoDbClient().getItem(
+        GetItemRequest.builder()
+            .tableName(NUMBER_CONSTRAINT_TABLE_NAME)
+            .key(Map.of(Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(number)))
+            .build());
+
+    assertThat(numberConstraintResponse.hasItem()).isFalse();
+  }
+
+  private void assertPhoneNumberIdentifierConstraintExists(final UUID phoneNumberIdentifier, final UUID uuid) {
+    final GetItemResponse pniConstraintResponse = dynamoDbExtension.getDynamoDbClient().getItem(
+        GetItemRequest.builder()
+            .tableName(PNI_CONSTRAINT_TABLE_NAME)
+            .key(Map.of(Accounts.ATTR_PNI_UUID, AttributeValues.fromUUID(phoneNumberIdentifier)))
+            .build());
+
+    assertThat(pniConstraintResponse.hasItem()).isTrue();
+    assertThat(AttributeValues.getUUID(pniConstraintResponse.item(), Accounts.KEY_ACCOUNT_UUID, null)).isEqualTo(uuid);
+  }
+
+  private void assertPhoneNumberIdentifierConstraintDoesNotExist(final UUID phoneNumberIdentifier) {
+    final GetItemResponse pniConstraintResponse = dynamoDbExtension.getDynamoDbClient().getItem(
+        GetItemRequest.builder()
+            .tableName(PNI_CONSTRAINT_TABLE_NAME)
+            .key(Map.of(Accounts.ATTR_PNI_UUID, AttributeValues.fromUUID(phoneNumberIdentifier)))
+            .build());
+
+    assertThat(pniConstraintResponse.hasItem()).isFalse();
+  }
+
+  private void verifyStoredState(String number, UUID uuid, UUID pni, Account expecting, boolean canonicallyDiscoverable) {
     final DynamoDbClient db = dynamoDbExtension.getDynamoDbClient();
 
     final GetItemResponse get = db.getItem(GetItemRequest.builder()
@@ -506,14 +965,15 @@ class AccountsTest {
           !canonicallyDiscoverable)).isEqualTo(canonicallyDiscoverable);
 
       Account result = Accounts.fromItem(get.item());
-      verifyStoredState(number, uuid, result, expecting);
+      verifyStoredState(number, uuid, pni, result, expecting);
     } else {
       throw new AssertionError("No data");
     }
   }
 
-  private void verifyStoredState(String number, UUID uuid, Account result, Account expecting) {
+  private void verifyStoredState(String number, UUID uuid, UUID pni, Account result, Account expecting) {
     assertThat(result.getNumber()).isEqualTo(number);
+    assertThat(result.getPhoneNumberIdentifier()).isEqualTo(Optional.ofNullable(pni));
     assertThat(result.getLastSeen()).isEqualTo(expecting.getLastSeen());
     assertThat(result.getUuid()).isEqualTo(uuid);
     assertThat(result.getVersion()).isEqualTo(expecting.getVersion());

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

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2013-2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.UUID;
+import java.util.function.Consumer;
+import org.junit.jupiter.api.Test;
+import org.mockito.stubbing.Answer;
+
+class AssignPhoneNumberIdentifierCrawlerListenerTest {
+
+  @Test
+  void onCrawlChunk() {
+    final UUID accountIdentifierWithPni = UUID.randomUUID();
+    final UUID accountIdentifierWithoutPni = UUID.randomUUID();
+
+    final String numberWithPni = "+18005551111";
+    final String numberWithoutPni = "+18005552222";
+
+    final Account accountWithPni = mock(Account.class);
+    when(accountWithPni.getUuid()).thenReturn(accountIdentifierWithPni);
+    when(accountWithPni.getNumber()).thenReturn(numberWithPni);
+    when(accountWithPni.getPhoneNumberIdentifier()).thenReturn(Optional.of(UUID.randomUUID()));
+
+    final Account accountWithoutPni = mock(Account.class);
+    when(accountWithoutPni.getUuid()).thenReturn(accountIdentifierWithoutPni);
+    when(accountWithoutPni.getNumber()).thenReturn(numberWithoutPni);
+    when(accountWithoutPni.getPhoneNumberIdentifier()).thenReturn(Optional.empty());
+
+    final AccountsManager accountsManager = mock(AccountsManager.class);
+    when(accountsManager.getByAccountIdentifier(accountIdentifierWithPni)).thenReturn(Optional.of(accountWithPni));
+    when(accountsManager.getByAccountIdentifier(accountIdentifierWithoutPni)).thenReturn(Optional.of(accountWithoutPni));
+
+    when(accountsManager.update(any(), any())).thenAnswer((Answer<Account>) invocation -> {
+      final Account account = invocation.getArgument(0, Account.class);
+      final Consumer<Account> updater = invocation.getArgument(1, Consumer.class);
+
+      updater.accept(account);
+
+      return account;
+    });
+
+    final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
+    when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString())).thenAnswer(
+        (Answer<UUID>) invocation -> UUID.randomUUID());
+
+    final AssignPhoneNumberIdentifierCrawlerListener crawler =
+        new AssignPhoneNumberIdentifierCrawlerListener(accountsManager, phoneNumberIdentifiers);
+
+    crawler.onCrawlChunk(Optional.empty(), List.of(accountWithPni, accountWithoutPni));
+
+    verify(accountsManager).update(eq(accountWithoutPni), any());
+    verify(accountWithoutPni).setNumber(eq(numberWithoutPni), any());
+
+    verify(accountsManager, never()).update(eq(accountWithPni), any());
+    verify(accountWithPni, never()).setNumber(any(), any());
+  }
+}

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

@@ -30,7 +30,7 @@ public class ContactDiscoveryWriterTest {
     when(acct.getUuid()).thenReturn(uuid);
     when(acct.isCanonicallyDiscoverable()).thenReturn(canonicallyDiscoverable);
     when(acct.shouldBeVisibleInDirectory()).thenReturn(shouldBeVisible);
-    when(mgr.get(uuid)).thenReturn(Optional.of(acct));
+    when(mgr.getByAccountIdentifier(uuid)).thenReturn(Optional.of(acct));
     ContactDiscoveryWriter writer = new ContactDiscoveryWriter(mgr);
     writer.onCrawlChunk(Optional.empty(), List.of(acct));
     verify(mgr, times(updateCalled ? 1 : 0)).update(acct, ContactDiscoveryWriter.NOOP_UPDATER);

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

@@ -82,7 +82,7 @@ class MessagePersisterIntegrationTest {
 
     when(account.getNumber()).thenReturn("+18005551234");
     when(account.getUuid()).thenReturn(accountUuid);
-    when(accountsManager.get(accountUuid)).thenReturn(Optional.of(account));
+    when(accountsManager.getByAccountIdentifier(accountUuid)).thenReturn(Optional.of(account));
     when(dynamicConfigurationManager.getConfiguration()).thenReturn(new DynamicConfiguration());
 
     messagesCache.start();

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

@@ -63,7 +63,7 @@ public class MessagePersisterTest extends AbstractRedisClusterTest {
 
         final Account account = mock(Account.class);
 
-        when(accountsManager.get(DESTINATION_ACCOUNT_UUID)).thenReturn(Optional.of(account));
+        when(accountsManager.getByAccountIdentifier(DESTINATION_ACCOUNT_UUID)).thenReturn(Optional.of(account));
         when(account.getNumber()).thenReturn(DESTINATION_ACCOUNT_NUMBER);
         when(dynamicConfigurationManager.getConfiguration()).thenReturn(new DynamicConfiguration());
 
@@ -98,7 +98,7 @@ public class MessagePersisterTest extends AbstractRedisClusterTest {
     public void testPersistNextQueuesNoQueues() {
         messagePersister.persistNextQueues(Instant.now());
 
-        verify(accountsManager, never()).get(any(UUID.class));
+        verify(accountsManager, never()).getByAccountIdentifier(any(UUID.class));
     }
 
     @Test
@@ -147,7 +147,7 @@ public class MessagePersisterTest extends AbstractRedisClusterTest {
 
             final Account account = mock(Account.class);
 
-            when(accountsManager.get(accountUuid)).thenReturn(Optional.of(account));
+            when(accountsManager.getByAccountIdentifier(accountUuid)).thenReturn(Optional.of(account));
             when(account.getNumber()).thenReturn(accountNumber);
 
             insertMessages(accountUuid, deviceId, messagesPerQueue, now);

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

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2013-2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+
+import java.util.UUID;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
+import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
+
+class PhoneNumberIdentifiersTest {
+
+  private static final String PNI_TABLE_NAME = "pni_test";
+
+  @RegisterExtension
+  static DynamoDbExtension DYNAMO_DB_EXTENSION = DynamoDbExtension.builder()
+      .tableName(PNI_TABLE_NAME)
+      .hashKey(PhoneNumberIdentifiers.KEY_E164)
+      .attributeDefinition(AttributeDefinition.builder()
+          .attributeName(PhoneNumberIdentifiers.KEY_E164)
+          .attributeType(ScalarAttributeType.S)
+          .build())
+      .build();
+
+  private PhoneNumberIdentifiers phoneNumberIdentifiers;
+
+  @BeforeEach
+  void setUp() {
+    phoneNumberIdentifiers = new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbClient(), PNI_TABLE_NAME);
+  }
+
+  @Test
+  void getPhoneNumberIdentifier() {
+    final String number = "+18005551234";
+    final String differentNumber = "+18005556789";
+
+    final UUID firstPni = phoneNumberIdentifiers.getPhoneNumberIdentifier(number);
+    final UUID secondPni = phoneNumberIdentifiers.getPhoneNumberIdentifier(number);
+
+    assertEquals(firstPni, secondPni);
+    assertNotEquals(firstPni, phoneNumberIdentifiers.getPhoneNumberIdentifier(differentNumber));
+  }
+
+  @Test
+  void generatePhoneNumberIdentifierIfNotExists() {
+    final String number = "+18005551234";
+
+    assertEquals(phoneNumberIdentifiers.generatePhoneNumberIdentifierIfNotExists(number),
+        phoneNumberIdentifiers.generatePhoneNumberIdentifierIfNotExists(number));
+  }
+}

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

@@ -34,7 +34,7 @@ class RefreshingAccountAndDeviceSupplierTest {
     when(initialDevice.getId()).thenReturn(deviceId);
     when(initialAccount.getDevice(deviceId)).thenReturn(Optional.of(initialDevice));
 
-    when(accountsManager.get(any(UUID.class))).thenAnswer(answer -> {
+    when(accountsManager.getByAccountIdentifier(any(UUID.class))).thenAnswer(answer -> {
       final Account account = mock(Account.class);
       final Device device = mock(Device.class);
 

service/src/test/java/org/whispersystems/textsecuregcm/tests/auth/ExternalServiceCredentialsGeneratorTest.java

@@ -5,17 +5,17 @@
 
 package org.whispersystems.textsecuregcm.tests.auth;
 
+import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
+
 import org.junit.Test;
 import org.whispersystems.textsecuregcm.auth.ExternalServiceCredentialGenerator;
 import org.whispersystems.textsecuregcm.auth.ExternalServiceCredentials;
 
-import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
-
 public class ExternalServiceCredentialsGeneratorTest {
 
   @Test
   public void testGenerateDerivedUsername() {
-    ExternalServiceCredentialGenerator generator = new ExternalServiceCredentialGenerator(new byte[32], new byte[32], true);
+    ExternalServiceCredentialGenerator generator = new ExternalServiceCredentialGenerator(new byte[32], new byte[32]);
     ExternalServiceCredentials credentials = generator.generateFor("+14152222222");
 
     assertThat(credentials.getUsername()).isNotEqualTo("+14152222222");

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

@@ -31,6 +31,7 @@ import java.security.SecureRandom;
 import java.time.Duration;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Locale;
 import java.util.Optional;
 import java.util.Set;
@@ -223,19 +224,20 @@ class AccountControllerTest {
     when(pendingAccountsManager.getCodeForNumber(SENDER_HAS_STORAGE)).thenReturn(Optional.of(new StoredVerificationCode("666666", System.currentTimeMillis(), null, null)));
     when(pendingAccountsManager.getCodeForNumber(SENDER_TRANSFER)).thenReturn(Optional.of(new StoredVerificationCode("1234", System.currentTimeMillis(), null, null)));
 
-    when(accountsManager.get(eq(SENDER_PIN))).thenReturn(Optional.of(senderPinAccount));
-    when(accountsManager.get(eq(SENDER_REG_LOCK))).thenReturn(Optional.of(senderRegLockAccount));
-    when(accountsManager.get(eq(SENDER_OVER_PIN))).thenReturn(Optional.of(senderPinAccount));
-    when(accountsManager.get(eq(SENDER))).thenReturn(Optional.empty());
-    when(accountsManager.get(eq(SENDER_OLD))).thenReturn(Optional.empty());
-    when(accountsManager.get(eq(SENDER_PREAUTH))).thenReturn(Optional.empty());
-    when(accountsManager.get(eq(SENDER_HAS_STORAGE))).thenReturn(Optional.of(senderHasStorage));
-    when(accountsManager.get(eq(SENDER_TRANSFER))).thenReturn(Optional.of(senderTransfer));
+    when(accountsManager.getByE164(eq(SENDER_PIN))).thenReturn(Optional.of(senderPinAccount));
+    when(accountsManager.getByE164(eq(SENDER_REG_LOCK))).thenReturn(Optional.of(senderRegLockAccount));
+    when(accountsManager.getByE164(eq(SENDER_OVER_PIN))).thenReturn(Optional.of(senderPinAccount));
+    when(accountsManager.getByE164(eq(SENDER))).thenReturn(Optional.empty());
+    when(accountsManager.getByE164(eq(SENDER_OLD))).thenReturn(Optional.empty());
+    when(accountsManager.getByE164(eq(SENDER_PREAUTH))).thenReturn(Optional.empty());
+    when(accountsManager.getByE164(eq(SENDER_HAS_STORAGE))).thenReturn(Optional.of(senderHasStorage));
+    when(accountsManager.getByE164(eq(SENDER_TRANSFER))).thenReturn(Optional.of(senderTransfer));
 
-    when(accountsManager.create(any(), any(), any(), any())).thenAnswer((Answer<Account>) invocation -> {
+    when(accountsManager.create(any(), any(), any(), any(), any())).thenAnswer((Answer<Account>) invocation -> {
       final Account account = mock(Account.class);
       when(account.getUuid()).thenReturn(UUID.randomUUID());
       when(account.getNumber()).thenReturn(invocation.getArgument(0, String.class));
+      when(account.getBadges()).thenReturn(invocation.getArgument(4, List.class));
 
       return account;
     });
@@ -1017,7 +1019,7 @@ class AccountControllerTest {
              .header("Authorization", AuthHelper.getProvisioningAuthHeader(SENDER, "bar"))
              .put(Entity.entity(new AccountAttributes(), MediaType.APPLICATION_JSON_TYPE), AccountCreationResult.class);
 
-    verify(accountsManager).create(eq(SENDER), eq("bar"), any(), any());
+    verify(accountsManager).create(eq(SENDER), eq("bar"), any(), any(), anyList());
 
     if (enrolledInVerifyExperiment) {
       verify(smsSender).reportVerificationSucceeded("VerificationSid");
@@ -1095,7 +1097,8 @@ class AccountControllerTest {
     verify(pinLimiter).validate(eq(SENDER_REG_LOCK));
 
     verify(accountsManager).create(eq(SENDER_REG_LOCK), eq("bar"), any(), argThat(
-        attributes -> Hex.toStringCondensed(registration_lock_key).equals(attributes.getRegistrationLock())));
+        attributes -> Hex.toStringCondensed(registration_lock_key).equals(attributes.getRegistrationLock())),
+        argThat(List::isEmpty));
   }
 
   @Test
@@ -1335,7 +1338,7 @@ class AccountControllerTest {
     when(existingAccount.getUuid()).thenReturn(UUID.randomUUID());
     when(existingAccount.getRegistrationLock()).thenReturn(existingRegistrationLock);
 
-    when(accountsManager.get(number)).thenReturn(Optional.of(existingAccount));
+    when(accountsManager.getByE164(number)).thenReturn(Optional.of(existingAccount));
 
     final Response response =
         resources.getJerseyTest()
@@ -1365,7 +1368,7 @@ class AccountControllerTest {
     when(existingAccount.getUuid()).thenReturn(UUID.randomUUID());
     when(existingAccount.getRegistrationLock()).thenReturn(existingRegistrationLock);
 
-    when(accountsManager.get(number)).thenReturn(Optional.of(existingAccount));
+    when(accountsManager.getByE164(number)).thenReturn(Optional.of(existingAccount));
 
     final Response response =
         resources.getJerseyTest()
@@ -1397,7 +1400,7 @@ class AccountControllerTest {
     when(existingAccount.getUuid()).thenReturn(UUID.randomUUID());
     when(existingAccount.getRegistrationLock()).thenReturn(existingRegistrationLock);
 
-    when(accountsManager.get(number)).thenReturn(Optional.of(existingAccount));
+    when(accountsManager.getByE164(number)).thenReturn(Optional.of(existingAccount));
 
     final Response response =
         resources.getJerseyTest()
@@ -1429,7 +1432,7 @@ class AccountControllerTest {
     when(existingAccount.getUuid()).thenReturn(UUID.randomUUID());
     when(existingAccount.getRegistrationLock()).thenReturn(existingRegistrationLock);
 
-    when(accountsManager.get(number)).thenReturn(Optional.of(existingAccount));
+    when(accountsManager.getByE164(number)).thenReturn(Optional.of(existingAccount));
 
     final Response response =
         resources.getJerseyTest()

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

@@ -7,7 +7,6 @@ package org.whispersystems.textsecuregcm.tests.controllers;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.mock;
@@ -132,8 +131,8 @@ class DeviceControllerTest {
     when(pendingDevicesManager.getCodeForNumber(AuthHelper.VALID_NUMBER)).thenReturn(
         Optional.of(new StoredVerificationCode("5678901", System.currentTimeMillis(), null, null)));
     when(pendingDevicesManager.getCodeForNumber(AuthHelper.VALID_NUMBER_TWO)).thenReturn(Optional.empty());
-    when(accountsManager.get(AuthHelper.VALID_NUMBER)).thenReturn(Optional.of(account));
-    when(accountsManager.get(AuthHelper.VALID_NUMBER_TWO)).thenReturn(Optional.of(maxedAccount));
+    when(accountsManager.getByE164(AuthHelper.VALID_NUMBER)).thenReturn(Optional.of(account));
+    when(accountsManager.getByE164(AuthHelper.VALID_NUMBER_TWO)).thenReturn(Optional.of(maxedAccount));
 
     AccountsHelper.setupMockUpdate(accountsManager);
   }
@@ -156,7 +155,7 @@ class DeviceControllerTest {
 
   @Test
   void validDeviceRegisterTest() {
-    when(accountsManager.get(AuthHelper.VALID_UUID)).thenReturn(Optional.of(AuthHelper.VALID_ACCOUNT));
+    when(accountsManager.getByAccountIdentifier(AuthHelper.VALID_UUID)).thenReturn(Optional.of(AuthHelper.VALID_ACCOUNT));
 
     final Device existingDevice = mock(Device.class);
     when(existingDevice.getId()).thenReturn(Device.MASTER_ID);

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

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.tests.controllers;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.time.Clock;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.util.UUID;
+import org.junit.jupiter.api.Test;
+import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
+import org.whispersystems.textsecuregcm.auth.ExternalServiceCredentialGenerator;
+import org.whispersystems.textsecuregcm.auth.ExternalServiceCredentials;
+import org.whispersystems.textsecuregcm.controllers.DirectoryV2Controller;
+import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.util.Pair;
+
+class DirectoryControllerV2Test {
+
+  @Test
+  void testAuthToken() {
+    final ExternalServiceCredentialGenerator credentialGenerator = new ExternalServiceCredentialGenerator(
+        new byte[]{0x1}, new byte[0], false, false,
+        Clock.fixed(Instant.ofEpochSecond(1633738643L), ZoneId.of("Etc/UTC")));
+
+    final DirectoryV2Controller controller = new DirectoryV2Controller(credentialGenerator);
+
+    final Account account = mock(Account.class);
+    final UUID uuid = UUID.fromString("11111111-1111-1111-1111-111111111111");
+    when(account.getUuid()).thenReturn(uuid);
+    when(account.getNumber()).thenReturn("+15055551111");
+
+    final ExternalServiceCredentials credentials = (ExternalServiceCredentials) controller.getAuthToken(
+        new AuthenticatedAccount(() -> new Pair<>(account, mock(Device.class)))).getEntity();
+
+    assertEquals(credentials.getUsername(), "EREREREREREREREREREREQAAAABZvPKn");
+    assertEquals(credentials.getPassword(), "1633738643:ff03669c64f3f938a279");
+    assertEquals(32, credentials.getUsername().length());
+    assertEquals(31, credentials.getPassword().length());
+  }
+
+}

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

@@ -210,7 +210,7 @@ class DonationControllerTest {
     when(receiptCredentialPresentation.getReceiptExpirationTime()).thenReturn(receiptExpiration);
     when(redeemedReceiptsManager.put(same(receiptSerial), eq(receiptExpiration), eq(receiptLevel), eq(AuthHelper.VALID_UUID))).thenReturn(
         CompletableFuture.completedFuture(Boolean.TRUE));
-    when(accountsManager.get(eq(AuthHelper.VALID_UUID))).thenReturn(Optional.of(AuthHelper.VALID_ACCOUNT));
+    when(accountsManager.getByAccountIdentifier(eq(AuthHelper.VALID_UUID))).thenReturn(Optional.of(AuthHelper.VALID_ACCOUNT));
 
     RedeemReceiptRequest request = new RedeemReceiptRequest(presentation, true, true);
     Response response = resources.getJerseyTest()

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

@@ -153,11 +153,11 @@ class KeysControllerTest {
     when(existsAccount.getNumber()).thenReturn(EXISTS_NUMBER);
     when(existsAccount.getUnidentifiedAccessKey()).thenReturn(Optional.of("1337".getBytes()));
 
-    when(accounts.get(EXISTS_NUMBER)).thenReturn(Optional.of(existsAccount));
-    when(accounts.get(EXISTS_UUID)).thenReturn(Optional.of(existsAccount));
+    when(accounts.getByE164(EXISTS_NUMBER)).thenReturn(Optional.of(existsAccount));
+    when(accounts.getByAccountIdentifier(EXISTS_UUID)).thenReturn(Optional.of(existsAccount));
 
-    when(accounts.get(NOT_EXISTS_NUMBER)).thenReturn(Optional.empty());
-    when(accounts.get(NOT_EXISTS_UUID)).thenReturn(Optional.empty());
+    when(accounts.getByE164(NOT_EXISTS_NUMBER)).thenReturn(Optional.empty());
+    when(accounts.getByAccountIdentifier(NOT_EXISTS_UUID)).thenReturn(Optional.empty());
 
     when(rateLimiters.getPreKeysLimiter()).thenReturn(rateLimiter);
 

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

@@ -160,13 +160,15 @@ class MessageControllerTest {
           false, false, false)));
     }};
 
-    Account singleDeviceAccount  = new Account(SINGLE_DEVICE_RECIPIENT, SINGLE_DEVICE_UUID, singleDeviceList, "1234".getBytes());
-    Account multiDeviceAccount   = new Account(MULTI_DEVICE_RECIPIENT, MULTI_DEVICE_UUID, multiDeviceList, "1234".getBytes());
-    internationalAccount = new Account(INTERNATIONAL_RECIPIENT, INTERNATIONAL_UUID, singleDeviceList, "1234".getBytes());
+    Account singleDeviceAccount  = new Account(SINGLE_DEVICE_RECIPIENT, SINGLE_DEVICE_UUID, UUID.randomUUID(),
+        singleDeviceList, "1234".getBytes());
+    Account multiDeviceAccount   = new Account(MULTI_DEVICE_RECIPIENT, MULTI_DEVICE_UUID, UUID.randomUUID(),
+        multiDeviceList, "1234".getBytes());
+    internationalAccount = new Account(INTERNATIONAL_RECIPIENT, INTERNATIONAL_UUID, UUID.randomUUID(), singleDeviceList, "1234".getBytes());
 
-    when(accountsManager.get(eq(SINGLE_DEVICE_UUID))).thenReturn(Optional.of(singleDeviceAccount));
-    when(accountsManager.get(eq(MULTI_DEVICE_UUID))).thenReturn(Optional.of(multiDeviceAccount));
-    when(accountsManager.get(INTERNATIONAL_UUID)).thenReturn(Optional.of(internationalAccount));
+    when(accountsManager.getByAccountIdentifier(eq(SINGLE_DEVICE_UUID))).thenReturn(Optional.of(singleDeviceAccount));
+    when(accountsManager.getByAccountIdentifier(eq(MULTI_DEVICE_UUID))).thenReturn(Optional.of(multiDeviceAccount));
+    when(accountsManager.getByAccountIdentifier(INTERNATIONAL_UUID)).thenReturn(Optional.of(internationalAccount));
 
     when(rateLimiters.getMessagesLimiter()).thenReturn(rateLimiter);
   }

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

@@ -28,7 +28,6 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.Set;
 import javax.ws.rs.client.Entity;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
@@ -71,29 +70,29 @@ import org.whispersystems.textsecuregcm.storage.VersionedProfile;
 import org.whispersystems.textsecuregcm.tests.util.AccountsHelper;
 import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
 import org.whispersystems.textsecuregcm.util.SystemMapper;
-import org.whispersystems.textsecuregcm.util.Util;
 import software.amazon.awssdk.services.s3.S3Client;
 import software.amazon.awssdk.services.s3.model.DeleteObjectRequest;
 
 @ExtendWith(DropwizardExtensionsSupport.class)
 class ProfileControllerTest {
 
-  private static Clock clock = mock(Clock.class);
-  private static AccountsManager  accountsManager     = mock(AccountsManager.class );
-  private static ProfilesManager  profilesManager     = mock(ProfilesManager.class);
-  private static UsernamesManager usernamesManager    = mock(UsernamesManager.class);
-  private static RateLimiters     rateLimiters        = mock(RateLimiters.class    );
-  private static RateLimiter      rateLimiter         = mock(RateLimiter.class     );
-  private static RateLimiter      usernameRateLimiter = mock(RateLimiter.class     );
+  private static final Clock clock = mock(Clock.class);
+  private static final AccountsManager accountsManager = mock(AccountsManager.class);
+  private static final ProfilesManager profilesManager = mock(ProfilesManager.class);
+  private static final UsernamesManager usernamesManager = mock(UsernamesManager.class);
+  private static final RateLimiters rateLimiters = mock(RateLimiters.class);
+  private static final RateLimiter rateLimiter = mock(RateLimiter.class);
+  private static final RateLimiter usernameRateLimiter = mock(RateLimiter.class);
 
-  private static S3Client                  s3client            = mock(S3Client.class);
-  private static PostPolicyGenerator       postPolicyGenerator = new PostPolicyGenerator("us-west-1", "profile-bucket", "accessKey");
-  private static PolicySigner              policySigner        = new PolicySigner("accessSecret", "us-west-1");
-  private static ServerZkProfileOperations zkProfileOperations = mock(ServerZkProfileOperations.class);
+  private static final S3Client s3client = mock(S3Client.class);
+  private static final PostPolicyGenerator postPolicyGenerator = new PostPolicyGenerator("us-west-1", "profile-bucket", "accessKey");
+  private static final PolicySigner policySigner = new PolicySigner("accessSecret", "us-west-1");
+  private static final ServerZkProfileOperations zkProfileOperations = mock(ServerZkProfileOperations.class);
+
+  @SuppressWarnings("unchecked")
+  private static final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager = mock(DynamicConfigurationManager.class);
 
-  private static DynamicConfigurationManager dynamicConfigurationManager = mock(DynamicConfigurationManager.class);
   private DynamicPaymentsConfiguration dynamicPaymentsConfiguration;
-
   private Account profileAccount;
 
   private static final ResourceExtension resources = ResourceExtension.builder()
@@ -138,7 +137,7 @@ class ProfileControllerTest {
 
     when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration);
     when(dynamicConfiguration.getPaymentsConfiguration()).thenReturn(dynamicPaymentsConfiguration);
-    when(dynamicPaymentsConfiguration.getAllowedCountryCodes()).thenReturn(Collections.emptySet());
+    when(dynamicPaymentsConfiguration.getDisallowedPrefixes()).thenReturn(Collections.emptyList());
 
     when(rateLimiters.getProfileLimiter()).thenReturn(rateLimiter);
     when(rateLimiters.getUsernameLookupLimiter()).thenReturn(usernameRateLimiter);
@@ -169,13 +168,13 @@ class ProfileControllerTest {
     when(capabilitiesAccount.isAnnouncementGroupSupported()).thenReturn(true);
     when(capabilitiesAccount.isChangeNumberSupported()).thenReturn(true);
 
-    when(accountsManager.get(AuthHelper.VALID_NUMBER_TWO)).thenReturn(Optional.of(profileAccount));
-    when(accountsManager.get(AuthHelper.VALID_UUID_TWO)).thenReturn(Optional.of(profileAccount));
+    when(accountsManager.getByE164(AuthHelper.VALID_NUMBER_TWO)).thenReturn(Optional.of(profileAccount));
+    when(accountsManager.getByAccountIdentifier(AuthHelper.VALID_UUID_TWO)).thenReturn(Optional.of(profileAccount));
     when(usernamesManager.get(AuthHelper.VALID_UUID_TWO)).thenReturn(Optional.of("n00bkiller"));
     when(usernamesManager.get("n00bkiller")).thenReturn(Optional.of(AuthHelper.VALID_UUID_TWO));
 
-    when(accountsManager.get(AuthHelper.VALID_NUMBER)).thenReturn(Optional.of(capabilitiesAccount));
-    when(accountsManager.get(AuthHelper.VALID_UUID)).thenReturn(Optional.of(capabilitiesAccount));
+    when(accountsManager.getByE164(AuthHelper.VALID_NUMBER)).thenReturn(Optional.of(capabilitiesAccount));
+    when(accountsManager.getByAccountIdentifier(AuthHelper.VALID_UUID)).thenReturn(Optional.of(capabilitiesAccount));
 
     when(profilesManager.get(eq(AuthHelper.VALID_UUID), eq("someversion"))).thenReturn(Optional.empty());
     when(profilesManager.get(eq(AuthHelper.VALID_UUID_TWO), eq("validversion"))).thenReturn(Optional.of(new VersionedProfile(
@@ -208,7 +207,7 @@ class ProfileControllerTest {
     assertThat(profile.getBadges()).hasSize(1).element(0).has(new Condition<>(
         badge -> "Test Badge".equals(badge.getName()), "has badge with expected name"));
 
-    verify(accountsManager).get(AuthHelper.VALID_UUID_TWO);
+    verify(accountsManager).getByAccountIdentifier(AuthHelper.VALID_UUID_TWO);
     verify(usernamesManager, times(1)).get(eq(AuthHelper.VALID_UUID_TWO));
     verify(rateLimiter, times(1)).validate(AuthHelper.VALID_UUID);
   }
@@ -229,7 +228,7 @@ class ProfileControllerTest {
     assertThat(profile.getBadges()).hasSize(1).element(0).has(new Condition<>(
         badge -> "Test Badge".equals(badge.getName()), "has badge with expected name"));
 
-    verify(accountsManager, times(1)).get(eq(AuthHelper.VALID_UUID_TWO));
+    verify(accountsManager, times(1)).getByAccountIdentifier(eq(AuthHelper.VALID_UUID_TWO));
     verify(usernamesManager, times(1)).get(eq("n00bkiller"));
     verify(usernameRateLimiter, times(1)).validate(eq(AuthHelper.VALID_UUID));
   }
@@ -424,12 +423,14 @@ class ProfileControllerTest {
   void testSetProfileWithAvatarUploadAndPreviousAvatar() throws InvalidInputException {
     ProfileKeyCommitment commitment = new ProfileKey(new byte[32]).getCommitment(AuthHelper.VALID_UUID_TWO);
 
-    ProfileAvatarUploadAttributes uploadAttributes= resources.getJerseyTest()
-                                                             .target("/v1/profile/")
-                                                             .request()
-                                                             .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID_TWO, AuthHelper.VALID_PASSWORD_TWO))
-                                                             .put(Entity.entity(new CreateProfileRequest(commitment, "validversion", "123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678", null, null,
-                                                                 null, true, List.of()), MediaType.APPLICATION_JSON_TYPE), ProfileAvatarUploadAttributes.class);
+    resources.getJerseyTest()
+        .target("/v1/profile/")
+        .request()
+        .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID_TWO, AuthHelper.VALID_PASSWORD_TWO))
+        .put(Entity.entity(new CreateProfileRequest(commitment, "validversion",
+            "123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678",
+            null, null,
+            null, true, List.of()), MediaType.APPLICATION_JSON_TYPE), ProfileAvatarUploadAttributes.class);
 
     ArgumentCaptor<VersionedProfile> profileArgumentCaptor = ArgumentCaptor.forClass(VersionedProfile.class);
 
@@ -511,9 +512,6 @@ class ProfileControllerTest {
 
   @Test
   void testSetProfilePaymentAddress() throws InvalidInputException {
-    when(dynamicPaymentsConfiguration.getAllowedCountryCodes())
-        .thenReturn(Set.of(Util.getCountryCode(AuthHelper.VALID_NUMBER_TWO)));
-
     ProfileKeyCommitment commitment = new ProfileKey(new byte[32]).getCommitment(AuthHelper.VALID_UUID);
 
     clearInvocations(AuthHelper.VALID_ACCOUNT_TWO);
@@ -552,6 +550,9 @@ class ProfileControllerTest {
 
   @Test
   void testSetProfilePaymentAddressCountryNotAllowed() throws InvalidInputException {
+    when(dynamicPaymentsConfiguration.getDisallowedPrefixes())
+        .thenReturn(List.of(AuthHelper.VALID_NUMBER_TWO.substring(0, 3)));
+
     ProfileKeyCommitment commitment = new ProfileKey(new byte[32]).getCommitment(AuthHelper.VALID_UUID);
 
     clearInvocations(AuthHelper.VALID_ACCOUNT_TWO);
@@ -591,7 +592,7 @@ class ProfileControllerTest {
     assertThat(profile.getBadges()).hasSize(1).element(0).has(new Condition<>(
         badge -> "Test Badge".equals(badge.getName()), "has badge with expected name"));
 
-    verify(accountsManager, times(1)).get(eq(AuthHelper.VALID_UUID_TWO));
+    verify(accountsManager, times(1)).getByAccountIdentifier(eq(AuthHelper.VALID_UUID_TWO));
     verify(usernamesManager, times(1)).get(eq(AuthHelper.VALID_UUID_TWO));
     verify(profilesManager, times(1)).get(eq(AuthHelper.VALID_UUID_TWO), eq("validversion"));
 
@@ -600,9 +601,6 @@ class ProfileControllerTest {
 
   @Test
   void testSetProfileUpdatesAccountCurrentVersion() throws InvalidInputException {
-    when(dynamicPaymentsConfiguration.getAllowedCountryCodes())
-        .thenReturn(Set.of(Util.getCountryCode(AuthHelper.VALID_NUMBER_TWO)));
-
     ProfileKeyCommitment commitment = new ProfileKey(new byte[32]).getCommitment(AuthHelper.VALID_UUID_TWO);
 
     clearInvocations(AuthHelper.VALID_ACCOUNT_TWO);

service/src/test/java/org/whispersystems/textsecuregcm/tests/push/APNSenderTest.java

@@ -51,7 +51,7 @@ public class APNSenderTest {
   public void setup() {
     when(destinationAccount.getDevice(1)).thenReturn(Optional.of(destinationDevice));
     when(destinationDevice.getApnId()).thenReturn(DESTINATION_APN_ID);
-    when(accountsManager.get(DESTINATION_UUID)).thenReturn(Optional.of(destinationAccount));
+    when(accountsManager.getByAccountIdentifier(DESTINATION_UUID)).thenReturn(Optional.of(destinationAccount));
   }
 
   @Test
@@ -158,7 +158,7 @@ public class APNSenderTest {
     assertThat(apnResult.getStatus()).isEqualTo(ApnResult.Status.NO_SUCH_USER);
 
     verifyNoMoreInteractions(apnsClient);
-    verify(accountsManager, times(1)).get(eq(DESTINATION_UUID));
+    verify(accountsManager, times(1)).getByAccountIdentifier(eq(DESTINATION_UUID));
     verify(destinationAccount, times(1)).getDevice(1);
     verify(destinationDevice, times(1)).getApnId();
     verify(destinationDevice, times(1)).getPushTimestamp();
@@ -261,7 +261,7 @@ public class APNSenderTest {
     assertThat(apnResult.getStatus()).isEqualTo(ApnResult.Status.NO_SUCH_USER);
 
     verifyNoMoreInteractions(apnsClient);
-    verify(accountsManager, times(1)).get(eq(DESTINATION_UUID));
+    verify(accountsManager, times(1)).getByAccountIdentifier(eq(DESTINATION_UUID));
     verify(destinationAccount, times(1)).getDevice(1);
     verify(destinationDevice, times(1)).getApnId();
     verify(destinationDevice, times(1)).getPushTimestamp();

service/src/test/java/org/whispersystems/textsecuregcm/tests/push/GCMSenderTest.java

@@ -72,7 +72,7 @@ public class GCMSenderTest {
     AccountsHelper.setupMockUpdate(accountsManager);
 
     when(destinationAccount.getDevice(1)).thenReturn(Optional.of(destinationDevice));
-    when(accountsManager.get(destinationUuid)).thenReturn(Optional.of(destinationAccount));
+    when(accountsManager.getByAccountIdentifier(destinationUuid)).thenReturn(Optional.of(destinationAccount));
     when(destinationDevice.getGcmId()).thenReturn(gcmId);
 
     when(invalidResult.isInvalidRegistrationId()).thenReturn(true);
@@ -90,7 +90,7 @@ public class GCMSenderTest {
     gcmSender.sendMessage(message);
 
     verify(sender, times(1)).send(any(Message.class));
-    verify(accountsManager, times(1)).get(eq(destinationUuid));
+    verify(accountsManager, times(1)).getByAccountIdentifier(eq(destinationUuid));
     verify(accountsManager, times(1)).updateDevice(eq(destinationAccount), eq(1L), any());
     verify(destinationDevice, times(1)).setUninstalledFeedbackTimestamp(eq(Util.todayInMillis()));
   }
@@ -110,7 +110,7 @@ public class GCMSenderTest {
     Device         destinationDevice  = mock(Device.class        );
 
     when(destinationAccount.getDevice(1)).thenReturn(Optional.of(destinationDevice));
-    when(accountsManager.get(destinationUuid)).thenReturn(Optional.of(destinationAccount));
+    when(accountsManager.getByAccountIdentifier(destinationUuid)).thenReturn(Optional.of(destinationAccount));
     when(destinationDevice.getGcmId()).thenReturn(gcmId);
 
     AccountsHelper.setupMockUpdate(accountsManager);
@@ -131,7 +131,7 @@ public class GCMSenderTest {
     gcmSender.sendMessage(message);
 
     verify(sender, times(1)).send(any(Message.class));
-    verify(accountsManager, times(1)).get(eq(destinationUuid));
+    verify(accountsManager, times(1)).getByAccountIdentifier(eq(destinationUuid));
     verify(accountsManager, times(1)).updateDevice(eq(destinationAccount), eq(1L), any());
     verify(destinationDevice, times(1)).setGcmId(eq(canonicalId));
   }

service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountTest.java

@@ -158,26 +158,30 @@ class AccountTest {
     when(disabledMasterDevice.getId()).thenReturn(1L);
     when(disabledLinkedDevice.getId()).thenReturn(2L);
 
-    assertTrue( new Account("+14151234567", UUID.randomUUID(), Set.of(enabledMasterDevice),                        new byte[0]).isEnabled());
-    assertTrue( new Account("+14151234567", UUID.randomUUID(), Set.of(enabledMasterDevice, enabledLinkedDevice),   new byte[0]).isEnabled());
-    assertTrue( new Account("+14151234567", UUID.randomUUID(), Set.of(enabledMasterDevice, disabledLinkedDevice),  new byte[0]).isEnabled());
-    assertFalse(new Account("+14151234567", UUID.randomUUID(), Set.of(disabledMasterDevice),                       new byte[0]).isEnabled());
-    assertFalse(new Account("+14151234567", UUID.randomUUID(), Set.of(disabledMasterDevice, enabledLinkedDevice),  new byte[0]).isEnabled());
-    assertFalse(new Account("+14151234567", UUID.randomUUID(), Set.of(disabledMasterDevice, disabledLinkedDevice), new byte[0]).isEnabled());
+    assertTrue( new Account("+14151234567", UUID.randomUUID(), UUID.randomUUID(), Set.of(enabledMasterDevice),                        new byte[0]).isEnabled());
+    assertTrue( new Account("+14151234567", UUID.randomUUID(), UUID.randomUUID(),
+        Set.of(enabledMasterDevice, enabledLinkedDevice),   new byte[0]).isEnabled());
+    assertTrue( new Account("+14151234567", UUID.randomUUID(), UUID.randomUUID(),
+        Set.of(enabledMasterDevice, disabledLinkedDevice),  new byte[0]).isEnabled());
+    assertFalse(new Account("+14151234567", UUID.randomUUID(), UUID.randomUUID(), Set.of(disabledMasterDevice),                       new byte[0]).isEnabled());
+    assertFalse(new Account("+14151234567", UUID.randomUUID(), UUID.randomUUID(),
+        Set.of(disabledMasterDevice, enabledLinkedDevice),  new byte[0]).isEnabled());
+    assertFalse(new Account("+14151234567", UUID.randomUUID(), UUID.randomUUID(),
+        Set.of(disabledMasterDevice, disabledLinkedDevice), new byte[0]).isEnabled());
   }
 
   @Test
   void testCapabilities() {
-    Account uuidCapable = new Account("+14152222222", UUID.randomUUID(), new HashSet<Device>() {{
+    Account uuidCapable = new Account("+14152222222", UUID.randomUUID(), UUID.randomUUID(), new HashSet<Device>() {{
       add(gv2CapableDevice);
     }}, "1234".getBytes());
 
-    Account uuidIncapable = new Account("+14152222222", UUID.randomUUID(), new HashSet<Device>() {{
+    Account uuidIncapable = new Account("+14152222222", UUID.randomUUID(), UUID.randomUUID(), new HashSet<Device>() {{
       add(gv2CapableDevice);
       add(gv2IncapableDevice);
     }}, "1234".getBytes());
 
-    Account uuidCapableWithExpiredIncapable = new Account("+14152222222", UUID.randomUUID(), new HashSet<Device>() {{
+    Account uuidCapableWithExpiredIncapable = new Account("+14152222222", UUID.randomUUID(), UUID.randomUUID(), new HashSet<Device>() {{
       add(gv2CapableDevice);
       add(gv2IncapableExpiredDevice);
     }}, "1234".getBytes());
@@ -213,20 +217,20 @@ class AccountTest {
 
     {
       final Account transferableMasterAccount =
-              new Account("+14152222222", UUID.randomUUID(), Collections.singleton(transferCapableMasterDevice), "1234".getBytes());
+              new Account("+14152222222", UUID.randomUUID(), UUID.randomUUID(), Collections.singleton(transferCapableMasterDevice), "1234".getBytes());
 
       assertTrue(transferableMasterAccount.isTransferSupported());
     }
 
     {
       final Account nonTransferableMasterAccount =
-              new Account("+14152222222", UUID.randomUUID(), Collections.singleton(nonTransferCapableMasterDevice), "1234".getBytes());
+              new Account("+14152222222", UUID.randomUUID(), UUID.randomUUID(), Collections.singleton(nonTransferCapableMasterDevice), "1234".getBytes());
 
       assertFalse(nonTransferableMasterAccount.isTransferSupported());
     }
 
     {
-      final Account transferableLinkedAccount = new Account("+14152222222", UUID.randomUUID(), new HashSet<>() {{
+      final Account transferableLinkedAccount = new Account("+14152222222", UUID.randomUUID(), UUID.randomUUID(), new HashSet<>() {{
         add(nonTransferCapableMasterDevice);
         add(transferCapableLinkedDevice);
       }}, "1234".getBytes());
@@ -237,7 +241,7 @@ class AccountTest {
 
   @Test
   void testDiscoverableByPhoneNumber() {
-    final Account account = new Account("+14152222222", UUID.randomUUID(), Collections.singleton(recentMasterDevice),
+    final Account account = new Account("+14152222222", UUID.randomUUID(), UUID.randomUUID(), Collections.singleton(recentMasterDevice),
         "1234".getBytes());
 
     assertTrue(account.isDiscoverableByPhoneNumber(),
@@ -252,66 +256,70 @@ class AccountTest {
 
   @Test
   void isGroupsV2Supported() {
-    assertTrue(new Account("+18005551234", UUID.randomUUID(), Set.of(gv2CapableDevice),
+    assertTrue(new Account("+18005551234", UUID.randomUUID(), UUID.randomUUID(), Set.of(gv2CapableDevice),
         "1234".getBytes(StandardCharsets.UTF_8)).isGroupsV2Supported());
-    assertTrue(new Account("+18005551234", UUID.randomUUID(), Set.of(gv2CapableDevice, gv2IncapableExpiredDevice),
+    assertTrue(new Account("+18005551234", UUID.randomUUID(), UUID.randomUUID(),
+        Set.of(gv2CapableDevice, gv2IncapableExpiredDevice),
         "1234".getBytes(StandardCharsets.UTF_8)).isGroupsV2Supported());
-    assertFalse(new Account("+18005551234", UUID.randomUUID(), Set.of(gv2CapableDevice, gv2IncapableDevice),
+    assertFalse(new Account("+18005551234", UUID.randomUUID(), UUID.randomUUID(),
+        Set.of(gv2CapableDevice, gv2IncapableDevice),
         "1234".getBytes(StandardCharsets.UTF_8)).isGroupsV2Supported());
   }
 
   @Test
   void isGv1MigrationSupported() {
-    assertTrue(new Account("+18005551234", UUID.randomUUID(), Set.of(gv1MigrationCapableDevice),
+    assertTrue(new Account("+18005551234", UUID.randomUUID(), UUID.randomUUID(), Set.of(gv1MigrationCapableDevice),
         "1234".getBytes(StandardCharsets.UTF_8)).isGv1MigrationSupported());
     assertFalse(
-        new Account("+18005551234", UUID.randomUUID(), Set.of(gv1MigrationCapableDevice, gv1MigrationIncapableDevice),
+        new Account("+18005551234", UUID.randomUUID(), UUID.randomUUID(),
+            Set.of(gv1MigrationCapableDevice, gv1MigrationIncapableDevice),
             "1234".getBytes(StandardCharsets.UTF_8)).isGv1MigrationSupported());
     assertTrue(new Account("+18005551234", UUID.randomUUID(),
-        Set.of(gv1MigrationCapableDevice, gv1MigrationIncapableExpiredDevice), "1234".getBytes(StandardCharsets.UTF_8))
+        UUID.randomUUID(), Set.of(gv1MigrationCapableDevice, gv1MigrationIncapableExpiredDevice), "1234".getBytes(StandardCharsets.UTF_8))
         .isGv1MigrationSupported());
   }
 
   @Test
   void isSenderKeySupported() {
-    assertThat(new Account("+18005551234", UUID.randomUUID(), Set.of(senderKeyCapableDevice),
+    assertThat(new Account("+18005551234", UUID.randomUUID(), UUID.randomUUID(), Set.of(senderKeyCapableDevice),
         "1234".getBytes(StandardCharsets.UTF_8)).isSenderKeySupported()).isTrue();
-    assertThat(new Account("+18005551234", UUID.randomUUID(), Set.of(senderKeyCapableDevice, senderKeyIncapableDevice),
+    assertThat(new Account("+18005551234", UUID.randomUUID(), UUID.randomUUID(),
+        Set.of(senderKeyCapableDevice, senderKeyIncapableDevice),
         "1234".getBytes(StandardCharsets.UTF_8)).isSenderKeySupported()).isFalse();
     assertThat(new Account("+18005551234", UUID.randomUUID(),
-        Set.of(senderKeyCapableDevice, senderKeyIncapableExpiredDevice),
+        UUID.randomUUID(), Set.of(senderKeyCapableDevice, senderKeyIncapableExpiredDevice),
         "1234".getBytes(StandardCharsets.UTF_8)).isSenderKeySupported()).isTrue();
   }
 
   @Test
   void isAnnouncementGroupSupported() {
     assertThat(new Account("+18005551234", UUID.randomUUID(),
-        Set.of(announcementGroupCapableDevice),
+        UUID.randomUUID(), Set.of(announcementGroupCapableDevice),
         "1234".getBytes(StandardCharsets.UTF_8)).isAnnouncementGroupSupported()).isTrue();
     assertThat(new Account("+18005551234", UUID.randomUUID(),
-        Set.of(announcementGroupCapableDevice, announcementGroupIncapableDevice),
+        UUID.randomUUID(), Set.of(announcementGroupCapableDevice, announcementGroupIncapableDevice),
         "1234".getBytes(StandardCharsets.UTF_8)).isAnnouncementGroupSupported()).isFalse();
     assertThat(new Account("+18005551234", UUID.randomUUID(),
-        Set.of(announcementGroupCapableDevice, announcementGroupIncapableExpiredDevice),
+        UUID.randomUUID(), Set.of(announcementGroupCapableDevice, announcementGroupIncapableExpiredDevice),
         "1234".getBytes(StandardCharsets.UTF_8)).isAnnouncementGroupSupported()).isTrue();
   }
 
   @Test
   void isChangeNumberSupported() {
     assertThat(new Account("+18005551234", UUID.randomUUID(),
-        Set.of(changeNumberCapableDevice),
+        UUID.randomUUID(), Set.of(changeNumberCapableDevice),
         "1234".getBytes(StandardCharsets.UTF_8)).isChangeNumberSupported()).isTrue();
     assertThat(new Account("+18005551234", UUID.randomUUID(),
-        Set.of(changeNumberCapableDevice, changeNumberIncapableDevice),
+        UUID.randomUUID(), Set.of(changeNumberCapableDevice, changeNumberIncapableDevice),
         "1234".getBytes(StandardCharsets.UTF_8)).isChangeNumberSupported()).isFalse();
     assertThat(new Account("+18005551234", UUID.randomUUID(),
-        Set.of(changeNumberCapableDevice, changeNumberIncapableExpiredDevice),
+        UUID.randomUUID(), Set.of(changeNumberCapableDevice, changeNumberIncapableExpiredDevice),
         "1234".getBytes(StandardCharsets.UTF_8)).isChangeNumberSupported()).isTrue();
   }
 
   @Test
   void stale() {
-    final Account account = new Account("+14151234567", UUID.randomUUID(), Collections.emptySet(), new byte[0]);
+    final Account account = new Account("+14151234567", UUID.randomUUID(), UUID.randomUUID(), Collections.emptySet(), new byte[0]);
 
     assertDoesNotThrow(account::getNumber);
 
@@ -327,7 +335,7 @@ class AccountTest {
     final Set<Device> devices = new HashSet<>();
     devices.add(createDevice(Device.MASTER_ID));
 
-    final Account account = new Account("+14151234567", UUID.randomUUID(), devices, new byte[0]);
+    final Account account = new Account("+14151234567", UUID.randomUUID(), UUID.randomUUID(), devices, new byte[0]);
 
     assertThat(account.getNextDeviceId()).isEqualTo(2L);
 
@@ -348,7 +356,7 @@ class AccountTest {
 
   @Test
   void addAndRemoveBadges() {
-    final Account account = new Account("+14151234567", UUID.randomUUID(), Set.of(createDevice(Device.MASTER_ID)), new byte[0]);
+    final Account account = new Account("+14151234567", UUID.randomUUID(), UUID.randomUUID(), Set.of(createDevice(Device.MASTER_ID)), new byte[0]);
     final Clock clock = mock(Clock.class);
     when(clock.instant()).thenReturn(Instant.ofEpochSecond(40));
 

service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java

@@ -25,7 +25,11 @@ import static org.mockito.Mockito.when;
 
 import io.lettuce.core.RedisException;
 import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
+import java.time.Clock;
+import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
@@ -54,6 +58,7 @@ import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.Device.DeviceCapabilities;
 import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
+import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
 import org.whispersystems.textsecuregcm.storage.ProfilesManager;
 import org.whispersystems.textsecuregcm.storage.StoredVerificationCodeManager;
 import org.whispersystems.textsecuregcm.storage.UsernamesManager;
@@ -94,11 +99,12 @@ class AccountsManagerTest {
     doAnswer((Answer<Void>) invocation -> {
       final Account account = invocation.getArgument(0, Account.class);
       final String number = invocation.getArgument(1, String.class);
+      final UUID phoneNumberIdentifier = invocation.getArgument(2, UUID.class);
 
-      account.setNumber(number);
+      account.setNumber(number, phoneNumberIdentifier);
 
       return null;
-    }).when(accounts).changeNumber(any(), anyString());
+    }).when(accounts).changeNumber(any(), anyString(), any());
 
     doAnswer(invocation -> {
       //noinspection unchecked
@@ -112,8 +118,17 @@ class AccountsManagerTest {
     final SecureBackupClient backupClient = mock(SecureBackupClient.class);
     when(backupClient.deleteBackups(any())).thenReturn(CompletableFuture.completedFuture(null));
 
+    final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
+    final Map<String, UUID> phoneNumberIdentifiersByE164 = new HashMap<>();
+
+    when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString())).thenAnswer((Answer<UUID>) invocation -> {
+      final String number = invocation.getArgument(0, String.class);
+      return phoneNumberIdentifiersByE164.computeIfAbsent(number, n -> UUID.randomUUID());
+    });
+
     accountsManager = new AccountsManager(
         accounts,
+        phoneNumberIdentifiers,
         RedisClusterHelper.buildMockRedisCluster(commands),
         deletedAccountsManager,
         directoryQueue,
@@ -124,7 +139,8 @@ class AccountsManagerTest {
         mock(StoredVerificationCodeManager.class),
         storageClient,
         backupClient,
-        mock(ClientPresenceManager.class));
+        mock(ClientPresenceManager.class),
+        mock(Clock.class));
   }
 
   @Test
@@ -132,13 +148,14 @@ class AccountsManagerTest {
     UUID uuid = UUID.randomUUID();
 
     when(commands.get(eq("AccountMap::+14152222222"))).thenReturn(uuid.toString());
-    when(commands.get(eq("Account3::" + uuid))).thenReturn("{\"number\": \"+14152222222\", \"name\": \"test\"}");
+    when(commands.get(eq("Account3::" + uuid))).thenReturn("{\"number\": \"+14152222222\", \"name\": \"test\", \"pni\": \"de24dc73-fbd8-41be-a7d5-764c70d9da7e\"}");
 
-    Optional<Account> account = accountsManager.get("+14152222222");
+    Optional<Account> account = accountsManager.getByE164("+14152222222");
 
     assertTrue(account.isPresent());
     assertEquals(account.get().getNumber(), "+14152222222");
     assertEquals(account.get().getProfileName(), "test");
+    assertEquals(Optional.of(UUID.fromString("de24dc73-fbd8-41be-a7d5-764c70d9da7e")), account.get().getPhoneNumberIdentifier());
 
     verify(commands, times(1)).get(eq("AccountMap::+14152222222"));
     verify(commands, times(1)).get(eq("Account3::" + uuid));
@@ -151,14 +168,15 @@ class AccountsManagerTest {
   void testGetAccountByUuidInCache() {
     UUID uuid = UUID.randomUUID();
 
-    when(commands.get(eq("Account3::" + uuid))).thenReturn("{\"number\": \"+14152222222\", \"name\": \"test\"}");
+    when(commands.get(eq("Account3::" + uuid))).thenReturn("{\"number\": \"+14152222222\", \"name\": \"test\", \"pni\": \"de24dc73-fbd8-41be-a7d5-764c70d9da7e\"}");
 
-    Optional<Account> account = accountsManager.get(uuid);
+    Optional<Account> account = accountsManager.getByAccountIdentifier(uuid);
 
     assertTrue(account.isPresent());
     assertEquals(account.get().getNumber(), "+14152222222");
     assertEquals(account.get().getUuid(), uuid);
     assertEquals(account.get().getProfileName(), "test");
+    assertEquals(Optional.of(UUID.fromString("de24dc73-fbd8-41be-a7d5-764c70d9da7e")), account.get().getPhoneNumberIdentifier());
 
     verify(commands, times(1)).get(eq("Account3::" + uuid));
     verifyNoMoreInteractions(commands);
@@ -166,110 +184,189 @@ class AccountsManagerTest {
     verifyNoInteractions(accounts);
   }
 
+  @Test
+  void testGetByPniInCache() {
+    UUID uuid = UUID.randomUUID();
+    UUID pni = UUID.randomUUID();
+
+    when(commands.get(eq("AccountMap::" + pni))).thenReturn(uuid.toString());
+    when(commands.get(eq("Account3::" + uuid))).thenReturn("{\"number\": \"+14152222222\", \"name\": \"test\", \"pni\": \"de24dc73-fbd8-41be-a7d5-764c70d9da7e\"}");
+
+    Optional<Account> account = accountsManager.getByPhoneNumberIdentifier(pni);
+
+    assertTrue(account.isPresent());
+    assertEquals(account.get().getNumber(), "+14152222222");
+    assertEquals(account.get().getProfileName(), "test");
+    assertEquals(Optional.of(UUID.fromString("de24dc73-fbd8-41be-a7d5-764c70d9da7e")), account.get().getPhoneNumberIdentifier());
+
+    verify(commands).get(eq("AccountMap::" + pni));
+    verify(commands).get(eq("Account3::" + uuid));
+    verifyNoMoreInteractions(commands);
+
+    verifyNoInteractions(accounts);
+  }
 
   @Test
   void testGetAccountByNumberNotInCache() {
     UUID uuid = UUID.randomUUID();
-    Account account = new Account("+14152222222", uuid, new HashSet<>(), new byte[16]);
+    UUID pni = UUID.randomUUID();
+    Account account = new Account("+14152222222", uuid, pni, new HashSet<>(), new byte[16]);
 
     when(commands.get(eq("AccountMap::+14152222222"))).thenReturn(null);
-    when(accounts.get(eq("+14152222222"))).thenReturn(Optional.of(account));
+    when(accounts.getByE164(eq("+14152222222"))).thenReturn(Optional.of(account));
 
-    Optional<Account> retrieved = accountsManager.get("+14152222222");
+    Optional<Account> retrieved = accountsManager.getByE164("+14152222222");
 
     assertTrue(retrieved.isPresent());
     assertSame(retrieved.get(), account);
 
     verify(commands, times(1)).get(eq("AccountMap::+14152222222"));
     verify(commands, times(1)).set(eq("AccountMap::+14152222222"), eq(uuid.toString()));
+    verify(commands, times(1)).set(eq("AccountMap::" + pni), eq(uuid.toString()));
     verify(commands, times(1)).set(eq("Account3::" + uuid), anyString());
     verifyNoMoreInteractions(commands);
 
-    verify(accounts, times(1)).get(eq("+14152222222"));
+    verify(accounts, times(1)).getByE164(eq("+14152222222"));
     verifyNoMoreInteractions(accounts);
   }
 
   @Test
   void testGetAccountByUuidNotInCache() {
     UUID uuid = UUID.randomUUID();
-    Account account = new Account("+14152222222", uuid, new HashSet<>(), new byte[16]);
+    UUID pni = UUID.randomUUID();
+    Account account = new Account("+14152222222", uuid, pni, new HashSet<>(), new byte[16]);
 
     when(commands.get(eq("Account3::" + uuid))).thenReturn(null);
-    when(accounts.get(eq(uuid))).thenReturn(Optional.of(account));
+    when(accounts.getByAccountIdentifier(eq(uuid))).thenReturn(Optional.of(account));
 
-    Optional<Account> retrieved = accountsManager.get(uuid);
+    Optional<Account> retrieved = accountsManager.getByAccountIdentifier(uuid);
 
     assertTrue(retrieved.isPresent());
     assertSame(retrieved.get(), account);
 
     verify(commands, times(1)).get(eq("Account3::" + uuid));
     verify(commands, times(1)).set(eq("AccountMap::+14152222222"), eq(uuid.toString()));
+    verify(commands, times(1)).set(eq("AccountMap::" + pni), eq(uuid.toString()));
     verify(commands, times(1)).set(eq("Account3::" + uuid), anyString());
     verifyNoMoreInteractions(commands);
 
-    verify(accounts, times(1)).get(eq(uuid));
+    verify(accounts, times(1)).getByAccountIdentifier(eq(uuid));
+    verifyNoMoreInteractions(accounts);
+  }
+
+  @Test
+  void testGetAccountByPniNotInCache() {
+    UUID uuid = UUID.randomUUID();
+    UUID pni = UUID.randomUUID();
+
+    Account account = new Account("+14152222222", uuid, pni, new HashSet<>(), new byte[16]);
+
+    when(commands.get(eq("AccountMap::" + pni))).thenReturn(null);
+    when(accounts.getByPhoneNumberIdentifier(pni)).thenReturn(Optional.of(account));
+
+    Optional<Account> retrieved = accountsManager.getByPhoneNumberIdentifier(pni);
+
+    assertTrue(retrieved.isPresent());
+    assertSame(retrieved.get(), account);
+
+    verify(commands).get(eq("AccountMap::" + pni));
+    verify(commands).set(eq("AccountMap::" + pni), eq(uuid.toString()));
+    verify(commands).set(eq("AccountMap::+14152222222"), eq(uuid.toString()));
+    verify(commands).set(eq("Account3::" + uuid), anyString());
+    verifyNoMoreInteractions(commands);
+
+    verify(accounts).getByPhoneNumberIdentifier(pni);
     verifyNoMoreInteractions(accounts);
   }
 
   @Test
   void testGetAccountByNumberBrokenCache() {
     UUID uuid = UUID.randomUUID();
-    Account account = new Account("+14152222222", uuid, new HashSet<>(), new byte[16]);
+    UUID pni = UUID.randomUUID();
+    Account account = new Account("+14152222222", uuid, pni, new HashSet<>(), new byte[16]);
 
     when(commands.get(eq("AccountMap::+14152222222"))).thenThrow(new RedisException("Connection lost!"));
-    when(accounts.get(eq("+14152222222"))).thenReturn(Optional.of(account));
+    when(accounts.getByE164(eq("+14152222222"))).thenReturn(Optional.of(account));
 
-    Optional<Account> retrieved = accountsManager.get("+14152222222");
+    Optional<Account> retrieved = accountsManager.getByE164("+14152222222");
 
     assertTrue(retrieved.isPresent());
     assertSame(retrieved.get(), account);
 
     verify(commands, times(1)).get(eq("AccountMap::+14152222222"));
     verify(commands, times(1)).set(eq("AccountMap::+14152222222"), eq(uuid.toString()));
+    verify(commands, times(1)).set(eq("AccountMap::" + pni), eq(uuid.toString()));
     verify(commands, times(1)).set(eq("Account3::" + uuid), anyString());
     verifyNoMoreInteractions(commands);
 
-    verify(accounts, times(1)).get(eq("+14152222222"));
+    verify(accounts, times(1)).getByE164(eq("+14152222222"));
     verifyNoMoreInteractions(accounts);
   }
 
   @Test
   void testGetAccountByUuidBrokenCache() {
     UUID uuid = UUID.randomUUID();
-    Account account = new Account("+14152222222", uuid, new HashSet<>(), new byte[16]);
+    UUID pni = UUID.randomUUID();
+    Account account = new Account("+14152222222", uuid, pni, new HashSet<>(), new byte[16]);
 
     when(commands.get(eq("Account3::" + uuid))).thenThrow(new RedisException("Connection lost!"));
-    when(accounts.get(eq(uuid))).thenReturn(Optional.of(account));
+    when(accounts.getByAccountIdentifier(eq(uuid))).thenReturn(Optional.of(account));
 
-    Optional<Account> retrieved = accountsManager.get(uuid);
+    Optional<Account> retrieved = accountsManager.getByAccountIdentifier(uuid);
 
     assertTrue(retrieved.isPresent());
     assertSame(retrieved.get(), account);
 
     verify(commands, times(1)).get(eq("Account3::" + uuid));
     verify(commands, times(1)).set(eq("AccountMap::+14152222222"), eq(uuid.toString()));
+    verify(commands, times(1)).set(eq("AccountMap::" + pni), eq(uuid.toString()));
     verify(commands, times(1)).set(eq("Account3::" + uuid), anyString());
     verifyNoMoreInteractions(commands);
 
-    verify(accounts, times(1)).get(eq(uuid));
+    verify(accounts, times(1)).getByAccountIdentifier(eq(uuid));
+    verifyNoMoreInteractions(accounts);
+  }
+
+  @Test
+  void testGetAccountByPniBrokenCache() {
+    UUID uuid = UUID.randomUUID();
+    UUID pni = UUID.randomUUID();
+
+    Account account = new Account("+14152222222", uuid, pni, new HashSet<>(), new byte[16]);
+
+    when(commands.get(eq("AccountMap::" + pni))).thenThrow(new RedisException("OH NO"));
+    when(accounts.getByPhoneNumberIdentifier(pni)).thenReturn(Optional.of(account));
+
+    Optional<Account> retrieved = accountsManager.getByPhoneNumberIdentifier(pni);
+
+    assertTrue(retrieved.isPresent());
+    assertSame(retrieved.get(), account);
+
+    verify(commands).get(eq("AccountMap::" + pni));
+    verify(commands).set(eq("AccountMap::" + pni), eq(uuid.toString()));
+    verify(commands).set(eq("AccountMap::+14152222222"), eq(uuid.toString()));
+    verify(commands).set(eq("Account3::" + uuid), anyString());
+    verifyNoMoreInteractions(commands);
+
+    verify(accounts).getByPhoneNumberIdentifier(pni);
     verifyNoMoreInteractions(accounts);
   }
 
   @Test
   void testUpdate_optimisticLockingFailure() {
     UUID uuid = UUID.randomUUID();
-    Account account = new Account("+14152222222", uuid, new HashSet<>(), new byte[16]);
+    Account account = new Account("+14152222222", uuid, UUID.randomUUID(), new HashSet<>(), new byte[16]);
 
     when(commands.get(eq("Account3::" + uuid))).thenReturn(null);
 
-    when(accounts.get(uuid)).thenReturn(
-        Optional.of(new Account("+14152222222", uuid, new HashSet<>(), new byte[16])));
+    when(accounts.getByAccountIdentifier(uuid)).thenReturn(
+        Optional.of(new Account("+14152222222", uuid, UUID.randomUUID(), new HashSet<>(), new byte[16])));
     doThrow(ContestedOptimisticLockException.class)
         .doAnswer(ACCOUNT_UPDATE_ANSWER)
         .when(accounts).update(any());
 
-    when(accounts.get(uuid)).thenReturn(
-        Optional.of(new Account("+14152222222", uuid, new HashSet<>(), new byte[16])));
+    when(accounts.getByAccountIdentifier(uuid)).thenReturn(
+        Optional.of(new Account("+14152222222", uuid, UUID.randomUUID(), new HashSet<>(), new byte[16])));
     doThrow(ContestedOptimisticLockException.class)
         .doAnswer(ACCOUNT_UPDATE_ANSWER)
         .when(accounts).update(any());
@@ -279,7 +376,7 @@ class AccountsManagerTest {
     assertEquals(1, account.getVersion());
     assertEquals("name", account.getProfileName());
 
-    verify(accounts, times(1)).get(uuid);
+    verify(accounts, times(1)).getByAccountIdentifier(uuid);
     verify(accounts, times(2)).update(any());
     verifyNoMoreInteractions(accounts);
   }
@@ -287,10 +384,10 @@ class AccountsManagerTest {
   @Test
   void testUpdate_dynamoOptimisticLockingFailureDuringCreate() {
     UUID uuid = UUID.randomUUID();
-    Account account = new Account("+14152222222", uuid, new HashSet<>(), new byte[16]);
+    Account account = new Account("+14152222222", uuid, UUID.randomUUID(), new HashSet<>(), new byte[16]);
 
     when(commands.get(eq("Account3::" + uuid))).thenReturn(null);
-    when(accounts.get(uuid)).thenReturn(Optional.empty())
+    when(accounts.getByAccountIdentifier(uuid)).thenReturn(Optional.empty())
         .thenReturn(Optional.of(account));
     when(accounts.create(any())).thenThrow(ContestedOptimisticLockException.class);
 
@@ -304,10 +401,10 @@ class AccountsManagerTest {
   @Test
   void testUpdateDevice() {
     final UUID uuid = UUID.randomUUID();
-    Account account = new Account("+14152222222", uuid, new HashSet<>(), new byte[16]);
+    Account account = new Account("+14152222222", uuid, UUID.randomUUID(), new HashSet<>(), new byte[16]);
 
-    when(accounts.get(uuid)).thenReturn(
-        Optional.of(new Account("+14152222222", uuid, new HashSet<>(), new byte[16])));
+    when(accounts.getByAccountIdentifier(uuid)).thenReturn(
+        Optional.of(new Account("+14152222222", uuid, UUID.randomUUID(), new HashSet<>(), new byte[16])));
 
     assertTrue(account.getDevices().isEmpty());
 
@@ -340,7 +437,7 @@ class AccountsManagerTest {
 
     final String e164 = "+18005550123";
     final AccountAttributes attributes = new AccountAttributes(false, 0, null, null, true, null);
-    accountsManager.create(e164, "password", null, attributes);
+    accountsManager.create(e164, "password", null, attributes, new ArrayList<>());
 
     verify(accounts).create(argThat(account -> e164.equals(account.getNumber())));
     verifyNoInteractions(keys);
@@ -359,7 +456,7 @@ class AccountsManagerTest {
 
     final String e164 = "+18005550123";
     final AccountAttributes attributes = new AccountAttributes(false, 0, null, null, true, null);
-    accountsManager.create(e164, "password", null, attributes);
+    accountsManager.create(e164, "password", null, attributes, new ArrayList<>());
 
     verify(accounts).create(
         argThat(account -> e164.equals(account.getNumber()) && existingUuid.equals(account.getUuid())));
@@ -382,7 +479,7 @@ class AccountsManagerTest {
 
     final String e164 = "+18005550123";
     final AccountAttributes attributes = new AccountAttributes(false, 0, null, null, true, null);
-    accountsManager.create(e164, "password", null, attributes);
+    accountsManager.create(e164, "password", null, attributes, new ArrayList<>());
 
     verify(accounts).create(
         argThat(account -> e164.equals(account.getNumber()) && recentlyDeletedUuid.equals(account.getUuid())));
@@ -395,7 +492,7 @@ class AccountsManagerTest {
   @ValueSource(booleans = {true, false})
   void testCreateWithDiscoverability(final boolean discoverable) throws InterruptedException {
     final AccountAttributes attributes = new AccountAttributes(false, 0, null, null, discoverable, null);
-    final Account account = accountsManager.create("+18005550123", "password", null, attributes);
+    final Account account = accountsManager.create("+18005550123", "password", null, attributes, new ArrayList<>());
 
     assertEquals(discoverable, account.isDiscoverableByPhoneNumber());
 
@@ -410,7 +507,7 @@ class AccountsManagerTest {
     final AccountAttributes attributes = new AccountAttributes(false, 0, null, null, true,
         new DeviceCapabilities(false, false, false, hasStorage, false, false, false, false, false));
 
-    final Account account = accountsManager.create("+18005550123", "password", null, attributes);
+    final Account account = accountsManager.create("+18005550123", "password", null, attributes, new ArrayList<>());
 
     assertEquals(hasStorage, account.isStorageSupported());
   }
@@ -419,7 +516,7 @@ class AccountsManagerTest {
   @MethodSource
   void testUpdateDirectoryQueue(final boolean visibleBeforeUpdate, final boolean visibleAfterUpdate,
       final boolean expectRefresh) {
-    final Account account = new Account("+14152222222", UUID.randomUUID(), new HashSet<>(), new byte[16]);
+    final Account account = new Account("+14152222222", UUID.randomUUID(), UUID.randomUUID(), new HashSet<>(), new byte[16]);
 
     // this sets up the appropriate result for Account#shouldBeVisibleInDirectory
     final Device device = new Device(Device.MASTER_ID, "device", "token", "salt", null, null, null, true, 1,
@@ -446,7 +543,7 @@ class AccountsManagerTest {
   @ParameterizedTest
   @MethodSource
   void testUpdateDeviceLastSeen(final boolean expectUpdate, final long initialLastSeen, final long updatedLastSeen) {
-    final Account account = new Account("+14152222222", UUID.randomUUID(), new HashSet<>(), new byte[16]);
+    final Account account = new Account("+14152222222", UUID.randomUUID(), UUID.randomUUID(), new HashSet<>(), new byte[16]);
     final Device device = new Device(Device.MASTER_ID, "device", "token", "salt", null, null, null, true, 1,
         new SignedPreKey(1, "key", "sig"), initialLastSeen, 0,
         "OWT", 0, new DeviceCapabilities());
@@ -476,7 +573,7 @@ class AccountsManagerTest {
     final String targetNumber = "+14153333333";
     final UUID uuid = UUID.randomUUID();
 
-    Account account = new Account(originalNumber, uuid, new HashSet<>(), new byte[16]);
+    Account account = new Account(originalNumber, uuid, UUID.randomUUID(), new HashSet<>(), new byte[16]);
     account = accountsManager.changeNumber(account, targetNumber);
 
     assertEquals(targetNumber, account.getNumber());
@@ -488,7 +585,7 @@ class AccountsManagerTest {
   void testChangePhoneNumberSameNumber() throws InterruptedException {
     final String number = "+14152222222";
 
-    Account account = new Account(number, UUID.randomUUID(), new HashSet<>(), new byte[16]);
+    Account account = new Account(number, UUID.randomUUID(), UUID.randomUUID(), new HashSet<>(), new byte[16]);
     account = accountsManager.changeNumber(account, number);
 
     assertEquals(number, account.getNumber());
@@ -506,10 +603,10 @@ class AccountsManagerTest {
     final UUID existingAccountUuid = UUID.randomUUID();
     final UUID uuid = UUID.randomUUID();
 
-    final Account existingAccount = new Account(targetNumber, existingAccountUuid, new HashSet<>(), new byte[16]);
-    when(accounts.get(targetNumber)).thenReturn(Optional.of(existingAccount));
+    final Account existingAccount = new Account(targetNumber, existingAccountUuid, UUID.randomUUID(), new HashSet<>(), new byte[16]);
+    when(accounts.getByE164(targetNumber)).thenReturn(Optional.of(existingAccount));
 
-    Account account = new Account(originalNumber, uuid, new HashSet<>(), new byte[16]);
+    Account account = new Account(originalNumber, uuid, UUID.randomUUID(), new HashSet<>(), new byte[16]);
     account = accountsManager.changeNumber(account, targetNumber);
 
     assertEquals(targetNumber, account.getNumber());
@@ -524,8 +621,8 @@ class AccountsManagerTest {
     final String targetNumber = "+14153333333";
     final UUID uuid = UUID.randomUUID();
 
-    final Account account = new Account(originalNumber, uuid, new HashSet<>(), new byte[16]);
+    final Account account = new Account(originalNumber, uuid, UUID.randomUUID(), new HashSet<>(), new byte[16]);
 
-    assertThrows(AssertionError.class, () -> accountsManager.update(account, a -> a.setNumber(targetNumber)));
+    assertThrows(AssertionError.class, () -> accountsManager.update(account, a -> a.setNumber(targetNumber, UUID.randomUUID())));
   }
 }

service/src/test/java/org/whispersystems/textsecuregcm/tests/util/AccountsHelper.java

@@ -62,7 +62,7 @@ public class AccountsHelper {
   }
 
   public static void setupMockGet(final AccountsManager mockAccountsManager, final Set<Account> mockAccounts) {
-    when(mockAccountsManager.get(any(UUID.class))).thenAnswer(answer -> {
+    when(mockAccountsManager.getByAccountIdentifier(any(UUID.class))).thenAnswer(answer -> {
 
       final UUID uuid = answer.getArgument(0, UUID.class);
 
@@ -176,7 +176,7 @@ public class AccountsHelper {
     } else {
       final ObjectMapper mapper = SystemMapper.getMapper();
       updatedAccount = mapper.readValue(mapper.writeValueAsBytes(account), Account.class);
-      updatedAccount.setNumber(account.getNumber());
+      updatedAccount.setNumber(account.getNumber(), account.getPhoneNumberIdentifier().orElse(null));
       account.markStale();
     }
 

service/src/test/java/org/whispersystems/textsecuregcm/tests/util/AuthHelper.java

@@ -137,17 +137,17 @@ public class AuthHelper {
 
     reset(ACCOUNTS_MANAGER);
 
-    when(ACCOUNTS_MANAGER.get(VALID_NUMBER)).thenReturn(Optional.of(VALID_ACCOUNT));
-    when(ACCOUNTS_MANAGER.get(VALID_UUID)).thenReturn(Optional.of(VALID_ACCOUNT));
+    when(ACCOUNTS_MANAGER.getByE164(VALID_NUMBER)).thenReturn(Optional.of(VALID_ACCOUNT));
+    when(ACCOUNTS_MANAGER.getByAccountIdentifier(VALID_UUID)).thenReturn(Optional.of(VALID_ACCOUNT));
 
-    when(ACCOUNTS_MANAGER.get(VALID_NUMBER_TWO)).thenReturn(Optional.of(VALID_ACCOUNT_TWO));
-    when(ACCOUNTS_MANAGER.get(VALID_UUID_TWO)).thenReturn(Optional.of(VALID_ACCOUNT_TWO));
+    when(ACCOUNTS_MANAGER.getByE164(VALID_NUMBER_TWO)).thenReturn(Optional.of(VALID_ACCOUNT_TWO));
+    when(ACCOUNTS_MANAGER.getByAccountIdentifier(VALID_UUID_TWO)).thenReturn(Optional.of(VALID_ACCOUNT_TWO));
 
-    when(ACCOUNTS_MANAGER.get(DISABLED_NUMBER)).thenReturn(Optional.of(DISABLED_ACCOUNT));
-    when(ACCOUNTS_MANAGER.get(DISABLED_UUID)).thenReturn(Optional.of(DISABLED_ACCOUNT));
+    when(ACCOUNTS_MANAGER.getByE164(DISABLED_NUMBER)).thenReturn(Optional.of(DISABLED_ACCOUNT));
+    when(ACCOUNTS_MANAGER.getByAccountIdentifier(DISABLED_UUID)).thenReturn(Optional.of(DISABLED_ACCOUNT));
 
-    when(ACCOUNTS_MANAGER.get(UNDISCOVERABLE_NUMBER)).thenReturn(Optional.of(UNDISCOVERABLE_ACCOUNT));
-    when(ACCOUNTS_MANAGER.get(UNDISCOVERABLE_UUID)).thenReturn(Optional.of(UNDISCOVERABLE_ACCOUNT));
+    when(ACCOUNTS_MANAGER.getByE164(UNDISCOVERABLE_NUMBER)).thenReturn(Optional.of(UNDISCOVERABLE_ACCOUNT));
+    when(ACCOUNTS_MANAGER.getByAccountIdentifier(UNDISCOVERABLE_UUID)).thenReturn(Optional.of(UNDISCOVERABLE_ACCOUNT));
 
     AccountsHelper.setupMockUpdateForAuthHelper(ACCOUNTS_MANAGER);
 
@@ -220,8 +220,8 @@ public class AuthHelper {
       when(account.getUuid()).thenReturn(uuid);
       when(account.getRelay()).thenReturn(Optional.empty());
       when(account.isEnabled()).thenReturn(true);
-      when(accountsManager.get(number)).thenReturn(Optional.of(account));
-      when(accountsManager.get(uuid)).thenReturn(Optional.of(account));
+      when(accountsManager.getByE164(number)).thenReturn(Optional.of(account));
+      when(accountsManager.getByAccountIdentifier(uuid)).thenReturn(Optional.of(account));
     }
   }
 

service/src/test/java/org/whispersystems/textsecuregcm/util/AttributeValuesTest.java

@@ -57,4 +57,10 @@ public class AttributeValuesTest {
     AttributeValue av = AttributeValues.fromByteBuffer(byteBuffer.flip());
     assertArrayEquals(new byte[]{0, 0, 0, 0, 0, 0, 0, 123}, AttributeValues.getByteArray(Map.of("foo", av), "foo", null));
   }
+
+  @Test
+  void testNullUuid() {
+    final Map<String, AttributeValue> item = Map.of("key", AttributeValue.builder().nul(true).build());
+    assertNull(AttributeValues.getUUID(item, "key", null));
+  }
 }

service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionTest.java

@@ -172,8 +172,8 @@ public class WebSocketConnectionTest {
     Account sender1 = mock(Account.class);
     when(sender1.getDevices()).thenReturn(sender1devices);
 
-    when(accountsManager.get("sender1")).thenReturn(Optional.of(sender1));
-    when(accountsManager.get("sender2")).thenReturn(Optional.empty());
+    when(accountsManager.getByE164("sender1")).thenReturn(Optional.of(sender1));
+    when(accountsManager.getByE164("sender2")).thenReturn(Optional.empty());
 
     String userAgent = "user-agent";
 
@@ -327,8 +327,8 @@ public class WebSocketConnectionTest {
     Account sender1 = mock(Account.class);
     when(sender1.getDevices()).thenReturn(sender1devices);
 
-    when(accountsManager.get("sender1")).thenReturn(Optional.of(sender1));
-    when(accountsManager.get("sender2")).thenReturn(Optional.<Account>empty());
+    when(accountsManager.getByE164("sender1")).thenReturn(Optional.of(sender1));
+    when(accountsManager.getByE164("sender2")).thenReturn(Optional.<Account>empty());
 
     String userAgent = "user-agent";
 
@@ -700,8 +700,8 @@ public class WebSocketConnectionTest {
     Account sender1 = mock(Account.class);
     when(sender1.getDevices()).thenReturn(sender1devices);
 
-    when(accountsManager.get("sender1")).thenReturn(Optional.of(sender1));
-    when(accountsManager.get("sender2")).thenReturn(Optional.empty());
+    when(accountsManager.getByE164("sender1")).thenReturn(Optional.of(sender1));
+    when(accountsManager.getByE164("sender2")).thenReturn(Optional.empty());
 
     String userAgent = "Signal-Desktop/1.2.3";
 
@@ -776,8 +776,8 @@ public class WebSocketConnectionTest {
     Account sender1 = mock(Account.class);
     when(sender1.getDevices()).thenReturn(sender1devices);
 
-    when(accountsManager.get("sender1")).thenReturn(Optional.of(sender1));
-    when(accountsManager.get("sender2")).thenReturn(Optional.empty());
+    when(accountsManager.getByE164("sender1")).thenReturn(Optional.of(sender1));
+    when(accountsManager.getByE164("sender2")).thenReturn(Optional.empty());
 
     String userAgent = "Signal-Android/4.68.3";