Fix

Author: jandro996

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

@@ -141,6 +141,7 @@ public class AppSecRequestContext implements DataBundle, Closeable {
   private boolean respDataPublished;
   private boolean pathParamsPublished;
   private volatile Map<String, Object> derivatives;
+  private final Object derivativesLock = new Object();
 
   private final AtomicBoolean rateLimited = new AtomicBoolean(false);
   private volatile boolean throttled;
@@ -649,9 +650,11 @@ public void close() {
       requestHeaders.clear();
       responseHeaders.clear();
       persistentData.clear();
-      if (derivatives != null) {
-        derivatives.clear();
-        derivatives = null;
+      synchronized (derivativesLock) {
+        if (derivatives != null) {
+          derivatives.clear();
+          derivatives = null;
+        }
       }
     }
   }
@@ -743,10 +746,7 @@ public void reportDerivatives(Map<String, Object> data) {
     log.debug("Reporting derivatives: {}", data);
     if (data == null || data.isEmpty()) return;
 
-    // Store raw derivatives
-    if (derivatives == null) {
-      derivatives = new HashMap<>();
-    }
+    Map<String, Object> newDerivatives = new LinkedHashMap<>();
 
     // Process each attribute according to the specification
     for (Map.Entry<String, Object> entry : data.entrySet()) {
@@ -762,7 +762,7 @@ public void reportDerivatives(Map<String, Object> data) {
           Object literalValue = config.get("value");
           if (literalValue != null) {
             // Preserve the original type - don't convert to string
-            derivatives.put(attributeKey, literalValue);
+            newDerivatives.put(attributeKey, literalValue);
             log.debug(
                 "Added literal attribute: {} = {} (type: {})",
                 attributeKey,
@@ -781,16 +781,25 @@ else if (config.containsKey("address")) {
           Object extractedValue = extractValueFromRequestData(address, keyPath, transformers);
           if (extractedValue != null) {
             // For extracted values, convert to string as they come from request data
-            derivatives.put(attributeKey, extractedValue.toString());
+            newDerivatives.put(attributeKey, extractedValue.toString());
             log.debug("Added extracted attribute: {} = {}", attributeKey, extractedValue);
           }
         }
       } else {
         // Handle plain string/numeric values
-        derivatives.put(attributeKey, attributeConfig);
+        newDerivatives.put(attributeKey, attributeConfig);
         log.debug("Added direct attribute: {} = {}", attributeKey, attributeConfig);
       }
     }
+
+    if (!newDerivatives.isEmpty()) {
+      synchronized (derivativesLock) {
+        if (derivatives == null) {
+          derivatives = new HashMap<>();
+        }
+        derivatives.putAll(newDerivatives);
+      }
+    }
   }
 
   /**
@@ -943,40 +952,48 @@ public boolean commitDerivatives(TraceSegment traceSegment) {
       return false;
     }
 
+    Map<String, Object> derivativesSnapshot;
+    synchronized (derivativesLock) {
+      if (derivatives == null || derivatives.isEmpty()) {
+        derivatives = null;
+        return true;
+      }
+      derivativesSnapshot = new LinkedHashMap<>(derivatives);
+      derivatives = null;
+    }
+
     // Process and commit derivatives directly
-    if (derivatives != null && !derivatives.isEmpty()) {
-      for (Map.Entry<String, Object> entry : derivatives.entrySet()) {
-        String key = entry.getKey();
-        Object value = entry.getValue();
-
-        // Handle different value types
-        if (value instanceof Number) {
-          traceSegment.setTagTop(key, (Number) value);
-        } else if (value instanceof String) {
-          // Try to parse as numeric, otherwise use as string
-          Number parsedNumber = convertToNumericAttribute((String) value);
-          if (parsedNumber != null) {
-            traceSegment.setTagTop(key, parsedNumber);
-          } else {
-            traceSegment.setTagTop(key, value);
-          }
-        } else if (value instanceof Boolean) {
-          traceSegment.setTagTop(key, value);
+    for (Map.Entry<String, Object> entry : derivativesSnapshot.entrySet()) {
+      String key = entry.getKey();
+      Object value = entry.getValue();
+
+      // Handle different value types
+      if (value instanceof Number) {
+        traceSegment.setTagTop(key, (Number) value);
+      } else if (value instanceof String) {
+        // Try to parse as numeric, otherwise use as string
+        Number parsedNumber = convertToNumericAttribute((String) value);
+        if (parsedNumber != null) {
+          traceSegment.setTagTop(key, parsedNumber);
         } else {
-          // Convert other types to string
-          traceSegment.setTagTop(key, value.toString());
+          traceSegment.setTagTop(key, value);
         }
+      } else if (value instanceof Boolean) {
+        traceSegment.setTagTop(key, value);
+      } else {
+        // Convert other types to string
+        traceSegment.setTagTop(key, value.toString());
       }
     }
 
-    // Clear all attribute maps
-    derivatives = null;
     return true;
   }
 
   // Mainly used for testing and logging
   Set<String> getDerivativeKeys() {
-    return derivatives == null ? emptySet() : new HashSet<>(derivatives.keySet());
+    synchronized (derivativesLock) {
+      return derivatives == null ? emptySet() : new HashSet<>(derivatives.keySet());
+    }
   }
 
   public boolean isThrottled(RateLimiter rateLimiter) {

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

@@ -11,11 +11,15 @@ import com.squareup.moshi.JsonAdapter
 import com.squareup.moshi.Moshi
 import com.squareup.moshi.Types
 import datadog.trace.api.Config
+import datadog.trace.api.internal.TraceSegment
 import datadog.trace.api.telemetry.LogCollector
 import datadog.trace.test.logging.TestLogCollector
 import datadog.trace.test.util.DDSpecification
 import datadog.trace.util.stacktrace.StackTraceEvent
 import datadog.trace.util.stacktrace.StackTraceFrame
+import java.util.concurrent.ConcurrentLinkedQueue
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
 import okio.Okio
 
 class AppSecRequestContextSpecification extends DDSpecification {
@@ -397,6 +401,83 @@ class AppSecRequestContextSpecification extends DDSpecification {
     keys.contains("_dd.appsec.s.res.content_type_upper")
   }
 
+  void testCommitDerivativesNormalizesAttributeTypes() {
+    given:
+    def context = new AppSecRequestContext()
+    context.reportDerivatives([
+      numeric: [value: "42"],
+      string: [value: "value"],
+      bool: true,
+      list: [value: [1, 2]]
+    ])
+    def traceSegment = Mock(TraceSegment)
+
+    when:
+    context.commitDerivatives(traceSegment)
+
+    then:
+    1 * traceSegment.setTagTop("numeric", 42L)
+    1 * traceSegment.setTagTop("string", "value")
+    1 * traceSegment.setTagTop("bool", true)
+    1 * traceSegment.setTagTop("list", "[1, 2]")
+    0 * _
+    context.getDerivativeKeys().isEmpty()
+  }
+
+  void testCommitDerivativesHandlesConcurrentUpdates() {
+    given:
+    def context = new AppSecRequestContext()
+    def start = new CountDownLatch(1)
+    def done = new CountDownLatch(6)
+    def errors = new ConcurrentLinkedQueue<Throwable>()
+    def traceSegment = TraceSegment.NoOp.INSTANCE
+
+    def reporters = []
+    for (int n = 0; n < 3; n++) {
+      reporters.add(Thread.start({
+        try {
+          start.await()
+          for (int idx = 0; idx < 100; idx++) {
+            def key = "key-${Thread.currentThread().id}-${idx}" as String
+            def val = [value: idx as String]
+            context.reportDerivatives([(key): val])
+          }
+        } catch (Throwable t) {
+          errors.add(t)
+        } finally {
+          done.countDown()
+        }
+      }))
+    }
+
+    def committers = []
+    for (int n = 0; n < 3; n++) {
+      committers.add(Thread.start({
+        try {
+          start.await()
+          for (int i = 0; i < 100; i++) {
+            context.commitDerivatives(traceSegment)
+          }
+        } catch (Throwable t) {
+          errors.add(t)
+        } finally {
+          done.countDown()
+        }
+      }))
+    }
+
+    when:
+    start.countDown()
+    def completed = done.await(10, TimeUnit.SECONDS)
+    (reporters + committers)*.join()
+    context.commitDerivatives(traceSegment)
+
+    then:
+    completed
+    errors.isEmpty()
+    context.getDerivativeKeys().isEmpty()
+  }
+
   def "test attribute handling with unknown address"() {
     given:
     def context = new AppSecRequestContext()