From ce5edbb7fc4d8f124edd635b5956530ea672f129 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Mon, 2 Aug 2021 14:52:31 -0400
Subject: [PATCH] Continue to verify rate limiters by e164 during UUID
 migration period

---
 .../textsecuregcm/limits/PreKeyRateLimiter.java |  3 +++
 .../limits/RateLimitChallengeManager.java       | 17 +++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/limits/PreKeyRateLimiter.java b/service/src/main/java/org/whispersystems/textsecuregcm/limits/PreKeyRateLimiter.java
index 9486561dc..1d42e8862 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/limits/PreKeyRateLimiter.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/limits/PreKeyRateLimiter.java
@@ -50,6 +50,9 @@ public class PreKeyRateLimiter {
 
     try {
       rateLimiters.getDailyPreKeysLimiter().validate(account.getUuid());
+
+      // TODO Remove after migration period
+      rateLimiters.getDailyPreKeysLimiter().validate(account.getNumber());
     } catch (final RateLimitExceededException e) {
 
       final boolean enforceLimit = dynamicConfigurationManager.getConfiguration()
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimitChallengeManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimitChallengeManager.java
index d756158bc..ede118b8e 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimitChallengeManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimitChallengeManager.java
@@ -56,10 +56,17 @@ public class RateLimitChallengeManager {
   public void answerPushChallenge(final Account account, final String challenge) throws RateLimitExceededException {
     rateLimiters.getPushChallengeAttemptLimiter().validate(account.getUuid());
 
+    // TODO Remove after migration period
+    rateLimiters.getPushChallengeAttemptLimiter().validate(account.getNumber());
+
     final boolean challengeSuccess = pushChallengeManager.answerChallenge(account, challenge);
 
     if (challengeSuccess) {
       rateLimiters.getPushChallengeSuccessLimiter().validate(account.getUuid());
+
+      // TODO Remove after migration period
+      rateLimiters.getPushChallengeSuccessLimiter().validate(account.getNumber());
+
       resetRateLimits(account);
     }
   }
@@ -69,6 +76,9 @@ public class RateLimitChallengeManager {
 
     rateLimiters.getRecaptchaChallengeAttemptLimiter().validate(account.getUuid());
 
+    // TODO Remove after migration period
+    rateLimiters.getRecaptchaChallengeAttemptLimiter().validate(account.getNumber());
+
     final boolean challengeSuccess = recaptchaClient.verify(captcha, mostRecentProxyIp);
 
     Metrics.counter(RECAPTCHA_ATTEMPT_COUNTER_NAME,
@@ -77,6 +87,10 @@ public class RateLimitChallengeManager {
 
     if (challengeSuccess) {
       rateLimiters.getRecaptchaChallengeSuccessLimiter().validate(account.getUuid());
+
+      // TODO Remove after migration period
+      rateLimiters.getRecaptchaChallengeSuccessLimiter().validate(account.getNumber());
+
       resetRateLimits(account);
     }
   }
@@ -84,6 +98,9 @@ public class RateLimitChallengeManager {
   private void resetRateLimits(final Account account) throws RateLimitExceededException {
     try {
       rateLimiters.getRateLimitResetLimiter().validate(account.getUuid());
+
+      // TODO Remove after migration period
+      rateLimiters.getRateLimitResetLimiter().validate(account.getNumber());
     } catch (final RateLimitExceededException e) {
       Metrics.counter(RESET_RATE_LIMIT_EXCEEDED_COUNTER_NAME,
           SOURCE_COUNTRY_TAG_NAME, Util.getCountryCode(account.getNumber())).increment();