From 1e1394560daf8027c352d44fdae1e6e2a7858fc2 Mon Sep 17 00:00:00 2001
From: Jon Chambers <63609320+jon-signal@users.noreply.github.com>
Date: Tue, 9 Nov 2021 11:42:44 -0500
Subject: [PATCH] Check length of cancellation reason list before getting
 reason codes

---
 .../textsecuregcm/storage/Accounts.java       |  2 +-
 .../textsecuregcm/storage/AccountsTest.java   | 24 +++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

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 8d5d8ded4..f07021966 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java
@@ -387,7 +387,7 @@ public class Accounts extends AbstractDynamoDbStore {
 
       } catch (final TransactionCanceledException e) {
 
-        if ("ConditionalCheckFailed".equals(e.cancellationReasons().get(1).code())) {
+        if (e.cancellationReasons().size() > 1 && "ConditionalCheckFailed".equals(e.cancellationReasons().get(1).code())) {
           log.error("Conflicting phone number mapping exists for account {}, PNI {}", account.getUuid(), account.getPhoneNumberIdentifier());
           throw 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 837d3844e..394f7fa09 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java
@@ -39,6 +39,7 @@ import org.whispersystems.textsecuregcm.util.SystemMapper;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+import software.amazon.awssdk.services.dynamodb.model.CancellationReason;
 import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
@@ -366,6 +367,29 @@ class AccountsTest {
     assertThatThrownBy(() -> accounts.update(account)).isInstanceOfAny(ContestedOptimisticLockException.class);
   }
 
+  @Test
+  // TODO Remove after initial migration is complete
+  void testUpdateWithMockTransactionCancellationException() {
+
+    final DynamoDbClient dynamoDbClient = mock(DynamoDbClient.class);
+    accounts = new Accounts(dynamoDbClient,
+        dynamoDbExtension.getTableName(), NUMBER_CONSTRAINT_TABLE_NAME, PNI_CONSTRAINT_TABLE_NAME, SCAN_PAGE_SIZE);
+
+    when(dynamoDbClient.transactWriteItems(any(TransactWriteItemsRequest.class)))
+        .thenThrow(TransactionCanceledException.builder()
+            .cancellationReasons(CancellationReason.builder()
+                .code("Test")
+                .build())
+            .build());
+
+    when(dynamoDbClient.getItem(any(GetItemRequest.class))).thenReturn(GetItemResponse.builder().build());
+
+    Device device = generateDevice(1);
+    Account account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), Collections.singleton(device));
+
+    assertThatThrownBy(() -> accounts.update(account)).isInstanceOfAny(TransactionCanceledException.class);
+  }
+
   @Test
   void testRetrieveFrom() {
     List<Account> users = new ArrayList<>();