Update Extended appsec request/response headers and request body collection

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

@@ -9,6 +9,7 @@
 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;
+import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_EXTENDED_DATA_COLLECTION;
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_HEADER_FINGERPRINT;
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_IP_BLOCKING;
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_NETWORK_FINGERPRINT;
@@ -168,7 +169,8 @@ private long getRulesAndDataCapabilities() {
             | CAPABILITY_ASM_SESSION_FINGERPRINT
             | CAPABILITY_ASM_NETWORK_FINGERPRINT
             | CAPABILITY_ASM_HEADER_FINGERPRINT
-            | CAPABILITY_ASM_TRACE_TAGGING_RULES;
+            | CAPABILITY_ASM_TRACE_TAGGING_RULES
+            | CAPABILITY_ASM_EXTENDED_DATA_COLLECTION;
     if (tracerConfig.isAppSecRaspEnabled()) {
       capabilities |= CAPABILITY_ASM_RASP_SQLI;
       capabilities |= CAPABILITY_ASM_RASP_SSRF;
@@ -554,7 +556,8 @@ public void close() {
             | CAPABILITY_ASM_SESSION_FINGERPRINT
             | CAPABILITY_ASM_NETWORK_FINGERPRINT
             | CAPABILITY_ASM_HEADER_FINGERPRINT
-            | CAPABILITY_ASM_TRACE_TAGGING_RULES);
+            | CAPABILITY_ASM_TRACE_TAGGING_RULES
+            | CAPABILITY_ASM_EXTENDED_DATA_COLLECTION);
     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

@@ -395,6 +395,26 @@ public void onDataAvailable(
             } else {
               log.debug("Ignoring action with type generate_stack (disabled by config)");
             }
+          } else if ("extended_data_collection".equals(actionInfo.type)) {
+            // Extended data collection is handled by the GatewayBridge
+            reqCtx.setExtendedDataCollection(true);
+            // Handle max_collected_headers parameter which can come as Number or String
+            // representation of a number
+            int maxHeaders = AppSecRequestContext.DEFAULT_EXTENDED_DATA_COLLECTION_MAX_HEADERS;
+            Object maxHeadersParam =
+                actionInfo.parameters.getOrDefault(
+                    "max_collected_headers",
+                    AppSecRequestContext.DEFAULT_EXTENDED_DATA_COLLECTION_MAX_HEADERS);
+            if (maxHeadersParam instanceof Number) {
+              maxHeaders = ((Number) maxHeadersParam).intValue();
+            } else if (maxHeadersParam instanceof String) {
+              try {
+                maxHeaders = Integer.parseInt((String) maxHeadersParam);
+              } catch (NumberFormatException e) {
+                log.debug("Failed to parse max_collected_headers value: {}", maxHeadersParam);
+              }
+            }
+            reqCtx.setExtendedDataCollectionMaxHeaders(maxHeaders);
           } else {
             log.info("Ignoring action with type {}", actionInfo.type);
             if (!gwCtx.isRasp) {

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

@@ -31,6 +31,8 @@
 public class AppSecRequestContext implements DataBundle, Closeable {
   private static final Logger log = LoggerFactory.getLogger(AppSecRequestContext.class);
 
+  public static final int DEFAULT_EXTENDED_DATA_COLLECTION_MAX_HEADERS = 50;
+
   // Values MUST be lowercase! Lookup with Ignore Case
   // was removed due performance reason
   // request headers that will always be set when appsec is enabled
@@ -77,6 +79,19 @@ public class AppSecRequestContext implements DataBundle, Closeable {
       new TreeSet<>(
           Arrays.asList("content-length", "content-type", "content-encoding", "content-language"));
 
+  // headers related with authorization
+  public static final Set<String> AUTHORIZATION_HEADERS =
+      new TreeSet<>(
+          Arrays.asList(
+              "authorization",
+              "proxy-authorization",
+              "www-authenticate",
+              "proxy-authenticate",
+              "authentication-info",
+              "proxy-authentication-info",
+              "cookie",
+              "set-cookie"));
+
   static {
     REQUEST_HEADERS_ALLOW_LIST.addAll(DEFAULT_REQUEST_HEADERS_ALLOW_LIST);
   }
@@ -99,6 +114,9 @@ public class AppSecRequestContext implements DataBundle, Closeable {
   private int peerPort;
   private String inferredClientIp;
 
+  private boolean extendedDataCollection = false;
+  private int extendedDataCollectionMaxHeaders = DEFAULT_EXTENDED_DATA_COLLECTION_MAX_HEADERS;
+
   private volatile StoredBodySupplier storedRequestBodySupplier;
   private String dbType;
 
@@ -133,6 +151,7 @@ public class AppSecRequestContext implements DataBundle, Closeable {
   private volatile int raspTimeouts;
 
   private volatile Object processedRequestBody;
+  private volatile boolean processedResponseBodySizeExceeded;
   private volatile boolean raspMatched;
 
   // keep a reference to the last published usr.id
@@ -266,6 +285,22 @@ public int getRaspTimeouts() {
     return raspTimeouts;
   }
 
+  public boolean isExtendedDataCollection() {
+    return extendedDataCollection;
+  }
+
+  public void setExtendedDataCollection(boolean extendedDataCollection) {
+    this.extendedDataCollection = extendedDataCollection;
+  }
+
+  public int getExtendedDataCollectionMaxHeaders() {
+    return extendedDataCollectionMaxHeaders;
+  }
+
+  public void setExtendedDataCollectionMaxHeaders(int extendedDataCollectionMaxHeaders) {
+    this.extendedDataCollectionMaxHeaders = extendedDataCollectionMaxHeaders;
+  }
+
   public WafContext getOrCreateWafContext(
       WafHandle wafHandle, boolean createMetrics, boolean isRasp) {
     if (createMetrics) {
@@ -964,6 +999,14 @@ public Object getProcessedRequestBody() {
     return processedRequestBody;
   }
 
+  public boolean isProcessedResponseBodySizeExceeded() {
+    return processedResponseBodySizeExceeded;
+  }
+
+  public void setProcessedResponseBodySizeExceeded(boolean processedResponseBodySizeExceeded) {
+    this.processedResponseBodySizeExceeded = processedResponseBodySizeExceeded;
+  }
+
   public boolean isRaspMatched() {
     return raspMatched;
   }

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

@@ -3,6 +3,7 @@
 import static com.datadog.appsec.event.data.MapDataBundle.Builder.CAPACITY_0_2;
 import static com.datadog.appsec.event.data.MapDataBundle.Builder.CAPACITY_3_4;
 import static com.datadog.appsec.event.data.MapDataBundle.Builder.CAPACITY_6_10;
+import static com.datadog.appsec.gateway.AppSecRequestContext.AUTHORIZATION_HEADERS;
 import static com.datadog.appsec.gateway.AppSecRequestContext.DEFAULT_REQUEST_HEADERS_ALLOW_LIST;
 import static com.datadog.appsec.gateway.AppSecRequestContext.REQUEST_HEADERS_ALLOW_LIST;
 import static com.datadog.appsec.gateway.AppSecRequestContext.RESPONSE_HEADERS_ALLOW_LIST;
@@ -705,14 +706,9 @@ private Flow<Void> onRequestBodyProcessed(RequestContext ctx_, Object obj) {
               obj,
               ctx,
               () -> {
-                if (Config.get().isAppSecRaspCollectRequestBody()) {
-                  ctx_.getTraceSegment()
-                      .setTagTop("_dd.appsec.rasp.request_body_size.exceeded", true);
-                }
+                ctx.setProcessedResponseBodySizeExceeded(true);
               });
-      if (Config.get().isAppSecRaspCollectRequestBody()) {
-        ctx.setProcessedRequestBody(converted);
-      }
+      ctx.setProcessedRequestBody(converted);
       DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.REQUEST_BODY_OBJECT, converted);
       try {
         GatewayContext gwCtx = new GatewayContext(false);
@@ -885,35 +881,34 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
         traceSeg.setDataTop("appsec", wrapper);
 
         // Report collected request and response headers based on allow list
-        boolean collectAll =
-            Config.get().isAppSecCollectAllHeaders()
-                // Until redaction is defined we don't want to collect all headers due to risk of
-                // leaking sensitive data
-                && !Config.get().isAppSecHeaderCollectionRedactionEnabled();
+        boolean collectAll = ctx.isExtendedDataCollection();
         writeRequestHeaders(
-            traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders(), collectAll);
+            ctx, traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders(), collectAll);
         writeResponseHeaders(
-            traceSeg, RESPONSE_HEADERS_ALLOW_LIST, ctx.getResponseHeaders(), collectAll);
+            ctx, traceSeg, RESPONSE_HEADERS_ALLOW_LIST, ctx.getResponseHeaders(), collectAll);
 
         // Report collected stack traces
         List<StackTraceEvent> stackTraces = ctx.getStackTraces();
         if (stackTraces != null && !stackTraces.isEmpty()) {
           StackUtils.addStacktraceEventsToMetaStruct(ctx_, METASTRUCT_EXPLOIT, stackTraces);
         }
 
-        // Report collected parsed request body if there is a RASP event
-        if (ctx.isRaspMatched() && ctx.getProcessedRequestBody() != null) {
+        if (ctx.isExtendedDataCollection() && ctx.getProcessedRequestBody() != null) {
           ctx_.getOrCreateMetaStructTop(
               METASTRUCT_REQUEST_BODY, k -> ctx.getProcessedRequestBody());
+          if (ctx.isProcessedResponseBodySizeExceeded()) {
+            traceSeg.setTagTop("_dd.appsec.request_body_size.exceeded", true);
+          }
         }
 
       } else if (hasUserInfo(traceSeg)) {
         // Report all collected request headers on user tracking event
-        writeRequestHeaders(traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders(), false);
+        writeRequestHeaders(
+            ctx, traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders(), false);
       } else {
         // Report minimum set of collected request headers
         writeRequestHeaders(
-            traceSeg, DEFAULT_REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders(), false);
+            ctx, traceSeg, DEFAULT_REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders(), false);
       }
       // If extracted any derivatives - commit them
       if (!ctx.commitDerivatives(traceSeg)) {
@@ -1026,42 +1021,60 @@ private static boolean hasUserCollectionEvent(final TraceSegment traceSeg) {
   }
 
   private static void writeRequestHeaders(
+      AppSecRequestContext ctx,
       final TraceSegment traceSeg,
       final Set<String> allowed,
       final Map<String, List<String>> headers,
       final boolean collectAll) {
     writeHeaders(
-        traceSeg, "http.request.headers.", "_dd.appsec.request.", allowed, headers, collectAll);
+        ctx,
+        traceSeg,
+        "http.request.headers.",
+        "_dd.appsec.request.",
+        allowed,
+        headers,
+        collectAll,
+        true);
   }
 
   private static void writeResponseHeaders(
+      AppSecRequestContext ctx,
       final TraceSegment traceSeg,
       final Set<String> allowed,
       final Map<String, List<String>> headers,
       final boolean collectAll) {
     writeHeaders(
-        traceSeg, "http.response.headers.", "_dd.appsec.response.", allowed, headers, collectAll);
+        ctx,
+        traceSeg,
+        "http.response.headers.",
+        "_dd.appsec.response.",
+        allowed,
+        headers,
+        collectAll,
+        false);
   }
 
   private static void writeHeaders(
+      AppSecRequestContext ctx,
       final TraceSegment traceSeg,
       final String prefix,
       final String discardedPrefix,
       final Set<String> allowed,
       final Map<String, List<String>> headers,
-      final boolean collectAll) {
+      final boolean collectAll,
+      final boolean checkCookie) {
 
     if (headers == null || headers.isEmpty()) {
       return;
     }
 
-    final int headerLimit = Config.get().getAppsecMaxCollectedHeaders();
+    final int headerLimit = ctx.getExtendedDataCollectionMaxHeaders();
     final Set<String> added = new HashSet<>();
     int excluded = 0;
 
     // Try to add allowed headers (prioritized)
     for (String name : allowed) {
-      if (added.size() >= headerLimit) {
+      if (collectAll && added.size() >= headerLimit) {
         break;
       }
       List<String> values = headers.get(name);
@@ -1086,14 +1099,20 @@ private static void writeHeaders(
           excluded++;
           continue;
         }
-
-        List<String> values = entry.getValue();
-        String joined = String.join(",", values);
+        String joined;
+        if (AUTHORIZATION_HEADERS.contains(name)) {
+          joined = String.join(",", "<redacted>");
+        } else {
+          joined = String.join(",", entry.getValue());
+        }
         if (!joined.isEmpty()) {
           traceSeg.setTagTop(prefix + name, joined);
           added.add(name);
         }
       }
+      if (checkCookie && !ctx.getCookies().isEmpty()) {
+        traceSeg.setTagTop(prefix + "cookie", "<redacted>");
+      }
 
       if (excluded > 0) {
         traceSeg.setTagTop(discardedPrefix + "header_collection.discarded", excluded);

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/config/AppSecConfigServiceImplSpecification.groovy

@@ -2,6 +2,7 @@ package com.datadog.appsec.config
 
 import com.datadog.appsec.AppSecSystem
 import com.datadog.appsec.util.AbortStartupException
+import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_EXTENDED_DATA_COLLECTION
 import datadog.remoteconfig.ConfigurationChangesTypedListener
 import datadog.remoteconfig.ConfigurationDeserializer
 import datadog.remoteconfig.ConfigurationEndListener
@@ -289,7 +290,8 @@ class AppSecConfigServiceImplSpecification extends DDSpecification {
       | CAPABILITY_ASM_SESSION_FINGERPRINT
       | CAPABILITY_ASM_NETWORK_FINGERPRINT
       | CAPABILITY_ASM_HEADER_FINGERPRINT
-      | CAPABILITY_ASM_TRACE_TAGGING_RULES)
+      | CAPABILITY_ASM_TRACE_TAGGING_RULES
+      | CAPABILITY_ASM_EXTENDED_DATA_COLLECTION)
     0 * poller._
 
     when:
@@ -444,7 +446,8 @@ class AppSecConfigServiceImplSpecification extends DDSpecification {
       | CAPABILITY_ASM_SESSION_FINGERPRINT
       | CAPABILITY_ASM_NETWORK_FINGERPRINT
       | CAPABILITY_ASM_HEADER_FINGERPRINT
-      | CAPABILITY_ASM_TRACE_TAGGING_RULES)
+      | CAPABILITY_ASM_TRACE_TAGGING_RULES
+      | CAPABILITY_ASM_EXTENDED_DATA_COLLECTION)
     0 * poller._
 
     when:
@@ -540,7 +543,8 @@ class AppSecConfigServiceImplSpecification extends DDSpecification {
       | CAPABILITY_ASM_SESSION_FINGERPRINT
       | CAPABILITY_ASM_NETWORK_FINGERPRINT
       | CAPABILITY_ASM_HEADER_FINGERPRINT
-      | CAPABILITY_ASM_TRACE_TAGGING_RULES)
+      | CAPABILITY_ASM_TRACE_TAGGING_RULES
+      | CAPABILITY_ASM_EXTENDED_DATA_COLLECTION)
     4 * poller.removeListeners(_)
     1 * poller.removeConfigurationEndListener(_)
     1 * poller.stop()