Add X-Signal-Timestamp response filter

Author: eager-signal

src/main/java/org/signal/storageservice/StorageService.java

@@ -5,9 +5,6 @@
 
 package org.signal.storageservice;
 
-import static com.codahale.metrics.MetricRegistry.name;
-
-import com.codahale.metrics.SharedMetricRegistries;
 import com.fasterxml.jackson.annotation.JsonAutoDetect;
 import com.fasterxml.jackson.annotation.PropertyAccessor;
 import com.fasterxml.jackson.databind.DeserializationFeature;
@@ -23,9 +20,6 @@
 import io.dropwizard.core.Application;
 import io.dropwizard.core.setup.Bootstrap;
 import io.dropwizard.core.setup.Environment;
-import io.micrometer.core.instrument.Metrics;
-import io.micrometer.core.instrument.Tags;
-import io.micrometer.datadog.DatadogMeterRegistry;
 import java.time.Clock;
 import java.util.Set;
 import org.signal.libsignal.zkgroup.ServerSecretParams;
@@ -41,14 +35,9 @@
 import org.signal.storageservice.controllers.HealthCheckController;
 import org.signal.storageservice.controllers.ReadinessController;
 import org.signal.storageservice.controllers.StorageController;
-import org.signal.storageservice.metrics.CpuUsageGauge;
-import org.signal.storageservice.metrics.FileDescriptorGauge;
-import org.signal.storageservice.metrics.FreeMemoryGauge;
+import org.signal.storageservice.filters.TimestampResponseFilter;
 import org.signal.storageservice.metrics.MetricsApplicationEventListener;
 import org.signal.storageservice.metrics.MetricsUtil;
-import org.signal.storageservice.metrics.NetworkReceivedGauge;
-import org.signal.storageservice.metrics.NetworkSentGauge;
-import org.signal.storageservice.metrics.StorageMetrics;
 import org.signal.storageservice.providers.CompletionExceptionMapper;
 import org.signal.storageservice.providers.InvalidProtocolBufferExceptionMapper;
 import org.signal.storageservice.providers.ProtocolBufferMessageBodyProvider;
@@ -58,7 +47,6 @@
 import org.signal.storageservice.storage.GroupsManager;
 import org.signal.storageservice.storage.StorageManager;
 import org.signal.storageservice.util.UncaughtExceptionHandler;
-import org.signal.storageservice.util.HostSupplier;
 import org.signal.storageservice.util.logging.LoggingUnhandledExceptionMapper;
 
 public class StorageService extends Application<StorageServiceConfiguration> {
@@ -105,6 +93,8 @@ public void run(StorageServiceConfiguration config, Environment environment) thr
     environment.jersey().register(new PolymorphicAuthDynamicFeature<>(ImmutableMap.of(User.class, userAuthFilter, GroupUser.class, groupUserAuthFilter)));
     environment.jersey().register(new PolymorphicAuthValueFactoryProvider.Binder<>(ImmutableSet.of(User.class, GroupUser.class)));
 
+    environment.jersey().register(new TimestampResponseFilter(Clock.systemUTC()));
+
     environment.jersey().register(new HealthCheckController());
     environment.jersey().register(new ReadinessController(bigtableDataClient,
         Set.of(config.getBigTableConfiguration().getGroupsTableId(),

src/main/java/org/signal/storageservice/filters/TimestampResponseFilter.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.signal.storageservice.filters;
+
+import java.io.IOException;
+import java.time.Clock;
+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.signal.storageservice.util.HeaderUtils;
+
+/**
+ * Injects a timestamp header into all outbound responses.
+ */
+public class TimestampResponseFilter implements Filter, ContainerResponseFilter {
+
+  private final Clock clock;
+
+  public TimestampResponseFilter(final Clock clock) {
+    this.clock = clock;
+  }
+
+  @Override
+  public void doFilter(final ServletRequest request, final ServletResponse response, final FilterChain chain)
+      throws ServletException, IOException {
+
+    if (response instanceof HttpServletResponse httpServletResponse) {
+      httpServletResponse.setHeader(HeaderUtils.TIMESTAMP_HEADER, String.valueOf(clock.millis()));
+    }
+
+    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, String.valueOf(clock.millis()));
+  }
+}

src/main/java/org/signal/storageservice/util/HeaderUtils.java

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.signal.storageservice.util;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public final class HeaderUtils {
+
+  private static final Logger logger = LoggerFactory.getLogger(HeaderUtils.class);
+
+  public static final String TIMESTAMP_HEADER = "X-Signal-Timestamp";
+  
+  private HeaderUtils() {
+    // utility class
+  }
+
+}

src/test/java/org/signal/storageservice/filters/TimestampResponseFilterTest.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.signal.storageservice.filters;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.time.Clock;
+import java.time.Instant;
+import java.time.ZoneId;
+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.junit.jupiter.api.Test;
+import org.signal.storageservice.util.HeaderUtils;
+
+class TimestampResponseFilterTest {
+
+  private static final long EPOCH_MILLIS = 1738182156000L;
+
+  private static final Clock CLOCK = Clock.fixed(Instant.ofEpochMilli(EPOCH_MILLIS), ZoneId.systemDefault());
+
+  @Test
+  void testJerseyFilter() {
+    final ContainerRequestContext requestContext = mock(ContainerRequestContext.class);
+    final ContainerResponseContext responseContext = mock(ContainerResponseContext.class);
+    final MultivaluedMap<String, Object> headers = org.glassfish.jersey.message.internal.HeaderUtils.createOutbound();
+    when(responseContext.getHeaders()).thenReturn(headers);
+
+    new TimestampResponseFilter(CLOCK).filter(requestContext, responseContext);
+
+    assertTrue(headers.containsKey(HeaderUtils.TIMESTAMP_HEADER));
+    assertEquals(1, headers.get(HeaderUtils.TIMESTAMP_HEADER).size());
+    assertEquals(String.valueOf(EPOCH_MILLIS), headers.get(HeaderUtils.TIMESTAMP_HEADER).get(0));
+  }
+
+  @Test
+  void testServletFilter() throws Exception {
+    final HttpServletRequest request = mock(HttpServletRequest.class);
+    final HttpServletResponse response = mock(HttpServletResponse.class);
+
+    new TimestampResponseFilter(CLOCK).doFilter(request, response, mock(FilterChain.class));
+
+    verify(response).setHeader(eq(HeaderUtils.TIMESTAMP_HEADER), eq(String.valueOf(EPOCH_MILLIS)));
+  }
+}