Add Thymeleaf support to IAST XSS vulnerability (#5901)

#What Does This Do
Instrument Thymeleaf classes that allow template generation avoiding escape functions

#Additional Notes
Is not possible to get tainted tex, templateName and line in the same class, so two instrumentation sharing context are needed.

org.thymeleaf.standard.processor.StandardUtextTagProcessor#doProcess -> Set IElementTagStructureHandler as potentially dangerous storing it with template name and line in the instrumentation context

org.thymeleaf.processor.element.IElementTagStructureHandler#setBody -> if is in the instrumentation context check if the input is tainted to report the vulnerability

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/Location.java

@@ -26,6 +26,11 @@ public static Location forSpanAndClassAndMethod(
     return new Location(spanId, clazz, -1, method);
   }
 
+  public static Location forSpanAndFileAndLine(
+      final long spanId, final String file, final int line) {
+    return new Location(spanId, file, line, null);
+  }
+
   public long getSpanId() {
     return spanId == null ? 0 : spanId;
   }

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/XssModuleImpl.java

@@ -56,7 +56,7 @@ public void onXss(@Nonnull String s, @Nonnull String clazz, @Nonnull String meth
     if (!overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, span)) {
       return;
     }
-    final Evidence evidence = new Evidence(s.toString(), notMarkedRanges);
+    final Evidence evidence = new Evidence(s, notMarkedRanges);
     reporter.report(
         span,
         new Vulnerability(
@@ -92,4 +92,35 @@ public void onXss(@Nonnull String format, @Nullable Object[] args) {
     checkInjection(
         span, VulnerabilityType.XSS, rangesProviderFor(to, format), rangesProviderFor(to, args));
   }
+
+  @Override
+  public void onXss(@Nonnull CharSequence s, @Nullable String file, int line) {
+    if (!canBeTainted(s) || file == null || file.isEmpty()) {
+      return;
+    }
+    final AgentSpan span = AgentTracer.activeSpan();
+    final IastRequestContext ctx = IastRequestContext.get(span);
+    if (ctx == null) {
+      return;
+    }
+    TaintedObject taintedObject = ctx.getTaintedObjects().get(s);
+    if (taintedObject == null) {
+      return;
+    }
+    Range[] notMarkedRanges =
+        Ranges.getNotMarkedRanges(taintedObject.getRanges(), VulnerabilityType.XSS.mark());
+    if (notMarkedRanges == null || notMarkedRanges.length == 0) {
+      return;
+    }
+    if (!overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, span)) {
+      return;
+    }
+    final Evidence evidence = new Evidence(s.toString(), notMarkedRanges);
+    reporter.report(
+        span,
+        new Vulnerability(
+            VulnerabilityType.XSS,
+            Location.forSpanAndFileAndLine(span.getSpanId(), file, line),
+            evidence));
+  }
 }

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/sink/XssModuleTest.groovy

@@ -119,6 +119,30 @@ class XssModuleTest extends IastModuleImplTestBase {
     '/==>var<==' | ['a', 'b']             | VulnerabilityMarks.SQL_INJECTION_MARK | "/==>var<== a b"
   }
 
+  void 'module detects Charsequence XSS with file and line'() {
+    setup:
+    final param = mapTainted(s, mark)
+
+    when:
+    module.onXss(param as CharSequence, file as String, line as int)
+
+    then:
+    if (expected != null) {
+      1 * reporter.report(_, _) >> { args -> assertEvidence(args[1] as Vulnerability, expected) }
+    } else {
+      0 * reporter.report(_, _)
+    }
+
+    where:
+    s     | file | line        | mark                                  | expected
+    null    | 'test' | 3      | NOT_MARKED                            | null
+    '/var'    | 'test' | 3    | NOT_MARKED                            | null
+    '/==>var<=='| 'test' | 3  | NOT_MARKED                            | "/==>var<=="
+    '/==>var<=='| 'test' | 3  | VulnerabilityMarks.XSS_MARK           | null
+    '/==>var<=='| 'test' | 3  | VulnerabilityMarks.SQL_INJECTION_MARK | "/==>var<=="
+    '/==>var<=='| null | 3  | VulnerabilityMarks.SQL_INJECTION_MARK | null
+  }
+
   void 'iast module detects String xss with class and method (#value)'() {
     setup:
     final param = mapTainted(value, mark)

dd-java-agent/instrumentation/thymeleaf/build.gradle

@@ -0,0 +1,23 @@
+muzzle {
+  pass {
+    group = 'org.thymeleaf'
+    module = 'thymeleaf'
+    versions = '[3.0.0.RELEASE,]'
+    assertInverse = true
+  }
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
+addTestSuiteForDir('latestDepTest', 'test')
+
+dependencies {
+
+  compileOnly group: 'org.thymeleaf', name: 'thymeleaf', version: '3.0.0.RELEASE'
+
+  testImplementation group: 'org.thymeleaf', name: 'thymeleaf', version: '3.0.0.RELEASE'
+
+  testRuntimeOnly project(':dd-java-agent:instrumentation:iast-instrumenter')
+
+  latestDepTestImplementation group: 'org.thymeleaf', name: 'thymeleaf', version: '+'
+}

dd-java-agent/instrumentation/thymeleaf/src/main/java/datadog/trace/instrumentation/thymeleaf/BodyAdvice.java

@@ -0,0 +1,28 @@
+package datadog.trace.instrumentation.thymeleaf;
+
+import datadog.trace.api.iast.InstrumentationBridge;
+import datadog.trace.api.iast.Sink;
+import datadog.trace.api.iast.VulnerabilityTypes;
+import datadog.trace.api.iast.sink.XssModule;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import net.bytebuddy.asm.Advice;
+import org.thymeleaf.engine.ElementTagStructureHandler;
+import org.thymeleaf.processor.element.IElementTagStructureHandler;
+
+public class BodyAdvice {
+  @Advice.OnMethodExit(suppress = Throwable.class)
+  @Sink(VulnerabilityTypes.XSS)
+  public static void setBody(
+      @Advice.This ElementTagStructureHandler self, @Advice.Argument(0) final CharSequence text) {
+    final XssModule module = InstrumentationBridge.XSS;
+    if (module != null) {
+      ContextStore<IElementTagStructureHandler, ThymeleafContext> contextStore =
+          InstrumentationContext.get(IElementTagStructureHandler.class, ThymeleafContext.class);
+      ThymeleafContext ctx = contextStore.get(self);
+      if (ctx != null) {
+        module.onXss(text, ctx.getTemplateName(), ctx.getLine());
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/thymeleaf/src/main/java/datadog/trace/instrumentation/thymeleaf/ElementTagStructureHandlerInstrumentation.java

@@ -0,0 +1,43 @@
+package datadog.trace.instrumentation.thymeleaf;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import java.util.Map;
+
+@AutoService(Instrumenter.class)
+public class ElementTagStructureHandlerInstrumentation extends Instrumenter.Iast
+    implements Instrumenter.ForSingleType {
+  public ElementTagStructureHandlerInstrumentation() {
+    super("thymeleaf");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "org.thymeleaf.engine.ElementTagStructureHandler";
+  }
+
+  @Override
+  public void adviceTransformations(AdviceTransformation transformation) {
+
+    transformation.applyAdvice(
+        isMethod().and(named("setBody")).and(takesArgument(0, CharSequence.class)),
+        packageName + ".BodyAdvice");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".ThymeleafContext"};
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap(
+        "org.thymeleaf.processor.element.IElementTagStructureHandler",
+        "datadog.trace.instrumentation.thymeleaf.ThymeleafContext");
+  }
+}

dd-java-agent/instrumentation/thymeleaf/src/main/java/datadog/trace/instrumentation/thymeleaf/ProcessAdvice.java

@@ -0,0 +1,24 @@
+package datadog.trace.instrumentation.thymeleaf;
+
+import datadog.trace.api.iast.InstrumentationBridge;
+import datadog.trace.api.iast.Propagation;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import net.bytebuddy.asm.Advice;
+import org.thymeleaf.model.IProcessableElementTag;
+import org.thymeleaf.processor.element.IElementTagStructureHandler;
+
+public class ProcessAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  @Propagation
+  public static void doProcess(
+      @Advice.Argument(1) final IProcessableElementTag tag,
+      @Advice.Argument(4) final IElementTagStructureHandler handler) {
+    if (InstrumentationBridge.XSS != null) {
+      ContextStore<IElementTagStructureHandler, ThymeleafContext> contextStore =
+          InstrumentationContext.get(IElementTagStructureHandler.class, ThymeleafContext.class);
+      contextStore.put(handler, new ThymeleafContext(tag.getTemplateName(), tag.getLine()));
+    }
+  }
+}

dd-java-agent/instrumentation/thymeleaf/src/main/java/datadog/trace/instrumentation/thymeleaf/StandardUtextTagProcessorInstrumentation.java

@@ -0,0 +1,40 @@
+package datadog.trace.instrumentation.thymeleaf;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import java.util.Map;
+
+@AutoService(Instrumenter.class)
+public class StandardUtextTagProcessorInstrumentation extends Instrumenter.Iast
+    implements Instrumenter.ForSingleType {
+
+  public StandardUtextTagProcessorInstrumentation() {
+    super("thymeleaf");
+  }
+
+  @Override
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(isMethod().and(named("doProcess")), packageName + ".ProcessAdvice");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".ThymeleafContext"};
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap(
+        "org.thymeleaf.processor.element.IElementTagStructureHandler",
+        "datadog.trace.instrumentation.thymeleaf.ThymeleafContext");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "org.thymeleaf.standard.processor.StandardUtextTagProcessor";
+  }
+}

dd-java-agent/instrumentation/thymeleaf/src/main/java/datadog/trace/instrumentation/thymeleaf/ThymeleafContext.java

@@ -0,0 +1,21 @@
+package datadog.trace.instrumentation.thymeleaf;
+
+public class ThymeleafContext {
+
+  private final String templateName;
+
+  private final int line;
+
+  public ThymeleafContext(final String file, final int line) {
+    this.templateName = file;
+    this.line = line;
+  }
+
+  public String getTemplateName() {
+    return templateName;
+  }
+
+  public int getLine() {
+    return line;
+  }
+}

dd-java-agent/instrumentation/thymeleaf/src/test/groovy/datadog/trace/instrumentation/thymeleaf/ThymeleafXssTest.groovy

@@ -0,0 +1,38 @@
+package datadog.trace.instrumentation.thymeleaf
+
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.iast.InstrumentationBridge
+import datadog.trace.api.iast.sink.XssModule
+import org.thymeleaf.TemplateEngine
+import org.thymeleaf.context.Context
+
+class ThymeleafXssTest extends AgentTestRunner {
+
+  @Override
+  protected void configurePreAgent() {
+    injectSysConfig("dd.iast.enabled", "true")
+  }
+
+  void 'test that XssModule is called' (){
+    given:
+    XssModule xssModule = Mock(XssModule)
+    if(hasModule){
+      InstrumentationBridge.registerIastModule(xssModule)
+    }
+    final engine = new TemplateEngine()
+    final context = new Context(Locale.getDefault(), [q:'this is vulnerable'])
+
+    when:
+    engine.process(template, context)
+
+    then:
+    expected * xssModule.onXss('this is vulnerable', template, 1)
+    0 * _
+
+    where:
+    hasModule | template | expected
+    false | '<span th:utext="${q}"></span>'  | 0
+    true | '<span th:utext="${q}"></span>'  | 1
+    true | '<span th:text="${q}"></span>'  | 0
+  }
+}

dd-smoke-tests/springboot-thymeleaf/build.gradle

@@ -0,0 +1,30 @@
+plugins {
+  id 'java'
+  id 'org.springframework.boot' version '2.7.15'
+  id 'io.spring.dependency-management' version '1.0.15.RELEASE'
+  id 'java-test-fixtures'
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+description = 'SpringBoot thymeleaf 3 Smoke Tests.'
+
+java {
+  sourceCompatibility = '1.8'
+}
+
+repositories {
+  mavenCentral()
+}
+
+dependencies {
+  implementation 'org.springframework.boot:spring-boot-starter-web'
+  implementation 'org.springframework.boot:spring-boot-starter-thymeleaf'
+
+  testImplementation project(':dd-smoke-tests')
+  testImplementation(testFixtures(project(":dd-smoke-tests:iast-util")))
+}
+
+tasks.withType(Test).configureEach {
+  dependsOn "bootJar"
+  jvmArgs "-Ddatadog.smoketest.springboot.shadowJar.path=${tasks.bootJar.archiveFile.get()}"
+}

dd-smoke-tests/springboot-thymeleaf/src/main/java/datadog/smoketest/springboot/SpringbootApplication.java

@@ -0,0 +1,11 @@
+package datadog.smoketest.springboot;
+
+import org.springframework.boot.SpringApplication;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+
+@SpringBootApplication
+public class SpringbootApplication {
+  public static void main(String[] args) {
+    SpringApplication.run(SpringbootApplication.class, args);
+  }
+}

dd-smoke-tests/springboot-thymeleaf/src/main/java/datadog/smoketest/springboot/XssController.java

@@ -0,0 +1,18 @@
+package datadog.smoketest.springboot;
+
+import org.springframework.stereotype.Controller;
+import org.springframework.ui.Model;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RequestParam;
+
+@Controller
+@RequestMapping("/xss")
+public class XssController {
+
+  @GetMapping("/utext")
+  public String utext(@RequestParam(name = "string") String name, Model model) {
+    model.addAttribute("xss", name);
+    return "utext";
+  }
+}

dd-smoke-tests/springboot-thymeleaf/src/main/resources/templates/utext.html

@@ -0,0 +1,16 @@
+<!DOCTYPE html>
+
+<html xmlns:th="http://www.thymeleaf.org">
+
+<head>
+  <title>Good Thymes Virtual Grocery</title>
+  <meta http-equiv="Content-Type" content="text/html; charset=UTF-8" />
+</head>
+
+<body>
+<p>Testing xss for utext...</p>
+<p th:utext="${xss}">Test!</p>
+
+</body>
+
+</html>