Prevent before callsites targeting calls to super in constructors

Author: manuel-alvarez-alvarez

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

@@ -11,6 +11,7 @@
 import static datadog.trace.plugin.csi.util.CallSiteConstants.METHOD_HANDLER_CLASS;
 import static datadog.trace.plugin.csi.util.CallSiteConstants.OPCODES_FQDN;
 import static datadog.trace.plugin.csi.util.CallSiteConstants.STACK_DUP_MODE_CLASS;
+import static datadog.trace.plugin.csi.util.CallSiteConstants.TYPE_CLASS;
 import static datadog.trace.plugin.csi.util.CallSiteConstants.TYPE_RESOLVER;
 import static datadog.trace.plugin.csi.util.CallSiteUtils.deleteFile;
 import static datadog.trace.plugin.csi.util.JavaParserUtils.getPrimaryType;
@@ -185,20 +186,24 @@ private void addAdviceLambda(
     final MethodType pointCut = spec.getPointcut();
     final BlockStmt adviceBody = new BlockStmt();
     final Expression advice;
+    final String type;
     if (spec.isInvokeDynamic()) {
       advice = invokeDynamicAdviceSignature(adviceBody);
     } else {
       advice = invokeAdviceSignature(adviceBody);
     }
     if (spec instanceof BeforeSpecification) {
+      type = "BEFORE";
       writeStackOperations(spec, adviceBody);
       writeAdviceMethodCall(spec, adviceBody);
       writeOriginalMethodCall(spec, adviceBody);
     } else if (spec instanceof AfterSpecification) {
+      type = "AFTER";
       writeStackOperations(spec, adviceBody);
       writeOriginalMethodCall(spec, adviceBody);
       writeAdviceMethodCall(spec, adviceBody);
     } else {
+      type = "AROUND";
       writeAdviceMethodCall(spec, adviceBody);
     }
     body.addStatement(
@@ -207,6 +212,9 @@ private void addAdviceLambda(
             .setName("addAdvice")
             .setArguments(
                 new NodeList<>(
+                    new FieldAccessExpr()
+                        .setScope(new TypeExpr(new ClassOrInterfaceType().setName(TYPE_CLASS)))
+                        .setName(type),
                     new StringLiteralExpr(pointCut.getOwner().getInternalName()),
                     new StringLiteralExpr(pointCut.getMethodName()),
                     new StringLiteralExpr(pointCut.getMethodType().getDescriptor()),

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

@@ -364,19 +364,19 @@ private static LambdaExpr findAdviceLambda(
     final MethodType pointcut = spec.getPointcut();
     for (final MethodCallExpr add : addAdvices) {
       final NodeList<Expression> arguments = add.getArguments();
-      final String owner = arguments.get(0).asStringLiteralExpr().asString();
+      final String owner = arguments.get(1).asStringLiteralExpr().asString();
       if (!owner.equals(pointcut.getOwner().getInternalName())) {
         continue;
       }
-      final String method = arguments.get(1).asStringLiteralExpr().asString();
+      final String method = arguments.get(2).asStringLiteralExpr().asString();
       if (!method.equals(pointcut.getMethodName())) {
         continue;
       }
-      final String description = arguments.get(2).asStringLiteralExpr().asString();
+      final String description = arguments.get(3).asStringLiteralExpr().asString();
       if (!description.equals(pointcut.getMethodType().getDescriptor())) {
         continue;
       }
-      return arguments.get(3).asLambdaExpr();
+      return arguments.get(4).asLambdaExpr();
     }
     throw new IllegalArgumentException("Cannot find lambda expression for pointcut " + pointcut);
   }

buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/util/CallSiteConstants.java

@@ -30,6 +30,8 @@ private CallSiteConstants() {}
 
   public static final String HAS_ENABLED_PROPERTY_CLASS = CALL_SITES_CLASS + ".HasEnabledProperty";
 
+  public static final String TYPE_CLASS = "Type";
+
   public static final String STACK_DUP_MODE_CLASS = "StackDupMode";
 
   public static final String METHOD_HANDLER_CLASS = "MethodHandler";

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

@@ -44,6 +44,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
       interfaces(CallSites)
       helpers(BeforeAdvice)
       advices(0) {
+        type("BEFORE")
         pointcut('java/security/MessageDigest', 'getInstance', '(Ljava/lang/String;)Ljava/security/MessageDigest;')
         statements(
           'handler.dupParameters(descriptor, StackDupMode.COPY);',
@@ -76,6 +77,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
       interfaces(CallSites)
       helpers(AroundAdvice)
       advices(0) {
+        type("AROUND")
         pointcut('java/lang/String', 'replaceAll', '(Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;')
         statements(
           'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AroundAdvice", "around", "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;");'
@@ -106,6 +108,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
       interfaces(CallSites)
       helpers(AfterAdvice)
       advices(0) {
+        type("AFTER")
         pointcut('java/lang/String', 'concat', '(Ljava/lang/String;)Ljava/lang/String;')
         statements(
           'handler.dupInvoke(owner, descriptor, StackDupMode.COPY);',

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

@@ -1,11 +1,16 @@
 package datadog.trace.plugin.csi.impl.assertion
 
 class AdviceAssert {
+  protected String type
   protected String owner
   protected String method
   protected String descriptor
   protected Collection<String> statements
 
+  void type(String type) {
+    assert type == this.type
+  }
+
   void pointcut(String owner, String method, String descriptor) {
     assert owner == this.owner
     assert method == this.method

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

@@ -83,10 +83,12 @@ class AssertBuilder<C extends CallSiteAssert> {
     return getMethodCalls(acceptMethod).findAll {
       it.nameAsString == 'addAdvice'
     }.collect {
-      def (owner, method, descriptor) = it.arguments.subList(0, 3)*.asStringLiteralExpr()*.asString()
-      final handlerLambda = it.arguments[3].asLambdaExpr()
+      final adviceType = it.arguments.get(0).asFieldAccessExpr().getName()
+      def (owner, method, descriptor) = it.arguments.subList(1, 4)*.asStringLiteralExpr()*.asString()
+      final handlerLambda = it.arguments[4].asLambdaExpr()
       final advice = handlerLambda.body.asBlockStmt().statements*.toString()
       return new AdviceAssert([
+        type      : adviceType,
         owner     : owner,
         method    : method,
         descriptor: descriptor,

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

@@ -217,8 +217,8 @@ class IastExtensionTest extends BaseCsiPluginTest {
       return getMethodCalls(acceptMethod).findAll {
         it.nameAsString == 'addAdvice'
       }.collect {
-        def (owner, method, descriptor) = it.arguments.subList(0, 3)*.asStringLiteralExpr()*.asString()
-        final handlerLambda = it.arguments[3].asLambdaExpr()
+        def (owner, method, descriptor) = it.arguments.subList(1, 4)*.asStringLiteralExpr()*.asString()
+        final handlerLambda = it.arguments[4].asLambdaExpr()
         final statements = handlerLambda.body.asBlockStmt().statements
         final instrumentedStmt = statements.get(0).asIfStmt()
         final executedStmt = statements.get(1).asIfStmt()

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/Advices.java

@@ -6,6 +6,8 @@
 import datadog.trace.agent.tooling.bytebuddy.ClassFileLocators;
 import datadog.trace.agent.tooling.csi.CallSiteAdvice;
 import datadog.trace.agent.tooling.csi.CallSites;
+import datadog.trace.agent.tooling.csi.InvokeAdvice;
+import datadog.trace.agent.tooling.csi.InvokeDynamicAdvice;
 import java.lang.reflect.Field;
 import java.util.Arrays;
 import java.util.Collections;
@@ -141,6 +143,11 @@ public CallSiteAdvice findAdvice(
     return methodAdvices.get(descriptor);
   }
 
+  /** Gets the type of advice we are dealing with */
+  public byte typeOf(final CallSiteAdvice advice) {
+    return ((HasType) advice).getType();
+  }
+
   public String[] getHelpers() {
     return helpers;
   }
@@ -176,15 +183,17 @@ public void addHelpers(final String... helperClassNames) {
 
     @Override
     public void addAdvice(
-        final String type,
+        final byte type,
+        final String owner,
         final String method,
         final String descriptor,
         final CallSiteAdvice advice) {
       final Map<String, Map<String, CallSiteAdvice>> typeAdvices =
-          advices.computeIfAbsent(type, k -> new HashMap<>());
+          advices.computeIfAbsent(owner, k -> new HashMap<>());
       final Map<String, CallSiteAdvice> methodAdvices =
           typeAdvices.computeIfAbsent(method, k -> new HashMap<>());
-      final CallSiteAdvice oldAdvice = methodAdvices.put(descriptor, advice);
+      final CallSiteAdvice oldAdvice =
+          methodAdvices.put(descriptor, HasType.withType(advice, type));
       if (oldAdvice != null) {
         throw new UnsupportedOperationException(
             String.format(
@@ -360,4 +369,67 @@ public interface Listener {
     void onConstantPool(
         @Nonnull TypeDescription type, @Nonnull ConstantPool pool, final byte[] classFile);
   }
+
+  private interface HasType {
+    byte getType();
+
+    static CallSiteAdvice withType(final CallSiteAdvice advice, final byte type) {
+      if (advice instanceof InvokeAdvice) {
+        return new InvokeWithType((InvokeAdvice) advice, type);
+      } else {
+        return new InvokeDynamicWithType((InvokeDynamicAdvice) advice, type);
+      }
+    }
+  }
+
+  private static class InvokeWithType implements InvokeAdvice, HasType {
+    private final InvokeAdvice advice;
+    private final byte type;
+
+    private InvokeWithType(InvokeAdvice advice, byte type) {
+      this.advice = advice;
+      this.type = type;
+    }
+
+    @Override
+    public byte getType() {
+      return type;
+    }
+
+    @Override
+    public void apply(
+        final MethodHandler handler,
+        final int opcode,
+        final String owner,
+        final String name,
+        final String descriptor,
+        final boolean isInterface) {
+      advice.apply(handler, opcode, owner, name, descriptor, isInterface);
+    }
+  }
+
+  private static class InvokeDynamicWithType implements InvokeDynamicAdvice, HasType {
+    private final InvokeDynamicAdvice advice;
+    private final byte type;
+
+    private InvokeDynamicWithType(final InvokeDynamicAdvice advice, final byte type) {
+      this.advice = advice;
+      this.type = type;
+    }
+
+    @Override
+    public byte getType() {
+      return type;
+    }
+
+    @Override
+    public void apply(
+        final MethodHandler handler,
+        final String name,
+        final String descriptor,
+        final Handle bootstrapMethodHandle,
+        final Object... bootstrapMethodArguments) {
+      advice.apply(handler, name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments);
+    }
+  }
 }