Add setUser to user monitoring SDK

Author: manuel-alvarez-alvarez

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/user/AppSecEventTrackerSpecification.groovy

@@ -3,6 +3,7 @@ package com.datadog.appsec.user
 import com.datadog.appsec.gateway.NoopFlow
 import datadog.appsec.api.blocking.BlockingContentType
 import datadog.appsec.api.blocking.BlockingException
+import datadog.appsec.api.user.User
 import datadog.trace.api.EventTracker
 import datadog.trace.api.GlobalTracer
 import datadog.trace.api.ProductTraceSource
@@ -31,6 +32,7 @@ import static datadog.trace.api.gateway.Events.EVENTS
 import static datadog.trace.api.telemetry.LoginEvent.LOGIN_FAILURE
 import static datadog.trace.api.telemetry.LoginEvent.LOGIN_SUCCESS
 import static datadog.trace.api.telemetry.LoginEvent.SIGN_UP
+import static datadog.appsec.api.user.User.setUser
 
 class AppSecEventTrackerSpecification extends DDSpecification {
 
@@ -72,6 +74,7 @@ class AppSecEventTrackerSpecification extends DDSpecification {
         }
       }
     GlobalTracer.setEventTracker(tracker)
+    User.setUserService(tracker)
     ActiveSubsystems.APPSEC_ACTIVE = true
   }
 
@@ -128,6 +131,20 @@ class AppSecEventTrackerSpecification extends DDSpecification {
     0 * _
   }
 
+  def 'test track user (SDK)'() {
+    when:
+    setUser(USER_ID, ['key1': 'value1', 'key2': 'value2'])
+
+    then:
+    1 * traceSegment.setTagTop('usr.id', USER_ID)
+    1 * traceSegment.setTagTop('usr', ['key1':'value1', 'key2':'value2'])
+    1 * traceSegment.setTagTop('_dd.appsec.user.collection_mode', SDK.fullName())
+    1 * traceSegment.setTagTop('asm.keep', true)
+    1 * traceSegment.setTagTop('_dd.p.ts', ProductTraceSource.ASM)
+    1 * user.apply(_ as RequestContext, USER_ID) >> NoopFlow.INSTANCE
+    0 * _
+  }
+
   def 'test wrong event argument validation (SDK)'() {
     when:
     GlobalTracer.getEventTracker().trackLoginSuccessEvent(null, null)
@@ -164,6 +181,12 @@ class AppSecEventTrackerSpecification extends DDSpecification {
 
     then:
     thrown IllegalArgumentException
+
+    when:
+    setUser(null, null)
+
+    then:
+    thrown IllegalArgumentException
   }
 
   def "test onSignup (#mode)"() {

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/user/EventTrackerAppSecDisabledForkedTest.groovy

@@ -1,6 +1,6 @@
 package com.datadog.appsec.user
 
-
+import datadog.appsec.api.user.User
 import datadog.trace.api.GlobalTracer
 import datadog.trace.api.UserIdCollectionMode
 import datadog.trace.api.appsec.AppSecEventTracker
@@ -26,7 +26,7 @@ class EventTrackerAppSecDisabledForkedTest extends DDSpecification {
   void setup() {
     tracker = new AppSecEventTracker()
     GlobalTracer.setEventTracker(tracker)
-
+    User.setUserService(tracker)
     traceSegment = Mock(TraceSegment)
     final tracer = Stub(AgentTracer.TracerAPI) {
       getTraceSegment() >> traceSegment

dd-java-agent/appsec/src/testFixtures/groovy/com/datadog/appsec/AppSecHttpServerTest.groovy

@@ -26,8 +26,7 @@ abstract class AppSecHttpServerTest<SERVER> extends WithHttpServer<SERVER> {
     sco.createRemaining(config)
     assert sco.configurationPoller(config) == null
     assert sco.monitoring instanceof Monitoring.DisabledMonitoring
-
-    GlobalTracer.setEventTracker(new AppSecEventTracker())
+    AppSecEventTracker.install()
 
     AppSecSystem.start(ss, sco)
   }

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

@@ -11,6 +11,7 @@ import datadog.trace.api.config.AppSecConfig
 import datadog.trace.core.DDSpan
 import okhttp3.FormBody
 import okhttp3.HttpUrl
+import okhttp3.OkHttpClient
 import okhttp3.Request
 import okhttp3.RequestBody
 import org.springframework.boot.SpringApplication
@@ -290,17 +291,7 @@ class SpringBootBasedTest extends AppSecHttpServerTest<ConfigurableApplicationCo
   void 'test user event'() {
     setup:
     def client = clientBuilder().cookieJar(cookieJar()).followRedirects(false).build()
-    def formBody = new FormBody.Builder()
-      .add("username", "admin")
-      .add("password", "admin")
-      .build()
-
-    def loginRequest = request(LOGIN, "POST", formBody).build()
-    def loginResponse = client.newCall(loginRequest).execute()
-    assert loginResponse.code() == LOGIN.status
-    assert loginResponse.body().string() == LOGIN.body
-    TEST_WRITER.waitForTraces(1)
-    TEST_WRITER.start() // clear all traces
+    doLogin(client, 'admin', 'admin')
 
     when:
     def request = request(SUCCESS, "GET", null).build()
@@ -322,13 +313,7 @@ class SpringBootBasedTest extends AppSecHttpServerTest<ConfigurableApplicationCo
     setup:
     def logMessagePrefix = 'Attempt to replace'
     def client = clientBuilder().cookieJar(cookieJar()).followRedirects(false).build()
-    def formBody = new FormBody.Builder()
-      .add('username', 'admin')
-      .add('password', 'admin')
-      .build()
-    def loginRequest = request(LOGIN, 'POST', formBody).build()
-    def loginResponse = client.newCall(loginRequest).execute()
-    assert loginResponse.code() == LOGIN.status
+    doLogin(client, 'admin', 'admin')
 
     when: 'sdk with different user'
     def sdkBody = new FormBody.Builder().add("sdkUser", "sdkUser").build()
@@ -350,4 +335,37 @@ class SpringBootBasedTest extends AppSecHttpServerTest<ConfigurableApplicationCo
       event.message.startsWith(logMessagePrefix)
     })
   }
+
+  void 'test automated user tracking and setUser SDK used simultaneously'() {
+    setup:
+    def client = clientBuilder().cookieJar(cookieJar()).followRedirects(false).build()
+    doLogin(client, 'admin', 'admin')
+
+    when:
+    def sdkBody = new FormBody.Builder().add("sdkEvent", "setUser").add("sdkUser", "sdkUser").build()
+    def sdkRequest = request(SDK, 'POST', sdkBody).build()
+    final response = client.newCall(sdkRequest).execute()
+    TEST_WRITER.waitForTraces(1)
+    def span = TEST_WRITER.flatten().first() as DDSpan
+
+    then:
+    response.code() == SDK.status
+    response.body().string() == SDK.body
+    span.getTag('_dd.appsec.usr.id') == 'admin' //
+    // SDK should take priority over automated login events
+    span.getTag('usr.id') == 'sdkUser'
+    span.getTag('_dd.appsec.user.collection_mode') == 'sdk'
+  }
+
+  private void doLogin(final OkHttpClient client, final String username, final String password) {
+    def formBody = new FormBody.Builder()
+      .add('username', username)
+      .add('password', password)
+      .build()
+    def loginRequest = request(LOGIN, 'POST', formBody).build()
+    def loginResponse = client.newCall(loginRequest).execute()
+    assert loginResponse.code() == LOGIN.status
+    TEST_WRITER.waitForTraces(1)
+    TEST_WRITER.start() // clear all traces
+  }
 }

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

@@ -10,6 +10,7 @@ import org.springframework.web.bind.annotation.RequestMapping
 import org.springframework.web.bind.annotation.RequestParam
 import org.springframework.web.bind.annotation.ResponseBody
 
+import static datadog.appsec.api.user.User.setUser
 import static datadog.trace.instrumentation.springsecurity5.TestEndpoint.REGISTER
 import static datadog.trace.instrumentation.springsecurity5.TestEndpoint.SUCCESS
 import static datadog.trace.instrumentation.springsecurity5.TestEndpoint.SDK
@@ -48,9 +49,22 @@ class UserController {
 
   @PostMapping("/sdk")
   @ResponseBody
-  String sdk(@RequestParam("sdkUser") String sdkUser) {
+  String sdk(@RequestParam(name = "sdkEvent", defaultValue = "login.success") String event, @RequestParam(name = "sdkUser", required = false) String sdkUser) {
     SpringBootBasedTest.controller(SDK) {
-      GlobalTracer.getEventTracker().trackLoginSuccessEvent(sdkUser, emptyMap())
+      switch (event) {
+        case "login.success":
+          GlobalTracer.getEventTracker().trackLoginSuccessEvent(sdkUser, emptyMap())
+          break
+        case "login.failure":
+          GlobalTracer.getEventTracker().trackLoginFailureEvent(sdkUser, false, emptyMap())
+          break
+        case "setUser":
+          setUser(sdkUser, emptyMap())
+          break
+        default:
+          GlobalTracer.getEventTracker().trackCustomEvent(event, emptyMap())
+          break
+      }
       return "OK"
     }
   }

dd-trace-api/build.gradle

@@ -36,6 +36,7 @@ excludedClassesCoverage += [
   'datadog.trace.api.experimental.DataStreamsContextCarrier',
   'datadog.trace.api.experimental.DataStreamsContextCarrier.NoOp',
   'datadog.appsec.api.blocking.*',
+  'datadog.appsec.api.user.*',
   // Default fallback methods to not break legacy API
   'datadog.trace.context.TraceScope',
   'datadog.trace.context.NoopTraceScope.NoopContinuation',

dd-trace-api/src/main/java/datadog/appsec/api/user/User.java

@@ -0,0 +1,28 @@
+package datadog.appsec.api.user;
+
+import java.util.Map;
+
+public class User {
+
+  private static volatile UserService SERVICE = UserService.NO_OP;
+
+  /**
+   * Controls the implementation for user service. The AppSec subsystem calls this method on
+   * startup. This can be called explicitly for e.g. testing purposes.
+   *
+   * @param service the implementation for the user service.
+   */
+  public static void setUserService(final UserService service) {
+    SERVICE = service;
+  }
+
+  /**
+   * Sets the user monitoring tags on the root span using the prefix {@code usr}
+   *
+   * @param id identifier of the user
+   * @param metadata custom metadata data represented as key/value map
+   */
+  public static void setUser(final String id, final Map<String, String> metadata) {
+    SERVICE.trackUserEvent(id, metadata);
+  }
+}

dd-trace-api/src/main/java/datadog/appsec/api/user/UserService.java

@@ -0,0 +1,14 @@
+package datadog.appsec.api.user;
+
+import java.util.Map;
+
+public interface UserService {
+
+  UserService NO_OP =
+      new UserService() {
+        @Override
+        public void trackUserEvent(final String userId, final Map<String, String> metadata) {}
+      };
+
+  void trackUserEvent(String userId, Map<String, String> metadata);
+}

internal-api/src/main/java/datadog/trace/api/appsec/AppSecEventTracker.java

@@ -8,8 +8,11 @@
 import static datadog.trace.api.telemetry.LoginEvent.LOGIN_SUCCESS;
 import static datadog.trace.api.telemetry.LoginEvent.SIGN_UP;
 import static datadog.trace.util.Strings.toHexString;
+import static java.util.Collections.emptyMap;
 
 import datadog.appsec.api.blocking.BlockingException;
+import datadog.appsec.api.user.User;
+import datadog.appsec.api.user.UserService;
 import datadog.trace.api.EventTracker;
 import datadog.trace.api.GlobalTracer;
 import datadog.trace.api.ProductTraceSource;
@@ -30,7 +33,7 @@
 import java.util.Map;
 import java.util.function.BiFunction;
 
-public class AppSecEventTracker extends EventTracker {
+public class AppSecEventTracker extends EventTracker implements UserService {
 
   private static final int HASH_SIZE_BYTES = 16; // 128 bits
   private static final String ANON_PREFIX = "anon_";
@@ -40,7 +43,9 @@ public class AppSecEventTracker extends EventTracker {
   private static final String SIGNUP_TAG = "users.signup";
 
   public static void install() {
-    GlobalTracer.setEventTracker(new AppSecEventTracker());
+    final AppSecEventTracker tracker = new AppSecEventTracker();
+    GlobalTracer.setEventTracker(tracker);
+    User.setUserService(tracker);
   }
 
   @Override
@@ -68,6 +73,14 @@ public final void trackCustomEvent(String eventName, Map<String, String> metadat
     onCustomEvent(SDK, eventName, metadata);
   }
 
+  @Override
+  public void trackUserEvent(final String userId, final Map<String, String> metadata) {
+    if (userId == null || userId.isEmpty()) {
+      throw new IllegalArgumentException("UserId is null or empty");
+    }
+    onUserEvent(SDK, userId, metadata);
+  }
+
   public void onUserNotFound(final UserIdCollectionMode mode) {
     if (!isEnabled(mode)) {
       return;
@@ -88,6 +101,11 @@ public void onUserNotFound(final UserIdCollectionMode mode) {
   }
 
   public void onUserEvent(final UserIdCollectionMode mode, final String userId) {
+    onUserEvent(mode, userId, emptyMap());
+  }
+
+  public void onUserEvent(
+      final UserIdCollectionMode mode, final String userId, final Map<String, String> metadata) {
     if (!isEnabled(mode)) {
       return;
     }
@@ -108,6 +126,9 @@ public void onUserEvent(final UserIdCollectionMode mode, final String userId) {
     }
     if (isNewUser(mode, segment)) {
       segment.setTagTop("usr.id", finalUserId);
+      if (metadata != null && !metadata.isEmpty()) {
+        segment.setTagTop("usr", metadata);
+      }
       segment.setTagTop("_dd.appsec.user.collection_mode", mode.fullName());
       segment.setTagTop(Tags.ASM_KEEP, true);
       segment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM);