Finish HTTP client spans and close scopes when AppSec blocks requests

dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java

@@ -9,6 +9,7 @@
 import static datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator.HTTP_REQUEST;
 import static datadog.trace.instrumentation.apachehttpclient.HttpHeadersInjectAdapter.SETTER;
 
+import datadog.appsec.api.blocking.BlockingException;
 import datadog.context.ContextScope;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -40,8 +41,16 @@ private static ContextScope activateHttpSpan(final HttpUriRequest request) {
     final AgentSpan span = startSpan(APACHE_HTTP_CLIENT.toString(), HTTP_REQUEST);
     final ContextScope scope = activateSpan(span);
 
-    DECORATE.afterStart(span);
-    DECORATE.onRequest(span, request);
+    try {
+      DECORATE.afterStart(span);
+      DECORATE.onRequest(span, request);
+    } catch (BlockingException e) {
+      DECORATE.onError(span, e);
+      DECORATE.beforeFinish(span);
+      scope.close();
+      span.finish();
+      throw e;
+    }
 
     return scope;
   }

dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-4.0/src/test/groovy/ApacheHttpClientBlockingForkedTest.groovy

@@ -0,0 +1,80 @@
+import datadog.appsec.api.blocking.BlockingContentType
+import datadog.appsec.api.blocking.BlockingException
+import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.api.gateway.Flow
+import datadog.trace.api.gateway.RequestContextSlot
+import datadog.trace.bootstrap.CallDepthThreadLocalMap
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer
+import datadog.trace.bootstrap.instrumentation.api.TagContext
+import org.apache.http.HttpHost
+import org.apache.http.client.HttpClient
+import org.apache.http.client.methods.HttpGet
+import org.apache.http.impl.client.DefaultHttpClient
+import org.apache.http.message.BasicHttpRequest
+
+import java.util.function.BiFunction
+
+import static datadog.trace.api.gateway.Events.EVENTS
+
+/** Forked so AppSec configuration is applied before instrumentation is installed. */
+class ApacheHttpClientBlockingForkedTest extends InstrumentationSpecification {
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig('appsec.enabled', 'true')
+    injectSysConfig('appsec.rasp.enabled', 'true')
+  }
+
+  def 'blocking restores the parent and finishes each client span'() {
+    given:
+    def subscription = AgentTracer.get().getSubscriptionService(RequestContextSlot.APPSEC)
+    def blockedSpans = []
+    def flow = Stub(Flow) {
+      getAction() >> new Flow.Action.RequestBlockingAction(403, BlockingContentType.JSON)
+    }
+    subscription.registerCallback(EVENTS.httpClientRequest(), { ctx, request ->
+      blockedSpans.add(AgentTracer.activeSpan())
+      flow
+    } as BiFunction)
+    def parent = TEST_TRACER.startSpan('test', 'parent',
+      new TagContext().withRequestContextDataAppSec(new Object()))
+    def parentScope = AgentTracer.activateSpan(parent)
+    def client = new DefaultHttpClient()
+
+    when:
+    2.times {
+      try {
+        if (hostRequest) {
+          client.execute(new HttpHost('localhost', 1), new BasicHttpRequest('GET', '/blocked'))
+        } else {
+          client.execute(new HttpGet('http://localhost:1/blocked'))
+        }
+        assert false: 'The request must be blocked before reaching the client'
+      } catch (BlockingException expected) {
+        assert AgentTracer.activeSpan().is(parent)
+        assert CallDepthThreadLocalMap.getCallDepth(HttpClient) == 0
+        assert blockedSpans.last().finished
+      }
+    }
+    parentScope.close()
+    parentScope = null
+    parent.finish()
+
+    then:
+    blockedSpans.size() == 2
+    blockedSpans.every { it.parentId == parent.spanId }
+    TEST_WRITER.waitForTraces(1)
+    TEST_WRITER.size() == 1
+    TEST_WRITER[0].size() == 3
+
+    cleanup:
+    parentScope?.close()
+    if (parent != null && !parent.finished) {
+      parent.finish()
+    }
+    subscription.reset()
+
+    where:
+    hostRequest << [false, true]
+  }
+}

dd-java-agent/instrumentation/commons-httpclient-2.0/src/main/java/datadog/trace/instrumentation/commonshttpclient/CommonsHttpClientInstrumentation.java

@@ -60,21 +60,32 @@ public static class ExecAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static ContextScope methodEnter(@Advice.Argument(1) final HttpMethod httpMethod) {
 
+      ContextScope scope = null;
       try {
         final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpClient.class);
         if (callDepth > 0) {
           return null;
         }
 
         final AgentSpan span = startSpan("commons-http-client", HTTP_REQUEST);
-        final ContextScope scope = activateSpan(span);
+        scope = activateSpan(span);
 
         DECORATE.afterStart(span);
         DECORATE.onRequest(span, httpMethod);
 
         return scope;
       } catch (BlockingException e) {
         CallDepthThreadLocalMap.reset(HttpClient.class);
+        if (scope != null) {
+          final AgentSpan span = spanFromScope(scope);
+          try {
+            DECORATE.onError(span, e);
+            DECORATE.beforeFinish(span);
+          } finally {
+            scope.close();
+            span.finish();
+          }
+        }
         // re-throw blocking exceptions
         throw e;
       }

dd-java-agent/instrumentation/commons-httpclient-2.0/src/test/groovy/CommonsHttpClientBlockingForkedTest.groovy

@@ -0,0 +1,70 @@
+import datadog.appsec.api.blocking.BlockingContentType
+import datadog.appsec.api.blocking.BlockingException
+import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.api.gateway.Flow
+import datadog.trace.api.gateway.RequestContextSlot
+import datadog.trace.bootstrap.CallDepthThreadLocalMap
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer
+import datadog.trace.bootstrap.instrumentation.api.TagContext
+import org.apache.commons.httpclient.HttpClient
+import org.apache.commons.httpclient.methods.GetMethod
+
+import java.util.function.BiFunction
+
+import static datadog.trace.api.gateway.Events.EVENTS
+
+/** Forked so AppSec configuration is applied before instrumentation is installed. */
+class CommonsHttpClientBlockingForkedTest extends InstrumentationSpecification {
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig('appsec.enabled', 'true')
+    injectSysConfig('appsec.rasp.enabled', 'true')
+  }
+
+  def 'blocking restores the parent and finishes each client span'() {
+    given:
+    def subscription = AgentTracer.get().getSubscriptionService(RequestContextSlot.APPSEC)
+    def blockedSpans = []
+    def flow = Stub(Flow) {
+      getAction() >> new Flow.Action.RequestBlockingAction(403, BlockingContentType.JSON)
+    }
+    subscription.registerCallback(EVENTS.httpClientRequest(), { ctx, request ->
+      blockedSpans.add(AgentTracer.activeSpan())
+      flow
+    } as BiFunction)
+    def parent = TEST_TRACER.startSpan('test', 'parent',
+      new TagContext().withRequestContextDataAppSec(new Object()))
+    def parentScope = AgentTracer.activateSpan(parent)
+    def client = new HttpClient()
+
+    when:
+    2.times {
+      try {
+        client.executeMethod(new GetMethod('http://localhost:1/blocked'))
+        assert false: 'The request must be blocked before reaching the client'
+      } catch (BlockingException expected) {
+        assert AgentTracer.activeSpan().is(parent)
+        assert CallDepthThreadLocalMap.getCallDepth(HttpClient) == 0
+        assert blockedSpans.last().finished
+      }
+    }
+    parentScope.close()
+    parentScope = null
+    parent.finish()
+
+    then:
+    blockedSpans.size() == 2
+    blockedSpans.every { it.parentId == parent.spanId }
+    TEST_WRITER.waitForTraces(1)
+    TEST_WRITER.size() == 1
+    TEST_WRITER[0].size() == 3
+
+    cleanup:
+    parentScope?.close()
+    if (parent != null && !parent.finished) {
+      parent.finish()
+    }
+    subscription.reset()
+  }
+}

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/SendAdvice.java

@@ -20,6 +20,7 @@ public class SendAdvice {
   @Advice.OnMethodEnter(suppress = Throwable.class)
   public static ContextScope methodEnter(
       @Advice.Argument(value = 0) final HttpRequest httpRequest) {
+    ContextScope scope = null;
     try {
       if (DECORATE.isAgentRequest(httpRequest)) {
         return null;
@@ -34,7 +35,7 @@ public static ContextScope methodEnter(
       }
       DECORATE.allowContextInjection();
       final AgentSpan span = startSpan(INSTRUMENTATION_NAME, OPERATION_NAME);
-      final ContextScope scope = activateSpan(span);
+      scope = activateSpan(span);
 
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, httpRequest);
@@ -44,6 +45,16 @@ public static ContextScope methodEnter(
     } catch (BlockingException e) {
       CallDepthThreadLocalMap.reset(HttpClient.class);
       DECORATE.blockContextInjection();
+      if (scope != null) {
+        final AgentSpan span = spanFromScope(scope);
+        try {
+          DECORATE.onError(span, e);
+          DECORATE.beforeFinish(span);
+        } finally {
+          scope.close();
+          span.finish();
+        }
+      }
       // re-throw blocking exceptions
       throw e;
     }

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/SendAsyncAdvice.java

@@ -22,6 +22,7 @@ public class SendAsyncAdvice {
   public static ContextScope methodEnter(
       @Advice.Argument(value = 0) final HttpRequest httpRequest,
       @Advice.Argument(value = 1, readOnly = false) HttpResponse.BodyHandler<?> bodyHandler) {
+    ContextScope scope = null;
     try {
       if (DECORATE.isAgentRequest(httpRequest)) {
         return null;
@@ -36,7 +37,7 @@ public static ContextScope methodEnter(
       }
       DECORATE.allowContextInjection();
       final AgentSpan span = startSpan(INSTRUMENTATION_NAME, OPERATION_NAME);
-      final ContextScope scope = activateSpan(span);
+      scope = activateSpan(span);
       if (bodyHandler != null) {
         // Pass span directly — BodyHandlerWrapper captures the continuation lazily in apply(),
         // only once response headers arrive. This avoids leaking a continuation when the
@@ -52,6 +53,16 @@ public static ContextScope methodEnter(
     } catch (BlockingException e) {
       CallDepthThreadLocalMap.reset(HttpClient.class);
       DECORATE.blockContextInjection();
+      if (scope != null) {
+        final AgentSpan span = spanFromScope(scope);
+        try {
+          DECORATE.onError(span, e);
+          DECORATE.beforeFinish(span);
+        } finally {
+          scope.close();
+          span.finish();
+        }
+      }
       // re-throw blocking exceptions
       throw e;
     }

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/groovy/datadog/trace/instrumentation/httpclient/JavaHttpClientBlockingForkedTest.groovy

@@ -0,0 +1,82 @@
+package datadog.trace.instrumentation.httpclient
+
+import datadog.appsec.api.blocking.BlockingContentType
+import datadog.appsec.api.blocking.BlockingException
+import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.api.gateway.Flow
+import datadog.trace.api.gateway.RequestContextSlot
+import datadog.trace.bootstrap.CallDepthThreadLocalMap
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer
+import datadog.trace.bootstrap.instrumentation.api.TagContext
+import java.net.http.HttpClient
+import java.net.http.HttpRequest
+import java.net.http.HttpResponse
+
+import java.util.function.BiFunction
+
+import static datadog.trace.api.gateway.Events.EVENTS
+
+/** Forked so AppSec configuration is applied before instrumentation is installed. */
+class JavaHttpClientBlockingForkedTest extends InstrumentationSpecification {
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig('appsec.enabled', 'true')
+    injectSysConfig('appsec.rasp.enabled', 'true')
+  }
+
+  def 'blocking restores the parent and finishes each client span'() {
+    given:
+    def subscription = AgentTracer.get().getSubscriptionService(RequestContextSlot.APPSEC)
+    def blockedSpans = []
+    def flow = Stub(Flow) {
+      getAction() >> new Flow.Action.RequestBlockingAction(403, BlockingContentType.JSON)
+    }
+    subscription.registerCallback(EVENTS.httpClientRequest(), { ctx, clientRequest ->
+      blockedSpans.add(AgentTracer.activeSpan())
+      flow
+    } as BiFunction)
+    def parent = TEST_TRACER.startSpan('test', 'parent',
+      new TagContext().withRequestContextDataAppSec(new Object()))
+    def parentScope = AgentTracer.activateSpan(parent)
+    def client = HttpClient.newHttpClient()
+    def request = HttpRequest.newBuilder(URI.create('http://localhost:1/blocked')).build()
+
+    when:
+    2.times {
+      try {
+        if (async) {
+          client.sendAsync(request, HttpResponse.BodyHandlers.discarding())
+        } else {
+          client.send(request, HttpResponse.BodyHandlers.discarding())
+        }
+        assert false: 'The request must be blocked before reaching the client'
+      } catch (BlockingException expected) {
+        assert AgentTracer.activeSpan().is(parent)
+        assert CallDepthThreadLocalMap.getCallDepth(HttpClient) == 0
+        assert !JavaNetClientDecorator.DECORATE.isContextInjectionAllowed()
+        assert blockedSpans.last().finished
+      }
+    }
+    parentScope.close()
+    parentScope = null
+    parent.finish()
+
+    then:
+    blockedSpans.size() == 2
+    blockedSpans.every { it.parentId == parent.spanId }
+    TEST_WRITER.waitForTraces(1)
+    TEST_WRITER.size() == 1
+    TEST_WRITER[0].size() == 3
+
+    cleanup:
+    parentScope?.close()
+    if (parent != null && !parent.finished) {
+      parent.finish()
+    }
+    subscription.reset()
+
+    where:
+    async << [false, true]
+  }
+}