Allow callers to set/retrieve keys by ACI or PNI

service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java

@@ -21,10 +21,12 @@ import javax.validation.Valid;
 import javax.ws.rs.Consumes;
 import javax.ws.rs.GET;
 import javax.ws.rs.HeaderParam;
+import javax.ws.rs.NotFoundException;
 import javax.ws.rs.PUT;
 import javax.ws.rs.Path;
 import javax.ws.rs.PathParam;
 import javax.ws.rs.Produces;
+import javax.ws.rs.QueryParam;
 import javax.ws.rs.WebApplicationException;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
@@ -77,8 +79,10 @@ public class KeysController {
 
   @GET
   @Produces(MediaType.APPLICATION_JSON)
-  public PreKeyCount getStatus(@Auth AuthenticatedAccount auth) {
+  public PreKeyCount getStatus(@Auth final AuthenticatedAccount auth,
-    int count = keys.getCount(auth.getAccount().getUuid(), auth.getAuthenticatedDevice().getId());
+      @QueryParam("identity") final Optional<String> identityType) {
+
+    int count = keys.getCount(getIdentifier(auth.getAccount(), identityType), auth.getAuthenticatedDevice().getId());
 
     if (count > 0) {
       count = count - 1;
@@ -90,7 +94,9 @@ public class KeysController {
   @Timed
   @PUT
   @Consumes(MediaType.APPLICATION_JSON)
-  public void setKeys(@Auth DisabledPermittedAuthenticatedAccount disabledPermittedAuth, @Valid PreKeyState preKeys) {
+  public void setKeys(@Auth final DisabledPermittedAuthenticatedAccount disabledPermittedAuth,
+      @Valid final PreKeyState preKeys,
+      @QueryParam("identity") final Optional<String> identityType) {
     Account account = disabledPermittedAuth.getAccount();
     Device device = disabledPermittedAuth.getAuthenticatedDevice();
     boolean updateAccount = false;
@@ -103,14 +109,27 @@ public class KeysController {
       updateAccount = true;
     }
 
+    final boolean usePhoneNumberIdentity = usePhoneNumberIdentity(identityType);
+
     if (updateAccount) {
       account = accounts.update(account, a -> {
-        a.getDevice(device.getId()).ifPresent(d -> d.setSignedPreKey(preKeys.getSignedPreKey()));
+        a.getDevice(device.getId()).ifPresent(d -> {
-        a.setIdentityKey(preKeys.getIdentityKey());
+          if (usePhoneNumberIdentity) {
+            d.setPhoneNumberIdentitySignedPreKey(preKeys.getSignedPreKey());
+          } else {
+            d.setSignedPreKey(preKeys.getSignedPreKey());
+          }
+        });
+
+        if (usePhoneNumberIdentity) {
+          a.setPhoneNumberIdentityKey(preKeys.getIdentityKey());
+        } else {
+          a.setIdentityKey(preKeys.getIdentityKey());
+        }
       });
     }
 
-    keys.store(account.getUuid(), device.getId(), preKeys.getPreKeys());
+    keys.store(getIdentifier(account, identityType), device.getId(), preKeys.getPreKeys());
   }
 
   @Timed
@@ -130,14 +149,19 @@ public class KeysController {
 
     final Optional<Account> account = auth.map(AuthenticatedAccount::getAccount);
 
-    Optional<Account> target = accounts.getByAccountIdentifier(targetUuid);
+    final Account target;
-    OptionalAccess.verify(account, accessKey, target, deviceId);
+    {
+      final Optional<Account> maybeTarget = accounts.getByAccountIdentifier(targetUuid)
+          .or(() -> accounts.getByPhoneNumberIdentifier(targetUuid));
 
-    assert (target.isPresent());
+      OptionalAccess.verify(account, accessKey, maybeTarget, deviceId);
+
+      target = maybeTarget.orElseThrow();
+    }
 
     {
       final String sourceCountryCode = account.map(a -> Util.getCountryCode(a.getNumber())).orElse("0");
-      final String targetCountryCode = target.map(a -> Util.getCountryCode(a.getNumber())).orElseThrow();
+      final String targetCountryCode = Util.getCountryCode(target.getNumber());
 
       Metrics.counter(PREKEY_REQUEST_COUNTER_NAME, Tags.of(
           SOURCE_COUNTRY_TAG_NAME, sourceCountryCode,
@@ -147,7 +171,7 @@ public class KeysController {
 
     if (account.isPresent()) {
       rateLimiters.getPreKeysLimiter().validate(
-          account.get().getUuid() + "." + auth.get().getAuthenticatedDevice().getId() + "__" + target.get().getUuid()
+          account.get().getUuid() + "." + auth.get().getAuthenticatedDevice().getId() + "__" + targetUuid
               + "." + deviceId);
 
       try {
@@ -168,12 +192,15 @@ public class KeysController {
       }
     }
 
-    Map<Long, PreKey>        preKeysByDeviceId = getLocalKeys(target.get(), deviceId);
+    final boolean usePhoneNumberIdentity =
+        target.getPhoneNumberIdentifier().map(pni -> pni.equals(targetUuid)).orElse(false);
+
+    Map<Long, PreKey>        preKeysByDeviceId = getLocalKeys(target, deviceId, usePhoneNumberIdentity);
     List<PreKeyResponseItem> responseItems     = new LinkedList<>();
 
-    for (Device device : target.get().getDevices()) {
+    for (Device device : target.getDevices()) {
       if (device.isEnabled() && (deviceId.equals("*") || device.getId() == Long.parseLong(deviceId))) {
-        SignedPreKey signedPreKey = device.getSignedPreKey();
+        SignedPreKey signedPreKey = usePhoneNumberIdentity ? device.getPhoneNumberIdentitySignedPreKey() : device.getSignedPreKey();
         PreKey       preKey       = preKeysByDeviceId.get(device.getId());
 
         if (signedPreKey != null || preKey != null) {
@@ -182,49 +209,73 @@ public class KeysController {
       }
     }
 
+    final String identityKey = usePhoneNumberIdentity ? target.getPhoneNumberIdentityKey() : target.getIdentityKey();
+
     if (responseItems.isEmpty()) return Response.status(404).build();
-    else                         return Response.ok().entity(new PreKeyResponse(target.get().getIdentityKey(), responseItems)).build();
+    else                         return Response.ok().entity(new PreKeyResponse(identityKey, responseItems)).build();
   }
 
   @Timed
   @PUT
   @Path("/signed")
   @Consumes(MediaType.APPLICATION_JSON)
-  public void setSignedKey(@Auth AuthenticatedAccount auth, @Valid SignedPreKey signedPreKey) {
+  public void setSignedKey(@Auth final AuthenticatedAccount auth,
+      @Valid final SignedPreKey signedPreKey,
+      @QueryParam("identity") final Optional<String> identityType) {
+
     Device device = auth.getAuthenticatedDevice();
 
-    accounts.updateDevice(auth.getAccount(), device.getId(), d -> d.setSignedPreKey(signedPreKey));
+    accounts.updateDevice(auth.getAccount(), device.getId(), d -> {
+      if (usePhoneNumberIdentity(identityType)) {
+        d.setPhoneNumberIdentitySignedPreKey(signedPreKey);
+      } else {
+        d.setSignedPreKey(signedPreKey);
+      }
+    });
   }
 
   @Timed
   @GET
   @Path("/signed")
   @Produces(MediaType.APPLICATION_JSON)
-  public Optional<SignedPreKey> getSignedKey(@Auth AuthenticatedAccount auth) {
+  public Optional<SignedPreKey> getSignedKey(@Auth final AuthenticatedAccount auth,
-    Device device = auth.getAuthenticatedDevice();
+      @QueryParam("identity") final Optional<String> identityType) {
-    SignedPreKey signedPreKey = device.getSignedPreKey();
 
-    if (signedPreKey != null) {
+    Device device = auth.getAuthenticatedDevice();
-      return Optional.of(signedPreKey);
+    SignedPreKey signedPreKey = usePhoneNumberIdentity(identityType) ?
-    } else {
+        device.getPhoneNumberIdentitySignedPreKey() : device.getSignedPreKey();
-      return Optional.empty();
+
-    }
+    return Optional.ofNullable(signedPreKey);
   }
 
-  private Map<Long, PreKey> getLocalKeys(Account destination, String deviceIdSelector) {
+  private static boolean usePhoneNumberIdentity(final Optional<String> identityType) {
+    return "pni".equals(identityType.map(String::toLowerCase).orElse("aci"));
+  }
+
+  private static UUID getIdentifier(final Account account, final Optional<String> identityType) {
+    return usePhoneNumberIdentity(identityType) ?
+        account.getPhoneNumberIdentifier().orElseThrow(NotFoundException::new) :
+        account.getUuid();
+  }
+
+  private Map<Long, PreKey> getLocalKeys(Account destination, String deviceIdSelector, final boolean usePhoneNumberIdentity) {
     final Map<Long, PreKey> preKeys;
 
+    final UUID identifier = usePhoneNumberIdentity ?
+        destination.getPhoneNumberIdentifier().orElseThrow(NotFoundException::new) :
+        destination.getUuid();
+
     if (deviceIdSelector.equals("*")) {
       preKeys = new HashMap<>();
 
       for (final Device device : destination.getDevices()) {
-        keys.take(destination.getUuid(), device.getId()).ifPresent(preKey -> preKeys.put(device.getId(), preKey));
+        keys.take(identifier, device.getId()).ifPresent(preKey -> preKeys.put(device.getId(), preKey));
       }
     } else {
       try {
         long deviceId = Long.parseLong(deviceIdSelector);
 
-        preKeys = keys.take(destination.getUuid(), deviceId)
+        preKeys = keys.take(identifier, deviceId)
             .map(preKey -> Map.of(deviceId, preKey))
             .orElse(Collections.emptyMap());
       } catch (NumberFormatException e) {

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

@@ -47,6 +47,9 @@ public class Account {
   @JsonProperty
   private String identityKey;
 
+  @JsonProperty("pniIdentityKey")
+  private String phoneNumberIdentityKey;
+
   @JsonProperty("cpv")
   private String currentProfileVersion;
 
@@ -290,6 +293,14 @@ public class Account {
     return identityKey;
   }
 
+  public String getPhoneNumberIdentityKey() {
+    return phoneNumberIdentityKey;
+  }
+
+  public void setPhoneNumberIdentityKey(final String phoneNumberIdentityKey) {
+    this.phoneNumberIdentityKey = phoneNumberIdentityKey;
+  }
+
   public long getLastSeen() {
     requireNotStale();
 
@@ -506,5 +517,4 @@ public class Account {
       logger.error("Accessor called on stale account", new RuntimeException());
     }
   }
-
 }

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

@@ -53,6 +53,9 @@ public class Device {
   @JsonProperty
   private SignedPreKey signedPreKey;
 
+  @JsonProperty("pniSignedPreKey")
+  private SignedPreKey phoneNumberIdentitySignedPreKey;
+
   @JsonProperty
   private long lastSeen;
 
@@ -215,6 +218,14 @@ public class Device {
     this.signedPreKey = signedPreKey;
   }
 
+  public SignedPreKey getPhoneNumberIdentitySignedPreKey() {
+    return phoneNumberIdentitySignedPreKey;
+  }
+
+  public void setPhoneNumberIdentitySignedPreKey(final SignedPreKey phoneNumberIdentitySignedPreKey) {
+    this.phoneNumberIdentitySignedPreKey = phoneNumberIdentitySignedPreKey;
+  }
+
   public long getPushTimestamp() {
     return pushTimestamp;
   }

service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java

@@ -8,9 +8,11 @@ package org.whispersystems.textsecuregcm.tests.controllers;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
@@ -68,6 +70,7 @@ class KeysControllerTest {
 
   private static final String EXISTS_NUMBER = "+14152222222";
   private static final UUID   EXISTS_UUID   = UUID.randomUUID();
+  private static final UUID   EXISTS_PNI    = UUID.randomUUID();
 
   private static final String NOT_EXISTS_NUMBER = "+14152222220";
   private static final UUID   NOT_EXISTS_UUID   = UUID.randomUUID();
@@ -81,11 +84,16 @@ class KeysControllerTest {
   private final PreKey SAMPLE_KEY3   = new PreKey(334, "test5");
   private final PreKey SAMPLE_KEY4   = new PreKey(336, "test6");
 
+  private final PreKey SAMPLE_KEY_PNI = new PreKey(7777, "test7");
 
   private final SignedPreKey SAMPLE_SIGNED_KEY       = new SignedPreKey( 1111, "foofoo", "sig11"    );
   private final SignedPreKey SAMPLE_SIGNED_KEY2      = new SignedPreKey( 2222, "foobar", "sig22"    );
   private final SignedPreKey SAMPLE_SIGNED_KEY3      = new SignedPreKey( 3333, "barfoo", "sig33"    );
+  private final SignedPreKey SAMPLE_SIGNED_PNI_KEY   = new SignedPreKey( 4444, "foofoopni", "sig44" );
+  private final SignedPreKey SAMPLE_SIGNED_PNI_KEY2  = new SignedPreKey( 5555, "foobarpni", "sig55" );
+  private final SignedPreKey SAMPLE_SIGNED_PNI_KEY3  = new SignedPreKey( 6666, "barfoopni", "sig66" );
   private final SignedPreKey VALID_DEVICE_SIGNED_KEY = new SignedPreKey(89898, "zoofarb", "sigvalid");
+  private final SignedPreKey VALID_DEVICE_PNI_SIGNED_KEY = new SignedPreKey(7777, "zoofarber", "sigvalidest");
 
   private final static Keys KEYS = mock(Keys.class               );
   private final static AccountsManager             accounts                    = mock(AccountsManager.class            );
@@ -135,12 +143,17 @@ class KeysControllerTest {
     when(sampleDevice2.getSignedPreKey()).thenReturn(SAMPLE_SIGNED_KEY2);
     when(sampleDevice3.getSignedPreKey()).thenReturn(SAMPLE_SIGNED_KEY3);
     when(sampleDevice4.getSignedPreKey()).thenReturn(null);
+    when(sampleDevice.getPhoneNumberIdentitySignedPreKey()).thenReturn(SAMPLE_SIGNED_PNI_KEY);
+    when(sampleDevice2.getPhoneNumberIdentitySignedPreKey()).thenReturn(SAMPLE_SIGNED_PNI_KEY2);
+    when(sampleDevice3.getPhoneNumberIdentitySignedPreKey()).thenReturn(SAMPLE_SIGNED_PNI_KEY3);
+    when(sampleDevice4.getPhoneNumberIdentitySignedPreKey()).thenReturn(null);
     when(sampleDevice.getId()).thenReturn(1L);
     when(sampleDevice2.getId()).thenReturn(2L);
     when(sampleDevice3.getId()).thenReturn(3L);
     when(sampleDevice4.getId()).thenReturn(4L);
 
     when(existsAccount.getUuid()).thenReturn(EXISTS_UUID);
+    when(existsAccount.getPhoneNumberIdentifier()).thenReturn(Optional.of(EXISTS_PNI));
     when(existsAccount.getDevice(1L)).thenReturn(Optional.of(sampleDevice));
     when(existsAccount.getDevice(2L)).thenReturn(Optional.of(sampleDevice2));
     when(existsAccount.getDevice(3L)).thenReturn(Optional.of(sampleDevice3));
@@ -149,11 +162,13 @@ class KeysControllerTest {
     when(existsAccount.getDevices()).thenReturn(allDevices);
     when(existsAccount.isEnabled()).thenReturn(true);
     when(existsAccount.getIdentityKey()).thenReturn("existsidentitykey");
+    when(existsAccount.getPhoneNumberIdentityKey()).thenReturn("existspniidentitykey");
     when(existsAccount.getNumber()).thenReturn(EXISTS_NUMBER);
     when(existsAccount.getUnidentifiedAccessKey()).thenReturn(Optional.of("1337".getBytes()));
 
     when(accounts.getByE164(EXISTS_NUMBER)).thenReturn(Optional.of(existsAccount));
     when(accounts.getByAccountIdentifier(EXISTS_UUID)).thenReturn(Optional.of(existsAccount));
+    when(accounts.getByPhoneNumberIdentifier(EXISTS_PNI)).thenReturn(Optional.of(existsAccount));
 
     when(accounts.getByE164(NOT_EXISTS_NUMBER)).thenReturn(Optional.empty());
     when(accounts.getByAccountIdentifier(NOT_EXISTS_UUID)).thenReturn(Optional.empty());
@@ -161,10 +176,12 @@ class KeysControllerTest {
     when(rateLimiters.getPreKeysLimiter()).thenReturn(rateLimiter);
 
     when(KEYS.take(EXISTS_UUID, 1)).thenReturn(Optional.of(SAMPLE_KEY));
+    when(KEYS.take(EXISTS_PNI, 1)).thenReturn(Optional.of(SAMPLE_KEY_PNI));
 
     when(KEYS.getCount(AuthHelper.VALID_UUID, 1)).thenReturn(5);
 
     when(AuthHelper.VALID_DEVICE.getSignedPreKey()).thenReturn(VALID_DEVICE_SIGNED_KEY);
+    when(AuthHelper.VALID_DEVICE.getPhoneNumberIdentitySignedPreKey()).thenReturn(VALID_DEVICE_PNI_SIGNED_KEY);
     when(AuthHelper.VALID_ACCOUNT.getIdentityKey()).thenReturn(null);
   }
 
@@ -179,6 +196,8 @@ class KeysControllerTest {
         rateLimiter,
         rateLimitChallengeManager
     );
+
+    clearInvocations(AuthHelper.VALID_DEVICE);
   }
 
   @Test
@@ -209,6 +228,20 @@ class KeysControllerTest {
     assertThat(result.getPublicKey()).isEqualTo(VALID_DEVICE_SIGNED_KEY.getPublicKey());
   }
 
+  @Test
+  void getPhoneNumberIdentifierSignedPreKeyV2() {
+    SignedPreKey result = resources.getJerseyTest()
+        .target("/v2/keys/signed")
+        .queryParam("identity", "pni")
+        .request()
+        .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+        .get(SignedPreKey.class);
+
+    assertThat(result.getSignature()).isEqualTo(VALID_DEVICE_PNI_SIGNED_KEY.getSignature());
+    assertThat(result.getKeyId()).isEqualTo(VALID_DEVICE_PNI_SIGNED_KEY.getKeyId());
+    assertThat(result.getPublicKey()).isEqualTo(VALID_DEVICE_PNI_SIGNED_KEY.getPublicKey());
+  }
+
   @Test
   void putSignedPreKeyV2() {
     SignedPreKey   test     = new SignedPreKey(9998, "fooozzz", "baaarzzz");
@@ -221,9 +254,27 @@ class KeysControllerTest {
     assertThat(response.getStatus()).isEqualTo(204);
 
     verify(AuthHelper.VALID_DEVICE).setSignedPreKey(eq(test));
+    verify(AuthHelper.VALID_DEVICE, never()).setPhoneNumberIdentitySignedPreKey(any());
     verify(accounts).updateDevice(eq(AuthHelper.VALID_ACCOUNT), anyLong(), any());
   }
 
+  @Test
+  void putPhoneNumberIdentitySignedPreKeyV2() {
+    final SignedPreKey replacementKey = new SignedPreKey(9998, "fooozzz", "baaarzzz");
+
+    Response response = resources.getJerseyTest()
+        .target("/v2/keys/signed")
+        .queryParam("identity", "pni")
+        .request()
+        .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+        .put(Entity.entity(replacementKey, MediaType.APPLICATION_JSON_TYPE));
+
+    assertThat(response.getStatus()).isEqualTo(204);
+
+    verify(AuthHelper.VALID_DEVICE).setPhoneNumberIdentitySignedPreKey(eq(replacementKey));
+    verify(AuthHelper.VALID_DEVICE, never()).setSignedPreKey(any());
+    verify(accounts).updateDevice(eq(AuthHelper.VALID_ACCOUNT), anyLong(), any());
+  }
 
   @Test
   void disabledPutSignedPreKeyV2() {
@@ -255,6 +306,24 @@ class KeysControllerTest {
     verifyNoMoreInteractions(KEYS);
   }
 
+  @Test
+  void validSingleRequestByPhoneNumberIdentifierTestV2() {
+    PreKeyResponse result = resources.getJerseyTest()
+        .target(String.format("/v2/keys/%s/1", EXISTS_PNI))
+        .request()
+        .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+        .get(PreKeyResponse.class);
+
+    assertThat(result.getIdentityKey()).isEqualTo(existsAccount.getPhoneNumberIdentityKey());
+    assertThat(result.getDevicesCount()).isEqualTo(1);
+    assertThat(result.getDevice(1).getPreKey().getKeyId()).isEqualTo(SAMPLE_KEY_PNI.getKeyId());
+    assertThat(result.getDevice(1).getPreKey().getPublicKey()).isEqualTo(SAMPLE_KEY_PNI.getPublicKey());
+    assertThat(result.getDevice(1).getSignedPreKey()).isEqualTo(existsAccount.getDevice(1).get().getPhoneNumberIdentitySignedPreKey());
+
+    verify(KEYS).take(EXISTS_PNI, 1);
+    verifyNoMoreInteractions(KEYS);
+  }
+
   @Test
   void testUnidentifiedRequest() {
     PreKeyResponse result = resources.getJerseyTest()
@@ -448,6 +517,38 @@ class KeysControllerTest {
     verify(accounts).update(eq(AuthHelper.VALID_ACCOUNT), any());
   }
 
+  @Test
+  void putKeysByPhoneNumberIdentifierTestV2() {
+    final SignedPreKey signedPreKey = new SignedPreKey(31338, "foobaz", "myvalidsig");
+    final String       identityKey  = "barbar";
+
+    List<PreKey> preKeys = List.of(new PreKey(31337, "foobar"));
+
+    PreKeyState preKeyState = new PreKeyState(identityKey, signedPreKey, preKeys);
+
+    Response response =
+        resources.getJerseyTest()
+            .target("/v2/keys")
+            .queryParam("identity", "pni")
+            .request()
+            .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+            .put(Entity.entity(preKeyState, MediaType.APPLICATION_JSON_TYPE));
+
+    assertThat(response.getStatus()).isEqualTo(204);
+
+    ArgumentCaptor<List<PreKey>> listCaptor = ArgumentCaptor.forClass(List.class);
+    verify(KEYS).store(eq(AuthHelper.VALID_PNI), eq(1L), listCaptor.capture());
+
+    List<PreKey> capturedList = listCaptor.getValue();
+    assertThat(capturedList.size()).isEqualTo(1);
+    assertThat(capturedList.get(0).getKeyId()).isEqualTo(31337);
+    assertThat(capturedList.get(0).getPublicKey()).isEqualTo("foobar");
+
+    verify(AuthHelper.VALID_ACCOUNT).setPhoneNumberIdentityKey(eq("barbar"));
+    verify(AuthHelper.VALID_DEVICE).setPhoneNumberIdentitySignedPreKey(eq(signedPreKey));
+    verify(accounts).update(eq(AuthHelper.VALID_ACCOUNT), any());
+  }
+
   @Test
   void disabledPutKeysTestV2() {
     final PreKey       preKey       = new PreKey(31337, "foobar");

service/src/test/java/org/whispersystems/textsecuregcm/tests/util/AccountsHelper.java

@@ -95,6 +95,10 @@ public class AccountsHelper {
             when(updatedAccount.getUuid()).thenAnswer(stubbing);
             break;
           }
+          case "getPhoneNumberIdentifier": {
+            when(updatedAccount.getPhoneNumberIdentifier()).thenAnswer(stubbing);
+            break;
+          }
           case "getNumber": {
             when(updatedAccount.getNumber()).thenAnswer(stubbing);
             break;

service/src/test/java/org/whispersystems/textsecuregcm/tests/util/AuthHelper.java

@@ -37,6 +37,7 @@ public class AuthHelper {
 
   public static final String VALID_NUMBER   = "+14150000000";
   public static final UUID   VALID_UUID     = UUID.randomUUID();
+  public static final UUID   VALID_PNI      = UUID.randomUUID();
   public static final String VALID_PASSWORD = "foo";
 
   public static final String VALID_NUMBER_TWO = "+201511111110";
@@ -112,6 +113,7 @@ public class AuthHelper {
 
     when(VALID_ACCOUNT.getNumber()).thenReturn(VALID_NUMBER);
     when(VALID_ACCOUNT.getUuid()).thenReturn(VALID_UUID);
+    when(VALID_ACCOUNT.getPhoneNumberIdentifier()).thenReturn(Optional.of(VALID_PNI));
     when(VALID_ACCOUNT_TWO.getNumber()).thenReturn(VALID_NUMBER_TWO);
     when(VALID_ACCOUNT_TWO.getUuid()).thenReturn(VALID_UUID_TWO);
     when(DISABLED_ACCOUNT.getNumber()).thenReturn(DISABLED_NUMBER);
@@ -139,6 +141,7 @@ public class AuthHelper {
 
     when(ACCOUNTS_MANAGER.getByE164(VALID_NUMBER)).thenReturn(Optional.of(VALID_ACCOUNT));
     when(ACCOUNTS_MANAGER.getByAccountIdentifier(VALID_UUID)).thenReturn(Optional.of(VALID_ACCOUNT));
+    when(ACCOUNTS_MANAGER.getByPhoneNumberIdentifier(VALID_PNI)).thenReturn(Optional.of(VALID_ACCOUNT));
 
     when(ACCOUNTS_MANAGER.getByE164(VALID_NUMBER_TWO)).thenReturn(Optional.of(VALID_ACCOUNT_TWO));
     when(ACCOUNTS_MANAGER.getByAccountIdentifier(VALID_UUID_TWO)).thenReturn(Optional.of(VALID_ACCOUNT_TWO));