From f0a6be32fcc618a98ce68d2f6d3c614a4267ed95 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Wed, 27 Oct 2021 18:09:03 -0400 Subject: [PATCH] Add a crawler to assign PNIs to existing accounts --- .../textsecuregcm/WhisperServerService.java | 2 + ...nPhoneNumberIdentifierCrawlerListener.java | 66 +++++++++++++++++ ...neNumberIdentifierCrawlerListenerTest.java | 71 +++++++++++++++++++ 3 files changed, 139 insertions(+) create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/storage/AssignPhoneNumberIdentifierCrawlerListener.java create mode 100644 service/src/test/java/org/whispersystems/textsecuregcm/storage/AssignPhoneNumberIdentifierCrawlerListenerTest.java diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java index 7009fd581..1c68d3b7c 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java @@ -170,6 +170,7 @@ import org.whispersystems.textsecuregcm.storage.AccountDatabaseCrawlerCache; import org.whispersystems.textsecuregcm.storage.AccountDatabaseCrawlerListener; import org.whispersystems.textsecuregcm.storage.Accounts; import org.whispersystems.textsecuregcm.storage.AccountsManager; +import org.whispersystems.textsecuregcm.storage.AssignPhoneNumberIdentifierCrawlerListener; import org.whispersystems.textsecuregcm.storage.ContactDiscoveryWriter; import org.whispersystems.textsecuregcm.storage.DeletedAccounts; import org.whispersystems.textsecuregcm.storage.DeletedAccountsDirectoryReconciler; @@ -545,6 +546,7 @@ public class WhisperServerService extends Application fromUuid) { + } + + @Override + protected void onCrawlChunk(final Optional fromUuid, final List chunkAccounts) { + // There are exactly two ways an account can get a phone number identifier (PNI): + // + // 1. It's assigned at construction time for accounts created after the introduction of PNIs + // 2. It's assigned by this crawler + // + // That means that we don't need to worry about accidentally overwriting a PNI assigned by another source; if an + // account doesn't have a PNI when it winds up in a crawled chunk, there's no danger that it will have one after a + // refresh, and so we can blindly assign a random PNI. + chunkAccounts.stream() + .filter(account -> account.getPhoneNumberIdentifier().isEmpty()) + .map(Account::getUuid) + .forEach(accountIdentifier -> { + // We must not update the accounts in the chunk directly; instead, we need to get a fresh copy we're free to + // update as needed. + accountsManager.getByAccountIdentifier(accountIdentifier).ifPresent(accountWithoutPni -> { + final String number = accountWithoutPni.getNumber(); + final UUID phoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(number); + + accountsManager.update(accountWithoutPni, a -> a.setNumber(number, phoneNumberIdentifier)); + }); + + ASSIGNED_PNI_COUNTER.increment(); + }); + } +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AssignPhoneNumberIdentifierCrawlerListenerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AssignPhoneNumberIdentifierCrawlerListenerTest.java new file mode 100644 index 000000000..aabcac82b --- /dev/null +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AssignPhoneNumberIdentifierCrawlerListenerTest.java @@ -0,0 +1,71 @@ +/* + * Copyright 2013-2021 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.storage; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.List; +import java.util.Optional; +import java.util.UUID; +import java.util.function.Consumer; +import org.junit.jupiter.api.Test; +import org.mockito.stubbing.Answer; + +class AssignPhoneNumberIdentifierCrawlerListenerTest { + + @Test + void onCrawlChunk() { + final UUID accountIdentifierWithPni = UUID.randomUUID(); + final UUID accountIdentifierWithoutPni = UUID.randomUUID(); + + final String numberWithPni = "+18005551111"; + final String numberWithoutPni = "+18005552222"; + + final Account accountWithPni = mock(Account.class); + when(accountWithPni.getUuid()).thenReturn(accountIdentifierWithPni); + when(accountWithPni.getNumber()).thenReturn(numberWithPni); + when(accountWithPni.getPhoneNumberIdentifier()).thenReturn(Optional.of(UUID.randomUUID())); + + final Account accountWithoutPni = mock(Account.class); + when(accountWithoutPni.getUuid()).thenReturn(accountIdentifierWithoutPni); + when(accountWithoutPni.getNumber()).thenReturn(numberWithoutPni); + when(accountWithoutPni.getPhoneNumberIdentifier()).thenReturn(Optional.empty()); + + final AccountsManager accountsManager = mock(AccountsManager.class); + when(accountsManager.getByAccountIdentifier(accountIdentifierWithPni)).thenReturn(Optional.of(accountWithPni)); + when(accountsManager.getByAccountIdentifier(accountIdentifierWithoutPni)).thenReturn(Optional.of(accountWithoutPni)); + + when(accountsManager.update(any(), any())).thenAnswer((Answer) invocation -> { + final Account account = invocation.getArgument(0, Account.class); + final Consumer updater = invocation.getArgument(1, Consumer.class); + + updater.accept(account); + + return account; + }); + + final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class); + when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString())).thenAnswer( + (Answer) invocation -> UUID.randomUUID()); + + final AssignPhoneNumberIdentifierCrawlerListener crawler = + new AssignPhoneNumberIdentifierCrawlerListener(accountsManager, phoneNumberIdentifiers); + + crawler.onCrawlChunk(Optional.empty(), List.of(accountWithPni, accountWithoutPni)); + + verify(accountsManager).update(eq(accountWithoutPni), any()); + verify(accountWithoutPni).setNumber(eq(numberWithoutPni), any()); + + verify(accountsManager, never()).update(eq(accountWithPni), any()); + verify(accountWithPni, never()).setNumber(any(), any()); + } +}