Refresh registration recovery password expirations before retrying an…

… insertion
signalapp · Nov 25, 2024 · ff4e2bd · ff4e2bd
1 parent ffed19d
commit ff4e2bd
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 9 deletions.
diff --git a/...src/main/java/org/whispersystems/textsecuregcm/storage/RegistrationRecoveryPasswords.java b/...src/main/java/org/whispersystems/textsecuregcm/storage/RegistrationRecoveryPasswords.java
@@ -17,6 +17,7 @@
 import org.whispersystems.textsecuregcm.auth.SaltedTokenHash;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
 import org.whispersystems.textsecuregcm.util.ExceptionUtils;
+import org.whispersystems.textsecuregcm.util.Pair;
 import org.whispersystems.textsecuregcm.util.Util;
 import reactor.core.publisher.Flux;
 import reactor.core.scheduler.Scheduler;
@@ -65,13 +66,26 @@ public CompletableFuture<Optional<SaltedTokenHash>> lookup(final String number)
             .tableName(tableName)
             .key(Map.of(KEY_E164, AttributeValues.fromString(number)))
             .consistentRead(true)
-        .build())
+            .build())
         .thenApply(getItemResponse -> Optional.ofNullable(getItemResponse.item())
             .filter(item -> item.containsKey(ATTR_SALT))
             .filter(item -> item.containsKey(ATTR_HASH))
             .map(RegistrationRecoveryPasswords::saltedTokenHashFromItem));
   }
 
+  CompletableFuture<Optional<Pair<SaltedTokenHash, Long>>> lookupWithExpiration(final String key) {
+    return asyncClient.getItem(GetItemRequest.builder()
+            .tableName(tableName)
+            .key(Map.of(KEY_E164, AttributeValues.fromString(key)))
+            .consistentRead(true)
+            .build())
+        .thenApply(getItemResponse -> Optional.ofNullable(getItemResponse.item())
+            .filter(item -> item.containsKey(ATTR_SALT))
+            .filter(item -> item.containsKey(ATTR_HASH))
+            .filter(item -> item.containsKey(ATTR_EXP))
+            .map(item -> new Pair<>(saltedTokenHashFromItem(item), Long.parseLong(item.get(ATTR_EXP).n()))));
+  }
+
   public CompletableFuture<Optional<SaltedTokenHash>> lookup(final UUID phoneNumberIdentifier) {
     return lookup(phoneNumberIdentifier.toString());
   }

diff --git a/...n/java/org/whispersystems/textsecuregcm/storage/RegistrationRecoveryPasswordsManager.java b/...n/java/org/whispersystems/textsecuregcm/storage/RegistrationRecoveryPasswordsManager.java
@@ -95,9 +95,9 @@ public CompletableFuture<Boolean> migrateE164Record(final String number,
         .exceptionallyCompose(throwable -> {
           if (ExceptionUtils.unwrap(throwable) instanceof ContestedOptimisticLockException) {
             // Something about the original record changed; refresh and retry
-            return registrationRecoveryPasswords.lookup(number)
-                .thenCompose(maybeSaltedTokenHash -> maybeSaltedTokenHash
-                    .map(refreshedSaltedTokenHash -> migrateE164Record(number, phoneNumberIdentifier, refreshedSaltedTokenHash, expirationSeconds, remainingAttempts - 1))
+            return registrationRecoveryPasswords.lookupWithExpiration(number)
+                .thenCompose(maybePair -> maybePair
+                    .map(pair -> migrateE164Record(number, phoneNumberIdentifier, pair.first(), pair.second(), remainingAttempts - 1))
                     .orElseGet(() -> {
                       // The original record was deleted, and we can declare victory
                       return CompletableFuture.completedFuture(false);

diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/RegistrationRecoveryTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/RegistrationRecoveryTest.java
@@ -39,6 +39,7 @@
 import org.whispersystems.textsecuregcm.util.AttributeValues;
 import org.whispersystems.textsecuregcm.util.MockUtils;
 import org.whispersystems.textsecuregcm.util.MutableClock;
+import org.whispersystems.textsecuregcm.util.Pair;
 import reactor.core.scheduler.Schedulers;
 import reactor.util.function.Tuples;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
@@ -308,8 +309,8 @@ void migrateE164RecordRetry() {
     when(registrationRecoveryPasswords.insertPniRecord(eq(NUMBER), eq(PNI), eq(ANOTHER_HASH), anyLong()))
         .thenReturn(CompletableFuture.completedFuture(true));
 
-    when(registrationRecoveryPasswords.lookup(NUMBER))
-        .thenReturn(CompletableFuture.completedFuture(Optional.of(ANOTHER_HASH)));
+    when(registrationRecoveryPasswords.lookupWithExpiration(NUMBER))
+        .thenReturn(CompletableFuture.completedFuture(Optional.of(new Pair<>(ANOTHER_HASH, 1234L))));
 
     final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
     when(phoneNumberIdentifiers.getPhoneNumberIdentifier(NUMBER)).thenReturn(CompletableFuture.completedFuture(PNI));
@@ -319,7 +320,7 @@ void migrateE164RecordRetry() {
 
     assertTrue(() -> migrationManager.migrateE164Record(NUMBER, ORIGINAL_HASH, 1234).join());
 
-    verify(registrationRecoveryPasswords).lookup(NUMBER);
+    verify(registrationRecoveryPasswords).lookupWithExpiration(NUMBER);
     verify(registrationRecoveryPasswords, times(2)).insertPniRecord(any(), any(), any(), anyLong());
   }
 
@@ -333,6 +334,9 @@ void migrateE164RecordOriginalDeleted() {
     when(registrationRecoveryPasswords.lookup(NUMBER))
         .thenReturn(CompletableFuture.completedFuture(Optional.empty()));
 
+    when(registrationRecoveryPasswords.lookupWithExpiration(NUMBER))
+        .thenReturn(CompletableFuture.completedFuture(Optional.empty()));
+
     final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
     when(phoneNumberIdentifiers.getPhoneNumberIdentifier(NUMBER)).thenReturn(CompletableFuture.completedFuture(PNI));
 
@@ -341,7 +345,7 @@ void migrateE164RecordOriginalDeleted() {
 
     assertFalse(() -> migrationManager.migrateE164Record(NUMBER, ORIGINAL_HASH, 1234).join());
 
-    verify(registrationRecoveryPasswords).lookup(NUMBER);
+    verify(registrationRecoveryPasswords).lookupWithExpiration(NUMBER);
     verify(registrationRecoveryPasswords).insertPniRecord(any(), any(), any(), anyLong());
   }
 
@@ -355,6 +359,9 @@ void migrateE164RecordRetryExhausted() {
     when(registrationRecoveryPasswords.lookup(NUMBER))
         .thenReturn(CompletableFuture.completedFuture(Optional.of(ORIGINAL_HASH)));
 
+    when(registrationRecoveryPasswords.lookupWithExpiration(NUMBER))
+        .thenReturn(CompletableFuture.completedFuture(Optional.of(new Pair<>(ORIGINAL_HASH, CLOCK.instant().getEpochSecond() + EXPIRATION.getSeconds()))));
+
     final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
     when(phoneNumberIdentifiers.getPhoneNumberIdentifier(NUMBER)).thenReturn(CompletableFuture.completedFuture(PNI));
 
@@ -366,7 +373,7 @@ void migrateE164RecordRetryExhausted() {
 
     assertInstanceOf(ContestedOptimisticLockException.class, completionException.getCause());
 
-    verify(registrationRecoveryPasswords, times(10)).lookup(NUMBER);
+    verify(registrationRecoveryPasswords, times(10)).lookupWithExpiration(NUMBER);
     verify(registrationRecoveryPasswords, times(10)).insertPniRecord(any(), any(), any(), anyLong());
   }
 }