Implement auto-inferring helper methods

Author: sarahchen6

dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java

@@ -130,7 +130,7 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati
 
     adviceShader = AdviceShader.with(module);
 
-    String[] helperClassNames = module.helperClassNames();
+    String[] helperClassNames = module.getAllHelperClassNames();
     if (module.injectHelperDependencies()) {
       helperClassNames = HelperScanner.withClassDependencies(helperClassNames);
     }

dd-java-agent/agent-tooling/build.gradle

@@ -49,6 +49,7 @@ dependencies {
 
   testImplementation project(':dd-java-agent:testing')
   testImplementation libs.bytebuddy
+  testImplementation libs.bundles.junit5
   testImplementation group: 'com.google.guava', name: 'guava-testlib', version: '20.0'
 
   jmhImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: '2.3.5.RELEASE'

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperScanner.java

@@ -24,8 +24,7 @@
 public final class HelperScanner extends ClassVisitor {
   static final int READER_OPTIONS = ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES;
 
-  static final ClassFileLocator locator =
-      ClassFileLocator.ForClassLoader.of(Utils.getAgentClassLoader());
+  final ClassFileLocator locator;
 
   final MethodScanner methodScanner = new MethodScanner();
 
@@ -41,14 +40,30 @@ public final class HelperScanner extends ClassVisitor {
   Set<String> uses;
 
   HelperScanner() {
+    this(ClassFileLocator.ForClassLoader.of(Utils.getAgentClassLoader()));
+  }
+
+  HelperScanner(ClassFileLocator locator) {
     super(Opcodes.ASM7, null);
+    this.locator = locator;
   }
 
-  /** Expands helper class names to include any non-bootstrap classes they depend on. */
+  /**
+   * Expands helper class names with their non-bootstrap dependencies, via the agent class loader.
+   */
   public static String[] withClassDependencies(String... helperClassNames) {
     return new HelperScanner().simulateClassLoading(helperClassNames);
   }
 
+  /**
+   * Same as above, but reads bytecode via the given locator (e.g. during build time where the agent
+   * loader is absent).
+   */
+  public static String[] withClassDependencies(
+      ClassFileLocator locator, String... helperClassNames) {
+    return new HelperScanner(locator).simulateClassLoading(helperClassNames);
+  }
+
   /**
    * Simulates class-loading by finding all classes required to load the helper classes as well as
    * optional classes used in method instructions that may be needed later when invoking the method.

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java

@@ -113,13 +113,31 @@ public static ReferenceMatcher loadStaticMuzzleReferences(
   }
 
   /**
-   * @return Class names of helpers to inject into the user's classloader.
-   *     <p><b>NOTE:</b> The order of the returned helper classes matters. If a muzzle check fails
-   *     with a NoClassDefFoundError, as logged in build/reports/muzzle-*.txt, it is likely that one
-   *     helper class depends on another that appears later in the list. In this case, the returned
-   *     list must be reordered so that the referred helper class appears before the one that refers
-   *     to it.
+   * @return the full, load-ordered set of helper classes to inject into the user's classloader: the
+   *     helpers inferred at build time from the advice/helper byte-code, unioned with any {@link
+   *     #helperClassNames() manually-declared} additions. Falls back to {@link #helperClassNames()}
+   *     when no generated {@code $Muzzle} is available.
    */
+  public final String[] getAllHelperClassNames() {
+    String[] generated =
+        loadStaticMuzzleHelperClassNames(getClass().getClassLoader(), getClass().getName());
+    return null != generated ? generated : helperClassNames();
+  }
+
+  static String[] loadStaticMuzzleHelperClassNames(
+      ClassLoader classLoader, String instrumentationClass) {
+    String muzzleClass = instrumentationClass + "$Muzzle";
+    try {
+      // helper class names captured at build-time; see MuzzleGenerator
+      return (String[])
+          classLoader.loadClass(muzzleClass).getMethod("helperClassNames").invoke(null);
+    } catch (Throwable e) {
+      // no generated helper list: caller falls back to helperClassNames()
+      return null;
+    }
+  }
+
+  /** Optional manual additions to the injected helper set. */
   public String[] helperClassNames() {
     return NO_HELPERS;
   }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/HelperClassPredicate.java

@@ -0,0 +1,65 @@
+package datadog.trace.agent.tooling.muzzle;
+
+import datadog.trace.bootstrap.Constants;
+import java.util.function.Predicate;
+
+/**
+ * Classifies a referenced class as an injectable tracer helper, a bootstrap class, or a library
+ * class — similar to OpenTelemetry's {@code HelperClassPredicate#isHelperClass}; however, the main
+ * signal here is whether the class was compiled from the instrumentation subproject's own output
+ * ({@code ownOutput}). {@link #HELPER_PREFIXES} additionally covers helpers that live in other
+ * tracer subprojects.
+ */
+public final class HelperClassPredicate {
+
+  static final String[] HELPER_PREFIXES = {
+    "datadog.trace.instrumentation.",
+    "datadog.opentelemetry.shim.",
+    "datadog.trace.agent.tooling.iast.",
+    "datadog.trace.agent.tooling.nativeimage.",
+  };
+
+  private final Predicate<String> ownOutput;
+
+  /**
+   * @param ownOutput whether a dotted class name was compiled from the instrumentation subproject's
+   *     own output; injected because agent-tooling cannot resolve build directories itself.
+   */
+  public HelperClassPredicate(final Predicate<String> ownOutput) {
+    this.ownOutput = ownOutput;
+  }
+
+  public boolean isHelperClass(final String className) {
+    return !isBootstrap(className) && (ownOutput.test(className) || matchesHelperPrefix(className));
+  }
+
+  private static boolean matchesHelperPrefix(final String className) {
+    for (final String prefix : HELPER_PREFIXES) {
+      if (className.startsWith(prefix)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  /**
+   * Whether the class is on the bootstrap class-path (or a JDK/SLF4J type) and so never injected.
+   */
+  public static boolean isBootstrap(final String className) {
+    if (className.startsWith("java.")
+        || className.startsWith("javax.")
+        || className.startsWith("jdk.")
+        || className.startsWith("com.sun.")
+        || className.startsWith("sun.")
+        || className.startsWith("org.slf4j.")
+        || className.startsWith("datadog.slf4j.")) {
+      return true;
+    }
+    for (final String prefix : Constants.BOOTSTRAP_PACKAGE_PREFIXES) {
+      if (className.startsWith(prefix)) {
+        return true;
+      }
+    }
+    return false;
+  }
+}