Add items and bytes limits to baggage extraction

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 <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor
   private class BaggageExtractor implements BiConsumer<String, String> {
     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<String, String> parseBaggageHeaders(String input) {
+    private Baggage parseBaggageHeaders(String input) {
       Map<String, String> 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<String, String> 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<String, String> baggage = parseBaggageHeaders(value);
-        if (!baggage.isEmpty()) {
-          this.extracted = Baggage.create(baggage, this.w3cHeader);
+        Baggage parsed = parseBaggageHeaders(value);
+        if (parsed != null) {
+          this.extracted = parsed;
         }
       }
     }

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
+  }
 }