Collect client IP tags for AI Guard requests

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;
@@ -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.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()

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);