Read registration recovery passwords exclusively by PNI

integration-tests/src/main/java/org/signal/integration/IntegrationTools.java

@@ -8,6 +8,7 @@ package org.signal.integration;
 import java.time.Clock;
 import java.time.Duration;
 import java.util.Optional;
+import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import org.signal.integration.config.Config;
 import org.whispersystems.textsecuregcm.metrics.NoopAwsSdkMetricPublisher;
@@ -20,7 +21,6 @@ import org.whispersystems.textsecuregcm.storage.VerificationSessions;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
-import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 
 public class IntegrationTools {
 
@@ -57,8 +57,8 @@ public class IntegrationTools {
     this.verificationSessionManager = verificationSessionManager;
   }
 
-  public CompletableFuture<Void> populateRecoveryPassword(final String e164, final byte[] password) {
+  public CompletableFuture<Void> populateRecoveryPassword(final String number, final byte[] password) {
-    return registrationRecoveryPasswordsManager.storeForCurrentNumber(e164, password);
+    return registrationRecoveryPasswordsManager.storeForCurrentNumber(number, password);
   }
 
   public CompletableFuture<Optional<String>> peekVerificationSessionPushChallenge(final String sessionId) {

integration-tests/src/main/java/org/signal/integration/Operations.java

@@ -78,7 +78,7 @@ public final class Operations {
     final TestUser user = TestUser.create(number, accountPassword, registrationPassword);
     final AccountAttributes accountAttributes = user.accountAttributes();
 
-    INTEGRATION_TOOLS.populateRecoveryPassword(number, registrationPassword).join();
+    INTEGRATION_TOOLS.populateRecoveryPassword(user.phoneNumber(), registrationPassword).join();
 
     final ECKeyPair aciIdentityKeyPair = Curve.generateKeyPair();
     final ECKeyPair pniIdentityKeyPair = Curve.generateKeyPair();

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

@@ -1076,7 +1076,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
 
 
     final PhoneVerificationTokenManager phoneVerificationTokenManager = new PhoneVerificationTokenManager(
-        registrationServiceClient, registrationRecoveryPasswordsManager, registrationRecoveryChecker);
+        phoneNumberIdentifiers, registrationServiceClient, registrationRecoveryPasswordsManager, registrationRecoveryChecker);
     final List<Object> commonControllers = Lists.newArrayList(
         new AccountController(accountsManager, rateLimiters, turnTokenGenerator, registrationRecoveryPasswordsManager,
             usernameHashZkProofVerifier),
@@ -1119,8 +1119,9 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
             config.getCdnConfiguration().credentials().secretAccessKey().value(), config.getCdnConfiguration().region(),
             config.getCdnConfiguration().bucket()),
         new VerificationController(registrationServiceClient, new VerificationSessionManager(verificationSessions),
-            pushNotificationManager, registrationCaptchaManager, registrationRecoveryPasswordsManager, rateLimiters,
+            pushNotificationManager, registrationCaptchaManager, registrationRecoveryPasswordsManager,
-            accountsManager, registrationFraudChecker, dynamicConfigurationManager, clock)
+            phoneNumberIdentifiers, rateLimiters, accountsManager, registrationFraudChecker,
+            dynamicConfigurationManager, clock)
     );
     if (config.getSubscription() != null && config.getOneTimeDonations() != null) {
       SubscriptionManager subscriptionManager = new SubscriptionManager(subscriptions,

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

@@ -25,6 +25,7 @@ import org.whispersystems.textsecuregcm.entities.PhoneVerificationRequest;
 import org.whispersystems.textsecuregcm.entities.RegistrationServiceSession;
 import org.whispersystems.textsecuregcm.registration.RegistrationServiceClient;
 import org.whispersystems.textsecuregcm.spam.RegistrationRecoveryChecker;
+import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
 import org.whispersystems.textsecuregcm.storage.RegistrationRecoveryPasswordsManager;
 
 public class PhoneVerificationTokenManager {
@@ -33,13 +34,17 @@ public class PhoneVerificationTokenManager {
   private static final Duration REGISTRATION_RPC_TIMEOUT = Duration.ofSeconds(15);
   private static final long VERIFICATION_TIMEOUT_SECONDS = REGISTRATION_RPC_TIMEOUT.plusSeconds(1).getSeconds();
 
+  private final PhoneNumberIdentifiers phoneNumberIdentifiers;
+
   private final RegistrationServiceClient registrationServiceClient;
   private final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager;
   private final RegistrationRecoveryChecker registrationRecoveryChecker;
 
-  public PhoneVerificationTokenManager(final RegistrationServiceClient registrationServiceClient,
+  public PhoneVerificationTokenManager(final PhoneNumberIdentifiers phoneNumberIdentifiers,
+      final RegistrationServiceClient registrationServiceClient,
       final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager,
       final RegistrationRecoveryChecker registrationRecoveryChecker) {
+    this.phoneNumberIdentifiers = phoneNumberIdentifiers;
     this.registrationServiceClient = registrationServiceClient;
     this.registrationRecoveryPasswordsManager = registrationRecoveryPasswordsManager;
     this.registrationRecoveryChecker = registrationRecoveryChecker;
@@ -109,7 +114,7 @@ public class PhoneVerificationTokenManager {
       throw new ForbiddenException("recoveryPassword couldn't be verified");
     }
     try {
-      final boolean verified = registrationRecoveryPasswordsManager.verify(number, recoveryPassword)
+      final boolean verified = registrationRecoveryPasswordsManager.verify(phoneNumberIdentifiers.getPhoneNumberIdentifier(number).join(), recoveryPassword)
           .get(VERIFICATION_TIMEOUT_SECONDS, TimeUnit.SECONDS);
       if (!verified) {
         throw new ForbiddenException("recoveryPassword couldn't be verified");

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

@@ -23,6 +23,7 @@ import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
 import org.whispersystems.textsecuregcm.entities.PhoneVerificationRequest;
 import org.whispersystems.textsecuregcm.entities.RegistrationLockFailure;
 import org.whispersystems.textsecuregcm.entities.Svr3Credentials;
+import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
 import org.whispersystems.textsecuregcm.push.NotPushRegisteredException;

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

@@ -83,6 +83,7 @@ import org.whispersystems.textsecuregcm.spam.RegistrationFraudChecker;
 import org.whispersystems.textsecuregcm.spam.RegistrationFraudChecker.VerificationCheck;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
+import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
 import org.whispersystems.textsecuregcm.storage.RegistrationRecoveryPasswordsManager;
 import org.whispersystems.textsecuregcm.storage.VerificationSessionManager;
 import org.whispersystems.textsecuregcm.util.ExceptionUtils;
@@ -116,6 +117,7 @@ public class VerificationController {
   private final PushNotificationManager pushNotificationManager;
   private final RegistrationCaptchaManager registrationCaptchaManager;
   private final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager;
+  private final PhoneNumberIdentifiers phoneNumberIdentifiers;
   private final RateLimiters rateLimiters;
   private final AccountsManager accountsManager;
   private final RegistrationFraudChecker registrationFraudChecker;
@@ -127,6 +129,7 @@ public class VerificationController {
       final PushNotificationManager pushNotificationManager,
       final RegistrationCaptchaManager registrationCaptchaManager,
       final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager,
+      final PhoneNumberIdentifiers phoneNumberIdentifiers,
       final RateLimiters rateLimiters,
       final AccountsManager accountsManager,
       final RegistrationFraudChecker registrationFraudChecker,
@@ -137,6 +140,7 @@ public class VerificationController {
     this.pushNotificationManager = pushNotificationManager;
     this.registrationCaptchaManager = registrationCaptchaManager;
     this.registrationRecoveryPasswordsManager = registrationRecoveryPasswordsManager;
+    this.phoneNumberIdentifiers = phoneNumberIdentifiers;
     this.rateLimiters = rateLimiters;
     this.accountsManager = accountsManager;
     this.registrationFraudChecker = registrationFraudChecker;
@@ -626,8 +630,8 @@ public class VerificationController {
   }
 
   /**
-   * @throws ClientErrorException          with {@code 422} status if the ID cannot be decoded
+   * @throws ClientErrorException with {@code 422} status if the ID cannot be decoded
-   * @throws javax.ws.rs.NotFoundException if the ID cannot be found
+   * @throws NotFoundException    if the ID cannot be found
    */
   private RegistrationServiceSession retrieveRegistrationServiceSession(final String encodedSessionId) {
     final byte[] sessionId;

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

@@ -9,8 +9,11 @@ import java.util.Base64;
 import javax.annotation.Nullable;
 import org.signal.registration.rpc.RegistrationSessionMetadata;
 
-public record RegistrationServiceSession(byte[] id, String number, boolean verified,
+public record RegistrationServiceSession(byte[] id,
-                                         @Nullable Long nextSms, @Nullable Long nextVoiceCall,
+                                         String number,
+                                         boolean verified,
+                                         @Nullable Long nextSms,
+                                         @Nullable Long nextVoiceCall,
                                          @Nullable Long nextVerificationAttempt,
                                          long expiration) {
 

service/src/main/java/org/whispersystems/textsecuregcm/grpc/AccountsGrpcService.java

@@ -47,6 +47,7 @@ import org.whispersystems.textsecuregcm.auth.grpc.AuthenticationUtil;
 import org.whispersystems.textsecuregcm.controllers.AccountController;
 import org.whispersystems.textsecuregcm.entities.EncryptedUsername;
 import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
+import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.identity.PniServiceIdentifier;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;

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

@@ -16,7 +16,6 @@ import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import org.whispersystems.textsecuregcm.auth.SaltedTokenHash;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
-import org.whispersystems.textsecuregcm.util.Pair;
 import org.whispersystems.textsecuregcm.util.Util;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
@@ -28,9 +27,8 @@ import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
 
 public class RegistrationRecoveryPasswords {
 
-  // As a temporary transitional measure, this can be either a string representation of an E164-formatted phone number
+  // For historical reasons, we record the PNI as a UUID string rather than a compact byte array
-  // or a UUID (PNI) string
+  static final String KEY_PNI = "P";
-  static final String KEY_E164 = "P";
   static final String ATTR_EXP = "E";
   static final String ATTR_SALT = "S";
   static final String ATTR_HASH = "H";
@@ -54,10 +52,10 @@ public class RegistrationRecoveryPasswords {
     this.clock = requireNonNull(clock);
   }
 
-  public CompletableFuture<Optional<SaltedTokenHash>> lookup(final String number) {
+  public CompletableFuture<Optional<SaltedTokenHash>> lookup(final UUID phoneNumberIdentifier) {
     return asyncClient.getItem(GetItemRequest.builder()
             .tableName(tableName)
-            .key(Map.of(KEY_E164, AttributeValues.fromString(number)))
+            .key(Map.of(KEY_PNI, AttributeValues.fromString(phoneNumberIdentifier.toString())))
             .consistentRead(true)
             .build())
         .thenApply(getItemResponse -> Optional.ofNullable(getItemResponse.item())
@@ -66,10 +64,6 @@ public class RegistrationRecoveryPasswords {
             .map(RegistrationRecoveryPasswords::saltedTokenHashFromItem));
   }
 
-  public CompletableFuture<Optional<SaltedTokenHash>> lookup(final UUID phoneNumberIdentifier) {
-    return lookup(phoneNumberIdentifier.toString());
-  }
-
   public CompletableFuture<Void> addOrReplace(final String number, final UUID phoneNumberIdentifier, final SaltedTokenHash data) {
     final long expirationSeconds = expirationSeconds();
 
@@ -90,7 +84,7 @@ public class RegistrationRecoveryPasswords {
         .put(Put.builder()
             .tableName(tableName)
             .item(Map.of(
-                KEY_E164, AttributeValues.fromString(key),
+                KEY_PNI, AttributeValues.fromString(key),
                 ATTR_EXP, AttributeValues.fromLong(expirationSeconds),
                 ATTR_SALT, AttributeValues.fromString(salt),
                 ATTR_HASH, AttributeValues.fromString(hash)))
@@ -111,7 +105,7 @@ public class RegistrationRecoveryPasswords {
     return TransactWriteItem.builder()
         .delete(Delete.builder()
             .tableName(tableName)
-            .key(Map.of(KEY_E164, AttributeValues.fromString(key)))
+            .key(Map.of(KEY_PNI, AttributeValues.fromString(key)))
             .build())
         .build();
   }

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

@@ -10,6 +10,7 @@ import static java.util.Objects.requireNonNull;
 import java.lang.invoke.MethodHandles;
 import java.util.HexFormat;
 import java.util.Optional;
+import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -30,8 +31,8 @@ public class RegistrationRecoveryPasswordsManager {
     this.phoneNumberIdentifiers = phoneNumberIdentifiers;
   }
 
-  public CompletableFuture<Boolean> verify(final String number, final byte[] password) {
+  public CompletableFuture<Boolean> verify(final UUID phoneNumberIdentifier, final byte[] password) {
-    return registrationRecoveryPasswords.lookup(number)
+    return registrationRecoveryPasswords.lookup(phoneNumberIdentifier)
         .thenApply(maybeHash -> maybeHash.filter(hash -> hash.verify(bytesToString(password))))
         .whenComplete((result, error) -> {
           if (error != null) {

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

@@ -31,11 +31,11 @@ import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
 import org.whispersystems.textsecuregcm.entities.PhoneVerificationRequest;
+import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.push.NotPushRegisteredException;
 import org.whispersystems.textsecuregcm.push.PushNotificationManager;
-import org.whispersystems.textsecuregcm.push.WebSocketConnectionEventManager;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
@@ -53,7 +53,7 @@ class RegistrationLockVerificationManagerTest {
       ExternalServiceCredentialsGenerator.class);
   private final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager = mock(
       RegistrationRecoveryPasswordsManager.class);
-  private static PushNotificationManager pushNotificationManager = mock(PushNotificationManager.class);
+  private final PushNotificationManager pushNotificationManager = mock(PushNotificationManager.class);
   private final RateLimiters rateLimiters = mock(RateLimiters.class);
   private final RegistrationLockVerificationManager registrationLockVerificationManager = new RegistrationLockVerificationManager(
       accountsManager, disconnectionRequestManager, svr2CredentialsGenerator,
@@ -105,12 +105,13 @@ class RegistrationLockVerificationManagerTest {
             if (!verificationType.equals(PhoneVerificationRequest.VerificationType.RECOVERY_PASSWORD) || clientRegistrationLock != null) {
               verify(registrationRecoveryPasswordsManager).removeForNumber(account.getNumber());
             } else {
-              verify(registrationRecoveryPasswordsManager, never()).removeForNumber(account.getNumber());
+              verify(registrationRecoveryPasswordsManager, never()).removeForNumber(any());
             }
             verify(disconnectionRequestManager).requestDisconnection(account.getUuid(), List.of(Device.PRIMARY_ID));
             try {
               verify(pushNotificationManager).sendAttemptLoginNotification(any(), eq("failedRegistrationLock"));
-            } catch (NotPushRegisteredException npre) {}
+            } catch (final NotPushRegisteredException ignored) {
+            }
             if (alreadyLocked) {
               verify(account, never()).lockAuthTokenHash();
             } else {
@@ -126,10 +127,13 @@ class RegistrationLockVerificationManagerTest {
         doThrow(RateLimitExceededException.class).when(pinLimiter).validate(anyString());
         yield new Pair<>(RateLimitExceededException.class, ignored -> {
           verify(account, never()).lockAuthTokenHash();
+
           try {
             verify(pushNotificationManager, never()).sendAttemptLoginNotification(any(), eq("failedRegistrationLock"));
-          } catch (NotPushRegisteredException npre) {}
+          } catch (final NotPushRegisteredException ignored2) {
-          verify(registrationRecoveryPasswordsManager, never()).removeForNumber(account.getNumber());
+          }
+
+          verify(registrationRecoveryPasswordsManager, never()).removeForNumber(any());
           verify(disconnectionRequestManager, never()).requestDisconnection(any(), any());
         });
       }
@@ -167,7 +171,7 @@ class RegistrationLockVerificationManagerTest {
             PhoneVerificationRequest.VerificationType.SESSION));
 
     verify(account, never()).lockAuthTokenHash();
-    verify(registrationRecoveryPasswordsManager, never()).removeForNumber(account.getNumber());
+    verify(registrationRecoveryPasswordsManager, never()).removeForNumber(any());
     verify(disconnectionRequestManager, never()).requestDisconnection(any(), any());
   }
 

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

@@ -86,6 +86,7 @@ import org.whispersystems.textsecuregcm.storage.AccountBadge;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.ChangeNumberManager;
 import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
 import org.whispersystems.textsecuregcm.storage.RegistrationRecoveryPasswordsManager;
 import org.whispersystems.textsecuregcm.tests.util.AccountsHelper;
 import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
@@ -105,6 +106,7 @@ class AccountControllerV2Test {
 
   private final AccountsManager accountsManager = mock(AccountsManager.class);
   private final ChangeNumberManager changeNumberManager = mock(ChangeNumberManager.class);
+  private final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
   private final RegistrationServiceClient registrationServiceClient = mock(RegistrationServiceClient.class);
   private final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager = mock(
       RegistrationRecoveryPasswordsManager.class);
@@ -125,8 +127,8 @@ class AccountControllerV2Test {
       .setTestContainerFactory(new GrizzlyWebTestContainerFactory())
       .addResource(
           new AccountControllerV2(accountsManager, changeNumberManager,
-              new PhoneVerificationTokenManager(registrationServiceClient, registrationRecoveryPasswordsManager,
+              new PhoneVerificationTokenManager(phoneNumberIdentifiers, registrationServiceClient,
-                  registrationRecoveryChecker),
+                  registrationRecoveryPasswordsManager, registrationRecoveryChecker),
               registrationLockVerificationManager, rateLimiters))
       .build();
 
@@ -401,6 +403,8 @@ class AccountControllerV2Test {
 
     @Test
     void recoveryPasswordManagerVerificationTrue() throws Exception {
+      when(phoneNumberIdentifiers.getPhoneNumberIdentifier(any()))
+          .thenReturn(CompletableFuture.completedFuture(UUID.randomUUID()));
       when(registrationRecoveryPasswordsManager.verify(any(), any()))
           .thenReturn(CompletableFuture.completedFuture(true));
       when(registrationRecoveryChecker.checkRegistrationRecoveryAttempt(any(), any()))
@@ -443,6 +447,8 @@ class AccountControllerV2Test {
 
     @Test
     void registrationRecoveryCheckerAllowsAttempt() {
+      when(phoneNumberIdentifiers.getPhoneNumberIdentifier(any()))
+          .thenReturn(CompletableFuture.completedFuture(UUID.randomUUID()));
       when(registrationRecoveryChecker.checkRegistrationRecoveryAttempt(any(), any())).thenReturn(true);
       when(registrationRecoveryPasswordsManager.verify(any(), any()))
           .thenReturn(CompletableFuture.completedFuture(true));

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

@@ -7,6 +7,7 @@ package org.whispersystems.textsecuregcm.controllers;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doThrow;
@@ -77,6 +78,7 @@ import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.DeviceCapability;
 import org.whispersystems.textsecuregcm.storage.DeviceSpec;
+import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
 import org.whispersystems.textsecuregcm.storage.RegistrationRecoveryPasswordsManager;
 import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
 import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
@@ -94,6 +96,7 @@ class RegistrationControllerTest {
   private static final String PASSWORD = "password";
 
   private final AccountsManager accountsManager = mock(AccountsManager.class);
+  private final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
   private final RegistrationServiceClient registrationServiceClient = mock(RegistrationServiceClient.class);
   private final RegistrationLockVerificationManager registrationLockVerificationManager = mock(
       RegistrationLockVerificationManager.class);
@@ -113,8 +116,8 @@ class RegistrationControllerTest {
       .setTestContainerFactory(new GrizzlyWebTestContainerFactory())
       .addResource(
           new RegistrationController(accountsManager,
-              new PhoneVerificationTokenManager(registrationServiceClient, registrationRecoveryPasswordsManager,
+              new PhoneVerificationTokenManager(phoneNumberIdentifiers, registrationServiceClient,
-                  registrationRecoveryChecker),
+                  registrationRecoveryPasswordsManager, registrationRecoveryChecker),
               registrationLockVerificationManager, rateLimiters))
       .build();
 
@@ -243,6 +246,8 @@ class RegistrationControllerTest {
 
   @Test
   void recoveryPasswordManagerVerificationFailureOrTimeout() {
+    when(phoneNumberIdentifiers.getPhoneNumberIdentifier(any()))
+        .thenReturn(CompletableFuture.completedFuture(UUID.randomUUID()));
     when(registrationRecoveryChecker.checkRegistrationRecoveryAttempt(any(), any())).thenReturn(true);
     when(registrationRecoveryPasswordsManager.verify(any(), any()))
         .thenReturn(CompletableFuture.failedFuture(new RuntimeException()));
@@ -289,6 +294,8 @@ class RegistrationControllerTest {
 
   @Test
   void recoveryPasswordManagerVerificationTrue() throws InterruptedException {
+    when(phoneNumberIdentifiers.getPhoneNumberIdentifier(any()))
+        .thenReturn(CompletableFuture.completedFuture(UUID.randomUUID()));
     when(registrationRecoveryChecker.checkRegistrationRecoveryAttempt(any(), any())).thenReturn(true);
     when(registrationRecoveryPasswordsManager.verify(any(), any()))
         .thenReturn(CompletableFuture.completedFuture(true));
@@ -325,6 +332,8 @@ class RegistrationControllerTest {
 
   @Test
   void registrationRecoveryCheckerAllowsAttempt() throws InterruptedException {
+    when(phoneNumberIdentifiers.getPhoneNumberIdentifier(any()))
+        .thenReturn(CompletableFuture.completedFuture(UUID.randomUUID()));
     when(registrationRecoveryChecker.checkRegistrationRecoveryAttempt(any(), any())).thenReturn(true);
     when(registrationRecoveryPasswordsManager.verify(any(), any()))
         .thenReturn(CompletableFuture.completedFuture(true));

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

@@ -39,6 +39,7 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.stream.Stream;
@@ -77,6 +78,7 @@ import org.whispersystems.textsecuregcm.spam.RegistrationFraudChecker;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
+import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
 import org.whispersystems.textsecuregcm.storage.RegistrationRecoveryPasswordsManager;
 import org.whispersystems.textsecuregcm.storage.VerificationSessionManager;
 import org.whispersystems.textsecuregcm.util.SystemMapper;
@@ -87,7 +89,11 @@ class VerificationControllerTest {
   private static final long SESSION_EXPIRATION_SECONDS = Duration.ofMinutes(10).toSeconds();
 
   private static final byte[] SESSION_ID = "session".getBytes(StandardCharsets.UTF_8);
-  private static final String NUMBER = "+18005551212";
+  private static final String NUMBER = PhoneNumberUtil.getInstance().format(
+      PhoneNumberUtil.getInstance().getExampleNumber("US"),
+      PhoneNumberUtil.PhoneNumberFormat.E164);
+
+  private static final UUID PNI = UUID.randomUUID();
 
   private final RegistrationServiceClient registrationServiceClient = mock(RegistrationServiceClient.class);
   private final VerificationSessionManager verificationSessionManager = mock(VerificationSessionManager.class);
@@ -95,6 +101,7 @@ class VerificationControllerTest {
   private final RegistrationCaptchaManager registrationCaptchaManager = mock(RegistrationCaptchaManager.class);
   private final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager = mock(
       RegistrationRecoveryPasswordsManager.class);
+  private final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
   private final RateLimiters rateLimiters = mock(RateLimiters.class);
   private final AccountsManager accountsManager = mock(AccountsManager.class);
   private final Clock clock = Clock.systemUTC();
@@ -115,7 +122,7 @@ class VerificationControllerTest {
       .setTestContainerFactory(new GrizzlyWebTestContainerFactory())
       .addResource(
           new VerificationController(registrationServiceClient, verificationSessionManager, pushNotificationManager,
-              registrationCaptchaManager, registrationRecoveryPasswordsManager, rateLimiters, accountsManager,
+              registrationCaptchaManager, registrationRecoveryPasswordsManager, phoneNumberIdentifiers, rateLimiters, accountsManager,
               RegistrationFraudChecker.noop(), dynamicConfigurationManager, clock))
       .build();
 
@@ -131,6 +138,8 @@ class VerificationControllerTest {
         .thenReturn(new DynamicRegistrationConfiguration(false));
     when(dynamicConfigurationManager.getConfiguration())
         .thenReturn(dynamicConfiguration);
+    when(phoneNumberIdentifiers.getPhoneNumberIdentifier(NUMBER))
+        .thenReturn(CompletableFuture.completedFuture(PNI));
   }
 
   @ParameterizedTest

service/src/test/java/org/whispersystems/textsecuregcm/grpc/AccountsGrpcServiceTest.java

@@ -65,6 +65,7 @@ import org.whispersystems.textsecuregcm.controllers.AccountController;
 import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
 import org.whispersystems.textsecuregcm.entities.EncryptedUsername;
 import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
+import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.identity.PniServiceIdentifier;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
@@ -113,7 +114,7 @@ class AccountsGrpcServiceTest extends SimpleBaseGrpcTest<AccountsGrpcService, Ac
     when(rateLimiter.validateReactive(any(UUID.class))).thenReturn(Mono.empty());
     when(rateLimiter.validateReactive(anyString())).thenReturn(Mono.empty());
 
-    when(registrationRecoveryPasswordsManager.storeForCurrentNumber(anyString(), any()))
+    when(registrationRecoveryPasswordsManager.storeForCurrentNumber(any(), any()))
         .thenReturn(CompletableFuture.completedFuture(null));
 
     return new AccountsGrpcService(accountsManager,
@@ -261,7 +262,7 @@ class AccountsGrpcServiceTest extends SimpleBaseGrpcTest<AccountsGrpcService, Ac
           final List<byte[]> usernameHashes = invocation.getArgument(1);
 
           return CompletableFuture.completedFuture(
-              new AccountsManager.UsernameReservation(invocation.getArgument(0), usernameHashes.get(0)));
+              new AccountsManager.UsernameReservation(invocation.getArgument(0), usernameHashes.getFirst()));
         });
 
     final ReserveUsernameHashResponse expectedResponse = ReserveUsernameHashResponse.newBuilder()
@@ -684,12 +685,10 @@ class AccountsGrpcServiceTest extends SimpleBaseGrpcTest<AccountsGrpcService, Ac
 
   @Test
   void setRegistrationRecoveryPassword() {
-    final String phoneNumber =
+    final UUID phoneNumberIdentifier = UUID.randomUUID();
-        PhoneNumberUtil.getInstance().format(PhoneNumberUtil.getInstance().getExampleNumber("US"),
-            PhoneNumberUtil.PhoneNumberFormat.E164);
 
     final Account account = mock(Account.class);
-    when(account.getNumber()).thenReturn(phoneNumber);
+    when(account.getIdentifier(IdentityType.PNI)).thenReturn(phoneNumberIdentifier);
 
     when(accountsManager.getByAccountIdentifierAsync(AUTHENTICATED_ACI))
         .thenReturn(CompletableFuture.completedFuture(Optional.of(account)));
@@ -701,7 +700,7 @@ class AccountsGrpcServiceTest extends SimpleBaseGrpcTest<AccountsGrpcService, Ac
             .setRegistrationRecoveryPassword(ByteString.copyFrom(registrationRecoveryPassword))
             .build()));
 
-    verify(registrationRecoveryPasswordsManager).storeForCurrentNumber(phoneNumber, registrationRecoveryPassword);
+    verify(registrationRecoveryPasswordsManager).storeForCurrentNumber(account.getNumber(), registrationRecoveryPassword);
   }
 
   @Test

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

@@ -224,7 +224,7 @@ class AccountsManagerTest {
     final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager =
         mock(RegistrationRecoveryPasswordsManager.class);
 
-    when(registrationRecoveryPasswordsManager.removeForNumber(anyString())).thenReturn(CompletableFuture.completedFuture(null));
+    when(registrationRecoveryPasswordsManager.removeForNumber(any())).thenReturn(CompletableFuture.completedFuture(null));
 
     when(keysManager.deleteSingleUsePreKeys(any())).thenReturn(CompletableFuture.completedFuture(null));
     when(messagesManager.clear(any())).thenReturn(CompletableFuture.completedFuture(null));

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

@@ -288,10 +288,10 @@ public final class DynamoDbExtensionSchema {
         List.of(), List.of()),
 
     REGISTRATION_RECOVERY_PASSWORDS("registration_recovery_passwords_test",
-        RegistrationRecoveryPasswords.KEY_E164,
+        RegistrationRecoveryPasswords.KEY_PNI,
         null,
         List.of(AttributeDefinition.builder()
-            .attributeName(RegistrationRecoveryPasswords.KEY_E164)
+            .attributeName(RegistrationRecoveryPasswords.KEY_PNI)
             .attributeType(ScalarAttributeType.S)
             .build()),
         List.of(), List.of()),

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

@@ -32,20 +32,17 @@ import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
 public class RegistrationRecoveryTest {
 
   private static final MutableClock CLOCK = MockUtils.mutableClock(0);
-
   private static final Duration EXPIRATION = Duration.ofSeconds(1000);
-
   private static final String NUMBER = "+18005555555";
-  private static final UUID PNI = UUID.randomUUID();
 
   private static final SaltedTokenHash ORIGINAL_HASH = SaltedTokenHash.generateFor("pass1");
-
   private static final SaltedTokenHash ANOTHER_HASH = SaltedTokenHash.generateFor("pass2");
 
   @RegisterExtension
-  private static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension(
+  private static final DynamoDbExtension DYNAMO_DB_EXTENSION =
-      Tables.PNI,
+      new DynamoDbExtension(Tables.PNI, Tables.REGISTRATION_RECOVERY_PASSWORDS);
-      Tables.REGISTRATION_RECOVERY_PASSWORDS);
+
+  private UUID pni;
 
   private RegistrationRecoveryPasswords registrationRecoveryPasswords;
 
@@ -60,89 +57,63 @@ public class RegistrationRecoveryTest {
         DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
         CLOCK
     );
-
     final PhoneNumberIdentifiers phoneNumberIdentifiers =
         new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(), Tables.PNI.tableName());
 
     manager = new RegistrationRecoveryPasswordsManager(registrationRecoveryPasswords, phoneNumberIdentifiers);
+
+    pni = phoneNumberIdentifiers.getPhoneNumberIdentifier(NUMBER).join();
   }
 
   @Test
   public void testLookupAfterWrite() throws Exception {
-    registrationRecoveryPasswords.addOrReplace(NUMBER, PNI, ORIGINAL_HASH).get();
+    registrationRecoveryPasswords.addOrReplace(NUMBER, pni, ORIGINAL_HASH).get();
     final long initialExp = fetchTimestamp(NUMBER);
     final long expectedExpiration = CLOCK.instant().getEpochSecond() + EXPIRATION.getSeconds();
     assertEquals(expectedExpiration, initialExp);
 
-    {
+    final Optional<SaltedTokenHash> saltedTokenHashByPni = registrationRecoveryPasswords.lookup(pni).get();
-      final Optional<SaltedTokenHash> saltedTokenHashByNumber = registrationRecoveryPasswords.lookup(NUMBER).get();
+    assertTrue(saltedTokenHashByPni.isPresent());
-      assertTrue(saltedTokenHashByNumber.isPresent());
+    assertEquals(ORIGINAL_HASH.salt(), saltedTokenHashByPni.get().salt());
-      assertEquals(ORIGINAL_HASH.salt(), saltedTokenHashByNumber.get().salt());
+    assertEquals(ORIGINAL_HASH.hash(), saltedTokenHashByPni.get().hash());
-      assertEquals(ORIGINAL_HASH.hash(), saltedTokenHashByNumber.get().hash());
-    }
-
-    {
-      final Optional<SaltedTokenHash> saltedTokenHashByPni = registrationRecoveryPasswords.lookup(PNI).get();
-      assertTrue(saltedTokenHashByPni.isPresent());
-      assertEquals(ORIGINAL_HASH.salt(), saltedTokenHashByPni.get().salt());
-      assertEquals(ORIGINAL_HASH.hash(), saltedTokenHashByPni.get().hash());
-    }
   }
 
   @Test
   public void testLookupAfterRefresh() throws Exception {
-    registrationRecoveryPasswords.addOrReplace(NUMBER, PNI, ORIGINAL_HASH).get();
+    registrationRecoveryPasswords.addOrReplace(NUMBER, pni, ORIGINAL_HASH).get();
 
     CLOCK.increment(50, TimeUnit.SECONDS);
-    registrationRecoveryPasswords.addOrReplace(NUMBER, PNI, ORIGINAL_HASH).get();
+    registrationRecoveryPasswords.addOrReplace(NUMBER, pni, ORIGINAL_HASH).get();
     final long updatedExp = fetchTimestamp(NUMBER);
     final long expectedExp = CLOCK.instant().getEpochSecond() + EXPIRATION.getSeconds();
     assertEquals(expectedExp, updatedExp);
 
-    {
+    final Optional<SaltedTokenHash> saltedTokenHashByPni = registrationRecoveryPasswords.lookup(pni).get();
-      final Optional<SaltedTokenHash> saltedTokenHashByNumber = registrationRecoveryPasswords.lookup(NUMBER).get();
+    assertTrue(saltedTokenHashByPni.isPresent());
-      assertTrue(saltedTokenHashByNumber.isPresent());
+    assertEquals(ORIGINAL_HASH.salt(), saltedTokenHashByPni.get().salt());
-      assertEquals(ORIGINAL_HASH.salt(), saltedTokenHashByNumber.get().salt());
+    assertEquals(ORIGINAL_HASH.hash(), saltedTokenHashByPni.get().hash());
-      assertEquals(ORIGINAL_HASH.hash(), saltedTokenHashByNumber.get().hash());
-    }
-
-    {
-      final Optional<SaltedTokenHash> saltedTokenHashByPni = registrationRecoveryPasswords.lookup(PNI).get();
-      assertTrue(saltedTokenHashByPni.isPresent());
-      assertEquals(ORIGINAL_HASH.salt(), saltedTokenHashByPni.get().salt());
-      assertEquals(ORIGINAL_HASH.hash(), saltedTokenHashByPni.get().hash());
-    }
   }
 
   @Test
   public void testReplace() throws Exception {
-    registrationRecoveryPasswords.addOrReplace(NUMBER, PNI, ORIGINAL_HASH).get();
+    registrationRecoveryPasswords.addOrReplace(NUMBER, pni, ORIGINAL_HASH).get();
-    registrationRecoveryPasswords.addOrReplace(NUMBER, PNI, ANOTHER_HASH).get();
+    registrationRecoveryPasswords.addOrReplace(NUMBER, pni, ANOTHER_HASH).get();
 
-    {
+    final Optional<SaltedTokenHash> saltedTokenHashByPni = registrationRecoveryPasswords.lookup(pni).get();
-      final Optional<SaltedTokenHash> saltedTokenHashByNumber = registrationRecoveryPasswords.lookup(NUMBER).get();
+    assertTrue(saltedTokenHashByPni.isPresent());
-      assertTrue(saltedTokenHashByNumber.isPresent());
+    assertEquals(ANOTHER_HASH.salt(), saltedTokenHashByPni.get().salt());
-      assertEquals(ANOTHER_HASH.salt(), saltedTokenHashByNumber.get().salt());
+    assertEquals(ANOTHER_HASH.hash(), saltedTokenHashByPni.get().hash());
-      assertEquals(ANOTHER_HASH.hash(), saltedTokenHashByNumber.get().hash());
-    }
-
-    {
-      final Optional<SaltedTokenHash> saltedTokenHashByPni = registrationRecoveryPasswords.lookup(PNI).get();
-      assertTrue(saltedTokenHashByPni.isPresent());
-      assertEquals(ANOTHER_HASH.salt(), saltedTokenHashByPni.get().salt());
-      assertEquals(ANOTHER_HASH.hash(), saltedTokenHashByPni.get().hash());
-    }
   }
 
   @Test
   public void testRemove() throws Exception {
-    assertDoesNotThrow(() -> registrationRecoveryPasswords.removeEntry(NUMBER, PNI).join());
+    assertDoesNotThrow(() -> registrationRecoveryPasswords.removeEntry(NUMBER, pni).join());
 
-    registrationRecoveryPasswords.addOrReplace(NUMBER, PNI, ORIGINAL_HASH).get();
+    registrationRecoveryPasswords.addOrReplace(NUMBER, pni, ORIGINAL_HASH).get();
-    assertTrue(registrationRecoveryPasswords.lookup(NUMBER).get().isPresent());
+    assertTrue(registrationRecoveryPasswords.lookup(pni).get().isPresent());
 
-    registrationRecoveryPasswords.removeEntry(NUMBER, PNI).get();
+    registrationRecoveryPasswords.removeEntry(NUMBER, pni).get();
-    assertTrue(registrationRecoveryPasswords.lookup(NUMBER).get().isEmpty());
+    assertTrue(registrationRecoveryPasswords.lookup(pni).get().isEmpty());
   }
 
   @Test
@@ -153,30 +124,30 @@ public class RegistrationRecoveryTest {
 
     // initial store
     manager.storeForCurrentNumber(NUMBER, password).get();
-    assertTrue(manager.verify(NUMBER, password).get());
+    assertTrue(manager.verify(pni, password).get());
-    assertFalse(manager.verify(NUMBER, wrongPassword).get());
+    assertFalse(manager.verify(pni, wrongPassword).get());
 
     // update
     manager.storeForCurrentNumber(NUMBER, password).get();
-    assertTrue(manager.verify(NUMBER, password).get());
+    assertTrue(manager.verify(pni, password).get());
-    assertFalse(manager.verify(NUMBER, wrongPassword).get());
+    assertFalse(manager.verify(pni, wrongPassword).get());
 
     // replace
     manager.storeForCurrentNumber(NUMBER, updatedPassword).get();
-    assertTrue(manager.verify(NUMBER, updatedPassword).get());
+    assertTrue(manager.verify(pni, updatedPassword).get());
-    assertFalse(manager.verify(NUMBER, password).get());
+    assertFalse(manager.verify(pni, password).get());
-    assertFalse(manager.verify(NUMBER, wrongPassword).get());
+    assertFalse(manager.verify(pni, wrongPassword).get());
 
     manager.removeForNumber(NUMBER).get();
-    assertFalse(manager.verify(NUMBER, updatedPassword).get());
+    assertFalse(manager.verify(pni, updatedPassword).get());
-    assertFalse(manager.verify(NUMBER, password).get());
+    assertFalse(manager.verify(pni, password).get());
-    assertFalse(manager.verify(NUMBER, wrongPassword).get());
+    assertFalse(manager.verify(pni, wrongPassword).get());
   }
 
   private static long fetchTimestamp(final String number) throws ExecutionException, InterruptedException {
     return DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient().getItem(GetItemRequest.builder()
             .tableName(Tables.REGISTRATION_RECOVERY_PASSWORDS.tableName())
-            .key(Map.of(RegistrationRecoveryPasswords.KEY_E164, AttributeValues.fromString(number)))
+            .key(Map.of(RegistrationRecoveryPasswords.KEY_PNI, AttributeValues.fromString(number)))
             .build())
         .thenApply(getItemResponse -> {
           final Map<String, AttributeValue> item = getItemResponse.item();

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

@@ -85,6 +85,7 @@ public class AuthHelper {
 
   public static final String UNDISCOVERABLE_NUMBER   = "+18005551234";
   public static final UUID   UNDISCOVERABLE_UUID     = UUID.randomUUID();
+  public static final UUID   UNDISCOVERABLE_PNI      = UUID.randomUUID();
   public static final String UNDISCOVERABLE_PASSWORD = "IT'S A SECRET TO EVERYBODY.";
 
   public static final ECKeyPair VALID_IDENTITY_KEY_PAIR = Curve.generateKeyPair();
@@ -169,7 +170,9 @@ public class AuthHelper {
     when(VALID_ACCOUNT_TWO.getPhoneNumberIdentifier()).thenReturn(VALID_PNI_TWO);
     when(UNDISCOVERABLE_ACCOUNT.getNumber()).thenReturn(UNDISCOVERABLE_NUMBER);
     when(UNDISCOVERABLE_ACCOUNT.getUuid()).thenReturn(UNDISCOVERABLE_UUID);
+    when(UNDISCOVERABLE_ACCOUNT.getPhoneNumberIdentifier()).thenReturn(UNDISCOVERABLE_PNI);
     when(UNDISCOVERABLE_ACCOUNT.getIdentifier(IdentityType.ACI)).thenReturn(UNDISCOVERABLE_UUID);
+    when(UNDISCOVERABLE_ACCOUNT.getIdentifier(IdentityType.PNI)).thenReturn(UNDISCOVERABLE_PNI);
     when(VALID_ACCOUNT_3.getNumber()).thenReturn(VALID_NUMBER_3);
     when(VALID_ACCOUNT_3.getUuid()).thenReturn(VALID_UUID_3);
     when(VALID_ACCOUNT_3.getPhoneNumberIdentifier()).thenReturn(VALID_PNI_3);