Fix concurrent access to AppSec derivatives during tracing

dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java

@@ -88,7 +88,7 @@ private void extractSchemas(
       GatewayContext gwCtx = new GatewayContext(false);
       producerService.publishDataEvent(sub, ctx, bundle, gwCtx);
       // NOTE: must be checked before committing derivatives, as that clears the derivatives map
-      if (hasSchemaDerivative(ctx)) {
+      if (ctx.hasDerivativeKeyStartingWith(SCHEMA_DERIVATIVE_PREFIX)) {
         WafMetricCollector.get().apiSecurityRequestSchema(framework);
       } else {
         WafMetricCollector.get().apiSecurityRequestNoSchema(framework);
@@ -98,13 +98,4 @@ private void extractSchemas(
       log.debug("Subscriber info expired", e);
     }
   }
-
-  private static boolean hasSchemaDerivative(final AppSecRequestContext ctx) {
-    for (String key : ctx.getDerivativeKeys()) {
-      if (key != null && key.startsWith(SCHEMA_DERIVATIVE_PREFIX)) {
-        return true;
-      }
-    }
-    return false;
-  }
 }

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

@@ -1,7 +1,9 @@
 package com.datadog.appsec.gateway;
 
 import static datadog.trace.api.telemetry.LogCollector.SEND_TELEMETRY;
+import static java.util.Collections.emptyMap;
 import static java.util.Collections.emptySet;
+import static java.util.Collections.unmodifiableMap;
 
 import com.datadog.appsec.event.data.Address;
 import com.datadog.appsec.event.data.DataBundle;
@@ -127,8 +129,16 @@ public class AppSecRequestContext implements DataBundle, Closeable, AppSecContex
   private String apiSecurityFramework;
   private String endpoint;
   private boolean endpointComputed = false;
-  private final Map<String, List<String>> requestHeaders = new LinkedHashMap<>();
-  private final Map<String, List<String>> responseHeaders = new LinkedHashMap<>();
+  // Handed out live to readers that may run on the trace-processing thread (API Security schema
+  // extraction), so close() must release these by nulling the reference, never by clearing the map:
+  // mutating a snapshot mid-iteration corrupts it (APPSEC-70134). Volatile because that release can
+  // happen on a different thread than the header writes.
+  //
+  // null also means "no header seen yet", so an untouched context allocates nothing. Only
+  // addRequestHeader/addResponseHeader materialize, which assumes the single writer the plain
+  // LinkedHashMap already assumes; the getters null-coalesce so no reader allocates.
+  private volatile Map<String, List<String>> requestHeaders;
+  private volatile Map<String, List<String>> responseHeaders;
   private volatile Map<String, List<String>> collectedCookies;
   private boolean finishedRequestHeaders;
   private boolean finishedResponseHeaders;
@@ -150,6 +160,13 @@ public class AppSecRequestContext implements DataBundle, Closeable, AppSecContex
   private boolean responseBodyPublished;
   private boolean respDataPublished;
   private boolean pathParamsPublished;
+
+  /**
+   * WAF-reported attributes, published copy-on-write: writers publish a fresh unmodifiable copy
+   * instead of mutating in place, so a reader can keep the instance it read without
+   * synchronization. Mutating a published map would break readers running on the trace-processing
+   * thread (APPSEC-70134).
+   */
   private final AtomicReference<Map<String, Object>> derivatives = new AtomicReference<>();
 
   private final AtomicBoolean rateLimited = new AtomicBoolean(false);
@@ -555,9 +572,12 @@ void addRequestHeader(String name, String value) {
       return;
     }
 
-    List<String> strings =
-        requestHeaders.computeIfAbsent(name.toLowerCase(Locale.ROOT), h -> new ArrayList<>(1));
-    strings.add(value);
+    Map<String, List<String>> headers = requestHeaders;
+    if (headers == null) {
+      headers = new LinkedHashMap<>();
+      requestHeaders = headers;
+    }
+    headers.computeIfAbsent(name.toLowerCase(Locale.ROOT), h -> new ArrayList<>(1)).add(value);
   }
 
   void finishRequestHeaders() {
@@ -569,7 +589,8 @@ boolean isFinishedRequestHeaders() {
   }
 
   Map<String, List<String>> getRequestHeaders() {
-    return requestHeaders;
+    Map<String, List<String>> headers = requestHeaders;
+    return headers != null ? headers : emptyMap();
   }
 
   void addResponseHeader(String name, String value) {
@@ -581,9 +602,12 @@ void addResponseHeader(String name, String value) {
       return;
     }
 
-    List<String> strings =
-        responseHeaders.computeIfAbsent(name.toLowerCase(Locale.ROOT), h -> new ArrayList<>(1));
-    strings.add(value);
+    Map<String, List<String>> headers = responseHeaders;
+    if (headers == null) {
+      headers = new LinkedHashMap<>();
+      responseHeaders = headers;
+    }
+    headers.computeIfAbsent(name.toLowerCase(Locale.ROOT), h -> new ArrayList<>(1)).add(value);
   }
 
   public void finishResponseHeaders() {
@@ -595,7 +619,8 @@ public boolean isFinishedResponseHeaders() {
   }
 
   Map<String, List<String>> getResponseHeaders() {
-    return responseHeaders;
+    Map<String, List<String>> headers = responseHeaders;
+    return headers != null ? headers : emptyMap();
   }
 
   void addCookies(Map<String, List<String>> cookies) {
@@ -761,13 +786,11 @@ public void close() {
       // a late/async caller of getOrCreateWafContext() cannot resurrect a context (APPSEC-69085).
       closeWafContext();
       collectedCookies = null;
-      requestHeaders.clear();
-      responseHeaders.clear();
       persistentData.clear();
-      final Map<String, Object> derivatives = this.derivatives.getAndSet(null);
-      if (derivatives != null) {
-        derivatives.clear();
-      }
+      // Null the reference, never clear in place (APPSEC-70134): see the field declarations.
+      this.requestHeaders = null;
+      this.responseHeaders = null;
+      this.derivatives.set(null);
     }
   }
 
@@ -888,7 +911,8 @@ else if (config.containsKey("address")) {
             }
           }
 
-          return updated;
+          // Unmodifiable: readers may hold this instance indefinitely.
+          return unmodifiableMap(updated);
         });
   }
 
@@ -932,9 +956,9 @@ private Object getDataForAddress(String address) {
     // Map common addresses to our data structures
     switch (address) {
       case "server.request.headers":
-        return requestHeaders;
+        return getRequestHeaders();
       case "server.response.headers":
-        return responseHeaders;
+        return getResponseHeaders();
       case "server.request.cookies":
         return collectedCookies;
       case "server.request.uri.raw":
@@ -1074,10 +1098,26 @@ public boolean commitDerivatives(TraceSegment traceSegment) {
     return true;
   }
 
-  // Mainly used for testing and logging
+  /**
+   * Mainly used for testing and debug logging. The published map is unmodifiable, so its key set is
+   * returned directly as a stable view.
+   */
   public Set<String> getDerivativeKeys() {
     Map<String, Object> current = derivatives.get();
-    return current == null ? emptySet() : new HashSet<>(current.keySet());
+    return current == null ? emptySet() : current.keySet();
+  }
+
+  /**
+   * Whether any currently reported derivative key starts with the given prefix. Must be called
+   * before {@link #commitDerivatives(TraceSegment)}, which detaches the derivatives map.
+   */
+  public boolean hasDerivativeKeyStartingWith(final String prefix) {
+    for (String key : getDerivativeKeys()) {
+      if (key != null && key.startsWith(prefix)) {
+        return true;
+      }
+    }
+    return false;
   }
 
   public boolean isThrottled(RateLimiter rateLimiter) {

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/AppSecSpanPostProcessorTest.groovy

@@ -13,7 +13,7 @@ import spock.lang.Shared
 
 class AppSecSpanPostProcessorTest extends DDSpecification {
 
-  private static final String SCHEMA_DERIVATIVE_KEY = '_dd.appsec.s.req.body'
+  private static final String SCHEMA_DERIVATIVE_PREFIX = '_dd.appsec.s.'
 
   private static final String FRAMEWORK = 'netty'
 
@@ -55,7 +55,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification {
     1 * subInfo.isEmpty() >> false
     1 * ctx.getApiSecurityFramework() >> FRAMEWORK
     1 * producer.publishDataEvent(_, ctx, _, _)
-    1 * ctx.getDerivativeKeys() >> ([SCHEMA_DERIVATIVE_KEY] as Set)
+    1 * ctx.hasDerivativeKeyStartingWith(SCHEMA_DERIVATIVE_PREFIX) >> true
     1 * wafMetricCollector.apiSecurityRequestSchema(FRAMEWORK)
     1 * ctx.commitDerivatives(traceSegment)
     1 * ctx.setKeepOpenForApiSecurityPostProcessing(false)
@@ -299,7 +299,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification {
     1 * producer.getDataSubscribers(KnownAddresses.WAF_CONTEXT_PROCESSOR) >> subInfo
     1 * subInfo.isEmpty() >> false
     1 * producer.publishDataEvent(_, ctx, _, _)
-    1 * ctx.getDerivativeKeys() >> (derivativeKeys as Set)
+    1 * ctx.hasDerivativeKeyStartingWith(SCHEMA_DERIVATIVE_PREFIX) >> hasSchemaDerivative
     schemaMetrics * wafMetricCollector.apiSecurityRequestSchema(framework)
     noSchemaMetrics * wafMetricCollector.apiSecurityRequestNoSchema(framework)
     1 * ctx.commitDerivatives(traceSegment)
@@ -310,9 +310,8 @@ class AppSecSpanPostProcessorTest extends DDSpecification {
     0 * _
 
     where:
-    scenario                        | derivativeKeys                | framework || schemaMetrics | noSchemaMetrics
-    'no derivatives at all'         | []                            | FRAMEWORK || 0             | 1
-    'no derivative is a schema'     | ['_dd.appsec.fp.http.header'] | FRAMEWORK || 0             | 1
-    'missing component tag'         | [SCHEMA_DERIVATIVE_KEY]       | null      || 1             | 0
+    scenario                        | hasSchemaDerivative | framework || schemaMetrics | noSchemaMetrics
+    'no schema derivative'          | false               | FRAMEWORK || 0             | 1
+    'missing component tag'         | true                | null      || 1             | 0
   }
 }