Write registration recovery passwords exclusively by PNI

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

@@ -13,7 +13,6 @@ import java.util.concurrent.CompletableFuture;
 import org.signal.integration.config.Config;
 import org.whispersystems.textsecuregcm.metrics.NoopAwsSdkMetricPublisher;
 import org.whispersystems.textsecuregcm.registration.VerificationSession;
-import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
 import org.whispersystems.textsecuregcm.storage.RegistrationRecoveryPasswords;
 import org.whispersystems.textsecuregcm.storage.RegistrationRecoveryPasswordsManager;
 import org.whispersystems.textsecuregcm.storage.VerificationSessionManager;
@@ -35,9 +34,6 @@ public class IntegrationTools {
     final DynamoDbAsyncClient dynamoDbAsyncClient =
         config.dynamoDbClient().buildAsyncClient(credentialsProvider, new NoopAwsSdkMetricPublisher());
 
-    final PhoneNumberIdentifiers phoneNumberIdentifiers =
-        new PhoneNumberIdentifiers(dynamoDbAsyncClient, config.dynamoDbTables().phoneNumberIdentifiers());
-
     final RegistrationRecoveryPasswords registrationRecoveryPasswords = new RegistrationRecoveryPasswords(
         config.dynamoDbTables().registrationRecovery(), Duration.ofDays(1), dynamoDbAsyncClient, Clock.systemUTC());
 
@@ -45,7 +41,7 @@ public class IntegrationTools {
         dynamoDbAsyncClient, config.dynamoDbTables().verificationSessions(), Clock.systemUTC());
 
     return new IntegrationTools(
-        new RegistrationRecoveryPasswordsManager(registrationRecoveryPasswords, phoneNumberIdentifiers),
+        new RegistrationRecoveryPasswordsManager(registrationRecoveryPasswords),
         new VerificationSessionManager(verificationSessions)
     );
   }
@@ -57,8 +53,8 @@ public class IntegrationTools {
     this.verificationSessionManager = verificationSessionManager;
   }
 
-  public CompletableFuture<Void> populateRecoveryPassword(final String number, final byte[] password) {
+  public CompletableFuture<Void> populateRecoveryPassword(final UUID phoneNumberIdentifier, final byte[] password) {
-    return registrationRecoveryPasswordsManager.storeForCurrentNumber(number, password);
+    return registrationRecoveryPasswordsManager.store(phoneNumberIdentifier, 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(user.phoneNumber(), registrationPassword).join();
+    INTEGRATION_TOOLS.populateRecoveryPassword(user.pniUuid(), registrationPassword).join();
 
     final ECKeyPair aciIdentityKeyPair = Curve.generateKeyPair();
     final ECKeyPair pniIdentityKeyPair = Curve.generateKeyPair();

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

@@ -586,7 +586,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     ExperimentEnrollmentManager experimentEnrollmentManager = new ExperimentEnrollmentManager(
         dynamicConfigurationManager);
     RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager =
-        new RegistrationRecoveryPasswordsManager(registrationRecoveryPasswords, phoneNumberIdentifiers);
+        new RegistrationRecoveryPasswordsManager(registrationRecoveryPasswords);
     UsernameHashZkProofVerifier usernameHashZkProofVerifier = new UsernameHashZkProofVerifier();
 
     RegistrationServiceClient registrationServiceClient = config.getRegistrationServiceConfiguration()

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

@@ -157,7 +157,7 @@ public class RegistrationLockVerificationManager {
       // This allows users to re-register via registration recovery password
       // instead of always being forced to fall back to SMS verification.
       if (!phoneVerificationType.equals(PhoneVerificationRequest.VerificationType.RECOVERY_PASSWORD) || clientRegistrationLock != null) {
-        registrationRecoveryPasswordsManager.removeForNumber(updatedAccount.getNumber());
+        registrationRecoveryPasswordsManager.remove(updatedAccount.getIdentifier(IdentityType.PNI));
       }
 
       final List<Byte> deviceIds = updatedAccount.getDevices().stream().map(Device::getId).toList();

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

@@ -54,6 +54,7 @@ import org.whispersystems.textsecuregcm.entities.ReserveUsernameHashResponse;
 import org.whispersystems.textsecuregcm.entities.UsernameHashResponse;
 import org.whispersystems.textsecuregcm.entities.UsernameLinkHandle;
 import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
+import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.identity.ServiceIdentifier;
 import org.whispersystems.textsecuregcm.limits.RateLimitedByIp;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
@@ -259,7 +260,7 @@ public class AccountController {
 
     // if registration recovery password was sent to us, store it (or refresh its expiration)
     attributes.recoveryPassword().ifPresent(registrationRecoveryPassword ->
-        registrationRecoveryPasswordsManager.storeForCurrentNumber(updatedAccount.getNumber(), registrationRecoveryPassword));
+        registrationRecoveryPasswordsManager.store(updatedAccount.getIdentifier(IdentityType.PNI), registrationRecoveryPassword));
   }
 
   @GET

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

@@ -602,7 +602,7 @@ public class VerificationController {
     }
 
     if (resultSession.verified()) {
-      registrationRecoveryPasswordsManager.removeForNumber(registrationServiceSession.number());
+      registrationRecoveryPasswordsManager.remove(phoneNumberIdentifiers.getPhoneNumberIdentifier(registrationServiceSession.number()).join());
     }
 
     Metrics.counter(VERIFIED_COUNTER_NAME, Tags.of(
@@ -648,7 +648,7 @@ public class VerificationController {
           .orElseThrow(NotFoundException::new);
 
       if (registrationServiceSession.verified()) {
-        registrationRecoveryPasswordsManager.removeForNumber(registrationServiceSession.number());
+        registrationRecoveryPasswordsManager.remove(phoneNumberIdentifiers.getPhoneNumberIdentifier(registrationServiceSession.number()).join());
       }
 
       return registrationServiceSession;

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

@@ -337,7 +337,7 @@ public class AccountsGrpcService extends ReactorAccountsGrpc.AccountsImplBase {
 
     return Mono.fromFuture(() -> accountsManager.getByAccountIdentifierAsync(authenticatedDevice.accountIdentifier()))
         .map(maybeAccount -> maybeAccount.orElseThrow(Status.UNAUTHENTICATED::asRuntimeException))
-        .flatMap(account -> Mono.fromFuture(() -> registrationRecoveryPasswordsManager.storeForCurrentNumber(account.getNumber(), request.getRegistrationRecoveryPassword().toByteArray())))
+        .flatMap(account -> Mono.fromFuture(() -> registrationRecoveryPasswordsManager.store(account.getIdentifier(IdentityType.PNI), request.getRegistrationRecoveryPassword().toByteArray())))
         .thenReturn(SetRegistrationRecoveryPasswordResponse.newBuilder().build());
   }
 }

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

@@ -385,7 +385,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
         Metrics.counter(CREATE_COUNTER_NAME, tags).increment();
 
         accountAttributes.recoveryPassword().ifPresent(registrationRecoveryPassword ->
-            registrationRecoveryPasswordsManager.storeForCurrentNumber(account.getNumber(),
+            registrationRecoveryPasswordsManager.store(account.getIdentifier(IdentityType.PNI),
                 registrationRecoveryPassword));
       }, accountLockExecutor);
 
@@ -1279,7 +1279,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
             keysManager.deleteSingleUsePreKeys(account.getPhoneNumberIdentifier()),
             messagesManager.clear(account.getUuid()),
             profilesManager.deleteAll(account.getUuid()),
-            registrationRecoveryPasswordsManager.removeForNumber(account.getNumber()))
+            registrationRecoveryPasswordsManager.remove(account.getIdentifier(IdentityType.PNI)))
         .thenCompose(ignored -> accounts.delete(account.getUuid(), additionalWriteItems))
         .thenCompose(ignored -> redisDeleteAsync(account))
         .thenRun(() -> disconnectionRequestManager.requestDisconnection(account.getUuid()));

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

@@ -19,11 +19,9 @@ import org.whispersystems.textsecuregcm.util.AttributeValues;
 import org.whispersystems.textsecuregcm.util.Util;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
-import software.amazon.awssdk.services.dynamodb.model.Delete;
+import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
-import software.amazon.awssdk.services.dynamodb.model.Put;
+import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
-import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
-import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
 
 public class RegistrationRecoveryPasswords {
 
@@ -64,50 +62,26 @@ public class RegistrationRecoveryPasswords {
             .map(RegistrationRecoveryPasswords::saltedTokenHashFromItem));
   }
 
-  public CompletableFuture<Void> addOrReplace(final String number, final UUID phoneNumberIdentifier, final SaltedTokenHash data) {
+  public CompletableFuture<Void> addOrReplace(final UUID phoneNumberIdentifier, final SaltedTokenHash data) {
     final long expirationSeconds = expirationSeconds();
 
-    return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder()
+    return asyncClient.putItem(PutItemRequest.builder()
-            .transactItems(
-                buildPutRecoveryPasswordWriteItem(number, expirationSeconds, data.salt(), data.hash()),
-                buildPutRecoveryPasswordWriteItem(phoneNumberIdentifier.toString(), expirationSeconds, data.salt(), data.hash()))
-        .build())
-        .thenRun(Util.NOOP);
-  }
-
-  private TransactWriteItem buildPutRecoveryPasswordWriteItem(final String key,
-      final long expirationSeconds,
-      final String salt,
-      final String hash) {
-
-    return TransactWriteItem.builder()
-        .put(Put.builder()
             .tableName(tableName)
             .item(Map.of(
-                KEY_PNI, AttributeValues.fromString(key),
+                KEY_PNI, AttributeValues.fromString(phoneNumberIdentifier.toString()),
                 ATTR_EXP, AttributeValues.fromLong(expirationSeconds),
-                ATTR_SALT, AttributeValues.fromString(salt),
+                ATTR_SALT, AttributeValues.fromString(data.salt()),
-                ATTR_HASH, AttributeValues.fromString(hash)))
+                ATTR_HASH, AttributeValues.fromString(data.hash())))
             .build())
-        .build();
-  }
-
-  public CompletableFuture<Void> removeEntry(final String number, final UUID phoneNumberIdentifier) {
-    return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder()
-        .transactItems(
-            buildDeleteRecoveryPasswordWriteItem(number),
-            buildDeleteRecoveryPasswordWriteItem(phoneNumberIdentifier.toString()))
-        .build())
         .thenRun(Util.NOOP);
   }
 
-  private TransactWriteItem buildDeleteRecoveryPasswordWriteItem(final String key) {
+  public CompletableFuture<Void> removeEntry(final UUID phoneNumberIdentifier) {
-    return TransactWriteItem.builder()
+    return asyncClient.deleteItem(DeleteItemRequest.builder()
-        .delete(Delete.builder()
             .tableName(tableName)
-            .key(Map.of(KEY_PNI, AttributeValues.fromString(key)))
+            .key(Map.of(KEY_PNI, AttributeValues.fromString(phoneNumberIdentifier.toString())))
             .build())
-        .build();
+        .thenRun(Util.NOOP);
   }
 
   @VisibleForTesting

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

@@ -22,13 +22,9 @@ public class RegistrationRecoveryPasswordsManager {
   private static final Logger logger = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
 
   private final RegistrationRecoveryPasswords registrationRecoveryPasswords;
-  private final PhoneNumberIdentifiers phoneNumberIdentifiers;
-
-  public RegistrationRecoveryPasswordsManager(final RegistrationRecoveryPasswords registrationRecoveryPasswords,
-      final PhoneNumberIdentifiers phoneNumberIdentifiers) {
 
+  public RegistrationRecoveryPasswordsManager(final RegistrationRecoveryPasswords registrationRecoveryPasswords) {
     this.registrationRecoveryPasswords = requireNonNull(registrationRecoveryPasswords);
-    this.phoneNumberIdentifiers = phoneNumberIdentifiers;
   }
 
   public CompletableFuture<Boolean> verify(final UUID phoneNumberIdentifier, final byte[] password) {
@@ -42,30 +38,28 @@ public class RegistrationRecoveryPasswordsManager {
         .thenApply(Optional::isPresent);
   }
 
-  public CompletableFuture<Void> storeForCurrentNumber(final String number, final byte[] password) {
+  public CompletableFuture<Void> store(final UUID phoneNumberIdentifier, final byte[] password) {
     final String token = bytesToString(password);
     final SaltedTokenHash tokenHash = SaltedTokenHash.generateFor(token);
 
-    return phoneNumberIdentifiers.getPhoneNumberIdentifier(number)
+    return registrationRecoveryPasswords.addOrReplace(phoneNumberIdentifier, tokenHash)
-        .thenCompose(phoneNumberIdentifier -> registrationRecoveryPasswords.addOrReplace(number, phoneNumberIdentifier, tokenHash)
+        .whenComplete((result, error) -> {
-            .whenComplete((result, error) -> {
+          if (error != null) {
-              if (error != null) {
+            logger.warn("Failed to store Registration Recovery Password", error);
-                logger.warn("Failed to store Registration Recovery Password", error);
+          }
-              }
+        });
-            }));
   }
 
-  public CompletableFuture<Void> removeForNumber(final String number) {
+  public CompletableFuture<Void> remove(final UUID phoneNumberIdentifier) {
-    return phoneNumberIdentifiers.getPhoneNumberIdentifier(number)
+    return registrationRecoveryPasswords.removeEntry(phoneNumberIdentifier)
-        .thenCompose(phoneNumberIdentifier -> registrationRecoveryPasswords.removeEntry(number, phoneNumberIdentifier)
+        .whenComplete((ignored, error) -> {
-            .whenComplete((ignored, error) -> {
+          if (error instanceof ResourceNotFoundException) {
-              if (error instanceof ResourceNotFoundException) {
+            // These will naturally happen if a recovery password is already deleted. Since we can remove
-                // These will naturally happen if a recovery password is already deleted. Since we can remove
+            // the recovery password through many flows, we avoid creating log messages for these exceptions
-                // the recovery password through many flows, we avoid creating log messages for these exceptions
+          } else if (error != null) {
-              } else if (error != null) {
+            logger.warn("Failed to remove Registration Recovery Password", error);
-                logger.warn("Failed to remove Registration Recovery Password", error);
+          }
-              }
+        });
-            }));
   }
 
   private static String bytesToString(final byte[] bytes) {

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

@@ -223,7 +223,7 @@ record CommandDependencies(
     ClientPublicKeysManager clientPublicKeysManager =
         new ClientPublicKeysManager(clientPublicKeys, accountLockManager, accountLockExecutor);
     RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager =
-        new RegistrationRecoveryPasswordsManager(registrationRecoveryPasswords, phoneNumberIdentifiers);
+        new RegistrationRecoveryPasswordsManager(registrationRecoveryPasswords);
     AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
         pubsubClient, accountLockManager, keys, messagesManager, profilesManager,
         secureStorageClient, secureValueRecovery2Client, disconnectionRequestManager,

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

@@ -103,9 +103,9 @@ class RegistrationLockVerificationManagerTest {
           if (e instanceof WebApplicationException wae) {
             assertEquals(RegistrationLockVerificationManager.FAILURE_HTTP_STATUS, wae.getResponse().getStatus());
             if (!verificationType.equals(PhoneVerificationRequest.VerificationType.RECOVERY_PASSWORD) || clientRegistrationLock != null) {
-              verify(registrationRecoveryPasswordsManager).removeForNumber(account.getNumber());
+              verify(registrationRecoveryPasswordsManager).remove(account.getIdentifier(IdentityType.PNI));
             } else {
-              verify(registrationRecoveryPasswordsManager, never()).removeForNumber(any());
+              verify(registrationRecoveryPasswordsManager, never()).remove(any());
             }
             verify(disconnectionRequestManager).requestDisconnection(account.getUuid(), List.of(Device.PRIMARY_ID));
             try {
@@ -133,7 +133,7 @@ class RegistrationLockVerificationManagerTest {
           } catch (final NotPushRegisteredException ignored2) {
           }
 
-          verify(registrationRecoveryPasswordsManager, never()).removeForNumber(any());
+          verify(registrationRecoveryPasswordsManager, never()).remove(any());
           verify(disconnectionRequestManager, never()).requestDisconnection(any(), any());
         });
       }
@@ -171,7 +171,7 @@ class RegistrationLockVerificationManagerTest {
             PhoneVerificationRequest.VerificationType.SESSION));
 
     verify(account, never()).lockAuthTokenHash();
-    verify(registrationRecoveryPasswordsManager, never()).removeForNumber(any());
+    verify(registrationRecoveryPasswordsManager, never()).remove(any());
     verify(disconnectionRequestManager, never()).requestDisconnection(any(), any());
   }
 

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

@@ -786,7 +786,7 @@ class AccountControllerTest {
             .withRecoveryPassword(recoveryPassword)))) {
 
       assertThat(response.getStatus()).isEqualTo(204);
-      verify(registrationRecoveryPasswordsManager).storeForCurrentNumber(eq(AuthHelper.UNDISCOVERABLE_NUMBER), eq(recoveryPassword));
+      verify(registrationRecoveryPasswordsManager).store(AuthHelper.UNDISCOVERABLE_PNI, recoveryPassword);
     }
   }
 

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

@@ -582,7 +582,7 @@ class VerificationControllerTest {
       assertTrue(verificationSessionResponse.verified());
       assertTrue(verificationSessionResponse.requestedInformation().isEmpty());
 
-      verify(registrationRecoveryPasswordsManager).removeForNumber(registrationServiceSession.number());
+      verify(registrationRecoveryPasswordsManager).remove(PNI);
     }
   }
 
@@ -879,7 +879,7 @@ class VerificationControllerTest {
     try (Response response = request.get()) {
       assertEquals(HttpStatus.SC_OK, response.getStatus());
 
-      verify(registrationRecoveryPasswordsManager).removeForNumber(registrationServiceSession.number());
+      verify(registrationRecoveryPasswordsManager).remove(PNI);
     }
   }
 
@@ -1204,7 +1204,7 @@ class VerificationControllerTest {
           VerificationSessionResponse.class);
       assertTrue(verificationSessionResponse.verified());
 
-      verify(registrationRecoveryPasswordsManager).removeForNumber(registrationServiceSession.number());
+      verify(registrationRecoveryPasswordsManager).remove(PNI);
     }
   }
 
@@ -1336,7 +1336,7 @@ class VerificationControllerTest {
 
       assertTrue(verificationSessionResponse.verified());
 
-      verify(registrationRecoveryPasswordsManager).removeForNumber(verifiedSession.number());
+      verify(registrationRecoveryPasswordsManager).remove(PNI);
     }
   }
 

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

@@ -114,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(any(), any()))
+    when(registrationRecoveryPasswordsManager.store(any(), any()))
         .thenReturn(CompletableFuture.completedFuture(null));
 
     return new AccountsGrpcService(accountsManager,
@@ -700,7 +700,7 @@ class AccountsGrpcServiceTest extends SimpleBaseGrpcTest<AccountsGrpcService, Ac
             .setRegistrationRecoveryPassword(ByteString.copyFrom(registrationRecoveryPassword))
             .build()));
 
-    verify(registrationRecoveryPasswordsManager).storeForCurrentNumber(account.getNumber(), registrationRecoveryPassword);
+    verify(registrationRecoveryPasswordsManager).store(account.getIdentifier(IdentityType.PNI), registrationRecoveryPassword);
   }
 
   @Test

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

@@ -137,7 +137,7 @@ public class AccountCreationDeletionIntegrationTest {
     final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager =
         mock(RegistrationRecoveryPasswordsManager.class);
 
-    when(registrationRecoveryPasswordsManager.removeForNumber(any()))
+    when(registrationRecoveryPasswordsManager.remove(any()))
         .thenReturn(CompletableFuture.completedFuture(null));
 
     disconnectionRequestManager = mock(DisconnectionRequestManager.class);

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

@@ -130,7 +130,7 @@ class AccountsManagerChangeNumberIntegrationTest {
       final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager =
           mock(RegistrationRecoveryPasswordsManager.class);
 
-      when(registrationRecoveryPasswordsManager.removeForNumber(any()))
+      when(registrationRecoveryPasswordsManager.remove(any()))
           .thenReturn(CompletableFuture.completedFuture(null));
 
       accountsManager = new AccountsManager(

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(any())).thenReturn(CompletableFuture.completedFuture(null));
+    when(registrationRecoveryPasswordsManager.remove(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/AddRemoveDeviceIntegrationTest.java

@@ -137,7 +137,7 @@ public class AddRemoveDeviceIntegrationTest {
     final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager =
         mock(RegistrationRecoveryPasswordsManager.class);
 
-    when(registrationRecoveryPasswordsManager.removeForNumber(any()))
+    when(registrationRecoveryPasswordsManager.remove(any()))
         .thenReturn(CompletableFuture.completedFuture(null));
 
     PUBSUB_SERVER_EXTENSION.getRedisClient().useConnection(connection -> {

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

@@ -33,16 +33,14 @@ 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(Tables.PNI, Tables.REGISTRATION_RECOVERY_PASSWORDS);
+      new DynamoDbExtension(Tables.REGISTRATION_RECOVERY_PASSWORDS);
-
-  private UUID pni;
 
   private RegistrationRecoveryPasswords registrationRecoveryPasswords;
 
@@ -57,22 +55,18 @@ 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);
+    manager = new RegistrationRecoveryPasswordsManager(registrationRecoveryPasswords);
-
-    pni = phoneNumberIdentifiers.getPhoneNumberIdentifier(NUMBER).join();
   }
 
   @Test
   public void testLookupAfterWrite() throws Exception {
-    registrationRecoveryPasswords.addOrReplace(NUMBER, pni, ORIGINAL_HASH).get();
+    registrationRecoveryPasswords.addOrReplace(PNI, ORIGINAL_HASH).get();
-    final long initialExp = fetchTimestamp(NUMBER);
+    final long initialExp = fetchTimestamp(PNI);
     final long expectedExpiration = CLOCK.instant().getEpochSecond() + EXPIRATION.getSeconds();
     assertEquals(expectedExpiration, initialExp);
 
-    final Optional<SaltedTokenHash> saltedTokenHashByPni = registrationRecoveryPasswords.lookup(pni).get();
+    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());
@@ -80,15 +74,15 @@ public class RegistrationRecoveryTest {
 
   @Test
   public void testLookupAfterRefresh() throws Exception {
-    registrationRecoveryPasswords.addOrReplace(NUMBER, pni, ORIGINAL_HASH).get();
+    registrationRecoveryPasswords.addOrReplace(PNI, ORIGINAL_HASH).get();
 
     CLOCK.increment(50, TimeUnit.SECONDS);
-    registrationRecoveryPasswords.addOrReplace(NUMBER, pni, ORIGINAL_HASH).get();
+    registrationRecoveryPasswords.addOrReplace(PNI, ORIGINAL_HASH).get();
-    final long updatedExp = fetchTimestamp(NUMBER);
+    final long updatedExp = fetchTimestamp(PNI);
     final long expectedExp = CLOCK.instant().getEpochSecond() + EXPIRATION.getSeconds();
     assertEquals(expectedExp, updatedExp);
 
-    final Optional<SaltedTokenHash> saltedTokenHashByPni = registrationRecoveryPasswords.lookup(pni).get();
+    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());
@@ -96,10 +90,10 @@ public class RegistrationRecoveryTest {
 
   @Test
   public void testReplace() throws Exception {
-    registrationRecoveryPasswords.addOrReplace(NUMBER, pni, ORIGINAL_HASH).get();
+    registrationRecoveryPasswords.addOrReplace(PNI, ORIGINAL_HASH).get();
-    registrationRecoveryPasswords.addOrReplace(NUMBER, pni, ANOTHER_HASH).get();
+    registrationRecoveryPasswords.addOrReplace(PNI, ANOTHER_HASH).get();
 
-    final Optional<SaltedTokenHash> saltedTokenHashByPni = registrationRecoveryPasswords.lookup(pni).get();
+    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());
@@ -107,13 +101,13 @@ public class RegistrationRecoveryTest {
 
   @Test
   public void testRemove() throws Exception {
-    assertDoesNotThrow(() -> registrationRecoveryPasswords.removeEntry(NUMBER, pni).join());
+    assertDoesNotThrow(() -> registrationRecoveryPasswords.removeEntry(PNI).join());
 
-    registrationRecoveryPasswords.addOrReplace(NUMBER, pni, ORIGINAL_HASH).get();
+    registrationRecoveryPasswords.addOrReplace(PNI, ORIGINAL_HASH).get();
-    assertTrue(registrationRecoveryPasswords.lookup(pni).get().isPresent());
+    assertTrue(registrationRecoveryPasswords.lookup(PNI).get().isPresent());
 
-    registrationRecoveryPasswords.removeEntry(NUMBER, pni).get();
+    registrationRecoveryPasswords.removeEntry(PNI).get();
-    assertTrue(registrationRecoveryPasswords.lookup(pni).get().isEmpty());
+    assertTrue(registrationRecoveryPasswords.lookup(PNI).get().isEmpty());
   }
 
   @Test
@@ -123,31 +117,31 @@ public class RegistrationRecoveryTest {
     final byte[] wrongPassword = "qwerty123".getBytes(StandardCharsets.UTF_8);
 
     // initial store
-    manager.storeForCurrentNumber(NUMBER, password).get();
+    manager.store(PNI, password).get();
-    assertTrue(manager.verify(pni, password).get());
+    assertTrue(manager.verify(PNI, password).get());
-    assertFalse(manager.verify(pni, wrongPassword).get());
+    assertFalse(manager.verify(PNI, wrongPassword).get());
 
     // update
-    manager.storeForCurrentNumber(NUMBER, password).get();
+    manager.store(PNI, password).get();
-    assertTrue(manager.verify(pni, password).get());
+    assertTrue(manager.verify(PNI, password).get());
-    assertFalse(manager.verify(pni, wrongPassword).get());
+    assertFalse(manager.verify(PNI, wrongPassword).get());
 
     // replace
-    manager.storeForCurrentNumber(NUMBER, updatedPassword).get();
+    manager.store(PNI, updatedPassword).get();
-    assertTrue(manager.verify(pni, updatedPassword).get());
+    assertTrue(manager.verify(PNI, updatedPassword).get());
-    assertFalse(manager.verify(pni, password).get());
+    assertFalse(manager.verify(PNI, password).get());
-    assertFalse(manager.verify(pni, wrongPassword).get());
+    assertFalse(manager.verify(PNI, wrongPassword).get());
 
-    manager.removeForNumber(NUMBER).get();
+    manager.remove(PNI).get();
-    assertFalse(manager.verify(pni, updatedPassword).get());
+    assertFalse(manager.verify(PNI, updatedPassword).get());
-    assertFalse(manager.verify(pni, password).get());
+    assertFalse(manager.verify(PNI, password).get());
-    assertFalse(manager.verify(pni, wrongPassword).get());
+    assertFalse(manager.verify(PNI, wrongPassword).get());
   }
 
-  private static long fetchTimestamp(final String number) throws ExecutionException, InterruptedException {
+  private static long fetchTimestamp(final UUID phoneNumberIdentifier) throws ExecutionException, InterruptedException {
     return DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient().getItem(GetItemRequest.builder()
             .tableName(Tables.REGISTRATION_RECOVERY_PASSWORDS.tableName())
-            .key(Map.of(RegistrationRecoveryPasswords.KEY_PNI, AttributeValues.fromString(number)))
+            .key(Map.of(RegistrationRecoveryPasswords.KEY_PNI, AttributeValues.fromString(phoneNumberIdentifier.toString())))
             .build())
         .thenApply(getItemResponse -> {
           final Map<String, AttributeValue> item = getItemResponse.item();