Return a Retry-After on rate-limited responses

Previously, only endpoints throwing a RetryLaterException would include a Retry-After header in the 413 response. Now, by default, all RateLimitExceededExceptions will be marshalled into a 413 with a Retry-After included if possible.
2026-04-21 14:28:05 +01:00 · 2022-02-16 14:34:25 -06:00
parent 43792e2426
commit ae3a5c5f5e
11 changed files with 103 additions and 77 deletions
--- a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ChallengeControllerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ChallengeControllerTest.java
@@ -27,7 +27,7 @@ import org.junit.jupiter.api.extension.ExtendWith;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
 import org.whispersystems.textsecuregcm.auth.DisabledPermittedAuthenticatedAccount;
 import org.whispersystems.textsecuregcm.limits.RateLimitChallengeManager;
-import org.whispersystems.textsecuregcm.mappers.RetryLaterExceptionMapper;
+import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper;
 import org.whispersystems.textsecuregcm.push.NotPushRegisteredException;
 import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
 import org.whispersystems.textsecuregcm.util.SystemMapper;
@@ -45,7 +45,7 @@ class ChallengeControllerTest {
          Set.of(AuthenticatedAccount.class, DisabledPermittedAuthenticatedAccount.class)))
      .setMapper(SystemMapper.getMapper())
      .setTestContainerFactory(new GrizzlyWebTestContainerFactory())
-      .addResource(new RetryLaterExceptionMapper())
+      .addResource(new RateLimitExceededExceptionMapper())
      .addResource(challengeController)
      .build();

--- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/AccountControllerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/AccountControllerTest.java
@@ -1804,6 +1804,38 @@ class AccountControllerTest {
        .getStatus()).isEqualTo(404);
  }

+  @Test
+  void testAccountExistsRateLimited() throws RateLimitExceededException {
+    final Account account = mock(Account.class);
+    final UUID accountIdentifier = UUID.randomUUID();
+    when(accountsManager.getByAccountIdentifier(accountIdentifier)).thenReturn(Optional.of(account));
+
+    final RateLimiter checkAccountLimiter = mock(RateLimiter.class);
+    when(rateLimiters.getCheckAccountExistenceLimiter()).thenReturn(checkAccountLimiter);
+    doThrow(new RateLimitExceededException(Duration.ofSeconds(13))).when(checkAccountLimiter).validate("127.0.0.1");
+
+    final Response response = resources.getJerseyTest()
+        .target(String.format("/v1/accounts/account/%s", accountIdentifier))
+        .request()
+        .header("X-Forwarded-For", "127.0.0.1")
+        .head();
+
+    assertThat(response.getStatus()).isEqualTo(413);
+    assertThat(response.getHeaderString("Retry-After")).isEqualTo(String.valueOf(Duration.ofSeconds(13).toSeconds()));
+  }
+
+  @Test
+  void testAccountExistsNoForwardedFor() throws RateLimitExceededException {
+    final Response response = resources.getJerseyTest()
+        .target(String.format("/v1/accounts/account/%s", UUID.randomUUID()))
+        .request()
+        .header("X-Forwarded-For", "")
+        .head();
+
+    assertThat(response.getStatus()).isEqualTo(413);
+    assertThat(response.getHeaderString("Retry-After")).isNull();
+  }
+
  @Test
  void testAccountExistsAuthenticated() {
    assertThat(resources.getJerseyTest()
--- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java
@@ -8,6 +8,7 @@ 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.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.eq;
@@ -38,8 +39,6 @@ import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
-import org.junit.jupiter.params.ParameterizedTest;
-import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.ArgumentCaptor;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
 import org.whispersystems.textsecuregcm.auth.DisabledPermittedAuthenticatedAccount;
@@ -50,11 +49,11 @@ import org.whispersystems.textsecuregcm.entities.PreKey;
 import org.whispersystems.textsecuregcm.entities.PreKeyCount;
 import org.whispersystems.textsecuregcm.entities.PreKeyResponse;
 import org.whispersystems.textsecuregcm.entities.PreKeyState;
-import org.whispersystems.textsecuregcm.entities.RateLimitChallenge;
 import org.whispersystems.textsecuregcm.entities.SignedPreKey;
-import org.whispersystems.textsecuregcm.limits.RateLimitChallengeManager;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
+import org.whispersystems.textsecuregcm.mappers.RateLimitChallengeExceptionMapper;
+import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.ServerRejectedExceptionMapper;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
@@ -107,6 +106,7 @@ class KeysControllerTest {
      .setTestContainerFactory(new GrizzlyWebTestContainerFactory())
      .addResource(new ServerRejectedExceptionMapper())
      .addResource(new KeysController(rateLimiters, KEYS, accounts))
+      .addResource(new RateLimitExceededExceptionMapper())
      .build();

  @BeforeEach
@@ -316,6 +316,21 @@ class KeysControllerTest {
    verifyNoMoreInteractions(KEYS);
  }

+  @Test
+  void testGetKeysRateLimited() throws RateLimitExceededException {
+    Duration retryAfter = Duration.ofSeconds(31);
+    doThrow(new RateLimitExceededException(retryAfter)).when(rateLimiter).validate(anyString());
+
+    Response result = resources.getJerseyTest()
+        .target(String.format("/v2/keys/%s/*", EXISTS_PNI))
+        .request()
+        .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+        .get();
+
+    assertThat(result.getStatus()).isEqualTo(413);
+    assertThat(result.getHeaderString("Retry-After")).isEqualTo(String.valueOf(retryAfter.toSeconds()));
+  }
+
  @Test
  void testUnidentifiedRequest() {
    PreKeyResponse result = resources.getJerseyTest()
--- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/ProfileControllerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/ProfileControllerTest.java
@@ -10,6 +10,7 @@ import static org.assertj.core.api.Assertions.assertThatNoException;
 import static org.mockito.ArgumentMatchers.refEq;
 import static org.mockito.Mockito.any;
 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;
@@ -26,6 +27,7 @@ import io.dropwizard.testing.junit5.ResourceExtension;
 import java.nio.charset.StandardCharsets;
 import java.security.SecureRandom;
 import java.time.Clock;
+import java.time.Duration;
 import java.time.Instant;
 import java.util.Base64;
 import java.util.Collections;
@@ -82,6 +84,7 @@ import org.whispersystems.textsecuregcm.entities.ProfileKeyCredentialProfileResp
 import org.whispersystems.textsecuregcm.entities.VersionedProfileResponse;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
+import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper;
 import org.whispersystems.textsecuregcm.s3.PolicySigner;
 import org.whispersystems.textsecuregcm.s3.PostPolicyGenerator;
 import org.whispersystems.textsecuregcm.storage.Account;
@@ -123,6 +126,7 @@ class ProfileControllerTest {
      .addProvider(AuthHelper.getAuthFilter())
      .addProvider(new PolymorphicAuthValueFactoryProvider.Binder<>(
          ImmutableSet.of(AuthenticatedAccount.class, DisabledPermittedAuthenticatedAccount.class)))
+      .addProvider(new RateLimitExceededExceptionMapper())
      .setMapper(SystemMapper.getMapper())
      .setTestContainerFactory(new GrizzlyWebTestContainerFactory())
      .addResource(new ProfileController(
@@ -210,6 +214,7 @@ class ProfileControllerTest {
  @AfterEach
  void teardown() {
    reset(accountsManager);
+    reset(rateLimiter);
  }

  @Test
@@ -228,6 +233,20 @@ class ProfileControllerTest {
    verify(rateLimiter, times(1)).validate(AuthHelper.VALID_UUID);
  }

+  @Test
+  void testProfileGetByUuidRateLimited() throws RateLimitExceededException {
+    doThrow(new RateLimitExceededException(Duration.ofSeconds(13))).when(rateLimiter).validate(AuthHelper.VALID_UUID);
+
+    Response response= resources.getJerseyTest()
+        .target("/v1/profile/" + AuthHelper.VALID_UUID_TWO)
+        .request()
+        .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+        .get();
+
+    assertThat(response.getStatus()).isEqualTo(413);
+    assertThat(response.getHeaderString("Retry-After")).isEqualTo(String.valueOf(Duration.ofSeconds(13).toSeconds()));
+  }
+
  @Test
  void testProfileGetByUuidUnidentified() throws RateLimitExceededException {
    BaseProfileResponse profile = resources.getJerseyTest()