Adding baggage limits for OT Baggage

dd-trace-core/src/main/java/datadog/trace/core/propagation/ContextInterpreter.java

@@ -31,12 +31,16 @@
 import java.util.Collections;
 import java.util.Map;
 import java.util.TreeMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * When adding new context fields to the ContextInterpreter class remember to clear them in the
  * reset() method.
  */
 public abstract class ContextInterpreter implements AgentPropagation.KeyClassifier {
+  private static final Logger LOG = LoggerFactory.getLogger(ContextInterpreter.class);
+
   private TraceConfig traceConfig;
 
   protected Map<String, String> headerTags;
@@ -48,6 +52,9 @@ public abstract class ContextInterpreter implements AgentPropagation.KeyClassifi
   protected TagMap.Ledger tagLedger;
   protected Map<String, String> baggage;
 
+  private int baggageItemCount;
+  private int baggageBytes;
+
   protected CharSequence lastParentId;
   protected CharSequence origin;
   protected long endToEndStartTime;
@@ -63,6 +70,8 @@ public abstract class ContextInterpreter implements AgentPropagation.KeyClassifi
   private final boolean aiGuardEnabled;
   private boolean collectIpHeaders;
   private final boolean requestHeaderTagsCommaAllowed;
+  private final int baggageMaxItems;
+  private final int baggageMaxBytes;
 
   protected static final boolean LOG_EXTRACT_HEADER_NAMES = Config.get().isLogExtractHeaderNames();
   private static final DDCache<String, String> CACHE = DDCaches.newFixedSizeCache(64);
@@ -78,6 +87,8 @@ protected ContextInterpreter(Config config) {
     this.aiGuardEnabled = config.isAiGuardEnabled();
     this.propagationTagsFactory = PropagationTags.factory(config);
     this.requestHeaderTagsCommaAllowed = config.isRequestHeaderTagsCommaAllowed();
+    this.baggageMaxItems = config.getTraceBaggageMaxItems();
+    this.baggageMaxBytes = config.getTraceBaggageMaxBytes();
   }
 
   final TagMap.Ledger tagLedger() {
@@ -216,15 +227,42 @@ protected final boolean handleMappedBaggage(String key, String value) {
     final String lowerCaseKey = toLowerCase(key);
     final String mappedKey = baggageMapping.get(lowerCaseKey);
     if (null != mappedKey) {
-      if (baggage.isEmpty()) {
-        baggage = new TreeMap<>();
-      }
-      baggage.put(mappedKey, HttpCodec.decode(value));
+      addBaggageItem(mappedKey, value);
       return true;
     }
     return false;
   }
 
+  protected final boolean addBaggageItem(String key, String value) {
+    if (key == null || value == null || baggageMaxItems == 0 || baggageMaxBytes == 0) {
+      return false;
+    }
+    final String oldValue = baggage.get(key);
+    if (oldValue == null && baggageItemCount >= baggageMaxItems) {
+      LOG.debug("Dropping baggage item {}: item limit {} reached", key, baggageMaxItems);
+      return false;
+    }
+
+    final long projectedBytes =
+        oldValue == null
+            ? (long) baggageBytes + key.length() + value.length()
+            : (long) baggageBytes + value.length() - oldValue.length();
+
+    if (projectedBytes > baggageMaxBytes) {
+      LOG.debug("Dropping baggage item {}: byte limit {} reached", key, baggageMaxBytes);
+      return false;
+    }
+    if (baggage.isEmpty()) {
+      baggage = new TreeMap<>();
+    }
+    baggage.put(key, HttpCodec.decode(value));
+    if (oldValue == null) {
+      baggageItemCount++;
+    }
+    baggageBytes = (int) projectedBytes;
+    return true;
+  }
+
   public ContextInterpreter reset(TraceConfig traceConfig) {
     this.traceConfig = traceConfig;
     traceId = DDTraceId.ZERO;
@@ -234,6 +272,8 @@ public ContextInterpreter reset(TraceConfig traceConfig) {
     endToEndStartTime = 0;
     if (tagLedger != null) tagLedger.reset();
     baggage = Collections.emptyMap();
+    baggageItemCount = 0;
+    baggageBytes = 0;
     valid = true;
     fullContext = true;
     httpHeaders = null;

dd-trace-core/src/main/java/datadog/trace/core/propagation/DatadogHttpCodec.java

@@ -19,7 +19,6 @@
 import datadog.trace.core.DDSpanContext;
 import datadog.trace.core.propagation.PropagationTags.HeaderType;
 import java.util.Map;
-import java.util.TreeMap;
 import java.util.function.Supplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -188,13 +187,7 @@ public boolean accept(String key, String value) {
                 propagationTags = propagationTagsFactory.fromHeaderValue(HeaderType.DATADOG, value);
                 break;
               case OT_BAGGAGE:
-                {
-                  if (baggage.isEmpty()) {
-                    baggage = new TreeMap<>();
-                  }
-                  baggage.put(
-                      lowerCaseKey.substring(OT_BAGGAGE_PREFIX.length()), HttpCodec.decode(value));
-                }
+                addBaggageItem(lowerCaseKey.substring(OT_BAGGAGE_PREFIX.length()), value);
                 break;
               default:
             }

dd-trace-core/src/main/java/datadog/trace/core/propagation/HaystackHttpCodec.java

@@ -40,7 +40,7 @@ class HaystackHttpCodec {
 
   static final String HAYSTACK_TRACE_ID_BAGGAGE_KEY = "Haystack-Trace-ID";
   static final String HAYSTACK_SPAN_ID_BAGGAGE_KEY = "Haystack-Span-ID";
-  private static final String HAYSTACK_PARENT_ID_BAGGAGE_KEY = "Haystack-Parent-ID";
+  static final String HAYSTACK_PARENT_ID_BAGGAGE_KEY = "Haystack-Parent-ID";
 
   // public static final long DATADOG = new BigInteger("Datadog!".getBytes()).longValue();
   public static final String DATADOG = "44617461-646f-6721";
@@ -131,6 +131,9 @@ private static class HaystackContextInterpreter extends ContextInterpreter {
 
     private static final String BAGGAGE_PREFIX_LC = "baggage-";
 
+    // Largest reserved value we accept. Only relevant for traceID/spanID
+    private static final int MAX_RESERVED_ID_LENGTH = 64;
+
     private static final int TRACE_ID = 0;
     private static final int SPAN_ID = 1;
     private static final int PARENT_ID = 2;
@@ -203,15 +206,15 @@ public boolean accept(String key, String value) {
           if (null != firstValue) {
             switch (classification) {
               case TRACE_ID:
-                traceId = DD64bTraceId.fromHex(convertUUIDToHexString(value));
-                addBaggageItem(HAYSTACK_TRACE_ID_BAGGAGE_KEY, value);
+                traceId = DD64bTraceId.fromHex(convertUUIDToHexString(firstValue));
+                addReservedBaggageItem(HAYSTACK_TRACE_ID_BAGGAGE_KEY, firstValue);
                 break;
               case SPAN_ID:
-                spanId = DDSpanId.fromHex(convertUUIDToHexString(value));
-                addBaggageItem(HAYSTACK_SPAN_ID_BAGGAGE_KEY, value);
+                spanId = DDSpanId.fromHex(convertUUIDToHexString(firstValue));
+                addReservedBaggageItem(HAYSTACK_SPAN_ID_BAGGAGE_KEY, firstValue);
                 break;
               case PARENT_ID:
-                addBaggageItem(HAYSTACK_PARENT_ID_BAGGAGE_KEY, value);
+                addBaggageItem(HAYSTACK_PARENT_ID_BAGGAGE_KEY, firstValue);
                 break;
               case BAGGAGE:
                 {
@@ -238,7 +241,18 @@ public boolean accept(String key, String value) {
       return true;
     }
 
-    private void addBaggageItem(String key, String value) {
+    /**
+     * Records the value of a reserved key, e.g. traceID/spanID. Ignores baggage item and byte
+     * limits to ensure propagation of key headers. However, if the header exceeds
+     * MAX_RESERVED_ID_LENGTH, value is rejected.
+     *
+     * @param key the reserved baggage key.
+     * @param value the id as it arrived, ignored when longer than {@link #MAX_RESERVED_ID_LENGTH}.
+     */
+    private void addReservedBaggageItem(String key, String value) {
+      if (value == null || value.length() > MAX_RESERVED_ID_LENGTH) {
+        return;
+      }
       if (baggage.isEmpty()) {
         baggage = new TreeMap<>();
       }

dd-trace-core/src/main/java/datadog/trace/core/propagation/W3CHttpCodec.java

@@ -24,7 +24,6 @@
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
 import datadog.trace.core.DDSpanContext;
 import java.util.Map;
-import java.util.TreeMap;
 import java.util.function.Supplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -194,13 +193,7 @@ public boolean accept(String key, String value) {
                 endToEndStartTime = extractEndToEndStartTime(firstHeaderValue(value));
                 break;
               case OT_BAGGAGE:
-                {
-                  if (baggage.isEmpty()) {
-                    baggage = new TreeMap<>();
-                  }
-                  baggage.put(
-                      lowerCaseKey.substring(OT_BAGGAGE_PREFIX.length()), HttpCodec.decode(value));
-                }
+                addBaggageItem(lowerCaseKey.substring(OT_BAGGAGE_PREFIX.length()), value);
                 break;
               default:
             }

dd-trace-core/src/main/java/datadog/trace/core/propagation/XRayHttpCodec.java

@@ -18,7 +18,6 @@
 import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.core.DDSpanContext;
 import java.util.Map;
-import java.util.TreeMap;
 import java.util.function.Supplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -184,7 +183,7 @@ public boolean accept(String key, String value) {
         if (!baggageMapping.isEmpty()) {
           String mappedKey = baggageMapping.get(toLowerCase(key));
           if (null != mappedKey) {
-            addBaggageItem(this, mappedKey, HttpCodec.decode(value));
+            addBaggageItem(mappedKey, value);
           }
         }
         return true;
@@ -235,7 +234,7 @@ static void handleXRayTraceHeader(ContextInterpreter interpreter, String value)
           } else {
             int eqIndex = part.indexOf('=');
             if (eqIndex > 0) {
-              addBaggageItem(interpreter, part.substring(0, eqIndex), part.substring(eqIndex + 1));
+              interpreter.addBaggageItem(part.substring(0, eqIndex), part.substring(eqIndex + 1));
             }
           }
           startPart = endPart + 1;
@@ -255,12 +254,5 @@ private static long extractEndToEndStartTime(String value) {
     private static int convertSamplingPriority(char samplingPriority) {
       return '1' == samplingPriority ? SAMPLER_KEEP : SAMPLER_DROP;
     }
-
-    private static void addBaggageItem(ContextInterpreter interpreter, String key, String value) {
-      if (interpreter.baggage.isEmpty()) {
-        interpreter.baggage = new TreeMap<>();
-      }
-      interpreter.baggage.put(key, HttpCodec.decode(value));
-    }
   }
 }

dd-trace-core/src/test/java/datadog/trace/core/propagation/DatadogHttpExtractorTest.java

@@ -1,6 +1,8 @@
 package datadog.trace.core.propagation;
 
 import static datadog.trace.api.config.TracerConfig.REQUEST_HEADER_TAGS_COMMA_ALLOWED;
+import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_MAX_BYTES;
+import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_MAX_ITEMS;
 import static datadog.trace.api.sampling.PrioritySampling.UNSET;
 import static datadog.trace.bootstrap.instrumentation.api.ContextVisitors.stringValuesMap;
 import static datadog.trace.core.propagation.DatadogHttpCodec.DATADOG_TAGS_KEY;
@@ -31,6 +33,7 @@
 import datadog.trace.test.junit.utils.converter.PrioritySamplingConverter;
 import datadog.trace.test.junit.utils.converter.TraceIdConverter;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.function.Supplier;
 import org.junit.jupiter.api.Test;
@@ -325,6 +328,119 @@ void baggageIsMappedOnContextCreation(
     }
   }
 
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_ITEMS, value = "3")
+  void extractOtBaggageStopsAtItemLimit() {
+    Map<String, String> headers = otBaggageHeaders(50);
+    headers.put(SOME_CUSTOM_BAGGAGE_HEADER, "mappedBaggageValue");
+
+    TagContext context = this.extractor.extract(headers, stringValuesMap());
+
+    assertEquals(3, context.getBaggage().size());
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_ITEMS, value = "1")
+  void extractMappedBaggageStopsAtItemLimit() {
+    Map<String, String> headers = otBaggageHeaders(50);
+    headers.put(SOME_CUSTOM_BAGGAGE_HEADER, "mappedBaggageValue");
+
+    TagContext context = this.extractor.extract(headers, stringValuesMap());
+
+    assertEquals(1, context.getBaggage().size());
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "24")
+  void extractOtBaggageStopsAtByteLimit() {
+    // with single digit indices each stored item is "keyN" + "valueN" = 10 bytes, so 2 fit in 24
+    // bytes and a third would take the total to 30
+    TagContext context = this.extractor.extract(otBaggageHeaders(10), stringValuesMap());
+
+    assertEquals(2, context.getBaggage().size());
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "24")
+  void extractOtBaggageChargesRepeatedKeyOnce() {
+    // headers are visited in insertion order, so the duplicate key is seen before the last item
+    Map<String, String> headers = new LinkedHashMap<>();
+    // "key0" + "val0" is 8 bytes, and the duplicate refunds the value it replaces, so the total
+    // stays at 8 rather than doubling
+    headers.put(OT_BAGGAGE_PREFIX + "key0", "val0");
+    headers.put(OT_BAGGAGE_PREFIX + "KEY0", "val0");
+    // leaving room for these 11 bytes, taking the total to 19
+    headers.put(OT_BAGGAGE_PREFIX + "a", "0123456789");
+
+    TagContext context = this.extractor.extract(headers, stringValuesMap());
+
+    Map<String, String> expected = new HashMap<>();
+    expected.put("key0", "val0");
+    expected.put("a", "0123456789");
+    assertEquals(expected, context.getBaggage());
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_ITEMS, value = "1")
+  void extractOtBaggageAllowsReplacementAtItemLimit() {
+    Map<String, String> headers = new LinkedHashMap<>();
+    headers.put(OT_BAGGAGE_PREFIX + "key0", "old");
+    headers.put(OT_BAGGAGE_PREFIX + "KEY0", "replacement");
+
+    TagContext context = this.extractor.extract(headers, stringValuesMap());
+
+    assertEquals(singletonMap("key0", "replacement"), context.getBaggage());
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "24")
+  void extractOtBaggageChargesOnlyTheDeltaWhenReplacingAValue() {
+    Map<String, String> headers = new LinkedHashMap<>();
+    headers.put(OT_BAGGAGE_PREFIX + "key0", "val0"); // 8 bytes
+    // replaces the value, charging the 8 byte difference rather than another 16 bytes
+    headers.put(OT_BAGGAGE_PREFIX + "KEY0", "012345678901");
+    headers.put(OT_BAGGAGE_PREFIX + "a", "0123456"); // 8 bytes, taking the total to exactly 24
+
+    TagContext context = this.extractor.extract(headers, stringValuesMap());
+
+    Map<String, String> expected = new HashMap<>();
+    expected.put("key0", "012345678901");
+    expected.put("a", "0123456");
+    assertEquals(expected, context.getBaggage());
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "8")
+  void extractOtBaggageChargesEncodedValueSize() {
+    Map<String, String> headers = new LinkedHashMap<>();
+    headers.put(OT_BAGGAGE_PREFIX + "a", "b"); // 2 characters
+    headers.put(OT_BAGGAGE_PREFIX + "c", "%E2%99%A5"); // 1 character key + 9 character raw value
+
+    TagContext context = this.extractor.extract(headers, stringValuesMap());
+
+    assertEquals(singletonMap("a", "b"), context.getBaggage());
+  }
+
+  @Test
+  @WithConfig(key = TRACE_BAGGAGE_MAX_BYTES, value = "3")
+  void extractOtBaggageChargesLiteralUtf8ByCharacterCount() {
+    Map<String, String> headers = new LinkedHashMap<>();
+    headers.put(OT_BAGGAGE_PREFIX + "a", "♥"); // 2 characters
+    headers.put(OT_BAGGAGE_PREFIX + "b", "c"); // 2 more characters, no longer fits
+
+    TagContext context = this.extractor.extract(headers, stringValuesMap());
+
+    assertEquals(singletonMap("a", "♥"), context.getBaggage());
+  }
+
+  private static Map<String, String> otBaggageHeaders(int count) {
+    Map<String, String> headers = new HashMap<>();
+    for (int i = 0; i < count; i++) {
+      headers.put(OT_BAGGAGE_PREFIX + "key" + i, "value" + i);
+    }
+    return headers;
+  }
+
   private static String asString(CharSequence cs) {
     return cs == null ? null : cs.toString();
   }