Add methods for migrating E164-mapped registration recovery passwords to PNI-mapped records

This commit is contained in:
Jon Chambers
2024-11-22 16:55:59 -05:00
committed by Jon Chambers
parent 3c8b2a82a3
commit af1d21c225
3 changed files with 311 additions and 11 deletions

View File

@@ -7,6 +7,7 @@ package org.whispersystems.textsecuregcm.storage;
import static java.util.Objects.requireNonNull;
import com.google.common.annotations.VisibleForTesting;
import java.time.Clock;
import java.time.Duration;
import java.util.Map;
@@ -15,15 +16,21 @@ 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.ExceptionUtils;
import org.whispersystems.textsecuregcm.util.Util;
import reactor.core.publisher.Flux;
import reactor.util.function.Tuple3;
import reactor.util.function.Tuples;
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.Delete;
import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
import software.amazon.awssdk.services.dynamodb.model.Put;
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledException;
public class RegistrationRecoveryPasswords extends AbstractDynamoDbStore {
@@ -69,15 +76,10 @@ public class RegistrationRecoveryPasswords extends AbstractDynamoDbStore {
.key(Map.of(KEY_E164, AttributeValues.fromString(number)))
.consistentRead(true)
.build())
.thenApply(getItemResponse -> {
final Map<String, AttributeValue> item = getItemResponse.item();
if (item == null || !item.containsKey(ATTR_SALT) || !item.containsKey(ATTR_HASH)) {
return Optional.empty();
}
final String salt = item.get(ATTR_SALT).s();
final String hash = item.get(ATTR_HASH).s();
return Optional.of(new SaltedTokenHash(hash, salt));
});
.thenApply(getItemResponse -> Optional.ofNullable(getItemResponse.item())
.filter(item -> item.containsKey(ATTR_SALT))
.filter(item -> item.containsKey(ATTR_HASH))
.map(RegistrationRecoveryPasswords::saltedTokenHashFromItem));
}
public CompletableFuture<Optional<SaltedTokenHash>> lookup(final UUID phoneNumberIdentifier) {
@@ -130,7 +132,91 @@ public class RegistrationRecoveryPasswords extends AbstractDynamoDbStore {
.build();
}
private long expirationSeconds() {
@VisibleForTesting
long expirationSeconds() {
return clock.instant().plus(expiration).getEpochSecond();
}
public Flux<Tuple3<String, SaltedTokenHash, Long>> getE164AssociatedRegistrationRecoveryPasswords() {
return Flux.from(asyncClient.scanPaginator(ScanRequest.builder()
.tableName(tableName)
.consistentRead(true)
.filterExpression("begins_with(#key, :e164Prefix)")
.expressionAttributeNames(Map.of("#key", KEY_E164))
.expressionAttributeValues(Map.of(":e164Prefix", AttributeValue.fromS("+")))
.build())
.items())
.map(item -> Tuples.of(item.get(KEY_E164).s(), saltedTokenHashFromItem(item), Long.parseLong(item.get(ATTR_EXP).n())));
}
public CompletableFuture<Boolean> insertPniRecord(final String phoneNumber,
final UUID phoneNumberIdentifier,
final SaltedTokenHash saltedTokenHash,
final long expirationSeconds) {
// We try to write both the old and new record inside a transaction, but with different conditions. For the
// E164-based record, we insist that the record be entirely unchanged. This prevents us from writing an out-of-sync
// record if we read one thing in the `Scan` pass, but then somebody updated the record before we tried to write
// the PNI-based record. We refresh and retry if this happens.
//
// For the PNI-based record, we only want to write the record if one doesn't already exist for the given PNI. If one
// already exists, we'll just leave it alone.
return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder()
.transactItems(
TransactWriteItem.builder()
.put(Put.builder()
.tableName(tableName)
.item(Map.of(
KEY_E164, AttributeValues.fromString(phoneNumber),
ATTR_EXP, AttributeValues.fromLong(expirationSeconds),
ATTR_SALT, AttributeValues.fromString(saltedTokenHash.salt()),
ATTR_HASH, AttributeValues.fromString(saltedTokenHash.hash())))
.conditionExpression("#key = :key AND #expiration = :expiration AND #salt = :salt AND #hash = :hash")
.expressionAttributeNames(Map.of(
"#key", KEY_E164,
"#expiration", ATTR_EXP,
"#salt", ATTR_SALT,
"#hash", ATTR_HASH))
.expressionAttributeValues(Map.of(
":key", AttributeValues.fromString(phoneNumber),
":expiration", AttributeValues.fromLong(expirationSeconds),
":salt", AttributeValues.fromString(saltedTokenHash.salt()),
":hash", AttributeValues.fromString(saltedTokenHash.hash())))
.build())
.build(),
TransactWriteItem.builder()
.put(Put.builder()
.tableName(tableName)
.item(Map.of(
KEY_E164, AttributeValues.fromString(phoneNumberIdentifier.toString()),
ATTR_EXP, AttributeValues.fromLong(expirationSeconds),
ATTR_SALT, AttributeValues.fromString(saltedTokenHash.salt()),
ATTR_HASH, AttributeValues.fromString(saltedTokenHash.hash())))
.conditionExpression("attribute_not_exists(#key)")
.expressionAttributeNames(Map.of("#key", KEY_E164))
.build())
.build())
.build())
.thenApply(ignored -> true)
.exceptionally(throwable -> {
if (ExceptionUtils.unwrap(throwable) instanceof TransactionCanceledException transactionCanceledException) {
if ("ConditionalCheckFailed".equals(transactionCanceledException.cancellationReasons().get(1).code())) {
// A PNI-associated record has already been stored; we can just treat this as success
return false;
}
if ("ConditionalCheckFailed".equals(transactionCanceledException.cancellationReasons().get(0).code())) {
// No PNI-associated record is present, but the original record has changed
throw new ContestedOptimisticLockException();
}
}
throw ExceptionUtils.wrap(throwable);
});
}
private static SaltedTokenHash saltedTokenHashFromItem(final Map<String, AttributeValue> item) {
return new SaltedTokenHash(item.get(ATTR_HASH).s(), item.get(ATTR_SALT).s());
}
}

View File

@@ -10,10 +10,12 @@ 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;
import org.whispersystems.textsecuregcm.auth.SaltedTokenHash;
import org.whispersystems.textsecuregcm.util.ExceptionUtils;
import software.amazon.awssdk.services.dynamodb.model.ResourceNotFoundException;
public class RegistrationRecoveryPasswordsManager {
@@ -67,6 +69,38 @@ public class RegistrationRecoveryPasswordsManager {
}));
}
public CompletableFuture<Boolean> migrateE164Record(final String number, final SaltedTokenHash saltedTokenHash, final long expirationSeconds) {
return phoneNumberIdentifiers.getPhoneNumberIdentifier(number)
.thenCompose(phoneNumberIdentifier -> migrateE164Record(number, phoneNumberIdentifier, saltedTokenHash, expirationSeconds, 10));
}
public CompletableFuture<Boolean> migrateE164Record(final String number,
final UUID phoneNumberIdentifier,
final SaltedTokenHash saltedTokenHash,
final long expirationSeconds,
final int remainingAttempts) {
if (remainingAttempts <= 0) {
return CompletableFuture.failedFuture(new ContestedOptimisticLockException());
}
return registrationRecoveryPasswords.insertPniRecord(number, phoneNumberIdentifier, saltedTokenHash, expirationSeconds)
.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))
.orElseGet(() -> {
// The original record was deleted, and we can declare victory
return CompletableFuture.completedFuture(false);
}));
}
return CompletableFuture.failedFuture(throwable);
});
}
private static String bytesToString(final byte[] bytes) {
return HexFormat.of().formatHex(bytes);
}