Free AppSecRequestContext resources when the request ends (#7535)

Author: manuel-alvarez-alvarez

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

@@ -11,6 +11,8 @@
 import datadog.trace.api.Config;
 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 io.sqreen.powerwaf.Additive;
 import io.sqreen.powerwaf.PowerwafContext;
 import io.sqreen.powerwaf.PowerwafMetrics;
@@ -89,7 +91,7 @@ public class AppSecRequestContext implements DataBundle, Closeable {
   private String savedRawURI;
   private final Map<String, List<String>> requestHeaders = new LinkedHashMap<>();
   private final Map<String, List<String>> responseHeaders = new LinkedHashMap<>();
-  private Map<String, List<String>> collectedCookies;
+  private volatile Map<String, List<String>> collectedCookies;
   private boolean finishedRequestHeaders;
   private boolean finishedResponseHeaders;
   private String peerAddress;
@@ -106,13 +108,13 @@ public class AppSecRequestContext implements DataBundle, Closeable {
   private boolean convertedReqBodyPublished;
   private boolean respDataPublished;
   private boolean pathParamsPublished;
-  private Map<String, String> derivatives;
+  private volatile Map<String, String> derivatives;
 
   private final AtomicBoolean rateLimited = new AtomicBoolean(false);
   private volatile boolean throttled;
 
   // should be guarded by this
-  private Additive additive;
+  private volatile Additive additive;
   // set after additive is set
   private volatile PowerwafMetrics wafMetrics;
   private volatile PowerwafMetrics raspMetrics;
@@ -197,12 +199,14 @@ public Additive getOrCreateAdditive(PowerwafContext ctx, boolean createMetrics,
   }
 
   public void closeAdditive() {
-    synchronized (this) {
-      if (additive != null) {
-        try {
-          additive.close();
-        } finally {
-          additive = null;
+    if (additive != null) {
+      synchronized (this) {
+        if (additive != null) {
+          try {
+            additive.close();
+          } finally {
+            additive = null;
+          }
         }
       }
     }
@@ -419,20 +423,33 @@ public void setRespDataPublished(boolean respDataPublished) {
 
   @Override
   public void close() {
-    synchronized (this) {
-      if (additive == null) {
-        return;
-      }
-    }
-
-    log.warn("WAF object had not been closed (probably missed request-end event)");
-    closeAdditive();
+    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.warn("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;
+    }
+
+    collectedCookies = null;
+    requestHeaders.clear();
+    responseHeaders.clear();
+    persistentData.clear();
+  }
+
   /** @return the portion of the body read so far, if any */
   public CharSequence getStoredRequestBody() {
     StoredBodySupplier storedRequestBodySupplier = this.storedRequestBodySupplier;
@@ -516,6 +533,7 @@ boolean commitDerivatives(TraceSegment traceSegment) {
       return false;
     }
     derivatives.forEach(traceSegment::setTagTop);
+    derivatives = null;
     return true;
   }
 

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

@@ -33,6 +33,7 @@
 import datadog.trace.api.internal.TraceSegment;
 import datadog.trace.api.telemetry.RuleType;
 import datadog.trace.api.telemetry.WafMetricCollector;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
 import java.net.URI;
@@ -495,10 +496,18 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
       }
     }
 
-    ctx.close();
+    ctx.close(requiresPostProcessing(spanInfo));
     return NoopFlow.INSTANCE;
   }
 
+  private boolean requiresPostProcessing(final IGSpanInfo spanInfo) {
+    if (!(spanInfo instanceof AgentSpan)) {
+      return true; // be conservative
+    }
+    final AgentSpan span = (AgentSpan) spanInfo;
+    return span.isRequiresPostProcessing();
+  }
+
   private Flow<Void> onRequestHeadersDone(RequestContext ctx_) {
     AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
     if (ctx == null || ctx.isReqDataPublished()) {

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

@@ -244,4 +244,32 @@ class AppSecRequestContextSpecification extends DDSpecification {
     0 * rateLimiter.isThrottled()
     result == result2
   }
+
+
+  void 'test that internal data is cleared on close'() {
+    setup:
+    final ctx = new AppSecRequestContext()
+    final fullCleanup = !postProcessing
+
+    when:
+    ctx.requestHeaders.put('Accept', ['*'])
+    ctx.responseHeaders.put('Content-Type', ['text/plain'])
+    ctx.collectedCookies = [cookie : ['test']]
+    ctx.persistentData.put(KnownAddresses.REQUEST_METHOD, 'GET')
+    ctx.derivatives = ['a': 'b']
+    ctx.additive = createAdditive()
+    ctx.close(postProcessing)
+
+    then:
+    ctx.additive == null
+    ctx.derivatives == null
+
+    ctx.requestHeaders.isEmpty() == fullCleanup
+    ctx.responseHeaders.isEmpty() == fullCleanup
+    ctx.cookies.isEmpty() == fullCleanup
+    ctx.persistentData.isEmpty() == fullCleanup
+
+    where:
+    postProcessing << [true, false]
+  }
 }

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

@@ -8,6 +8,7 @@ import com.datadog.appsec.event.data.DataBundle
 import com.datadog.appsec.event.data.KnownAddresses
 import com.datadog.appsec.report.AppSecEvent
 import com.datadog.appsec.report.AppSecEventWrapper
+import datadog.trace.api.Config
 import datadog.trace.api.function.TriConsumer
 import datadog.trace.api.function.TriFunction
 import datadog.trace.api.gateway.BlockResponseFunction
@@ -137,7 +138,7 @@ class GatewayBridgeSpecification extends DDSpecification {
     1 * spanInfo.getTags() >> ['http.client_ip':'1.1.1.1']
     1 * mockAppSecCtx.transferCollectedEvents() >> [event]
     1 * mockAppSecCtx.peerAddress >> '2001::1'
-    1 * mockAppSecCtx.close()
+    1 * mockAppSecCtx.close(false)
     1 * traceSegment.setTagTop("_dd.appsec.enabled", 1)
     1 * traceSegment.setTagTop("_dd.runtime_family", "jvm")
     1 * traceSegment.setTagTop('appsec.event', true)

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

@@ -832,6 +832,11 @@ public void addLink(AgentSpanLink link) {
     }
   }
 
+  @Override
+  public boolean isRequiresPostProcessing() {
+    return context.isRequiresPostProcessing();
+  }
+
   // to be accessible in Spock spies, which the field wouldn't otherwise be
   public long getStartTimeNano() {
     return startTimeNano;

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java

@@ -141,6 +141,8 @@ public interface AgentSpan extends MutableSpan, IGSpanInfo, ImplicitContextKeyed
 
   void addLink(AgentSpanLink link);
 
+  boolean isRequiresPostProcessing();
+
   @Override
   default ScopedContext storeInto(ScopedContext context) {
     return context.with(ScopedContextKey.SPAN_KEY, this);

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java

@@ -855,6 +855,11 @@ public void addLink(AgentSpanLink link) {}
     public AgentSpan setMetaStruct(String field, Object value) {
       return this;
     }
+
+    @Override
+    public boolean isRequiresPostProcessing() {
+      return false;
+    }
   }
 
   public static final class NoopAgentScope implements AgentScope {