Refactor remote address/X-Forwarded-For handling

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -44,6 +44,7 @@ import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ThreadPoolExecutor;
 import javax.servlet.DispatcherType;
+import javax.servlet.Filter;
 import javax.servlet.FilterRegistration;
 import javax.servlet.ServletRegistration;
 import org.eclipse.jetty.servlets.CrossOriginFilter;
@@ -115,6 +116,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.RemoteAddressFilter;
 import org.whispersystems.textsecuregcm.filters.RemoteDeprecationFilter;
 import org.whispersystems.textsecuregcm.filters.RequestStatisticsFilter;
 import org.whispersystems.textsecuregcm.filters.TimestampResponseFilter;
@@ -212,8 +214,8 @@ import org.whispersystems.textsecuregcm.util.DynamoDbFromConfig;
 import org.whispersystems.textsecuregcm.util.ManagedAwsCrt;
 import org.whispersystems.textsecuregcm.util.SystemMapper;
 import org.whispersystems.textsecuregcm.util.UsernameHashZkProofVerifier;
-import org.whispersystems.textsecuregcm.util.VirtualThreadPinEventMonitor;
 import org.whispersystems.textsecuregcm.util.VirtualExecutorServiceProvider;
+import org.whispersystems.textsecuregcm.util.VirtualThreadPinEventMonitor;
 import org.whispersystems.textsecuregcm.util.logging.LoggingUnhandledExceptionMapper;
 import org.whispersystems.textsecuregcm.util.logging.UncaughtExceptionHandler;
 import org.whispersystems.textsecuregcm.websocket.AuthenticatedConnectListener;
@@ -718,10 +720,16 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
                 config.getBadges(), asyncCdnS3Client, profileCdnPolicyGenerator, profileCdnPolicySigner, profileBadgeConverter, rateLimiters, zkProfileOperations, config.getCdnConfiguration().bucket()), basicCredentialAuthenticationInterceptor))
         .addService(new ProfileAnonymousGrpcService(accountsManager, profilesManager, profileBadgeConverter, zkProfileOperations));
 
-    RemoteDeprecationFilter remoteDeprecationFilter = new RemoteDeprecationFilter(dynamicConfigurationManager);
-    environment.servlets()
-        .addFilter("RemoteDeprecationFilter", remoteDeprecationFilter)
-        .addMappingForUrlPatterns(EnumSet.of(DispatcherType.REQUEST), false, "/*");
+    final List<Filter> filters = new ArrayList<>();
+    final RemoteDeprecationFilter remoteDeprecationFilter = new RemoteDeprecationFilter(dynamicConfigurationManager);
+    filters.add(remoteDeprecationFilter);
+    filters.add(new RemoteAddressFilter(useRemoteAddress));
+
+    for (Filter filter : filters) {
+      environment.servlets()
+          .addFilter(filter.getClass().getSimpleName(), filter)
+          .addMappingForUrlPatterns(EnumSet.of(DispatcherType.REQUEST), false, "/*");
+    }
 
     // Note: interceptors run in the reverse order they are added; the remote deprecation filter
     // depends on the user-agent context so it has to come first here!
@@ -832,7 +840,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         new CertificateController(new CertificateGenerator(config.getDeliveryCertificate().certificate().value(),
             config.getDeliveryCertificate().ecPrivateKey(), config.getDeliveryCertificate().expiresDays()),
             zkAuthOperations, callingGenericZkSecretParams, clock),
-        new ChallengeController(rateLimitChallengeManager, useRemoteAddress),
+        new ChallengeController(rateLimitChallengeManager),
         new DeviceController(config.getLinkDeviceSecretConfiguration().secret().value(), accountsManager,
             rateLimiters, rateLimitersCluster, config.getMaxDevices(), clock),
         new DirectoryV2Controller(directoryV2CredentialsGenerator),
@@ -859,7 +867,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
             config.getCdnConfiguration().bucket()),
         new VerificationController(registrationServiceClient, new VerificationSessionManager(verificationSessions),
             pushNotificationManager, registrationCaptchaManager, registrationRecoveryPasswordsManager, rateLimiters,
-            accountsManager, useRemoteAddress, dynamicConfigurationManager, clock)
+            accountsManager, dynamicConfigurationManager, clock)
     );
     if (config.getSubscription() != null && config.getOneTimeDonations() != null) {
       commonControllers.add(new SubscriptionController(clock, config.getSubscription(), config.getOneTimeDonations(),
@@ -890,9 +898,11 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     JettyWebSocketServletContainerInitializer.configure(environment.getApplicationContext(), null);
 
     WebSocketResourceProviderFactory<AuthenticatedAccount> webSocketServlet = new WebSocketResourceProviderFactory<>(
-        webSocketEnvironment, AuthenticatedAccount.class, config.getWebSocketConfiguration());
+        webSocketEnvironment, AuthenticatedAccount.class, config.getWebSocketConfiguration(),
+        RemoteAddressFilter.REMOTE_ADDRESS_ATTRIBUTE_NAME);
     WebSocketResourceProviderFactory<AuthenticatedAccount> provisioningServlet = new WebSocketResourceProviderFactory<>(
-        provisioningEnvironment, AuthenticatedAccount.class, config.getWebSocketConfiguration());
+        provisioningEnvironment, AuthenticatedAccount.class, config.getWebSocketConfiguration(),
+        RemoteAddressFilter.REMOTE_ADDRESS_ATTRIBUTE_NAME);
 
     ServletRegistration.Dynamic websocket = environment.servlets().addServlet("WebSocket", webSocketServlet);
     ServletRegistration.Dynamic provisioning = environment.servlets().addServlet("Provisioning", provisioningServlet);

service/src/main/java/org/whispersystems/textsecuregcm/controllers/ChallengeController.java

@@ -19,15 +19,14 @@ import io.swagger.v3.oas.annotations.parameters.RequestBody;
 import io.swagger.v3.oas.annotations.responses.ApiResponse;
 import io.swagger.v3.oas.annotations.tags.Tag;
 import java.io.IOException;
-import javax.servlet.http.HttpServletRequest;
 import javax.validation.Valid;
-import javax.ws.rs.BadRequestException;
 import javax.ws.rs.Consumes;
 import javax.ws.rs.HeaderParam;
 import javax.ws.rs.POST;
 import javax.ws.rs.PUT;
 import javax.ws.rs.Path;
 import javax.ws.rs.Produces;
+import javax.ws.rs.container.ContainerRequestContext;
 import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
@@ -35,6 +34,7 @@ import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
 import org.whispersystems.textsecuregcm.entities.AnswerChallengeRequest;
 import org.whispersystems.textsecuregcm.entities.AnswerPushChallengeRequest;
 import org.whispersystems.textsecuregcm.entities.AnswerRecaptchaChallengeRequest;
+import org.whispersystems.textsecuregcm.filters.RemoteAddressFilter;
 import org.whispersystems.textsecuregcm.limits.RateLimitChallengeManager;
 import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
 import org.whispersystems.textsecuregcm.push.NotPushRegisteredException;
@@ -42,8 +42,6 @@ import org.whispersystems.textsecuregcm.spam.Extract;
 import org.whispersystems.textsecuregcm.spam.FilterSpam;
 import org.whispersystems.textsecuregcm.spam.PushChallengeConfig;
 import org.whispersystems.textsecuregcm.spam.ScoreThreshold;
-import org.whispersystems.textsecuregcm.util.HeaderUtils;
-import org.whispersystems.textsecuregcm.util.HttpServletRequestUtil;
 
 @Path("/v1/challenge")
 @Tag(name = "Challenge")
@@ -51,15 +49,12 @@ import org.whispersystems.textsecuregcm.util.HttpServletRequestUtil;
 public class ChallengeController {
 
   private final RateLimitChallengeManager rateLimitChallengeManager;
-  private final boolean useRemoteAddress;
 
   private static final String CHALLENGE_RESPONSE_COUNTER_NAME = name(ChallengeController.class, "challengeResponse");
   private static final String CHALLENGE_TYPE_TAG = "type";
 
-  public ChallengeController(final RateLimitChallengeManager rateLimitChallengeManager,
-      final boolean useRemoteAddress) {
+  public ChallengeController(final RateLimitChallengeManager rateLimitChallengeManager) {
     this.rateLimitChallengeManager = rateLimitChallengeManager;
-    this.useRemoteAddress = useRemoteAddress;
   }
 
   @PUT
@@ -84,8 +79,7 @@ public class ChallengeController {
       description = "If present, an positive integer indicating the number of seconds before a subsequent attempt could succeed"))
   public Response handleChallengeResponse(@Auth final AuthenticatedAccount auth,
       @Valid final AnswerChallengeRequest answerRequest,
-      @HeaderParam(HttpHeaders.X_FORWARDED_FOR) final String forwardedFor,
-      @Context HttpServletRequest request,
+      @Context ContainerRequestContext requestContext,
       @HeaderParam(HttpHeaders.USER_AGENT) final String userAgent,
       @Extract final ScoreThreshold captchaScoreThreshold,
       @Extract final PushChallengeConfig pushChallengeConfig) throws RateLimitExceededException, IOException {
@@ -103,9 +97,8 @@ public class ChallengeController {
       } else if (answerRequest instanceof AnswerRecaptchaChallengeRequest recaptchaChallengeRequest) {
         tags = tags.and(CHALLENGE_TYPE_TAG, "recaptcha");
 
-        final String remoteAddress = useRemoteAddress
-            ? HttpServletRequestUtil.getRemoteAddress(request)
-            : HeaderUtils.getMostRecentProxy(forwardedFor).orElseThrow(BadRequestException::new);
+        final String remoteAddress = (String) requestContext.getProperty(
+            RemoteAddressFilter.REMOTE_ADDRESS_ATTRIBUTE_NAME);
         boolean success = rateLimitChallengeManager.answerRecaptchaChallenge(
             auth.getAccount(),
             recaptchaChallengeRequest.getCaptcha(),

service/src/main/java/org/whispersystems/textsecuregcm/controllers/VerificationController.java

@@ -31,7 +31,6 @@ import java.util.Optional;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.TimeUnit;
-import javax.servlet.http.HttpServletRequest;
 import javax.validation.Valid;
 import javax.validation.constraints.NotNull;
 import javax.ws.rs.BadRequestException;
@@ -49,6 +48,7 @@ import javax.ws.rs.PathParam;
 import javax.ws.rs.Produces;
 import javax.ws.rs.ServerErrorException;
 import javax.ws.rs.WebApplicationException;
+import javax.ws.rs.container.ContainerRequestContext;
 import javax.ws.rs.core.Context;
 import javax.ws.rs.core.HttpHeaders;
 import javax.ws.rs.core.MediaType;
@@ -66,6 +66,7 @@ import org.whispersystems.textsecuregcm.entities.SubmitVerificationCodeRequest;
 import org.whispersystems.textsecuregcm.entities.UpdateVerificationSessionRequest;
 import org.whispersystems.textsecuregcm.entities.VerificationCodeRequest;
 import org.whispersystems.textsecuregcm.entities.VerificationSessionResponse;
+import org.whispersystems.textsecuregcm.filters.RemoteAddressFilter;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
@@ -88,8 +89,6 @@ import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 import org.whispersystems.textsecuregcm.storage.RegistrationRecoveryPasswordsManager;
 import org.whispersystems.textsecuregcm.storage.VerificationSessionManager;
 import org.whispersystems.textsecuregcm.util.ExceptionUtils;
-import org.whispersystems.textsecuregcm.util.HeaderUtils;
-import org.whispersystems.textsecuregcm.util.HttpServletRequestUtil;
 import org.whispersystems.textsecuregcm.util.Pair;
 import org.whispersystems.textsecuregcm.util.Util;
 
@@ -123,7 +122,6 @@ public class VerificationController {
   private final RateLimiters rateLimiters;
   private final AccountsManager accountsManager;
 
-  private final boolean useRemoteAddress;
   private final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager;
   private final Clock clock;
 
@@ -134,7 +132,6 @@ public class VerificationController {
       final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager,
       final RateLimiters rateLimiters,
       final AccountsManager accountsManager,
-      final boolean useRemoteAddress,
       final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager,
       final Clock clock) {
     this.registrationServiceClient = registrationServiceClient;
@@ -144,7 +141,6 @@ public class VerificationController {
     this.registrationRecoveryPasswordsManager = registrationRecoveryPasswordsManager;
     this.rateLimiters = rateLimiters;
     this.accountsManager = accountsManager;
-    this.useRemoteAddress = useRemoteAddress;
     this.dynamicConfigurationManager = dynamicConfigurationManager;
     this.clock = clock;
   }
@@ -205,16 +201,13 @@ public class VerificationController {
   @Consumes(MediaType.APPLICATION_JSON)
   @Produces(MediaType.APPLICATION_JSON)
   public VerificationSessionResponse updateSession(@PathParam("sessionId") final String encodedSessionId,
-      @HeaderParam(com.google.common.net.HttpHeaders.X_FORWARDED_FOR) String forwardedFor,
       @HeaderParam(HttpHeaders.USER_AGENT) final String userAgent,
-      @Context HttpServletRequest request,
+      @Context ContainerRequestContext requestContext,
       @NotNull @Valid final UpdateVerificationSessionRequest updateVerificationSessionRequest,
       @NotNull @Extract final ScoreThreshold scoreThreshold,
       @NotNull @Extract final SenderOverride senderOverride) {
 
-    final String sourceHost = useRemoteAddress
-        ? HttpServletRequestUtil.getRemoteAddress(request)
-        : HeaderUtils.getMostRecentProxy(forwardedFor).orElseThrow();
+    final String sourceHost = (String) requestContext.getProperty(RemoteAddressFilter.REMOTE_ADDRESS_ATTRIBUTE_NAME);
 
     final Pair<String, PushNotification.TokenType> pushTokenAndType = validateAndExtractPushToken(
         updateVerificationSessionRequest);

service/src/main/java/org/whispersystems/textsecuregcm/filters/RemoteAddressFilter.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.filters;
+
+import javax.servlet.Filter;
+import javax.servlet.FilterChain;
+import javax.servlet.ServletException;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletRequest;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.util.HeaderUtils;
+import org.whispersystems.textsecuregcm.util.HttpServletRequestUtil;
+import java.io.IOException;
+
+/**
+ * Sets a {@link HttpServletRequest} attribute (that will also be available as a
+ * {@link javax.ws.rs.container.ContainerRequestContext} property) with the remote address of the connection, using
+ * either the {@link HttpServletRequest#getRemoteAddr()} or the {@code X-Forwarded-For} HTTP header value, depending on
+ * whether {@link #preferRemoteAddress} is {@code true}.
+ */
+public class RemoteAddressFilter implements Filter {
+
+  public static final String REMOTE_ADDRESS_ATTRIBUTE_NAME = RemoteAddressFilter.class.getName() + ".remoteAddress";
+  private static final Logger logger = LoggerFactory.getLogger(RemoteAddressFilter.class);
+
+  private final boolean preferRemoteAddress;
+
+
+  public RemoteAddressFilter(boolean preferRemoteAddress) {
+    this.preferRemoteAddress = preferRemoteAddress;
+  }
+
+  @Override
+  public void doFilter(final ServletRequest request, final ServletResponse response, final FilterChain chain)
+      throws ServletException, IOException {
+
+    if (request instanceof HttpServletRequest httpServletRequest) {
+
+      final String remoteAddress;
+
+      if (preferRemoteAddress) {
+        remoteAddress = HttpServletRequestUtil.getRemoteAddress(httpServletRequest);
+      } else {
+        final String forwardedFor = httpServletRequest.getHeader(com.google.common.net.HttpHeaders.X_FORWARDED_FOR);
+        remoteAddress = HeaderUtils.getMostRecentProxy(forwardedFor)
+            .orElseGet(() -> HttpServletRequestUtil.getRemoteAddress(httpServletRequest));
+      }
+
+      request.setAttribute(REMOTE_ADDRESS_ATTRIBUTE_NAME, remoteAddress);
+
+    } else {
+      logger.warn("request was of unexpected type: {}", request.getClass());
+    }
+
+    chain.doFilter(request, response);
+  }
+}

service/src/main/java/org/whispersystems/textsecuregcm/filters/RemoteDeprecationFilter.java

@@ -7,15 +7,12 @@ package org.whispersystems.textsecuregcm.filters;
 
 import static com.codahale.metrics.MetricRegistry.name;
 
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.net.HttpHeaders;
 import com.vdurmont.semver4j.Semver;
-
 import io.grpc.Metadata;
 import io.grpc.ServerCall;
 import io.grpc.ServerCallHandler;
 import io.grpc.ServerInterceptor;
-import io.grpc.Status;
 import io.micrometer.core.instrument.Metrics;
 import java.io.IOException;
 import java.util.Map;
@@ -30,7 +27,6 @@ import javax.servlet.http.HttpServletResponse;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicRemoteDeprecationConfiguration;
 import org.whispersystems.textsecuregcm.grpc.StatusConstants;
-import org.whispersystems.textsecuregcm.grpc.UserAgentInterceptor;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 import org.whispersystems.textsecuregcm.util.ua.ClientPlatform;
 import org.whispersystems.textsecuregcm.util.ua.UnrecognizedUserAgentException;

service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimitByIpFilter.java

@@ -8,33 +8,25 @@ package org.whispersystems.textsecuregcm.limits;
 import static java.util.Objects.requireNonNull;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.net.HttpHeaders;
 import java.io.IOException;
 import java.time.Duration;
 import java.util.Optional;
-import javax.inject.Provider;
-import javax.servlet.http.HttpServletRequest;
 import javax.ws.rs.ClientErrorException;
 import javax.ws.rs.container.ContainerRequestContext;
 import javax.ws.rs.container.ContainerRequestFilter;
-import javax.ws.rs.core.Context;
 import javax.ws.rs.core.Response;
 import javax.ws.rs.ext.ExceptionMapper;
 import org.glassfish.jersey.server.ExtendedUriInfo;
 import org.slf4j.Logger;
 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.util.HeaderUtils;
-import org.whispersystems.textsecuregcm.util.HttpServletRequestUtil;
 
 public class RateLimitByIpFilter implements ContainerRequestFilter {
 
   private static final Logger logger = LoggerFactory.getLogger(RateLimitByIpFilter.class);
 
-  @Context
-  private Provider<HttpServletRequest> httpServletRequestProvider;
-
   @VisibleForTesting
   static final RateLimitExceededException INVALID_HEADER_EXCEPTION = new RateLimitExceededException(Duration.ofHours(1),
       true);
@@ -42,12 +34,10 @@ public class RateLimitByIpFilter implements ContainerRequestFilter {
   private static final ExceptionMapper<RateLimitExceededException> EXCEPTION_MAPPER = new RateLimitExceededExceptionMapper();
 
   private final RateLimiters rateLimiters;
-  private final boolean useRemoteAddress;
 
 
-  public RateLimitByIpFilter(final RateLimiters rateLimiters, final boolean useRemoteAddress) {
+  public RateLimitByIpFilter(final RateLimiters rateLimiters) {
     this.rateLimiters = requireNonNull(rateLimiters);
-    this.useRemoteAddress = useRemoteAddress;
   }
 
   @Override
@@ -70,18 +60,14 @@ public class RateLimitByIpFilter implements ContainerRequestFilter {
     final RateLimiters.For handle = annotation.value();
 
     try {
-      final String xffHeader = requestContext.getHeaders().getFirst(HttpHeaders.X_FORWARDED_FOR);
-      final Optional<String> remoteAddress = useRemoteAddress
-          ? Optional.of(HttpServletRequestUtil.getRemoteAddress(httpServletRequestProvider.get()))
-          : Optional.ofNullable(xffHeader)
-              .flatMap(HeaderUtils::getMostRecentProxy);
+      final Optional<String> remoteAddress = Optional.ofNullable(
+          (String) requestContext.getProperty(RemoteAddressFilter.REMOTE_ADDRESS_ATTRIBUTE_NAME));
 
-      // checking if we failed to extract the most recent IP from the X-Forwarded-For header
-      // for any reason
+      // checking if we failed to extract the most recent IP for any reason
       if (remoteAddress.isEmpty()) {
         // checking if annotation is configured to fail when the most recent IP is not resolved
         if (annotation.failOnUnresolvedIp()) {
-          logger.error("Missing/bad X-Forwarded-For: {}", xffHeader);
+          logger.error("Remote address was null");
           throw INVALID_HEADER_EXCEPTION;
         }
         // otherwise, allow request