diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java index 8deae41da..ba2397233 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java @@ -6,7 +6,6 @@ package org.whispersystems.textsecuregcm.controllers; import com.google.common.net.HttpHeaders; import io.dropwizard.auth.Auth; -import io.micrometer.core.instrument.DistributionSummary; import io.micrometer.core.instrument.Metrics; import io.micrometer.core.instrument.Tag; import io.micrometer.core.instrument.Tags; @@ -371,9 +370,8 @@ public class KeysController { final Account target = maybeTarget.orElseThrow(NotFoundException::new); if (account.isPresent()) { - rateLimiters.getPreKeysLimiter().validate( - account.get().getUuid() + "." + maybeAuthenticatedDevice.get().deviceId() + "__" + targetIdentifier.uuid() - + "." + deviceId); + rateLimiters.getPreKeysLimiter().validate(getPreKeysLimiterKey(account.get(), maybeAuthenticatedDevice.get(), + targetIdentifier, target, deviceId)); } final List devices = parseDeviceId(deviceId, target); @@ -409,4 +407,24 @@ public class KeysController { throw new WebApplicationException(Response.status(422).build()); } } + + private String getPreKeysLimiterKey( + final Account account, + final AuthenticatedDevice authenticatedDevice, + final ServiceIdentifier targetIdentifier, + final Account targetAccount, + final String targetDeviceId) { + final String targetRegistrationId = targetDeviceId.equals("*") + ? "*" + : String.valueOf( + parseDeviceId(targetDeviceId, targetAccount).getFirst().getRegistrationId(targetIdentifier.identityType())); + + return String.format("%s.%s__%s.%s.%s", + account.getUuid(), + authenticatedDevice.deviceId(), + targetIdentifier.uuid(), + targetDeviceId, + targetRegistrationId + ); + } } diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/KeysControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/KeysControllerTest.java index e2ad61184..c4406fabe 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/KeysControllerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/KeysControllerTest.java @@ -371,8 +371,16 @@ class KeysControllerTest { .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD)) .get(); + final String expectedRatelimitKey = String.format("%s.%s__%s.%s.%s", + AuthHelper.VALID_UUID, + AuthHelper.VALID_DEVICE.getId(), + EXISTS_PNI, + "*", + "*"); + assertThat(result.getStatus()).isEqualTo(429); assertThat(result.getHeaderString("Retry-After")).isEqualTo(String.valueOf(retryAfter.toSeconds())); + verify(rateLimiter).validate(expectedRatelimitKey); } @Test