From 9577d552c6411071dd0c9cff4b3227fcc72531ac Mon Sep 17 00:00:00 2001
From: Jonathan Klabunde Tomer <125505367+jkt-signal@users.noreply.github.com>
Date: Tue, 29 Aug 2023 15:19:49 -0700
Subject: [PATCH] pass challenge type to rate limit reset listeners

---
 .../limits/RateLimitChallengeManager.java             |  9 +++++----
 .../textsecuregcm/spam/ChallengeType.java             | 11 +++++++++++
 .../spam/RateLimitChallengeListener.java              |  2 +-
 .../limits/RateLimitChallengeManagerTest.java         |  5 +++--
 4 files changed, 20 insertions(+), 7 deletions(-)
 create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/spam/ChallengeType.java

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 052576cbe..c2d2b4bda 100644
--- 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.captcha.CaptchaChecker;
 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 class RateLimitChallengeManager {
 
     if (challengeSuccess) {
       rateLimiters.getPushChallengeSuccessLimiter().validate(account.getUuid());
-      resetRateLimits(account);
+      resetRateLimits(account, ChallengeType.PUSH);
     }
   }
 
@@ -80,12 +81,12 @@ public class RateLimitChallengeManager {
 
     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 @@ public class RateLimitChallengeManager {
       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
new file mode 100644
index 000000000..500d65e47
--- /dev/null
+++ 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
index b15fa171d..d4e58e752 100644
--- 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 @@ import java.io.IOException;
 
 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/service/src/test/java/org/whispersystems/textsecuregcm/limits/RateLimitChallengeManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/limits/RateLimitChallengeManagerTest.java
index 938ff757e..7665455e6 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/limits/RateLimitChallengeManagerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/limits/RateLimitChallengeManagerTest.java
@@ -21,6 +21,7 @@ import org.whispersystems.textsecuregcm.captcha.Action;
 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 @@ class RateLimitChallengeManagerTest {
     rateLimitChallengeManager.answerPushChallenge(account, "challenge");
 
     if (successfulChallenge) {
-      verify(rateLimitChallengeListener).handleRateLimitChallengeAnswered(account);
+      verify(rateLimitChallengeListener).handleRateLimitChallengeAnswered(account, ChallengeType.PUSH);
     } else {
       verifyNoInteractions(rateLimitChallengeListener);
     }
@@ -88,7 +89,7 @@ class RateLimitChallengeManagerTest {
     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);
     }