Handle split contacts in storage service when in PNP mode.

app/src/androidTest/java/org/thoughtcrime/securesms/database/RecipientTableTest.kt

@@ -183,6 +183,25 @@ class RecipientTableTest {
     assertNotEquals(byAci, byE164)
   }
 
+  @Test
+  fun givenARecipientWithPniAndAci_whenISplitItForStorageSync_thenIExpectItToBeSplit() {
+    FeatureFlagsAccessor.forceValue(FeatureFlags.PHONE_NUMBER_PRIVACY, true)
+
+    val mainId = SignalDatabase.recipients.getAndPossiblyMerge(ACI_A, PNI_A, E164_A)
+    val mainRecord = SignalDatabase.recipients.getRecord(mainId)
+
+    SignalDatabase.recipients.splitForStorageSync(mainRecord.storageId!!)
+
+    val byAci: RecipientId = SignalDatabase.recipients.getByServiceId(ACI_A).get()
+
+    val byE164: RecipientId = SignalDatabase.recipients.getByE164(E164_A).get()
+    val byPni: RecipientId = SignalDatabase.recipients.getByServiceId(PNI_A).get()
+
+    assertEquals(mainId, byAci)
+    assertEquals(byE164, byPni)
+    assertNotEquals(byAci, byE164)
+  }
+
   companion object {
     val ACI_A = ACI.from(UUID.fromString("aaaa0000-5a76-47fa-a98a-7e72c948a82e"))
     val PNI_A = PNI.from(UUID.fromString("aaaa1111-c960-4f6c-8385-671ad2ffb999"))

app/src/androidTest/java/org/thoughtcrime/securesms/database/RecipientTableTest_getAndPossiblyMerge.kt

@@ -58,6 +58,33 @@ class RecipientTableTest_getAndPossiblyMerge {
     FeatureFlagsAccessor.forceValue(FeatureFlags.PHONE_NUMBER_PRIVACY, true)
   }
 
+  @Test
+  fun allNonMergeTests() {
+    test("e164-only insert") {
+      val id = process(E164_A, null, null)
+      expect(E164_A, null, null)
+
+      val record = SignalDatabase.recipients.getRecord(id)
+      assertEquals(RecipientTable.RegisteredState.UNKNOWN, record.registered)
+    }
+
+    test("pni-only insert") {
+      val id = process(null, PNI_A, null)
+      expect(null, PNI_A, null)
+
+      val record = SignalDatabase.recipients.getRecord(id)
+      assertEquals(RecipientTable.RegisteredState.REGISTERED, record.registered)
+    }
+
+    test("aci-only insert") {
+      val id = process(null, null, ACI_A)
+      expect(null, null, ACI_A)
+
+      val record = SignalDatabase.recipients.getRecord(id)
+      assertEquals(RecipientTable.RegisteredState.REGISTERED, record.registered)
+    }
+  }
+
   @Test
   fun allSimpleTests() {
     test("no match, e164-only") {
@@ -841,9 +868,10 @@ class RecipientTableTest_getAndPossiblyMerge {
       return id
     }
 
-    fun process(e164: String?, pni: PNI?, aci: ACI?, changeSelf: Boolean = false, pniVerified: Boolean = false) {
+    fun process(e164: String?, pni: PNI?, aci: ACI?, changeSelf: Boolean = false, pniVerified: Boolean = false): RecipientId {
       outputRecipientId = SignalDatabase.recipients.getAndPossiblyMerge(serviceId = aci ?: pni, pni = pni, e164 = e164, pniVerified = pniVerified, changeSelf = changeSelf)
       generatedIds += outputRecipientId
+      return outputRecipientId
     }
 
     fun expect(e164: String?, pni: PNI?, aci: ACI?) {

app/src/androidTest/java/org/thoughtcrime/securesms/storage/ContactRecordProcessorTest.kt

@@ -0,0 +1,127 @@
+package org.thoughtcrime.securesms.storage
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertNotEquals
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.signal.core.util.update
+import org.thoughtcrime.securesms.database.RecipientTable
+import org.thoughtcrime.securesms.database.SignalDatabase
+import org.thoughtcrime.securesms.keyvalue.SignalStore
+import org.thoughtcrime.securesms.recipients.RecipientId
+import org.thoughtcrime.securesms.util.Base64
+import org.thoughtcrime.securesms.util.FeatureFlags
+import org.thoughtcrime.securesms.util.FeatureFlagsAccessor
+import org.whispersystems.signalservice.api.push.ACI
+import org.whispersystems.signalservice.api.push.PNI
+import org.whispersystems.signalservice.api.storage.SignalContactRecord
+import org.whispersystems.signalservice.api.storage.StorageId
+import org.whispersystems.signalservice.internal.storage.protos.ContactRecord
+import java.util.UUID
+
+@RunWith(AndroidJUnit4::class)
+class ContactRecordProcessorTest {
+
+  @Before
+  fun setup() {
+    SignalStore.account().setE164(E164_SELF)
+    SignalStore.account().setAci(ACI_SELF)
+    SignalStore.account().setPni(PNI_SELF)
+    FeatureFlagsAccessor.forceValue(FeatureFlags.PHONE_NUMBER_PRIVACY, true)
+  }
+
+  @Test
+  fun process_splitContact_normalSplit() {
+    // GIVEN
+    val originalId = SignalDatabase.recipients.getAndPossiblyMerge(ACI_A, PNI_A, E164_A)
+    setStorageId(originalId, STORAGE_ID_A)
+
+    val remote1 = buildRecord(STORAGE_ID_B) {
+      setServiceId(ACI_A.toString())
+      setUnregisteredAtTimestamp(100)
+    }
+
+    val remote2 = buildRecord(STORAGE_ID_C) {
+      setServiceId(PNI_A.toString())
+      setServicePni(PNI_A.toString())
+      setServiceE164(E164_A)
+    }
+
+    // WHEN
+    val subject = ContactRecordProcessor()
+    subject.process(listOf(remote1, remote2), StorageSyncHelper.KEY_GENERATOR)
+
+    // THEN
+    val byAci: RecipientId = SignalDatabase.recipients.getByServiceId(ACI_A).get()
+
+    val byE164: RecipientId = SignalDatabase.recipients.getByE164(E164_A).get()
+    val byPni: RecipientId = SignalDatabase.recipients.getByServiceId(PNI_A).get()
+
+    assertEquals(originalId, byAci)
+    assertEquals(byE164, byPni)
+    assertNotEquals(byAci, byE164)
+  }
+
+  @Test
+  fun process_splitContact_doNotSplitIfAciRecordIsRegistered() {
+    // GIVEN
+    val originalId = SignalDatabase.recipients.getAndPossiblyMerge(ACI_A, PNI_A, E164_A)
+    setStorageId(originalId, STORAGE_ID_A)
+
+    val remote1 = buildRecord(STORAGE_ID_B) {
+      setServiceId(ACI_A.toString())
+      setUnregisteredAtTimestamp(0)
+    }
+
+    val remote2 = buildRecord(STORAGE_ID_C) {
+      setServiceId(PNI_A.toString())
+      setServicePni(PNI_A.toString())
+      setServiceE164(E164_A)
+    }
+
+    // WHEN
+    val subject = ContactRecordProcessor()
+    subject.process(listOf(remote1, remote2), StorageSyncHelper.KEY_GENERATOR)
+
+    // THEN
+    val byAci: RecipientId = SignalDatabase.recipients.getByServiceId(ACI_A).get()
+    val byE164: RecipientId = SignalDatabase.recipients.getByE164(E164_A).get()
+    val byPni: RecipientId = SignalDatabase.recipients.getByPni(PNI_A).get()
+
+    assertEquals(originalId, byAci)
+    assertEquals(byE164, byPni)
+    assertEquals(byAci, byE164)
+  }
+
+  private fun buildRecord(id: StorageId, applyParams: ContactRecord.Builder.() -> ContactRecord.Builder): SignalContactRecord {
+    return SignalContactRecord(id, ContactRecord.getDefaultInstance().toBuilder().applyParams().build())
+  }
+
+  private fun setStorageId(recipientId: RecipientId, storageId: StorageId) {
+    SignalDatabase.rawDatabase
+      .update(RecipientTable.TABLE_NAME)
+      .values(RecipientTable.STORAGE_SERVICE_ID to Base64.encodeBytes(storageId.raw))
+      .where("${RecipientTable.ID} = ?", recipientId)
+      .run()
+  }
+
+  companion object {
+    val ACI_A = ACI.from(UUID.fromString("aaaa0000-5a76-47fa-a98a-7e72c948a82e"))
+    val ACI_B = ACI.from(UUID.fromString("bbbb0000-0b60-4a68-9cd9-ed2f8453f9ed"))
+    val ACI_SELF = ACI.from(UUID.fromString("77770000-b477-4f35-a824-d92987a63641"))
+
+    val PNI_A = PNI.from(UUID.fromString("aaaa1111-c960-4f6c-8385-671ad2ffb999"))
+    val PNI_B = PNI.from(UUID.fromString("bbbb1111-cd55-40bf-adda-c35a85375533"))
+    val PNI_SELF = PNI.from(UUID.fromString("77771111-b014-41fb-bf73-05cb2ec52910"))
+
+    const val E164_A = "+12222222222"
+    const val E164_B = "+13333333333"
+    const val E164_SELF = "+10000000000"
+
+    val STORAGE_ID_A: StorageId = StorageId.forContact(byteArrayOf(1, 2, 3, 4))
+    val STORAGE_ID_B: StorageId = StorageId.forContact(byteArrayOf(5, 6, 7, 8))
+    val STORAGE_ID_C: StorageId = StorageId.forContact(byteArrayOf(9, 10, 11, 12))
+  }
+}

app/src/main/java/org/thoughtcrime/securesms/database/RecipientTable.kt

@@ -2236,6 +2236,7 @@ open class RecipientTable(context: Context, databaseHelper: SignalDatabase) : Da
 
     if (update(id, contentValues)) {
       Log.i(TAG, "[WithSplit] Newly marked $id as unregistered.")
+      markNeedsSync(id)
       ApplicationDependencies.getDatabaseObserver().notifyRecipientChanged(id)
     }
 
@@ -2254,10 +2255,33 @@ open class RecipientTable(context: Context, databaseHelper: SignalDatabase) : Da
 
     if (update(id, contentValues)) {
       Log.i(TAG, "[WithoutSplit] Newly marked $id as unregistered.")
+      markNeedsSync(id)
       ApplicationDependencies.getDatabaseObserver().notifyRecipientChanged(id)
     }
   }
 
+  /**
+   * Removes the target recipient's E164+PNI, then creates a new recipient with that E164+PNI.
+   * Done so we can match a split contact during storage sync.
+   */
+  fun splitForStorageSync(storageId: ByteArray) {
+    check(FeatureFlags.phoneNumberPrivacy())
+
+    val record = getByStorageId(storageId)!!
+    check(record.serviceId != null && record.pni != null && record.serviceId != record.pni)
+
+    writableDatabase
+      .update(TABLE_NAME)
+      .values(
+        PNI_COLUMN to null,
+        PHONE to null
+      )
+      .where("$ID = ?", record.id)
+      .run()
+
+    getAndPossiblyMerge(record.pni, record.pni, record.e164)
+  }
+
   fun bulkUpdatedRegisteredStatus(registered: Map<RecipientId, ServiceId?>, unregistered: Collection<RecipientId>) {
     writableDatabase.withinTransaction {
       val registeredWithServiceId: Set<RecipientId> = getRegisteredWithServiceIds()

app/src/main/java/org/thoughtcrime/securesms/storage/ContactRecordProcessor.java

@@ -19,9 +19,15 @@ import org.whispersystems.signalservice.api.storage.SignalContactRecord;
 import org.whispersystems.signalservice.api.util.OptionalUtil;
 import org.whispersystems.signalservice.internal.storage.protos.ContactRecord.IdentityState;
 
+import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.TreeMap;
+import java.util.TreeSet;
 
 public class ContactRecordProcessor extends DefaultStorageRecordProcessor<SignalContactRecord> {
 
@@ -47,6 +53,69 @@ public class ContactRecordProcessor extends DefaultStorageRecordProcessor<Signal
     this.selfE164       = selfE164;
   }
 
+  /**
+   * For contact records specifically, we have some extra work that needs to be done before we process all of the records.
+   *
+   * We have to look and see if there is an unregistered ACI-only record and another E164/PNI-only record that points to the
+   * same local contact row.
+   *
+   * If so, we actually want to mimic the split and turn them into two separate contact rows locally. The reasons are nuanced,
+   * but the TL;DR is that we want to split unregistered users into separate rows so that a user could re-register and get a
+   * different ACI.
+   */
+  @Override
+  public void process(@NonNull Collection<SignalContactRecord> remoteRecords, @NonNull StorageKeyGenerator keyGenerator) throws IOException {
+    if (!FeatureFlags.phoneNumberPrivacy()) {
+      super.process(remoteRecords, keyGenerator);
+      return;
+    }
+
+    List<SignalContactRecord> unregisteredAciOnly = new ArrayList<>();
+    List<SignalContactRecord> pniE164Only         = new ArrayList<>();
+
+    for (SignalContactRecord remoteRecord : remoteRecords) {
+      if (isInvalid(remoteRecord)) {
+        continue;
+      }
+
+      if (remoteRecord.getUnregisteredTimestamp() > 0 && remoteRecord.getServiceId() != null && !remoteRecord.getPni().isPresent() && !remoteRecord.getNumber().isPresent()) {
+        unregisteredAciOnly.add(remoteRecord);
+      } else if (remoteRecord.getServiceId() != null && remoteRecord.getServiceId().equals(remoteRecord.getPni().orElse(null))) {
+        pniE164Only.add(remoteRecord);
+      }
+    }
+
+    if (unregisteredAciOnly.isEmpty() || pniE164Only.isEmpty()) {
+      super.process(remoteRecords, keyGenerator);
+      return;
+    }
+
+    Log.i(TAG, "We have some unregistered ACI-only contacts as well as some PNI-only contacts. Need to do an intersection to detect any possible required splits.");
+
+    TreeSet<SignalContactRecord> localMatches = new TreeSet<>(this);
+
+    for (SignalContactRecord aciOnly : unregisteredAciOnly) {
+      Optional<SignalContactRecord> localMatch = getMatching(aciOnly, keyGenerator);
+
+      if (localMatch.isPresent()) {
+        localMatches.add(localMatch.get());
+      }
+    }
+
+    for (SignalContactRecord pniOnly : pniE164Only) {
+      Optional<SignalContactRecord> localMatch = getMatching(pniOnly, keyGenerator);
+
+      if (localMatch.isPresent() && localMatches.contains(localMatch.get())) {
+        Log.w(TAG, "Found a situation where we need to split our local record in two in order to match the remote state.");
+
+        SignalDatabase.recipients().splitForStorageSync(localMatch.get().getId().getRaw());
+      }
+    }
+
+
+    super.process(remoteRecords, keyGenerator);
+  }
+
   /**
    * Error cases:
    * - You can't have a contact record without an address component.

app/src/test/java/org/thoughtcrime/securesms/storage/ContactRecordProcessorTest.kt

@@ -311,9 +311,9 @@ class ContactRecordProcessorTest {
   }
 
   companion object {
-    val STORAGE_ID_A: StorageId = StorageId.forStoryDistributionList(byteArrayOf(1, 2, 3, 4))
+    val STORAGE_ID_A: StorageId = StorageId.forContact(byteArrayOf(1, 2, 3, 4))
-    val STORAGE_ID_B: StorageId = StorageId.forStoryDistributionList(byteArrayOf(5, 6, 7, 8))
+    val STORAGE_ID_B: StorageId = StorageId.forContact(byteArrayOf(5, 6, 7, 8))
-    val STORAGE_ID_C: StorageId = StorageId.forStoryDistributionList(byteArrayOf(5, 6, 7, 8))
+    val STORAGE_ID_C: StorageId = StorageId.forContact(byteArrayOf(9, 10, 11, 12))
 
     val ACI_A = ACI.from(UUID.fromString("3436efbe-5a76-47fa-a98a-7e72c948a82e"))
     val ACI_B = ACI.from(UUID.fromString("8de7f691-0b60-4a68-9cd9-ed2f8453f9ed"))