diff --git a/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java b/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java index 1e67f15..acae38e 100644 --- a/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java +++ b/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java @@ -3,6 +3,7 @@ import static io.grpc.Metadata.ASCII_STRING_MARSHALLER; import static java.util.Objects.requireNonNull; import static org.hypertrace.core.grpcutils.context.RequestContextConstants.CACHE_MEANINGFUL_HEADERS; +import static org.hypertrace.core.grpcutils.context.RequestContextConstants.SUPPRESS_USER_TRACKING_HEADER_KEY; import static org.hypertrace.core.grpcutils.context.RequestContextConstants.TENANT_ID_HEADER_KEY; import com.google.common.collect.ListMultimap; @@ -120,6 +121,21 @@ public Optional getRequestId() { return this.getHeaderValue(RequestContextConstants.REQUEST_ID_HEADER_KEY); } + public boolean isUserTrackingSuppressed() { + return this.getHeaderValue(SUPPRESS_USER_TRACKING_HEADER_KEY) + .map(Boolean::parseBoolean) + .orElse(false); + } + + public RequestContext withUserTrackingSuppressed(boolean suppressed) { + this.removeHeader(SUPPRESS_USER_TRACKING_HEADER_KEY); + return this.put(SUPPRESS_USER_TRACKING_HEADER_KEY, String.valueOf(suppressed)); + } + + public RequestContext withUserTrackingSuppressed() { + return this.withUserTrackingSuppressed(true); + } + private Optional getJwt() { return this.getHeaderValue(RequestContextConstants.AUTHORIZATION_HEADER) .flatMap(JWT_PARSER::fromAuthHeader); diff --git a/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContextConstants.java b/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContextConstants.java index 1951043..9dfd39e 100644 --- a/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContextConstants.java +++ b/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContextConstants.java @@ -13,6 +13,7 @@ public class RequestContextConstants { public static final String TENANT_ID_HEADER_KEY = "x-tenant-id"; public static final String REQUEST_ID_HEADER_KEY = "request-id"; public static final String CONTEXT_ID_HEADER_KEY = "context-id"; + public static final String SUPPRESS_USER_TRACKING_HEADER_KEY = "x-suppress-user-tracking"; public static final Metadata.Key TENANT_ID_METADATA_KEY = Metadata.Key.of(TENANT_ID_HEADER_KEY, ASCII_STRING_MARSHALLER); @@ -24,6 +25,7 @@ public class RequestContextConstants { Set.of( TENANT_ID_HEADER_KEY, CONTEXT_ID_HEADER_KEY, + SUPPRESS_USER_TRACKING_HEADER_KEY, "X-B3-", "grpc-trace-bin", "traceparent", diff --git a/grpc-context-utils/src/test/java/org/hypertrace/core/grpcutils/context/RequestContextTest.java b/grpc-context-utils/src/test/java/org/hypertrace/core/grpcutils/context/RequestContextTest.java index 6119837..165a946 100644 --- a/grpc-context-utils/src/test/java/org/hypertrace/core/grpcutils/context/RequestContextTest.java +++ b/grpc-context-utils/src/test/java/org/hypertrace/core/grpcutils/context/RequestContextTest.java @@ -2,7 +2,9 @@ import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; import com.google.common.collect.ImmutableList; import io.grpc.Metadata; @@ -251,4 +253,86 @@ void backwardsCompatibilityForAdd() { assertEquals(Optional.of("f-v2"), requestContext.getHeaderValue("first")); assertEquals(List.of("f-v2"), requestContext.getAllHeaderValues("first")); } + + @Test + void testUserTrackingSuppressedDefault() { + RequestContext requestContext = new RequestContext(); + assertFalse(requestContext.isUserTrackingSuppressed()); + } + + @Test + void testUserTrackingSuppressedWithTrueValue() { + RequestContext requestContext = + new RequestContext().put(RequestContextConstants.SUPPRESS_USER_TRACKING_HEADER_KEY, "true"); + assertTrue(requestContext.isUserTrackingSuppressed()); + } + + @Test + void testUserTrackingSuppressedWithFalseValue() { + RequestContext requestContext = + new RequestContext() + .put(RequestContextConstants.SUPPRESS_USER_TRACKING_HEADER_KEY, "false"); + assertFalse(requestContext.isUserTrackingSuppressed()); + } + + @Test + void testWithUserTrackingSuppressedBoolean() { + RequestContext originalContext1 = new RequestContext(); + RequestContext suppressedContext = originalContext1.withUserTrackingSuppressed(true); + assertTrue(suppressedContext.isUserTrackingSuppressed()); + + RequestContext originalContext2 = new RequestContext(); + RequestContext unsuppressedContext = originalContext2.withUserTrackingSuppressed(false); + assertFalse(unsuppressedContext.isUserTrackingSuppressed()); + } + + @Test + void testWithUserTrackingSuppressedDefault() { + RequestContext originalContext = new RequestContext(); + + RequestContext suppressedContext = originalContext.withUserTrackingSuppressed(); + assertTrue(suppressedContext.isUserTrackingSuppressed()); + } + + @Test + void testWithUserTrackingSuppressedPreservesOtherHeaders() { + RequestContext originalContext = + new RequestContext() + .put(RequestContextConstants.TENANT_ID_HEADER_KEY, "tenant-123") + .put("other-header", "other-value"); + + RequestContext suppressedContext = originalContext.withUserTrackingSuppressed(); + + assertTrue(suppressedContext.isUserTrackingSuppressed()); + assertEquals(Optional.of("tenant-123"), suppressedContext.getTenantId()); + assertEquals(Optional.of("other-value"), suppressedContext.getHeaderValue("other-header")); + } + + @Test + void testWithUserTrackingSuppressedOverridesExisting() { + RequestContext originalContext = + new RequestContext() + .put(RequestContextConstants.SUPPRESS_USER_TRACKING_HEADER_KEY, "false"); + + assertFalse(originalContext.isUserTrackingSuppressed()); + + RequestContext suppressedContext = originalContext.withUserTrackingSuppressed(); + assertTrue(suppressedContext.isUserTrackingSuppressed()); + } + + @Test + void testUserTrackingSuppressedHeaderPropagation() { + Metadata metadata = new Metadata(); + metadata.put( + Metadata.Key.of( + RequestContextConstants.SUPPRESS_USER_TRACKING_HEADER_KEY, + Metadata.ASCII_STRING_MARSHALLER), + "true"); + metadata.put(RequestContextConstants.TENANT_ID_METADATA_KEY, "test-tenant"); + + RequestContext requestContext = RequestContext.fromMetadata(metadata); + + assertTrue(requestContext.isUserTrackingSuppressed()); + assertEquals(Optional.of("test-tenant"), requestContext.getTenantId()); + } }