Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -120,6 +121,21 @@ public Optional<String> 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<Jwt> getJwt() {
return this.getHeaderValue(RequestContextConstants.AUTHORIZATION_HEADER)
.flatMap(JWT_PARSER::fromAuthHeader);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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<String> TENANT_ID_METADATA_KEY =
Metadata.Key.of(TENANT_ID_HEADER_KEY, ASCII_STRING_MARSHALLER);
Expand All @@ -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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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());
}
}
Loading