Skip to content

Commit

Permalink
Update URI and URL call sites for precise taint tracking (#7299)
Browse files Browse the repository at this point in the history
Improve propagation on URL and URI call sites
  • Loading branch information
manuel-alvarez-alvarez committed Jul 16, 2024
1 parent b417127 commit c31955f
Show file tree
Hide file tree
Showing 13 changed files with 360 additions and 90 deletions.
Original file line number Diff line number Diff line change
@@ -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;

Expand All @@ -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,
Expand All @@ -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());
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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 {

Expand Down Expand Up @@ -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:password@test.com:80/index?name=value#fragment'
['==>http<==', 'user:password', 'test.com', 80, '/index', 'name=value', 'fragment'] | '==>http<==://user:password@test.com:80/index?name=value#fragment'
['http', 'user:password', 'test.com', 80, '/index', '==>name=value<==', 'fragment'] | 'http://user:password@test.com: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 + '\'' +
'}'
}
}

}
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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);
}
Expand All @@ -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);
}
Expand All @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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);
}
Expand All @@ -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);
}
Expand Down
Loading

0 comments on commit c31955f

Please sign in to comment.