From 9a19ef82fd43a90adca19578dfd17e0a02ef61b1 Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Fri, 22 Nov 2024 15:24:48 -0600
Subject: [PATCH] Use pni in DynamoDB account put condition expression

---
 .../textsecuregcm/storage/Accounts.java       |  6 +--
 .../textsecuregcm/storage/AccountsTest.java   | 39 +++++++++++++++++++
 2 files changed, 42 insertions(+), 3 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 8e6adf473..66b7e6761 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java
@@ -1419,9 +1419,9 @@ public class Accounts extends AbstractDynamoDbStore {
 
     return TransactWriteItem.builder()
         .put(Put.builder()
-            .conditionExpression("attribute_not_exists(#number) OR #number = :number")
-            .expressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164))
-            .expressionAttributeValues(Map.of(":number", numberAttr))
+            .conditionExpression("attribute_not_exists(#pni) OR #pni = :pni")
+            .expressionAttributeNames(Map.of("#pni", ATTR_PNI_UUID))
+            .expressionAttributeValues(Map.of(":pni", pniUuidAttr))
             .tableName(accountsTableName)
             .item(item)
             .build())
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 ce58c046d..4b407cd7f 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java
@@ -235,6 +235,45 @@ class AccountsTest {
     assertPhoneNumberIdentifierConstraintExists(account.getPhoneNumberIdentifier(), account.getUuid());
   }
 
+  @Test
+  void testStoreAciCollisionFails() {
+    Device device = generateDevice(DEVICE_ID_1);
+    Account account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), List.of(device));
+
+    boolean freshUser = createAccount(account);
+
+    assertThat(freshUser).isTrue();
+    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier(), null, account, true);
+
+    assertPhoneNumberConstraintExists("+14151112222", account.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(account.getPhoneNumberIdentifier(), account.getUuid());
+
+    account.setNumber("+14153334444", UUID.randomUUID());
+    assertThrows(IllegalArgumentException.class, () -> createAccount(account),
+        "Reusing ACI with different PNI should fail");
+  }
+
+  @Test
+  void testStorePniCollisionFails() {
+    Device device1 = generateDevice(DEVICE_ID_1);
+    Account account1 = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), List.of(device1));
+
+    boolean freshUser = createAccount(account1);
+
+    assertThat(freshUser).isTrue();
+    verifyStoredState("+14151112222", account1.getUuid(), account1.getPhoneNumberIdentifier(), null, account1, true);
+
+    assertPhoneNumberConstraintExists("+14151112222", account1.getUuid());
+    assertPhoneNumberIdentifierConstraintExists(account1.getPhoneNumberIdentifier(), account1.getUuid());
+
+    Device device2 = generateDevice(DEVICE_ID_1);
+    Account account2 = generateAccount("+14151112222", UUID.randomUUID(), account1.getPhoneNumberIdentifier(),
+        List.of(device2));
+
+    assertThrows(AccountAlreadyExistsException.class, () -> accounts.create(account2, Collections.emptyList()),
+        "New ACI with same PNI should fail");
+  }
+
   @Test
   void testRetrieve() {
     final List<Device> devicesFirst = List.of(generateDevice(DEVICE_ID_1), generateDevice(DEVICE_ID_2));