Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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",
};
}

Expand Down
Original file line number Diff line number Diff line change
@@ -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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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));
}
}
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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 {
Expand Down
Loading