feat(servlet): add jakarta-servlet-6.0 instrumentation module

dd-java-agent/instrumentation/servlet/jakarta-servlet-6.0/build.gradle

@@ -0,0 +1,82 @@
+import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar
+
+plugins {
+  id 'java-test-fixtures'
+}
+
+muzzle {
+  pass {
+    group = "jakarta.servlet"
+    module = 'jakarta.servlet-api'
+    versions = "[6.0,)"
+  }
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+apply plugin: 'dd-trace-java.call-site-instrumentation'
+
+// jakarta.servlet-api dependencies are compiled with Java 11 and
+// the gradle muzzle tasks uses the JVM gradle is running with
+if (!JavaVersion.current().java11Compatible) {
+  project.afterEvaluate {
+    tasks.findAll {it.group == 'Muzzle' }.each {
+      logger.info("Disabling task $it.path (requires Java 11)")
+      it.enabled = false
+    }
+  }
+}
+
+configurations {
+  javaxClassesToRelocate
+}
+
+tasks.register('relocatedJavaxJar', ShadowJar) {
+  relocate 'javax.servlet', 'jakarta.servlet'
+  relocate 'datadog.trace.instrumentation.servlet3', 'datadog.trace.instrumentation.servlet6'
+  relocate 'datadog.trace.instrumentation.servlet', 'datadog.trace.instrumentation.servlet6'
+
+  archiveClassifier.set('relocated-javax')
+
+  configurations = [project.configurations.javaxClassesToRelocate]
+
+  include '**/*.jar'
+  include '**/Servlet31InputStreamWrapper.class'
+  include '**/HttpServletGetInputStreamAdvice.class'
+  include '**/HttpServletGetReaderAdvice.class'
+  include '**/BufferedReaderWrapper.class'
+  include '**/ServletBlockingHelper.class'
+  include '**/AbstractServletInputStreamWrapper.class'
+
+  includeEmptyDirs = false
+}
+
+dependencies {
+  implementation files(relocatedJavaxJar.outputs.files)
+  compileOnly group: 'jakarta.servlet', name: 'jakarta.servlet-api', version: '6.1.0'
+  testImplementation group: 'jakarta.servlet', name: 'jakarta.servlet-api', version: '6.1.0'
+  testImplementation group: 'jakarta.servlet.jsp', name: 'jakarta.servlet.jsp-api', version: '3.0.0'
+  testRuntimeOnly project(':dd-java-agent:instrumentation:datadog:asm:iast-instrumenter')
+
+  javaxClassesToRelocate project(':dd-java-agent:instrumentation:servlet:javax-servlet:javax-servlet-iast'), {
+    transitive = false
+  }
+  javaxClassesToRelocate project(':dd-java-agent:instrumentation:servlet:javax-servlet:javax-servlet-3.0'), {
+    transitive = false
+  }
+
+  testFixturesApi(project(':dd-java-agent:instrumentation-testing')) {
+    exclude group: 'org.eclipse.jetty', module: 'jetty-server'
+  }
+  testFixturesCompileOnly group: 'jakarta.servlet', name: 'jakarta.servlet-api', version: '6.1.0'
+
+  testImplementation libs.bundles.mockito
+
+  testFixturesCompileOnly(libs.bundles.groovy)
+  testFixturesCompileOnly(libs.bundles.spock)
+
+  // tested against jakarta.servlet-api 6.0+
+}
+
+tasks.named("jar", Jar) {
+  from zipTree(relocatedJavaxJar.outputs.files.asPath)
+}

dd-java-agent/instrumentation/servlet/jakarta-servlet-6.0/src/main/java/datadog/trace/instrumentation/servlet6/IastJakartaServletInstrumentation.java

@@ -0,0 +1,81 @@
+package datadog.trace.instrumentation.servlet6;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasSuperType;
+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.api.iast.InstrumentationBridge;
+import datadog.trace.api.iast.sink.ApplicationModule;
+import datadog.trace.bootstrap.InstrumentationContext;
+import jakarta.servlet.ServletContext;
+import jakarta.servlet.http.HttpServlet;
+import java.util.Collections;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumenterModule.class)
+public class IastJakartaServletInstrumentation extends InstrumenterModule.Iast
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
+  public IastJakartaServletInstrumentation() {
+    super("servlet", "servlet-6");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "jakarta.servlet.http.HttpServlet";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return hasSuperType(named(hierarchyMarkerType()));
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return Collections.singletonMap("jakarta.servlet.ServletContext", Boolean.class.getName());
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("service"))
+            .and(isPublic())
+            .and(takesArguments(2))
+            .and(takesArgument(0, named("jakarta.servlet.ServletRequest")))
+            .and(takesArgument(1, named("jakarta.servlet.ServletResponse"))),
+        getClass().getName() + "$IastAdvice");
+  }
+
+  @Override
+  protected boolean isOptOutEnabled() {
+    return true;
+  }
+
+  public static class IastAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void after(@Advice.This final HttpServlet servlet) {
+      final ApplicationModule applicationModule = InstrumentationBridge.APPLICATION;
+      if (applicationModule == null) {
+        return;
+      }
+      final ServletContext context = servlet.getServletContext();
+      if (InstrumentationContext.get(ServletContext.class, Boolean.class).get(context) != null) {
+        return;
+      }
+      InstrumentationContext.get(ServletContext.class, Boolean.class).put(context, true);
+      if (applicationModule != null) {
+        applicationModule.onRealPath(context.getRealPath("/"));
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/servlet/jakarta-servlet-6.0/src/main/java/datadog/trace/instrumentation/servlet6/IastOptOutJakartaHttpServletRequestInstrumentation.java

@@ -0,0 +1,104 @@
+package datadog.trace.instrumentation.servlet6;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
+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.*;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.iast.*;
+import datadog.trace.api.iast.sink.ApplicationModule;
+import datadog.trace.bootstrap.InstrumentationContext;
+import jakarta.servlet.ServletContext;
+import jakarta.servlet.SessionTrackingMode;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpSession;
+import java.util.*;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@SuppressWarnings("unused")
+@AutoService(InstrumenterModule.class)
+public class IastOptOutJakartaHttpServletRequestInstrumentation extends InstrumenterModule.Iast
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
+
+  private static final String CLASS_NAME =
+      IastOptOutJakartaHttpServletRequestInstrumentation.class.getName();
+  private static final ElementMatcher.Junction<? super TypeDescription> WRAPPER_CLASS =
+      named("jakarta.servlet.http.HttpServletRequestWrapper");
+
+  public IastOptOutJakartaHttpServletRequestInstrumentation() {
+    super("servlet", "servlet-6", "servlet-request");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "jakarta.servlet.http.HttpServletRequest";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named(hierarchyMarkerType()))
+        .and(not(WRAPPER_CLASS))
+        .and(not(extendsClass(WRAPPER_CLASS)));
+  }
+
+  @Override
+  protected boolean isOptOutEnabled() {
+    return true;
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("getSession").and(returns(named("jakarta.servlet.http.HttpSession"))).and(isPublic()),
+        CLASS_NAME + "$GetHttpSessionAdvice");
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return Collections.singletonMap(
+        "jakarta.servlet.ServletContext", "jakarta.servlet.SessionTrackingMode");
+  }
+
+  public static class GetHttpSessionAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Sink(VulnerabilityTypes.SESSION_REWRITING)
+    public static void onExit(
+        @Advice.This final HttpServletRequest request, @Advice.Return final HttpSession session) {
+      if (session == null) {
+        return;
+      }
+      final ApplicationModule module = InstrumentationBridge.APPLICATION;
+      if (module == null) {
+        return;
+      }
+      final ServletContext context = request.getServletContext();
+
+      if (InstrumentationContext.get(ServletContext.class, SessionTrackingMode.class).get(context)
+          != null) {
+        return;
+      }
+      // We only want to report it once per application
+      InstrumentationContext.get(ServletContext.class, SessionTrackingMode.class)
+          .put(context, SessionTrackingMode.URL);
+      if (context.getEffectiveSessionTrackingModes() != null
+          && !context.getEffectiveSessionTrackingModes().isEmpty()) {
+        Set<String> sessionTrackingModes = new HashSet<>();
+        for (SessionTrackingMode mode : context.getEffectiveSessionTrackingModes()) {
+          sessionTrackingModes.add(mode.name());
+        }
+        module.checkSessionTrackingModes(sessionTrackingModes);
+      }
+    }
+  }
+
+  @Override
+  public int order() {
+    // apply this instrumentation after the regular servlet one.
+    return 1;
+  }
+}

dd-java-agent/instrumentation/servlet/jakarta-servlet-6.0/src/main/java/datadog/trace/instrumentation/servlet6/JakartaHttpServletRequestCallSite.java

@@ -0,0 +1,67 @@
+package datadog.trace.instrumentation.servlet6;
+
+import datadog.trace.agent.tooling.csi.CallSite;
+import datadog.trace.api.iast.IastCallSites;
+import datadog.trace.api.iast.IastContext;
+import datadog.trace.api.iast.InstrumentationBridge;
+import datadog.trace.api.iast.Source;
+import datadog.trace.api.iast.SourceTypes;
+import datadog.trace.api.iast.propagation.PropagationModule;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import jakarta.servlet.http.HttpServletRequest;
+
+/**
+ * Calls to these methods are often triggered outside of customer code, we use call sites to avoid
+ * all these unwanted tainting operations
+ */
+@CallSite(spi = IastCallSites.class)
+public class JakartaHttpServletRequestCallSite {
+
+  @Source(SourceTypes.REQUEST_PATH)
+  @CallSite.After("java.lang.String jakarta.servlet.http.HttpServletRequest.getRequestURI()")
+  @CallSite.After("java.lang.String jakarta.servlet.http.HttpServletRequestWrapper.getRequestURI()")
+  @CallSite.After("java.lang.String jakarta.servlet.http.HttpServletRequest.getPathInfo()")
+  @CallSite.After("java.lang.String jakarta.servlet.http.HttpServletRequestWrapper.getPathInfo()")
+  @CallSite.After("java.lang.String jakarta.servlet.http.HttpServletRequest.getPathTranslated()")
+  @CallSite.After(
+      "java.lang.String jakarta.servlet.http.HttpServletRequestWrapper.getPathTranslated()")
+  public static String afterPath(
+      @CallSite.This final HttpServletRequest self, @CallSite.Return final String retValue) {
+    if (null != retValue && !retValue.isEmpty()) {
+      final PropagationModule module = InstrumentationBridge.PROPAGATION;
+      if (module != null) {
+        try {
+          final IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan());
+          if (ctx != null) {
+            module.taintString(ctx, retValue, SourceTypes.REQUEST_PATH);
+          }
+        } catch (final Throwable e) {
+          module.onUnexpectedException("afterPath threw", e);
+        }
+      }
+    }
+    return retValue;
+  }
+
+  @Source(SourceTypes.REQUEST_URI)
+  @CallSite.After("java.lang.StringBuffer jakarta.servlet.http.HttpServletRequest.getRequestURL()")
+  @CallSite.After(
+      "java.lang.StringBuffer jakarta.servlet.http.HttpServletRequestWrapper.getRequestURL()")
+  public static StringBuffer afterGetRequestURL(
+      @CallSite.This final HttpServletRequest self, @CallSite.Return final StringBuffer retValue) {
+    if (null != retValue && retValue.length() > 0) {
+      final PropagationModule module = InstrumentationBridge.PROPAGATION;
+      if (module != null) {
+        try {
+          final IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan());
+          if (ctx != null) {
+            module.taintObject(ctx, retValue, SourceTypes.REQUEST_URI);
+          }
+        } catch (final Throwable e) {
+          module.onUnexpectedException("afterGetRequestURL threw", e);
+        }
+      }
+    }
+    return retValue;
+  }
+}