Add request headers to span when there's a user tracking event

DataDog · May 10, 2024 · 46284b6 · 46284b6
1 parent ab74adc
commit 46284b6
Show file tree

Hide file tree

Showing 5 changed files with 124 additions and 3 deletions.
diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java b/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 List<String> USER_TRACKING_TAGS =
+      Arrays.asList(
+          "appsec.events.users.login.success.track", "appsec.events.users.login.failure.track");
+
   private final SubscriptionService subscriptionService;
   private final EventProducerService producerService;
   private final RateLimiter rateLimiter;
@@ -151,12 +156,14 @@ public void init() {
                 AppSecEventWrapper wrapper = new AppSecEventWrapper(collectedEvents);
                 traceSeg.setDataTop("appsec", wrapper);
 
-                // Report collected request and response headers based on allow list
-                writeRequestHeaders(traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders());
+                // Report collected response headers based on allow list
                 writeResponseHeaders(
                     traceSeg, RESPONSE_HEADERS_ALLOW_LIST, ctx.getResponseHeaders());
+              }
+
+              if (hasAppSecEvents(traceSeg)) {
+                writeRequestHeaders(traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders());
               } else {
-                // Report minimum set of collected request headers
                 writeRequestHeaders(
                     traceSeg, DEFAULT_REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders());
               }
@@ -411,6 +418,22 @@ public void stop() {
     subscriptionService.reset();
   }
 
+  private static boolean hasAppSecEvents(final TraceSegment traceSeg) {
+    if (isTruthy(traceSeg.getDataTop("appsec.event"))) {
+      return true;
+    }
+    for (String tagName : USER_TRACKING_TAGS) {
+      if (isTruthy(traceSeg.getTagTop(tagName))) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  private static boolean isTruthy(final Object value) {
+    return value != null && "true".equals(value.toString());
+  }
+
   private static void writeRequestHeaders(
       final TraceSegment traceSeg,
       final Set<String> allowed,

diff --git a/...agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/...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
+  }
 }
diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/AgentTestRunner.groovy b/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()

diff --git a/dd-trace-api/src/main/java/datadog/trace/api/internal/TraceSegment.java b/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) {}
 

diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java b/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);