Skip to content

Commit 8e21b90

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
authored
Ensure StringBuilder.append(String, Object) is properly instrumented (DataDog#5274)
1 parent 2f89f00 commit 8e21b90

10 files changed

Lines changed: 119 additions & 24 deletions

File tree

‎buildSrc/call-site-instrumentation-plugin/build.gradle.kts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ dependencies {
4343
testImplementation("org.objenesis", "objenesis", "3.0.1")
4444
testImplementation("org.codehaus.groovy", "groovy-all", "3.0.15")
4545
testImplementation("javax.servlet", "javax.servlet-api", "3.0.1")
46+
testImplementation("com.github.spotbugs", "spotbugs-annotations", "4.2.0")
4647
}
4748

4849
sourceSets {

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

Lines changed: 23 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -182,9 +182,8 @@ private static class AdviceMethodVisitor extends MethodVisitor {
182182
private final SpecificationVisitor spec;
183183
private final MethodType advice;
184184
private final Map<Integer, ParameterSpecification> parameters = new HashMap<>();
185-
private final List<String> signatures = new ArrayList<>();
186-
private boolean inokeDynamic;
187-
private AdviceSpecificationCtor adviceCtor;
185+
private final Map<AdviceSpecificationCtor, List<AdviceSpecificationData>> adviceData =
186+
new HashMap<>();
188187

189188
public AdviceMethodVisitor(
190189
@Nonnull final SpecificationVisitor spec,
@@ -197,15 +196,19 @@ public AdviceMethodVisitor(
197196

198197
@Override
199198
public AnnotationVisitor visitAnnotation(final String descriptor, final boolean visible) {
200-
adviceCtor = ADVICE_BUILDERS.get(descriptor);
201-
if (adviceCtor != null) {
199+
AdviceSpecificationCtor ctor = ADVICE_BUILDERS.get(descriptor);
200+
if (ctor != null) {
201+
final List<AdviceSpecificationData> list =
202+
adviceData.computeIfAbsent(ctor, c -> new ArrayList<>());
203+
final AdviceSpecificationData data = new AdviceSpecificationData();
204+
list.add(data);
202205
return new AnnotationVisitor(ASM_API_VERSION) {
203206
@Override
204207
public void visit(final String key, final Object value) {
205208
if ("value".equals(key)) {
206-
signatures.add((String) value);
209+
data.signature = (String) value;
207210
} else if ("invokeDynamic".equals(key)) {
208-
inokeDynamic = (boolean) value;
211+
data.invokeDynamic = (boolean) value;
209212
}
210213
}
211214
};
@@ -233,7 +236,7 @@ public AnnotationVisitor visitAnnotation(
233236
@Override
234237
public AnnotationVisitor visitParameterAnnotation(
235238
final int parameter, final String descriptor, final boolean visible) {
236-
if (adviceCtor != null) {
239+
if (!adviceData.isEmpty()) {
237240
final ParameterSpecificationCtor parameterCtor = PARAMETER_BUILDERS.get(descriptor);
238241
if (parameterCtor != null) {
239242
ParameterSpecification parameterSpec = parameterCtor.build();
@@ -265,11 +268,13 @@ public void visit(final String key, final Object value) {
265268

266269
@Override
267270
public void visitEnd() {
268-
if (adviceCtor != null) {
269-
signatures.stream()
270-
.map(sig -> adviceCtor.build(advice, parameters, sig, inokeDynamic))
271-
.forEach(spec.advices::add);
272-
}
271+
adviceData.forEach(
272+
(adviceCtor, list) ->
273+
list.stream()
274+
.map(
275+
data ->
276+
adviceCtor.build(advice, parameters, data.signature, data.invokeDynamic))
277+
.forEach(spec.advices::add));
273278
}
274279
}
275280

@@ -286,4 +291,9 @@ AdviceSpecification build(
286291
private interface ParameterSpecificationCtor {
287292
ParameterSpecification build();
288293
}
294+
295+
private static class AdviceSpecificationData {
296+
private String signature;
297+
private boolean invokeDynamic;
298+
}
289299
}

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,12 @@ import datadog.trace.plugin.csi.impl.CallSiteSpecification.AdviceSpecification
55
import datadog.trace.plugin.csi.impl.CallSiteSpecification.AfterSpecification
66
import datadog.trace.plugin.csi.impl.CallSiteSpecification.AroundSpecification
77
import datadog.trace.plugin.csi.impl.CallSiteSpecification.BeforeSpecification
8+
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
89
import groovy.transform.CompileDynamic
910
import org.objectweb.asm.Type
1011

12+
import javax.annotation.Nonnull
13+
import javax.annotation.Nullable
1114
import javax.servlet.ServletRequest
1215
import java.lang.invoke.MethodHandles
1316
import java.lang.invoke.MethodType
@@ -465,6 +468,32 @@ final class AsmSpecificationBuilderTest extends BaseCsiPluginTest {
465468
result.minJavaVersion == 9
466469
}
467470

471+
@CallSite
472+
static class TestWithOtherAnnotations {
473+
@CallSite.Around("java.lang.StringBuilder java.lang.StringBuilder.append(java.lang.Object)")
474+
@CallSite.Around("java.lang.StringBuffer java.lang.StringBuffer.append(java.lang.Object)")
475+
@Nonnull
476+
@SuppressFBWarnings(
477+
"NP_PARAMETER_MUST_BE_NONNULL_BUT_MARKED_AS_NULLABLE") // we do check for null on self
478+
// parameter
479+
static Appendable aroundAppend(@CallSite.This @Nullable final Appendable self, @CallSite.Argument(0) @Nullable final Object param) throws Throwable {
480+
return self.append(param.toString())
481+
}
482+
}
483+
484+
void 'test specification builder with multiple method annotations'() {
485+
setup:
486+
final advice = fetchClass(TestWithOtherAnnotations)
487+
final specificationBuilder = new AsmSpecificationBuilder()
488+
489+
when:
490+
final result = specificationBuilder.build(advice).orElseThrow(RuntimeException::new)
491+
492+
then:
493+
result.clazz.className == TestWithOtherAnnotations.name
494+
result.advices.size() == 2
495+
}
496+
468497
private static List<Integer> getArguments(final AdviceSpecification advice) {
469498
return advice.arguments.map(it -> it.index).collect(Collectors.toList())
470499
}

‎dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,9 @@ public static Appendable aroundAppend(
7979
}
8080
return result;
8181
} catch (final Throwable e) {
82-
throw StackUtils.removeLast(e);
82+
final String clazz = StringBuilderCallSite.class.getName();
83+
throw StackUtils.filterUntil(
84+
e, s -> s.getClassName().equals(clazz) && s.getMethodName().equals("aroundAppend"));
8385
}
8486
}
8587

‎dd-java-agent/instrumentation/java-lang/src/test/groovy/datadog/trace/instrumentation/java/lang/StringBuilderCallSiteTest.groovy‎

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@ class StringBuilderCallSiteTest extends AgentTestRunner {
4646
when:
4747
if (param.class == String) {
4848
suite.append(target, (String) param)
49+
} else if (param.class == CharSequence) {
50+
suite.append(target, (CharSequence) param)
4951
} else {
5052
suite.append(target, param)
5153
}
@@ -55,19 +57,39 @@ class StringBuilderCallSiteTest extends AgentTestRunner {
5557
if (param.class == String) {
5658
1 * iastModule.onStringBuilderAppend(target, (String) param)
5759
} else {
58-
1 * iastModule.onStringBuilderAppend(target, param)
60+
1 * iastModule.onStringBuilderAppend(target, param.toString())
5961
}
6062
_ * TEST_CHECKPOINTER._
6163
0 * _
6264

6365
where:
6466
suite | target | param | expected
67+
new TestStringBuilderSuite() | new StringBuilder('Hello ') | 23.5F | 'Hello 23.5'
6568
new TestStringBuilderSuite() | new StringBuilder('Hello ') | new StringBuffer('World!') | 'Hello World!'
6669
new TestStringBuilderSuite() | new StringBuilder('Hello ') | 'World!' | 'Hello World!'
70+
new TestStringBufferSuite() | new StringBuffer('Hello ') | 23.5F | 'Hello 23.5'
6771
new TestStringBufferSuite() | new StringBuffer('Hello ') | new StringBuilder('World!') | 'Hello World!'
6872
new TestStringBufferSuite() | new StringBuffer('Hello ') | 'World!' | 'Hello World!'
6973
}
7074

75+
void 'test string builder append object throwing exceptions'() {
76+
setup:
77+
final iastModule = Mock(StringModule)
78+
InstrumentationBridge.registerIastModule(iastModule)
79+
80+
when:
81+
suite.append(target, new BrokenToString())
82+
83+
then:
84+
final ex = thrown(NuclearException)
85+
ex.stackTrace.find { it.className == StringBuilderCallSite.name } == null
86+
87+
where:
88+
suite | target
89+
new TestStringBuilderSuite() | new StringBuilder('Hello ')
90+
new TestStringBufferSuite() | new StringBuffer('Hello ')
91+
}
92+
7193
void 'test string builder toString call site'() {
7294
setup:
7395
final iastModule = Mock(StringModule)
@@ -118,16 +140,20 @@ class StringBuilderCallSiteTest extends AgentTestRunner {
118140
result == expected
119141

120142
1 * iastModule.onStringBuilderAppend(_ as StringBuilder, '')
143+
1 * iastModule.onStringBuilderAppend(_ as StringBuilder, args[0])
121144
1 * iastModule.onStringBuilderToString(_ as StringBuilder, args[0])
122145

123146
1 * iastModule.onStringBuilderAppend(_ as StringBuilder, args[0])
147+
1 * iastModule.onStringBuilderAppend(_ as StringBuilder, args[1].toString())
124148
1 * iastModule.onStringBuilderToString(_ as StringBuilder, args[0..1].join())
125149

126150
1 * iastModule.onStringBuilderAppend(_ as StringBuilder, args[0..1].join())
151+
1 * iastModule.onStringBuilderAppend(_ as StringBuilder, args[2])
127152
1 * iastModule.onStringBuilderToString(_ as StringBuilder, args[0..2].join())
128153

129154
1 * iastModule.onStringBuilderAppend(_ as StringBuilder, args[0..2].join())
130-
1 * iastModule.onStringBuilderToString(_ as StringBuilder, expected)
155+
1 * iastModule.onStringBuilderAppend(_ as StringBuilder, args[3].toString())
156+
1 * iastModule.onStringBuilderToString(_ as StringBuilder, args[0..3].join())
131157

132158
0 * _
133159
}

‎dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestAbstractStringBuilderSuite.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,5 +10,7 @@ public interface TestAbstractStringBuilderSuite<E> {
1010

1111
void append(final E target, final CharSequence param);
1212

13+
void append(final E target, final Object param);
14+
1315
String toString(final E target);
1416
}

‎dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringBufferSuite.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,13 @@ public void append(final StringBuffer buffer, final CharSequence param) {
3737
LOGGER.debug("After string buffer append {}", result);
3838
}
3939

40+
@Override
41+
public void append(final StringBuffer buffer, final Object param) {
42+
LOGGER.debug("Before string buffer append {}", param);
43+
final StringBuffer result = buffer.append(param);
44+
LOGGER.debug("After string buffer append {}", result);
45+
}
46+
4047
@Override
4148
public String toString(final StringBuffer buffer) {
4249
LOGGER.debug("Before string buffer toString {}", buffer);

‎dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringBuilderSuite.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,13 @@ public void append(final StringBuilder builder, final CharSequence param) {
3838
LOGGER.debug("After string builder append {}", result);
3939
}
4040

41+
@Override
42+
public void append(final StringBuilder builder, final Object param) {
43+
LOGGER.debug("Before string builder append {}", param);
44+
final StringBuilder result = builder.append(param);
45+
LOGGER.debug("After string builder append {}", result);
46+
}
47+
4148
@Override
4249
public String toString(final StringBuilder builder) {
4350
LOGGER.debug("Before string builder toString {}", builder);

‎internal-api/src/main/java/datadog/trace/util/stacktrace/StackUtils.java‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,14 +32,20 @@ public static <E extends Throwable> E filterFirstDatadog(final E exception) {
3232
return filterFirst(exception, AbstractStackWalker::isNotDatadogTraceStackElement);
3333
}
3434

35-
public static <E extends Throwable> E removeLast(final E exception) {
35+
public static <E extends Throwable> E filterUntil(
36+
final E exception, final Predicate<StackTraceElement> trace) {
3637
return update(
3738
exception,
3839
stack -> {
3940
final StackTraceElement[] source = exception.getStackTrace();
40-
final StackTraceElement[] result = new StackTraceElement[source.length - 1];
41-
System.arraycopy(source, 0, result, 0, result.length);
42-
return result;
41+
for (int i = 0; i < source.length; i++) {
42+
if (trace.test(source[i])) {
43+
final StackTraceElement[] result = new StackTraceElement[source.length - i - 1];
44+
System.arraycopy(source, i + 1, result, 0, result.length);
45+
return result;
46+
}
47+
}
48+
return source;
4349
});
4450
}
4551

‎internal-api/src/test/java/datadog/trace/util/stacktrace/StackUtilsTest.java‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ public void test_filter_first_datadog() {
6060
}
6161

6262
@Test
63-
public void test_remove_last() {
63+
public void test_filter_until() {
6464
final StackTraceElement[] stack =
6565
new StackTraceElement[] {
6666
stack().className("org.junit.jupiter.api.Test").build(),
@@ -69,11 +69,16 @@ public void test_remove_last() {
6969
stack().className("datadog.trace.util.stacktrace.StackUtils").build(),
7070
stack().className("com.google.common.truth.Truth").build()
7171
};
72-
final StackTraceElement[] expected =
73-
new StackTraceElement[] {stack[0], stack[1], stack[2], stack[3]};
7472

75-
final Throwable removed = StackUtils.removeLast(withStack(stack));
73+
final StackTraceElement[] expected = new StackTraceElement[] {stack[4]};
74+
final Throwable removed =
75+
StackUtils.filterUntil(
76+
withStack(stack),
77+
entry -> entry.getClassName().equals("datadog.trace.util.stacktrace.StackUtils"));
7678
assertThat(removed.getStackTrace()).isEqualTo(expected);
79+
80+
final Throwable noRemoval = StackUtils.filterUntil(withStack(stack), entry -> false);
81+
assertThat(noRemoval.getStackTrace()).isEqualTo(stack);
7782
}
7883

7984
private static Throwable withStack(final StackTraceElement... stack) {

0 commit comments

Comments
 (0)