From ad66ae831ff35f39b40065180db609c003dda470 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Wed, 15 May 2024 16:00:05 +0200 Subject: [PATCH] Collect WAF headers on user sdk events (#7014) Add request headers to span when there's a user tracking event --- .../datadog/appsec/gateway/GatewayBridge.java | 18 ++++++++ .../gateway/GatewayBridgeSpecification.groovy | 28 ++++++++++++ .../trace/agent/test/AgentTestRunner.groovy | 12 +++++ .../trace/api/internal/TraceSegment.java | 45 +++++++++++++++++++ .../datadog/trace/core/DDSpanContext.java | 13 ++++++ 5 files changed, 116 insertions(+) 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 index d9a3e95d1ce..040f3e8527d 100644 --- 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> 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 allowed, diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index f04b48042d2..1b6ba5d3973 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/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 + } } 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 index b323128d638..55535d8c30d 100644 --- 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 index cf1ce90fa15..61753fe00ae 100644 --- 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 index eb11bb93b40..d867aabee1e 100644 --- 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);