Prevent publishing the same usr.id to the WAF twice (#7699)

Author: manuel-alvarez-alvarez

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java

@@ -124,6 +124,9 @@ public class AppSecRequestContext implements DataBundle, Closeable {
   private volatile boolean blocked;
   private volatile int timeouts;
 
+  // keep a reference to the last published usr.id
+  private volatile String userId;
+
   private static final AtomicIntegerFieldUpdater<AppSecRequestContext> TIMEOUTS_UPDATER =
       AtomicIntegerFieldUpdater.newUpdater(AppSecRequestContext.class, "timeouts");
 
@@ -423,6 +426,14 @@ public void setRespDataPublished(boolean respDataPublished) {
     this.respDataPublished = respDataPublished;
   }
 
+  public String getUserId() {
+    return userId;
+  }
+
+  public void setUserId(String userId) {
+    this.userId = userId;
+  }
+
   @Override
   public void close() {
     final AgentSpan span = AgentTracer.activeSpan();

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java

@@ -156,21 +156,35 @@ private TriFunction<RequestContext, UserIdCollectionMode, String, Flow<Void>> on
         return NoopFlow.INSTANCE;
       }
       final TraceSegment segment = ctx_.getTraceSegment();
+      // user id can be set by the SDK overriding the auto event, always update the segment
       segment.setTagTop("usr.id", userId);
       segment.setTagTop("_dd.appsec.user.collection_mode", mode.shortName());
-      final Address<?>[] addresses =
-          address == KnownAddresses.USER_ID
-              ? new Address[] {KnownAddresses.USER_ID}
-              : new Address[] {KnownAddresses.USER_ID, address};
+      final List<Address<?>> addresses = new ArrayList<>(2);
+      final boolean newUserId = !userId.equals(ctx.getUserId());
+      if (newUserId) {
+        // unlikely that multiple threads will update the value at the same time
+        ctx.setUserId(userId);
+        addresses.add(KnownAddresses.USER_ID);
+      }
+      if (address != KnownAddresses.USER_ID) {
+        addresses.add(address);
+      }
+      if (addresses.isEmpty()) {
+        // nothing to publish so short-circuit here
+        return NoopFlow.INSTANCE;
+      }
+      final Address<?>[] addressArray = addresses.toArray(new Address[0]);
       while (true) {
         DataSubscriberInfo subInfo =
             userIdSubInfo.computeIfAbsent(
-                address, k -> producerService.getDataSubscribers(addresses));
+                address, k -> producerService.getDataSubscribers(addressArray));
         if (subInfo == null || subInfo.isEmpty()) {
           return NoopFlow.INSTANCE;
         }
-        MapDataBundle.Builder bundle =
-            new MapDataBundle.Builder(CAPACITY_0_2).add(KnownAddresses.USER_ID, userId);
+        MapDataBundle.Builder bundle = new MapDataBundle.Builder(CAPACITY_0_2);
+        if (newUserId) {
+          bundle.add(KnownAddresses.USER_ID, userId);
+        }
         if (address != KnownAddresses.USER_ID) {
           // we don't support null values for the address so we use an invalid placeholder here
           bundle.add(address, "invalid");

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy

@@ -1020,4 +1020,20 @@ class GatewayBridgeSpecification extends DDSpecification {
     'loginSuccessCB' | KnownAddresses.LOGIN_SUCCESS
     'loginFailureCB' | KnownAddresses.LOGIN_FAILURE
   }
+
+  void 'ensure that the same user id is not published twice'() {
+    setup:
+    eventDispatcher.getDataSubscribers(_) >> nonEmptyDsInfo
+    final userId = 'admin'
+
+    when:
+    userIdCB.apply(ctx, UserIdCollectionMode.IDENTIFICATION, userId)
+    userIdCB.apply(ctx, UserIdCollectionMode.SDK, userId)
+
+    then:
+    2 * traceSegment.setTagTop('usr.id', userId)
+    1 * traceSegment.setTagTop('_dd.appsec.user.collection_mode', UserIdCollectionMode.IDENTIFICATION.shortName())
+    1 * traceSegment.setTagTop('_dd.appsec.user.collection_mode', UserIdCollectionMode.SDK.shortName())
+    1 * eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext)
+  }
 }

dd-java-agent/instrumentation/spring-security-5/src/test/groovy/datadog/trace/instrumentation/springsecurity5/SpringBootBasedTest.groovy

@@ -315,6 +315,7 @@ class SpringBootBasedTest extends AppSecHttpServerTest<ConfigurableApplicationCo
 
   void 'test multiple user ids do not cause warn messages'() {
     setup:
+    def logMessagePrefix = 'Attempt to replace'
     def client = clientBuilder().cookieJar(cookieJar()).followRedirects(false).build()
     def formBody = new FormBody.Builder()
       .add('username', 'admin')
@@ -324,14 +325,24 @@ class SpringBootBasedTest extends AppSecHttpServerTest<ConfigurableApplicationCo
     def loginResponse = client.newCall(loginRequest).execute()
     assert loginResponse.code() == LOGIN.status
 
-    when:
+    when: 'sdk with different user'
     def sdkBody = new FormBody.Builder().add("sdkUser", "sdkUser").build()
     def sdkRequest = request(SDK, 'POST', sdkBody).build()
     client.newCall(sdkRequest).execute()
 
     then:
     1 * reqCtxLogAppender.doAppend({ LoggingEvent event ->
-      event.level.levelInt == Level.DEBUG_INT && event.message.startsWith('Attempt to replace')
+      event.level.levelInt == Level.DEBUG_INT && event.message.startsWith(logMessagePrefix)
+    })
+
+    when: 'sdk with same user'
+    sdkBody = new FormBody.Builder().add("sdkUser", "admin").build()
+    sdkRequest = request(SDK, 'POST', sdkBody).build()
+    client.newCall(sdkRequest).execute()
+
+    then:
+    0 * reqCtxLogAppender.doAppend({ LoggingEvent event ->
+      event.message.startsWith(logMessagePrefix)
     })
   }
 }