fix : added fallback support for host headers (#279)

sarthak77 · Sarthak Singhal · web-flow · commit 0b21b01f0d97 · 2021-11-03T12:47:07.000+05:30
* added gethost method

* modified

* added unit test

* resolved PR comments

* added unit test

* nit

* nit

Co-authored-by: Sarthak Singhal &lt;sarthaksinghal@Sarthaks-MacBook-Pro.local&gt;
diff --git a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/enrichers/ApiBoundaryTypeAttributeEnricher.java b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/enrichers/ApiBoundaryTypeAttributeEnricher.java
@@ -128,7 +128,7 @@ public void enrichEvent(StructuredTrace trace, Event event) {
   private void enrichHostHeader(Event event) {
     Protocol protocol = EnrichedSpanUtils.getProtocol(event);
     if (protocol == Protocol.PROTOCOL_GRPC) {
-      Optional<String> host = getGrpcAuthority(event);
+      Optional<String> host = getGrpcHostHeader(event);
       if (host.isPresent()) {
         addEnrichedAttribute(event, HOST_HEADER, AttributeValueCreator.create(host.get()));
         return;
@@ -144,11 +144,18 @@ private void enrichHostHeader(Event event) {
     }
   }
 
-  private Optional<String> getGrpcAuthority(Event event) {
+  private Optional<String> getGrpcHostHeader(Event event) {
     Optional<String> grpcAuthority = RpcSemanticConventionUtils.getGrpcAuthority(event);
     if (grpcAuthority.isPresent()) {
       return getSanitizedHostValue(grpcAuthority.get());
     }
+
+    Optional<String> grpcRequestMetadataHost =
+        RpcSemanticConventionUtils.getGrpcRequestMetadataHost(event);
+    if (grpcRequestMetadataHost.isPresent()) {
+      return getSanitizedHostValue(grpcRequestMetadataHost.get());
+    }
+
     return Optional.empty();
   }
 
diff --git a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/test/java/org/hypertrace/traceenricher/enrichment/enrichers/ApiBoundaryTypeAttributeEnricherTest.java b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/test/java/org/hypertrace/traceenricher/enrichment/enrichers/ApiBoundaryTypeAttributeEnricherTest.java
@@ -2,6 +2,7 @@
 
 import static org.hypertrace.core.span.normalizer.constants.OTelSpanTag.OTEL_SPAN_TAG_RPC_SYSTEM;
 import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_AUTHORITY;
+import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_HOST;
 import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
@@ -273,6 +274,10 @@ public void testEnrichEventWithGrpcAuthority() {
         innerEntrySpan,
         RPC_REQUEST_METADATA_AUTHORITY.getValue(),
         AttributeValue.newBuilder().setValue("testHost").build());
+    addAttributeToEvent(
+        innerEntrySpan,
+        RPC_REQUEST_METADATA_HOST.getValue(),
+        AttributeValue.newBuilder().setValue("testHost2").build());
     addAttributeToEvent(
         innerEntrySpan,
         OTEL_SPAN_TAG_RPC_SYSTEM.getValue(),
@@ -300,6 +305,25 @@ public void testEnrichEventWithGrpcNoAuthority() {
     Assertions.assertEquals(EnrichedSpanUtils.getHostHeader(innerEntrySpan), "testHost");
   }
 
+  @Test
+  public void testEnrichEventWithGrpcNoAuthorityButRequestMetadataHost() {
+    mockProtocol(innerEntrySpan, Protocol.PROTOCOL_GRPC);
+    addEnrichedAttributeToEvent(
+        innerEntrySpan, X_FORWARDED_HOST_METADATA, AttributeValueCreator.create("testHost"));
+
+    addAttributeToEvent(
+        innerEntrySpan,
+        RPC_REQUEST_METADATA_HOST.getValue(),
+        AttributeValue.newBuilder().setValue("testHost2").build());
+    addAttributeToEvent(
+        innerEntrySpan,
+        OTEL_SPAN_TAG_RPC_SYSTEM.getValue(),
+        AttributeValue.newBuilder().setValue("grpc").build());
+
+    target.enrichEvent(trace, innerEntrySpan);
+    Assertions.assertEquals(EnrichedSpanUtils.getHostHeader(innerEntrySpan), "testHost2");
+  }
+
   private void mockStructuredGraph() {
     when(graph.getParentEvent(innerExitSpan)).thenReturn(innerEntrySpan);
     when(graph.getChildrenEvents(innerExitSpan)).thenReturn(null);
diff --git a/semantic-convention-utils/src/main/java/org/hypertrace/semantic/convention/utils/rpc/RpcSemanticConventionUtils.java b/semantic-convention-utils/src/main/java/org/hypertrace/semantic/convention/utils/rpc/RpcSemanticConventionUtils.java
@@ -17,6 +17,7 @@
 import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_BODY_TRUNCATED;
 import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_AUTHORITY;
 import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_CONTENT_LENGTH;
+import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_HOST;
 import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_PATH;
 import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_USER_AGENT;
 import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_X_FORWARDED_FOR;
@@ -68,6 +69,7 @@ public class RpcSemanticConventionUtils {
   private static final String OTEL_RPC_SYSTEM_GRPC =
       OTelRpcSemanticConventions.RPC_SYSTEM_VALUE_GRPC.getValue();
   private static final String OTEL_SPAN_TAG_RPC_SYSTEM_ATTR = OTEL_SPAN_TAG_RPC_SYSTEM.getValue();
+  private static final String RPC_REQUEST_METADATA_HOST_ATTR = RPC_REQUEST_METADATA_HOST.getValue();
 
   private static final String OTHER_GRPC_HOST_PORT = RawSpanConstants.getValue(Grpc.GRPC_HOST_PORT);
   private static final String OTHER_GRPC_METHOD = RawSpanConstants.getValue(Grpc.GRPC_METHOD);
@@ -430,6 +432,20 @@ public static Optional<String> getGrpcXForwardedFor(Event event) {
     return Optional.empty();
   }
 
+  public static Optional<String> getGrpcRequestMetadataHost(Event event) {
+    if (event.getAttributes() == null || event.getAttributes().getAttributeMap() == null) {
+      return Optional.empty();
+    }
+
+    Map<String, AttributeValue> attributeValueMap = event.getAttributes().getAttributeMap();
+
+    if (!isRpcSystemGrpc(attributeValueMap)) {
+      return Optional.empty();
+    }
+    return Optional.ofNullable(attributeValueMap.get(RPC_REQUEST_METADATA_HOST_ATTR))
+        .map(AttributeValue::getValue);
+  }
+
   public static Optional<String> getRpcPath(Event event) {
     String service = getRpcService(event).orElse("");
     String method = getRpcMethod(event).orElse("");
diff --git a/semantic-convention-utils/src/test/java/org/hypertrace/semantic/convention/utils/rpc/RpcSemanticConventionUtilsTest.java b/semantic-convention-utils/src/test/java/org/hypertrace/semantic/convention/utils/rpc/RpcSemanticConventionUtilsTest.java
@@ -17,6 +17,7 @@
 import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_BODY_TRUNCATED;
 import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_AUTHORITY;
 import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_CONTENT_LENGTH;
+import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_HOST;
 import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_PATH;
 import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_USER_AGENT;
 import static org.hypertrace.core.span.normalizer.constants.RpcSpanTag.RPC_REQUEST_METADATA_X_FORWARDED_FOR;
@@ -165,6 +166,37 @@ public void testGetGrpcXForwardedFor() {
     assertEquals("198.12.34.1", RpcSemanticConventionUtils.getGrpcXForwardedFor(event).get());
   }
 
+  @Test
+  public void testGetGrpcRequestMetadataHostForGrpcSystem() {
+    Event event = mock(Event.class);
+    when(event.getAttributes())
+        .thenReturn(
+            Attributes.newBuilder()
+                .setAttributeMap(
+                    Map.of(
+                        OTEL_SPAN_TAG_RPC_SYSTEM.getValue(),
+                        AttributeValue.newBuilder().setValue("grpc").build(),
+                        RPC_REQUEST_METADATA_HOST.getValue(),
+                        AttributeValue.newBuilder().setValue("webhost:9011").build()))
+                .build());
+    assertEquals(
+        "webhost:9011", RpcSemanticConventionUtils.getGrpcRequestMetadataHost(event).get());
+  }
+
+  @Test
+  public void testGetGrpcRequestMetadataHostForNotGrpcSystem() {
+    Event event = mock(Event.class);
+    when(event.getAttributes())
+        .thenReturn(
+            Attributes.newBuilder()
+                .setAttributeMap(
+                    Map.of(
+                        RPC_REQUEST_METADATA_HOST.getValue(),
+                        AttributeValue.newBuilder().setValue("webhost:9011").build()))
+                .build());
+    assertEquals(Optional.empty(), RpcSemanticConventionUtils.getGrpcRequestMetadataHost(event));
+  }
+
   @Test
   public void testGetGrpcStatusMsg() {
     Event event =
diff --git a/span-normalizer/span-normalizer-constants/src/main/java/org/hypertrace/core/span/normalizer/constants/RpcSpanTag.java b/span-normalizer/span-normalizer-constants/src/main/java/org/hypertrace/core/span/normalizer/constants/RpcSpanTag.java
@@ -5,6 +5,7 @@ public enum RpcSpanTag {
   RPC_ERROR_NAME("rpc.error_name"),
   RPC_ERROR_MESSAGE("rpc.error_message"),
   RPC_REQUEST_METADATA("rpc.request.metadata"),
+  RPC_REQUEST_METADATA_HOST("rpc.request.metadata.host"),
   RPC_RESPONSE_METADATA("rpc.response.metadata"),
   RPC_REQUEST_BODY("rpc.request.body"),
   RPC_RESPONSE_BODY("rpc.response.body"),
