From 738ec2a38e72c11ea814a1af97c234c0bfb09788 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Tue, 9 Mar 2021 16:51:54 -0500
Subject: [PATCH] Use ForwardedIpUtil everywhere we're handling X-Forwarded-For
 values.

---
 .../controllers/AccountController.java             |  6 ++----
 .../controllers/MessageController.java             | 14 ++++++++++----
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
index 501cfb9fd..577dd2dc2 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
@@ -71,6 +71,7 @@ import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.PendingAccountsManager;
 import org.whispersystems.textsecuregcm.storage.UsernamesManager;
 import org.whispersystems.textsecuregcm.util.Constants;
+import org.whispersystems.textsecuregcm.util.ForwardedIpUtil;
 import org.whispersystems.textsecuregcm.util.Hex;
 import org.whispersystems.textsecuregcm.util.Util;
 import org.whispersystems.textsecuregcm.util.VerificationCode;
@@ -198,10 +199,7 @@ public class AccountController {
       transport = "voice";
     }
 
-    String requester = Arrays.stream(forwardedFor.split(","))
-                             .map(String::trim)
-                             .reduce((a, b) -> b)
-                             .orElseThrow();
+    String requester = ForwardedIpUtil.getMostRecentProxy(forwardedFor).orElseThrow();
 
     Optional<StoredVerificationCode> storedChallenge = pendingAccounts.getCodeForNumber(number);
     CaptchaRequirement               requirement     = requiresCaptcha(number, transport, forwardedFor, requester, captcha, storedChallenge, pushChallenge);
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 5612b936a..648a6190a 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
@@ -74,6 +74,7 @@ import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.util.Base64;
 import org.whispersystems.textsecuregcm.util.Constants;
+import org.whispersystems.textsecuregcm.util.ForwardedIpUtil;
 import org.whispersystems.textsecuregcm.util.Util;
 import org.whispersystems.textsecuregcm.util.ua.UnrecognizedUserAgentException;
 import org.whispersystems.textsecuregcm.util.ua.UserAgentUtil;
@@ -265,7 +266,12 @@ public class MessageController {
 
           if (StringUtils.isAllBlank(masterDevice.getApnId(), masterDevice.getVoipApnId(), masterDevice.getGcmId()) || masterDevice.getUninstalledFeedbackTimestamp() > 0) {
             if (dynamicConfigurationManager.getConfiguration().getMessageRateConfiguration().getRateLimitedCountryCodes().contains(senderCountryCode)) {
-              if (dynamicConfigurationManager.getConfiguration().getMessageRateConfiguration().getRateLimitedHosts().contains(forwardedFor)) {
+
+              final boolean isRateLimitedHost = ForwardedIpUtil.getMostRecentProxy(forwardedFor)
+                  .map(proxy -> dynamicConfigurationManager.getConfiguration().getMessageRateConfiguration().getRateLimitedHosts().contains(proxy))
+                  .orElse(false);
+
+              if (isRateLimitedHost) {
                 return declineDelivery(messages, source.get(), destination.get());
               }
             }
@@ -566,15 +572,15 @@ public class MessageController {
   }
 
   @VisibleForTesting
-  void recordInternationalUnsealedSenderMetrics(final String senderIp, final String senderCountryCode, final String destinationNumber) {
-    {
+  void recordInternationalUnsealedSenderMetrics(final String forwardedFor, final String senderCountryCode, final String destinationNumber) {
+    ForwardedIpUtil.getMostRecentProxy(forwardedFor).ifPresent(senderIp -> {
       final String destinationSetKey = getDestinationSetKey(senderIp);
       final String messageCountKey = getMessageCountKey(senderIp);
 
       recordInternationalUnsealedSenderMetricsScript.execute(
           List.of(destinationSetKey, messageCountKey),
           List.of(destinationNumber));
-    }
+    });
 
     Metrics.counter(INTERNATIONAL_UNSEALED_SENDER_COUNTER_NAME, SENDER_COUNTRY_TAG_NAME, senderCountryCode).increment();
   }