Update URI and URL call sites for precise taint tracking

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

@@ -1,9 +1,16 @@
 package com.datadog.iast.propagation;
 
 import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
+import static datadog.trace.api.iast.VulnerabilityMarks.XSS_MARK;
 
+import com.datadog.iast.taint.TaintedObject;
+import com.datadog.iast.taint.TaintedObjects;
+import com.datadog.iast.util.RangeBuilder;
+import datadog.trace.api.iast.IastContext;
 import datadog.trace.api.iast.propagation.CodecModule;
 import datadog.trace.api.iast.propagation.PropagationModule;
+import java.net.URI;
+import java.net.URL;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
@@ -25,6 +32,13 @@ public void onUrlDecode(
     propagationModule.taintStringIfTainted(result, value);
   }
 
+  @Override
+  public void onUrlEncode(
+      @Nonnull final String value, @Nullable final String encoding, @Nonnull final String result) {
+    // the new string should be safe to be used in
+    propagationModule.taintStringIfTainted(result, value, false, XSS_MARK);
+  }
+
   @Override
   public void onStringFromBytes(
       @Nonnull final byte[] value,
@@ -51,4 +65,63 @@ public void onBase64Encode(@Nullable byte[] value, @Nullable byte[] result) {
   public void onBase64Decode(@Nullable byte[] value, @Nullable byte[] result) {
     propagationModule.taintObjectIfTainted(result, value);
   }
+
+  @Override
+  public void onUriCreate(@Nonnull final URI result, final Object... args) {
+    final IastContext ctx = IastContext.Provider.get();
+    if (ctx == null) {
+      return;
+    }
+    taintUrlIfAnyTainted(ctx.getTaintedObjects(), result, args);
+  }
+
+  @Override
+  public void onUrlCreate(@Nonnull final URL result, final Object... args) {
+    final IastContext ctx = IastContext.Provider.get();
+    if (ctx == null) {
+      return;
+    }
+    final TaintedObjects to = ctx.getTaintedObjects();
+    if (args.length > 0 && args[0] instanceof URL) {
+      final TaintedObject tainted = to.get(args[0]);
+      if (tainted != null) {
+        final URL context = (URL) args[0];
+        // TODO improve matching when using a URL context
+        // Checking if the toString representation of the context is present in the final URL is a
+        // bit fragile, the spec might override some of it's parts (e.g. the final file). We can get
+        // away with this as having a context with tainted values is probably pretty rare.
+        if (!context.getProtocol().equals(result.getProtocol())
+            || !context.getHost().equals(result.getHost())) {
+          args[0] = null;
+        }
+      }
+    }
+    taintUrlIfAnyTainted(to, result, args);
+  }
+
+  private void taintUrlIfAnyTainted(
+      @Nonnull final TaintedObjects to, @Nonnull final Object url, @Nonnull final Object... args) {
+    final String toString = url.toString();
+    final RangeBuilder builder = new RangeBuilder();
+    boolean hasTainted = false;
+    for (final Object arg : args) {
+      final TaintedObject tainted = to.get(arg);
+      if (tainted != null) {
+        hasTainted = true;
+        final int offset = toString.indexOf(arg.toString());
+        if (offset >= 0) {
+          builder.add(tainted.getRanges(), offset);
+        }
+      }
+    }
+    if (!hasTainted) {
+      return;
+    }
+    if (builder.isEmpty()) {
+      // no mappings of tainted values in the URL, resort to fully tainting it
+      propagationModule.taintObjectIfAnyTainted(url, args);
+    } else {
+      to.taint(url, builder.toArray());
+    }
+  }
 }

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/CodecModuleTest.groovy

@@ -13,6 +13,9 @@ import datadog.trace.bootstrap.instrumentation.api.AgentTracer
 import java.nio.charset.StandardCharsets
 
 import static com.datadog.iast.taint.TaintUtils.addFromTaintFormat
+import static com.datadog.iast.taint.TaintUtils.fromTaintFormat
+import static com.datadog.iast.taint.TaintUtils.getStringFromTaintFormat
+import static com.datadog.iast.taint.TaintUtils.taintFormat
 
 class CodecModuleTest extends IastModuleImplTestBase {
 
@@ -306,4 +309,101 @@ class CodecModuleTest extends IastModuleImplTestBase {
       assert it.source.value == 'World!'
     }
   }
+
+  void 'test uri creation'() {
+    given:
+    final to = ctx.getTaintedObjects()
+    final params = args.collect {
+      return it instanceof String ? addFromTaintFormat(to, it as String) : it
+    }.toArray()
+    final uri = URI.metaClass.&invokeConstructor(params) as URI
+
+    when:
+    module.onUriCreate(uri, params)
+
+    then:
+    final rangeCount = fromTaintFormat(expected)?.length
+    assert uri.toString() == getStringFromTaintFormat(expected)
+    if (rangeCount > 0) {
+      final tainted = to.get(uri)
+      assert taintFormat(uri.toString(), tainted.ranges) == expected
+    } else {
+      assert to.get(uri) == null
+    }
+
+    where:
+    args                                                                                | expected
+    ['http://test.com/index?name=value#fragment']                                       | 'http://test.com/index?name=value#fragment'
+    ['==>http<==://test.com/index?name=value#fragment']                                 | '==>http<==://test.com/index?name=value#fragment'
+    ['http://test.com/index?==>name=value<==#fragment']                                 | 'http://test.com/index?==>name=value<==#fragment'
+    ['http', 'user:password', 'test.com', 80, '/index', 'name=value', 'fragment']       | 'http://user:[email protected]:80/index?name=value#fragment'
+    ['==>http<==', 'user:password', 'test.com', 80, '/index', 'name=value', 'fragment'] | '==>http<==://user:[email protected]:80/index?name=value#fragment'
+    ['http', 'user:password', 'test.com', 80, '/index', '==>name=value<==', 'fragment'] | 'http://user:[email protected]:80/index?==>name=value<==#fragment'
+  }
+
+  void 'test url creation'() {
+    given:
+    final to = ctx.getTaintedObjects()
+    final params = args.collect {
+      def result = it
+      if (it instanceof String) {
+        result = addFromTaintFormat(to, it as String)
+      } else if (it instanceof TaintedURL) {
+        final format = (it as TaintedURL).url
+        result = new URL(getStringFromTaintFormat(format))
+        final ranges = fromTaintFormat(format)
+        if (ranges?.length > 0) {
+          to.taint(result, ranges)
+        }
+      }
+      return result
+    }.toArray()
+    final url = URL.metaClass.&invokeConstructor(params) as URL
+
+    when:
+    module.onUrlCreate(url, params)
+
+    then:
+    final rangeCount = fromTaintFormat(expected)?.length
+    assert url.toString() == getStringFromTaintFormat(expected)
+    if (rangeCount > 0) {
+      final tainted = to.get(url)
+      assert taintFormat(url.toString(), tainted.ranges) == expected
+    } else {
+      assert to.get(url) == null
+    }
+
+    where:
+    args                                                                                            | expected
+    ['http://test.com/index?name=value#fragment']                                                   | 'http://test.com/index?name=value#fragment'
+    ['==>http<==://test.com/index?name=value#fragment']                                             | '==>http<==://test.com/index?name=value#fragment'
+    ['http://test.com/index?==>name=value<==#fragment']                                             | 'http://test.com/index?==>name=value<==#fragment'
+    ['http', 'test.com', 80, '/index?name=value#fragment']                                          | 'http://test.com:80/index?name=value#fragment'
+    ['==>http<==', 'test.com', 80, '/index?name=value#fragment']                                    | '==>http<==://test.com:80/index?name=value#fragment'
+    ['http', 'test.com', 80, '/index?==>name=value<==#fragment']                                    | 'http://test.com:80/index?==>name=value<==#fragment'
+    [new TaintedURL('http://test.com'), '/index?name=value#fragment']                               | 'http://test.com/index?name=value#fragment'
+    [new TaintedURL('==>http<==://test.com'), '/index?name=value#fragment']                         | '==>http<==://test.com/index?name=value#fragment'
+    [new TaintedURL('http://test.com'), '/index?==>name=value<==#fragment']                         | 'http://test.com/index?==>name=value<==#fragment'
+    [new TaintedURL('==>http<==://test.com'), '/index?==>name=value<==#fragment']                   | '==>http<==://test.com/index?==>name=value<==#fragment'
+    [null, 'http://test.com/index?==>name=value<==#fragment']                                       | 'http://test.com/index?==>name=value<==#fragment'
+    [new TaintedURL('==>http<==://ignored.com'), 'http://test.com/index?name=value#fragment']       | 'http://test.com/index?name=value#fragment'
+    [new TaintedURL('==>http<==://ignored.com'), '==>http<==://test.com/index?name=value#fragment'] | '==>http<==://test.com/index?name=value#fragment'
+    [new TaintedURL('==>http<==://ignored.com'), 'http://test.com/index?==>name=value<==#fragment'] | 'http://test.com/index?==>name=value<==#fragment'
+  }
+
+  protected static class TaintedURL {
+    private final String url
+
+    protected TaintedURL(String url) {
+      this.url = url
+    }
+
+    @Override
+    String toString() {
+      return "TaintedURL{" +
+        "url='" + url + '\'' +
+        '}'
+    }
+  }
+
 }

dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URICallSite.java

@@ -1,9 +1,12 @@
 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;
 import datadog.trace.api.iast.Propagation;
+import datadog.trace.api.iast.propagation.CodecModule;
 import datadog.trace.api.iast.propagation.PropagationModule;
 import java.net.URI;
 import javax.annotation.Nonnull;
@@ -17,10 +20,10 @@ public class URICallSite {
   public static URI afterCreate(
       @CallSite.Argument @Nullable final String value, @CallSite.Return @Nonnull final URI result) {
     if (value != null) {
-      final PropagationModule module = InstrumentationBridge.PROPAGATION;
+      final CodecModule module = InstrumentationBridge.CODEC;
       if (module != null) {
         try {
-          module.taintObjectIfTainted(result, value);
+          module.onUriCreate(result, value);
         } catch (final Throwable e) {
           module.onUnexpectedException("create threw", e);
         }
@@ -40,10 +43,10 @@ public static URI afterCreate(
   public static URI afterCtor(
       @CallSite.AllArguments final Object[] args, @CallSite.Return @Nonnull final URI result) {
     if (args != null && args.length > 0) {
-      final PropagationModule module = InstrumentationBridge.PROPAGATION;
+      final CodecModule module = InstrumentationBridge.CODEC;
       if (module != null) {
         try {
-          module.taintObjectIfAnyTainted(result, args);
+          module.onUriCreate(result, args);
         } catch (final Throwable e) {
           module.onUnexpectedException("ctor threw", e);
         }
@@ -52,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);
       }

dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URLCallSite.java

@@ -1,11 +1,14 @@
 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;
 import datadog.trace.api.iast.Propagation;
 import datadog.trace.api.iast.Sink;
 import datadog.trace.api.iast.VulnerabilityTypes;
+import datadog.trace.api.iast.propagation.CodecModule;
 import datadog.trace.api.iast.propagation.PropagationModule;
 import datadog.trace.api.iast.sink.SsrfModule;
 import java.net.Proxy;
@@ -29,10 +32,10 @@ public class URLCallSite {
   public static URL afterCtor(
       @CallSite.AllArguments final Object[] args, @CallSite.Return @Nonnull final URL result) {
     if (args != null && args.length > 0) {
-      final PropagationModule module = InstrumentationBridge.PROPAGATION;
+      final CodecModule module = InstrumentationBridge.CODEC;
       if (module != null) {
         try {
-          module.taintObjectIfAnyTainted(result, args);
+          module.onUrlCreate(result, args);
         } catch (final Throwable e) {
           module.onUnexpectedException("ctor threw", e);
         }
@@ -41,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);
       }