Added test cases

Author: Mohit Agarwal

grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/RequestContextAsCreds.java

@@ -33,7 +33,7 @@ protected void applyRequestContext(MetadataApplier applier, RequestContext reque
         // Exclude null headers
         if (entry.getValue() != null) {
           String key = entry.getKey();
-          if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
+          if (key.toLowerCase().endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
             metadata.put(Metadata.Key.of(entry.getKey(), BINARY_BYTE_MARSHALLER), entry.getValue().getBytes());
           } else {
             metadata.put(Metadata.Key.of(entry.getKey(), ASCII_STRING_MARSHALLER), entry.getValue());

grpc-client-utils/src/test/java/org/hypertrace/core/grpcutils/client/ContextKeyBasedCredsTest.java

@@ -12,6 +12,7 @@
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
+import static io.grpc.Metadata.BINARY_BYTE_MARSHALLER;
 import static org.mockito.Mockito.mock;
 
 public class ContextKeyBasedCredsTest {
@@ -29,6 +30,8 @@ public void testApplyRequestMetadata_shouldApplyAllHeaders() {
     requestContext.add(RequestContextConstants.TENANT_ID_HEADER_KEY, TENANT_ID);
     requestContext.add(RequestContextConstants.AUTHORIZATION_HEADER, TEST_AUTH_HEADER);
     requestContext.add("x-some-tenant-header", "v1");
+    requestContext.add("grpc-trace-bin", "AAARf5ZpQwlN/8FVe1axOPlaAQIdRU/Y8j0LAgE");
+
 
     Context ctx = Context.current().withValue(RequestContext.CURRENT, requestContext);
 
@@ -38,14 +41,23 @@ public void testApplyRequestMetadata_shouldApplyAllHeaders() {
         @Override
         public void apply(Metadata headers) {
           Map<String, String> headersMap = headers.keys().stream()
-              .collect(Collectors.toUnmodifiableMap(k -> k, k -> headers.get(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER))));
+              .collect(Collectors.toUnmodifiableMap(k -> k, k -> {
+                String value;
+                if (k.toLowerCase().endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
+                  value = new String(headers.get(Metadata.Key.of(k, Metadata.BINARY_BYTE_MARSHALLER)));
+                } else {
+                  value = headers.get(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER));
+                }
+                return value;
+              } ));
 
           // Should filter out the TENANT_ID_HEADER_KEY
           Assertions.assertEquals(
               Map.of(
                   RequestContextConstants.TENANT_ID_HEADER_KEY, TENANT_ID,
                   RequestContextConstants.AUTHORIZATION_HEADER, TEST_AUTH_HEADER,
-                  "x-some-tenant-header", "v1"
+                  "x-some-tenant-header", "v1",
+                  "grpc-trace-bin", "AAARf5ZpQwlN/8FVe1axOPlaAQIdRU/Y8j0LAgE"
               ),
               headersMap);
         }

grpc-client-utils/src/test/java/org/hypertrace/core/grpcutils/client/GrpcClientRequestContextUtilTest.java

@@ -61,6 +61,26 @@ public void testExecuteWithHeadersContextJustAuthHeader_shouldPropagateAllHeader
     }
   }
 
+  @Test
+  public void testExecuteWithHeadersContextAuthAndTracing_shouldPropagateAllHeaders() {
+    Context ctx = Context.current();
+    Context previous = ctx.attach();
+    try {
+      Map<String, String> requestHeaders = Map.of("authorization", "v1", "a2", "v2",
+        "grpc-trace-bin", "AAARf5ZpQwlN/8FVe1axOPlaAQIdRU/Y8j0LAgE");
+
+      GrpcClientRequestContextUtil.executeWithHeadersContext(requestHeaders, () -> {
+        RequestContext requestContext = RequestContext.CURRENT.get();
+        Assertions.assertEquals(Map.of("authorization", "v1", "a2", "v2",
+          "grpc-trace-bin", "AAARf5ZpQwlN/8FVe1axOPlaAQIdRU/Y8j0LAgE"), requestContext.getRequestHeaders());
+        return new Object();
+      });
+    } finally {
+      ctx.detach(previous);
+    }
+  }
+
+
   @Test
   public void testExecuteInTenantIdContext() {
     Assertions.assertNull(RequestContext.CURRENT.get());

grpc-server-utils/src/main/java/org/hypertrace/core/grpcutils/server/RequestContextServerInterceptor.java

@@ -45,7 +45,7 @@ RequestContext createRequestContextFromMetadata(Metadata metadata) {
         .forEach(k -> {
           String value;
           //check if key ends with binary suffix
-          if (k.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
+          if (k.toLowerCase().endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
             byte[] bytes = metadata.get(Metadata.Key.of(k, Metadata.BINARY_BYTE_MARSHALLER));
             value = new String(bytes, StandardCharsets.UTF_8);
           } else {

grpc-server-utils/src/test/java/org/hypertrace/core/grpcutils/server/RequestContextServerInterceptorTests.java

@@ -43,5 +43,19 @@ public void testMetadataKeys() {
     Assertions.assertEquals(2, requestContext.getAll().size());
     Assertions.assertEquals("Bearer Some-bearer-auth-2", requestContext.get("authorization").get());
     Assertions.assertEquals("test-tenant-id-2", requestContext.get(RequestContextConstants.TENANT_ID_METADATA_KEY.name()).get());
+
+
+    metadata = new Metadata();
+    metadata.put(Metadata.Key.of("Authorization", Metadata.ASCII_STRING_MARSHALLER), "Bearer Some-bearer-auth-3");
+    metadata.put(Metadata.Key.of("X-tenant-Id", Metadata.ASCII_STRING_MARSHALLER), "test-tenant-id-3");
+    metadata.put(Metadata.Key.of("X-Some-Other-Header", Metadata.ASCII_STRING_MARSHALLER), "Some-other-header-val-3");
+    metadata.put(Metadata.Key.of("grpc-trace-bin", Metadata.BINARY_BYTE_MARSHALLER), "AAARf5ZpQwlN/8FVe1axOPlaAQIdRU/Y8j0LAgE".getBytes());
+
+    requestContext = interceptor.createRequestContextFromMetadata(metadata);
+
+    Assertions.assertEquals(3, requestContext.getAll().size());
+    Assertions.assertEquals("Bearer Some-bearer-auth-3", requestContext.get("authorization").get());
+    Assertions.assertEquals("test-tenant-id-3", requestContext.get(RequestContextConstants.TENANT_ID_METADATA_KEY.name()).get());
+    Assertions.assertEquals("AAARf5ZpQwlN/8FVe1axOPlaAQIdRU/Y8j0LAgE", requestContext.get("grpc-trace-bin").get());
   }
 }