diff --git a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java index 2057c1a4532..e68e2ac85ad 100644 --- a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java +++ b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java @@ -24,10 +24,12 @@ import datadog.trace.api.aiguard.AIGuard.ToolCall.Function; import datadog.trace.api.aiguard.Evaluator; import datadog.trace.api.aiguard.noop.NoOpEvaluator; +import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.telemetry.WafMetricCollector; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData; import datadog.trace.bootstrap.instrumentation.api.Tags; import java.io.IOException; import java.lang.annotation.Annotation; @@ -213,6 +215,32 @@ private boolean isBlockingEnabled(final Options options, final Object isBlocking return options.block() && "true".equalsIgnoreCase(isBlockingEnabled.toString()); } + /** + * Applies the {@link ClientIpAddressData} captured during HTTP server request decoration to the + * local root span. This is the lazy half of AI Guard client IP collection: {@code + * HttpServerDecorator} resolves the IP eagerly and stashes it on the {@link RequestContext}; we + * consume it here once an {@code ai_guard} span is created, so IP tags are not added to spans of + * non-AI requests in services that have AI Guard enabled. + */ + private static void applyClientIpTags(final AgentSpan localRootSpan) { + final RequestContext requestContext = localRootSpan.getRequestContext(); + if (requestContext == null) { + return; + } + final ClientIpAddressData clientIpAddressData = requestContext.getClientIpAddressData(); + if (clientIpAddressData == null) { + return; + } + final String peerIp = clientIpAddressData.getPeerIp(); + if (peerIp != null && localRootSpan.getTag(Tags.NETWORK_CLIENT_IP) == null) { + localRootSpan.setTag(Tags.NETWORK_CLIENT_IP, peerIp); + } + final String inferredClientIp = clientIpAddressData.getInferredClientIp(); + if (inferredClientIp != null && localRootSpan.getTag(Tags.HTTP_CLIENT_IP) == null) { + localRootSpan.setTag(Tags.HTTP_CLIENT_IP, inferredClientIp); + } + } + @Override public Evaluation evaluate(final List messages, final Options options) { if (messages == null || messages.isEmpty()) { @@ -229,6 +257,7 @@ public Evaluation evaluate(final List messages, final Options options) if (localRootSpan != null) { localRootSpan.setTag(Tags.AI_GUARD_KEEP, true); localRootSpan.setTag(EVENT_TAG, true); + applyClientIpTags(localRootSpan); } try (final AgentScope scope = tracer.activateSpan(span)) { final Message last = messages.get(messages.size() - 1); diff --git a/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy b/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy index c017a2ac0ce..8a9fdc297de 100644 --- a/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy +++ b/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy @@ -7,9 +7,11 @@ import com.squareup.moshi.Moshi import datadog.common.version.VersionInfo import datadog.trace.api.Config import datadog.trace.api.aiguard.AIGuard +import datadog.trace.api.gateway.RequestContext import datadog.trace.api.telemetry.WafMetricCollector import datadog.trace.bootstrap.instrumentation.api.AgentSpan import datadog.trace.bootstrap.instrumentation.api.AgentTracer +import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData import datadog.trace.bootstrap.instrumentation.api.Tags import datadog.trace.test.util.DDSpecification import okhttp3.Call @@ -272,6 +274,68 @@ class AIGuardInternalTests extends DDSpecification { new AIGuard.Options().block(false) | true | false } + void 'test evaluate applies captured client ip tags to local root span'() { + given: + final requestContext = Mock(RequestContext) + localRootSpan.getRequestContext() >> requestContext + requestContext.getClientIpAddressData() >> new ClientIpAddressData('4.4.4.4', '2.3.4.5') + final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'It is fine']]]) + + when: + aiguard.evaluate(TOOL_CALL, AIGuard.Options.DEFAULT) + + then: + 1 * localRootSpan.getTag(Tags.NETWORK_CLIENT_IP) >> null + 1 * localRootSpan.setTag(Tags.NETWORK_CLIENT_IP, '4.4.4.4') + 1 * localRootSpan.getTag(Tags.HTTP_CLIENT_IP) >> null + 1 * localRootSpan.setTag(Tags.HTTP_CLIENT_IP, '2.3.4.5') + } + + void 'test evaluate does not overwrite existing client ip tags'() { + given: + final requestContext = Mock(RequestContext) + localRootSpan.getRequestContext() >> requestContext + requestContext.getClientIpAddressData() >> new ClientIpAddressData('4.4.4.4', '2.3.4.5') + final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'It is fine']]]) + + when: + aiguard.evaluate(TOOL_CALL, AIGuard.Options.DEFAULT) + + then: + 1 * localRootSpan.getTag(Tags.NETWORK_CLIENT_IP) >> '9.9.9.9' + 0 * localRootSpan.setTag(Tags.NETWORK_CLIENT_IP, _) + 1 * localRootSpan.getTag(Tags.HTTP_CLIENT_IP) >> '8.8.8.8' + 0 * localRootSpan.setTag(Tags.HTTP_CLIENT_IP, _) + } + + void 'test evaluate is a noop for client ip tags when no data captured'() { + given: + final requestContext = Mock(RequestContext) + localRootSpan.getRequestContext() >> requestContext + requestContext.getClientIpAddressData() >> null + final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'It is fine']]]) + + when: + aiguard.evaluate(TOOL_CALL, AIGuard.Options.DEFAULT) + + then: + 0 * localRootSpan.setTag(Tags.NETWORK_CLIENT_IP, _) + 0 * localRootSpan.setTag(Tags.HTTP_CLIENT_IP, _) + } + + void 'test evaluate is a noop for client ip tags when no request context'() { + given: + localRootSpan.getRequestContext() >> null + final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'It is fine']]]) + + when: + aiguard.evaluate(TOOL_CALL, AIGuard.Options.DEFAULT) + + then: + 0 * localRootSpan.setTag(Tags.NETWORK_CLIENT_IP, _) + 0 * localRootSpan.setTag(Tags.HTTP_CLIENT_IP, _) + } + void 'test evaluate with API errors'() { given: final errors = [[status: 400, title: 'Bad request']] diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index 8b01e5d9206..a60401f5811 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -30,6 +30,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData; import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities; @@ -248,10 +249,31 @@ public AgentSpan onRequest( } AgentSpanContext.Extracted extracted = getExtractedSpanContext(parentContext); - boolean clientIpResolverEnabled = + // Whether to attach IP tags to all requests or not. + // This should be enabled if: + // - DD_TRACE_CLIENT_IP_ENABLED=true, or + // - DD_APPSEC_ENABLED=true (or AppSec enabled at runtime) + // This applies to tags: + // - http.client_ip (IP resolved from proxy tags) + // - network.client.ip (peer IP) + // - tags with proxy header values + // For backwards compatibility, it does not apply to: + // - peer.ipv4 + // - peer.ipv6 + final boolean shouldTagIps = config.isClientIpEnabled() || traceClientIpResolverEnabled && APPSEC_ACTIVE; + // Whether to stash IP data for later tagging or not. + // AI Guard requires client IP tags on the local root span when an ai_guard span is created. + // Resolve the IPs eagerly but do not tag the span yet; stash them on the request context so + // AIGuardInternal can apply them lazily, only on requests that actually create an ai_guard + // span. + final boolean shouldStashIps = + !shouldTagIps && traceClientIpResolverEnabled && config.isAiGuardEnabled(); + // Whether to resolve client IP based on proxy headers or no. + final boolean shouldResolveIp = shouldTagIps || shouldStashIps; + if (extracted != null) { - if (clientIpResolverEnabled) { + if (shouldTagIps) { String forwarded = extracted.getForwarded(); if (forwarded != null) { span.setTag(Tags.HTTP_FORWARDED, forwarded); @@ -332,7 +354,7 @@ public AgentSpan onRequest( } String inferredAddressStr = null; - if (clientIpResolverEnabled && extracted != null) { + if (shouldResolveIp && extracted != null) { InetAddress inferredAddress = ClientIpAddressResolver.resolve(extracted, span); // the peer address should be used if: // 1. the headers yield nothing, regardless of whether it is public or not @@ -349,9 +371,11 @@ public AgentSpan onRequest( } if (inferredAddress != null) { inferredAddressStr = inferredAddress.getHostAddress(); - span.setTag(Tags.HTTP_CLIENT_IP, inferredAddressStr); + if (shouldTagIps) { + span.setTag(Tags.HTTP_CLIENT_IP, inferredAddressStr); + } } - } else if (clientIpResolverEnabled && span.getLocalRootSpan() != span) { + } else if (shouldTagIps && span.getLocalRootSpan() != span) { // in this case extracted == null // If there is no extracted we can't do anything but use the peer addr. // Additionally, extracted == null arises on subspans for which the resolution @@ -370,10 +394,16 @@ public AgentSpan onRequest( } else { span.setTag(Tags.PEER_HOST_IPV4, peerIp); } - if (clientIpResolverEnabled) { + if (shouldTagIps) { span.setTag(Tags.NETWORK_CLIENT_IP, peerIp); } } + if (shouldStashIps && (peerIp != null || inferredAddressStr != null)) { + RequestContext requestContext = span.getRequestContext(); + if (requestContext != null && requestContext.getClientIpAddressData() == null) { + requestContext.setClientIpAddressData(new ClientIpAddressData(peerIp, inferredAddressStr)); + } + } setPeerPort(span, peerPort); Flow flow = callIGCallbackAddressAndPort(span, peerIp, peerPort, inferredAddressStr); if (flow.getAction() instanceof RequestBlockingAction) { diff --git a/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecoratorTest.groovy b/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecoratorTest.groovy index b58f6f10f0a..da411dc2431 100644 --- a/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecoratorTest.groovy +++ b/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecoratorTest.groovy @@ -15,6 +15,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext import datadog.trace.bootstrap.instrumentation.api.AgentTracer import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI +import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData import datadog.trace.bootstrap.instrumentation.api.ContextVisitors import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities @@ -300,6 +301,62 @@ class HttpServerDecoratorTest extends ServerDecoratorTest { 0 * this.span.setTag(Tags.NETWORK_CLIENT_IP, _) } + void 'enabling ai guard captures client ip data without tagging the span during request decoration'() { + setup: + injectSysConfig('dd.ai_guard.enabled', 'true') + ActiveSubsystems.APPSEC_ACTIVE = false + + def extracted = Mock(AgentSpanContext.Extracted) + def context = AgentSpan.fromSpanContext(extracted) + def requestContext = Mock(RequestContext) + def decorator = newDecorator() + + when: + decorator.onRequest(this.span, [peerIp: '4.4.4.4'], null, context) + + then: + _ * extracted.getXForwardedFor() >> '2.3.4.5' + _ * extracted.getXClusterClientIp() >> null + _ * extracted.getXRealIp() >> null + _ * extracted.getXClientIp() >> null + _ * extracted.getCustomIpHeader() >> null + _ * extracted.getTrueClientIp() >> null + _ * extracted.getFastlyClientIp() >> null + _ * extracted.getCfConnectingIp() >> null + _ * extracted.getCfConnectingIpv6() >> null + _ * this.span.getRequestContext() >> requestContext + _ * requestContext.getClientIpAddressData() >> null + 1 * requestContext.setClientIpAddressData({ ClientIpAddressData data -> + data.peerIp == '4.4.4.4' && data.inferredClientIp == '2.3.4.5' + }) + 0 * this.span.setTag(Tags.HTTP_CLIENT_IP, _) + 0 * this.span.setTag(Tags.NETWORK_CLIENT_IP, _) + 0 * this.span.setTag(Tags.HTTP_FORWARDED_IP, _) + } + + void 'enabling ai guard does not override client_ip_without_appsec tagging behavior'() { + setup: + injectSysConfig('dd.ai_guard.enabled', 'true') + injectSysConfig('dd.trace.client-ip.enabled', 'true') + ActiveSubsystems.APPSEC_ACTIVE = false + + def extracted = Mock(AgentSpanContext.Extracted) + def context = AgentSpan.fromSpanContext(extracted) + def requestContext = Mock(RequestContext) + def decorator = newDecorator() + + when: + decorator.onRequest(this.span, [peerIp: '4.4.4.4'], null, context) + + then: + _ * this.span.getRequestContext() >> requestContext + 2 * extracted.getXForwardedFor() >> '2.3.4.5' + 1 * this.span.setTag(Tags.HTTP_CLIENT_IP, '2.3.4.5') + 1 * this.span.setTag(Tags.NETWORK_CLIENT_IP, '4.4.4.4') + // ai guard capture is skipped because tags were already set + 0 * requestContext.setClientIpAddressData(_) + } + void 'disabling appsec but enabling client_ip_without_appsec enables header collection and ip address resolution'() { setup: injectSysConfig('dd.trace.client-ip.enabled', 'true') diff --git a/dd-java-agent/appsec/src/jmh/java/datadog/appsec/benchmark/AppSecBenchmark.java b/dd-java-agent/appsec/src/jmh/java/datadog/appsec/benchmark/AppSecBenchmark.java index 0d863a54474..3238723bf51 100644 --- a/dd-java-agent/appsec/src/jmh/java/datadog/appsec/benchmark/AppSecBenchmark.java +++ b/dd-java-agent/appsec/src/jmh/java/datadog/appsec/benchmark/AppSecBenchmark.java @@ -17,6 +17,7 @@ import datadog.trace.api.gateway.RequestContextSlot; import datadog.trace.api.gateway.SubscriptionService; import datadog.trace.api.internal.TraceSegment; +import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData; import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter; import datadog.trace.bootstrap.instrumentation.api.URIDefaultDataAdapter; import java.io.IOException; @@ -254,6 +255,14 @@ public T getOrCreateMetaStructTop(String key, Function defaultVal return null; } + @Override + public void setClientIpAddressData(ClientIpAddressData clientIpAddressData) {} + + @Override + public ClientIpAddressData getClientIpAddressData() { + return null; + } + @Override public void close() throws IOException {} } 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 223da6678d2..36a8ebc861d 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 @@ -28,6 +28,7 @@ import datadog.trace.api.gateway.RequestContextSlot import datadog.trace.api.gateway.SubscriptionService import datadog.trace.api.http.StoredBodySupplier import datadog.trace.api.internal.TraceSegment +import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData import datadog.trace.api.telemetry.LoginEvent import datadog.trace.api.telemetry.RuleType import datadog.trace.api.telemetry.WafMetricCollector @@ -79,6 +80,14 @@ class GatewayBridgeSpecification extends DDSpecification { return null } + @Override + void setClientIpAddressData(ClientIpAddressData clientIpAddressData) {} + + @Override + ClientIpAddressData getClientIpAddressData() { + return null + } + @Override void close() throws IOException {} } @@ -1253,6 +1262,14 @@ class GatewayBridgeSpecification extends DDSpecification { return null } + @Override + void setClientIpAddressData(ClientIpAddressData clientIpAddressData) {} + + @Override + ClientIpAddressData getClientIpAddressData() { + return null + } + @Override void close() throws IOException {} } diff --git a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/ValidatingRequestContextDecorator.groovy b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/ValidatingRequestContextDecorator.groovy index 13c7ab59c3c..68640c244cb 100644 --- a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/ValidatingRequestContextDecorator.groovy +++ b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/ValidatingRequestContextDecorator.groovy @@ -4,6 +4,7 @@ import datadog.trace.api.gateway.BlockResponseFunction import datadog.trace.api.gateway.RequestContext import datadog.trace.api.gateway.RequestContextSlot import datadog.trace.api.internal.TraceSegment +import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData import java.util.function.Function @@ -56,6 +57,16 @@ class ValidatingRequestContextDecorator implements RequestContext { return delegate.getOrCreateMetaStructTop(key, defaultValue) } + @Override + void setClientIpAddressData(ClientIpAddressData clientIpAddressData) { + delegate.setClientIpAddressData(clientIpAddressData) + } + + @Override + ClientIpAddressData getClientIpAddressData() { + return delegate.getClientIpAddressData() + } + @Override void close() throws IOException { delegate.close() diff --git a/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/SpringbootApplication.java b/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/SpringbootApplication.java index ae0bd8277c4..d424242784b 100644 --- a/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/SpringbootApplication.java +++ b/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/SpringbootApplication.java @@ -10,10 +10,14 @@ public class SpringbootApplication { public static void main(final String[] args) { - try { - activateAppSec(); - } catch (Exception e) { - System.out.println("Could not activate appSec: " + e.getMessage()); + if (!Boolean.getBoolean("smoketest.skipAppSecActivation")) { + try { + activateAppSec(); + } catch (Exception e) { + System.out.println("Could not activate appSec: " + e.getMessage()); + } + } else { + System.out.println("AppSec activation skipped"); } SpringApplication.run(SpringbootApplication.class, args); diff --git a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy index 3bc7a30481f..d1a3f84ec18 100644 --- a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy +++ b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy @@ -10,6 +10,10 @@ class AIGuardSmokeTest extends AbstractAppSecServerSmokeTest { @Shared protected String[] defaultAIGuardProperties = [ '-Ddd.ai_guard.enabled=true', + // Make sure AI Guard features (e.g. client IP tags collection) do not depend on AppSec. + '-Ddd.appsec.enabled=false', + '-Ddd.trace.client-ip.enabled=false', + '-Dsmoketest.skipAppSecActivation=true', "-Ddd.ai_guard.endpoint=http://localhost:${httpPort}/aiguard".toString(), ] @@ -29,7 +33,6 @@ class AIGuardSmokeTest extends AbstractAppSecServerSmokeTest { final springBootShadowJar = System.getProperty("datadog.smoketest.appsec.springboot.shadowJar.path") final command = [javaPath()] command.addAll(defaultJavaProperties) - command.addAll(defaultAppSecProperties) command.addAll(defaultAIGuardProperties) command.addAll(['-jar', springBootShadowJar, "--server.port=${httpPort}".toString()]) final builder = new ProcessBuilder(command).directory(new File(buildDirectory)) @@ -134,6 +137,36 @@ class AIGuardSmokeTest extends AbstractAppSecServerSmokeTest { assert span.meta.get('ai_guard.blocked') == 'true' } + void 'client ip tags are added to the local root span when an ai_guard span is created'() { + given: + final publicIp = '5.6.7.9' + final request = new Request.Builder() + .url("http://localhost:${httpPort}/aiguard/allow") + .header('X-Forwarded-For', publicIp) + .get() + .build() + + when: + final response = client.newCall(request).execute() + + then: + response.code() == 200 + + and: + waitForTraceCount(2) // /aiguard/allow + internal /aiguard/evaluate mock + final aiGuardSpan = traces*.spans + ?.flatten() + ?.find { it.resource == 'ai_guard' } as DecodedSpan + aiGuardSpan != null + final rootSpan = traces*.spans + ?.flatten() + ?.find { it.traceId == aiGuardSpan.traceId && it.parentId == 0 } as DecodedSpan + rootSpan != null + rootSpan.meta.get('http.client_ip') == publicIp + rootSpan.meta.get('network.client.ip') != null + rootSpan.meta.get('network.client.ip') != publicIp + } + void 'test multimodal content parts evaluation'() { given: def request = new Request.Builder() 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 a3e90a37960..f2eb17fe8a2 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 @@ -26,6 +26,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AppendableSpanLinks; import datadog.trace.bootstrap.instrumentation.api.Baggage; +import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData; import datadog.trace.bootstrap.instrumentation.api.ProfilerContext; import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities; @@ -182,6 +183,8 @@ public class DDSpanContext private volatile BlockResponseFunction blockResponseFunction; + private volatile ClientIpAddressData clientIpAddressData; + private final ProfilingContextIntegration profilingContextIntegration; private final boolean injectBaggageAsTags; private final boolean injectLinksAsTags; @@ -1354,6 +1357,16 @@ public BlockResponseFunction getBlockResponseFunction() { return getRootSpanContextOrThis().blockResponseFunction; } + @Override + public void setClientIpAddressData(final ClientIpAddressData clientIpAddressData) { + getRootSpanContextOrThis().clientIpAddressData = clientIpAddressData; + } + + @Override + public ClientIpAddressData getClientIpAddressData() { + return getRootSpanContextOrThis().clientIpAddressData; + } + public PropagationTags getPropagationTags() { return getRootSpanContextOrThis().propagationTags; } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ContextInterpreter.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ContextInterpreter.java index e94097db2a4..5afc104c675 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ContextInterpreter.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ContextInterpreter.java @@ -60,6 +60,7 @@ public abstract class ContextInterpreter implements AgentPropagation.KeyClassifi private final String customIpHeaderName; private final boolean clientIpResolutionEnabled; private final boolean clientIpWithoutAppSec; + private final boolean aiGuardEnabled; private boolean collectIpHeaders; private final boolean requestHeaderTagsCommaAllowed; @@ -74,6 +75,7 @@ protected ContextInterpreter(Config config) { this.customIpHeaderName = config.getTraceClientIpHeader(); this.clientIpResolutionEnabled = config.isTraceClientIpResolverEnabled(); this.clientIpWithoutAppSec = config.isClientIpEnabled(); + this.aiGuardEnabled = config.isAiGuardEnabled(); this.propagationTagsFactory = PropagationTags.factory(config); this.requestHeaderTagsCommaAllowed = config.isRequestHeaderTagsCommaAllowed(); } @@ -237,7 +239,8 @@ public ContextInterpreter reset(TraceConfig traceConfig) { httpHeaders = null; collectIpHeaders = this.clientIpWithoutAppSec - || this.clientIpResolutionEnabled && ActiveSubsystems.APPSEC_ACTIVE; + || this.clientIpResolutionEnabled + && (ActiveSubsystems.APPSEC_ACTIVE || this.aiGuardEnabled); headerTags = traceConfig.getRequestHeaderTags(); baggageMapping = traceConfig.getBaggageMapping(); propagationTags = null; diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java index 31519ce4ca7..32bfa7fad6e 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java @@ -17,6 +17,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData; import datadog.trace.bootstrap.instrumentation.api.TagContext; import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter; import datadog.trace.bootstrap.instrumentation.api.URIDataAdapterBase; @@ -802,6 +803,16 @@ public T getOrCreateMetaStructTop(String key, Function defaultVal return null; } + @Override + public void setClientIpAddressData(ClientIpAddressData clientIpAddressData) { + // No-op for temporary context + } + + @Override + public ClientIpAddressData getClientIpAddressData() { + return null; + } + @Override public void close() { // No-op for temporary context diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/RequestContext.java b/internal-api/src/main/java/datadog/trace/api/gateway/RequestContext.java index 08a6cadc98e..af9e2f8f493 100644 --- a/internal-api/src/main/java/datadog/trace/api/gateway/RequestContext.java +++ b/internal-api/src/main/java/datadog/trace/api/gateway/RequestContext.java @@ -1,6 +1,7 @@ package datadog.trace.api.gateway; import datadog.trace.api.internal.TraceSegment; +import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData; import java.io.Closeable; import java.io.IOException; import java.util.function.Function; @@ -20,6 +21,18 @@ public interface RequestContext extends Closeable { T getOrCreateMetaStructTop(String key, Function defaultValue); + /** + * Stores the client IP address information resolved during HTTP server request decoration so that + * later consumers (such as AI Guard) can apply it to the local root span without re-running the + * resolver. Implementations should store the data on the root request context. + */ + void setClientIpAddressData(ClientIpAddressData clientIpAddressData); + + /** + * Returns the previously stored {@link ClientIpAddressData}, or {@code null} if none was stored. + */ + ClientIpAddressData getClientIpAddressData(); + class Noop implements RequestContext { public static final RequestContext INSTANCE = new Noop(); @@ -48,6 +61,14 @@ public T getOrCreateMetaStructTop(String key, Function defaultVal return null; } + @Override + public void setClientIpAddressData(ClientIpAddressData clientIpAddressData) {} + + @Override + public ClientIpAddressData getClientIpAddressData() { + return null; + } + @Override public void close() throws IOException {} } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ClientIpAddressData.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ClientIpAddressData.java new file mode 100644 index 00000000000..d22bcde4618 --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ClientIpAddressData.java @@ -0,0 +1,23 @@ +package datadog.trace.bootstrap.instrumentation.api; + +/** + * Holds the client IP information resolved during HTTP server request decoration so that consumers + * (such as AI Guard) can apply it lazily to the local root span without re-running the resolver. + */ +public final class ClientIpAddressData { + private final String peerIp; + private final String inferredClientIp; + + public ClientIpAddressData(final String peerIp, final String inferredClientIp) { + this.peerIp = peerIp; + this.inferredClientIp = inferredClientIp; + } + + public String getPeerIp() { + return peerIp; + } + + public String getInferredClientIp() { + return inferredClientIp; + } +} diff --git a/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java b/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java index 3ae4aba9a71..4767a0051b9 100644 --- a/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java +++ b/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java @@ -15,6 +15,7 @@ import datadog.trace.api.http.StoredBodySupplier; import datadog.trace.api.internal.TraceSegment; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData; import java.util.Collections; import java.util.function.BiConsumer; import java.util.function.BiFunction; @@ -64,6 +65,14 @@ public BlockResponseFunction getBlockResponseFunction() { public T getOrCreateMetaStructTop(String key, Function defaultValue) { return null; } + + @Override + public void setClientIpAddressData(ClientIpAddressData clientIpAddressData) {} + + @Override + public ClientIpAddressData getClientIpAddressData() { + return null; + } }; flow = new Flow.ResultFlow<>(null); callback = new Callback(context, flow);