Missing http.response.headers.content-type span tag on blocking responses (#10711) (#10884)

Backport #10711 to release/v1.60.x

Author: jandro996

dd-java-agent/appsec/build.gradle

@@ -19,6 +19,8 @@ dependencies {
   implementation group: 'io.sqreen', name: 'libsqreen', version: '17.3.0'
   implementation libs.moshi
 
+  compileOnly project(':dd-java-agent:agent-bootstrap')
+  testImplementation project(':dd-java-agent:agent-bootstrap')
   testImplementation libs.bytebuddy
   testImplementation project(':remote-config:remote-config-core')
   testImplementation project(':utils:test-utils')

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java

@@ -162,6 +162,8 @@ public class AppSecRequestContext implements DataBundle, Closeable {
   private final AtomicInteger raspMetricsCounter = new AtomicInteger(0);
 
   private volatile boolean wafBlocked;
+  private volatile String blockingResponseContentType;
+  private volatile Integer blockingResponseContentLength;
   private volatile boolean wafErrors;
   private volatile boolean wafTruncated;
   private volatile boolean wafRequestBlockFailure;
@@ -237,6 +239,22 @@ public boolean isWafBlocked() {
     return wafBlocked;
   }
 
+  public void setBlockingResponseContentType(String contentType) {
+    this.blockingResponseContentType = contentType;
+  }
+
+  public String getBlockingResponseContentType() {
+    return blockingResponseContentType;
+  }
+
+  public void setBlockingResponseContentLength(Integer contentLength) {
+    this.blockingResponseContentLength = contentLength;
+  }
+
+  public Integer getBlockingResponseContentLength() {
+    return blockingResponseContentLength;
+  }
+
   public void setWafErrors() {
     this.wafErrors = true;
   }

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java

@@ -26,6 +26,7 @@
 import com.datadog.appsec.report.AppSecEvent;
 import com.datadog.appsec.report.AppSecEventWrapper;
 import com.datadog.appsec.util.BodyParser;
+import datadog.appsec.api.blocking.BlockingContentType;
 import datadog.trace.api.Config;
 import datadog.trace.api.ProductTraceSource;
 import datadog.trace.api.appsec.HttpClientPayload;
@@ -42,6 +43,7 @@
 import datadog.trace.api.telemetry.LoginEvent;
 import datadog.trace.api.telemetry.RuleType;
 import datadog.trace.api.telemetry.WafMetricCollector;
+import datadog.trace.bootstrap.blocking.BlockingActionHelper;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
 import datadog.trace.util.stacktrace.StackTraceEvent;
@@ -929,6 +931,18 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
         writeResponseHeaders(
             ctx, traceSeg, RESPONSE_HEADERS_ALLOW_LIST, ctx.getResponseHeaders(), false);
       }
+      // For blocking responses the normal response-header collection is bypassed; write the
+      // content-type and content-length that were determined when the blocking action was raised.
+      String blockingContentType = ctx.getBlockingResponseContentType();
+      if (blockingContentType != null) {
+        traceSeg.setTagTop("http.response.headers.content-type", blockingContentType);
+        Integer blockingContentLength = ctx.getBlockingResponseContentLength();
+        if (blockingContentLength != null) {
+          traceSeg.setTagTop(
+              "http.response.headers.content-length", String.valueOf(blockingContentLength));
+        }
+      }
+
       // If extracted any derivatives - commit them
       if (!ctx.commitDerivatives(traceSeg)) {
         log.debug("Unable to commit, derivatives will be skipped {}", ctx.getDerivativeKeys());
@@ -1230,7 +1244,9 @@ private Flow<Void> maybePublishRequestData(AppSecRequestContext ctx) {
 
       try {
         GatewayContext gwCtx = new GatewayContext(false);
-        return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
+        Flow<Void> flow = producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
+        maybeRecordBlockingContentType(ctx, flow);
+        return flow;
       } catch (ExpiredSubscriberInfoException e) {
         this.initialReqDataSubInfo = null;
       }
@@ -1263,7 +1279,9 @@ private Flow<Void> maybePublishResponseData(AppSecRequestContext ctx) {
 
       try {
         GatewayContext gwCtx = new GatewayContext(false);
-        return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
+        Flow<Void> flow = producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
+        maybeRecordBlockingContentType(ctx, flow);
+        return flow;
       } catch (ExpiredSubscriberInfoException e) {
         respDataSubInfo = null;
       }
@@ -1277,6 +1295,26 @@ private ApiSecurityDownstreamSampler downstreamSampler() {
     return downstreamSampler;
   }
 
+  private static void maybeRecordBlockingContentType(AppSecRequestContext ctx, Flow<?> flow) {
+    Flow.Action action = flow.getAction();
+    if (!(action instanceof Flow.Action.RequestBlockingAction)) {
+      return;
+    }
+    Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
+    BlockingContentType bct = rba.getBlockingContentType();
+    if (bct == BlockingContentType.NONE) {
+      return; // redirect — no response body
+    }
+    List<String> acceptValues = ctx.getRequestHeaders().get("accept");
+    String acceptHeader =
+        (acceptValues == null || acceptValues.isEmpty()) ? null : acceptValues.get(0);
+    BlockingActionHelper.TemplateType tt =
+        BlockingActionHelper.determineTemplateType(bct, acceptHeader);
+    byte[] template = BlockingActionHelper.getTemplate(tt, rba.getSecurityResponseId());
+    ctx.setBlockingResponseContentType(BlockingActionHelper.getContentType(tt));
+    ctx.setBlockingResponseContentLength(template.length);
+  }
+
   private static Map<String, List<String>> parseQueryStringParams(
       String queryString, Charset uriEncoding) {
     if (queryString == null) {

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy

@@ -18,6 +18,8 @@ import datadog.trace.api.appsec.MediaType
 import datadog.trace.api.config.GeneralConfig
 import datadog.trace.api.function.TriConsumer
 import datadog.trace.api.function.TriFunction
+import datadog.appsec.api.blocking.BlockingContentType
+import datadog.trace.bootstrap.blocking.BlockingActionHelper
 import datadog.trace.api.gateway.BlockResponseFunction
 import datadog.trace.api.gateway.Flow
 import datadog.trace.api.gateway.IGSpanInfo
@@ -1637,6 +1639,60 @@ class GatewayBridgeSpecification extends DDSpecification {
     }
   }
 
+  void 'blocking response content-type span tag is written for AUTO bct resolved to application/json'() {
+    setup:
+    def rba = new Flow.Action.RequestBlockingAction(403, BlockingContentType.AUTO)
+    def blockingFlow = [getAction: {
+        rba
+      }, getResult: {
+        null
+      }] as Flow
+    IGSpanInfo spanInfo = Stub(AgentSpan) {
+      getTags() >> TagMap.fromMap([:])
+    }
+
+    when:
+    requestMethodURICB.apply(ctx, 'GET', TestURIDataAdapter.create('/'))
+    reqHeaderCB.accept(ctx, 'accept', 'application/json')
+    reqHeadersDoneCB.apply(ctx)
+    eventDispatcher.getDataSubscribers(_) >> nonEmptyDsInfo
+    eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >> blockingFlow
+    requestSocketAddressCB.apply(ctx, '0.0.0.0', 5555)
+    requestEndedCB.apply(ctx, spanInfo)
+
+    then:
+    1 * traceSegment.setTagTop('http.response.headers.content-type', 'application/json')
+    1 * traceSegment.setTagTop('http.response.headers.content-length',
+    String.valueOf(BlockingActionHelper.getTemplate(BlockingActionHelper.TemplateType.JSON, null).length))
+  }
+
+  void 'blocking response content-type span tag is written for AUTO bct resolved to text/html'() {
+    setup:
+    def rba = new Flow.Action.RequestBlockingAction(403, BlockingContentType.AUTO)
+    def blockingFlow = [getAction: {
+        rba
+      }, getResult: {
+        null
+      }] as Flow
+    IGSpanInfo spanInfo = Stub(AgentSpan) {
+      getTags() >> TagMap.fromMap([:])
+    }
+
+    when:
+    requestMethodURICB.apply(ctx, 'GET', TestURIDataAdapter.create('/'))
+    reqHeaderCB.accept(ctx, 'accept', 'text/html,application/xhtml+xml')
+    reqHeadersDoneCB.apply(ctx)
+    eventDispatcher.getDataSubscribers(_) >> nonEmptyDsInfo
+    eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >> blockingFlow
+    requestSocketAddressCB.apply(ctx, '0.0.0.0', 5555)
+    requestEndedCB.apply(ctx, spanInfo)
+
+    then:
+    1 * traceSegment.setTagTop('http.response.headers.content-type', 'text/html;charset=utf-8')
+    1 * traceSegment.setTagTop('http.response.headers.content-length',
+    String.valueOf(BlockingActionHelper.getTemplate(BlockingActionHelper.TemplateType.HTML, null).length))
+  }
+
   static toLowerCaseHeaders(final Map<String, List<String>> headers) {
     return headers.collectEntries {
       [(it.key.toLowerCase(Locale.ROOT)): it.value]