Validate device message sizes when distributing PNI keys

This commit is contained in:
Jon Chambers 2025-03-04 11:33:06 -05:00 committed by Jon Chambers
parent 1346fcb59e
commit df56c65b54
4 changed files with 101 additions and 29 deletions

View File

@ -51,6 +51,7 @@ import org.whispersystems.textsecuregcm.entities.RegistrationLockFailure;
import org.whispersystems.textsecuregcm.entities.StaleDevices;
import org.whispersystems.textsecuregcm.limits.RateLimiters;
import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
import org.whispersystems.textsecuregcm.push.MessageTooLargeException;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.ChangeNumberManager;
@ -94,6 +95,7 @@ public class AccountControllerV2 {
@ApiResponse(responseCode = "403", description = "Verification failed for the provided Registration Recovery Password")
@ApiResponse(responseCode = "409", description = "Mismatched number of devices or device ids in 'devices to notify' list", content = @Content(schema = @Schema(implementation = MismatchedDevices.class)))
@ApiResponse(responseCode = "410", description = "Mismatched registration ids in 'devices to notify' list", content = @Content(schema = @Schema(implementation = StaleDevices.class)))
@ApiResponse(responseCode = "413", description = "One or more device messages was too large")
@ApiResponse(responseCode = "422", description = "The request did not pass validation")
@ApiResponse(responseCode = "423", content = @Content(schema = @Schema(implementation = RegistrationLockFailure.class)))
@ApiResponse(responseCode = "429", description = "Too many attempts", headers = @Header(
@ -143,7 +145,8 @@ public class AccountControllerV2 {
request.devicePniSignedPrekeys(),
request.devicePniPqLastResortPrekeys(),
request.deviceMessages(),
request.pniRegistrationIds());
request.pniRegistrationIds(),
userAgentString);
return AccountIdentityResponseBuilder.fromAccount(updatedAccount);
} catch (MismatchedDevicesException e) {
@ -159,6 +162,8 @@ public class AccountControllerV2 {
.build());
} catch (IllegalArgumentException e) {
throw new BadRequestException(e);
} catch (MessageTooLargeException e) {
throw new WebApplicationException(Response.Status.REQUEST_ENTITY_TOO_LARGE);
}
}
@ -176,6 +181,7 @@ public class AccountControllerV2 {
content = @Content(schema = @Schema(implementation = MismatchedDevices.class)))
@ApiResponse(responseCode = "410", description = "The registration IDs provided for some devices do not match those stored on the server.",
content = @Content(schema = @Schema(implementation = StaleDevices.class)))
@ApiResponse(responseCode = "413", description = "One or more device messages was too large")
public AccountIdentityResponse distributePhoneNumberIdentityKeys(
@Mutable @Auth final AuthenticatedDevice authenticatedDevice,
@HeaderParam(HttpHeaders.USER_AGENT) @Nullable final String userAgentString,
@ -196,7 +202,8 @@ public class AccountControllerV2 {
request.devicePniSignedPrekeys(),
request.devicePniPqLastResortPrekeys(),
request.deviceMessages(),
request.pniRegistrationIds());
request.pniRegistrationIds(),
userAgentString);
return AccountIdentityResponseBuilder.fromAccount(updatedAccount);
} catch (MismatchedDevicesException e) {
@ -212,6 +219,8 @@ public class AccountControllerV2 {
.build());
} catch (IllegalArgumentException e) {
throw new BadRequestException(e);
} catch (MessageTooLargeException e) {
throw new WebApplicationException(Response.Status.REQUEST_ENTITY_TOO_LARGE);
}
}

View File

@ -22,6 +22,7 @@ import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
import org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope;
import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
import org.whispersystems.textsecuregcm.push.MessageSender;
import org.whispersystems.textsecuregcm.push.MessageTooLargeException;
import org.whispersystems.textsecuregcm.util.DestinationDeviceValidator;
public class ChangeNumberManager {
@ -42,8 +43,9 @@ public class ChangeNumberManager {
@Nullable final Map<Byte, ECSignedPreKey> deviceSignedPreKeys,
@Nullable final Map<Byte, KEMSignedPreKey> devicePqLastResortPreKeys,
@Nullable final List<IncomingMessage> deviceMessages,
@Nullable final Map<Byte, Integer> pniRegistrationIds)
throws InterruptedException, MismatchedDevicesException, StaleDevicesException {
@Nullable final Map<Byte, Integer> pniRegistrationIds,
@Nullable final String senderUserAgent)
throws InterruptedException, MismatchedDevicesException, StaleDevicesException, MessageTooLargeException {
if (ObjectUtils.allNotNull(pniIdentityKey, deviceSignedPreKeys, deviceMessages, pniRegistrationIds)) {
// AccountsManager validates the device set on deviceSignedPreKeys and pniRegistrationIds
@ -63,14 +65,14 @@ public class ChangeNumberManager {
if (pniIdentityKey == null) {
return account;
}
return updatePniKeys(account, pniIdentityKey, deviceSignedPreKeys, devicePqLastResortPreKeys, deviceMessages, pniRegistrationIds);
return updatePniKeys(account, pniIdentityKey, deviceSignedPreKeys, devicePqLastResortPreKeys, deviceMessages, pniRegistrationIds, senderUserAgent);
}
final Account updatedAccount = accountsManager.changeNumber(
account, number, pniIdentityKey, deviceSignedPreKeys, devicePqLastResortPreKeys, pniRegistrationIds);
if (deviceMessages != null) {
sendDeviceMessages(updatedAccount, deviceMessages);
sendDeviceMessages(updatedAccount, deviceMessages, senderUserAgent);
}
return updatedAccount;
@ -81,7 +83,9 @@ public class ChangeNumberManager {
final Map<Byte, ECSignedPreKey> deviceSignedPreKeys,
@Nullable final Map<Byte, KEMSignedPreKey> devicePqLastResortPreKeys,
final List<IncomingMessage> deviceMessages,
final Map<Byte, Integer> pniRegistrationIds) throws MismatchedDevicesException, StaleDevicesException {
final Map<Byte, Integer> pniRegistrationIds,
final String senderUserAgent) throws MismatchedDevicesException, StaleDevicesException, MessageTooLargeException {
validateDeviceMessages(account, deviceMessages);
// Don't try to be smart about ignoring unnecessary retries. If we make literally no change we will skip the ddb
@ -89,7 +93,7 @@ public class ChangeNumberManager {
final Account updatedAccount = accountsManager.updatePniKeys(
account, pniIdentityKey, deviceSignedPreKeys, devicePqLastResortPreKeys, pniRegistrationIds);
sendDeviceMessages(updatedAccount, deviceMessages);
sendDeviceMessages(updatedAccount, deviceMessages, senderUserAgent);
return updatedAccount;
}
@ -110,7 +114,18 @@ public class ChangeNumberManager {
false);
}
private void sendDeviceMessages(final Account account, final List<IncomingMessage> deviceMessages) {
private void sendDeviceMessages(final Account account,
final List<IncomingMessage> deviceMessages,
final String senderUserAgent) throws MessageTooLargeException {
for (final IncomingMessage message : deviceMessages) {
MessageSender.validateContentLength(message.content().length,
false,
true,
false,
senderUserAgent);
}
try {
final long serverTimestamp = System.currentTimeMillis();

View File

@ -16,6 +16,7 @@ import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doThrow;
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.when;
@ -79,6 +80,7 @@ import org.whispersystems.textsecuregcm.limits.RateLimiters;
import org.whispersystems.textsecuregcm.mappers.ImpossiblePhoneNumberExceptionMapper;
import org.whispersystems.textsecuregcm.mappers.NonNormalizedPhoneNumberExceptionMapper;
import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper;
import org.whispersystems.textsecuregcm.push.MessageTooLargeException;
import org.whispersystems.textsecuregcm.registration.RegistrationServiceClient;
import org.whispersystems.textsecuregcm.spam.RegistrationRecoveryChecker;
import org.whispersystems.textsecuregcm.storage.Account;
@ -139,7 +141,7 @@ class AccountControllerV2Test {
void setUp() throws Exception {
when(rateLimiters.getRegistrationLimiter()).thenReturn(registrationLimiter);
when(changeNumberManager.changeNumber(any(), any(), any(), any(), any(), any(), any())).thenAnswer(
when(changeNumberManager.changeNumber(any(), any(), any(), any(), any(), any(), any(), any())).thenAnswer(
(Answer<Account>) invocation -> {
final Account account = invocation.getArgument(0);
final String number = invocation.getArgument(1);
@ -189,7 +191,7 @@ class AccountControllerV2Test {
MediaType.APPLICATION_JSON_TYPE), AccountIdentityResponse.class);
verify(changeNumberManager).changeNumber(eq(AuthHelper.VALID_ACCOUNT), eq(NEW_NUMBER), any(), any(), any(),
any(), any());
any(), any(), any());
assertEquals(AuthHelper.VALID_UUID, accountIdentityResponse.uuid());
assertEquals(NEW_NUMBER, accountIdentityResponse.number());
@ -212,7 +214,7 @@ class AccountControllerV2Test {
MediaType.APPLICATION_JSON_TYPE), AccountIdentityResponse.class);
verify(changeNumberManager).changeNumber(eq(AuthHelper.VALID_ACCOUNT), eq(AuthHelper.VALID_NUMBER), any(), any(), any(),
any(), any());
any(), any(), any());
assertEquals(AuthHelper.VALID_UUID, accountIdentityResponse.uuid());
assertEquals(AuthHelper.VALID_NUMBER, accountIdentityResponse.number());
@ -220,7 +222,7 @@ class AccountControllerV2Test {
}
@Test
void changeNumberNonNormalizedNumber() throws Exception {
void changeNumberNonNormalizedNumber() {
try (Response response = resources.getJerseyTest()
.target("/v2/accounts/number")
.request()
@ -293,13 +295,15 @@ class AccountControllerV2Test {
null, NEW_NUMBER, "123", new IdentityKey(Curve.generateKeyPair().getPublicKey()),
Collections.emptyList(), Collections.emptyMap(), null, Map.of((byte) 1, pniRegistrationId));
final Response response = resources.getJerseyTest()
try (final Response response = resources.getJerseyTest()
.target("/v2/accounts/number")
.request()
.header(HttpHeaders.AUTHORIZATION,
AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
.put(Entity.entity(changeNumberRequest, MediaType.APPLICATION_JSON_TYPE));
assertEquals(expectedStatusCode, response.getStatus());
.put(Entity.entity(changeNumberRequest, MediaType.APPLICATION_JSON_TYPE))) {
assertEquals(expectedStatusCode, response.getStatus());
}
}
private static Stream<Arguments> invalidRegistrationId() {
@ -422,7 +426,7 @@ class AccountControllerV2Test {
final AccountIdentityResponse accountIdentityResponse = response.readEntity(AccountIdentityResponse.class);
verify(changeNumberManager).changeNumber(eq(AuthHelper.VALID_ACCOUNT), eq(NEW_NUMBER), any(), any(), any(),
any(), any());
any(), any(), any());
assertEquals(AuthHelper.VALID_UUID, accountIdentityResponse.uuid());
assertEquals(NEW_NUMBER, accountIdentityResponse.number());
@ -481,6 +485,33 @@ class AccountControllerV2Test {
}
}
@Test
void deviceMessageTooLarge() throws Exception {
when(registrationServiceClient.getSession(any(), any()))
.thenReturn(CompletableFuture.completedFuture(
Optional.of(new RegistrationServiceSession(new byte[16], NEW_NUMBER, true, null, null, null,
SESSION_EXPIRATION_SECONDS))));
reset(changeNumberManager);
when(changeNumberManager.changeNumber(any(), any(), any(), any(), any(), any(), any(), any()))
.thenThrow(MessageTooLargeException.class);
try (final Response response = resources.getJerseyTest()
.target("/v2/accounts/number")
.request()
.header(HttpHeaders.AUTHORIZATION,
AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
.put(Entity.entity(
new ChangeNumberRequest(encodeSessionId("session"), null, NEW_NUMBER, "123", new IdentityKey(Curve.generateKeyPair().getPublicKey()),
Collections.emptyList(),
Collections.emptyMap(), null, Collections.emptyMap()),
MediaType.APPLICATION_JSON_TYPE))) {
assertEquals(413, response.getStatus());
}
}
/**
* Valid request JSON with the given Recovery Password
*/
@ -558,7 +589,7 @@ class AccountControllerV2Test {
@BeforeEach
void setUp() throws Exception {
when(changeNumberManager.updatePniKeys(any(), any(), any(), any(), any(), any())).thenAnswer(
when(changeNumberManager.updatePniKeys(any(), any(), any(), any(), any(), any(), any())).thenAnswer(
(Answer<Account>) invocation -> {
final Account account = invocation.getArgument(0);
final IdentityKey pniIdentityKey = invocation.getArgument(1);
@ -594,7 +625,7 @@ class AccountControllerV2Test {
AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
.put(Entity.json(requestJson()), AccountIdentityResponse.class);
verify(changeNumberManager).updatePniKeys(eq(AuthHelper.VALID_ACCOUNT), eq(IDENTITY_KEY), any(), any(), any(), any());
verify(changeNumberManager).updatePniKeys(eq(AuthHelper.VALID_ACCOUNT), eq(IDENTITY_KEY), any(), any(), any(), any(), any());
assertEquals(AuthHelper.VALID_UUID, accountIdentityResponse.uuid());
assertEquals(AuthHelper.VALID_NUMBER, accountIdentityResponse.number());
@ -646,6 +677,23 @@ class AccountControllerV2Test {
}
}
@Test
void deviceMessageTooLarge() throws Exception {
reset(changeNumberManager);
when(changeNumberManager.updatePniKeys(any(), any(), any(), any(), any(), any(), any()))
.thenThrow(MessageTooLargeException.class);
try (final Response response = resources.getJerseyTest()
.target("/v2/accounts/phone_number_identity_key_distribution")
.request()
.header(HttpHeaders.AUTHORIZATION,
AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
.put(Entity.json(requestJson()))) {
assertEquals(413, response.getStatus());
}
}
/**
* Valid request JSON for a {@link org.whispersystems.textsecuregcm.entities.PhoneNumberIdentityKeyDistributionRequest}
*/

View File

@ -102,7 +102,7 @@ public class ChangeNumberManagerTest {
void changeNumberNoMessages() throws Exception {
Account account = mock(Account.class);
when(account.getNumber()).thenReturn("+18005551234");
changeNumberManager.changeNumber(account, "+18025551234", null, null, null, null, null);
changeNumberManager.changeNumber(account, "+18025551234", null, null, null, null, null, null);
verify(accountsManager).changeNumber(account, "+18025551234", null, null, null, null);
verify(accountsManager, never()).updateDevice(any(), anyByte(), any());
verify(messageSender, never()).sendMessages(eq(account), any());
@ -117,7 +117,7 @@ public class ChangeNumberManagerTest {
final Map<Byte, ECSignedPreKey> prekeys = Map.of(Device.PRIMARY_ID,
KeysHelper.signedECPreKey(1, pniIdentityKeyPair));
changeNumberManager.changeNumber(account, "+18025551234", pniIdentityKey, prekeys, null, Collections.emptyList(), Collections.emptyMap());
changeNumberManager.changeNumber(account, "+18025551234", pniIdentityKey, prekeys, null, Collections.emptyList(), Collections.emptyMap(), null);
verify(accountsManager).changeNumber(account, "+18025551234", pniIdentityKey, prekeys, null, Collections.emptyMap());
verify(messageSender, never()).sendMessages(eq(account), any());
}
@ -152,7 +152,7 @@ public class ChangeNumberManagerTest {
when(msg.destinationDeviceId()).thenReturn(deviceId2);
when(msg.content()).thenReturn(new byte[]{1});
changeNumberManager.changeNumber(account, changedE164, pniIdentityKey, prekeys, null, List.of(msg), registrationIds);
changeNumberManager.changeNumber(account, changedE164, pniIdentityKey, prekeys, null, List.of(msg), registrationIds, null);
verify(accountsManager).changeNumber(account, changedE164, pniIdentityKey, prekeys, null, registrationIds);
@ -205,7 +205,7 @@ public class ChangeNumberManagerTest {
when(msg.destinationDeviceId()).thenReturn(deviceId2);
when(msg.content()).thenReturn(new byte[]{1});
changeNumberManager.changeNumber(account, changedE164, pniIdentityKey, prekeys, pqPrekeys, List.of(msg), registrationIds);
changeNumberManager.changeNumber(account, changedE164, pniIdentityKey, prekeys, pqPrekeys, List.of(msg), registrationIds, null);
verify(accountsManager).changeNumber(account, changedE164, pniIdentityKey, prekeys, pqPrekeys, registrationIds);
@ -256,7 +256,7 @@ public class ChangeNumberManagerTest {
when(msg.destinationDeviceId()).thenReturn(deviceId2);
when(msg.content()).thenReturn(new byte[]{1});
changeNumberManager.changeNumber(account, originalE164, pniIdentityKey, prekeys, pqPrekeys, List.of(msg), registrationIds);
changeNumberManager.changeNumber(account, originalE164, pniIdentityKey, prekeys, pqPrekeys, List.of(msg), registrationIds, null);
verify(accountsManager).updatePniKeys(account, pniIdentityKey, prekeys, pqPrekeys, registrationIds);
@ -303,7 +303,7 @@ public class ChangeNumberManagerTest {
when(msg.destinationDeviceId()).thenReturn(deviceId2);
when(msg.content()).thenReturn(new byte[]{1});
changeNumberManager.updatePniKeys(account, pniIdentityKey, prekeys, null, List.of(msg), registrationIds);
changeNumberManager.updatePniKeys(account, pniIdentityKey, prekeys, null, List.of(msg), registrationIds, null);
verify(accountsManager).updatePniKeys(account, pniIdentityKey, prekeys, null, registrationIds);
@ -352,7 +352,7 @@ public class ChangeNumberManagerTest {
when(msg.destinationDeviceId()).thenReturn(deviceId2);
when(msg.content()).thenReturn(new byte[]{1});
changeNumberManager.updatePniKeys(account, pniIdentityKey, prekeys, pqPrekeys, List.of(msg), registrationIds);
changeNumberManager.updatePniKeys(account, pniIdentityKey, prekeys, pqPrekeys, List.of(msg), registrationIds, null);
verify(accountsManager).updatePniKeys(account, pniIdentityKey, prekeys, pqPrekeys, registrationIds);
@ -407,7 +407,7 @@ public class ChangeNumberManagerTest {
destinationDeviceId3, 89);
assertThrows(StaleDevicesException.class,
() -> changeNumberManager.changeNumber(account, "+18005559876", new IdentityKey(Curve.generateKeyPair().getPublicKey()), preKeys, null, messages, registrationIds));
() -> changeNumberManager.changeNumber(account, "+18005559876", new IdentityKey(Curve.generateKeyPair().getPublicKey()), preKeys, null, messages, registrationIds, null));
}
@Test
@ -445,7 +445,7 @@ public class ChangeNumberManagerTest {
destinationDeviceId3, 89);
assertThrows(StaleDevicesException.class,
() -> changeNumberManager.updatePniKeys(account, new IdentityKey(Curve.generateKeyPair().getPublicKey()), preKeys, null, messages, registrationIds));
() -> changeNumberManager.updatePniKeys(account, new IdentityKey(Curve.generateKeyPair().getPublicKey()), preKeys, null, messages, registrationIds, null));
}
@Test
@ -475,6 +475,6 @@ public class ChangeNumberManagerTest {
final Map<Byte, Integer> registrationIds = Map.of((byte) 1, 17, destinationDeviceId2, 47, destinationDeviceId3, 89);
assertThrows(IllegalArgumentException.class,
() -> changeNumberManager.changeNumber(account, "+18005559876", new IdentityKey(Curve.generateKeyPair().getPublicKey()), null, null, messages, registrationIds));
() -> changeNumberManager.changeNumber(account, "+18005559876", new IdentityKey(Curve.generateKeyPair().getPublicKey()), null, null, messages, registrationIds, null));
}
}