Skip to content

Commit c830d6c

Browse files
MariovidoMariovido
authored
Merge branch 'master' into mario.vidal/untrusted_deserialization_improvements
2 parents 5c1f37a + 594a2a4 commit c830d6c

39 files changed

Lines changed: 694 additions & 173 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)