From aa60fae3b166eb6c4bf98d48e47cd0b1fc796a92 Mon Sep 17 00:00:00 2001
From: Jonathan Klabunde Tomer <125505367+jkt-signal@users.noreply.github.com>
Date: Wed, 18 Sep 2024 18:38:21 -0400
Subject: [PATCH] install RateLimitByIpFilter in soft-enforcement mode

---
 .../textsecuregcm/WhisperServerService.java   |  4 +++
 .../limits/RateLimitByIpFilter.java           | 31 ++++++++++++++++---
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 27516266b..db981a1eb 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -167,6 +167,7 @@ import org.whispersystems.textsecuregcm.keytransparency.KeyTransparencyServiceCl
 import org.whispersystems.textsecuregcm.limits.CardinalityEstimator;
 import org.whispersystems.textsecuregcm.limits.MessageDeliveryLoopMonitor;
 import org.whispersystems.textsecuregcm.limits.PushChallengeManager;
+import org.whispersystems.textsecuregcm.limits.RateLimitByIpFilter;
 import org.whispersystems.textsecuregcm.limits.RateLimitChallengeManager;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.mappers.CompletionExceptionMapper;
@@ -1003,6 +1004,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
 
     environment.jersey().register(new BufferingInterceptor());
     environment.jersey().register(new VirtualExecutorServiceProvider("managed-async-virtual-thread-"));
+    environment.jersey().register(new RateLimitByIpFilter(rateLimiters, true));
     environment.jersey().register(new RequestStatisticsFilter(TrafficSource.HTTP));
     environment.jersey().register(MultiRecipientMessageProvider.class);
     environment.jersey().register(new AuthDynamicFeature(accountAuthFilter));
@@ -1021,11 +1023,13 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
             clientReleaseManager, messageDeliveryLoopMonitor));
     webSocketEnvironment.jersey()
         .register(new WebsocketRefreshApplicationEventListener(accountsManager, clientPresenceManager));
+    webSocketEnvironment.jersey().register(new RateLimitByIpFilter(rateLimiters, true));
     webSocketEnvironment.jersey().register(new RequestStatisticsFilter(TrafficSource.WEBSOCKET));
     webSocketEnvironment.jersey().register(MultiRecipientMessageProvider.class);
     webSocketEnvironment.jersey().register(new MetricsApplicationEventListener(TrafficSource.WEBSOCKET, clientReleaseManager));
     webSocketEnvironment.jersey().register(new KeepAliveController(clientPresenceManager));
 
+
     final List<SpamFilter> spamFilters = ServiceLoader.load(SpamFilter.class)
         .stream()
         .map(ServiceLoader.Provider::get)
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimitByIpFilter.java b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimitByIpFilter.java
index c7605833b..c008bb89b 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimitByIpFilter.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimitByIpFilter.java
@@ -22,6 +22,11 @@ import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
 import org.whispersystems.textsecuregcm.filters.RemoteAddressFilter;
 import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper;
+import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
+
+import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.Tag;
+import io.micrometer.core.instrument.Tags;
 
 public class RateLimitByIpFilter implements ContainerRequestFilter {
 
@@ -33,11 +38,18 @@ public class RateLimitByIpFilter implements ContainerRequestFilter {
 
   private static final ExceptionMapper<RateLimitExceededException> EXCEPTION_MAPPER = new RateLimitExceededExceptionMapper();
 
-  private final RateLimiters rateLimiters;
+  private static final String NO_IP_COUNTER_NAME = MetricsUtil.name(RateLimitByIpFilter.class, "noIpAddress");
 
+  private final RateLimiters rateLimiters;
+  private final boolean softEnforcement;
+
+  public RateLimitByIpFilter(final RateLimiters rateLimiters, final boolean softEnforcement) {
+    this.rateLimiters = requireNonNull(rateLimiters);
+    this.softEnforcement = softEnforcement;
+  }
 
   public RateLimitByIpFilter(final RateLimiters rateLimiters) {
-    this.rateLimiters = requireNonNull(rateLimiters);
+    this(rateLimiters, false);
   }
 
   @Override
@@ -65,10 +77,19 @@ public class RateLimitByIpFilter implements ContainerRequestFilter {
 
       // checking if we failed to extract the most recent IP for any reason
       if (remoteAddress.isEmpty()) {
+        Metrics.counter(
+            NO_IP_COUNTER_NAME,
+            Tags.of(
+                Tag.of("limiter", handle.id()),
+                Tag.of("fail", String.valueOf(annotation.failOnUnresolvedIp()))))
+            .increment();
+
         // checking if annotation is configured to fail when the most recent IP is not resolved
         if (annotation.failOnUnresolvedIp()) {
           logger.error("Remote address was null");
-          throw INVALID_HEADER_EXCEPTION;
+          if (!softEnforcement) {
+            throw INVALID_HEADER_EXCEPTION;
+          }
         }
         // otherwise, allow request
         return;
@@ -78,7 +99,9 @@ public class RateLimitByIpFilter implements ContainerRequestFilter {
       rateLimiter.validate(remoteAddress.get());
     } catch (RateLimitExceededException e) {
       final Response response = EXCEPTION_MAPPER.toResponse(e);
-      throw new ClientErrorException(response);
+      if (!softEnforcement) {
+        throw new ClientErrorException(response);
+      }
     }
   }
 }