Use non-legacy rate limit error in BackupAuthManager

signalapp · Jul 11, 2024 · 0ef3e00 · 0ef3e00
1 parent 2408590
commit 0ef3e00
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupAuthManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupAuthManager.java
@@ -15,6 +15,7 @@
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Stream;
+import javax.annotation.Nullable;
 import org.signal.libsignal.zkgroup.GenericServerSecretParams;
 import org.signal.libsignal.zkgroup.InvalidInputException;
 import org.signal.libsignal.zkgroup.VerificationFailedException;
@@ -28,12 +29,12 @@
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
 import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
+import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.RedeemedReceiptsManager;
 import org.whispersystems.textsecuregcm.util.Util;
-import javax.annotation.Nullable;
 
 /**
  * Issues ZK backup auth credentials for authenticated accounts
@@ -101,8 +102,9 @@ public CompletableFuture<Void> commitBackupId(final Account account,
       return CompletableFuture.completedFuture(null);
     }
 
-    return rateLimiters.forDescriptor(RateLimiters.For.SET_BACKUP_ID)
-        .validateAsync(account.getUuid())
+    return RateLimiter.adaptLegacyException(rateLimiters
+            .forDescriptor(RateLimiters.For.SET_BACKUP_ID)
+            .validateAsync(account.getUuid()))
         .thenCompose(ignored -> this.accountsManager
             .updateAsync(account, acc -> acc.setBackupCredentialRequest(serializedRequest))
             .thenRun(Util.NOOP))

diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupAuthManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupAuthManagerTest.java
@@ -406,9 +406,10 @@ void testRateLimits() {
     when(accountsManager.updateAsync(any(), any())).thenReturn(CompletableFuture.completedFuture(account));
 
     // Should be rate limited
-    CompletableFutureTestUtil.assertFailsWithCause(
+    final RateLimitExceededException ex = CompletableFutureTestUtil.assertFailsWithCause(
         RateLimitExceededException.class,
         authManager.commitBackupId(account, credentialRequest));
+    assertThat(ex.isLegacy()).isFalse();
 
     // If we don't change the request, shouldn't be rate limited
     when(account.getBackupCredentialRequest()).thenReturn(credentialRequest.serialize());