Add support for user tracking in spring security (#7633)

dd-java-agent/appsec/src/main/java/com/datadog/appsec/user/AppSecEventTrackerImpl.java

@@ -96,6 +96,15 @@ public void onLoginFailureEvent(
     }
   }
 
+  @Override
+  public void onUserEvent(UserIdCollectionMode mode, String userId) {
+    TraceSegment segment = beforeEvent(mode, userId);
+    if (segment == null) {
+      return;
+    }
+    onUserId(mode, segment, userId, EVENTS.userId());
+  }
+
   @Override
   public void onCustomEvent(
       final UserIdCollectionMode mode, final String eventName, final Map<String, String> metadata) {

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

@@ -237,6 +237,25 @@ class AppSecEventTrackerSpecification extends DDSpecification {
     'ident' | 'identification' | 'user doesn\'t exist' | USER_ID
   }
 
+  def "test onUserEvent (#mode)"() {
+    setup:
+    final collectionMode = UserIdCollectionMode.fromString(mode, null)
+
+    when:
+    tracker.onUserEvent(collectionMode, USER_ID)
+
+    then:
+    1 * traceSegment.getTagTop('_dd.appsec.user.collection_mode') >> null
+    1 * userCallback.apply(_ as RequestContext, collectionMode, expectedUserId) >> NoopFlow.INSTANCE
+    0 * _
+
+    where:
+    mode    | modeTag          | expectedUserId
+    'anon'  | 'anonymization'  | ANONYMIZED_USER_ID
+    'ident' | 'identification' | USER_ID
+  }
+
+
   def "test onUserNotFound (#mode)"() {
     setup:
     final collectionMode = UserIdCollectionMode.fromString(mode, null)

dd-java-agent/instrumentation/spring-security-5/build.gradle

@@ -3,39 +3,27 @@ muzzle {
     group = 'org.springframework.security'
     module = 'spring-security-core'
     versions = "[5.5.0,)"
-    // assertInverse = true
   }
 }
 
-ext {
-  minJavaVersionForTests = JavaVersion.VERSION_17
-}
-
 apply from: "$rootDir/gradle/java.gradle"
+addTestSuiteForDir('latestDepTest', 'test')
 
-[compileMain_java17Java, compileTestJava].each {
-  it.configure {
-    setJavaVersion(it, 17)
-    sourceCompatibility = JavaVersion.VERSION_1_8
-    targetCompatibility = JavaVersion.VERSION_1_8
-  }
-}
-
-compileTestGroovy {
-  javaLauncher = getJavaLauncherFor(17)
-}
+final springSecurityVersion = '5.8.2'
+final springBootVersion = '2.6.0'
 
 dependencies {
-  main_java17CompileOnly group: 'org.springframework.security', name: 'spring-security-core', version: '5.5.0'
-  compileOnly group: 'org.springframework.security', name: 'spring-security-core', version: '5.5.0'
-
-  testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-test', version: '3.0.0'
-  testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: '3.0.0'
-  testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-security', version: '3.0.0'
-  testImplementation 'org.springframework.boot:spring-boot-starter-jdbc:3.0.0'
+  compileOnly group: 'org.springframework.security', name: 'spring-security-core', version: springSecurityVersion
 
   testImplementation testFixtures(project(':dd-java-agent:appsec'))
-  testImplementation project(':dd-java-agent:instrumentation:tomcat-appsec-6')
-  testImplementation project(':dd-java-agent:instrumentation:tomcat-5.5')
-  testImplementation 'com.h2database:h2:2.1.212'
+  testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-test', version: springBootVersion
+  testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: springBootVersion
+  testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-security', version: springBootVersion
+
+  latestDepTestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-test', version: '2.+'
+  latestDepTestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: '2.+'
+  latestDepTestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-security', version: '2.+'
+
+  testRuntimeOnly project(':dd-java-agent:instrumentation:tomcat-appsec-6')
+  testRuntimeOnly project(':dd-java-agent:instrumentation:tomcat-5.5')
 }

dd-java-agent/instrumentation/spring-security-5/src/main/java/datadog/trace/instrumentation/springsecurity5/AppSecDeferredContext.java

@@ -0,0 +1,22 @@
+package datadog.trace.instrumentation.springsecurity5;
+
+import java.util.function.Supplier;
+import org.springframework.security.core.context.SecurityContext;
+
+public class AppSecDeferredContext implements Supplier<SecurityContext> {
+
+  private final Supplier<SecurityContext> delegate;
+
+  public AppSecDeferredContext(final Supplier<SecurityContext> delegate) {
+    this.delegate = delegate;
+  }
+
+  @Override
+  public SecurityContext get() {
+    SecurityContext context = delegate.get();
+    if (context != null) {
+      SpringSecurityUserEventDecorator.DECORATE.onUser(context.getAuthentication());
+    }
+    return context;
+  }
+}

dd-java-agent/instrumentation/spring-security-5/src/main/java/datadog/trace/instrumentation/springsecurity5/AuthenticationProviderInstrumentation.java

@@ -10,8 +10,12 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.ActiveSubsystems;
+import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.AuthenticationException;
 
 @AutoService(InstrumenterModule.class)
 public class AuthenticationProviderInstrumentation extends InstrumenterModule.AppSec
@@ -46,6 +50,19 @@ public void methodAdvice(MethodTransformer transformer) {
             .and(takesArgument(0, named("org.springframework.security.core.Authentication")))
             .and(returns(named("org.springframework.security.core.Authentication")))
             .and(isPublic()),
-        packageName + ".AuthenticationProviderAdvice");
+        getClass().getName() + "$AuthenticationProviderAdvice");
+  }
+
+  public static class AuthenticationProviderAdvice {
+
+    @Advice.OnMethodExit(onThrowable = AuthenticationException.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Argument(value = 0, readOnly = false) Authentication authentication,
+        @Advice.Return final Authentication result,
+        @Advice.Thrown final AuthenticationException throwable) {
+      if (ActiveSubsystems.APPSEC_ACTIVE) {
+        SpringSecurityUserEventDecorator.DECORATE.onLogin(authentication, throwable, result);
+      }
+    }
   }
 }

dd-java-agent/instrumentation/spring-security-5/src/main/java/datadog/trace/instrumentation/springsecurity5/SecurityContextHolderInstrumentation.java

@@ -0,0 +1,90 @@
+package datadog.trace.instrumentation.springsecurity5;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.ActiveSubsystems;
+import java.util.function.Supplier;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.springframework.security.core.context.SecurityContext;
+
+@AutoService(InstrumenterModule.class)
+public class SecurityContextHolderInstrumentation extends InstrumenterModule.AppSec
+    implements Instrumenter.ForTypeHierarchy {
+
+  public SecurityContextHolderInstrumentation() {
+    super("spring-security");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "org.springframework.security.core.context.SecurityContextHolderStrategy";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named(hierarchyMarkerType()));
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.instrumentation.springsecurity5.SpringSecurityUserEventDecorator",
+      "datadog.trace.instrumentation.springsecurity5.AppSecDeferredContext"
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("setContext"))
+            .and(takesArguments(1))
+            .and(
+                takesArgument(
+                    0, named("org.springframework.security.core.context.SecurityContext")))
+            .and(isPublic()),
+        getClass().getName() + "$SetSecurityContextAdvice");
+    transformer.applyAdvice(
+        isMethod().and(named("setDeferredContext")).and(takesArguments(1)).and(isPublic()),
+        getClass().getName() + "$SetDeferredSecurityContextAdvice");
+  }
+
+  public static class SetSecurityContextAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(@Advice.Argument(0) final SecurityContext context) {
+      if (context == null) {
+        return;
+      }
+      if (!ActiveSubsystems.APPSEC_ACTIVE) {
+        return;
+      }
+      SpringSecurityUserEventDecorator.DECORATE.onUser(context.getAuthentication());
+    }
+  }
+
+  public static class SetDeferredSecurityContextAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.Argument(value = 0, readOnly = false) Supplier<SecurityContext> deferred) {
+      if (deferred == null) {
+        return;
+      }
+      if (!ActiveSubsystems.APPSEC_ACTIVE) {
+        return;
+      }
+      deferred = new AppSecDeferredContext(deferred);
+    }
+  }
+}

dd-java-agent/instrumentation/spring-security-5/src/main/java17/datadog/trace/instrumentation/springsecurity5/SpringSecurityUserEventDecorator.java → dd-java-agent/instrumentation/spring-security-5/src/main/java/datadog/trace/instrumentation/springsecurity5/SpringSecurityUserEventDecorator.java

@@ -111,6 +111,25 @@ public void onLogin(Authentication authentication, Throwable throwable, Authenti
     }
   }
 
+  public void onUser(final Authentication authentication) {
+    if (authentication == null) {
+      return;
+    }
+
+    final AppSecEventTracker tracker = AppSecEventTracker.getEventTracker();
+    if (tracker == null) {
+      return;
+    }
+
+    if (shouldSkipAuthentication(authentication)) {
+      return;
+    }
+
+    UserIdCollectionMode mode = UserIdCollectionMode.get();
+    String userId = authentication.getName();
+    tracker.onUserEvent(mode, userId);
+  }
+
   private static boolean shouldSkipAuthentication(final Authentication authentication) {
     if (authentication instanceof UsernamePasswordAuthenticationToken) {
       return false;

dd-java-agent/instrumentation/spring-security-5/src/main/java/datadog/trace/instrumentation/springsecurity5/UserDetailsManagerInstrumentation.java

@@ -9,8 +9,11 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.ActiveSubsystems;
+import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
+import org.springframework.security.core.userdetails.UserDetails;
 
 @AutoService(InstrumenterModule.class)
 public class UserDetailsManagerInstrumentation extends InstrumenterModule.AppSec
@@ -46,6 +49,18 @@ public void methodAdvice(MethodTransformer transformer) {
                 takesArgument(
                     0, named("org.springframework.security.core.userdetails.UserDetails")))
             .and(isPublic()),
-        packageName + ".UserDetailsManagerAdvice");
+        getClass().getName() + "$UserDetailsManagerAdvice");
+  }
+
+  public static class UserDetailsManagerAdvice {
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Argument(value = 0, readOnly = false) UserDetails user,
+        @Advice.Thrown Throwable throwable) {
+      if (ActiveSubsystems.APPSEC_ACTIVE) {
+        SpringSecurityUserEventDecorator.DECORATE.onSignup(user, throwable);
+      }
+    }
   }
 }

dd-java-agent/instrumentation/spring-security-5/src/main/java/datadog/trace/instrumentation/springsecurity5/UsernameNotFoundExceptionInstrumentation.java

@@ -9,6 +9,7 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -41,6 +42,14 @@ public String[] helperClassNames() {
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
         isConstructor().and(takesArgument(0, named("java.lang.String"))).and(isPublic()),
-        packageName + ".UsernameNotFoundExceptionAdvice");
+        getClass().getName() + "$UsernameNotFoundExceptionAdvice");
+  }
+
+  public static class UsernameNotFoundExceptionAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter() {
+      SpringSecurityUserEventDecorator.DECORATE.onUserNotFound();
+    }
   }
 }