Skip to content

Commit 594a2a4

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
authored
Exploit prevention for SSRF (in java.net.URL) (#7373)
Adds support for SSRF protection via RASP to network connections started via java.net.URL: Extend support of IAST call sites to RASP (usually both IAST and RASP share the same targets as IAST detects the vulnerability and RASP performs the protection) Add SSRF support to the WAF via the server.io.net.url address
1 parent 54884d1 commit 594a2a4

32 files changed

Lines changed: 674 additions & 151 deletions

File tree

‎buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
import datadog.trace.plugin.csi.util.ErrorCode;
5656
import datadog.trace.plugin.csi.util.MethodType;
5757
import java.io.File;
58+
import java.util.ArrayList;
5859
import java.util.Arrays;
5960
import java.util.List;
6061
import java.util.stream.Collectors;
@@ -156,9 +157,11 @@ private static ClassOrInterfaceDeclaration callSitesType(
156157
type.setModifier(PUBLIC, true);
157158
type.setName(getClassName(advice));
158159
type.addImplementedType(CALL_SITES_CLASS);
159-
if (!CALL_SITES_FQCN.equals(callSite.getSpi().getClassName())) {
160-
javaClass.addImport(callSite.getSpi().getClassName());
161-
type.addImplementedType(getClassName(callSite.getSpi(), false));
160+
for (final Type spi : callSite.getSpi()) {
161+
if (!CALL_SITES_FQCN.equals(spi.getClassName())) {
162+
javaClass.addImport(spi.getClassName());
163+
type.addImplementedType(getClassName(spi, false));
164+
}
162165
}
163166
javaClass.addType(type);
164167
return type;
@@ -168,9 +171,12 @@ private static void addAutoServiceAnnotation(
168171
final ClassOrInterfaceDeclaration javaClass, final CallSiteSpecification callSite) {
169172
final NormalAnnotationExpr autoService = new NormalAnnotationExpr();
170173
autoService.setName(AUTO_SERVICE_FQDN);
171-
autoService.addPair(
172-
"value",
173-
new ClassExpr(new ClassOrInterfaceType().setName(getClassName(callSite.getSpi(), false))));
174+
final Type[] spiTypes = callSite.getSpi();
175+
final List<Expression> spiExprs = new ArrayList<>(spiTypes.length);
176+
for (final Type spi : spiTypes) {
177+
spiExprs.add(new ClassExpr(new ClassOrInterfaceType().setName(getClassName(spi, false))));
178+
}
179+
autoService.addPair("value", new ArrayInitializerExpr().setValues(new NodeList<>(spiExprs)));
174180
javaClass.addAnnotation(autoService);
175181
}
176182

‎buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AsmSpecificationBuilder.java‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ private static class SpecificationVisitor extends ClassVisitor {
7676
private boolean isCallSite;
7777
private final List<AdviceSpecification> advices = new ArrayList<>();
7878
private final Set<Type> helpers = new HashSet<>();
79-
private Type spi;
79+
private final Set<Type> spi = new HashSet<>();
8080
private List<String> enabled = new ArrayList<>();
8181
private CallSiteSpecification result;
8282

@@ -101,16 +101,17 @@ public AnnotationVisitor visitAnnotation(final String descriptor, final boolean
101101
if (isCallSite) {
102102
helpers.add(clazz);
103103
return new AnnotationVisitor(ASM_API_VERSION) {
104-
@Override
105-
public void visit(final String key, final Object value) {
106-
if ("spi".equals(key)) {
107-
spi = (Type) value;
108-
}
109-
}
110104

111105
@Override
112106
public AnnotationVisitor visitArray(final String name) {
113-
if ("helpers".equals(name)) {
107+
if ("spi".equals(name)) {
108+
return new AnnotationVisitor(ASM_API_VERSION) {
109+
@Override
110+
public void visit(final String name, final Object value) {
111+
spi.add((Type) value);
112+
}
113+
};
114+
} else if ("helpers".equals(name)) {
114115
return new AnnotationVisitor(ASM_API_VERSION) {
115116
@Override
116117
public void visit(final String name, final Object value) {

‎buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/CallSiteSpecification.java‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -32,19 +32,19 @@ public class CallSiteSpecification implements Validatable {
3232

3333
private final Type clazz;
3434
private final List<AdviceSpecification> advices;
35-
private final Type spi;
35+
private final Type[] spi;
3636
private final Enabled enabled;
3737
private final Type[] helpers;
3838

3939
public CallSiteSpecification(
4040
@Nonnull final Type clazz,
4141
@Nonnull final List<AdviceSpecification> advices,
42-
@Nonnull final Type spi,
42+
@Nonnull final Set<Type> spi,
4343
@Nonnull final List<String> enabled,
4444
@Nonnull final Set<Type> helpers) {
4545
this.clazz = clazz;
4646
this.advices = advices;
47-
this.spi = spi;
47+
this.spi = spi.toArray(new Type[0]);
4848
this.enabled = enabled.isEmpty() ? null : new Enabled(enabled);
4949
this.helpers = helpers.toArray(new Type[0]);
5050
}
@@ -53,12 +53,14 @@ public CallSiteSpecification(
5353
public void validate(@Nonnull final ValidationContext context) {
5454
final TypeResolver typeResolver = context.getContextProperty(TYPE_RESOLVER);
5555
try {
56-
Class<?> spiClass = typeResolver.resolveType(spi);
57-
if (!spiClass.isInterface()) {
58-
context.addError(ErrorCode.CALL_SITE_SPI_SHOULD_BE_AN_INTERFACE, spiClass);
59-
} else {
60-
if (spiClass.getDeclaredMethods().length > 0) {
61-
context.addError(ErrorCode.CALL_SITE_SPI_SHOULD_BE_EMPTY, spiClass);
56+
for (Type spiType : spi) {
57+
Class<?> spiClass = typeResolver.resolveType(spiType);
58+
if (!spiClass.isInterface()) {
59+
context.addError(ErrorCode.CALL_SITE_SPI_SHOULD_BE_AN_INTERFACE, spiClass);
60+
} else {
61+
if (spiClass.getDeclaredMethods().length > 0) {
62+
context.addError(ErrorCode.CALL_SITE_SPI_SHOULD_BE_EMPTY, spiClass);
63+
}
6264
}
6365
}
6466
} catch (ResolutionException e) {
@@ -84,7 +86,7 @@ public Type getClazz() {
8486
return clazz;
8587
}
8688

87-
public Type getSpi() {
89+
public Type[] getSpi() {
8890
return spi;
8991
}
9092

‎buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/ext/IastExtension.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
import java.util.stream.Collectors;
5454
import java.util.stream.Stream;
5555
import javax.annotation.Nonnull;
56+
import org.objectweb.asm.Type;
5657

5758
@SuppressWarnings("OptionalGetWithoutIsPresent")
5859
public class IastExtension implements Extension {
@@ -79,7 +80,12 @@ public class IastExtension implements Extension {
7980

8081
@Override
8182
public boolean appliesTo(@Nonnull final CallSiteSpecification spec) {
82-
return IAST_CALL_SITES_FQCN.equals(spec.getSpi().getClassName());
83+
for (Type spi : spec.getSpi()) {
84+
if (IAST_CALL_SITES_FQCN.equals(spi.getClassName())) {
85+
return true;
86+
}
87+
}
88+
return false;
8389
}
8490

8591
@Override

‎buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ import datadog.trace.agent.tooling.csi.CallSites
66
import datadog.trace.plugin.csi.AdviceGenerator
77
import datadog.trace.plugin.csi.impl.assertion.AssertBuilder
88
import datadog.trace.plugin.csi.impl.assertion.CallSiteAssert
9+
import datadog.trace.plugin.csi.impl.ext.tests.IastCallSites
10+
import datadog.trace.plugin.csi.impl.ext.tests.RaspCallSites
911
import groovy.transform.CompileDynamic
1012
import spock.lang.Requires
1113
import spock.lang.TempDir
@@ -448,6 +450,29 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
448450
}
449451
}
450452

453+
@CallSite(spi = [IastCallSites, RaspCallSites])
454+
class MultipleSpiClassesAdvice {
455+
@CallSite.After("void java.lang.StringBuilder.<init>(java.lang.String)")
456+
static Object after(@CallSite.AllArguments Object[] args, @CallSite.Return Object result) {
457+
return result
458+
}
459+
}
460+
461+
void 'test multiple spi classes'() {
462+
setup:
463+
final spec = buildClassSpecification(MultipleSpiClassesAdvice)
464+
final generator = buildAdviceGenerator(buildDir)
465+
466+
when:
467+
final result = generator.generate(spec)
468+
469+
then:
470+
assertNoErrors result
471+
assertCallSites(result.file) {
472+
spi(IastCallSites, RaspCallSites)
473+
}
474+
}
475+
451476
private static AdviceGenerator buildAdviceGenerator(final File targetFolder) {
452477
return new AdviceGeneratorImpl(targetFolder, pointcutParser())
453478
}

‎buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AsmSpecificationBuilderTest.groovy‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ final class AsmSpecificationBuilderTest extends BaseCsiPluginTest {
4949
final result = specificationBuilder.build(advice).orElseThrow(RuntimeException::new)
5050

5151
then:
52-
result.spi == Type.getType(WithSpiClass.Spi)
52+
result.spi == [Type.getType(WithSpiClass.Spi)] as Type[]
5353
}
5454

5555
@CallSite(spi = CallSites, helpers = [SampleHelper1.class, SampleHelper2.class])

‎buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/CallSiteSpecificationTest.groovy‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ class CallSiteSpecificationTest extends BaseCsiPluginTest {
1010
def 'test call site spi should be an interface'() {
1111
setup:
1212
final context = mockValidationContext()
13-
final spec = new CallSiteSpecification(Type.getType(String), [Mock(AdviceSpecification)], Type.getType(String), [] as List<String>, [] as Set<Type>)
13+
final spec = new CallSiteSpecification(Type.getType(String), [Mock(AdviceSpecification)], [Type.getType(String)] as Set<Type>, [] as List<String>, [] as Set<Type>)
1414

1515
when:
1616
spec.validate(context)
@@ -22,7 +22,7 @@ class CallSiteSpecificationTest extends BaseCsiPluginTest {
2222
def 'test call site spi should not define any methods'() {
2323
setup:
2424
final context = mockValidationContext()
25-
final spec = new CallSiteSpecification(Type.getType(String), [Mock(AdviceSpecification)], Type.getType(Comparable), [] as List<String>, [] as Set<Type>)
25+
final spec = new CallSiteSpecification(Type.getType(String), [Mock(AdviceSpecification)], [Type.getType(Comparable)] as Set<Type>, [] as List<String>, [] as Set<Type>)
2626

2727
when:
2828
spec.validate(context)
@@ -34,7 +34,7 @@ class CallSiteSpecificationTest extends BaseCsiPluginTest {
3434
def 'test call site should have advices'() {
3535
setup:
3636
final context = mockValidationContext()
37-
final spec = new CallSiteSpecification(Type.getType(String), [], Type.getType(CallSiteAdvice), [] as List<String>, [] as Set<Type>)
37+
final spec = new CallSiteSpecification(Type.getType(String), [], [Type.getType(CallSiteAdvice)] as Set<Type>, [] as List<String>, [] as Set<Type>)
3838

3939
when:
4040
spec.validate(context)

‎buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/assertion/AssertBuilder.groovy‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -31,33 +31,42 @@ class AssertBuilder<C extends CallSiteAssert> {
3131
def (enabled, enabledArgs) = getEnabledDeclaration(targetType, interfaces)
3232
return (C) new CallSiteAssert([
3333
interfaces : getInterfaces(targetType),
34+
spi : getSpi(targetType),
3435
helpers : getHelpers(targetType),
3536
advices : getAdvices(targetType),
3637
enabled : enabled,
3738
enabledArgs: enabledArgs
3839
])
3940
}
4041

41-
protected List<Class<?>> getInterfaces(final ClassOrInterfaceDeclaration type) {
42+
protected Set<Class<?>> getSpi(final ClassOrInterfaceDeclaration type) {
43+
return type.getAnnotationByName('AutoService').get().asNormalAnnotationExpr()
44+
.collect { it.pairs.find { it.name.toString() == 'value' }.value.asArrayInitializerExpr() }
45+
.collectMany { it.getValues() }
46+
.collect { it.asClassExpr().getType().resolve().typeDeclaration.get().clazz }
47+
.toSet()
48+
}
49+
50+
protected Set<Class<?>> getInterfaces(final ClassOrInterfaceDeclaration type) {
4251
return type.asClassOrInterfaceDeclaration().implementedTypes.collect {
4352
final resolved = it.asClassOrInterfaceType().resolve()
4453
return resolved.typeDeclaration.get().clazz
45-
}
54+
}.toSet()
4655
}
4756

48-
protected def getEnabledDeclaration(final ClassOrInterfaceDeclaration type, final List<Class<?>> interfaces) {
57+
protected def getEnabledDeclaration(final ClassOrInterfaceDeclaration type, final Set<Class<?>> interfaces) {
4958
if (!interfaces.contains(CallSites.HasEnabledProperty)) {
5059
return [null, null]
5160
}
5261
final isEnabled = type.getMethodsByName('isEnabled').first()
5362
final returnStatement = isEnabled.body.get().statements.first.get().asReturnStmt()
5463
final enabledMethodCall = returnStatement.expression.get().asMethodCallExpr()
5564
final enabled = resolveMethod(enabledMethodCall)
56-
final enabledArgs = enabledMethodCall.getArguments().collect { it.asStringLiteralExpr().asString() }
65+
final enabledArgs = enabledMethodCall.getArguments().collect { it.asStringLiteralExpr().asString() }.toSet()
5766
return [enabled, enabledArgs]
5867
}
5968

60-
protected List<Class<?>> getHelpers(final ClassOrInterfaceDeclaration type) {
69+
protected Set<Class<?>> getHelpers(final ClassOrInterfaceDeclaration type) {
6170
final acceptMethod = type.getMethodsByName('accept').first()
6271
final methodCalls = getMethodCalls(acceptMethod)
6372
return methodCalls.findAll {
@@ -66,15 +75,15 @@ class AssertBuilder<C extends CallSiteAssert> {
6675
it.arguments
6776
}.collect {
6877
typeResolver().resolveType(classNameToType(it.asStringLiteralExpr().asString()))
69-
}
78+
}.toSet()
7079
}
7180

7281
protected List<AdviceAssert> getAdvices(final ClassOrInterfaceDeclaration type) {
7382
final acceptMethod = type.getMethodsByName('accept').first()
7483
return getMethodCalls(acceptMethod).findAll {
7584
it.nameAsString == 'addAdvice'
7685
}.collect {
77-
def (owner, method, descriptor) = it.arguments.subList(0, 3)*.asStringLiteralExpr()*.asString()
86+
def (owner, method, descriptor) = it.arguments.subList(0, 3)*.asStringLiteralExpr()*.asString()
7887
final handlerLambda = it.arguments[3].asLambdaExpr()
7988
final advice = handlerLambda.body.asBlockStmt().statements*.toString()
8089
return new AdviceAssert([

‎buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/assertion/CallSiteAssert.groovy‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,20 +2,27 @@ package datadog.trace.plugin.csi.impl.assertion
22

33
import java.lang.reflect.Method
44

5+
import static java.util.Arrays.asList
6+
57
class CallSiteAssert {
68

7-
protected Collection<Class<?>> interfaces
8-
protected Collection<Class<?>> helpers
9+
protected Set<Class<?>> interfaces
10+
protected Set<Class<?>> spi
11+
protected Set<Class<?>> helpers
912
protected Collection<AdviceAssert> advices
1013
protected Method enabled
11-
protected Collection<String> enabledArgs
14+
protected Set<String> enabledArgs
1215

1316
void interfaces(Class<?>... values) {
14-
assert values.toList() == interfaces
17+
assertSameElements(interfaces, values)
1518
}
1619

1720
void helpers(Class<?>... values) {
18-
assert values.toList() == helpers
21+
assertSameElements(helpers, values)
22+
}
23+
24+
void spi(Class<?>...values) {
25+
assertSameElements(spi, values)
1926
}
2027

2128
void advices(int index, @DelegatesTo(AdviceAssert) Closure closure) {
@@ -26,6 +33,10 @@ class CallSiteAssert {
2633

2734
void enabled(Method method, String... args) {
2835
assert method == enabled
29-
assert args.toList() == enabledArgs
36+
assertSameElements(enabledArgs, args)
37+
}
38+
39+
private static <E> void assertSameElements(final Set<E> expected, final E...received) {
40+
assert received.length == expected.size() && expected.containsAll(asList(received))
3041
}
3142
}

‎buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/ext/IastExtensionTest.groovy‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,7 @@ class IastExtensionTest extends BaseCsiPluginTest {
151151

152152
IastExtensionCallSiteAssert(CallSiteAssert base) {
153153
interfaces = base.interfaces
154+
spi = base.spi
154155
helpers = base.helpers
155156
advices = base.advices
156157
enabled = base.enabled

0 commit comments

Comments
 (0)