Add safe extraction for malformed base64 kafka headers

dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TextMapExtractAdapter.java

@@ -1,13 +1,16 @@
 package datadog.trace.instrumentation.kafka_clients;
 
+import static datadog.trace.api.Functions.BASE64_DECODE;
+import static datadog.trace.api.Functions.UTF8_BYTES_TO_STRING;
+import static datadog.trace.api.telemetry.LogCollector.EXCLUDE_TELEMETRY;
 import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.KAFKA_PRODUCED_KEY;
-import static java.nio.charset.StandardCharsets.UTF_8;
 
 import datadog.trace.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation.ContextVisitor;
 import java.nio.ByteBuffer;
 import java.util.Base64;
+import java.util.function.Function;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.Headers;
 import org.slf4j.Logger;
@@ -19,10 +22,17 @@ public class TextMapExtractAdapter implements ContextVisitor<Headers> {
   public static final TextMapExtractAdapter GETTER =
       new TextMapExtractAdapter(Config.get().isKafkaClientBase64DecodingEnabled());
 
-  private final Base64.Decoder base64;
+  private final Function<byte[], String> headerValueTransformer;
+  private final Base64.Decoder decoder;
 
-  public TextMapExtractAdapter(boolean base64DecodeHeaders) {
-    this.base64 = base64DecodeHeaders ? Base64.getDecoder() : null;
+  public TextMapExtractAdapter(boolean decodeBase64Headers) {
+    if (decodeBase64Headers) {
+      this.headerValueTransformer = BASE64_DECODE;
+      this.decoder = Base64.getDecoder();
+    } else {
+      this.headerValueTransformer = UTF8_BYTES_TO_STRING;
+      this.decoder = null;
+    }
   }
 
   @Override
@@ -33,10 +43,12 @@ public void forEachKey(Headers carrier, AgentPropagation.KeyClassifier classifie
       if (null == value) {
         continue;
       }
-      if (base64 != null) {
-        value = base64.decode(value);
+      String decoded = headerValueTransformer.apply(value);
+      if (decoded == null) {
+        log.debug(EXCLUDE_TELEMETRY, "Failed to Base64-decode Kafka header '{}', skipping", key);
+        continue;
       }
-      if (!classifier.accept(key, new String(value, UTF_8))) {
+      if (!classifier.accept(key, decoded)) {
         return;
       }
     }
@@ -47,11 +59,11 @@ public long extractTimeInQueueStart(Headers carrier) {
     if (null != header) {
       try {
         ByteBuffer buf = ByteBuffer.allocate(8);
-        buf.put(base64 != null ? base64.decode(header.value()) : header.value());
+        buf.put(decoder != null ? decoder.decode(header.value()) : header.value());
         buf.flip();
         return buf.getLong();
       } catch (Exception e) {
-        log.debug("Unable to get kafka produced time", e);
+        log.debug(EXCLUDE_TELEMETRY, "Unable to get kafka produced time", e);
       }
     }
     return 0;

dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/KafkaClientTestBase.groovy

@@ -1,13 +1,15 @@
 import datadog.trace.api.datastreams.DataStreamsTags
 import datadog.trace.api.datastreams.DataStreamsTransactionExtractor
 import datadog.trace.api.config.TraceInstrumentationConfig
+import datadog.trace.api.config.TracerConfig
 import datadog.trace.instrumentation.kafka_common.ClusterIdHolder
 
 import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled
 
+import datadog.trace.agent.test.InstrumentationSpecification
 import datadog.trace.agent.test.asserts.TraceAssert
 import datadog.trace.agent.test.naming.VersionedNamingTestBase
 import datadog.trace.api.Config
@@ -1544,9 +1546,48 @@ class KafkaClientDataStreamsDisabledForkedTest extends KafkaClientTestBase {
 }
 
 class KafkaClientContextSwapForkedTest extends KafkaClientV0ForkedTest {
-  @Override
   void configurePreAgent() {
     super.configurePreAgent()
     injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false")
   }
 }
+
+class KafkaClientBadBase64HeaderForkedTest extends InstrumentationSpecification {
+  KafkaEmbedded embeddedKafka
+
+  def setup() {
+    embeddedKafka = new KafkaEmbedded(1, true, KafkaClientTestBase.SHARED_TOPIC)
+    embeddedKafka.before()
+  }
+
+  def cleanup() {
+    embeddedKafka?.after()
+  }
+
+  @Override
+  void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig(TraceInstrumentationConfig.KAFKA_CLIENT_BASE64_DECODING_ENABLED, "true")
+    injectSysConfig(TracerConfig.HEADER_TAGS, "x-custom-header:my.custom.tag")
+  }
+
+  def "producer span is created when message carries non-Base64 headers and base64 decoding is enabled"() {
+    setup:
+    def senderProps = KafkaTestUtils.senderProps(embeddedKafka.getBrokersAsString())
+    def producer = new KafkaProducer<String, String>(senderProps, new StringSerializer(), new StringSerializer())
+
+    when:
+    def headers = new RecordHeaders([
+      new RecordHeader("x-custom-header", "not-valid-base64!@#".getBytes(StandardCharsets.UTF_8)),
+      new RecordHeader("x-another-header", "also-not-base64!!".getBytes(StandardCharsets.UTF_8))
+    ])
+    producer.send(new ProducerRecord<>(KafkaClientTestBase.SHARED_TOPIC, 0, null, "hello", headers)).get()
+
+    then:
+    TEST_WRITER.waitForTraces(1)
+    !TEST_WRITER.isEmpty()
+
+    cleanup:
+    producer?.close()
+  }
+}

dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/TextMapExtractAdapterTest.groovy

@@ -2,12 +2,11 @@ import com.google.common.io.BaseEncoding
 import datadog.trace.agent.test.InstrumentationSpecification
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation
 import datadog.trace.instrumentation.kafka_clients.TextMapExtractAdapter
+import java.nio.charset.StandardCharsets
 import org.apache.kafka.common.header.Headers
 import org.apache.kafka.common.header.internals.RecordHeader
 import org.apache.kafka.common.header.internals.RecordHeaders
 
-import java.nio.charset.StandardCharsets
-
 class TextMapExtractAdapterTest extends InstrumentationSpecification {
 
   def "check can decode base64 mangled headers"() {
@@ -32,4 +31,27 @@ class TextMapExtractAdapterTest extends InstrumentationSpecification {
     where:
     base64Decode << [true, false]
   }
+
+  def "invalid base64 header is skipped and subsequent valid headers are still processed"() {
+    given:
+    def validBase64 = BaseEncoding.base64().encode("bar".getBytes(StandardCharsets.UTF_8))
+    Headers headers = new RecordHeaders([
+      new RecordHeader("bad-key", "not-valid-base64!@#".getBytes(StandardCharsets.UTF_8)),
+      new RecordHeader("good-key", validBase64.getBytes(StandardCharsets.UTF_8))
+    ])
+    TextMapExtractAdapter adapter = new TextMapExtractAdapter(true)
+    when:
+    Map<String, String> extracted = [:]
+    adapter.forEachKey(headers, new AgentPropagation.KeyClassifier() {
+        @Override
+        boolean accept(String key, String value) {
+          extracted[key] = value
+          return true
+        }
+      })
+    then:
+    noExceptionThrown()
+    !extracted.containsKey("bad-key")
+    extracted["good-key"] == "bar"
+  }
 }

dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TextMapExtractAdapter.java

@@ -1,12 +1,15 @@
 package datadog.trace.instrumentation.kafka_clients38;
 
-import static java.nio.charset.StandardCharsets.UTF_8;
+import static datadog.trace.api.Functions.BASE64_DECODE;
+import static datadog.trace.api.Functions.UTF8_BYTES_TO_STRING;
+import static datadog.trace.api.telemetry.LogCollector.EXCLUDE_TELEMETRY;
 
 import datadog.trace.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation.ContextVisitor;
 import java.nio.ByteBuffer;
 import java.util.Base64;
+import java.util.function.Function;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.Headers;
 import org.slf4j.Logger;
@@ -18,10 +21,17 @@ public class TextMapExtractAdapter implements ContextVisitor<Headers> {
   public static final TextMapExtractAdapter GETTER =
       new TextMapExtractAdapter(Config.get().isKafkaClientBase64DecodingEnabled());
 
-  private final Base64.Decoder base64;
+  private final Function<byte[], String> headerValueTransformer;
+  private final Base64.Decoder decoder;
 
-  public TextMapExtractAdapter(boolean base64DecodeHeaders) {
-    this.base64 = base64DecodeHeaders ? Base64.getDecoder() : null;
+  public TextMapExtractAdapter(boolean decodeBase64Headers) {
+    if (decodeBase64Headers) {
+      this.headerValueTransformer = BASE64_DECODE;
+      this.decoder = Base64.getDecoder();
+    } else {
+      this.headerValueTransformer = UTF8_BYTES_TO_STRING;
+      this.decoder = null;
+    }
   }
 
   @Override
@@ -32,10 +42,12 @@ public void forEachKey(Headers carrier, AgentPropagation.KeyClassifier classifie
       if (null == value) {
         continue;
       }
-      if (base64 != null) {
-        value = base64.decode(value);
+      String decoded = headerValueTransformer.apply(value);
+      if (decoded == null) {
+        log.debug(EXCLUDE_TELEMETRY, "Failed to Base64-decode Kafka header '{}', skipping", key);
+        continue;
       }
-      if (!classifier.accept(key, new String(value, UTF_8))) {
+      if (!classifier.accept(key, decoded)) {
         return;
       }
     }
@@ -46,11 +58,11 @@ public long extractTimeInQueueStart(Headers carrier) {
     if (null != header) {
       try {
         ByteBuffer buf = ByteBuffer.allocate(8);
-        buf.put(base64 != null ? base64.decode(header.value()) : header.value());
+        buf.put(decoder != null ? decoder.decode(header.value()) : header.value());
         buf.flip();
         return buf.getLong();
       } catch (Exception e) {
-        log.debug("Unable to get kafka produced time", e);
+        log.debug(EXCLUDE_TELEMETRY, "Unable to get kafka produced time", e);
       }
     }
     return 0;

dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/groovy/KafkaClientTestBase.groovy

@@ -1,7 +1,9 @@
+import datadog.trace.agent.test.InstrumentationSpecification
 import datadog.trace.agent.test.asserts.TraceAssert
 import datadog.trace.agent.test.naming.VersionedNamingTestBase
 import datadog.trace.api.Config
 import datadog.trace.api.config.TraceInstrumentationConfig
+import datadog.trace.api.config.TracerConfig
 import datadog.trace.api.DDTags
 import datadog.trace.api.datastreams.DataStreamsTags
 import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags
@@ -1210,9 +1212,48 @@ class KafkaClientDataStreamsDisabledForkedTest extends KafkaClientTestBase {
 }
 
 class KafkaClientContextSwapForkedTest extends KafkaClientV0ForkedTest {
-  @Override
   void configurePreAgent() {
     super.configurePreAgent()
     injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false")
   }
 }
+
+class KafkaClientBadBase64HeaderForkedTest extends InstrumentationSpecification {
+  EmbeddedKafkaBroker embeddedKafka
+
+  def setup() {
+    embeddedKafka = new EmbeddedKafkaKraftBroker(1, 2, KafkaClientTestBase.SHARED_TOPIC)
+    embeddedKafka.afterPropertiesSet()
+  }
+
+  def cleanup() {
+    embeddedKafka.destroy()
+  }
+
+  @Override
+  void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig(TraceInstrumentationConfig.KAFKA_CLIENT_BASE64_DECODING_ENABLED, "true")
+    injectSysConfig(TracerConfig.HEADER_TAGS, "x-custom-header:my.custom.tag")
+  }
+
+  def "producer span is created when message carries non-Base64 headers and base64 decoding is enabled"() {
+    setup:
+    def producerProps = KafkaTestUtils.producerProps(embeddedKafka.getBrokersAsString())
+    def producer = new KafkaProducer<String, String>(producerProps, new StringSerializer(), new StringSerializer())
+
+    when:
+    def headers = new RecordHeaders([
+      new RecordHeader("x-custom-header", "not-valid-base64!@#".getBytes(StandardCharsets.UTF_8)),
+      new RecordHeader("x-another-header", "also-not-base64!!".getBytes(StandardCharsets.UTF_8))
+    ])
+    producer.send(new ProducerRecord<>(KafkaClientTestBase.SHARED_TOPIC, 0, null, "hello", headers)).get()
+
+    then:
+    TEST_WRITER.waitForTraces(1)
+    !TEST_WRITER.isEmpty()
+
+    cleanup:
+    producer?.close()
+  }
+}

dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/groovy/TextMapExtractAdapterTest.groovy

@@ -2,12 +2,11 @@ import com.google.common.io.BaseEncoding
 import datadog.trace.agent.test.InstrumentationSpecification
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation
 import datadog.trace.instrumentation.kafka_clients38.TextMapExtractAdapter
+import java.nio.charset.StandardCharsets
 import org.apache.kafka.common.header.Headers
 import org.apache.kafka.common.header.internals.RecordHeader
 import org.apache.kafka.common.header.internals.RecordHeaders
 
-import java.nio.charset.StandardCharsets
-
 class TextMapExtractAdapterTest extends InstrumentationSpecification {
 
   def "check can decode base64 mangled headers"() {
@@ -32,4 +31,27 @@ class TextMapExtractAdapterTest extends InstrumentationSpecification {
     where:
     base64Decode << [true, false]
   }
+
+  def "invalid base64 header is skipped and subsequent valid headers are still processed"() {
+    given:
+    def validBase64 = BaseEncoding.base64().encode("bar".getBytes(StandardCharsets.UTF_8))
+    Headers headers = new RecordHeaders([
+      new RecordHeader("bad-key", "not-valid-base64!@#".getBytes(StandardCharsets.UTF_8)),
+      new RecordHeader("good-key", validBase64.getBytes(StandardCharsets.UTF_8))
+    ])
+    TextMapExtractAdapter adapter = new TextMapExtractAdapter(true)
+    when:
+    Map<String, String> extracted = [:]
+    adapter.forEachKey(headers, new AgentPropagation.KeyClassifier() {
+        @Override
+        boolean accept(String key, String value) {
+          extracted[key] = value
+          return true
+        }
+      })
+    then:
+    noExceptionThrown()
+    !extracted.containsKey("bad-key")
+    extracted["good-key"] == "bar"
+  }
 }

internal-api/src/main/java/datadog/trace/api/Functions.java

@@ -1,11 +1,13 @@
 package datadog.trace.api;
 
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.function.Function.identity;
 
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
+import java.util.Base64;
 import java.util.Locale;
 import java.util.function.BiFunction;
 import java.util.function.Function;
@@ -174,4 +176,16 @@ public T apply(Object input) {
       }
     }
   }
+
+  public static final Function<byte[], String> UTF8_BYTES_TO_STRING =
+      bytes -> new String(bytes, UTF_8);
+
+  public static final Function<byte[], String> BASE64_DECODE =
+      bytes -> {
+        try {
+          return new String(Base64.getDecoder().decode(bytes), UTF_8);
+        } catch (final Exception ignored) {
+          return null;
+        }
+      };
 }