Apache http client 4: do not copy all request headers on redirect (#7483)

* Apache httpclient: do not copy all request headers on redirect

* Break the propagation if we cannot safely copy

* Update internal-api/src/main/java/datadog/trace/util/PropagationUtils.java

Co-authored-by: Stuart McCulloch <[email protected]>

* review

* document

* add propagationutils to jacoco excludes

---------

Co-authored-by: Stuart McCulloch <[email protected]>

Author: amarziali

dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java

@@ -11,46 +11,24 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan.Context.Extracted;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
+import datadog.trace.util.PropagationUtils;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.TraceState;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.propagation.TextMapGetter;
 import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.context.propagation.TextMapSetter;
-import java.util.Arrays;
 import java.util.Collection;
 import javax.annotation.Nullable;
 import javax.annotation.ParametersAreNonnullByDefault;
 
 @ParametersAreNonnullByDefault
 public class AgentTextMapPropagator implements TextMapPropagator {
-  private final Collection<String> fields;
-
-  public AgentTextMapPropagator() {
-    // Cannot suppose the current injector from AgentTracer so declare them all
-    this.fields =
-        Arrays.asList(
-            // W3C headers
-            "traceparent",
-            "tracestate",
-            // DD headers
-            "x-datadog-trace-id",
-            "x-datadog-parent-id",
-            "x-datadog-sampling-priority",
-            "x-datadog-origin",
-            "x-datadog-tags",
-            // B3 single headers
-            "X-B3-TraceId",
-            "X-B3-SpanId",
-            "X-B3-Sampled",
-            // B3 multi header
-            "b3");
-  }
 
   @Override
   public Collection<String> fields() {
-    return this.fields;
+    return PropagationUtils.KNOWN_PROPAGATION_HEADERS;
   }
 
   @Override

dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpClientRedirectInstrumentation.java

@@ -8,13 +8,16 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.util.PropagationUtils;
 import java.util.Locale;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.implementation.bytecode.assign.Assigner;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.apache.http.Header;
 import org.apache.http.HttpRequest;
+import org.apache.http.client.methods.HttpRequestWrapper;
+import org.apache.http.protocol.HttpContext;
 
 /**
  * Early versions don't copy headers over on redirect. This instrumentation copies our headers over
@@ -26,7 +29,7 @@ public class ApacheHttpClientRedirectInstrumentation extends InstrumenterModule.
     implements Instrumenter.ForTypeHierarchy {
 
   public ApacheHttpClientRedirectInstrumentation() {
-    super("httpasyncclient", "apache-httpasyncclient");
+    super("httpasyncclient", "apache-httpasyncclient", "httpclient-redirect");
   }
 
   @Override
@@ -44,18 +47,24 @@ public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
         isMethod()
             .and(named("getRedirect"))
-            .and(takesArgument(0, named("org.apache.http.HttpRequest"))),
+            .and(takesArgument(2, named("org.apache.http.protocol.HttpContext"))),
         ApacheHttpClientRedirectInstrumentation.class.getName() + "$ClientRedirectAdvice");
   }
 
   public static class ClientRedirectAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     private static void onAfterExecute(
-        @Advice.Argument(value = 0) final HttpRequest original,
+        @Advice.Argument(value = 2) final HttpContext context,
         @Advice.Return(typing = Assigner.Typing.DYNAMIC) final HttpRequest redirect) {
       if (redirect == null) {
         return;
       }
+      Object originalRequest = context.getAttribute("http.request");
+      if (!(originalRequest instanceof HttpRequest)) {
+        return;
+      }
+      HttpRequest original = (HttpRequest) originalRequest;
+
       // Apache HttpClient 4.0.1+ copies headers from original to redirect only
       // if redirect headers are empty. Because we add headers
       // "x-datadog-" and "x-b3-" to redirect: it means redirect headers never
@@ -65,11 +74,17 @@ private static void onAfterExecute(
       if (!redirect.headerIterator().hasNext()) {
         // redirect didn't have other headers besides tracing, so we need to do copy
         // (same work as Apache HttpClient 4.0.1+ does w/o instrumentation)
-        redirect.setHeaders(original.getAllHeaders());
+        if (original instanceof HttpRequestWrapper) {
+          // We should use the initial request because the wrapped one might contain more headers
+          // (i.e. Host) we do not want to copy
+          // if we cannot access the original request we cannot safely copy.
+          // At this point we break the propagation not to corrupt the customer request
+          redirect.setHeaders(((HttpRequestWrapper) original).getOriginal().getAllHeaders());
+        }
       } else {
         for (final Header header : original.getAllHeaders()) {
-          final String name = header.getName().toLowerCase(Locale.ROOT);
-          if (name.startsWith("x-datadog-") || name.startsWith("x-b3-")) {
+          if (PropagationUtils.KNOWN_PROPAGATION_HEADERS.contains(
+              header.getName().toLowerCase(Locale.ROOT))) {
             if (!redirect.containsHeader(header.getName())) {
               redirect.setHeader(header.getName(), header.getValue());
             }

dd-java-agent/instrumentation/apache-httpclient-4/build.gradle

@@ -34,7 +34,7 @@ dependencies {
   compileOnly group: 'org.apache.httpcomponents', name: 'httpclient', version: '4.0'
   testImplementation(testFixtures(project(':dd-java-agent:agent-iast')))
   testImplementation group: 'org.apache.httpcomponents', name: 'httpclient', version: '4.0'
-
+  testImplementation(project(':dd-java-agent:instrumentation:apache-httpasyncclient-4'))
   // to instrument the integration test
   iastIntegrationTestImplementation(testFixtures(project(':dd-java-agent:agent-iast')))
   iastIntegrationTestImplementation group: 'org.apache.httpcomponents', name: 'httpclient', version: '4.0'

internal-api/build.gradle

@@ -140,6 +140,7 @@ excludedClassesCoverage += [
   "datadog.trace.util.PidHelper",
   "datadog.trace.util.PidHelper.Fallback",
   "datadog.trace.util.ProcessUtils",
+  "datadog.trace.util.PropagationUtils",
   "datadog.trace.util.UnsafeUtils",
   "datadog.trace.api.ConfigCollector",
   "datadog.trace.api.Config.HostNameHolder",

internal-api/src/main/java/datadog/trace/util/PropagationUtils.java

@@ -0,0 +1,33 @@
+package datadog.trace.util;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.LinkedHashSet;
+
+public class PropagationUtils {
+  private PropagationUtils() {
+    // avoid constructing instances of this class.
+  }
+
+  /** The list of the known propagation headers. They must be lowercased. */
+  public static final Collection<String> KNOWN_PROPAGATION_HEADERS =
+      Collections.unmodifiableCollection(
+          new LinkedHashSet<>(
+              Arrays.asList(
+                  // W3C headers
+                  "traceparent",
+                  "tracestate",
+                  // DD headers
+                  "x-datadog-trace-id",
+                  "x-datadog-parent-id",
+                  "x-datadog-sampling-priority",
+                  "x-datadog-origin",
+                  "x-datadog-tags",
+                  // B3 single headers
+                  "x-b3-traceid",
+                  "x-b3-spanid",
+                  "x-b3-sampled",
+                  // B3 multi header
+                  "b3")));
+}