diff --git a/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggagePropagator.java b/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggagePropagator.java index 0e345296b32..9872f3859d8 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggagePropagator.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggagePropagator.java @@ -1,7 +1,5 @@ package datadog.trace.core.baggage; -import static java.util.Collections.emptyMap; - import datadog.context.Context; import datadog.context.propagation.CarrierSetter; import datadog.context.propagation.CarrierVisitor; @@ -19,6 +17,7 @@ import java.util.HashMap; import java.util.Map; import java.util.function.BiConsumer; +import javax.annotation.Nullable; import javax.annotation.ParametersAreNonnullByDefault; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -142,8 +141,7 @@ public Context extract(Context context, C carrier, CarrierVisitor visitor private class BaggageExtractor implements BiConsumer { private static final char KEY_VALUE_SEPARATOR = '='; private static final char PAIR_SEPARATOR = ','; - private Baggage extracted; - private String w3cHeader; + @Nullable private Baggage extracted; /** URL decode value */ private String decode(final String value) { @@ -156,41 +154,42 @@ private String decode(final String value) { return decoded; } - private Map parseBaggageHeaders(String input) { + private Baggage parseBaggageHeaders(String input) { Map baggage = new HashMap<>(); int start = 0; - boolean truncatedCache = false; + String w3cHeader = input; int pairSeparatorInd = input.indexOf(PAIR_SEPARATOR); pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd; int kvSeparatorInd = input.indexOf(KEY_VALUE_SEPARATOR); while (kvSeparatorInd != -1) { int end = pairSeparatorInd; + boolean limitReached = baggage.size() >= maxItems || end > maxBytes; + if (limitReached) { + // if header was not invalidated already, and we go out of range: + // - fully invalidate if it's after the first k/v pair, + // - otherwise ignore from the current k/v separator + w3cHeader = (w3cHeader == null || start == 0) ? null : w3cHeader.substring(0, start - 1); + break; + } if (kvSeparatorInd > end) { LOG.debug( "Dropping baggage headers due to key with no value {}", input.substring(start, end)); BAGGAGE_METRICS.onBaggageMalformed(); - return emptyMap(); + return null; } String key = decode(input.substring(start, kvSeparatorInd).trim()); String value = decode(input.substring(kvSeparatorInd + 1, end).trim()); if (key.isEmpty() || value.isEmpty()) { LOG.debug("Dropping baggage headers due to empty k/v {}:{}", key, value); BAGGAGE_METRICS.onBaggageMalformed(); - return emptyMap(); + return null; } baggage.put(key, value); // need to percent-encode non-ascii headers we pass down - if (UTF_ESCAPER.keyNeedsEncoding(key) || UTF_ESCAPER.valNeedsEncoding(value)) { - truncatedCache = true; - this.w3cHeader = null; - } else if (!truncatedCache && (end > maxBytes || baggage.size() > maxItems)) { - if (start == 0) { // if we go out of range after first k/v pair, there is no cache - this.w3cHeader = null; - } else { - this.w3cHeader = input.substring(0, start - 1); // -1 to ignore the k/v separator - } - truncatedCache = true; + if (w3cHeader != null + && (UTF_ESCAPER.keyNeedsEncoding(key) || UTF_ESCAPER.valNeedsEncoding(value))) { + w3cHeader = null; } kvSeparatorInd = input.indexOf(KEY_VALUE_SEPARATOR, pairSeparatorInd + 1); @@ -199,20 +198,16 @@ private Map parseBaggageHeaders(String input) { start = end + 1; } - if (!truncatedCache) { - this.w3cHeader = input; - } - - return baggage; + return baggage.isEmpty() ? null : Baggage.create(baggage, w3cHeader); } @Override public void accept(String key, String value) { // Only process tags that are relevant to baggage if (BAGGAGE_KEY.equalsIgnoreCase(key)) { - Map baggage = parseBaggageHeaders(value); - if (!baggage.isEmpty()) { - this.extracted = Baggage.create(baggage, this.w3cHeader); + Baggage parsed = parseBaggageHeaders(value); + if (parsed != null) { + this.extracted = parsed; } } } diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/baggage/BaggagePropagatorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/baggage/BaggagePropagatorTest.groovy index 288eec8d15a..29eb3519aea 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/baggage/BaggagePropagatorTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/core/baggage/BaggagePropagatorTest.groovy @@ -222,4 +222,74 @@ class BaggagePropagatorTest extends DDSpecification { "key1=val1,key2=val2" | "key1=val1,key2=val2" "key1=val1,key2=val2,key3=val3" | "key1=val1,key2=val2" } + + def "test baggage extract items limit"() { + setup: + propagator = new BaggagePropagator(true, true, 2, DEFAULT_TRACE_BAGGAGE_MAX_BYTES) //creating a new instance after injecting config + def headers = [ + (BAGGAGE_KEY) : baggageHeader, + ] + + when: + context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap()) + + then: 'parsing stops once the item limit is exceeded' + Baggage.fromContext(context).asMap() == baggageMap + + where: + baggageHeader | baggageMap + "key1=val1" | [key1: "val1"] + "key1=val1,key2=val2" | [key1: "val1", key2: "val2"] + "key1=val1,key2=val2,key3=val3" | [key1: "val1", key2: "val2"] + } + + def "test baggage extract bytes limit"() { + setup: + propagator = new BaggagePropagator(true, true, DEFAULT_TRACE_BAGGAGE_MAX_ITEMS, 20) //creating a new instance after injecting config + def headers = [ + (BAGGAGE_KEY) : baggageHeader, + ] + + when: + context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap()) + + then: 'parsing stops once the byte limit is exceeded' + Baggage.fromContext(context).asMap() == baggageMap + + where: + baggageHeader | baggageMap + "key1=val1" | [key1: "val1"] + "key1=val1,key2=val2" | [key1: "val1", key2: "val2"] + "key1=val1,key2=val2,key3=val3" | [key1: "val1", key2: "val2"] + } + + def "test baggage extract 0 item limit"() { + setup: + propagator = new BaggagePropagator(true, true, 0, DEFAULT_TRACE_BAGGAGE_MAX_BYTES) //creating a new instance after injecting config + def headers = [ + (BAGGAGE_KEY) : "key1=value1", + ] + + when: + context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap()) + + then: + Baggage.fromContext(context) == null + } + + + + def "test baggage extract 0 byte limit"() { + setup: + propagator = new BaggagePropagator(true, true, DEFAULT_TRACE_BAGGAGE_MAX_ITEMS, 0) //creating a new instance after injecting config + def headers = [ + (BAGGAGE_KEY) : "key1=value1", + ] + + when: + context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap()) + + then: + Baggage.fromContext(context) == null + } }