Migrate to `429` for all ratelimit responses

This commit is contained in:
Katherine 2024-08-05 12:02:11 -07:00 committed by GitHub
parent 10d559bbb5
commit 0e4625ef88
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
33 changed files with 63 additions and 110 deletions

View File

@ -102,9 +102,8 @@ public class BackupAuthManager {
return CompletableFuture.completedFuture(null); return CompletableFuture.completedFuture(null);
} }
return RateLimiter.adaptLegacyException(rateLimiters return rateLimiters.forDescriptor(RateLimiters.For.SET_BACKUP_ID)
.forDescriptor(RateLimiters.For.SET_BACKUP_ID) .validateAsync(account.getUuid())
.validateAsync(account.getUuid()))
.thenCompose(ignored -> this.accountsManager .thenCompose(ignored -> this.accountsManager
.updateAsync(account, acc -> acc.setBackupCredentialRequest(serializedRequest)) .updateAsync(account, acc -> acc.setBackupCredentialRequest(serializedRequest))
.thenRun(Util.NOOP)) .thenRun(Util.NOOP))

View File

@ -156,9 +156,8 @@ public class BackupManager {
final AuthenticatedBackupUser backupUser) { final AuthenticatedBackupUser backupUser) {
checkBackupLevel(backupUser, BackupLevel.MEDIA); checkBackupLevel(backupUser, BackupLevel.MEDIA);
return RateLimiter.adaptLegacyException(rateLimiters return rateLimiters.forDescriptor(RateLimiters.For.BACKUP_ATTACHMENT)
.forDescriptor(RateLimiters.For.BACKUP_ATTACHMENT) .validateAsync(rateLimitKey(backupUser)).thenApply(ignored -> {
.validateAsync(rateLimitKey(backupUser))).thenApply(ignored -> {
final byte[] bytes = new byte[15]; final byte[] bytes = new byte[15];
secureRandom.nextBytes(bytes); secureRandom.nextBytes(bytes);
final String attachmentKey = Base64.getUrlEncoder().encodeToString(bytes); final String attachmentKey = Base64.getUrlEncoder().encodeToString(bytes);

View File

@ -134,7 +134,7 @@ public class AccountControllerV2 {
// Only verify and check reglock if there's a data change to be made... // Only verify and check reglock if there's a data change to be made...
if (!authenticatedAccount.getAccount().getNumber().equals(number)) { if (!authenticatedAccount.getAccount().getNumber().equals(number)) {
RateLimiter.adaptLegacyException(() -> rateLimiters.getRegistrationLimiter().validate(number)); rateLimiters.getRegistrationLimiter().validate(number);
final PhoneVerificationRequest.VerificationType verificationType = phoneVerificationTokenManager.verify(number, final PhoneVerificationRequest.VerificationType verificationType = phoneVerificationTokenManager.verify(number,
request); request);

View File

@ -28,28 +28,21 @@ public class RateLimitExceededException extends Exception implements Convertible
@Nullable @Nullable
private final Duration retryDuration; private final Duration retryDuration;
private final boolean legacy;
/** /**
* Constructs a new exception indicating when it may become safe to retry * Constructs a new exception indicating when it may become safe to retry
* *
* @param retryDuration A duration to wait before retrying, null if no duration can be indicated * @param retryDuration A duration to wait before retrying, null if no duration can be indicated
* @param legacy whether to use a legacy status code when mapping the exception to an HTTP response
*/ */
public RateLimitExceededException(@Nullable final Duration retryDuration, final boolean legacy) { public RateLimitExceededException(@Nullable final Duration retryDuration) {
super(null, null, true, false); super(null, null, true, false);
this.retryDuration = retryDuration; this.retryDuration = retryDuration;
this.legacy = legacy;
} }
public Optional<Duration> getRetryDuration() { public Optional<Duration> getRetryDuration() {
return Optional.ofNullable(retryDuration); return Optional.ofNullable(retryDuration);
} }
public boolean isLegacy() {
return legacy;
}
@Override @Override
public Status grpcStatus() { public Status grpcStatus() {
return Status.RESOURCE_EXHAUSTED; return Status.RESOURCE_EXHAUSTED;

View File

@ -111,7 +111,7 @@ public class RegistrationController {
throw new WebApplicationException("Invalid signature", 422); throw new WebApplicationException("Invalid signature", 422);
} }
RateLimiter.adaptLegacyException(() -> rateLimiters.getRegistrationLimiter().validate(number)); rateLimiters.getRegistrationLimiter().validate(number);
final PhoneVerificationRequest.VerificationType verificationType = phoneVerificationTokenManager.verify(number, final PhoneVerificationRequest.VerificationType verificationType = phoneVerificationTokenManager.verify(number,
registrationRequest); registrationRequest);

View File

@ -173,9 +173,7 @@ public class VerificationController {
} catch (final CompletionException e) { } catch (final CompletionException e) {
if (ExceptionUtils.unwrap(e) instanceof RateLimitExceededException re) { if (ExceptionUtils.unwrap(e) instanceof RateLimitExceededException re) {
RateLimiter.adaptLegacyException(() -> { throw re;
throw re;
});
} }
throw new ServerErrorException(Response.Status.INTERNAL_SERVER_ERROR, e); throw new ServerErrorException(Response.Status.INTERNAL_SERVER_ERROR, e);
@ -318,9 +316,8 @@ public class VerificationController {
final boolean pushChallengePresent = updateVerificationSessionRequest.pushChallenge() != null; final boolean pushChallengePresent = updateVerificationSessionRequest.pushChallenge() != null;
if (pushChallengePresent) { if (pushChallengePresent) {
RateLimiter.adaptLegacyException( rateLimiters.getVerificationPushChallengeLimiter()
() -> rateLimiters.getVerificationPushChallengeLimiter() .validate(registrationServiceSession.encodedSessionId());
.validate(registrationServiceSession.encodedSessionId()));
} }
final boolean pushChallengeMatches; final boolean pushChallengeMatches;
@ -383,8 +380,7 @@ public class VerificationController {
return verificationSession; return verificationSession;
} }
RateLimiter.adaptLegacyException( rateLimiters.getVerificationCaptchaLimiter().validate(registrationServiceSession.encodedSessionId());
() -> rateLimiters.getVerificationCaptchaLimiter().validate(registrationServiceSession.encodedSessionId()));
final AssessmentResult assessmentResult; final AssessmentResult assessmentResult;
try { try {
@ -507,7 +503,7 @@ public class VerificationController {
throw new ClientErrorException(response); throw new ClientErrorException(response);
} }
throw new RateLimitExceededException(rateLimitExceededException.getRetryDuration().orElse(null), false); throw new RateLimitExceededException(rateLimitExceededException.getRetryDuration().orElse(null));
} else if (unwrappedException instanceof RegistrationServiceException registrationServiceException) { } else if (unwrappedException instanceof RegistrationServiceException registrationServiceException) {
throw registrationServiceException.getRegistrationSession() throw registrationServiceException.getRegistrationSession()
@ -584,7 +580,7 @@ public class VerificationController {
throw new ClientErrorException(response); throw new ClientErrorException(response);
} }
throw new RateLimitExceededException(rateLimitExceededException.getRetryDuration().orElse(null), false); throw new RateLimitExceededException(rateLimitExceededException.getRetryDuration().orElse(null));
} else if (unwrappedException instanceof RegistrationServiceException registrationServiceException) { } else if (unwrappedException instanceof RegistrationServiceException registrationServiceException) {

View File

@ -24,7 +24,7 @@ public class VerificationSessionRateLimitExceededException extends RateLimitExce
public VerificationSessionRateLimitExceededException( public VerificationSessionRateLimitExceededException(
final RegistrationServiceSession registrationServiceSession, @Nullable final Duration retryDuration, final RegistrationServiceSession registrationServiceSession, @Nullable final Duration retryDuration,
final boolean legacy) { final boolean legacy) {
super(retryDuration, legacy); super(retryDuration);
this.registrationServiceSession = registrationServiceSession; this.registrationServiceSession = registrationServiceSession;
} }

View File

@ -28,8 +28,8 @@ public class RateLimitByIpFilter implements ContainerRequestFilter {
private static final Logger logger = LoggerFactory.getLogger(RateLimitByIpFilter.class); private static final Logger logger = LoggerFactory.getLogger(RateLimitByIpFilter.class);
@VisibleForTesting @VisibleForTesting
static final RateLimitExceededException INVALID_HEADER_EXCEPTION = new RateLimitExceededException(Duration.ofHours(1), static final RateLimitExceededException INVALID_HEADER_EXCEPTION = new RateLimitExceededException(Duration.ofHours(1)
true); );
private static final ExceptionMapper<RateLimitExceededException> EXCEPTION_MAPPER = new RateLimitExceededExceptionMapper(); private static final ExceptionMapper<RateLimitExceededException> EXCEPTION_MAPPER = new RateLimitExceededExceptionMapper();

View File

@ -78,32 +78,4 @@ public interface RateLimiter {
default CompletionStage<Void> clearAsync(final UUID accountUuid) { default CompletionStage<Void> clearAsync(final UUID accountUuid) {
return clearAsync(accountUuid.toString()); return clearAsync(accountUuid.toString());
} }
/**
* If the future throws a {@link RateLimitExceededException}, it will adapt it to ensure that
* {@link RateLimitExceededException#isLegacy()} returns {@code false}
*/
static CompletionStage<Void> adaptLegacyException(final CompletionStage<Void> rateLimitFuture) {
return rateLimitFuture.exceptionally(ExceptionUtils.exceptionallyHandler(RateLimitExceededException.class, e -> {
throw ExceptionUtils.wrap(new RateLimitExceededException(e.getRetryDuration().orElse(null), false));
}));
}
/**
* If the wrapped {@code validate()} call throws a {@link RateLimitExceededException}, it will adapt it to ensure that
* {@link RateLimitExceededException#isLegacy()} returns {@code false}
*/
static void adaptLegacyException(final RateLimitValidator validator) throws RateLimitExceededException {
try {
validator.validate();
} catch (final RateLimitExceededException e) {
throw new RateLimitExceededException(e.getRetryDuration().orElse(null), false);
}
}
@FunctionalInterface
interface RateLimitValidator {
void validate() throws RateLimitExceededException;
}
} }

View File

@ -65,7 +65,7 @@ public class StaticRateLimiter implements RateLimiter {
counter.increment(); counter.increment();
final Duration retryAfter = Duration.ofMillis( final Duration retryAfter = Duration.ofMillis(
(long) Math.ceil((double) deficitPermitsAmount / config.leakRatePerMillis())); (long) Math.ceil((double) deficitPermitsAmount / config.leakRatePerMillis()));
throw new RateLimitExceededException(retryAfter, true); throw new RateLimitExceededException(retryAfter);
} }
} catch (RedisException e) { } catch (RedisException e) {
if (!failOpen()) { if (!failOpen()) {
@ -84,7 +84,7 @@ public class StaticRateLimiter implements RateLimiter {
counter.increment(); counter.increment();
final Duration retryAfter = Duration.ofMillis( final Duration retryAfter = Duration.ofMillis(
(long) Math.ceil((double) deficitPermitsAmount / config.leakRatePerMillis())); (long) Math.ceil((double) deficitPermitsAmount / config.leakRatePerMillis()));
return failedFuture(new RateLimitExceededException(retryAfter, true)); return failedFuture(new RateLimitExceededException(retryAfter));
}) })
.exceptionally(throwable -> { .exceptionally(throwable -> {
if (ExceptionUtils.unwrap(throwable) instanceof RedisException && failOpen()) { if (ExceptionUtils.unwrap(throwable) instanceof RedisException && failOpen()) {

View File

@ -16,11 +16,8 @@ public class RateLimitExceededExceptionMapper implements ExceptionMapper<RateLim
private static final Logger logger = LoggerFactory.getLogger(RateLimitExceededExceptionMapper.class); private static final Logger logger = LoggerFactory.getLogger(RateLimitExceededExceptionMapper.class);
private static final int LEGACY_STATUS_CODE = 413;
private static final int STATUS_CODE = 429;
/** /**
* Convert a RateLimitExceededException to a {@value STATUS_CODE} (or legacy {@value LEGACY_STATUS_CODE}) response * Convert a RateLimitExceededException to a 429 response
* with a Retry-After header. * with a Retry-After header.
* *
* @param e A RateLimitExceededException potentially containing a recommended retry duration * @param e A RateLimitExceededException potentially containing a recommended retry duration
@ -28,7 +25,6 @@ public class RateLimitExceededExceptionMapper implements ExceptionMapper<RateLim
*/ */
@Override @Override
public Response toResponse(RateLimitExceededException e) { public Response toResponse(RateLimitExceededException e) {
final int statusCode = e.isLegacy() ? LEGACY_STATUS_CODE : STATUS_CODE;
return e.getRetryDuration() return e.getRetryDuration()
.filter(d -> { .filter(d -> {
if (d.isNegative()) { if (d.isNegative()) {
@ -38,7 +34,7 @@ public class RateLimitExceededExceptionMapper implements ExceptionMapper<RateLim
// only include non-negative durations in retry headers // only include non-negative durations in retry headers
return !d.isNegative(); return !d.isNegative();
}) })
.map(d -> Response.status(statusCode).header("Retry-After", d.toSeconds())) .map(d -> Response.status(Response.Status.TOO_MANY_REQUESTS).header("Retry-After", d.toSeconds()))
.orElseGet(() -> Response.status(statusCode)).build(); .orElseGet(() -> Response.status(Response.Status.TOO_MANY_REQUESTS)).build();
} }
} }

View File

@ -97,8 +97,8 @@ public class RegistrationServiceClient implements Managed {
case CREATE_REGISTRATION_SESSION_ERROR_TYPE_RATE_LIMITED -> throw new CompletionException( case CREATE_REGISTRATION_SESSION_ERROR_TYPE_RATE_LIMITED -> throw new CompletionException(
new RateLimitExceededException(response.getError().getMayRetry() new RateLimitExceededException(response.getError().getMayRetry()
? Duration.ofSeconds(response.getError().getRetryAfterSeconds()) ? Duration.ofSeconds(response.getError().getRetryAfterSeconds())
: null, : null
true)); ));
case CREATE_REGISTRATION_SESSION_ERROR_TYPE_ILLEGAL_PHONE_NUMBER -> throw new IllegalArgumentException(); case CREATE_REGISTRATION_SESSION_ERROR_TYPE_ILLEGAL_PHONE_NUMBER -> throw new IllegalArgumentException();
default -> throw new RuntimeException( default -> throw new RuntimeException(
"Unrecognized error type from registration service: " + response.getError().getErrorType()); "Unrecognized error type from registration service: " + response.getError().getErrorType());

View File

@ -7,7 +7,7 @@ package org.whispersystems.textsecuregcm.auth;
public enum RegistrationLockError { public enum RegistrationLockError {
MISMATCH(RegistrationLockVerificationManager.FAILURE_HTTP_STATUS), MISMATCH(RegistrationLockVerificationManager.FAILURE_HTTP_STATUS),
RATE_LIMITED(413) // This will be changed to 429 in a future revision RATE_LIMITED(429)
; ;
private final int expectedStatus; private final int expectedStatus;

View File

@ -409,7 +409,6 @@ public class BackupAuthManagerTest {
final RateLimitExceededException ex = CompletableFutureTestUtil.assertFailsWithCause( final RateLimitExceededException ex = CompletableFutureTestUtil.assertFailsWithCause(
RateLimitExceededException.class, RateLimitExceededException.class,
authManager.commitBackupId(account, credentialRequest)); authManager.commitBackupId(account, credentialRequest));
assertThat(ex.isLegacy()).isFalse();
// If we don't change the request, shouldn't be rate limited // If we don't change the request, shouldn't be rate limited
when(account.getBackupCredentialRequest()).thenReturn(credentialRequest.serialize()); when(account.getBackupCredentialRequest()).thenReturn(credentialRequest.serialize());
@ -436,7 +435,7 @@ public class BackupAuthManagerTest {
final RateLimiters limiters = mock(RateLimiters.class); final RateLimiters limiters = mock(RateLimiters.class);
final RateLimiter limiter = mock(RateLimiter.class); final RateLimiter limiter = mock(RateLimiter.class);
when(limiter.validateAsync(aci)) when(limiter.validateAsync(aci))
.thenReturn(CompletableFuture.failedFuture(new RateLimitExceededException(null, false))); .thenReturn(CompletableFuture.failedFuture(new RateLimitExceededException(null)));
when(limiters.forDescriptor(RateLimiters.For.SET_BACKUP_ID)).thenReturn(limiter); when(limiters.forDescriptor(RateLimiters.For.SET_BACKUP_ID)).thenReturn(limiter);
return limiters; return limiters;
} }

View File

@ -151,11 +151,10 @@ public class BackupManagerTest {
public void createTemporaryMediaAttachmentRateLimited() { public void createTemporaryMediaAttachmentRateLimited() {
final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA); final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupLevel.MEDIA);
when(mediaUploadLimiter.validateAsync(eq(BackupManager.rateLimitKey(backupUser)))) when(mediaUploadLimiter.validateAsync(eq(BackupManager.rateLimitKey(backupUser))))
.thenReturn(CompletableFuture.failedFuture(new RateLimitExceededException(null, true))); .thenReturn(CompletableFuture.failedFuture(new RateLimitExceededException(null)));
final RateLimitExceededException e = CompletableFutureTestUtil.assertFailsWithCause( final RateLimitExceededException e = CompletableFutureTestUtil.assertFailsWithCause(
RateLimitExceededException.class, RateLimitExceededException.class,
backupManager.createTemporaryAttachmentUploadDescriptor(backupUser).toCompletableFuture()); backupManager.createTemporaryAttachmentUploadDescriptor(backupUser).toCompletableFuture());
assertThat(e.isLegacy()).isFalse();
} }
@Test @Test

View File

@ -912,7 +912,7 @@ class AccountControllerTest {
.request() .request()
.head()) { .head()) {
assertThat(response.getStatus()).isEqualTo(413); assertThat(response.getStatus()).isEqualTo(429);
assertThat(response.getHeaderString("Retry-After")).isEqualTo(String.valueOf(expectedRetryAfter.toSeconds())); assertThat(response.getHeaderString("Retry-After")).isEqualTo(String.valueOf(expectedRetryAfter.toSeconds()));
} }
} }
@ -963,7 +963,7 @@ class AccountControllerTest {
.request() .request()
.get(); .get();
assertThat(response.getStatus()).isEqualTo(413); assertThat(response.getStatus()).isEqualTo(429);
assertThat(response.getHeaderString("Retry-After")).isEqualTo(String.valueOf(expectedRetryAfter.toSeconds())); assertThat(response.getHeaderString("Retry-After")).isEqualTo(String.valueOf(expectedRetryAfter.toSeconds()));
} }

View File

@ -292,7 +292,7 @@ class AccountControllerV2Test {
@Test @Test
void rateLimitedNumber() throws Exception { void rateLimitedNumber() throws Exception {
doThrow(new RateLimitExceededException(null, true)) doThrow(new RateLimitExceededException(null))
.when(registrationLimiter).validate(anyString()); .when(registrationLimiter).validate(anyString());
final Invocation.Builder request = resources.getJerseyTest() final Invocation.Builder request = resources.getJerseyTest()
@ -364,7 +364,7 @@ class AccountControllerV2Test {
final Exception e = switch (error) { final Exception e = switch (error) {
case MISMATCH -> new WebApplicationException(error.getExpectedStatus()); case MISMATCH -> new WebApplicationException(error.getExpectedStatus());
case RATE_LIMITED -> new RateLimitExceededException(null, true); case RATE_LIMITED -> new RateLimitExceededException(null);
}; };
doThrow(e) doThrow(e)
.when(registrationLockVerificationManager).verifyRegistrationLock(any(), any(), any(), any(), any()); .when(registrationLockVerificationManager).verifyRegistrationLock(any(), any(), any(), any(), any());

View File

@ -252,7 +252,7 @@ public class ArchiveControllerTest {
public static Stream<Arguments> setBackupIdException() { public static Stream<Arguments> setBackupIdException() {
return Stream.of( return Stream.of(
Arguments.of(new RateLimitExceededException(null, false), false, 429), Arguments.of(new RateLimitExceededException(null), false, 429),
Arguments.of(Status.INVALID_ARGUMENT.withDescription("async").asRuntimeException(), false, 400), Arguments.of(Status.INVALID_ARGUMENT.withDescription("async").asRuntimeException(), false, 400),
Arguments.of(Status.INVALID_ARGUMENT.withDescription("sync").asRuntimeException(), true, 400) Arguments.of(Status.INVALID_ARGUMENT.withDescription("sync").asRuntimeException(), true, 400)
); );
@ -529,7 +529,7 @@ public class ArchiveControllerTest {
// rate limit // rate limit
when(backupManager.createTemporaryAttachmentUploadDescriptor(any())) when(backupManager.createTemporaryAttachmentUploadDescriptor(any()))
.thenReturn(CompletableFuture.failedFuture(new RateLimitExceededException(null, false))); .thenReturn(CompletableFuture.failedFuture(new RateLimitExceededException(null)));
final Response response = resources.getJerseyTest() final Response response = resources.getJerseyTest()
.target("v1/archives/media/upload/form") .target("v1/archives/media/upload/form")
.request() .request()

View File

@ -122,7 +122,7 @@ public class CallLinkControllerTest {
@Test @Test
void testGetCreateAuthRatelimited() throws RateLimitExceededException{ void testGetCreateAuthRatelimited() throws RateLimitExceededException{
doThrow(new RateLimitExceededException(null, false)) doThrow(new RateLimitExceededException(null))
.when(createCallLinkLimiter).validate(AuthHelper.VALID_UUID); .when(createCallLinkLimiter).validate(AuthHelper.VALID_UUID);
try (Response response = resources.getJerseyTest() try (Response response = resources.getJerseyTest()

View File

@ -170,7 +170,7 @@ class CallRoutingControllerTest {
@Test @Test
void testGetTurnEndpointRateLimited() throws RateLimitExceededException { void testGetTurnEndpointRateLimited() throws RateLimitExceededException {
doThrow(new RateLimitExceededException(null, false)) doThrow(new RateLimitExceededException(null))
.when(getCallEndpointLimiter).validate(AuthHelper.VALID_UUID); .when(getCallEndpointLimiter).validate(AuthHelper.VALID_UUID);
try (final Response response = resources.getJerseyTest() try (final Response response = resources.getJerseyTest()

View File

@ -100,7 +100,7 @@ class ChallengeControllerTest {
"""; """;
final Duration retryAfter = Duration.ofMinutes(17); final Duration retryAfter = Duration.ofMinutes(17);
doThrow(new RateLimitExceededException(retryAfter, true)).when(rateLimitChallengeManager) doThrow(new RateLimitExceededException(retryAfter)).when(rateLimitChallengeManager)
.answerPushChallenge(any(), any()); .answerPushChallenge(any(), any());
final Response response = EXTENSION.target("/v1/challenge") final Response response = EXTENSION.target("/v1/challenge")
@ -108,7 +108,7 @@ class ChallengeControllerTest {
.header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD)) .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
.put(Entity.json(pushChallengeJson)); .put(Entity.json(pushChallengeJson));
assertEquals(413, response.getStatus()); assertEquals(429, response.getStatus());
assertEquals(String.valueOf(retryAfter.toSeconds()), response.getHeaderString("Retry-After")); assertEquals(String.valueOf(retryAfter.toSeconds()), response.getHeaderString("Retry-After"));
} }
@ -175,7 +175,7 @@ class ChallengeControllerTest {
"""; """;
final Duration retryAfter = Duration.ofMinutes(17); final Duration retryAfter = Duration.ofMinutes(17);
doThrow(new RateLimitExceededException(retryAfter, true)).when(rateLimitChallengeManager) doThrow(new RateLimitExceededException(retryAfter)).when(rateLimitChallengeManager)
.answerCaptchaChallenge(any(), any(), any(), any(), any()); .answerCaptchaChallenge(any(), any(), any(), any(), any());
final Response response = EXTENSION.target("/v1/challenge") final Response response = EXTENSION.target("/v1/challenge")
@ -183,7 +183,7 @@ class ChallengeControllerTest {
.header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD)) .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
.put(Entity.json(captchaChallengeJson)); .put(Entity.json(captchaChallengeJson));
assertEquals(413, response.getStatus()); assertEquals(429, response.getStatus());
assertEquals(String.valueOf(retryAfter.toSeconds()), response.getHeaderString("Retry-After")); assertEquals(String.valueOf(retryAfter.toSeconds()), response.getHeaderString("Retry-After"));
} }

View File

@ -462,7 +462,7 @@ class KeysControllerTest {
@Test @Test
void testGetKeysRateLimited() throws RateLimitExceededException { void testGetKeysRateLimited() throws RateLimitExceededException {
Duration retryAfter = Duration.ofSeconds(31); Duration retryAfter = Duration.ofSeconds(31);
doThrow(new RateLimitExceededException(retryAfter, true)).when(rateLimiter).validate(anyString()); doThrow(new RateLimitExceededException(retryAfter)).when(rateLimiter).validate(anyString());
Response result = resources.getJerseyTest() Response result = resources.getJerseyTest()
.target(String.format("/v2/keys/PNI:%s/*", EXISTS_PNI)) .target(String.format("/v2/keys/PNI:%s/*", EXISTS_PNI))
@ -470,7 +470,7 @@ class KeysControllerTest {
.header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD)) .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
.get(); .get();
assertThat(result.getStatus()).isEqualTo(413); assertThat(result.getStatus()).isEqualTo(429);
assertThat(result.getHeaderString("Retry-After")).isEqualTo(String.valueOf(retryAfter.toSeconds())); assertThat(result.getHeaderString("Retry-After")).isEqualTo(String.valueOf(retryAfter.toSeconds()));
} }

View File

@ -1695,10 +1695,10 @@ class MessageControllerTest {
.header(HeaderUtils.UNIDENTIFIED_ACCESS_KEY, Base64.getEncoder().encodeToString(UNIDENTIFIED_ACCESS_BYTES)); .header(HeaderUtils.UNIDENTIFIED_ACCESS_KEY, Base64.getEncoder().encodeToString(UNIDENTIFIED_ACCESS_BYTES));
when(rateLimiter.validateAsync(any(UUID.class))) when(rateLimiter.validateAsync(any(UUID.class)))
.thenReturn(CompletableFuture.failedFuture(new RateLimitExceededException(Duration.ofSeconds(77), true))); .thenReturn(CompletableFuture.failedFuture(new RateLimitExceededException(Duration.ofSeconds(77))));
try (final Response response = invocationBuilder.put(entity)) { try (final Response response = invocationBuilder.put(entity)) {
assertEquals(413, response.getStatus()); assertEquals(429, response.getStatus());
} }
} }

View File

@ -263,7 +263,7 @@ class ProfileControllerTest {
@Test @Test
void testProfileGetByAciRateLimited() throws RateLimitExceededException { void testProfileGetByAciRateLimited() throws RateLimitExceededException {
doThrow(new RateLimitExceededException(Duration.ofSeconds(13), true)).when(rateLimiter) doThrow(new RateLimitExceededException(Duration.ofSeconds(13))).when(rateLimiter)
.validate(AuthHelper.VALID_UUID); .validate(AuthHelper.VALID_UUID);
final Response response = resources.getJerseyTest() final Response response = resources.getJerseyTest()
@ -272,7 +272,7 @@ class ProfileControllerTest {
.header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD)) .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
.get(); .get();
assertThat(response.getStatus()).isEqualTo(413); assertThat(response.getStatus()).isEqualTo(429);
assertThat(response.getHeaderString("Retry-After")).isEqualTo(String.valueOf(Duration.ofSeconds(13).toSeconds())); assertThat(response.getHeaderString("Retry-After")).isEqualTo(String.valueOf(Duration.ofSeconds(13).toSeconds()));
} }
@ -392,7 +392,7 @@ class ProfileControllerTest {
@Test @Test
void testProfileGetByPniRateLimited() throws RateLimitExceededException { void testProfileGetByPniRateLimited() throws RateLimitExceededException {
doThrow(new RateLimitExceededException(Duration.ofSeconds(13), true)).when(rateLimiter) doThrow(new RateLimitExceededException(Duration.ofSeconds(13))).when(rateLimiter)
.validate(AuthHelper.VALID_UUID); .validate(AuthHelper.VALID_UUID);
final Response response = resources.getJerseyTest() final Response response = resources.getJerseyTest()
@ -401,7 +401,7 @@ class ProfileControllerTest {
.header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD)) .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
.get(); .get();
assertThat(response.getStatus()).isEqualTo(413); assertThat(response.getStatus()).isEqualTo(429);
assertThat(response.getHeaderString("Retry-After")).isEqualTo(String.valueOf(Duration.ofSeconds(13).toSeconds())); assertThat(response.getHeaderString("Retry-After")).isEqualTo(String.valueOf(Duration.ofSeconds(13).toSeconds()));
} }

View File

@ -101,7 +101,7 @@ class ProvisioningControllerTest {
final String destination = UUID.randomUUID().toString(); final String destination = UUID.randomUUID().toString();
final byte[] messageBody = "test".getBytes(StandardCharsets.UTF_8); final byte[] messageBody = "test".getBytes(StandardCharsets.UTF_8);
doThrow(new RateLimitExceededException(Duration.ZERO, true)) doThrow(new RateLimitExceededException(Duration.ZERO))
.when(messagesRateLimiter).validate(AuthHelper.VALID_UUID); .when(messagesRateLimiter).validate(AuthHelper.VALID_UUID);
try (final Response response = RESOURCE_EXTENSION.getJerseyTest() try (final Response response = RESOURCE_EXTENSION.getJerseyTest()
@ -111,7 +111,7 @@ class ProvisioningControllerTest {
.put(Entity.entity(new ProvisioningMessage(Base64.getMimeEncoder().encodeToString(messageBody)), .put(Entity.entity(new ProvisioningMessage(Base64.getMimeEncoder().encodeToString(messageBody)),
MediaType.APPLICATION_JSON))) { MediaType.APPLICATION_JSON))) {
assertEquals(413, response.getStatus()); assertEquals(429, response.getStatus());
verify(provisioningManager, never()).sendProvisioningMessage(any(), any()); verify(provisioningManager, never()).sendProvisioningMessage(any(), any());
} }

View File

@ -339,7 +339,7 @@ class RegistrationControllerTest {
} else if (error != null) { } else if (error != null) {
final Exception e = switch (error) { final Exception e = switch (error) {
case MISMATCH -> new WebApplicationException(error.getExpectedStatus()); case MISMATCH -> new WebApplicationException(error.getExpectedStatus());
case RATE_LIMITED -> new RateLimitExceededException(null, true); case RATE_LIMITED -> new RateLimitExceededException(null);
}; };
doThrow(e) doThrow(e)
.when(registrationLockVerificationManager).verifyRegistrationLock(any(), any(), any(), any(), any()); .when(registrationLockVerificationManager).verifyRegistrationLock(any(), any(), any(), any(), any());

View File

@ -180,7 +180,7 @@ class VerificationControllerTest {
@Test @Test
void createSessionRateLimited() { void createSessionRateLimited() {
when(registrationServiceClient.createRegistrationSession(any(), anyBoolean(), any())) when(registrationServiceClient.createRegistrationSession(any(), anyBoolean(), any()))
.thenReturn(CompletableFuture.failedFuture(new RateLimitExceededException(null, true))); .thenReturn(CompletableFuture.failedFuture(new RateLimitExceededException(null)));
final Invocation.Builder request = resources.getJerseyTest() final Invocation.Builder request = resources.getJerseyTest()
.target("/v1/verification/session") .target("/v1/verification/session")

View File

@ -121,7 +121,7 @@ class AccountsAnonymousGrpcServiceTest extends
final Duration retryAfter = Duration.ofSeconds(11); final Duration retryAfter = Duration.ofSeconds(11);
when(rateLimiter.validateReactive(anyString())) when(rateLimiter.validateReactive(anyString()))
.thenReturn(Mono.error(new RateLimitExceededException(retryAfter, false))); .thenReturn(Mono.error(new RateLimitExceededException(retryAfter)));
//noinspection ResultOfMethodCallIgnored //noinspection ResultOfMethodCallIgnored
GrpcTestUtils.assertRateLimitExceeded(retryAfter, GrpcTestUtils.assertRateLimitExceeded(retryAfter,
@ -186,7 +186,7 @@ class AccountsAnonymousGrpcServiceTest extends
final Duration retryAfter = Duration.ofSeconds(13); final Duration retryAfter = Duration.ofSeconds(13);
when(rateLimiter.validateReactive(anyString())) when(rateLimiter.validateReactive(anyString()))
.thenReturn(Mono.error(new RateLimitExceededException(retryAfter, false))); .thenReturn(Mono.error(new RateLimitExceededException(retryAfter)));
//noinspection ResultOfMethodCallIgnored //noinspection ResultOfMethodCallIgnored
GrpcTestUtils.assertRateLimitExceeded(retryAfter, GrpcTestUtils.assertRateLimitExceeded(retryAfter,
@ -254,7 +254,7 @@ class AccountsAnonymousGrpcServiceTest extends
final Duration retryAfter = Duration.ofSeconds(17); final Duration retryAfter = Duration.ofSeconds(17);
when(rateLimiter.validateReactive(anyString())) when(rateLimiter.validateReactive(anyString()))
.thenReturn(Mono.error(new RateLimitExceededException(retryAfter, false))); .thenReturn(Mono.error(new RateLimitExceededException(retryAfter)));
//noinspection ResultOfMethodCallIgnored //noinspection ResultOfMethodCallIgnored
GrpcTestUtils.assertRateLimitExceeded(retryAfter, GrpcTestUtils.assertRateLimitExceeded(retryAfter,

View File

@ -342,7 +342,7 @@ class AccountsGrpcServiceTest extends SimpleBaseGrpcTest<AccountsGrpcService, Ac
final Duration retryAfter = Duration.ofMinutes(3); final Duration retryAfter = Duration.ofMinutes(3);
when(rateLimiter.validateReactive(any(UUID.class))) when(rateLimiter.validateReactive(any(UUID.class)))
.thenReturn(Mono.error(new RateLimitExceededException(retryAfter, false))); .thenReturn(Mono.error(new RateLimitExceededException(retryAfter)));
//noinspection ResultOfMethodCallIgnored //noinspection ResultOfMethodCallIgnored
GrpcTestUtils.assertRateLimitExceeded(retryAfter, GrpcTestUtils.assertRateLimitExceeded(retryAfter,
@ -577,7 +577,7 @@ class AccountsGrpcServiceTest extends SimpleBaseGrpcTest<AccountsGrpcService, Ac
final Duration retryAfter = Duration.ofSeconds(97); final Duration retryAfter = Duration.ofSeconds(97);
when(rateLimiter.validateReactive(any(UUID.class))) when(rateLimiter.validateReactive(any(UUID.class)))
.thenReturn(Mono.error(new RateLimitExceededException(retryAfter, false))); .thenReturn(Mono.error(new RateLimitExceededException(retryAfter)));
final byte[] usernameCiphertext = TestRandomUtil.nextBytes(EncryptedUsername.MAX_SIZE); final byte[] usernameCiphertext = TestRandomUtil.nextBytes(EncryptedUsername.MAX_SIZE);
@ -607,7 +607,7 @@ class AccountsGrpcServiceTest extends SimpleBaseGrpcTest<AccountsGrpcService, Ac
final Duration retryAfter = Duration.ofSeconds(11); final Duration retryAfter = Duration.ofSeconds(11);
when(rateLimiter.validateReactive(any(UUID.class))) when(rateLimiter.validateReactive(any(UUID.class)))
.thenReturn(Mono.error(new RateLimitExceededException(retryAfter, false))); .thenReturn(Mono.error(new RateLimitExceededException(retryAfter)));
//noinspection ResultOfMethodCallIgnored //noinspection ResultOfMethodCallIgnored
GrpcTestUtils.assertRateLimitExceeded(retryAfter, GrpcTestUtils.assertRateLimitExceeded(retryAfter,

View File

@ -628,7 +628,7 @@ class KeysGrpcServiceTest extends SimpleBaseGrpcTest<KeysGrpcService, KeysGrpc.K
final Duration retryAfterDuration = Duration.ofMinutes(7); final Duration retryAfterDuration = Duration.ofMinutes(7);
when(preKeysRateLimiter.validateReactive(anyString())) when(preKeysRateLimiter.validateReactive(anyString()))
.thenReturn(Mono.error(new RateLimitExceededException(retryAfterDuration, false))); .thenReturn(Mono.error(new RateLimitExceededException(retryAfterDuration)));
assertRateLimitExceeded(retryAfterDuration, () -> authenticatedServiceStub().getPreKeys(GetPreKeysRequest.newBuilder() assertRateLimitExceeded(retryAfterDuration, () -> authenticatedServiceStub().getPreKeys(GetPreKeysRequest.newBuilder()
.setTargetIdentifier(ServiceIdentifier.newBuilder() .setTargetIdentifier(ServiceIdentifier.newBuilder()

View File

@ -477,7 +477,7 @@ public class ProfileGrpcServiceTest extends SimpleBaseGrpcTest<ProfileGrpcServic
final Duration retryAfterDuration = Duration.ofMinutes(7); final Duration retryAfterDuration = Duration.ofMinutes(7);
when(accountsManager.getByServiceIdentifierAsync(any())).thenReturn(CompletableFuture.completedFuture(Optional.of(account))); when(accountsManager.getByServiceIdentifierAsync(any())).thenReturn(CompletableFuture.completedFuture(Optional.of(account)));
when(rateLimiter.validateReactive(any(UUID.class))) when(rateLimiter.validateReactive(any(UUID.class)))
.thenReturn(Mono.error(new RateLimitExceededException(retryAfterDuration, false))); .thenReturn(Mono.error(new RateLimitExceededException(retryAfterDuration)));
final GetUnversionedProfileRequest request = GetUnversionedProfileRequest.newBuilder() final GetUnversionedProfileRequest request = GetUnversionedProfileRequest.newBuilder()
.setServiceIdentifier(ServiceIdentifier.newBuilder() .setServiceIdentifier(ServiceIdentifier.newBuilder()

View File

@ -69,11 +69,11 @@ public class RateLimitedByIpTest {
public void testRateLimits() throws Exception { public void testRateLimits() throws Exception {
doNothing().when(RATE_LIMITER).validate(eq(IP)); doNothing().when(RATE_LIMITER).validate(eq(IP));
validateSuccess("/test/strict"); validateSuccess("/test/strict");
doThrow(new RateLimitExceededException(RETRY_AFTER, true)).when(RATE_LIMITER).validate(eq(IP)); doThrow(new RateLimitExceededException(RETRY_AFTER)).when(RATE_LIMITER).validate(eq(IP));
validateFailure("/test/strict", RETRY_AFTER); validateFailure("/test/strict", RETRY_AFTER);
doNothing().when(RATE_LIMITER).validate(eq(IP)); doNothing().when(RATE_LIMITER).validate(eq(IP));
validateSuccess("/test/strict"); validateSuccess("/test/strict");
doThrow(new RateLimitExceededException(RETRY_AFTER, true)).when(RATE_LIMITER).validate(eq(IP)); doThrow(new RateLimitExceededException(RETRY_AFTER)).when(RATE_LIMITER).validate(eq(IP));
validateFailure("/test/strict", RETRY_AFTER); validateFailure("/test/strict", RETRY_AFTER);
} }
@ -92,7 +92,7 @@ public class RateLimitedByIpTest {
.request() .request()
.get(); .get();
assertEquals(413, response.getStatus()); assertEquals(429, response.getStatus());
assertEquals("" + expectedRetryAfter.getSeconds(), response.getHeaderString(HttpHeaders.RETRY_AFTER)); assertEquals("" + expectedRetryAfter.getSeconds(), response.getHeaderString(HttpHeaders.RETRY_AFTER));
} }
} }

View File

@ -101,7 +101,7 @@ public final class MockUtils {
final Duration retryAfter, final Duration retryAfter,
final boolean legacyStatusCode) { final boolean legacyStatusCode) {
try { try {
final RateLimitExceededException exception = new RateLimitExceededException(retryAfter, legacyStatusCode); final RateLimitExceededException exception = new RateLimitExceededException(retryAfter);
doThrow(exception).when(mockRateLimiter).validate(eq(input)); doThrow(exception).when(mockRateLimiter).validate(eq(input));
doReturn(CompletableFuture.failedFuture(exception)).when(mockRateLimiter).validateAsync(eq(input)); doReturn(CompletableFuture.failedFuture(exception)).when(mockRateLimiter).validateAsync(eq(input));
doReturn(Mono.fromFuture(CompletableFuture.failedFuture(exception))).when(mockRateLimiter).validateReactive(eq(input)); doReturn(Mono.fromFuture(CompletableFuture.failedFuture(exception))).when(mockRateLimiter).validateReactive(eq(input));
@ -117,7 +117,7 @@ public final class MockUtils {
final Duration retryAfter, final Duration retryAfter,
final boolean legacyStatusCode) { final boolean legacyStatusCode) {
try { try {
final RateLimitExceededException exception = new RateLimitExceededException(retryAfter, legacyStatusCode); final RateLimitExceededException exception = new RateLimitExceededException(retryAfter);
doThrow(exception).when(mockRateLimiter).validate(eq(input)); doThrow(exception).when(mockRateLimiter).validate(eq(input));
doReturn(CompletableFuture.failedFuture(exception)).when(mockRateLimiter).validateAsync(eq(input)); doReturn(CompletableFuture.failedFuture(exception)).when(mockRateLimiter).validateAsync(eq(input));
doReturn(Mono.fromFuture(CompletableFuture.failedFuture(exception))).when(mockRateLimiter).validateReactive(eq(input)); doReturn(Mono.fromFuture(CompletableFuture.failedFuture(exception))).when(mockRateLimiter).validateReactive(eq(input));