Skip to content

Commit d1aac21

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
committed
Free AppSecRequestContext resources when the request ends
1 parent e024c38 commit d1aac21

7 files changed

Lines changed: 87 additions & 19 deletions

File tree

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

Lines changed: 35 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
import datadog.trace.api.Config;
1212
import datadog.trace.api.http.StoredBodySupplier;
1313
import datadog.trace.api.internal.TraceSegment;
14+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
15+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
1416
import io.sqreen.powerwaf.Additive;
1517
import io.sqreen.powerwaf.PowerwafContext;
1618
import io.sqreen.powerwaf.PowerwafMetrics;
@@ -89,7 +91,7 @@ public class AppSecRequestContext implements DataBundle, Closeable {
8991
private String savedRawURI;
9092
private final Map<String, List<String>> requestHeaders = new LinkedHashMap<>();
9193
private final Map<String, List<String>> responseHeaders = new LinkedHashMap<>();
92-
private Map<String, List<String>> collectedCookies;
94+
private volatile Map<String, List<String>> collectedCookies;
9395
private boolean finishedRequestHeaders;
9496
private boolean finishedResponseHeaders;
9597
private String peerAddress;
@@ -106,13 +108,13 @@ public class AppSecRequestContext implements DataBundle, Closeable {
106108
private boolean convertedReqBodyPublished;
107109
private boolean respDataPublished;
108110
private boolean pathParamsPublished;
109-
private Map<String, String> derivatives;
111+
private volatile Map<String, String> derivatives;
110112

111113
private final AtomicBoolean rateLimited = new AtomicBoolean(false);
112114
private volatile boolean throttled;
113115

114116
// should be guarded by this
115-
private Additive additive;
117+
private volatile Additive additive;
116118
// set after additive is set
117119
private volatile PowerwafMetrics wafMetrics;
118120
private volatile PowerwafMetrics raspMetrics;
@@ -197,12 +199,14 @@ public Additive getOrCreateAdditive(PowerwafContext ctx, boolean createMetrics,
197199
}
198200

199201
public void closeAdditive() {
200-
synchronized (this) {
201-
if (additive != null) {
202-
try {
203-
additive.close();
204-
} finally {
205-
additive = null;
202+
if (additive != null) {
203+
synchronized (this) {
204+
if (additive != null) {
205+
try {
206+
additive.close();
207+
} finally {
208+
additive = null;
209+
}
206210
}
207211
}
208212
}
@@ -419,20 +423,33 @@ public void setRespDataPublished(boolean respDataPublished) {
419423

420424
@Override
421425
public void close() {
422-
synchronized (this) {
423-
if (additive == null) {
424-
return;
425-
}
426-
}
427-
428-
log.warn("WAF object had not been closed (probably missed request-end event)");
429-
closeAdditive();
426+
final AgentSpan span = AgentTracer.activeSpan();
427+
close(span != null && span.isRequiresPostProcessing());
430428
}
431429

432430
/* end interface for GatewayBridge */
433431

434432
/* Should be accessible from the modules */
435433

434+
public void close(boolean requiresPostProcessing) {
435+
if (additive != null || derivatives != null) {
436+
log.warn("WAF object had not been closed (probably missed request-end event)");
437+
closeAdditive();
438+
derivatives = null;
439+
}
440+
441+
// check if we might need to further post process data related to the span in order to not free
442+
// related data
443+
if (requiresPostProcessing) {
444+
return;
445+
}
446+
447+
collectedCookies = null;
448+
requestHeaders.clear();
449+
responseHeaders.clear();
450+
persistentData.clear();
451+
}
452+
436453
/** @return the portion of the body read so far, if any */
437454
public CharSequence getStoredRequestBody() {
438455
StoredBodySupplier storedRequestBodySupplier = this.storedRequestBodySupplier;
@@ -516,6 +533,7 @@ boolean commitDerivatives(TraceSegment traceSegment) {
516533
return false;
517534
}
518535
derivatives.forEach(traceSegment::setTagTop);
536+
derivatives = null;
519537
return true;
520538
}
521539

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import datadog.trace.api.internal.TraceSegment;
3434
import datadog.trace.api.telemetry.RuleType;
3535
import datadog.trace.api.telemetry.WafMetricCollector;
36+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
3637
import datadog.trace.bootstrap.instrumentation.api.Tags;
3738
import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
3839
import java.net.URI;
@@ -495,10 +496,18 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
495496
}
496497
}
497498

498-
ctx.close();
499+
ctx.close(requiresPostProcessing(spanInfo));
499500
return NoopFlow.INSTANCE;
500501
}
501502

503+
private boolean requiresPostProcessing(final IGSpanInfo spanInfo) {
504+
if (!(spanInfo instanceof AgentSpan)) {
505+
return true; // be conservative
506+
}
507+
final AgentSpan span = (AgentSpan) spanInfo;
508+
return span.isRequiresPostProcessing();
509+
}
510+
502511
private Flow<Void> onRequestHeadersDone(RequestContext ctx_) {
503512
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
504513
if (ctx == null || ctx.isReqDataPublished()) {

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

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -244,4 +244,32 @@ class AppSecRequestContextSpecification extends DDSpecification {
244244
0 * rateLimiter.isThrottled()
245245
result == result2
246246
}
247+
248+
249+
void 'test that internal data is cleared on close'() {
250+
setup:
251+
final ctx = new AppSecRequestContext()
252+
final fullCleanup = !postProcessing
253+
254+
when:
255+
ctx.requestHeaders.put('Accept', ['*'])
256+
ctx.responseHeaders.put('Content-Type', ['text/plain'])
257+
ctx.collectedCookies = [cookie : ['test']]
258+
ctx.persistentData.put(KnownAddresses.REQUEST_METHOD, 'GET')
259+
ctx.derivatives = ['a': 'b']
260+
ctx.additive = createAdditive()
261+
ctx.close(postProcessing)
262+
263+
then:
264+
ctx.additive == null
265+
ctx.derivatives == null
266+
267+
ctx.requestHeaders.isEmpty() == fullCleanup
268+
ctx.responseHeaders.isEmpty() == fullCleanup
269+
ctx.cookies.isEmpty() == fullCleanup
270+
ctx.persistentData.isEmpty() == fullCleanup
271+
272+
where:
273+
postProcessing << [true, false]
274+
}
247275
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import com.datadog.appsec.event.data.DataBundle
88
import com.datadog.appsec.event.data.KnownAddresses
99
import com.datadog.appsec.report.AppSecEvent
1010
import com.datadog.appsec.report.AppSecEventWrapper
11+
import datadog.trace.api.Config
1112
import datadog.trace.api.function.TriConsumer
1213
import datadog.trace.api.function.TriFunction
1314
import datadog.trace.api.gateway.BlockResponseFunction
@@ -137,7 +138,7 @@ class GatewayBridgeSpecification extends DDSpecification {
137138
1 * spanInfo.getTags() >> ['http.client_ip':'1.1.1.1']
138139
1 * mockAppSecCtx.transferCollectedEvents() >> [event]
139140
1 * mockAppSecCtx.peerAddress >> '2001::1'
140-
1 * mockAppSecCtx.close()
141+
1 * mockAppSecCtx.close(false)
141142
1 * traceSegment.setTagTop("_dd.appsec.enabled", 1)
142143
1 * traceSegment.setTagTop("_dd.runtime_family", "jvm")
143144
1 * traceSegment.setTagTop('appsec.event', true)

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -832,6 +832,11 @@ public void addLink(AgentSpanLink link) {
832832
}
833833
}
834834

835+
@Override
836+
public boolean isRequiresPostProcessing() {
837+
return context.isRequiresPostProcessing();
838+
}
839+
835840
// to be accessible in Spock spies, which the field wouldn't otherwise be
836841
public long getStartTimeNano() {
837842
return startTimeNano;

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,8 @@ public interface AgentSpan extends MutableSpan, IGSpanInfo, ImplicitContextKeyed
141141

142142
void addLink(AgentSpanLink link);
143143

144+
boolean isRequiresPostProcessing();
145+
144146
@Override
145147
default ScopedContext storeInto(ScopedContext context) {
146148
return context.with(ScopedContextKey.SPAN_KEY, this);

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -855,6 +855,11 @@ public void addLink(AgentSpanLink link) {}
855855
public AgentSpan setMetaStruct(String field, Object value) {
856856
return this;
857857
}
858+
859+
@Override
860+
public boolean isRequiresPostProcessing() {
861+
return false;
862+
}
858863
}
859864

860865
public static final class NoopAgentScope implements AgentScope {

0 commit comments

Comments
 (0)