Delete registration recovery passwords by both phone number and PNI

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

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

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

@@ -27,7 +27,6 @@ import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
 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;

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

@@ -11,6 +11,7 @@ import java.time.Clock;
 import java.time.Duration;
 import java.util.Map;
 import java.util.Optional;
+import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import org.whispersystems.textsecuregcm.auth.SaltedTokenHash;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
@@ -18,12 +19,16 @@ import org.whispersystems.textsecuregcm.util.Util;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
-import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
+import software.amazon.awssdk.services.dynamodb.model.Delete;
 import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
-import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
+import software.amazon.awssdk.services.dynamodb.model.Put;
+import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
+import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
 
 public class RegistrationRecoveryPasswords extends AbstractDynamoDbStore {
 
+  // As a temporary transitional measure, this can be either a string representation of an E164-formatted phone number
+  // or a UUID (PNI) string
   static final String KEY_E164 = "P";
   static final String ATTR_EXP = "E";
   static final String ATTR_SALT = "S";
@@ -75,26 +80,57 @@ public class RegistrationRecoveryPasswords extends AbstractDynamoDbStore {
         });
   }
 
-  public CompletableFuture<Void> addOrReplace(final String number, final SaltedTokenHash data) {
-    return asyncClient.putItem(PutItemRequest.builder()
-            .tableName(tableName)
-            .item(Map.of(
-                KEY_E164, AttributeValues.fromString(number),
-                ATTR_EXP, AttributeValues.fromLong(expirationSeconds()),
-                ATTR_SALT, AttributeValues.fromString(data.salt()),
-                ATTR_HASH, AttributeValues.fromString(data.hash())))
-            .build())
+  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();
+
+    return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder()
+            .transactItems(
+                buildPutRecoveryPasswordWriteItem(number, expirationSeconds, data.salt(), data.hash())
+                // buildPutRecoveryPasswordWriteItem(phoneNumberIdentifier.toString(), expirationSeconds, data.salt(), data.hash())
+            )
+        .build())
         .thenRun(Util.NOOP);
   }
 
-  public CompletableFuture<Void> removeEntry(final String number) {
-    return asyncClient.deleteItem(DeleteItemRequest.builder()
+  private TransactWriteItem buildPutRecoveryPasswordWriteItem(final String key,
+      final long expirationSeconds,
+      final String salt,
+      final String hash) {
+
+    return TransactWriteItem.builder()
+        .put(Put.builder()
             .tableName(tableName)
-            .key(Map.of(KEY_E164, AttributeValues.fromString(number)))
+            .item(Map.of(
+                KEY_E164, AttributeValues.fromString(key),
+                ATTR_EXP, AttributeValues.fromLong(expirationSeconds),
+                ATTR_SALT, AttributeValues.fromString(salt),
+                ATTR_HASH, AttributeValues.fromString(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) {
+    return TransactWriteItem.builder()
+        .delete(Delete.builder()
+            .tableName(tableName)
+            .key(Map.of(KEY_E164, AttributeValues.fromString(key)))
+            .build())
+        .build();
+  }
+
   private long expirationSeconds() {
     return clock.instant().plus(expiration).getEpochSecond();
   }

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

@@ -21,10 +21,13 @@ 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 String number, final byte[] password) {
@@ -41,26 +44,27 @@ public class RegistrationRecoveryPasswordsManager {
   public CompletableFuture<Void> storeForCurrentNumber(final String number, final byte[] password) {
     final String token = bytesToString(password);
     final SaltedTokenHash tokenHash = SaltedTokenHash.generateFor(token);
-    return registrationRecoveryPasswords.addOrReplace(number, tokenHash)
-        .whenComplete((result, error) -> {
-          if (error != null) {
-            logger.warn("Failed to store Registration Recovery Password", error);
-          }
-        });
+
+    return phoneNumberIdentifiers.getPhoneNumberIdentifier(number)
+        .thenCompose(phoneNumberIdentifier -> registrationRecoveryPasswords.addOrReplace(number, phoneNumberIdentifier, tokenHash)
+            .whenComplete((result, error) -> {
+              if (error != null) {
+                logger.warn("Failed to store Registration Recovery Password", error);
+              }
+            }));
   }
 
   public CompletableFuture<Void> removeForNumber(final String number) {
-    // remove is a "fire-and-forget" operation,
-    // there is no action to be taken on its completion
-    return registrationRecoveryPasswords.removeEntry(number)
-        .whenComplete((ignored, error) -> {
-          if (error instanceof ResourceNotFoundException) {
-            // 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
-          } else if (error != null) {
-            logger.warn("Failed to remove Registration Recovery Password", error);
-          }
-        });
+    return phoneNumberIdentifiers.getPhoneNumberIdentifier(number)
+        .thenCompose(phoneNumberIdentifier -> registrationRecoveryPasswords.removeEntry(number, phoneNumberIdentifier)
+            .whenComplete((ignored, error) -> {
+              if (error instanceof ResourceNotFoundException) {
+                // 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
+              } else if (error != null) {
+                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

@@ -170,9 +170,6 @@ record CommandDependencies(
         dynamoDbAsyncClient
     );
 
-    RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager = new RegistrationRecoveryPasswordsManager(
-        registrationRecoveryPasswords);
-
     ClientPublicKeys clientPublicKeys =
         new ClientPublicKeys(dynamoDbAsyncClient, configuration.getDynamoDbTables().getClientPublicKeys().getTableName());
 
@@ -225,6 +222,8 @@ record CommandDependencies(
         configuration.getDynamoDbTables().getDeletedAccountsLock().getTableName());
     ClientPublicKeysManager clientPublicKeysManager =
         new ClientPublicKeysManager(clientPublicKeys, accountLockManager, accountLockExecutor);
+    RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager =
+        new RegistrationRecoveryPasswordsManager(registrationRecoveryPasswords, phoneNumberIdentifiers);
     AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
         pubsubClient, accountLockManager, keys, messagesManager, profilesManager,
         secureStorageClient, secureValueRecovery2Client, disconnectionRequestManager,