Collect client IP tags for AI Guard requests

When DD_AI_GUARD_ENABLED=true, resolve the client IP eagerly during HTTP
server request decoration and stash it on the request context. Apply the
tags (http.client_ip and network.client.ip) on the local root span only
when an ai_guard span is actually created, so non-AI requests of an
AI-Guard-enabled service do not get IP tags.

This honors the RFC scope (PII collection limited to AI interactions) and
makes the feature work without DD_TRACE_CLIENT_IP_ENABLED or AppSec.

APPSEC-62199

Author: smola

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.getAndResetClientIpAddressData();
+    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<Message> messages, final Options options) {
     if (messages == null || messages.isEmpty()) {
@@ -229,6 +257,7 @@ public Evaluation evaluate(final List<Message> 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);

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.getAndResetClientIpAddressData() >> 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.getAndResetClientIpAddressData() >> 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.getAndResetClientIpAddressData() >> 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']]

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;
@@ -250,6 +251,13 @@ public AgentSpan onRequest(
     AgentSpanContext.Extracted extracted = getExtractedSpanContext(parentContext);
     boolean clientIpResolverEnabled =
         config.isClientIpEnabled() || traceClientIpResolverEnabled && APPSEC_ACTIVE;
+    // AI Guard requires client IP tags on the local root span when an ai_guard span is created.
+    // To respect the RFC's PII scope, 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.
+    boolean aiGuardClientIpCaptureEnabled =
+        !clientIpResolverEnabled && config.isAiGuardEnabled() && traceClientIpResolverEnabled;
+    boolean resolveClientIp = clientIpResolverEnabled || aiGuardClientIpCaptureEnabled;
     if (extracted != null) {
       if (clientIpResolverEnabled) {
         String forwarded = extracted.getForwarded();
@@ -332,7 +340,7 @@ public AgentSpan onRequest(
     }
 
     String inferredAddressStr = null;
-    if (clientIpResolverEnabled && extracted != null) {
+    if (resolveClientIp && 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,7 +357,9 @@ public AgentSpan onRequest(
       }
       if (inferredAddress != null) {
         inferredAddressStr = inferredAddress.getHostAddress();
-        span.setTag(Tags.HTTP_CLIENT_IP, inferredAddressStr);
+        if (clientIpResolverEnabled) {
+          span.setTag(Tags.HTTP_CLIENT_IP, inferredAddressStr);
+        }
       }
     } else if (clientIpResolverEnabled && span.getLocalRootSpan() != span) {
       // in this case extracted == null
@@ -374,6 +384,12 @@ public AgentSpan onRequest(
         span.setTag(Tags.NETWORK_CLIENT_IP, peerIp);
       }
     }
+    if (aiGuardClientIpCaptureEnabled && (peerIp != null || inferredAddressStr != null)) {
+      RequestContext requestContext = span.getRequestContext();
+      if (requestContext != null) {
+        requestContext.setClientIpAddressData(new ClientIpAddressData(peerIp, inferredAddressStr));
+      }
+    }
     setPeerPort(span, peerPort);
     Flow<Void> flow = callIGCallbackAddressAndPort(span, peerIp, peerPort, inferredAddressStr);
     if (flow.getAction() instanceof RequestBlockingAction) {

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,61 @@ 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
+    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')

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> T getOrCreateMetaStructTop(String key, Function<String, T> defaultVal
       return null;
     }
 
+    @Override
+    public void setClientIpAddressData(ClientIpAddressData clientIpAddressData) {}
+
+    @Override
+    public ClientIpAddressData getAndResetClientIpAddressData() {
+      return null;
+    }
+
     @Override
     public void close() throws IOException {}
   }

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 getAndResetClientIpAddressData() {
+      return null
+    }
+
     @Override
     void close() throws IOException {}
   }
@@ -1253,6 +1262,14 @@ class GatewayBridgeSpecification extends DDSpecification {
         return null
       }
 
+      @Override
+      void setClientIpAddressData(ClientIpAddressData clientIpAddressData) {}
+
+      @Override
+      ClientIpAddressData getAndResetClientIpAddressData() {
+        return null
+      }
+
       @Override
       void close() throws IOException {}
     }

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 getAndResetClientIpAddressData() {
+    return delegate.getAndResetClientIpAddressData()
+  }
+
   @Override
   void close() throws IOException {
     delegate.close()