Allow callers to set/retrieve keys by ACI or PNI

This commit is contained in:
Jon Chambers 2021-11-03 18:51:03 -04:00 committed by Jon Chambers
parent 3a4c5a2bfb
commit f89a20dbc7
6 changed files with 209 additions and 29 deletions

View File

@ -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) {
int count = keys.getCount(auth.getAccount().getUuid(), auth.getAuthenticatedDevice().getId());
public PreKeyCount getStatus(@Auth final AuthenticatedAccount auth,
@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.setIdentityKey(preKeys.getIdentityKey());
a.getDevice(device.getId()).ifPresent(d -> {
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);
OptionalAccess.verify(account, accessKey, target, deviceId);
final Account target;
{
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) {
Device device = auth.getAuthenticatedDevice();
SignedPreKey signedPreKey = device.getSignedPreKey();
public Optional<SignedPreKey> getSignedKey(@Auth final AuthenticatedAccount auth,
@QueryParam("identity") final Optional<String> identityType) {
if (signedPreKey != null) {
return Optional.of(signedPreKey);
} else {
return Optional.empty();
}
Device device = auth.getAuthenticatedDevice();
SignedPreKey signedPreKey = usePhoneNumberIdentity(identityType) ?
device.getPhoneNumberIdentitySignedPreKey() : device.getSignedPreKey();
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) {

View File

@ -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());
}
}
}

View File

@ -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;
}

View File

@ -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");

View File

@ -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;

View File

@ -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));