WIP

Author: jandro996

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/DatabaseClientDecorator.java

@@ -140,6 +140,7 @@ public void onRawStatement(AgentSpan span, String sql) {
               brf.tryCommitBlockingResponse(
                   ctx.getTraceSegment(),
                   rba.getStatusCode(),
+                  rba.getBlockId(),
                   rba.getBlockingContentType(),
                   rba.getExtraHeaders());
             }

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpClientDecorator.java

@@ -210,6 +210,7 @@ protected void onHttpClientRequest(final AgentSpan span, final String url) {
         brf.tryCommitBlockingResponse(
             ctx.getTraceSegment(),
             rba.getStatusCode(),
+            rba.getBlockId(),
             rba.getBlockingContentType(),
             rba.getExtraHeaders());
       }

dd-java-agent/appsec/build.gradle

@@ -15,7 +15,7 @@ dependencies {
   implementation project(':internal-api')
   implementation project(':communication')
   implementation project(':telemetry')
-  implementation group: 'io.sqreen', name: 'libsqreen', version: '17.1.0'
+  implementation group: 'io.sqreen', name: 'libsqreen', version: '18.0.0'
   implementation libs.moshi
 
   testImplementation libs.bytebuddy

dd-java-agent/appsec/src/main/java/com/datadog/appsec/blocking/BlockingServiceImpl.java

@@ -68,6 +68,7 @@ public BlockingDetails shouldBlockUser(@Nonnull String userId) {
   @Override
   public boolean tryCommitBlockingResponse(
       int statusCode,
+      String blockId,
       @Nonnull BlockingContentType templateType,
       @Nonnull Map<String, String> extraHeaders) {
     log.info(
@@ -89,7 +90,7 @@ public boolean tryCommitBlockingResponse(
     log.debug("About to call block response function: {}", blockResponseFunction);
     boolean res =
         blockResponseFunction.tryCommitBlockingResponse(
-            reqCtx.getTraceSegment(), statusCode, templateType, extraHeaders);
+            reqCtx.getTraceSegment(), statusCode, blockId, templateType, extraHeaders);
     if (res) {
       TraceSegment traceSegment = reqCtx.getTraceSegment();
       if (traceSegment != null) {

dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java

@@ -477,7 +477,11 @@ private Flow.Action.RequestBlockingAction createBlockRequestAction(
         } catch (IllegalArgumentException iae) {
           log.warn("Unknown content type: {}; using auto", contentType);
         }
-        return new Flow.Action.RequestBlockingAction(statusCode, blockingContentType);
+        String blockId = (String) actionInfo.parameters.get("block_id");
+        if (blockId == null) {
+          throw new RuntimeException("block_request action has no block_id");
+        }
+        return new Flow.Action.RequestBlockingAction(statusCode, blockId, blockingContentType);
       } catch (RuntimeException cce) {
         log.warn("Invalid blocking action data", cce);
         if (!isRasp) {
@@ -506,7 +510,11 @@ private Flow.Action.RequestBlockingAction createRedirectRequestAction(
         if (location == null) {
           throw new RuntimeException("redirect_request action has no location");
         }
-        return Flow.Action.RequestBlockingAction.forRedirect(statusCode, location);
+        String blockId = (String) actionInfo.parameters.get("block_id");
+        if (blockId == null) {
+          throw new RuntimeException("redirect_request action has no block_id");
+        }
+        return Flow.Action.RequestBlockingAction.forRedirect(statusCode, blockId, location);
       } catch (RuntimeException cce) {
         log.warn("Invalid blocking action data", cce);
         if (!isRasp) {

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/BlockingResponseHelper.java

@@ -45,6 +45,7 @@ public static HttpResponse handleFinishForWaf(final AgentSpan span, final HttpRe
         brf.tryCommitBlockingResponse(
             requestContext.getTraceSegment(),
             rba.getStatusCode(),
+            rba.getBlockId(),
             rba.getBlockingContentType(),
             rba.getExtraHeaders());
         HttpResponse altResponse =

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/UnmarshallerHelpers.java

@@ -131,6 +131,7 @@ private static void executeCallback(
             blockResponseFunction.tryCommitBlockingResponse(
                 reqCtx.getTraceSegment(),
                 rba.getStatusCode(),
+                rba.getBlockId(),
                 rba.getBlockingContentType(),
                 rba.getExtraHeaders());
         if (success) {

dd-java-agent/instrumentation/jakarta-rs-annotations-3/src/main/java/datadog/trace/instrumentation/jakarta3/MessageBodyWriterInstrumentation.java

@@ -77,6 +77,7 @@ static void before(
         blockResponseFunction.tryCommitBlockingResponse(
             reqCtx.getTraceSegment(),
             rba.getStatusCode(),
+            rba.getBlockId(),
             rba.getBlockingContentType(),
             rba.getExtraHeaders());
 

dd-java-agent/instrumentation/jax-rs-annotations-2/src/main/java/datadog/trace/instrumentation/jaxrs2/MessageBodyWriterInstrumentation.java

@@ -82,6 +82,7 @@ static void before(
         blockResponseFunction.tryCommitBlockingResponse(
             reqCtx.getTraceSegment(),
             rba.getStatusCode(),
+            rba.getBlockId(),
             rba.getBlockingContentType(),
             rba.getExtraHeaders());
 

dd-java-agent/instrumentation/jersey-2-appsec/src/main/java/datadog/trace/instrumentation/jersey2/MessageBodyReaderInstrumentation.java

@@ -89,6 +89,7 @@ static void after(
           blockResponseFunction.tryCommitBlockingResponse(
               reqCtx.getTraceSegment(),
               rba.getStatusCode(),
+              rba.getBlockId(),
               rba.getBlockingContentType(),
               rba.getExtraHeaders());
           t = new BlockingException("Blocked request (for ReaderInterceptorExecutor/proceed)");