Add IAST telemetry annotations to bytebuddy advices (#5430)

Author: manuel-alvarez-alvarez

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

@@ -26,6 +26,7 @@
 import com.github.javaparser.ast.stmt.BlockStmt;
 import com.github.javaparser.ast.stmt.IfStmt;
 import com.github.javaparser.ast.stmt.Statement;
+import com.github.javaparser.resolution.declarations.ResolvedFieldDeclaration;
 import com.github.javaparser.symbolsolver.JavaSymbolSolver;
 import datadog.trace.plugin.csi.AdviceGenerator.CallSiteResult;
 import datadog.trace.plugin.csi.Extension;
@@ -37,6 +38,7 @@
 import datadog.trace.plugin.csi.util.MethodType;
 import java.io.File;
 import java.io.FileNotFoundException;
+import java.lang.reflect.Field;
 import java.net.URL;
 import java.net.URLClassLoader;
 import java.nio.file.Files;
@@ -116,7 +118,7 @@ private void addTelemetry(final Configuration configuration, final CallSiteResul
         if (metaData != null) {
           final List<LambdaExpr> adviceLambdas = filterAdviceLambdas(advices, callSiteMethod);
           for (final LambdaExpr advice : adviceLambdas) {
-            addTelemetryToAdvice(advice, metaData);
+            addTelemetryToAdvice(resolver, advice, metaData);
             hasTelemetry = true;
           }
         }
@@ -166,21 +168,24 @@ private CompilationUnit findOriginalCallSite(
     return parseSourceFile(configuration, resolver, javaFile);
   }
 
-  private void addTelemetryToAdvice(final LambdaExpr adviceLambda, final AdviceMetadata metaData) {
+  private void addTelemetryToAdvice(
+      final TypeResolver resolver, final LambdaExpr adviceLambda, final AdviceMetadata metaData) {
     final BlockStmt lambdaBody = adviceLambda.getBody().asBlockStmt();
     final String metric = getMetricName(metaData);
+    final String tagValue = getMetricTagValue(resolver, metaData);
     final String instrumentedMetric = "INSTRUMENTED_" + metric;
     final IfStmt instrumentedStatement =
         new IfStmt()
             .setCondition(isEnabledCondition(instrumentedMetric))
             .setThenStmt(
-                new BlockStmt().addStatement(addTelemetryCollectorMethod(instrumentedMetric)));
+                new BlockStmt()
+                    .addStatement(addTelemetryCollectorMethod(instrumentedMetric, tagValue)));
     lambdaBody.addStatement(0, instrumentedStatement);
     final String executedMetric = "EXECUTED_" + metric;
     final IfStmt executedStatement =
         new IfStmt()
             .setCondition(isEnabledCondition(executedMetric))
-            .setThenStmt(addTelemetryCollectorByteCode(executedMetric));
+            .setThenStmt(addTelemetryCollectorByteCode(executedMetric, tagValue));
     lambdaBody.addStatement(1, executedStatement);
   }
 
@@ -191,19 +196,27 @@ private static Expression isEnabledCondition(final String metric) {
         .addArgument(accessLocalField("verbosity"));
   }
 
-  private static MethodCallExpr addTelemetryCollectorMethod(final String metric) {
-    return new MethodCallExpr()
-        .setScope(new NameExpr(IAST_METRIC_COLLECTOR_CLASS))
-        .setName("add")
-        .addArgument(
-            new FieldAccessExpr().setScope(new NameExpr(IAST_METRIC_CLASS)).setName(metric))
-        .addArgument(intLiteral(1));
+  private static MethodCallExpr addTelemetryCollectorMethod(
+      final String metric, final String tagValue) {
+    final MethodCallExpr method =
+        new MethodCallExpr()
+            .setScope(new NameExpr(IAST_METRIC_COLLECTOR_CLASS))
+            .setName("add")
+            .addArgument(
+                new FieldAccessExpr().setScope(new NameExpr(IAST_METRIC_CLASS)).setName(metric));
+    if (tagValue != null) {
+      method.addArgument(new StringLiteralExpr(tagValue));
+    }
+    method.addArgument(intLiteral(1));
+    return method;
   }
 
-  private static BlockStmt addTelemetryCollectorByteCode(final String metric) {
+  private static BlockStmt addTelemetryCollectorByteCode(
+      final String metric, final String tagValue) {
     final BlockStmt stmt = new BlockStmt();
     // this code generates the java source code needed to provide the bytecode for the statement
-    // IastTelemetryCollector.add($"{metric}, 1);
+    // IastTelemetryCollector.add(${metric}, 1); or IastTelemetryCollector.add(${metric}, ${tag},
+    // 1);
     stmt.addStatement(
         new MethodCallExpr()
             .setScope(new NameExpr("handler"))
@@ -213,12 +226,23 @@ private static BlockStmt addTelemetryCollectorByteCode(final String metric) {
             .addArgument(new StringLiteralExpr(IAST_METRIC_INTERNAL_NAME))
             .addArgument(new StringLiteralExpr(metric))
             .addArgument(new StringLiteralExpr("L" + IAST_METRIC_INTERNAL_NAME + ";")));
+    if (tagValue != null) {
+      stmt.addStatement(
+          new MethodCallExpr()
+              .setScope(new NameExpr("handler"))
+              .setName("loadConstant")
+              .addArgument(new StringLiteralExpr(tagValue)));
+    }
     stmt.addStatement(
         new MethodCallExpr()
             .setScope(new NameExpr("handler"))
             .setName("instruction")
             .addArgument(
-                new FieldAccessExpr().setScope(new NameExpr(OPCODES_FQDN)).setName("LCONST_1")));
+                new FieldAccessExpr().setScope(new NameExpr(OPCODES_FQDN)).setName("ICONST_1")));
+    final String descriptor =
+        tagValue != null
+            ? "(L" + IAST_METRIC_INTERNAL_NAME + ";Ljava/lang/String;I)V"
+            : "(L" + IAST_METRIC_INTERNAL_NAME + ";I)V";
     stmt.addStatement(
         new MethodCallExpr()
             .setScope(new NameExpr("handler"))
@@ -227,19 +251,49 @@ private static BlockStmt addTelemetryCollectorByteCode(final String metric) {
                 new FieldAccessExpr().setScope(new NameExpr(OPCODES_FQDN)).setName("INVOKESTATIC"))
             .addArgument(new StringLiteralExpr(IAST_METRIC_COLLECTOR_INTERNAL_NAME))
             .addArgument(new StringLiteralExpr("add"))
-            .addArgument(new StringLiteralExpr("(L" + IAST_METRIC_INTERNAL_NAME + ";J)V"))
+            .addArgument(new StringLiteralExpr(descriptor))
             .addArgument(new BooleanLiteralExpr(false)));
     return stmt;
   }
 
   private static String getMetricName(final AdviceMetadata metaData) {
-    final StringBuilder metric = new StringBuilder(metaData.getKind());
-    if (metaData.getTag() != null) {
-      // Uses the name of the field to compose the name of the metric
-      final Expression tag = metaData.getTag();
-      metric.append("_").append(tag.asFieldAccessExpr().getName());
+    final AnnotationExpr kind = metaData.getKind();
+    return kind.getName().getId().toUpperCase();
+  }
+
+  private static String getMetricTagValue(
+      final TypeResolver resolver, final AdviceMetadata metadata) {
+    if (metadata.getTag() == null) {
+      return null;
+    }
+    final Expression tag = metadata.getTag();
+    if (tag.isStringLiteralExpr()) {
+      return tag.asStringLiteralExpr().getValue();
+    } else {
+      return getFieldValue(resolver, tag.asFieldAccessExpr());
+    }
+  }
+
+  private static String getFieldValue(final TypeResolver resolver, final FieldAccessExpr tag) {
+    final FieldAccessExpr fieldAccessExpr = tag.asFieldAccessExpr();
+    final ResolvedFieldDeclaration value = fieldAccessExpr.resolve().asField();
+    try {
+      final Field field = getField(value);
+      field.setAccessible(true);
+      return (String) field.get(field.getDeclaringClass());
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private static Field getField(final ResolvedFieldDeclaration resolved) {
+    try {
+      final Field field = resolved.getClass().getDeclaredField("field");
+      field.setAccessible(true);
+      return (Field) field.get(resolved);
+    } catch (Exception e) {
+      throw new RuntimeException(e);
     }
-    return metric.toString();
   }
 
   /** Find all advice lambdas in the generated call site provider */
@@ -367,10 +421,10 @@ private static Expression getAnnotationExpression(final AnnotationExpr expr) {
   }
 
   private static class AdviceMetadata {
-    private final String kind;
+    private final AnnotationExpr kind;
     private final Expression tag;
 
-    private AdviceMetadata(final String kind, final Expression tag) {
+    private AdviceMetadata(final AnnotationExpr kind, final Expression tag) {
       this.kind = kind;
       this.tag = tag;
     }
@@ -381,8 +435,7 @@ private static AdviceMetadata findAdviceMetadata(final BodyDeclaration<?> target
           .map(
               annotation -> {
                 final Expression tag = getAnnotationExpression(annotation);
-                final String typeName = annotation.getName().getId();
-                return new AdviceMetadata(typeName.toUpperCase(), tag);
+                return new AdviceMetadata(annotation, tag);
               })
           .findFirst()
           .orElse(null);
@@ -395,7 +448,7 @@ private static boolean isAdviceAnnotation(final AnnotationExpr expr) {
           || identifier.equals("Sink");
     }
 
-    public String getKind() {
+    public AnnotationExpr getKind() {
       return kind;
     }
 

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

@@ -1,32 +1,20 @@
 package datadog.trace.plugin.csi.impl
 
-import com.github.javaparser.JavaParser
-import com.github.javaparser.ParserConfiguration
-import com.github.javaparser.ast.CompilationUnit
-import com.github.javaparser.ast.Node
-import com.github.javaparser.ast.body.MethodDeclaration
-import com.github.javaparser.ast.body.TypeDeclaration
-import com.github.javaparser.ast.expr.ConditionalExpr
-import com.github.javaparser.ast.expr.MethodCallExpr
-import com.github.javaparser.symbolsolver.JavaSymbolSolver
+
 import datadog.trace.agent.tooling.csi.CallSite
 import datadog.trace.agent.tooling.csi.CallSites
 import datadog.trace.plugin.csi.AdviceGenerator
+import datadog.trace.plugin.csi.impl.assertion.AssertBuilder
+import datadog.trace.plugin.csi.impl.assertion.CallSiteAssert
 import groovy.transform.CompileDynamic
-import org.objectweb.asm.Type
 import spock.lang.Requires
 import spock.lang.TempDir
 
 import javax.servlet.ServletRequest
 import java.lang.invoke.MethodHandles
 import java.lang.invoke.MethodType
-import java.lang.reflect.Executable
-import java.lang.reflect.Method
-import java.security.MessageDigest
 
 import static CallSiteFactory.pointcutParser
-import static datadog.trace.plugin.csi.impl.CallSiteFactory.typeResolver
-import static datadog.trace.plugin.csi.util.CallSiteUtils.classNameToType
 
 @CompileDynamic
 final class AdviceGeneratorTest extends BaseCsiPluginTest {
@@ -464,127 +452,9 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
     return new AdviceGeneratorImpl(targetFolder, pointcutParser())
   }
 
-  private static CompilationUnit parseJavaFile(final File file)
-    throws FileNotFoundException {
-    final JavaSymbolSolver solver = new JavaSymbolSolver(typeResolver());
-    final JavaParser parser = new JavaParser(new ParserConfiguration().setSymbolResolver(solver));
-    return parser.parse(file).getResult().get();
-  }
-
-  private static Map<String, MethodDeclaration> groupMethods(final TypeDeclaration<?> classNode) {
-    return classNode.methods.groupBy { it.name.asString() }
-      .collectEntries { key, value -> [key, value.get(0)] }
-  }
-
-  private static List<Class<?>> getImplementedTypes(final TypeDeclaration<?> type) {
-    return type.asClassOrInterfaceDeclaration().implementedTypes.collect {
-      final resolved = it.asClassOrInterfaceType().resolve()
-      return resolved.typeDeclaration.get().clazz
-    }
-  }
-
-  private static Executable resolveMethod(final MethodCallExpr methodCallExpr) {
-    final resolved = methodCallExpr.resolve()
-    return resolved.@method as Method
-  }
-
-  private static List<MethodCallExpr> getMethodCalls(final MethodDeclaration method) {
-    return method.body.get().statements.findAll {
-      it.isExpressionStmt() && it.asExpressionStmt().expression.isMethodCallExpr()
-    }.collect {
-      it.asExpressionStmt().expression.asMethodCallExpr()
-    }
-  }
-
   private static void assertCallSites(final File generated, @DelegatesTo(CallSiteAssert) final Closure closure) {
-    final asserter = buildAsserter(generated)
+    final asserter = new AssertBuilder(generated).build()
     closure.delegate = asserter
     closure(asserter)
   }
-
-  private static CallSiteAssert buildAsserter(final File file) {
-    final javaFile = parseJavaFile(file)
-    assert javaFile.parsed == Node.Parsedness.PARSED
-    final adviceClass = javaFile.primaryType.get()
-    final interfaces = getImplementedTypes(adviceClass)
-    final methods = groupMethods(adviceClass)
-    Executable enabled = null
-    List<String> enabledArgs = null
-    if (interfaces.contains(CallSites.HasEnabledProperty)) {
-      final isEnabled = methods['isEnabled']
-      final returnStatement = isEnabled.body.get().statements.first.get().asReturnStmt()
-      final enabledMethodCall = returnStatement.expression.get().asMethodCallExpr()
-      enabled = resolveMethod(enabledMethodCall)
-      enabledArgs = enabledMethodCall.getArguments().collect { it.asStringLiteralExpr().asString() }
-    }
-    final accept = methods['accept']
-    final methodCalls = getMethodCalls(accept)
-    final addHelpers = methodCalls.find { it.name.toString() == 'addHelpers' }
-    assert addHelpers.scope.get().toString() == 'container'
-    assert addHelpers.name.toString() == 'addHelpers'
-    final helpers = addHelpers.getArguments().collect { typeResolver().resolveType(classNameToType(it.asStringLiteralExpr().asString())) }
-    final addAdvices = methodCalls.findAll { it.name.toString() == 'addAdvice' }
-    final advices = addAdvices.collect {
-      def (owner, method, descriptor) =  it.arguments.subList(0, 3)*.asStringLiteralExpr()*.asString()
-      final handlerLambda = it.arguments[3].asLambdaExpr()
-      final advice = handlerLambda.body.asBlockStmt().statements*.toString()
-      return new AdviceAssert([
-        owner     : owner,
-        method    : method,
-        descriptor: descriptor,
-        statements: advice
-      ])
-    }
-    return new CallSiteAssert([
-      interfaces : interfaces,
-      helpers    : helpers,
-      advices    : advices,
-      enabled    : enabled,
-      enabledArgs: enabledArgs
-    ])
-  }
-
-  static class CallSiteAssert {
-    private Collection<Class<?>> interfaces
-    private Collection<Class<?>> helpers
-    private Collection<AdviceAssert> advices
-    private Method enabled
-    private Collection<String> enabledArgs
-
-    void interfaces(Class<?>... values) {
-      assert values.toList() == interfaces
-    }
-
-    void helpers(Class<?>... values) {
-      assert values.toList() == helpers
-    }
-
-    void advices(int index, @DelegatesTo(AdviceAssert) Closure closure) {
-      final asserter = advices[index]
-      closure.delegate = asserter
-      closure(asserter)
-    }
-
-    void enabled(Method method, String... args) {
-      assert method == enabled
-      assert args.toList() == enabledArgs
-    }
-  }
-
-  static class AdviceAssert {
-    private String owner
-    private String method
-    private String descriptor
-    private Collection<String> statements
-
-    void pointcut(String owner, String method, String descriptor) {
-      assert owner == this.owner
-      assert method == this.method
-      assert descriptor == this.descriptor
-    }
-
-    void statements(String... values) {
-      assert values.toList() == statements
-    }
-  }
 }

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

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