WIP

Author: jandro996

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

@@ -704,8 +704,11 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
         traceSeg.setDataTop("appsec", wrapper);
 
         // Report collected request and response headers based on allow list
-        writeRequestHeaders(traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders());
-        writeResponseHeaders(traceSeg, RESPONSE_HEADERS_ALLOW_LIST, ctx.getResponseHeaders());
+        boolean collectAll = Config.get().isAppSecCollectAllHeaders();
+        writeRequestHeaders(
+            traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders(), collectAll);
+        writeResponseHeaders(
+            traceSeg, RESPONSE_HEADERS_ALLOW_LIST, ctx.getResponseHeaders(), collectAll);
 
         // Report collected stack traces
         List<StackTraceEvent> stackTraces = ctx.getStackTraces();
@@ -715,10 +718,11 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
 
       } else if (hasUserInfo(traceSeg)) {
         // Report all collected request headers on user tracking event
-        writeRequestHeaders(traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders());
+        writeRequestHeaders(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());
+        writeRequestHeaders(
+            traceSeg, DEFAULT_REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders(), false);
       }
       // If extracted any derivatives - commit them
       if (!ctx.commitDerivatives(traceSeg)) {
@@ -832,29 +836,33 @@ private static boolean hasUserCollectionEvent(final TraceSegment traceSeg) {
   private static void writeRequestHeaders(
       final TraceSegment traceSeg,
       final Set<String> allowed,
-      final Map<String, List<String>> headers) {
-    writeHeaders(traceSeg, "http.request.headers.", "_dd.appsec.request.", allowed, headers);
+      final Map<String, List<String>> headers,
+      final boolean collectAll) {
+    writeHeaders(
+        traceSeg, "http.request.headers.", "_dd.appsec.request.", allowed, headers, collectAll);
   }
 
   private static void writeResponseHeaders(
       final TraceSegment traceSeg,
       final Set<String> allowed,
-      final Map<String, List<String>> headers) {
-    writeHeaders(traceSeg, "http.response.headers.", "_dd.appsec.response.", allowed, headers);
+      final Map<String, List<String>> headers,
+      final boolean collectAll) {
+    writeHeaders(
+        traceSeg, "http.response.headers.", "_dd.appsec.response.", allowed, headers, collectAll);
   }
 
   private static void writeHeaders(
       final TraceSegment traceSeg,
       final String prefix,
       final String discardedPrefix,
       final Set<String> allowed,
-      final Map<String, List<String>> headers) {
+      final Map<String, List<String>> headers,
+      final boolean collectAll) {
 
     if (headers == null || headers.isEmpty()) {
       return;
     }
 
-    final boolean collectAll = Config.get().isAppSecCollectAllHeaders();
     final int headerLimit = Config.get().getAppsecMaxCollectedHeaders();
     final Set<String> added = new HashSet<>();
     int excluded = 0;

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

@@ -1175,7 +1175,7 @@ class GatewayBridgeSpecification extends DDSpecification {
     ]
 
     when:
-    GatewayBridge.writeRequestHeaders(traceSegment, allowedHeaders, headers)
+    GatewayBridge.writeRequestHeaders(traceSegment, allowedHeaders, headers, false)
 
     then:
     1 * traceSegment.setTagTop('http.request.headers.x-allowed-header', 'value1')
@@ -1195,7 +1195,7 @@ class GatewayBridgeSpecification extends DDSpecification {
     ]
 
     when:
-    GatewayBridge.writeResponseHeaders(traceSegment, allowedHeaders, headers)
+    GatewayBridge.writeResponseHeaders(traceSegment, allowedHeaders, headers, false)
 
     then:
     1 * traceSegment.setTagTop('http.response.headers.x-allowed-header', 'value1')
@@ -1205,7 +1205,6 @@ class GatewayBridgeSpecification extends DDSpecification {
 
   void 'test  writeRequestHeaders collecting all headers '(){
     setup:
-    injectEnvConfig('DD_APPSEC_COLLECT_ALL_HEADERS' , 'true')
     injectEnvConfig('DD_APPSEC_MAX_COLLECTED_HEADERS', '4')
 
     def allowedHeaders = ['x-allowed-header', 'x-multiple-allowed-header', 'x-always-included'] as Set
@@ -1218,7 +1217,7 @@ class GatewayBridgeSpecification extends DDSpecification {
     ]
 
     when:
-    GatewayBridge.writeRequestHeaders(traceSegment, allowedHeaders, headers)
+    GatewayBridge.writeRequestHeaders(traceSegment, allowedHeaders, headers, true)
 
     then:
     1 * traceSegment.setTagTop('http.request.headers.x-allowed-header', 'value1')
@@ -1244,7 +1243,7 @@ class GatewayBridgeSpecification extends DDSpecification {
     ]
 
     when:
-    GatewayBridge.writeResponseHeaders(traceSegment, allowedHeaders, headers)
+    GatewayBridge.writeResponseHeaders(traceSegment, allowedHeaders, headers, true)
 
     then:
     1 * traceSegment.setTagTop('http.response.headers.x-allowed-header', 'value1')