Add new features for trace tagging rules (#9131)

Signed-off-by: sezen.leblay <[email protected]>

Author: sezen-datadog

dd-java-agent/appsec/build.gradle

@@ -14,7 +14,7 @@ dependencies {
   implementation project(':internal-api')
   implementation project(':communication')
   implementation project(':telemetry')
-  implementation group: 'io.sqreen', name: 'libsqreen', version: '15.0.1'
+  implementation group: 'io.sqreen', name: 'libsqreen', version: '16.0.0'
   implementation libs.moshi
 
   testImplementation libs.bytebuddy

dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecConfigServiceImpl.java

@@ -5,6 +5,7 @@
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_AUTO_USER_INSTRUM_MODE;
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_CUSTOM_BLOCKING_RESPONSE;
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_CUSTOM_RULES;
+import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_DD_MULTICONFIG;
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_DD_RULES;
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_EXCLUSIONS;
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_EXCLUSION_DATA;
@@ -18,6 +19,7 @@
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_RASP_SSRF;
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_REQUEST_BLOCKING;
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_SESSION_FINGERPRINT;
+import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_TRACE_TAGGING_RULES;
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_TRUSTED_IPS;
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_USER_BLOCKING;
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ENDPOINT_FINGERPRINT;
@@ -37,8 +39,9 @@
 import com.datadog.ddwaf.exception.InvalidRuleSetException;
 import com.datadog.ddwaf.exception.UnclassifiedWafException;
 import com.squareup.moshi.JsonAdapter;
+import com.squareup.moshi.JsonReader;
+import com.squareup.moshi.JsonWriter;
 import com.squareup.moshi.Moshi;
-import com.squareup.moshi.Types;
 import datadog.remoteconfig.ConfigurationEndListener;
 import datadog.remoteconfig.ConfigurationPoller;
 import datadog.remoteconfig.PollingRateHinter;
@@ -65,6 +68,7 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
+import javax.annotation.Nullable;
 import okio.Okio;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -92,10 +96,25 @@ public class AppSecConfigServiceImpl implements AppSecConfigService {
       new WAFInitializationResultReporter();
   private final WAFStatsReporter statsReporter = new WAFStatsReporter();
 
-  private static final JsonAdapter<Map<String, Object>> ADAPTER =
+  private static final JsonAdapter<Object> ADAPTER =
       new Moshi.Builder()
+          .add(
+              Double.class,
+              new JsonAdapter<Number>() {
+                @Override
+                public Number fromJson(JsonReader reader) throws IOException {
+                  double value = reader.nextDouble();
+                  long longValue = (long) value;
+                  return value % 1 == 0 ? longValue : value;
+                }
+
+                @Override
+                public void toJson(JsonWriter writer, @Nullable Number value) throws IOException {
+                  throw new UnsupportedOperationException();
+                }
+              })
           .build()
-          .adapter(Types.newParameterizedType(Map.class, String.class, Object.class));
+          .adapter(Object.class);
 
   private boolean hasUserWafConfig;
   private boolean defaultConfigActivated;
@@ -104,7 +123,7 @@ public class AppSecConfigServiceImpl implements AppSecConfigService {
       Collections.newSetFromMap(new ConcurrentHashMap<>());
   private final Set<String> ignoredConfigKeys =
       Collections.newSetFromMap(new ConcurrentHashMap<>());
-  private final String DEFAULT_WAF_CONFIG_RULE = "DEFAULT_WAF_CONFIG";
+  private final String DEFAULT_WAF_CONFIG_RULE = "ASM_DD/default";
   private String currentRuleVersion;
   private List<AppSecModule> modulesToUpdateVersionIn;
 
@@ -137,6 +156,7 @@ private void subscribeConfigurationPoller() {
   private long getRulesAndDataCapabilities() {
     long capabilities =
         CAPABILITY_ASM_DD_RULES
+            | CAPABILITY_ASM_DD_MULTICONFIG
             | CAPABILITY_ASM_IP_BLOCKING
             | CAPABILITY_ASM_EXCLUSIONS
             | CAPABILITY_ASM_EXCLUSION_DATA
@@ -148,7 +168,8 @@ private long getRulesAndDataCapabilities() {
             | CAPABILITY_ENDPOINT_FINGERPRINT
             | CAPABILITY_ASM_SESSION_FINGERPRINT
             | CAPABILITY_ASM_NETWORK_FINGERPRINT
-            | CAPABILITY_ASM_HEADER_FINGERPRINT;
+            | CAPABILITY_ASM_HEADER_FINGERPRINT
+            | CAPABILITY_ASM_TRACE_TAGGING_RULES;
     if (tracerConfig.isAppSecRaspEnabled()) {
       capabilities |= CAPABILITY_ASM_RASP_SQLI;
       capabilities |= CAPABILITY_ASM_RASP_SSRF;
@@ -210,7 +231,8 @@ public void accept(ConfigKey configKey, byte[] content, PollingRateHinter pollin
       }
       final String key = configKey.toString();
       Map<String, Object> contentMap =
-          ADAPTER.fromJson(Okio.buffer(Okio.source(new ByteArrayInputStream(content))));
+          (Map<String, Object>)
+              ADAPTER.fromJson(Okio.buffer(Okio.source(new ByteArrayInputStream(content))));
       if (contentMap == null || contentMap.isEmpty()) {
         ignoredConfigKeys.add(key);
       } else {
@@ -255,7 +277,7 @@ private class AppSecConfigChangesDDListener extends AppSecConfigChangesListener
     @Override
     protected void beforeApply(final String key, final Map<String, Object> config) {
       if (defaultConfigActivated) { // if we get any config, remove the default one
-        log.debug("Removing default config");
+        log.debug("Removing default config ASM_DD/default");
         try {
           wafBuilder.removeConfig(DEFAULT_WAF_CONFIG_RULE);
         } catch (UnclassifiedWafException e) {
@@ -466,7 +488,8 @@ private static Map<String, Object> loadDefaultWafConfig() throws IOException {
         throw new IOException("Resource " + DEFAULT_CONFIG_LOCATION + " not found");
       }
 
-      Map<String, Object> ret = ADAPTER.fromJson(Okio.buffer(Okio.source(is)));
+      Map<String, Object> ret =
+          (Map<String, Object>) ADAPTER.fromJson(Okio.buffer(Okio.source(is)));
 
       StandardizedLogging._initialConfigSourceAndLibddwafVersion(log, "<bundled config>");
       if (log.isInfoEnabled()) {
@@ -483,7 +506,8 @@ private static Map<String, Object> loadUserWafConfig(Config tracerConfig) throws
       return null;
     }
     try (InputStream is = new FileInputStream(filename)) {
-      Map<String, Object> ret = ADAPTER.fromJson(Okio.buffer(Okio.source(is)));
+      Map<String, Object> ret =
+          (Map<String, Object>) ADAPTER.fromJson(Okio.buffer(Okio.source(is)));
 
       StandardizedLogging._initialConfigSourceAndLibddwafVersion(log, filename);
       if (log.isInfoEnabled()) {
@@ -512,6 +536,7 @@ public void close() {
     this.configurationPoller.removeCapabilities(
         CAPABILITY_ASM_ACTIVATION
             | CAPABILITY_ASM_DD_RULES
+            | CAPABILITY_ASM_DD_MULTICONFIG
             | CAPABILITY_ASM_IP_BLOCKING
             | CAPABILITY_ASM_EXCLUSIONS
             | CAPABILITY_ASM_EXCLUSION_DATA
@@ -529,7 +554,8 @@ public void close() {
             | CAPABILITY_ENDPOINT_FINGERPRINT
             | CAPABILITY_ASM_SESSION_FINGERPRINT
             | CAPABILITY_ASM_NETWORK_FINGERPRINT
-            | CAPABILITY_ASM_HEADER_FINGERPRINT);
+            | CAPABILITY_ASM_HEADER_FINGERPRINT
+            | CAPABILITY_ASM_TRACE_TAGGING_RULES);
     this.configurationPoller.removeListeners(Product.ASM_DD);
     this.configurationPoller.removeListeners(Product.ASM_DATA);
     this.configurationPoller.removeListeners(Product.ASM);

dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java

@@ -34,6 +34,7 @@
 import datadog.trace.api.ProductActivation;
 import datadog.trace.api.ProductTraceSource;
 import datadog.trace.api.gateway.Flow;
+import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.api.telemetry.LogCollector;
 import datadog.trace.api.telemetry.WafMetricCollector;
 import datadog.trace.api.time.SystemTimeSource;
@@ -402,12 +403,13 @@ public void onDataAvailable(
           }
         }
         Collection<AppSecEvent> events = buildEvents(resultWithData);
+        boolean isThrottled = reqCtx.isThrottled(rateLimiter);
 
-        if (!events.isEmpty()) {
-          if (!reqCtx.isThrottled(rateLimiter)) {
+        if (resultWithData.keep) {
+          if (!isThrottled) {
             AgentSpan activeSpan = AgentTracer.get().activeSpan();
             if (activeSpan != null) {
-              log.debug("Setting force-keep tag on the current span");
+              log.debug("Setting force-keep tag and manual keep tag on the current span");
               // Keep event related span, because it could be ignored in case of
               // reduced datadog sampling rate.
               activeSpan.getLocalRootSpan().setTag(Tags.ASM_KEEP, true);
@@ -418,18 +420,19 @@ public void onDataAvailable(
                   .getLocalRootSpan()
                   .setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM);
             } else {
-              // If active span is not available the ASM_KEEP tag will be set in the GatewayBridge
-              // when the request ends
+              // If active span is not available then we need to set manual keep in GatewayBridge
               log.debug("There is no active span available");
             }
-            reqCtx.reportEvents(events);
           } else {
             log.debug("Rate limited WAF events");
             if (!gwCtx.isRasp) {
               reqCtx.setWafRateLimited();
             }
           }
         }
+        if (resultWithData.events && !events.isEmpty() && !isThrottled) {
+          reqCtx.reportEvents(events);
+        }
 
         if (flow.isBlocking()) {
           if (!gwCtx.isRasp) {
@@ -438,8 +441,11 @@ public void onDataAvailable(
         }
       }
 
-      if (resultWithData.derivatives != null) {
-        reqCtx.reportDerivatives(resultWithData.derivatives);
+      reqCtx.setKeepType(
+          resultWithData.keep ? PrioritySampling.USER_KEEP : PrioritySampling.USER_DROP);
+
+      if (resultWithData.attributes != null && !resultWithData.attributes.isEmpty()) {
+        reqCtx.reportDerivatives(resultWithData.attributes);
       }
     }
 
@@ -564,6 +570,10 @@ private Collection<AppSecEvent> buildEvents(Waf.ResultWithData actionWithData) {
     }
     Collection<WAFResultData> listResults;
     try {
+      if (actionWithData.data == null || actionWithData.data.isEmpty()) {
+        log.debug("WAF returned no data");
+        return emptyList();
+      }
       listResults = RES_JSON_ADAPTER.fromJson(actionWithData.data);
     } catch (IOException e) {
       throw new UndeclaredThrowableException(e);

dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/ObjectIntrospection.java

@@ -16,6 +16,7 @@
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Date;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
@@ -184,6 +185,11 @@ private static Object doConversion(Object obj, int depth, State state) {
       return obj.toString();
     }
 
+    // Date objects - avoid accessing private fastTime field
+    if (obj instanceof Date) {
+      return ((Date) obj).getTime();
+    }
+
     // Jackson databind nodes (via reflection)
     Class<?> clazz = obj.getClass();
     if (clazz.getName().startsWith("com.fasterxml.jackson.databind.node.")) {