Allow multivalue multiline header extraction for ApacheHttpClient

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

@@ -82,10 +82,14 @@ public AgentSpan onRequest(final AgentSpan span, final REQUEST request) {
       String method = method(request);
       span.setTag(Tags.HTTP_METHOD, method);
 
+      System.out.println("In onRequest method in HttpClientDecorator");
       if (CLIENT_TAG_HEADERS) {
+        System.out.println("Tagging client headers...");
         for (Map.Entry<String, String> headerTag :
             traceConfig(span).getRequestHeaderTags().entrySet()) {
           String headerValue = getRequestHeader(request, headerTag.getKey());
+          System.out.println("headerTag: " + headerTag.getValue());
+          System.out.println("headerValue: " + headerValue);
           if (null != headerValue) {
             span.setTag(headerTag.getValue(), headerValue);
           }

dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpAsyncClientDecorator.java

@@ -55,9 +55,13 @@ protected int status(final HttpContext context) {
 
   @Override
   protected String getRequestHeader(HttpUriRequest request, String headerName) {
-    Header header = request.getFirstHeader(headerName);
-    if (header != null) {
-      return header.getValue();
+    Header[] headers = request.getHeaders(headerName);
+    if (headers.length > 0) {
+      StringBuilder result = new StringBuilder(headers[0].getValue());
+      for (int i = 1; i < headers.length; i++) {
+        result.append(",").append(headers[i].getValue());
+      }
+      return result.toString();
     }
     return null;
   }
@@ -66,9 +70,13 @@ protected String getRequestHeader(HttpUriRequest request, String headerName) {
   protected String getResponseHeader(HttpContext context, String headerName) {
     final Object responseObject = context.getAttribute(HttpCoreContext.HTTP_RESPONSE);
     if (responseObject instanceof HttpResponse) {
-      Header header = ((HttpResponse) responseObject).getFirstHeader(headerName);
-      if (header != null) {
-        return header.getValue();
+      Header[] headers = ((HttpResponse) responseObject).getHeaders(headerName);
+      if (headers.length > 0) {
+        StringBuilder result = new StringBuilder(headers[0].getValue());
+        for (int i = 1; i < headers.length; i++) {
+          result.append(",").append(headers[i].getValue());
+        }
+        return result.toString();
       }
     }
     return null;

dd-java-agent/instrumentation/apache-httpclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/ApacheHttpClientDecorator.java

@@ -46,18 +46,28 @@ protected int status(final HttpResponse httpResponse) {
 
   @Override
   protected String getRequestHeader(HttpUriRequest request, String headerName) {
-    Header header = request.getFirstHeader(headerName);
-    if (null != header) {
-      return header.getValue();
+    Header[] headers = request.getHeaders(headerName);
+    if (headers.length > 0) {
+      StringBuilder result = new StringBuilder(headers[0].getValue());
+      for (int i = 1; i < headers.length; i++) {
+        result.append(",").append(headers[i].getValue());
+      }
+      System.out.println("returning " + result);
+      return result.toString();
     }
     return null;
   }
 
   @Override
   protected String getResponseHeader(HttpResponse response, String headerName) {
-    Header header = response.getFirstHeader(headerName);
-    if (null != header) {
-      return header.getValue();
+    Header[] headers = response.getHeaders(headerName);
+    if (headers.length > 0) {
+      StringBuilder result = new StringBuilder(headers[0].getValue());
+      for (int i = 1; i < headers.length; i++) {
+        result.append(",").append(headers[i].getValue());
+      }
+      System.out.println("returning " + result);
+      return result.toString();
     }
     return null;
   }

dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientDecorator.java

@@ -4,6 +4,7 @@
 import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.util.Arrays;
 import org.apache.hc.core5.http.Header;
 import org.apache.hc.core5.http.HttpRequest;
 import org.apache.hc.core5.http.HttpResponse;
@@ -49,18 +50,31 @@ protected int status(final HttpResponse httpResponse) {
 
   @Override
   protected String getRequestHeader(HttpRequest request, String headerName) {
-    Header header = request.getFirstHeader(headerName);
-    if (null != header) {
-      return header.getValue();
+    System.out.println("headerName: " + headerName);
+    System.out.println("headers: " + Arrays.toString(request.getHeaders(headerName)));
+    Header[] headers = request.getHeaders(headerName);
+    if (headers.length > 0) {
+      StringBuilder result = new StringBuilder(headers[0].getValue());
+      for (int i = 1; i < headers.length; i++) {
+        result.append(",").append(headers[i].getValue());
+      }
+      System.out.println("returning " + result);
+      return result.toString();
     }
+    System.out.println("returning null");
     return null;
   }
 
   @Override
   protected String getResponseHeader(HttpResponse response, String headerName) {
-    Header header = response.getFirstHeader(headerName);
-    if (null != header) {
-      return header.getValue();
+    Header[] headers = response.getHeaders(headerName);
+    if (headers.length > 0) {
+      StringBuilder result = new StringBuilder(headers[0].getValue());
+      for (int i = 1; i < headers.length; i++) {
+        result.append(",").append(headers[i].getValue());
+      }
+      System.out.println("returning " + result);
+      return result.toString();
     }
     return null;
   }

dd-java-agent/instrumentation/apache-httpclient-5/src/test/groovy/ApacheHttpAsyncClient5Test.groovy

@@ -45,10 +45,12 @@ abstract class ApacheHttpAsyncClient5Test<T extends HttpRequest> extends HttpCli
   }
 
   @Override
-  int doRequest(String method, URI uri, Map<String, String> headers, String body, Closure callback) {
+  int doRequest(String method, URI uri, List<List<String>> headers, String body, Closure callback) {
     def request = SimpleHttpRequests.create(method, uri)
     request.setConfig(RequestConfig.custom().setConnectTimeout(CONNECT_TIMEOUT_MS, TimeUnit.MILLISECONDS).build())
-    headers.each { request.addHeader(it.key, it.value) }
+    for (String header : headers) {
+      request.addHeader(header[0], header[1])
+    }
 
     def future = client.execute(request, null)
     def response = future.get(READ_TIMEOUT_MS, TimeUnit.MILLISECONDS)

dd-java-agent/instrumentation/apache-httpclient-5/src/test/groovy/ApacheHttpClientResponseHandlerTest.groovy

@@ -38,10 +38,10 @@ class ApacheHttpClientResponseHandlerTest extends HttpClientTest implements Test
   }
 
   @Override
-  int doRequest(String method, URI uri, Map<String, String> headers, String body, Closure callback) {
+  int doRequest(String method, URI uri, List<List<String>> headers, String body, Closure callback) {
     def request = new BasicClassicHttpRequest(method, uri)
-    headers.entrySet().each {
-      request.addHeader(new BasicHeader(it.key, it.value))
+    for (List<String> header : headers) {
+      request.addHeader(new BasicHeader(header[0], header[1]))
     }
 
     CloseableHttpResponse response = null

dd-java-agent/instrumentation/apache-httpclient-5/src/test/groovy/ApacheHttpClientTest.groovy

@@ -31,10 +31,10 @@ abstract class ApacheHttpClientTest<T extends HttpRequest> extends HttpClientTes
   }
 
   @Override
-  int doRequest(String method, URI uri, Map<String, String> headers, String body, Closure callback) {
+  int doRequest(String method, URI uri, List<List<String>> headers, String body, Closure callback) {
     def request = createRequest(method, uri)
-    headers.entrySet().each {
-      request.addHeader(new BasicHeader(it.key, it.value))
+    for (List<String> header : headers) {
+      request.addHeader(new BasicHeader(header[0], header[1]))
     }
 
     CloseableHttpResponse response = null

dd-java-agent/instrumentation/okhttp-3/src/test/groovy/OkHttp3Test.groovy

@@ -35,12 +35,23 @@ abstract class OkHttp3Test extends HttpClientTest {
   .build()
 
   @Override
-  int doRequest(String method, URI uri, Map<String, String> headers, String body, Closure callback) {
+  int doRequest(String method, URI uri, List<List<String>> headers, String body, Closure callback) {
     def reqBody = HttpMethod.requiresRequestBody(method) ? RequestBody.create(MediaType.parse("text/plain"), body) : null
+    Map<String, String> headersMap = new HashMap<>()
+    for (List<String> header : headers) {
+      String key = header.get(0)
+      String val = header.get(1)
+      if (headersMap.containsKey(key)) {
+        String originalVal = headersMap.get(key)
+        headersMap.put(key, originalVal + "," + val)
+      } else {
+        headersMap.put(key, val)
+      }
+    }
     def request = new Request.Builder()
       .url(uri.toURL())
       .method(method, reqBody)
-      .headers(Headers.of(headers)).build()
+      .headers(Headers.of(headersMap)).build()
     def response = client.newCall(request).execute()
     callback?.call()
     return response.code()
@@ -58,7 +69,7 @@ abstract class OkHttp3Test extends HttpClientTest {
 
   def "request to agent not traced"() {
     when:
-    def status = doRequest(method, url, ["Datadog-Meta-Lang": "java"])
+    def status = doRequest(method, url, [["Datadog-Meta-Lang", "java"]])
 
     then:
     status == 200

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy

@@ -142,7 +142,7 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
    * @param method
    * @return
    */
-  abstract int doRequest(String method, URI uri, Map<String, String> headers = [:], String body = "", Closure callback = null)
+  abstract int doRequest(String method, URI uri, List<List<String>> headers = [], String body = "", Closure callback = null)
 
   String keyStorePath() {
     server.keystorePath
@@ -243,7 +243,7 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
 
     when:
     def status = runUnderTrace("parent") {
-      doRequest(method, url, [:], body)
+      doRequest(method, url, [], body)
     }
     println("RESPONSE: $status")
     if (isDataStreamsEnabled()) {
@@ -283,7 +283,7 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
   def "basic #method request with parent"() {
     when:
     def status = runUnderTrace("parent") {
-      doRequest(method, server.address.resolve("/success"), [:], body)
+      doRequest(method, server.address.resolve("/success"), [], body)
     }
     if (isDataStreamsEnabled()) {
       TEST_DATA_STREAMS_WRITER.waitForGroups(1)
@@ -426,7 +426,7 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
     when:
     injectSysConfig(HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN, "$renameService")
     def status = runUnderTrace("parent") {
-      doRequest(method, server.address.resolve("/success"), ["is-dd-server": "false"])
+      doRequest(method, server.address.resolve("/success"), [["is-dd-server", "false"]])
     }
     if (isDataStreamsEnabled()) {
       TEST_DATA_STREAMS_WRITER.waitForGroups(1)
@@ -464,7 +464,7 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
 
     when:
     def status = runUnderTrace("parent") {
-      doRequest(method, server.address.resolve("/success"), ["is-dd-server": "false"], "") {
+      doRequest(method, server.address.resolve("/success"), [["is-dd-server", "false"]], "") {
         runUnderTrace("child") {
           blockUntilChildSpansFinished(1)
         }
@@ -498,7 +498,7 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
 
   def "trace request with callback and no parent"() {
     when:
-    def status = doRequest(method, server.address.resolve("/success"), ["is-dd-server": "false"], "") {
+    def status = doRequest(method, server.address.resolve("/success"), [["is-dd-server", "false"]], "") {
       runUnderTrace("callback") {
         // FIXME: since in async we may not have the other trace report until the callback is done
         //  we should add a test method to detect that the other trace is finished but waiting for
@@ -650,7 +650,7 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
     def uri = server.address.resolve("/to-secured")
 
     when:
-    def status = doRequest(method, uri, [(BASIC_AUTH_KEY): BASIC_AUTH_VAL])
+    def status = doRequest(method, uri, [[BASIC_AUTH_KEY, BASIC_AUTH_VAL]])
     if (isDataStreamsEnabled()) {
       TEST_DATA_STREAMS_WRITER.waitForGroups(1)
     }
@@ -765,7 +765,9 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
   def "test request header #header tag mapping"() {
     when:
     def url = server.address.resolve("/success")
-    def status = doRequest(method, url, [(header): value])
+    System.out.println("testing header, value, value2: " + header + ", " + value + ", " + value2)
+    String[] headerVals = [header + ":" + value , header + ":" + value2]
+    def status = (value2 == null) ? doRequest(method, url, [[header, value]]) : doRequest(method, url, [[header, value], [header, value2]])
     if (isDataStreamsEnabled()) {
       TEST_DATA_STREAMS_WRITER.waitForGroups(1)
     }
@@ -788,11 +790,13 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
     }
 
     where:
-    method | header                           | value     | tags
-    'GET'  | 'X-Datadog-Test-Both-Header'     | 'foo'     | [ 'both_header_tag': 'foo' ]
-    'GET'  | 'X-Datadog-Test-Request-Header'  | 'bar'     | [ 'request_header_tag': 'bar' ]
-    'GET'  | 'X-Datadog-Test-Both-Header'     | 'bar,baz' | [ 'both_header_tag': 'bar,baz' ]
-    'GET'  | 'X-Datadog-Test-Request-Header'  | 'foo,bar' | [ 'request_header_tag': 'foo,bar' ]
+    method | header                           | value     | value2 | tags
+    'GET'  | 'X-Datadog-Test-Both-Header'     | 'foo'     | null   | [ 'both_header_tag': 'foo' ]
+    'GET'  | 'X-Datadog-Test-Request-Header'  | 'bar'     | null   | [ 'request_header_tag': 'bar' ]
+    'GET'  | 'X-Datadog-Test-Both-Header'     | 'bar,baz' | null   | [ 'both_header_tag': 'bar,baz' ]
+    'GET'  | 'X-Datadog-Test-Request-Header'  | 'foo,bar' | null   | [ 'request_header_tag': 'foo,bar' ]
+    'GET'  | 'X-Datadog-Test-Both-Header'     | 'bar,baz' | 'foo'  | [ 'both_header_tag': 'bar,baz,foo' ]
+    'GET'  | 'X-Datadog-Test-Request-Header'  | 'foo,bar' | 'baz'  | [ 'request_header_tag': 'foo,bar,baz' ]
   }
 
   def "test response header #header tag mapping"() {