Add a crawler to assign PNIs to existing accounts

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<WhisperServerConfiguration
     }
     accountDatabaseCrawlerListeners.add(new NonNormalizedAccountCrawlerListener(accountsManager, metricsCluster));
     accountDatabaseCrawlerListeners.add(new ContactDiscoveryWriter(accountsManager));
+    accountDatabaseCrawlerListeners.add(new AssignPhoneNumberIdentifierCrawlerListener(accountsManager, phoneNumberIdentifiers));
     // PushFeedbackProcessor may update device properties
     accountDatabaseCrawlerListeners.add(new PushFeedbackProcessor(accountsManager));
     // delete accounts last

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

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2013-2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import io.micrometer.core.instrument.Counter;
+import io.micrometer.core.instrument.Metrics;
+import java.util.List;
+import java.util.Optional;
+import java.util.UUID;
+
+import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
+
+// TODO Remove this crawler when PNIs have been assigned to all existing accounts
+public class AssignPhoneNumberIdentifierCrawlerListener extends AccountDatabaseCrawlerListener {
+
+  private final AccountsManager accountsManager;
+  private final PhoneNumberIdentifiers phoneNumberIdentifiers;
+
+  private static final Counter ASSIGNED_PNI_COUNTER =
+      Metrics.counter(name(AssignPhoneNumberIdentifierCrawlerListener.class, "assignPni"));
+
+  public AssignPhoneNumberIdentifierCrawlerListener(final AccountsManager accountsManager,
+      final PhoneNumberIdentifiers phoneNumberIdentifiers) {
+
+    this.accountsManager = accountsManager;
+    this.phoneNumberIdentifiers = phoneNumberIdentifiers;
+  }
+
+  @Override
+  public void onCrawlStart() {
+  }
+
+  @Override
+  public void onCrawlEnd(final Optional<UUID> fromUuid) {
+  }
+
+  @Override
+  protected void onCrawlChunk(final Optional<UUID> fromUuid, final List<Account> 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();
+        });
+  }
+}

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<Account>) invocation -> {
+      final Account account = invocation.getArgument(0, Account.class);
+      final Consumer<Account> updater = invocation.getArgument(1, Consumer.class);
+
+      updater.accept(account);
+
+      return account;
+    });
+
+    final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
+    when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString())).thenAnswer(
+        (Answer<UUID>) 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());
+  }
+}