From 46284b65f1338f62602633c5f95521c5ffb1f1eb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Fri, 10 May 2024 13:42:45 +0200 Subject: [PATCH] Add request headers to span when there's a user tracking event --- .../datadog/appsec/gateway/GatewayBridge.java | 29 ++++++++++-- .../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, 124 insertions(+), 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 index d9a3e95d1ce8..f6d0e6e4fe1b 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 List 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 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 f04b48042d2d..1b6ba5d3973f 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 b323128d6385..55535d8c30d4 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 cf1ce90fa15a..61753fe00ae4 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 eb11bb93b407..d867aabee1e1 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);