Add timestamp header to all responses

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

@@ -113,7 +113,6 @@ import org.whispersystems.textsecuregcm.controllers.ArtController;
 import org.whispersystems.textsecuregcm.controllers.AttachmentControllerV2;
 import org.whispersystems.textsecuregcm.controllers.AttachmentControllerV3;
 import org.whispersystems.textsecuregcm.controllers.AttachmentControllerV4;
-import org.whispersystems.textsecuregcm.controllers.OneTimeDonationController;
 import org.whispersystems.textsecuregcm.controllers.CallLinkController;
 import org.whispersystems.textsecuregcm.controllers.CallRoutingController;
 import org.whispersystems.textsecuregcm.controllers.CertificateController;
@@ -125,6 +124,7 @@ import org.whispersystems.textsecuregcm.controllers.KeepAliveController;
 import org.whispersystems.textsecuregcm.controllers.KeyTransparencyController;
 import org.whispersystems.textsecuregcm.controllers.KeysController;
 import org.whispersystems.textsecuregcm.controllers.MessageController;
+import org.whispersystems.textsecuregcm.controllers.OneTimeDonationController;
 import org.whispersystems.textsecuregcm.controllers.PaymentsController;
 import org.whispersystems.textsecuregcm.controllers.ProfileController;
 import org.whispersystems.textsecuregcm.controllers.ProvisioningController;
@@ -192,12 +192,12 @@ import org.whispersystems.textsecuregcm.metrics.TrafficSource;
 import org.whispersystems.textsecuregcm.providers.MultiRecipientMessageProvider;
 import org.whispersystems.textsecuregcm.providers.RedisClusterHealthCheck;
 import org.whispersystems.textsecuregcm.push.APNSender;
-import org.whispersystems.textsecuregcm.push.PushNotificationScheduler;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
 import org.whispersystems.textsecuregcm.push.FcmSender;
 import org.whispersystems.textsecuregcm.push.MessageSender;
 import org.whispersystems.textsecuregcm.push.ProvisioningManager;
 import org.whispersystems.textsecuregcm.push.PushNotificationManager;
+import org.whispersystems.textsecuregcm.push.PushNotificationScheduler;
 import org.whispersystems.textsecuregcm.push.ReceiptSender;
 import org.whispersystems.textsecuregcm.redis.ConnectionEventLogger;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
@@ -959,6 +959,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     final List<Filter> filters = new ArrayList<>();
     filters.add(remoteDeprecationFilter);
     filters.add(new RemoteAddressFilter());
+    filters.add(new TimestampResponseFilter());
 
     for (Filter filter : filters) {
       environment.servlets()
@@ -1013,7 +1014,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     webSocketEnvironment.jersey().register(MultiRecipientMessageProvider.class);
     webSocketEnvironment.jersey().register(new MetricsApplicationEventListener(TrafficSource.WEBSOCKET, clientReleaseManager));
     webSocketEnvironment.jersey().register(new KeepAliveController(clientPresenceManager));
-
+    webSocketEnvironment.jersey().register(new TimestampResponseFilter());
 
     final List<SpamFilter> spamFilters = ServiceLoader.load(SpamFilter.class)
         .stream()
@@ -1146,6 +1147,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     provisioningEnvironment.setConnectListener(new ProvisioningConnectListener(provisioningManager));
     provisioningEnvironment.jersey().register(new MetricsApplicationEventListener(TrafficSource.WEBSOCKET, clientReleaseManager));
     provisioningEnvironment.jersey().register(new KeepAliveController(clientPresenceManager));
+    provisioningEnvironment.jersey().register(new TimestampResponseFilter());
 
     registerCorsFilter(environment);
     registerExceptionMappers(environment, webSocketEnvironment, provisioningEnvironment);

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

@@ -5,18 +5,37 @@
 
 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.HttpServletResponse;
 import javax.ws.rs.container.ContainerRequestContext;
 import javax.ws.rs.container.ContainerResponseContext;
 import javax.ws.rs.container.ContainerResponseFilter;
 import org.whispersystems.textsecuregcm.util.HeaderUtils;
+import java.io.IOException;
 
 /**
  * Injects a timestamp header into all outbound responses.
  */
-public class TimestampResponseFilter implements ContainerResponseFilter {
+public class TimestampResponseFilter implements Filter, ContainerResponseFilter {
 
-    @Override
+  @Override
-    public void filter(final ContainerRequestContext requestContext, final ContainerResponseContext responseContext) {
+  public void doFilter(final ServletRequest request, final ServletResponse response, final FilterChain chain)
-        responseContext.getHeaders().add(HeaderUtils.TIMESTAMP_HEADER, System.currentTimeMillis());
+      throws ServletException, IOException {
+
+    if (response instanceof HttpServletResponse httpServletResponse) {
+      httpServletResponse.setHeader(HeaderUtils.TIMESTAMP_HEADER, String.valueOf(System.currentTimeMillis()));
     }
+
+    chain.doFilter(request, response);
+  }
+
+  @Override
+  public void filter(ContainerRequestContext requestContext, ContainerResponseContext responseContext) {
+    // not using add() - it's ok to overwrite any existing header, and we don't want a multi-value
+    responseContext.getHeaders().putSingle(HeaderUtils.TIMESTAMP_HEADER, System.currentTimeMillis());
+  }
 }

service/src/test/java/org/whispersystems/textsecuregcm/WhisperServerServiceTest.java

@@ -13,9 +13,11 @@ import io.dropwizard.testing.junit5.DropwizardAppExtension;
 import io.dropwizard.testing.junit5.DropwizardExtensionsSupport;
 import io.dropwizard.util.Resources;
 import java.net.URI;
+import java.util.List;
 import java.util.Map;
 import javax.ws.rs.client.Client;
 import javax.ws.rs.core.Response;
+import org.eclipse.jetty.websocket.api.Session;
 import org.eclipse.jetty.websocket.client.WebSocketClient;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
@@ -27,6 +29,7 @@ import org.whispersystems.textsecuregcm.storage.DynamoDbExtension;
 import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema;
 import org.whispersystems.textsecuregcm.tests.util.TestWebsocketListener;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
+import org.whispersystems.textsecuregcm.util.HeaderUtils;
 import org.whispersystems.websocket.messages.WebSocketResponseMessage;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
@@ -92,15 +95,42 @@ class WhisperServerServiceTest {
   @Test
   void websocket() throws Exception {
     // test unauthenticated websocket
+    final long start = System.currentTimeMillis();
 
     final TestWebsocketListener testWebsocketListener = new TestWebsocketListener();
-    webSocketClient.connect(testWebsocketListener,
+
+    final Session session = webSocketClient.connect(testWebsocketListener,
             URI.create(String.format("ws://localhost:%d/v1/websocket/", EXTENSION.getLocalPort())))
         .join();
+    final long sessionTimestamp = Long.parseLong(session.getUpgradeResponse().getHeader(HeaderUtils.TIMESTAMP_HEADER));
+    assertTrue(sessionTimestamp >= start);
 
     final WebSocketResponseMessage keepAlive = testWebsocketListener.doGet("/v1/keepalive").join();
-
     assertEquals(200, keepAlive.getStatus());
+    final long keepAliveTimestamp = Long.parseLong(
+        keepAlive.getHeaders().get(HeaderUtils.TIMESTAMP_HEADER.toLowerCase()));
+    assertTrue(keepAliveTimestamp >= start);
+
+    final WebSocketResponseMessage whoami = testWebsocketListener.doGet("/v1/accounts/whoami").join();
+    assertEquals(401, whoami.getStatus());
+    final long whoamiTimestamp = Long.parseLong(whoami.getHeaders().get(HeaderUtils.TIMESTAMP_HEADER.toLowerCase()));
+    assertTrue(whoamiTimestamp >= start);
+  }
+
+  @Test
+  void rest() throws Exception {
+    // test unauthenticated rest
+    final long start = System.currentTimeMillis();
+
+    final Response whoami = EXTENSION.client().target(
+        "http://localhost:%d/v1/accounts/whoami".formatted(EXTENSION.getLocalPort())).request().get();
+
+    assertEquals(401, whoami.getStatus());
+    final List<Object> timestampValues = whoami.getHeaders().get(HeaderUtils.TIMESTAMP_HEADER.toLowerCase());
+    assertEquals(1, timestampValues.size());
+
+    final long whoamiTimestamp = Long.parseLong(timestampValues.getFirst().toString());
+    assertTrue(whoamiTimestamp >= start);
   }
 
   @Test
@@ -140,7 +170,6 @@ class WhisperServerServiceTest {
           .key(Map.of(numbers.hashKeyName(), numberAV))
           .build());
     }
-
   }
 
 }

service/src/test/java/org/whispersystems/textsecuregcm/filters/TimestampResponseFilterTest.java

@@ -6,28 +6,42 @@
 package org.whispersystems.textsecuregcm.filters;
 
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.matches;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import javax.servlet.FilterChain;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
 import javax.ws.rs.container.ContainerRequestContext;
 import javax.ws.rs.container.ContainerResponseContext;
 import javax.ws.rs.core.MultivaluedMap;
-import org.glassfish.jersey.message.internal.HeaderUtils;
 import org.junit.jupiter.api.Test;
+import org.whispersystems.textsecuregcm.util.HeaderUtils;
 
 class TimestampResponseFilterTest {
 
-    @Test
+  @Test
-    void testFilter() {
+  void testJerseyFilter() {
-        final ContainerRequestContext        requestContext  = mock(ContainerRequestContext.class);
+    final ContainerRequestContext requestContext = mock(ContainerRequestContext.class);
-        final ContainerResponseContext       responseContext = mock(ContainerResponseContext.class);
+    final ContainerResponseContext responseContext = mock(ContainerResponseContext.class);
+    final MultivaluedMap<String, Object> headers = org.glassfish.jersey.message.internal.HeaderUtils.createOutbound();
+    when(responseContext.getHeaders()).thenReturn(headers);
 
-        final MultivaluedMap<String, Object> headers         = HeaderUtils.createOutbound();
+    new TimestampResponseFilter().filter(requestContext, responseContext);
 
-        when(responseContext.getHeaders()).thenReturn(headers);
+    assertTrue(headers.containsKey(org.whispersystems.textsecuregcm.util.HeaderUtils.TIMESTAMP_HEADER));
+  }
 
-        new TimestampResponseFilter().filter(requestContext, responseContext);
+  @Test
+  void testServletFilter() throws Exception {
+    final HttpServletRequest request = mock(HttpServletRequest.class);
+    final HttpServletResponse response = mock(HttpServletResponse.class);
 
-        assertTrue(headers.containsKey(org.whispersystems.textsecuregcm.util.HeaderUtils.TIMESTAMP_HEADER));
+    new TimestampResponseFilter().doFilter(request, response, mock(FilterChain.class));
-    }
+
+    verify(response).setHeader(eq(HeaderUtils.TIMESTAMP_HEADER), matches("\\d{10,}"));
+  }
 }