instrument inner requests in apache http client

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

@@ -14,7 +14,7 @@ addTestSuiteForDir('latestDepTest', 'test')
 dependencies {
   compileOnly group: 'org.apache.httpcomponents.client5', name: 'httpclient5', version: '5.0'
 
-  testImplementation group: 'org.apache.httpcomponents.client5', name: 'httpclient5', version: '5.0'
+  testImplementation group: 'org.apache.httpcomponents.client5', name: 'httpclient5', version: '5.0.2' // we need this fix https://github.com/apache/httpcomponents-client/commit/a22d889807eb74a0d4b8dd8d6360802a391a6eb3
 
   latestDepTestImplementation group: 'org.apache.httpcomponents.client5', name: 'httpclient5', version: '+'
 }

dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientInstrumentation.java

@@ -7,10 +7,12 @@
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
-import datadog.appsec.api.blocking.BlockingException;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import java.util.Collections;
+import java.util.Map;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.implementation.bytecode.assign.Assigner;
@@ -68,6 +70,15 @@ public String[] helperClassNames() {
     };
   }
 
+  @Override
+  public Map<String, String> contextStore() {
+    // used to mark when a request has been instrumented.
+    // We don't count depth like we usually do, because sub-requests can be spawned and need to be
+    // traced
+    return Collections.singletonMap(
+        "org.apache.hc.core5.http.ClassicHttpRequest", "java.lang.Boolean");
+  }
+
   @Override
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
@@ -116,17 +127,13 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class RequestAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope methodEnter(@Advice.Argument(0) final ClassicHttpRequest request) {
-      try {
-        return HelperMethods.doMethodEnter(request);
-      } catch (BlockingException e) {
-        HelperMethods.onBlockingRequest();
-        // re-throw blocking exceptions
-        throw e;
-      }
+      return HelperMethods.doMethodEnter(
+          InstrumentationContext.get(ClassicHttpRequest.class, Boolean.class), request);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
+        @Advice.Argument(0) final ClassicHttpRequest request,
         @Advice.Enter final AgentScope scope,
         @Advice.Return final Object result,
         @Advice.Thrown final Throwable throwable) {
@@ -139,17 +146,13 @@ public static class HostRequestAdvice {
     public static AgentScope methodEnter(
         @Advice.Argument(0) final HttpHost host,
         @Advice.Argument(1) final ClassicHttpRequest request) {
-      try {
-        return HelperMethods.doMethodEnter(host, request);
-      } catch (BlockingException e) {
-        HelperMethods.onBlockingRequest();
-        // re-throw blocking exceptions
-        throw e;
-      }
+      return HelperMethods.doMethodEnter(
+          InstrumentationContext.get(ClassicHttpRequest.class, Boolean.class), host, request);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
+        @Advice.Argument(1) final ClassicHttpRequest request,
         @Advice.Enter final AgentScope scope,
         @Advice.Return final Object result,
         @Advice.Thrown final Throwable throwable) {
@@ -170,24 +173,21 @@ public static AgentScope methodEnter(
                 typing = Assigner.Typing.DYNAMIC,
                 readOnly = false)
             Object handler) {
-      try {
-        final AgentScope scope = HelperMethods.doMethodEnter(host, request);
-        // Wrap the handler so we capture the status code
-        if (null != scope && handler instanceof HttpClientResponseHandler) {
-          handler =
-              new WrappingStatusSettingResponseHandler(
-                  scope.span(), (HttpClientResponseHandler) handler);
-        }
-        return scope;
-      } catch (BlockingException e) {
-        HelperMethods.onBlockingRequest();
-        // re-throw blocking exceptions
-        throw e;
+      final AgentScope scope =
+          HelperMethods.doMethodEnter(
+              InstrumentationContext.get(ClassicHttpRequest.class, Boolean.class), host, request);
+      // Wrap the handler so we capture the status code
+      if (null != scope && handler instanceof HttpClientResponseHandler) {
+        handler =
+            new WrappingStatusSettingResponseHandler(
+                scope.span(), (HttpClientResponseHandler) handler);
       }
+      return scope;
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
+        @Advice.Argument(1) final ClassicHttpRequest request,
         @Advice.Enter final AgentScope scope,
         @Advice.Return final Object result,
         @Advice.Thrown final Throwable throwable) {

dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java

@@ -7,33 +7,52 @@
 import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.HTTP_REQUEST;
 import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER;
 
-import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import org.apache.hc.client5.http.classic.HttpClient;
+import org.apache.hc.core5.http.ClassicHttpRequest;
 import org.apache.hc.core5.http.HttpHost;
 import org.apache.hc.core5.http.HttpRequest;
 import org.apache.hc.core5.http.HttpResponse;
 
 public class HelperMethods {
-  public static AgentScope doMethodEnter(final HttpRequest request) {
-    final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpClient.class);
-    if (callDepth > 0) {
+
+  public static AgentScope doMethodEnter(
+      final ContextStore<ClassicHttpRequest, Boolean> instrumentationMarker,
+      final ClassicHttpRequest request) {
+    if (testAndSet(instrumentationMarker, request)) {
       return null;
     }
-
     return activateHttpSpan(request);
   }
 
-  public static AgentScope doMethodEnter(HttpHost host, HttpRequest request) {
-    final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpClient.class);
-    if (callDepth > 0) {
+  public static AgentScope doMethodEnter(
+      final ContextStore<ClassicHttpRequest, Boolean> instrumentationMarker,
+      HttpHost host,
+      final ClassicHttpRequest request) {
+    if (testAndSet(instrumentationMarker, request)) {
       return null;
     }
-
     return activateHttpSpan(new HostAndRequestAsHttpUriRequest(host, request));
   }
 
+  // checks current value in context store,
+  // and ensures it's set to true when this method exists
+  private static boolean testAndSet(
+      final ContextStore<ClassicHttpRequest, Boolean> instrumentationMarker,
+      final ClassicHttpRequest request) {
+    if (request == null) {
+      // we probably don't want to instrument a call with a null request ?
+      return true;
+    }
+    Boolean instrumented = instrumentationMarker.get(request);
+    if (instrumented == Boolean.TRUE) {
+      return true;
+    }
+    instrumentationMarker.put(request, Boolean.TRUE);
+    return false;
+  }
+
   private static AgentScope activateHttpSpan(final HttpRequest request) {
     final AgentSpan span = startSpan(HTTP_REQUEST);
     final AgentScope scope = activateSpan(span);
@@ -55,22 +74,18 @@ public static void doMethodExit(
     if (scope == null) {
       return;
     }
-    final AgentSpan span = scope.span();
     try {
+      final AgentSpan span = scope.span();
       if (result instanceof HttpResponse) {
         DECORATE.onResponse(span, (HttpResponse) result);
       } // else they probably provided a ResponseHandler.
 
       DECORATE.onError(span, throwable);
       DECORATE.beforeFinish(span);
     } finally {
+      AgentSpan span = scope.span();
       scope.close();
       span.finish();
-      CallDepthThreadLocalMap.reset(HttpClient.class);
     }
   }
-
-  public static void onBlockingRequest() {
-    CallDepthThreadLocalMap.reset(HttpClient.class);
-  }
 }

dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/test/groovy/ApacheHttpClientTest.groovy

@@ -215,3 +215,72 @@ class ApacheClientResponseHandlerAllV0Test extends ApacheClientResponseHandlerAl
 class ApacheClientResponseHandlerAllV1ForkedTest extends ApacheClientResponseHandlerAll implements TestingGenericHttpNamingConventions.ClientV1 {
 }
 
+/**
+ * Tests that HTTP calls made from within an ExecChainHandler (exec interceptor) are instrumented.
+ * Reproduces the scenario from https://github.com/DataDog/dd-trace-java/issues/10383: an
+ * interceptor fetches a token via a separate HttpClient, adds it as a header, then proceeds.
+ */
+@Timeout(5)
+class ApacheClientNestedExecuteTest extends ApacheHttpClientTest<ClassicHttpRequest> {
+
+  @Override
+  ClassicHttpRequest createRequest(String method, URI uri) {
+    return new BasicClassicHttpRequest(method, uri)
+  }
+
+  @Override
+  CloseableHttpResponse executeRequest(ClassicHttpRequest request, URI uri) {
+    return client.execute(request)
+  }
+
+  def "HTTP call from ExecChainHandler (e.g. token fetch) is instrumented"() {
+    when:
+    def tokenUri = server.address.resolve("/success")
+    def mainUri = server.address.resolve("/success")
+    def requestConfig = RequestConfig.custom()
+      .setConnectTimeout(CONNECT_TIMEOUT_MS, TimeUnit.MILLISECONDS)
+      .build()
+    def tokenClient = HttpClients.custom()
+      .setConnectionManager(new BasicHttpClientConnectionManager())
+      .setDefaultRequestConfig(requestConfig)
+      .build()
+    def tokenUriFinal = tokenUri
+    def clientWithInterceptor = HttpClients.custom()
+      .setConnectionManager(new BasicHttpClientConnectionManager())
+      .setDefaultRequestConfig(requestConfig)
+      .addExecInterceptorLast("token", { request, scope, chain ->
+        String token = tokenClient.execute(
+          new BasicClassicHttpRequest("GET", tokenUriFinal), { resp ->
+            String.valueOf(resp.code)
+          }
+          )
+        request.addHeader(new BasicHeader("x-token", token))
+        return chain.proceed(request, scope)
+      })
+      .build()
+    def response = clientWithInterceptor.execute(
+      new BasicClassicHttpRequest("GET", mainUri), { r ->
+        r
+      }
+      )
+
+    then:
+    response != null
+    assertTraces(3) {
+      sortSpansByStart()
+      trace(size(2)) {
+        sortSpansByStart()
+        // Token request runs first (inside interceptor), then main request
+        clientSpan(it, null, "GET", false, false, tokenUri)
+        clientSpan(it, span(0), "GET", false, false, mainUri)
+      }
+      server.distributedRequestTrace(it, trace(0)[0])
+      server.distributedRequestTrace(it, trace(0)[1])
+    }
+
+    cleanup:
+    tokenClient?.close()
+    clientWithInterceptor?.close()
+  }
+}
+