Avoid duplicate headers when injecting on java.net http client

Author: amarziali

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java/datadog/trace/instrumentation/httpclient/HttpHeadersInstrumentation.java

@@ -48,7 +48,9 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".HttpHeadersInjectAdapter", packageName + ".JavaNetClientDecorator",
+      packageName + ".CaseInsensitiveKey",
+      packageName + ".HttpHeadersInjectAdapter",
+      packageName + ".JavaNetClientDecorator",
     };
   }
 

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/CaseInsensitiveKey.java

@@ -0,0 +1,36 @@
+package datadog.trace.instrumentation.httpclient;
+
+import java.util.Locale;
+import java.util.Objects;
+
+/**
+ * A class usable as key in a Set or Map that matches case-insensitively but still contains the
+ * original value.
+ */
+public final class CaseInsensitiveKey {
+  private final String value;
+  private final String lowercase;
+
+  public CaseInsensitiveKey(final String value) {
+    this.value = value;
+    this.lowercase = value != null ? value.toLowerCase(Locale.ROOT) : null;
+  }
+
+  public String getValue() {
+    return value;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (!(o instanceof CaseInsensitiveKey)) {
+      return false;
+    }
+    CaseInsensitiveKey that = (CaseInsensitiveKey) o;
+    return Objects.equals(lowercase, that.lowercase);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(lowercase);
+  }
+}

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/HeadersAdvice.java

@@ -6,16 +6,27 @@
 import static datadog.trace.instrumentation.httpclient.JavaNetClientDecorator.DECORATE;
 
 import java.net.http.HttpHeaders;
-import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
 
 public class HeadersAdvice {
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
   public static void methodExit(@Advice.Return(readOnly = false) HttpHeaders headers) {
-    final Map<String, List<String>> headerMap = new HashMap<>(headers.map());
+    final Map<CaseInsensitiveKey, List<String>> headerMap = new LinkedHashMap<>();
+    // Note: we don't want to modify the case of the current headers
+    // However adding duplicate keys will throw an IllegalArgumentException so we need to dedupe
+    // case insensitively
+    for (final Map.Entry<String, List<String>> entry : headers.map().entrySet()) {
+      headerMap.put(new CaseInsensitiveKey(entry.getKey()), entry.getValue());
+    }
     DECORATE.injectContext(getCurrentContext(), headerMap, SETTER);
-    headers = HttpHeaders.of(headerMap, KEEP);
+    // convert back
+    final Map<String, List<String>> finalMap = new LinkedHashMap<>();
+    for (final Map.Entry<CaseInsensitiveKey, List<String>> entry : headerMap.entrySet()) {
+      finalMap.put(entry.getKey().getValue(), entry.getValue());
+    }
+    headers = HttpHeaders.of(finalMap, KEEP);
   }
 }

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/HttpHeadersInjectAdapter.java

@@ -6,17 +6,18 @@
 import java.util.Map;
 import java.util.function.BiPredicate;
 
-public class HttpHeadersInjectAdapter implements CarrierSetter<Map<String, List<String>>> {
+public class HttpHeadersInjectAdapter
+    implements CarrierSetter<Map<CaseInsensitiveKey, List<String>>> {
 
   public static final HttpHeadersInjectAdapter SETTER = new HttpHeadersInjectAdapter();
   public static final BiPredicate<String, String> KEEP = HttpHeadersInjectAdapter::keep;
 
-  @Override
-  public void set(final Map<String, List<String>> carrier, final String key, final String value) {
-    carrier.put(key, Collections.singletonList(value));
-  }
-
   public static boolean keep(String key, String value) {
     return true;
   }
+
+  @Override
+  public void set(Map<CaseInsensitiveKey, List<String>> carrier, String key, String value) {
+    carrier.put(new CaseInsensitiveKey(key), Collections.singletonList(value));
+  }
 }

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/groovy/datadog/trace/instrumentation/httpclient/JavaHttpClientTest.groovy

@@ -1,8 +1,9 @@
 package datadog.trace.instrumentation.httpclient
 
-import datadog.trace.agent.test.base.HttpClientTest
+
 import static datadog.trace.instrumentation.httpclient.JavaNetClientDecorator.DECORATE
 
+import datadog.trace.agent.test.base.HttpClientTest
 import java.net.http.HttpClient
 import java.net.http.HttpRequest
 import java.net.http.HttpResponse
@@ -42,6 +43,22 @@ abstract class JavaHttpClientTest extends HttpClientTest {
   boolean testRedirects() {
     false
   }
+
+  def 'should not inject duplicate headers'() {
+    when:
+    def status = doRequest("GET", server.address.resolve("/success"),
+      // our codec inject names all lowercase
+      ["X-Datadog-Trace-ID": "0"])
+
+    then:
+    status == 200
+    assertTraces(2) {
+      trace(size(1)) {
+        clientSpan(it, null)
+      }
+      server.distributedRequestTrace(it, trace(0).last())
+    }
+  }
 }
 
 class JavaHttpClientV0Test extends JavaHttpClientTest {