From 23a076a204d1f8a2efd6915a663995ca0bc40234 Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Thu, 9 Sep 2021 17:52:45 -0700
Subject: [PATCH] =?UTF-8?q?Update=20`Account#getNextDeviceId`=20to=20not?=
 =?UTF-8?q?=20reuse=20disable=20device=E2=80=99s=20IDs?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../textsecuregcm/storage/Account.java        | 12 +++------
 .../tests/storage/AccountTest.java            | 27 +++++++++++++++++++
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Account.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Account.java
index 3b7b8d41a..fc04c2ad9 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Account.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Account.java
@@ -212,17 +212,13 @@ public class Account {
   public long getNextDeviceId() {
     requireNotStale();
 
-    long highestDevice = Device.MASTER_ID;
+    long candidateId = Device.MASTER_ID + 1;
 
-    for (Device device : devices) {
-      if (!device.isEnabled()) {
-        return device.getId();
-      } else if (device.getId() > highestDevice) {
-        highestDevice = device.getId();
-      }
+    while (getDevice(candidateId).isPresent()) {
+      candidateId++;
     }
 
-    return highestDevice + 1;
+    return candidateId;
   }
 
   public int getEnabledDeviceCount() {
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountTest.java
index 93bf16ec9..6e051dc08 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountTest.java
@@ -12,6 +12,8 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
+import static org.whispersystems.textsecuregcm.tests.util.DevicesHelper.createDevice;
+import static org.whispersystems.textsecuregcm.tests.util.DevicesHelper.setEnabled;
 
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
@@ -315,4 +317,29 @@ class AccountTest {
     assertThrows(AssertionError.class, account::getNumber);
     assertDoesNotThrow(account::getUuid);
   }
+
+  @Test
+  void getNextDeviceId() {
+
+    final Set<Device> devices = new HashSet<>();
+    devices.add(createDevice(Device.MASTER_ID));
+
+    final Account account = new Account("+14151234567", UUID.randomUUID(), devices, new byte[0]);
+
+    assertThat(account.getNextDeviceId()).isEqualTo(2L);
+
+    account.addDevice(createDevice(2L));
+
+    assertThat(account.getNextDeviceId()).isEqualTo(3L);
+
+    account.addDevice(createDevice(3L));
+
+    setEnabled(account.getDevice(2L).orElseThrow(), false);
+
+    assertThat(account.getNextDeviceId()).isEqualTo(4L);
+
+    account.removeDevice(2L);
+
+    assertThat(account.getNextDeviceId()).isEqualTo(2L);
+  }
 }