Merge pull request #2 from hypertrace/update_bin_headers

Added parsing of binaryheaders in RequestContext

Author: tim-mwangi

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

@@ -9,6 +9,7 @@
 import org.slf4j.LoggerFactory;
 
 import static io.grpc.Metadata.ASCII_STRING_MARSHALLER;
+import static io.grpc.Metadata.BINARY_BYTE_MARSHALLER;
 
 public abstract class RequestContextAsCreds extends CallCredentials {
   private static Logger LOGGER = LoggerFactory.getLogger(RequestContextAsCreds.class);
@@ -31,7 +32,12 @@ protected void applyRequestContext(MetadataApplier applier, RequestContext reque
       for (Map.Entry<String, String> entry : requestContext.getAll().entrySet()) {
         // Exclude null headers
         if (entry.getValue() != null) {
-          metadata.put(Metadata.Key.of(entry.getKey(), ASCII_STRING_MARSHALLER), entry.getValue());
+          String key = entry.getKey();
+          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

@@ -9,6 +9,8 @@
 import org.hypertrace.core.grpcutils.context.RequestContext;
 import org.hypertrace.core.grpcutils.context.RequestContextConstants;
 
+import java.nio.charset.StandardCharsets;
+
 /**
  * Interceptor which intercepts the request headers to extract request context and sets it in the context so that the
  * server logic can use the context and they can be passed onto other downstream services.
@@ -41,7 +43,14 @@ RequestContext createRequestContextFromMetadata(Metadata metadata) {
             RequestContextConstants.HEADER_PREFIXES_TO_BE_PROPAGATED.stream()
                 .anyMatch(prefix -> k.toLowerCase().startsWith(prefix.toLowerCase())))
         .forEach(k -> {
-          String value = metadata.get(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER));
+          String value;
+          //check if key ends with binary 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 {
+            value = metadata.get(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER));
+          }
           // The value could be null or empty for some keys so validate that.
           if (value != null && !value.isEmpty()) {
             requestContext.add(k, value);

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());
   }
 }