diff --git a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/CallSiteSpecification.java b/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/CallSiteSpecification.java index ac92aeec886..7ce110d33e8 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/CallSiteSpecification.java +++ b/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/CallSiteSpecification.java @@ -570,11 +570,20 @@ public AfterSpecification( @Override protected void validateAdvice(@Nonnull final ValidationContext context) { - if (advice.isVoidReturn()) { - context.addError(ErrorCode.ADVICE_AFTER_SHOULD_NOT_RETURN_VOID); - } - if (findReturn() == null) { - context.addError(ErrorCode.ADVICE_AFTER_SHOULD_HAVE_RETURN); + if (shouldNotUseReturn(pointcut)) { + if (!advice.isVoidReturn()) { + context.addError(ErrorCode.ADVICE_AFTER_VOID_METHOD_SHOULD_RETURN_VOID); + } + if (findReturn() != null) { + context.addError(ErrorCode.ADVICE_AFTER_VOID_METHOD_SHOULD_NOT_HAVE_RETURN); + } + } else { + if (advice.isVoidReturn()) { + context.addError(ErrorCode.ADVICE_AFTER_SHOULD_NOT_RETURN_VOID); + } + if (findReturn() == null) { + context.addError(ErrorCode.ADVICE_AFTER_SHOULD_HAVE_RETURN); + } } if (!pointcut.isConstructor()) { if (!isStaticPointcut() && !includeThis()) { @@ -588,6 +597,10 @@ protected void validateAdvice(@Nonnull final ValidationContext context) { super.validateAdvice(context); } + private boolean shouldNotUseReturn(final MethodType type) { + return !type.isConstructor() && type.isVoidReturn(); + } + @Override public String toString() { return "@CallSite.After(" + signature + ")"; diff --git a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/util/ErrorCode.java b/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/util/ErrorCode.java index e4fc54f7dbc..293125a76ee 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/util/ErrorCode.java +++ b/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/util/ErrorCode.java @@ -359,6 +359,20 @@ public String apply(final Object[] objects) { } }, + ADVICE_AFTER_VOID_METHOD_SHOULD_RETURN_VOID { + @Override + public String apply(final Object[] objects) { + return "After advice for void method should return void"; + } + }, + + ADVICE_AFTER_VOID_METHOD_SHOULD_NOT_HAVE_RETURN { + @Override + public String apply(final Object[] objects) { + return "After advice for void method should not contain @Return annotated parameters"; + } + }, + ADVICE_AFTER_SHOULD_HAVE_RETURN { @Override public String apply(final Object[] objects) { diff --git a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy b/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy index a0b3e15e2cc..3ec2a6da5c5 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy +++ b/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy @@ -473,6 +473,36 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest { } } + + @CallSite(spi = CallSites) + class AfterAdviceWithVoidReturn { + @CallSite.After("void java.lang.StringBuilder.setLength(int)") + static void after(@CallSite.This StringBuilder self, @CallSite.Argument(0) int length) { + } + } + + void 'test after advice with void return'() { + setup: + final spec = buildClassSpecification(AfterAdviceWithVoidReturn) + final generator = buildAdviceGenerator(buildDir) + + when: + final result = generator.generate(spec) + + then: + assertNoErrors result + assertCallSites(result.file) { + advices(0) { + pointcut('java/lang/StringBuilder', 'setLength', '(I)V') + statements( + 'handler.dupInvoke(owner, descriptor, StackDupMode.COPY);', + 'handler.method(opcode, owner, name, descriptor, isInterface);', + 'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AfterAdviceWithVoidReturn", "after", "(Ljava/lang/StringBuilder;I)V");', + ) + } + } + } + private static AdviceGenerator buildAdviceGenerator(final File targetFolder) { return new AdviceGeneratorImpl(targetFolder, pointcutParser()) } diff --git a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceSpecificationTest.groovy b/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceSpecificationTest.groovy index 12ef7348145..f3407dd786f 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceSpecificationTest.groovy +++ b/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceSpecificationTest.groovy @@ -542,4 +542,26 @@ class AdviceSpecificationTest extends BaseCsiPluginTest { then: 0 * context.addError(_, _) } + + + @CallSite(spi = CallSites) + class AfterWithVoidWrongAdvice { + @CallSite.After("void java.lang.String.getChars(int, int, char[], int)") + static String after(@CallSite.AllArguments final Object[] args, @CallSite.Return final String result) { + return result; + } + } + + void 'test after advice with void should not use @Return'() { + setup: + final context = mockValidationContext() + final spec = buildClassSpecification(AfterWithVoidWrongAdvice) + + when: + spec.advices.each { it.validate(context) } + + then: + 1 * context.addError(ErrorCode.ADVICE_AFTER_VOID_METHOD_SHOULD_RETURN_VOID, _) + 1 * context.addError(ErrorCode.ADVICE_AFTER_VOID_METHOD_SHOULD_NOT_HAVE_RETURN, _) + } } diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy index 2ee1e555758..416092536e9 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy @@ -83,6 +83,10 @@ class BaseCallSiteTest extends DDSpecification { return buildPointcut(String.getDeclaredMethod('concat', String)) } + protected static Pointcut stringBuilderSetLengthPointcut() { + return buildPointcut(StringBuilder.getDeclaredMethod('setLength', int)) + } + protected static Pointcut stringReaderPointcut() { return buildPointcut(StringReader.getDeclaredConstructor(String)) } diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteTransformerTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteTransformerTest.groovy index 69a107dd2c6..801319cc6db 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteTransformerTest.groovy +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteTransformerTest.groovy @@ -10,6 +10,7 @@ import net.bytebuddy.dynamic.DynamicType import net.bytebuddy.jar.asm.Opcodes import net.bytebuddy.jar.asm.Type +import java.util.function.BiConsumer import java.util.function.BiFunction import java.util.function.Consumer @@ -227,6 +228,38 @@ class CallSiteTransformerTest extends BaseCallSiteTest { helperCLass != null } + void 'test after advice with void return'() { + setup: + final source = Type.getType(StringBuilderSetLengthExample) + final target = renameType(source, 'Test') + final pointcut = stringBuilderSetLengthPointcut() + final advice = Mock(InvokeAdvice) + final callSite = Type.getType(StringBuilderSetLengthCallSite) + final callSiteTransformer = new CallSiteTransformer(mockAdvices([mockCallSites(advice, pointcut)])) + final sb = new StringBuilder("Hello World!") + final int length = 5 + StringBuilderSetLengthCallSite.LAST_CALL = null + + when: + final transformedClass = transformType(source, target, callSiteTransformer) + final instance = loadType(target, transformedClass) as BiConsumer + instance.accept(sb, length) + + then: + 1 * advice.apply(_ as MethodHandler, Opcodes.INVOKEVIRTUAL, pointcut.type, pointcut.method, pointcut.descriptor, false) >> { params -> + final args = params as Object[] + final handler = args[0] as MethodHandler + final owner = args[2] as String + final descriptor = args[4] as String + handler.dupInvoke(owner, descriptor, COPY) + handler.method(args[1] as int, owner, args[3] as String, descriptor, args[5] as Boolean) + handler.advice(callSite.getInternalName(), "after", "(Ljava/lang/StringBuilder;I)V") + } + sb.toString() == 'Hello' + StringBuilderSetLengthCallSite.LAST_CALL[0] == sb + StringBuilderSetLengthCallSite.LAST_CALL[1] == length + } + static class InstrumentationHelper { private static Consumer callback = null // codenarc forces the lowercase name diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/csi/StringBuilderSetLengthCallSite.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/csi/StringBuilderSetLengthCallSite.java new file mode 100644 index 00000000000..fb10143fbea --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/csi/StringBuilderSetLengthCallSite.java @@ -0,0 +1,10 @@ +package datadog.trace.agent.tooling.csi; + +public class StringBuilderSetLengthCallSite { + + public static volatile Object[] LAST_CALL; + + public static void after(final StringBuilder builder, int length) { + LAST_CALL = new Object[] {builder, length}; + } +} diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/csi/StringBuilderSetLengthExample.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/csi/StringBuilderSetLengthExample.java new file mode 100644 index 00000000000..90bdc0dc606 --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/csi/StringBuilderSetLengthExample.java @@ -0,0 +1,16 @@ +package datadog.trace.agent.tooling.csi; + +import java.util.function.BiConsumer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class StringBuilderSetLengthExample implements BiConsumer { + + private static final Logger LOGGER = LoggerFactory.getLogger(StringBuilderSetLengthExample.class); + + public void accept(final StringBuilder builder, final Integer length) { + LOGGER.debug("Before apply {} {}", builder, length); + builder.setLength(length); + LOGGER.debug("After apply {} {}", builder, length); + } +}