From c279fa5ed359e694e7185adbf908ba6c6e39749d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Fri, 30 Aug 2024 21:27:50 +0200 Subject: [PATCH] Free AppSecRequestContext resources when the request ends (#7535) --- .../appsec/gateway/AppSecRequestContext.java | 52 +++++++++++++------ .../datadog/appsec/gateway/GatewayBridge.java | 11 +++- .../AppSecRequestContextSpecification.groovy | 28 ++++++++++ .../gateway/GatewayBridgeSpecification.groovy | 3 +- .../main/java/datadog/trace/core/DDSpan.java | 5 ++ .../instrumentation/api/AgentSpan.java | 2 + .../instrumentation/api/AgentTracer.java | 5 ++ 7 files changed, 87 insertions(+), 19 deletions(-) diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java index 98e64d8593f..dc4e3c28f70 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java +++ b/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> requestHeaders = new LinkedHashMap<>(); private final Map> responseHeaders = new LinkedHashMap<>(); - private Map> collectedCookies; + private volatile Map> 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 derivatives; + private volatile Map 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; } diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java index 30a755f062c..7e318d689ed 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java +++ b/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 onRequestHeadersDone(RequestContext ctx_) { AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); if (ctx == null || ctx.isReqDataPublished()) { diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/AppSecRequestContextSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/AppSecRequestContextSpecification.groovy index 9cd74259c45..dd2d8c69949 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/AppSecRequestContextSpecification.groovy +++ b/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] + } } diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index 2a975790af9..3f247a6b32a 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/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) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java index da63ebf0c27..2a0e6711a63 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java +++ b/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; diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java index 6b55ad91b9d..73431389ae1 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java +++ b/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); diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java index 203eb664bd0..c5ea96a35d4 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java +++ b/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 {