From 38958714624cd0cfcb1d6be45fc1b503b85fe298 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Fri, 10 Dec 2021 13:08:11 -0500 Subject: [PATCH] Repair missing PNIs in JSON blobs on account load --- .../textsecuregcm/storage/Accounts.java | 28 +++++++-- .../textsecuregcm/storage/AccountsTest.java | 62 ++++++++++++++++++- 2 files changed, 83 insertions(+), 7 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java index 7db4a0835..47fa51046 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java @@ -498,7 +498,7 @@ public class Accounts extends AbstractDynamoDbStore { return Optional.ofNullable(response.item()) .map(item -> item.get(KEY_ACCOUNT_UUID)) .map(this::accountByUuid) - .map(Accounts::fromItem); + .map(this::fromItem); }); } @@ -513,7 +513,7 @@ public class Accounts extends AbstractDynamoDbStore { return Optional.ofNullable(response.item()) .map(item -> item.get(KEY_ACCOUNT_UUID)) .map(this::accountByUuid) - .map(Accounts::fromItem); + .map(this::fromItem); }); } @@ -528,7 +528,7 @@ public class Accounts extends AbstractDynamoDbStore { return Optional.ofNullable(response.item()) .map(item -> item.get(KEY_ACCOUNT_UUID)) .map(this::accountByUuid) - .map(Accounts::fromItem); + .map(this::fromItem); }); } @@ -544,7 +544,7 @@ public class Accounts extends AbstractDynamoDbStore { public Optional getByAccountIdentifier(UUID uuid) { return GET_BY_UUID_TIMER.record(() -> Optional.ofNullable(accountByUuid(AttributeValues.fromUUID(uuid))) - .map(Accounts::fromItem)); + .map(this::fromItem)); } public void delete(UUID uuid) { @@ -611,7 +611,7 @@ public class Accounts extends AbstractDynamoDbStore { final List accounts = timer.record(() -> scan(scanRequestBuilder.build(), maxCount) .stream() - .map(Accounts::fromItem) + .map(this::fromItem) .collect(Collectors.toList())); return new AccountCrawlChunk(accounts, accounts.size() > 0 ? accounts.get(accounts.size() - 1).getUuid() : null); @@ -623,8 +623,9 @@ public class Accounts extends AbstractDynamoDbStore { .collect(Collectors.joining(", ")); } + // TODO Make this static once PNI repairs are complete and the call to #update(Account) is no longer needed @VisibleForTesting - static Account fromItem(Map item) { + Account fromItem(Map item) { if (!item.containsKey(ATTR_ACCOUNT_DATA) || !item.containsKey(ATTR_ACCOUNT_E164) || // TODO: eventually require ATTR_CANONICALLY_DISCOVERABLE @@ -644,12 +645,27 @@ public class Accounts extends AbstractDynamoDbStore { accountIdentifier, account.getPhoneNumberIdentifier(), phoneNumberIdentifierFromAttribute); } + boolean shouldResetPni = false; + + if (account.getPhoneNumberIdentifier() == null && phoneNumberIdentifierFromAttribute != null) { + log.info("Account {} has a PNI in its attributes, but not JSON blob; will repair", accountIdentifier); + shouldResetPni = true; + } + account.setNumber(item.get(ATTR_ACCOUNT_E164).s(), phoneNumberIdentifierFromAttribute); account.setUuid(accountIdentifier); account.setUsername(AttributeValues.getString(item, ATTR_USERNAME, null)); account.setVersion(Integer.parseInt(item.get(ATTR_VERSION).n())); account.setCanonicallyDiscoverable(Optional.ofNullable(item.get(ATTR_CANONICALLY_DISCOVERABLE)).map(av -> av.bool()).orElse(false)); + if (shouldResetPni) { + try { + update(account); + } catch (final Exception e) { + log.warn("Failed to reset PNI for account {}", accountIdentifier, e); + } + } + return account; } catch (IOException e) { diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java index 85d8b21d0..3c9b05c23 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java @@ -769,6 +769,66 @@ class AccountsTest { assertThat(account.getUsername()).hasValueSatisfying(u -> assertThat(u).isEqualTo(username)); } + // TODO Remove when PNI repair pass is complete + @Test + void testRepairMissingPni() throws JsonProcessingException { + final String number = "+18005551234"; + final UUID accountIdentifier = UUID.randomUUID(); + final UUID phoneNumberIdentifier = UUID.randomUUID(); + + { + final Account account = generateAccount(number, accountIdentifier, phoneNumberIdentifier); + accounts.create(account); + + // Artificially inject data with a null PNI in the JSON blob into the backing table + account.setNumber(number, null); + + dynamoDbExtension.getDynamoDbClient().updateItem(UpdateItemRequest.builder() + .tableName(ACCOUNTS_TABLE_NAME) + .key(Map.of(Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(accountIdentifier))) + .updateExpression("SET #data = :data") + .expressionAttributeNames(Map.of("#data", Accounts.ATTR_ACCOUNT_DATA)) + .expressionAttributeValues( + Map.of(":data", AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)))) + .build()); + } + + { + final GetItemResponse response = dynamoDbExtension.getDynamoDbClient().getItem(GetItemRequest.builder() + .tableName(ACCOUNTS_TABLE_NAME) + .key(Map.of(Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(accountIdentifier))) + .consistentRead(true) + .build()); + + final String accountJson = new String( + AttributeValues.getByteArray(response.item(), Accounts.ATTR_ACCOUNT_DATA, null), StandardCharsets.UTF_8); + + final Account accountFromJson = SystemMapper.getMapper().readValue(accountJson, Account.class); + assertThat(accountFromJson.getPhoneNumberIdentifier()).isNull(); + } + + { + // Loading the account should auto-repair the missing PNI + final Account account = accounts.getByAccountIdentifier(accountIdentifier).orElseThrow(); + assertThat(account.getPhoneNumberIdentifier()).isEqualTo(phoneNumberIdentifier); + assertThat(account.isStale()).isFalse(); + } + + { + final GetItemResponse response = dynamoDbExtension.getDynamoDbClient().getItem(GetItemRequest.builder() + .tableName(ACCOUNTS_TABLE_NAME) + .key(Map.of(Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(accountIdentifier))) + .consistentRead(true) + .build()); + + final String accountJson = new String( + AttributeValues.getByteArray(response.item(), Accounts.ATTR_ACCOUNT_DATA, null), StandardCharsets.UTF_8); + + final Account accountFromJson = SystemMapper.getMapper().readValue(accountJson, Account.class); + assertThat(accountFromJson.getPhoneNumberIdentifier()).isEqualTo(phoneNumberIdentifier); + } + } + private Device generateDevice(long id) { Random random = new Random(System.currentTimeMillis()); SignedPreKey signedPreKey = new SignedPreKey(random.nextInt(), "testPublicKey-" + random.nextInt(), "testSignature-" + random.nextInt()); @@ -851,7 +911,7 @@ class AccountsTest { assertThat(AttributeValues.getBool(get.item(), Accounts.ATTR_CANONICALLY_DISCOVERABLE, !canonicallyDiscoverable)).isEqualTo(canonicallyDiscoverable); - Account result = Accounts.fromItem(get.item()); + Account result = accounts.fromItem(get.item()); verifyStoredState(number, uuid, pni, result, expecting); } else { throw new AssertionError("No data");