pass challenge type to rate limit reset listeners

signalapp · Aug 29, 2023 · 9577d55 · 9577d55
1 parent 093f17d
commit 9577d55
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 7 deletions.
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimitChallengeManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimitChallengeManager.java
@@ -19,6 +19,7 @@
 import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
 import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
 import org.whispersystems.textsecuregcm.push.NotPushRegisteredException;
+import org.whispersystems.textsecuregcm.spam.ChallengeType;
 import org.whispersystems.textsecuregcm.spam.RateLimitChallengeListener;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.util.Util;
@@ -59,7 +60,7 @@ public void answerPushChallenge(final Account account, final String challenge) t
 
     if (challengeSuccess) {
       rateLimiters.getPushChallengeSuccessLimiter().validate(account.getUuid());
-      resetRateLimits(account);
+      resetRateLimits(account, ChallengeType.PUSH);
     }
   }
 
@@ -80,12 +81,12 @@ public boolean answerRecaptchaChallenge(final Account account, final String capt
 
     if (challengeSuccess) {
       rateLimiters.getRecaptchaChallengeSuccessLimiter().validate(account.getUuid());
-      resetRateLimits(account);
+      resetRateLimits(account, ChallengeType.CAPTCHA);
     }
     return challengeSuccess;
   }
 
-  private void resetRateLimits(final Account account) throws RateLimitExceededException {
+  private void resetRateLimits(final Account account, final ChallengeType type) throws RateLimitExceededException {
     try {
       rateLimiters.getRateLimitResetLimiter().validate(account.getUuid());
     } catch (final RateLimitExceededException e) {
@@ -95,7 +96,7 @@ private void resetRateLimits(final Account account) throws RateLimitExceededExce
       throw e;
     }
 
-    rateLimitChallengeListeners.forEach(listener -> listener.handleRateLimitChallengeAnswered(account));
+    rateLimitChallengeListeners.forEach(listener -> listener.handleRateLimitChallengeAnswered(account, type));
   }
 
   public void sendPushChallenge(final Account account) throws NotPushRegisteredException {

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/spam/ChallengeType.java b/service/src/main/java/org/whispersystems/textsecuregcm/spam/ChallengeType.java
@@ -0,0 +1,11 @@
+/*
+ * Copyright 2013-2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.spam;
+
+public enum ChallengeType {
+    PUSH,
+    CAPTCHA
+}
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/spam/RateLimitChallengeListener.java b/service/src/main/java/org/whispersystems/textsecuregcm/spam/RateLimitChallengeListener.java
@@ -11,7 +11,7 @@
 
 public interface RateLimitChallengeListener {
 
-  void handleRateLimitChallengeAnswered(Account account);
+  void handleRateLimitChallengeAnswered(Account account, ChallengeType type);
 
   /**
    * Configures this rate limit challenge listener. This method will be called before the service begins processing any

diff --git a/.../src/test/java/org/whispersystems/textsecuregcm/limits/RateLimitChallengeManagerTest.java b/.../src/test/java/org/whispersystems/textsecuregcm/limits/RateLimitChallengeManagerTest.java
@@ -21,6 +21,7 @@
 import org.whispersystems.textsecuregcm.captcha.AssessmentResult;
 import org.whispersystems.textsecuregcm.captcha.CaptchaChecker;
 import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
+import org.whispersystems.textsecuregcm.spam.ChallengeType;
 import org.whispersystems.textsecuregcm.spam.RateLimitChallengeListener;
 import org.whispersystems.textsecuregcm.storage.Account;
 
@@ -63,7 +64,7 @@ void answerPushChallenge(final boolean successfulChallenge) throws RateLimitExce
     rateLimitChallengeManager.answerPushChallenge(account, "challenge");
 
     if (successfulChallenge) {
-      verify(rateLimitChallengeListener).handleRateLimitChallengeAnswered(account);
+      verify(rateLimitChallengeListener).handleRateLimitChallengeAnswered(account, ChallengeType.PUSH);
     } else {
       verifyNoInteractions(rateLimitChallengeListener);
     }
@@ -88,7 +89,7 @@ void answerRecaptchaChallenge(final boolean successfulChallenge) throws RateLimi
     rateLimitChallengeManager.answerRecaptchaChallenge(account, "captcha", "10.0.0.1", "Test User-Agent");
 
     if (successfulChallenge) {
-      verify(rateLimitChallengeListener).handleRateLimitChallengeAnswered(account);
+      verify(rateLimitChallengeListener).handleRateLimitChallengeAnswered(account, ChallengeType.CAPTCHA);
     } else {
       verifyNoInteractions(rateLimitChallengeListener);
     }