From 836307b0c74e01b0d52dbe59764b69bf5a2076fc Mon Sep 17 00:00:00 2001
From: Sergey Skrobotov <sergey@signal.org>
Date: Mon, 14 Nov 2022 14:58:11 -0800
Subject: [PATCH] adding a metric for ipv4/ipv6 requests count

---
 .../textsecuregcm/WhisperServerService.java   |  6 +-
 .../filters/ContentLengthFilter.java          | 37 ---------
 .../filters/RequestStatisticsFilter.java      | 78 +++++++++++++++++++
 ....java => RequestStatisticsFilterTest.java} | 10 +--
 4 files changed, 86 insertions(+), 45 deletions(-)
 delete mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/filters/ContentLengthFilter.java
 create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/filters/RequestStatisticsFilter.java
 rename service/src/test/java/org/whispersystems/textsecuregcm/filters/{ContentLengthFilterTest.java => RequestStatisticsFilterTest.java} (70%)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index b6a5902d2..a9d8c5f8d 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -110,7 +110,7 @@ import org.whispersystems.textsecuregcm.currency.CoinMarketCapClient;
 import org.whispersystems.textsecuregcm.currency.CurrencyConversionManager;
 import org.whispersystems.textsecuregcm.currency.FixerClient;
 import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
-import org.whispersystems.textsecuregcm.filters.ContentLengthFilter;
+import org.whispersystems.textsecuregcm.filters.RequestStatisticsFilter;
 import org.whispersystems.textsecuregcm.filters.RemoteDeprecationFilter;
 import org.whispersystems.textsecuregcm.filters.TimestampResponseFilter;
 import org.whispersystems.textsecuregcm.limits.DynamicRateLimiters;
@@ -624,7 +624,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         .addFilter("RemoteDeprecationFilter", new RemoteDeprecationFilter(dynamicConfigurationManager))
         .addMappingForUrlPatterns(EnumSet.of(DispatcherType.REQUEST), false, "/*");
 
-    environment.jersey().register(new ContentLengthFilter(TrafficSource.HTTP));
+    environment.jersey().register(new RequestStatisticsFilter(TrafficSource.HTTP));
     environment.jersey().register(MultiRecipientMessageProvider.class);
     environment.jersey().register(new MetricsApplicationEventListener(TrafficSource.HTTP));
     environment.jersey()
@@ -646,7 +646,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
             clientPresenceManager, websocketScheduledExecutor));
     webSocketEnvironment.jersey()
         .register(new WebsocketRefreshApplicationEventListener(accountsManager, clientPresenceManager));
-    webSocketEnvironment.jersey().register(new ContentLengthFilter(TrafficSource.WEBSOCKET));
+    webSocketEnvironment.jersey().register(new RequestStatisticsFilter(TrafficSource.WEBSOCKET));
     webSocketEnvironment.jersey().register(MultiRecipientMessageProvider.class);
     webSocketEnvironment.jersey().register(new MetricsApplicationEventListener(TrafficSource.WEBSOCKET));
     webSocketEnvironment.jersey().register(new KeepAliveController(clientPresenceManager));
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/filters/ContentLengthFilter.java b/service/src/main/java/org/whispersystems/textsecuregcm/filters/ContentLengthFilter.java
deleted file mode 100644
index 36fbc9023..000000000
--- a/service/src/main/java/org/whispersystems/textsecuregcm/filters/ContentLengthFilter.java
+++ /dev/null
@@ -1,37 +0,0 @@
-/*
- * Copyright 2021 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.whispersystems.textsecuregcm.filters;
-
-import static com.codahale.metrics.MetricRegistry.name;
-
-import io.micrometer.core.instrument.Metrics;
-import java.io.IOException;
-import javax.ws.rs.container.ContainerRequestContext;
-import javax.ws.rs.container.ContainerRequestFilter;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import org.whispersystems.textsecuregcm.metrics.TrafficSource;
-
-public class ContentLengthFilter implements ContainerRequestFilter {
-
-  private static final String CONTENT_LENGTH_DISTRIBUTION_NAME = name(ContentLengthFilter.class, "contentLength");
-
-  private static final Logger logger = LoggerFactory.getLogger(ContentLengthFilter.class);
-  private final TrafficSource trafficSource;
-
-  public ContentLengthFilter(final TrafficSource trafficeSource) {
-    this.trafficSource = trafficeSource;
-  }
-
-  @Override
-  public void filter(final ContainerRequestContext requestContext) throws IOException {
-    try {
-      Metrics.summary(CONTENT_LENGTH_DISTRIBUTION_NAME, "trafficSource", trafficSource.name().toLowerCase()).record(requestContext.getLength());
-    } catch (final Exception e) {
-      logger.warn("Error recording content length", e);
-    }
-  }
-}
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/filters/RequestStatisticsFilter.java b/service/src/main/java/org/whispersystems/textsecuregcm/filters/RequestStatisticsFilter.java
new file mode 100644
index 000000000..078fb6a1b
--- /dev/null
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/filters/RequestStatisticsFilter.java
@@ -0,0 +1,78 @@
+/*
+ * Copyright 2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.filters;
+
+import static com.codahale.metrics.MetricRegistry.name;
+import static java.util.Objects.requireNonNull;
+
+import com.google.common.net.HttpHeaders;
+import com.google.common.net.InetAddresses;
+import io.micrometer.core.instrument.Metrics;
+import java.io.IOException;
+import java.net.Inet4Address;
+import java.net.Inet6Address;
+import java.net.InetAddress;
+import javax.annotation.Nonnull;
+import javax.ws.rs.container.ContainerRequestContext;
+import javax.ws.rs.container.ContainerRequestFilter;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.metrics.TrafficSource;
+import org.whispersystems.textsecuregcm.util.HeaderUtils;
+
+public class RequestStatisticsFilter implements ContainerRequestFilter {
+
+  private static final Logger logger = LoggerFactory.getLogger(RequestStatisticsFilter.class);
+
+  private static final String CONTENT_LENGTH_DISTRIBUTION_NAME = name(RequestStatisticsFilter.class, "contentLength");
+
+  private static final String IP_VERSION_METRIC = name(RequestStatisticsFilter.class, "ipVersion");
+
+  private static final String TRAFFIC_SOURCE_TAG = "trafficSource";
+
+  private static final String IP_VERSION_TAG = "ipVersion";
+
+  @Nonnull
+  private final String trafficSourceTag;
+
+
+  public RequestStatisticsFilter(@Nonnull final TrafficSource trafficeSource) {
+    this.trafficSourceTag = requireNonNull(trafficeSource).name().toLowerCase();
+  }
+
+  @Override
+  public void filter(final ContainerRequestContext requestContext) throws IOException {
+    try {
+      Metrics.summary(CONTENT_LENGTH_DISTRIBUTION_NAME, TRAFFIC_SOURCE_TAG, trafficSourceTag)
+          .record(requestContext.getLength());
+      Metrics.counter(IP_VERSION_METRIC, TRAFFIC_SOURCE_TAG, trafficSourceTag, IP_VERSION_TAG, resolveIpVersion(requestContext))
+          .increment();
+    } catch (final Exception e) {
+      logger.warn("Error recording request statistics", e);
+    }
+  }
+
+  @Nonnull
+  private static String resolveIpVersion(@Nonnull final ContainerRequestContext ctx) {
+    return HeaderUtils.getMostRecentProxy(ctx.getHeaderString(HttpHeaders.X_FORWARDED_FOR))
+        .map(ipString -> {
+              try {
+                //noinspection UnstableApiUsage
+                final InetAddress addr = InetAddresses.forString(ipString);
+                if (addr instanceof Inet4Address) {
+                  return "IPv4";
+                }
+                if (addr instanceof Inet6Address) {
+                  return "IPv6";
+                }
+              } catch (IllegalArgumentException e) {
+                // ignore illegal argument exception
+              }
+              return null;
+            })
+        .orElse("unresolved");
+  }
+}
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/filters/ContentLengthFilterTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/filters/RequestStatisticsFilterTest.java
similarity index 70%
rename from service/src/test/java/org/whispersystems/textsecuregcm/filters/ContentLengthFilterTest.java
rename to service/src/test/java/org/whispersystems/textsecuregcm/filters/RequestStatisticsFilterTest.java
index 7a9b2ca09..32038cc8d 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/filters/ContentLengthFilterTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/filters/RequestStatisticsFilterTest.java
@@ -12,12 +12,12 @@ import javax.ws.rs.container.ContainerRequestContext;
 import org.junit.jupiter.api.Test;
 import org.whispersystems.textsecuregcm.metrics.TrafficSource;
 
-class ContentLengthFilterTest {
+class RequestStatisticsFilterTest {
 
   @Test
   void testFilter() throws Exception {
 
-    final ContentLengthFilter contentLengthFilter = new ContentLengthFilter(TrafficSource.WEBSOCKET);
+    final RequestStatisticsFilter requestStatisticsFilter = new RequestStatisticsFilter(TrafficSource.WEBSOCKET);
 
     final ContainerRequestContext requestContext = mock(ContainerRequestContext.class);
 
@@ -25,8 +25,8 @@ class ContentLengthFilterTest {
     when(requestContext.getLength()).thenReturn(Integer.MAX_VALUE);
     when(requestContext.getLength()).thenThrow(RuntimeException.class);
 
-    contentLengthFilter.filter(requestContext);
-    contentLengthFilter.filter(requestContext);
-    contentLengthFilter.filter(requestContext);
+    requestStatisticsFilter.filter(requestContext);
+    requestStatisticsFilter.filter(requestContext);
+    requestStatisticsFilter.filter(requestContext);
   }
 }