APPSEC-55380 String taint tracking: translateEscapes

Author: sezen-datadog

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/StringModuleImpl.java

@@ -33,6 +33,7 @@
 import java.util.stream.Stream;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import org.jetbrains.annotations.NotNull;
 
 public class StringModuleImpl implements StringModule {
 
@@ -293,6 +294,34 @@ public void onStringJoin(
     }
   }
 
+  @Override
+  public void onStringTranslateEscapes(
+      @NotNull String self, @org.jetbrains.annotations.Nullable String result) {
+    if (!canBeTainted(result)) {
+      return;
+    }
+    if (self == result) { // same ref, no change in taint status
+      return;
+    }
+    final IastContext ctx = IastContext.Provider.get();
+    if (ctx == null) {
+      return;
+    }
+    final TaintedObjects taintedObjects = ctx.getTaintedObjects();
+    final TaintedObject taintedSelf = taintedObjects.get(self);
+    if (taintedSelf == null) {
+      return; // original string is not tainted
+    }
+    final Range[] rangesSelf = taintedSelf.getRanges();
+    if (rangesSelf.length == 0) {
+      return; // original string is not tainted
+    }
+    final Range[] newRanges = Ranges.forSubstring(0, result.length(), rangesSelf);
+    if (newRanges != null) {
+      taintedObjects.taint(result, newRanges); // only possibility left
+    }
+  }
+
   @Override
   @SuppressFBWarnings("ES_COMPARING_PARAMETER_STRING_WITH_EQ")
   public void onStringRepeat(@Nonnull String self, int count, @Nonnull String result) {

dd-java-agent/instrumentation/java-lang/java-lang-15/build.gradle

@@ -0,0 +1,43 @@
+plugins {
+  id 'idea'
+}
+
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_15
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+apply plugin: 'call-site-instrumentation'
+
+muzzle {
+  pass {
+    coreJdk()
+  }
+}
+
+idea {
+  module {
+    jdkName = '17'
+  }
+}
+
+csi {
+  javaVersion = JavaLanguageVersion.of(15)
+}
+
+addTestSuiteForDir('latestDepTest', 'test')
+
+dependencies {
+  testRuntimeOnly project(':dd-java-agent:instrumentation:iast-instrumenter')
+}
+
+project.tasks.withType(AbstractCompile).configureEach {
+  setJavaVersion(it, 17)
+  if (it.name != 'compileCsiJava') {
+    sourceCompatibility = JavaVersion.VERSION_15
+    targetCompatibility = JavaVersion.VERSION_15
+    if (it instanceof JavaCompile) {
+      it.options.release.set(15)
+    }
+  }
+}

dd-java-agent/instrumentation/java-lang/java-lang-15/gradle.lockfile

@@ -0,0 +1,28 @@
+package datadog.trace.instrumentation.java.lang.jdk15;
+
+import datadog.trace.agent.tooling.csi.CallSite;
+import datadog.trace.api.iast.IastCallSites;
+import datadog.trace.api.iast.InstrumentationBridge;
+import datadog.trace.api.iast.Propagation;
+import datadog.trace.api.iast.propagation.StringModule;
+
+@Propagation
+@CallSite(
+    spi = IastCallSites.class,
+    enabled = {"datadog.trace.api.iast.IastEnabledChecks", "isMajorJavaVersionAtLeast", "15"})
+public class StringCallSite {
+  @CallSite.After("java.lang.String java.lang.String.translateEscapes()")
+  public static String afterTranslateEscapes(
+      @CallSite.This final String self,
+      @CallSite.Return final String result) {
+    final StringModule module = InstrumentationBridge.STRING;
+    try {
+      if (module != null) {
+        module.onTranslateEscapes(self, result);
+      }
+    } catch (final Throwable e) {
+      module.onUnexpectedException("afterTranslateEscapes threw", e);
+    }
+    return result;
+  }
+}

dd-java-agent/instrumentation/java-lang/java-lang-15/src/main/java/datadog/trace/instrumentation/java/lang/jdk15/StringCallSite.java

@@ -0,0 +1,35 @@
+package datadog.trace.instrumentation.java.lang.jdk15
+
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.iast.InstrumentationBridge
+import datadog.trace.api.iast.propagation.StringModule
+import foo.bar.TestStringJDK15Suite
+import spock.lang.Requires
+
+@Requires({
+  jvm.java15Compatible
+})
+class StringCallSiteTest extends AgentTestRunner {
+
+  @Override
+  protected void configurePreAgent() {
+    injectSysConfig("dd.iast.enabled", "true")
+  }
+
+  def 'test string translate escapes call site'() {
+    setup:
+    final iastModule = Mock(StringModule)
+    InstrumentationBridge.registerIastModule(iastModule)
+
+    when:
+    final result = TestStringJDK15Suite.stringTranslateEscapes(input)
+
+    then:
+    result == output
+    1 * iastModule.onStringTranslateEscapes(input, output)
+
+    where:
+    input                   | output
+    'Hello\tThis is a line' | 'Hello\n   This is a line\n'
+  }
+}

dd-java-agent/instrumentation/java-lang/java-lang-15/src/test/groovy/datadog/trace/instrumentation/java/lang/jdk15/StringCallSiteTest.groovy

@@ -0,0 +1,18 @@
+package foo.bar;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public abstract class TestStringJDK15Suite {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(TestStringJDK15Suite.class);
+
+  private TestStringJDK15Suite() {}
+
+  public static String stringTranslateEscapes(String self) {
+    LOGGER.debug("Before string translate escapes {}", self);
+    final String result = self.translateEscapes();
+    LOGGER.debug("After string translate escapes {}", result);
+    return result;
+  }
+}

dd-java-agent/instrumentation/java-lang/java-lang-15/src/test/java/foo/bar/TestStringJDK15Suite.java

@@ -35,6 +35,8 @@ void onStringJoin(
 
   void onStringToUpperCase(@Nonnull String self, @Nullable String result);
 
+  void onStringTranslateEscapes(@Nonnull String self, @Nullable String result);
+
   void onStringToLowerCase(@Nonnull String self, @Nullable String result);
 
   void onStringTrim(@Nonnull String self, @Nullable String result);