Adding wildcard feature for DD_TRACE_HEADER_TAGS and enabling for Http Response headers

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java

@@ -642,14 +642,19 @@ static ResponseHeaderTagClassifier create(AgentSpan span, Map<String, String> he
 
     private final AgentSpan span;
     private final Map<String, String> headerTags;
+    private final String wildcardHeaderPrefix;
 
     public ResponseHeaderTagClassifier(AgentSpan span, Map<String, String> headerTags) {
       this.span = span;
       this.headerTags = headerTags;
+      this.wildcardHeaderPrefix = this.headerTags.getOrDefault("*", null);
     }
 
     @Override
     public boolean accept(String key, String value) {
+      if (wildcardHeaderPrefix != null) {
+        span.setTag((wildcardHeaderPrefix + key).toLowerCase(Locale.ROOT), value);
+      }
       String mappedKey = headerTags.get(key.toLowerCase(Locale.ROOT));
       if (mappedKey != null) {
         span.setTag(mappedKey, value);

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecoratorTest.groovy

@@ -7,6 +7,7 @@ import datadog.trace.api.gateway.Flow
 import datadog.trace.api.gateway.InstrumentationGateway
 import datadog.trace.api.gateway.RequestContext
 import datadog.trace.api.gateway.RequestContextSlot
+import datadog.trace.api.TraceConfig
 import datadog.trace.bootstrap.ActiveSubsystems
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan
@@ -33,6 +34,7 @@ import static datadog.trace.api.gateway.Events.EVENTS
 class HttpServerDecoratorTest extends ServerDecoratorTest {
 
   def span = Mock(AgentSpan)
+  def respHeaders = ['X-Custom-Header': 'custom-value', 'Content-Type': 'application/json']
 
   boolean origAppSecActive
 
@@ -350,12 +352,43 @@ class HttpServerDecoratorTest extends ServerDecoratorTest {
     null   | null           | false
   }
 
+  def "test response headers with trace.header.tags"() {
+    setup:
+    def traceConfig = Mock(TraceConfig)
+    traceConfig.getResponseHeaderTags() >> headerTags
+
+    def tags = [:]
+
+    def responseSpan = Mock(AgentSpan)
+    responseSpan.traceConfig() >> traceConfig
+    responseSpan.setTag(_, _) >> { String k, String v ->
+      tags[k] = v
+      return responseSpan
+    }
+
+    def decorator = newDecorator(null, false)
+
+    when:
+    decorator.onResponse(responseSpan, resp)
+
+    then:
+    for (Map.Entry<String, String> entry : expectedTag.entrySet()) {
+      assert tags[entry.getKey()] == entry.getValue()
+    }
+
+    where:
+    headerTags                         | resp                                                                                             | expectedTag
+    [:]                                | [status: 200, headers: ['X-Custom-Header': 'custom-value', 'Content-Type': 'application/json']]  | [:]
+    ["x-custom-header": "abc"]         | [status: 200, headers: ['X-Custom-Header': 'custom-value', 'Content-Type': 'application/json']]  | [abc:"custom-value"]
+    ["*": "datadog.response.headers."] | [status: 200, headers: ['X-Custom-Header': 'custom-value', 'Content-Type': 'application/json']]  | ["datadog.response.headers.x-custom-header":"custom-value", "datadog.response.headers.content-type":"application/json"]
+  }
+
   @Override
   def newDecorator() {
-    return newDecorator(null)
+    return newDecorator(null, true)
   }
 
-  def newDecorator(TracerAPI tracer) {
+  def newDecorator(TracerAPI tracer, boolean noopResponseGetter) {
     if (!tracer) {
       tracer = AgentTracer.NOOP_TRACER
     }
@@ -383,7 +416,10 @@ class HttpServerDecoratorTest extends ServerDecoratorTest {
 
         @Override
         protected AgentPropagation.ContextVisitor<Map> responseGetter() {
-          return null
+          if (noopResponseGetter){
+            return null
+          }
+          return new MapCarrierVisitor()
         }
 
         @Override
@@ -415,6 +451,19 @@ class HttpServerDecoratorTest extends ServerDecoratorTest {
         protected int status(Map m) {
           return m.status == null ? 0 : m.status
         }
+
+        static class MapCarrierVisitor
+        implements AgentPropagation.ContextVisitor<Map> {
+          @Override
+          void forEachKey(Map carrier, AgentPropagation.KeyClassifier classifier) {
+            Map<String, String> headers = carrier.headers
+            if (headers != null) {
+              for (Map.Entry<String, String> entry : headers.entrySet()) {
+                classifier.accept(entry.key, entry.value)
+              }
+            }
+          }
+        }
       }
   }
 
@@ -449,7 +498,7 @@ class HttpServerDecoratorTest extends ServerDecoratorTest {
       getUniversalCallbackProvider() >> cbpAppSec // no iast callbacks, so this is equivalent
       getDataStreamsMonitoring() >> Mock(DataStreamsMonitoring)
     }
-    def decorator = newDecorator(mTracer)
+    def decorator = newDecorator(mTracer, true)
 
     when:
     decorator.startSpan("test", headers, null)

internal-api/src/main/java/datadog/trace/bootstrap/config/provider/ConfigConverter.java

@@ -318,6 +318,10 @@ private static void loadMapWithOptionalMapping(
                     "Illegal tag starting with non letter for key '" + key + "'");
               }
             } else {
+              if (key.charAt(0) == '*') {
+                map.put(key, defaultPrefix);
+                return;
+              }
               if (Character.isLetter(key.charAt(0))) {
                 value = defaultPrefix + Strings.normalizedHeaderTag(key);
               } else {

internal-api/src/test/groovy/datadog/trace/api/ConfigTest.groovy

@@ -165,6 +165,7 @@ class ConfigTest extends DDSpecification {
   private static final DD_JMXFETCH_METRICS_CONFIGS_ENV = "DD_JMXFETCH_METRICS_CONFIGS"
   private static final DD_TRACE_AGENT_PORT_ENV = "DD_TRACE_AGENT_PORT"
   private static final DD_AGENT_PORT_LEGACY_ENV = "DD_AGENT_PORT"
+  private static final DD_TRACE_HEADER_TAGS = "DD_TRACE_HEADER_TAGS"
   private static final DD_TRACE_REPORT_HOSTNAME = "DD_TRACE_REPORT_HOSTNAME"
   private static final DD_RUNTIME_METRICS_ENABLED_ENV = "DD_RUNTIME_METRICS_ENABLED"
   private static final DD_TRACE_LONG_RUNNING_ENABLED = "DD_TRACE_EXPERIMENTAL_LONG_RUNNING_ENABLED"
@@ -568,6 +569,7 @@ class ConfigTest extends DDSpecification {
     environmentVariables.set(DD_TRACE_X_DATADOG_TAGS_MAX_LENGTH, "42")
     environmentVariables.set(DD_TRACE_LONG_RUNNING_ENABLED, "true")
     environmentVariables.set(DD_TRACE_LONG_RUNNING_FLUSH_INTERVAL, "81")
+    environmentVariables.set(DD_TRACE_HEADER_TAGS, "*")
 
     when:
     def config = new Config()
@@ -587,6 +589,8 @@ class ConfigTest extends DDSpecification {
     config.xDatadogTagsMaxLength == 42
     config.isLongRunningTraceEnabled()
     config.getLongRunningTraceFlushInterval() == 81
+    config.requestHeaderTags == ["*":"http.request.headers."]
+    config.responseHeaderTags == ["*":"http.response.headers."]
   }
 
   def "sys props override env vars"() {

internal-api/src/test/groovy/datadog/trace/bootstrap/config/provider/ConfigConverterTest.groovy

@@ -135,32 +135,35 @@ class ConfigConverterTest extends DDSpecification {
 
   def "test parseMapWithOptionalMappings"() {
     when:
-    def result = ConfigConverter.parseMapWithOptionalMappings(mapString, "test", "", lowercaseKeys)
+    def result = ConfigConverter.parseMapWithOptionalMappings(mapString, "test", defaultPrefix, lowercaseKeys)
 
     then:
     result == expected
 
     where:
-    mapString                     | expected                               | lowercaseKeys
-    "header1:one,header2:two"     | [header1: "one", header2: "two"]       | false
-    "header1:one, header2:two"    | [header1: "one", header2: "two"]       | false
-    "header1,header2:two"         | [header1: "header1", header2: "two"]   | false
-    "Header1:one,header2:two"     | [header1: "one", header2: "two"]       | true
-    "\"header1:one,header2:two\"" | ["\"header1": "one", header2: "two\""] | true
-    "header1"                     | [header1: "header1"]                   | true
-    ",header1:tag"                | [header1: "tag"]                       | true
-    "header1:tag,"                | [header1: "tag"]                       | true
-    "header:tag:value"            | [header: "tag:value"]                  | true
-    ""                            | [:]                                    | true
-    null                          | [:]                                    | true
+    mapString                     | expected                               | lowercaseKeys | defaultPrefix
+    "header1:one,header2:two"     | [header1: "one", header2: "two"]       | false         | ""
+    "header1:one, header2:two"    | [header1: "one", header2: "two"]       | false         | ""
+    "header1,header2:two"         | [header1: "header1", header2: "two"]   | false         | ""
+    "Header1:one,header2:two"     | [header1: "one", header2: "two"]       | true          | ""
+    "\"header1:one,header2:two\"" | ["\"header1": "one", header2: "two\""] | true          | ""
+    "header1"                     | [header1: "header1"]                   | true          | ""
+    ",header1:tag"                | [header1: "tag"]                       | true          | ""
+    "header1:tag,"                | [header1: "tag"]                       | true          | ""
+    "header:tag:value"            | [header: "tag:value"]                  | true          | ""
+    ""                            | [:]                                    | true          | ""
+    null                          | [:]                                    | true          | ""
+    // Test for wildcard header tags
+    "*"                           | ["*":"datadog.response.headers."]      | true          | "datadog.response.headers"
+    "*:"                          | [:]                                    | true          | "datadog.response.headers"
 
     // logs warning: Illegal key only tag starting with non letter '1header'
-    "1header,header2:two"         | [:]                                    | true
+    "1header,header2:two"         | [:]                                    | true          | ""
     // logs warning: Illegal tag starting with non letter for key 'header'
-    "header::tag"                 | [:]                                    | true
+    "header::tag"                 | [:]                                    | true          | ""
     // logs warning: Illegal empty key at position 0
-    ":tag"                        | [:]                                    | true
+    ":tag"                        | [:]                                    | true          | ""
     // logs warning: Illegal empty key at position 11
-    "header:tag,:tag"             | [:]                                    | true
+    "header:tag,:tag"             | [:]                                    | true          | ""
   }
 }