Add new features for trace tagging rules

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

Author: sezen-datadog

dd-java-agent/appsec/build.gradle

@@ -15,7 +15,7 @@ dependencies {
   implementation project(':internal-api')
   implementation project(':communication')
   implementation project(':telemetry')
-  implementation group: 'io.sqreen', name: 'libsqreen', version: '15.0.0'
+  implementation group: 'io.sqreen', name: 'libsqreen', version: '15.0.2'
   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;
@@ -36,8 +38,8 @@
 import com.datadog.ddwaf.exception.InvalidRuleSetException;
 import com.datadog.ddwaf.exception.UnclassifiedWafException;
 import com.squareup.moshi.JsonAdapter;
-import com.squareup.moshi.Moshi;
-import com.squareup.moshi.Types;
+import com.squareup.moshi.JsonReader;
+import com.squareup.moshi.JsonWriter;
 import datadog.remoteconfig.ConfigurationEndListener;
 import datadog.remoteconfig.ConfigurationPoller;
 import datadog.remoteconfig.PollingRateHinter;
@@ -58,6 +60,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -90,15 +93,12 @@ public class AppSecConfigServiceImpl implements AppSecConfigService {
       new WAFInitializationResultReporter();
   private final WAFStatsReporter statsReporter = new WAFStatsReporter();
 
-  private static final JsonAdapter<Map<String, Object>> ADAPTER =
-      new Moshi.Builder()
-          .build()
-          .adapter(Types.newParameterizedType(Map.class, String.class, Object.class));
+  private static final JsonAdapter<Object> ADAPTER = new SafeMapAdapter();
 
   private boolean hasUserWafConfig;
   private boolean defaultConfigActivated;
   private final Set<String> usedDDWafConfigKeys = new HashSet<>();
-  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;
 
@@ -129,6 +129,7 @@ private void subscribeConfigurationPoller() {
 
     long capabilities =
         CAPABILITY_ASM_DD_RULES
+            | CAPABILITY_ASM_DD_MULTICONFIG
             | CAPABILITY_ASM_IP_BLOCKING
             | CAPABILITY_ASM_EXCLUSIONS
             | CAPABILITY_ASM_EXCLUSION_DATA
@@ -140,7 +141,8 @@ private void subscribeConfigurationPoller() {
             | 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;
@@ -185,7 +187,8 @@ public void accept(ConfigKey configKey, byte[] content, PollingRateHinter pollin
         }
       } else {
         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))));
         try {
           handleWafUpdateResultReport(configKey.toString(), contentMap);
         } catch (AppSecModule.AppSecModuleActivationException e) {
@@ -211,7 +214,7 @@ private class AppSecConfigChangesDDListener extends AppSecConfigChangesListener
     public void accept(ConfigKey configKey, byte[] content, PollingRateHinter pollingRateHinter)
         throws IOException {
       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) {
@@ -306,7 +309,10 @@ private void subscribeAsmFeatures() {
   private void distributeSubConfigurations(
       String key, AppSecModuleConfigurer.Reconfiguration reconfiguration) {
     if (usedDDWafConfigKeys.isEmpty() && !defaultConfigActivated && !hasUserWafConfig) {
-      // no config left in the WAF builder, add the default config
+      // ASM_DD Failure Fallback: If none of the configurations obtained through ASM_DD were loaded
+      // successfully,
+      // libraries must revert back to the default configuration
+      log.debug("No ASM_DD configurations loaded, falling back to default configuration");
       init();
     }
     for (Map.Entry<String, SubconfigListener> entry : subconfigListeners.entrySet()) {
@@ -427,7 +433,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()) {
@@ -444,7 +451,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()) {
@@ -473,6 +481,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
@@ -490,7 +499,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);
@@ -558,4 +568,48 @@ private static WafConfig createWafConfig(Config config) {
     }
     return wafConfig;
   }
+
+  private static class SafeMapAdapter extends JsonAdapter<Object> {
+    @Override
+    public Object fromJson(JsonReader reader) throws IOException {
+      switch (reader.peek()) {
+        case BEGIN_OBJECT:
+          Map<String, Object> map = new LinkedHashMap<>();
+          reader.beginObject();
+          while (reader.hasNext()) {
+            map.put(reader.nextName(), fromJson(reader));
+          }
+          reader.endObject();
+          return map;
+
+        case BEGIN_ARRAY:
+          List<Object> list = new ArrayList<>();
+          reader.beginArray();
+          while (reader.hasNext()) {
+            list.add(fromJson(reader));
+          }
+          reader.endArray();
+          return list;
+
+        case STRING:
+        case NUMBER:
+          return reader.nextString();
+
+        case BOOLEAN:
+          return reader.nextBoolean();
+
+        case NULL:
+          reader.nextNull();
+          return null;
+
+        default:
+          throw new IllegalStateException("Unexpected token: " + reader.peek());
+      }
+    }
+
+    @Override
+    public void toJson(JsonWriter writer, Object value) throws IOException {
+      throw new UnsupportedOperationException("Serialization not supported");
+    }
+  }
 }

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

@@ -30,6 +30,7 @@
 import datadog.communication.monitor.Counter;
 import datadog.communication.monitor.Monitoring;
 import datadog.trace.api.Config;
+import datadog.trace.api.DDTags;
 import datadog.trace.api.ProductActivation;
 import datadog.trace.api.ProductTraceSource;
 import datadog.trace.api.gateway.Flow;
@@ -223,6 +224,11 @@ private void initOrUpdateWafHandle(AppSecModuleConfigurer.Reconfiguration reconf
     reconf.reloadSubscriptions();
   }
 
+  /**
+   * Creates a rate limiter for AppSec events. The rate limiter accounts for when libddwaf returns
+   * keep with a value of true, rather than when events are present, as specified in the technical
+   * specification.
+   */
   private static RateLimiter getRateLimiter(Monitoring monitoring) {
     if (monitoring == null) {
       return null;
@@ -401,15 +407,18 @@ 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);
+              // Set manual keep tag as required by the technical spec
+              activeSpan.getLocalRootSpan().setTag(DDTags.MANUAL_KEEP, true);
               // If APM is disabled, inform downstream services that the current
               // distributed trace contains at least one ASM event and must inherit
               // the given force-keep priority
@@ -421,14 +430,16 @@ public void onDataAvailable(
               // when the request ends
               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) {
@@ -437,8 +448,8 @@ public void onDataAvailable(
         }
       }
 
-      if (resultWithData.derivatives != null) {
-        reqCtx.reportDerivatives(resultWithData.derivatives);
+      if (resultWithData.attributes != null && !resultWithData.attributes.isEmpty()) {
+        reqCtx.reportDerivatives(resultWithData.attributes);
       }
     }
 
@@ -559,6 +570,10 @@ private Waf.ResultWithData runWafTransient(
   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);