Api Security schema computation moved to post-processing stage in serialization thread

Author: ValentinZakharov

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

@@ -11,8 +11,6 @@
 import datadog.trace.api.UserIdCollectionMode;
 import datadog.trace.api.http.StoredBodySupplier;
 import datadog.trace.api.internal.TraceSegment;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.util.stacktrace.StackTraceEvent;
 import io.sqreen.powerwaf.Additive;
 import io.sqreen.powerwaf.PowerwafContext;
@@ -485,32 +483,15 @@ public String getSessionId() {
 
   @Override
   public void close() {
-    final AgentSpan span = AgentTracer.activeSpan();
-    close(span != null && span.isRequiresPostProcessing());
-  }
-
-  /* end interface for GatewayBridge */
-
-  /* Should be accessible from the modules */
-
-  public void close(boolean requiresPostProcessing) {
-    if (additive != null || derivatives != null) {
-      log.debug(
-          SEND_TELEMETRY, "WAF object had not been closed (probably missed request-end event)");
-      closeAdditive();
-      derivatives = null;
-    }
-
-    // check if we might need to further post process data related to the span in order to not free
-    // related data
-    if (requiresPostProcessing) {
-      return;
-    }
-
+    closeAdditive();
     collectedCookies = null;
     requestHeaders.clear();
     responseHeaders.clear();
     persistentData.clear();
+    if (derivatives != null) {
+      derivatives.clear();
+      derivatives = null;
+    }
   }
 
   /** @return the portion of the body read so far, if any */

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

@@ -766,11 +766,6 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
       return NoopFlow.INSTANCE;
     }
 
-    maybeExtractSchemas(ctx);
-
-    // WAF call
-    ctx.closeAdditive();
-
     TraceSegment traceSeg = ctx_.getTraceSegment();
 
     // AppSec report metric and events for web span only
@@ -833,7 +828,9 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
       }
     }
 
-    ctx.close(spanInfo.isRequiresPostProcessing());
+    if (!spanInfo.isRequiresPostProcessing()) {
+      ctx.close();
+    }
     return NoopFlow.INSTANCE;
   }
 
@@ -893,7 +890,13 @@ private void onRequestHeader(RequestContext ctx_, String name, String value) {
   }
 
   private void onPostProcessing(RequestContext ctx_) {
-    // Do AppSec post-processing
+    AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
+    if (ctx == null) {
+      return;
+    }
+
+    maybeExtractSchemas(ctx);
+    ctx.close();
   }
 
   public void stop() {

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

@@ -249,7 +249,6 @@ class AppSecRequestContextSpecification extends DDSpecification {
   void 'test that internal data is cleared on close'() {
     setup:
     final ctx = new AppSecRequestContext()
-    final fullCleanup = !postProcessing
 
     when:
     ctx.requestHeaders.put('Accept', ['*'])
@@ -258,19 +257,17 @@ class AppSecRequestContextSpecification extends DDSpecification {
     ctx.persistentData.put(KnownAddresses.REQUEST_METHOD, 'GET')
     ctx.derivatives = ['a': 'b']
     ctx.additive = createAdditive()
-    ctx.close(postProcessing)
+    ctx.close()
 
     then:
     ctx.additive == null
     ctx.derivatives == null
+    ctx.additive == null
 
-    ctx.requestHeaders.isEmpty() == fullCleanup
-    ctx.responseHeaders.isEmpty() == fullCleanup
-    ctx.cookies.isEmpty() == fullCleanup
-    ctx.persistentData.isEmpty() == fullCleanup
-
-    where:
-    postProcessing << [true, false]
+    ctx.requestHeaders.isEmpty()
+    ctx.responseHeaders.isEmpty()
+    ctx.cookies.isEmpty()
+    ctx.persistentData.isEmpty()
   }
 
   def "test increase and get WafTimeouts"() {

dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java

@@ -276,7 +276,9 @@ void onRootSpanPublished(final AgentSpan root) {
     RequestContext requestContext = root.getRequestContext();
     if (requestContext != null) {
       try {
-        requestContext.close();
+        if (!root.isRequiresPostProcessing()) {
+          requestContext.close();
+        }
       } catch (IOException e) {
         log.warn("Error closing request context data", e);
       }