From 5354104128eb26c22e70aa118dcb026499b980a5 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Tue, 23 Feb 2021 11:49:09 -0500
Subject: [PATCH] Only apply unsealed sender rate limits to targeted country
 codes.

---
 .../DynamicMessageRateConfiguration.java      | 10 +++++++++
 .../controllers/MessageController.java        | 21 ++++++++++---------
 .../controllers/MessageControllerTest.java    |  3 +++
 3 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicMessageRateConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicMessageRateConfiguration.java
index fdf8d145d..6043765a2 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicMessageRateConfiguration.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicMessageRateConfiguration.java
@@ -7,12 +7,22 @@ package org.whispersystems.textsecuregcm.configuration.dynamic;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 
+import java.util.Collections;
+import java.util.Set;
+
 public class DynamicMessageRateConfiguration {
 
   @JsonProperty
   private boolean enforceUnsealedSenderRateLimit = false;
 
+  @JsonProperty
+  private Set<String> rateLimitedCountryCodes = Collections.emptySet();
+
   public boolean isEnforceUnsealedSenderRateLimit() {
     return enforceUnsealedSenderRateLimit;
   }
+
+  public Set<String> getRateLimitedCountryCodes() {
+    return rateLimitedCountryCodes;
+  }
 }
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
index b43f26b04..b7eded59d 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
@@ -12,7 +12,6 @@ import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.SharedMetricRegistries;
 import com.codahale.metrics.Timer;
 import com.codahale.metrics.annotation.Timed;
-import com.google.i18n.phonenumbers.PhoneNumberUtil;
 import com.google.protobuf.ByteString;
 import io.dropwizard.auth.Auth;
 import io.dropwizard.util.DataSize;
@@ -172,16 +171,18 @@ public class MessageController {
         });
       });
 
-      try {
-        rateLimiters.getUnsealedSenderLimiter().validate(source.get().getUuid().toString(), destinationName.toString());
-      } catch (RateLimitExceededException e) {
-        Metrics.counter(REJECT_UNSEALED_SENDER_COUNTER_NAME, SENDER_COUNTRY_TAG_NAME, Util.getCountryCode(source.get().getNumber())).increment();
+      if (dynamicConfigurationManager.getConfiguration().getMessageRateConfiguration().getRateLimitedCountryCodes().contains(senderCountryCode)) {
+        try {
+          rateLimiters.getUnsealedSenderLimiter().validate(source.get().getUuid().toString(), destinationName.toString());
+        } catch (RateLimitExceededException e) {
+          Metrics.counter(REJECT_UNSEALED_SENDER_COUNTER_NAME, SENDER_COUNTRY_TAG_NAME, senderCountryCode).increment();
 
-        if (dynamicConfigurationManager.getConfiguration().getMessageRateConfiguration().isEnforceUnsealedSenderRateLimit()) {
-          logger.debug("Rejected unsealed sender limit from: {}", source.get().getNumber());
-          throw e;
-        } else {
-          logger.debug("Would reject unsealed sender limit from: {}", source.get().getNumber());
+          if (dynamicConfigurationManager.getConfiguration().getMessageRateConfiguration().isEnforceUnsealedSenderRateLimit()) {
+            logger.debug("Rejected unsealed sender limit from: {}", source.get().getNumber());
+            throw e;
+          } else {
+            logger.debug("Would reject unsealed sender limit from: {}", source.get().getNumber());
+          }
         }
       }
     }
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/MessageControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/MessageControllerTest.java
index d285e5270..cfce99395 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/MessageControllerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/MessageControllerTest.java
@@ -51,6 +51,7 @@ import org.mockito.ArgumentMatcher;
 import org.whispersystems.textsecuregcm.auth.AmbiguousIdentifier;
 import org.whispersystems.textsecuregcm.auth.DisabledPermittedAccount;
 import org.whispersystems.textsecuregcm.auth.OptionalAccess;
+import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.controllers.MessageController;
 import org.whispersystems.textsecuregcm.entities.IncomingMessageList;
 import org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope;
@@ -132,6 +133,8 @@ public class MessageControllerTest {
 
     when(rateLimiters.getMessagesLimiter()).thenReturn(rateLimiter);
     when(rateLimiters.getUnsealedSenderLimiter()).thenReturn(unsealedSenderLimiter);
+
+    when(dynamicConfigurationManager.getConfiguration()).thenReturn(new DynamicConfiguration());
   }
 
   @Test