Honour request header tags mapped onto X-Amzn-Trace-Id

dd-trace-core/src/main/java/datadog/trace/core/propagation/DatadogHttpCodec.java

@@ -135,8 +135,9 @@ public boolean accept(String key, String value) {
           } else if (ORIGIN_KEY.equalsIgnoreCase(key)) {
             classification = ORIGIN;
           } else if (isAwsPropagationEnabled && X_AMZN_TRACE_ID.equalsIgnoreCase(key)) {
+            // Must not return here: the header also has to reach handleTags below so an explicit
+            // DD_TRACE_REQUEST_HEADER_TAGS mapping on it is still honoured.
             handleXRayTraceHeader(this, value);
-            return true;
           } else if (handledXForwarding(key, value)) {
             return true;
           } else if (DATADOG_TAGS_KEY.equalsIgnoreCase(key)) {

dd-trace-core/src/main/java/datadog/trace/core/propagation/XRayHttpCodec.java

@@ -155,8 +155,9 @@ public boolean accept(String key, String value) {
         switch (first) {
           case 'x':
             if (X_AMZN_TRACE_ID.equalsIgnoreCase(key)) {
+              // Must not return here: the header also has to reach handleTags below so an explicit
+              // DD_TRACE_REQUEST_HEADER_TAGS mapping on it is still honoured.
               handleXRayTraceHeader(this, value);
-              return true;
             } else if (handledXForwarding(key, value)) {
               return true;
             }

dd-trace-core/src/test/java/datadog/trace/core/propagation/DatadogHttpExtractorTest.java

@@ -12,6 +12,7 @@
 import static datadog.trace.core.propagation.DatadogHttpCodec.TRACE_ID_KEY;
 import static datadog.trace.core.propagation.HttpCodecTestHelper.headers;
 import static datadog.trace.core.propagation.HttpCodecTestHelper.otBaggageHeaders;
+import static datadog.trace.core.propagation.XRayHttpCodec.X_AMZN_TRACE_ID;
 import static datadog.trace.test.junit.utils.converter.TraceIdConverter.TRACE_ID_MAX_PLUS_1;
 import static java.util.Collections.singletonMap;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -119,6 +120,35 @@ void extractHttpHeadersWithoutComma() {
     assertEquals(expectedHeader, context.getTags().getString(SOME_TAG));
   }
 
+  @Test
+  void extractHeaderTagMappedOnAwsTraceHeader() {
+    // X-Amzn-Trace-Id is consumed as AWS X-Ray context, but a header tag mapped onto it must still
+    // be honoured. This is the load-balancer header example from the AWS documentation: it carries
+    // a Self field and is not in the padded format this codec extracts ids from, so only the tag
+    // is expected from it.
+    this.extractor.cleanup();
+    DynamicConfig<DynamicConfig.Snapshot> dynamicConfig =
+        DynamicConfig.create().setHeaderTags(singletonMap(X_AMZN_TRACE_ID, SOME_TAG)).apply();
+    this.extractor = DatadogHttpCodec.newExtractor(Config.get(), dynamicConfig::captureTraceConfig);
+
+    String awsTraceHeader =
+        "Self=1-67891233-12456789abcdef012345678;Root=1-67891233-abcdef012345678912345678";
+    // spotless:off
+    Map<String, String> headers = headers(
+        TRACE_ID_KEY, "1",
+        SPAN_ID_KEY, "2",
+        X_AMZN_TRACE_ID, awsTraceHeader
+    );
+    // spotless:on
+
+    ExtractedContext context =
+        (ExtractedContext) this.extractor.extract(headers, stringValuesMap());
+
+    assertEquals(awsTraceHeader, context.getTags().getString(SOME_TAG));
+    assertEquals(DDTraceId.from("1"), context.getTraceId());
+    assertEquals(DDSpanId.from("2"), context.getSpanId());
+  }
+
   @ParameterizedTest
   @ValueSource(booleans = {false, true})
   void extractHeaderTagsWithNoPropagation(boolean withOrigin) {

dd-trace-core/src/test/java/datadog/trace/core/propagation/XRayHttpExtractorTest.java

@@ -15,6 +15,7 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.DDSpanId;
 import datadog.trace.api.DDTraceId;
+import datadog.trace.api.DynamicConfig;
 import datadog.trace.api.TraceConfig;
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
 import datadog.trace.test.junit.utils.config.WithConfig;
@@ -36,6 +37,25 @@ protected HttpCodec.Extractor newExtractor(
     return XRayHttpCodec.newExtractor(config, traceConfigSupplier);
   }
 
+  @Test
+  void extractTraceHeaderAlsoCapturesMappedHeaderTag() {
+    // the trace header is consumed as X-Ray context, and must also honour a header tag mapped onto
+    // it without that costing the extracted ids
+    this.extractor.cleanup();
+    DynamicConfig<DynamicConfig.Snapshot> dynamicConfig =
+        DynamicConfig.create().setHeaderTags(singletonMap(X_AMZN_TRACE_ID, SOME_TAG)).apply();
+    this.extractor = XRayHttpCodec.newExtractor(Config.get(), dynamicConfig::captureTraceConfig);
+
+    String traceHeader = "Root=1-00000000-00000000" + zeroPadId("1") + ";Parent=" + zeroPadId("2");
+
+    TagContext context =
+        this.extractor.extract(headers(X_AMZN_TRACE_ID, traceHeader), stringValuesMap());
+
+    assertEquals(traceHeader, context.getTags().getString(SOME_TAG));
+    assertEquals(zeroPadId("1"), context.getTraceId().toHexStringPadded(16));
+    assertEquals(zeroPadId("2"), DDSpanId.toHexStringPadded(context.getSpanId()));
+  }
+
   @Nested
   class BaggageLimits extends AbstractOTBaggageTest {
     @Override