Improve propagation on URL and URI call sites

DataDog · Jul 10, 2024 · f35169a · f35169a
1 parent c399041
commit f35169a
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 19 deletions.
diff --git a/...umentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URICallSite.java b/...umentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URICallSite.java
@@ -1,5 +1,7 @@
 package datadog.trace.instrumentation.java.net;
 
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
+
 import datadog.trace.agent.tooling.csi.CallSite;
 import datadog.trace.api.iast.IastCallSites;
 import datadog.trace.api.iast.InstrumentationBridge;
@@ -53,28 +55,50 @@ public static URI afterCtor(
     return result;
   }
 
+  /**
+   * Internally the URI is tainted following the <code>toString</code> representation
+   *
+   * @see CodecModule#onUriCreate(URI, Object...)
+   */
   @CallSite.After("java.lang.String java.net.URI.toString()")
-  @CallSite.After("java.lang.String java.net.URI.toASCIIString()")
   public static String afterToString(
       @CallSite.This final URI url, @CallSite.Return final String result) {
     final PropagationModule module = InstrumentationBridge.PROPAGATION;
     if (module != null) {
       try {
-        module.taintStringIfTainted(result, url);
+        module.taintStringIfTainted(result, url, true, NOT_MARKED);
       } catch (final Throwable e) {
         module.onUnexpectedException("After toString threw", e);
       }
     }
     return result;
   }
 
+  /** @see #afterToString(URI, String) */
+  @CallSite.After("java.lang.String java.net.URI.toASCIIString()")
+  public static String afterToASCIIString(
+      @CallSite.This final URI url, @CallSite.Return final String result) {
+    final PropagationModule module = InstrumentationBridge.PROPAGATION;
+    if (module != null && result != null) {
+      try {
+        boolean keepRanges = url.toString().equals(result);
+        module.taintStringIfTainted(result, url, keepRanges, NOT_MARKED);
+      } catch (final Throwable e) {
+        module.onUnexpectedException("After toASCIIString threw", e);
+      }
+    }
+    return result;
+  }
+
+  /** @see #afterToString(URI, String) */
   @CallSite.After("java.net.URI java.net.URI.normalize()")
   public static URI afterNormalize(
       @CallSite.This final URI url, @CallSite.Return final URI result) {
     final PropagationModule module = InstrumentationBridge.PROPAGATION;
-    if (module != null) {
+    if (module != null && result != null) {
       try {
-        module.taintObjectIfTainted(result, url);
+        boolean keepRanges = url.toString().equals(result.toString());
+        module.taintObjectIfTainted(result, url, keepRanges, NOT_MARKED);
       } catch (final Throwable e) {
         module.onUnexpectedException("After toString threw", e);
       }

diff --git a/...umentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URLCallSite.java b/...umentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URLCallSite.java
@@ -1,5 +1,7 @@
 package datadog.trace.instrumentation.java.net;
 
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
+
 import datadog.trace.agent.tooling.csi.CallSite;
 import datadog.trace.api.iast.IastCallSites;
 import datadog.trace.api.iast.InstrumentationBridge;
@@ -42,29 +44,52 @@ public static URL afterCtor(
     return result;
   }
 
+  /**
+   * Internally the URL is tainted following the <code>toString</code> representation
+   *
+   * @see CodecModule#onUrlCreate(URL, Object...)
+   */
   @Propagation
   @CallSite.After("java.lang.String java.net.URL.toString()")
-  @CallSite.After("java.lang.String java.net.URL.toExternalForm()")
   public static String afterToString(
       @CallSite.This final URL url, @CallSite.Return final String result) {
     final PropagationModule module = InstrumentationBridge.PROPAGATION;
-    if (module != null) {
+    if (module != null && result != null) {
       try {
-        module.taintStringIfTainted(result, url);
+        module.taintStringIfTainted(result, url, true, NOT_MARKED);
       } catch (final Throwable e) {
         module.onUnexpectedException("After toString threw", e);
       }
     }
     return result;
   }
 
+  /** @see #afterToString(URL, String) */
+  @Propagation
+  @CallSite.After("java.lang.String java.net.URL.toExternalForm()")
+  public static String afterToExternalForm(
+      @CallSite.This final URL url, @CallSite.Return final String result) {
+    final PropagationModule module = InstrumentationBridge.PROPAGATION;
+    if (module != null && result != null) {
+      try {
+        boolean keepRanges = url.toString().equals(result);
+        module.taintStringIfTainted(result, url, keepRanges, NOT_MARKED);
+      } catch (final Throwable e) {
+        module.onUnexpectedException("After toExternalForm threw", e);
+      }
+    }
+    return result;
+  }
+
+  /** @see #afterToString(URL, String) */
   @Propagation
   @CallSite.After("java.net.URI java.net.URL.toURI()")
   public static URI afterToURI(@CallSite.This final URL url, @CallSite.Return final URI result) {
     final PropagationModule module = InstrumentationBridge.PROPAGATION;
-    if (module != null) {
+    if (module != null && result != null) {
       try {
-        module.taintObjectIfTainted(result, url);
+        boolean keepRanges = url.toString().equals(result.toString());
+        module.taintObjectIfTainted(result, url, keepRanges, NOT_MARKED);
       } catch (final Throwable e) {
         module.onUnexpectedException("After toURI threw", e);
       }

diff --git a/...on/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URICallSIteTest.groovy b/...on/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URICallSIteTest.groovy
@@ -60,12 +60,14 @@ class URICallSIteTest extends AgentTestRunner {
     TestURICallSiteSuite.&"$method".call(args as Object[])
 
     then:
-    1 * module."taint${target}IfTainted"(_, _ as URI)
+    1 * module."taint${target}IfTainted"(_, _ as URI, keepRanges, _)
 
     where:
-    method          | target   | args
-    'normalize'     | 'Object' | [new URI('http://test.com/index?name=value#fragment')]
-    'toString'      | 'String' | [new URI('http://test.com/index?name=value#fragment')]
-    'toASCIIString' | 'String' | [new URI('http://test.com/index?name=value#fragment')]
+    method          | target   | args                                                           | keepRanges
+    'normalize'     | 'Object' | [new URI('http://test.com/index?name=value#fragment')]         | true
+    'normalize'     | 'Object' | [new URI('http://test.com/test/../index?name=value#fragment')] | false
+    'toASCIIString' | 'String' | [new URI('http://test.com/index?name=value#fragment')]         | true
+    'toASCIIString' | 'String' | [new URI('http://test.com/漢/index?name=value#fragment')]      | false
+    'toString'      | 'String' | [new URI('http://test.com/index?name=value#fragment')]         | true
   }
 }
diff --git a/...on/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URLCallSiteTest.groovy b/...on/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URLCallSiteTest.groovy
@@ -45,13 +45,13 @@ class URLCallSiteTest extends AgentTestRunner {
     TestURLCallSiteSuite.&"$method".call(args as Object[])
 
     then:
-    1 * module."taint${target}IfTainted"(_, _ as URL)
+    1 * module."taint${target}IfTainted"(_, _ as URL, keepRanges, _)
 
     where:
-    method           | target   | args
-    'toURI'          | 'Object' | [new URL('http://test.com/index?name=value#fragment')]
-    'toString'       | 'String' | [new URL('http://test.com/index?name=value#fragment')]
-    'toExternalForm' | 'String' | [new URL('http://test.com/index?name=value#fragment')]
+    method           | target   | args                                                   | keepRanges
+    'toURI'          | 'Object' | [new URL('http://test.com/index?name=value#fragment')] | true
+    'toExternalForm' | 'String' | [new URL('http://test.com/index?name=value#fragment')] | true
+    'toString'       | 'String' | [new URL('http://test.com/index?name=value#fragment')] | true
   }
 
   void 'test ssrf endpoints'() {