Collect WAF headers on user sdk events

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

@@ -62,6 +62,11 @@ public class GatewayBridge {
   private static final Pattern QUERY_PARAM_SPLITTER = Pattern.compile("&");
   private static final Map<String, List<String>> EMPTY_QUERY_PARAMS = Collections.emptyMap();
 
+  /** User tracking tags that will force the collection of request headers */
+  private static final String[] USER_TRACKING_TAGS = {
+    "appsec.events.users.login.success.track", "appsec.events.users.login.failure.track"
+  };
+
   private final SubscriptionService subscriptionService;
   private final EventProducerService producerService;
   private final RateLimiter rateLimiter;
@@ -155,6 +160,9 @@ public void init() {
                 writeRequestHeaders(traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders());
                 writeResponseHeaders(
                     traceSeg, RESPONSE_HEADERS_ALLOW_LIST, ctx.getResponseHeaders());
+              } else if (hasUserTrackingEvent(traceSeg)) {
+                // Report all collected request headers on user tracking event
+                writeRequestHeaders(traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders());
               } else {
                 // Report minimum set of collected request headers
                 writeRequestHeaders(
@@ -411,6 +419,16 @@ public void stop() {
     subscriptionService.reset();
   }
 
+  private static boolean hasUserTrackingEvent(final TraceSegment traceSeg) {
+    for (String tagName : USER_TRACKING_TAGS) {
+      final Object value = traceSeg.getTagTop(tagName);
+      if (value != null && "true".equalsIgnoreCase(value.toString())) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   private static void writeRequestHeaders(
       final TraceSegment traceSeg,
       final Set<String> allowed,

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

@@ -804,4 +804,32 @@ class GatewayBridgeSpecification extends DDSpecification {
     1 * traceSegment.setTagTop('http.request.headers.x-sigsci-tags', 'SQLI, XSS')
     1 * traceSegment.setTagTop('http.request.headers.akamai-user-risk', 'uuid=913c4545-757b-4d8d-859d-e1361a828361;status=0')
   }
+
+  void 'request headers are always set when there are user tracking events'() {
+    given:
+    final mockAppSecCtx = Stub(AppSecRequestContext) {
+      transferCollectedEvents() >> []
+      getRequestHeaders() >> [
+        'host': ['localhost']
+      ]
+    }
+    final mockCtx = Stub(RequestContext) {
+      getData(RequestContextSlot.APPSEC) >> mockAppSecCtx
+      getTraceSegment() >> traceSegment
+    }
+    final spanInfo = Stub(AgentSpan)
+    traceSegment.getTagTop(tag) >> true
+
+    when:
+    requestEndedCB.apply(mockCtx, spanInfo)
+
+    then:
+    (userTracking ? 1 : 0) * traceSegment.setTagTop('http.request.headers.host', 'localhost')
+
+    where:
+    tag                                       | userTracking
+    'appsec.events.users.login.success.track' | true
+    'appsec.events.users.login.failure.track' | true
+    'appsec.another.unrelated.tag'            | false
+  }
 }

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/AgentTestRunner.groovy

@@ -522,6 +522,18 @@ abstract class AgentTestRunner extends DDSpecification implements AgentBuilder.L
       delegate.setTagCurrent(key, value, sanitize)
     }
 
+    @Override
+    Object getTagTop(String key, boolean sanitize) {
+      check()
+      return delegate.getTagTop(key, sanitize)
+    }
+
+    @Override
+    Object getTagCurrent(String key, boolean sanitize) {
+      check()
+      return delegate.getTagCurrent(key, sanitize)
+    }
+
     @Override
     void setDataTop(String key, Object value) {
       check()

dd-trace-api/src/main/java/datadog/trace/api/internal/TraceSegment.java

@@ -27,6 +27,23 @@ default void setTagTop(String key, Object value) {
    */
   void setTagTop(String key, Object value, boolean sanitize);
 
+  /**
+   * Get the tag value from the top of this {@code TraceSegment}.
+   *
+   * @param key key of the tag
+   */
+  default Object getTagTop(String key) {
+    return getTagTop(key, false);
+  }
+
+  /**
+   * Get the tag value from the top of this {@code TraceSegment} with optional key sanitization.
+   *
+   * @param key key of the tag
+   * @param sanitize indicates is key need to be sanitized
+   */
+  Object getTagTop(String key, boolean sanitize);
+
   /**
    * Add a tag to the current span in this {@code TraceSegment}.
    *
@@ -46,6 +63,24 @@ default void setTagCurrent(String key, Object value) {
    */
   void setTagCurrent(String key, Object value, boolean sanitize);
 
+  /**
+   * Get the tag value from the current span in this {@code TraceSegment}.
+   *
+   * @param key key of the tag
+   */
+  default Object getTagCurrent(String key) {
+    return getTagCurrent(key, false);
+  }
+
+  /**
+   * Get the tag value from the current span in this {@code TraceSegment}. with optional key
+   * sanitization.
+   *
+   * @param key key of the tag
+   * @param sanitize indicates is key need to be sanitized
+   */
+  Object getTagCurrent(String key, boolean sanitize);
+
   /**
    * Add data to the top of this {@code TraceSegment}. The {@code toString} representation of the
    * {@code value} must be valid top level JSON, i.e. an {@code Object} or an {@code Array}.
@@ -94,6 +129,16 @@ public void setTagTop(String key, Object value, boolean sanitize) {}
     @Override
     public void setTagCurrent(String key, Object value, boolean sanitize) {}
 
+    @Override
+    public Object getTagTop(String key, boolean sanitize) {
+      return null;
+    }
+
+    @Override
+    public Object getTagCurrent(String key, boolean sanitize) {
+      return null;
+    }
+
     @Override
     public void setDataTop(String key, Object value) {}
 

dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java

@@ -908,6 +908,11 @@ public void setTagTop(String key, Object value, boolean sanitize) {
     getRootSpanContextOrThis().setTagCurrent(key, value, sanitize);
   }
 
+  @Override
+  public Object getTagTop(String key, boolean sanitize) {
+    return getRootSpanContextOrThis().getTagCurrent(key, sanitize);
+  }
+
   @Override
   public void setTagCurrent(String key, Object value, boolean sanitize) {
     if (sanitize) {
@@ -916,6 +921,14 @@ public void setTagCurrent(String key, Object value, boolean sanitize) {
     this.setTag(key, value);
   }
 
+  @Override
+  public Object getTagCurrent(String key, boolean sanitize) {
+    if (sanitize) {
+      key = TagsHelper.sanitize(key);
+    }
+    return this.getTag(key);
+  }
+
   @Override
   public void setDataTop(String key, Object value) {
     getRootSpanContextOrThis().setDataCurrent(key, value);