Fix class crawl logic

Author: sarahchen6

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

@@ -18,7 +18,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.function.Predicate;
 import net.bytebuddy.asm.AsmVisitorWrapper;
 import net.bytebuddy.description.field.FieldDescription;
 import net.bytebuddy.description.field.FieldList;
@@ -88,8 +87,7 @@ public ClassVisitor wrap(
   private static Reference[] generateReferences(
       Instrumenter.HasMethodAdvice instrumenter,
       AdviceShader adviceShader,
-      Set<String> allAdviceClasses,
-      Predicate<String> shouldRecurse) {
+      Set<String> allAdviceClasses) {
     // track sources we've generated references from to avoid recursion
     final Set<String> referenceSources = new HashSet<>();
     final Map<String, Reference> references = new LinkedHashMap<>();
@@ -107,8 +105,7 @@ private static Reference[] generateReferences(
     for (String adviceClass : adviceClasses) {
       if (referenceSources.add(adviceClass)) {
         for (Map.Entry<String, Reference> entry :
-            ReferenceCreator.createReferencesFrom(
-                    adviceClass, adviceShader, contextClassLoader, shouldRecurse)
+            ReferenceCreator.createReferencesFrom(adviceClass, adviceShader, contextClassLoader)
                 .entrySet()) {
           Reference toMerge = references.get(entry.getKey());
           if (null == toMerge) {
@@ -128,22 +125,19 @@ private byte[] generateMuzzleClass(InstrumenterModule module) {
     AdviceShader adviceShader = AdviceShader.with(module.adviceShading());
     HelperClassPredicate helperPredicate = new HelperClassPredicate(this::isOwnOutput);
 
-    // Crawl advice, following references into our own helper classes (even in library packages).
+    // Crawl advice for muzzle references (only recursing into the instrumentation package).
     Set<String> adviceClasses = new HashSet<>();
     List<Reference> allReferences = new ArrayList<>();
     for (Instrumenter instrumenter : module.typeInstrumentations()) {
       if (instrumenter instanceof Instrumenter.HasMethodAdvice) {
         Collections.addAll(
             allReferences,
             generateReferences(
-                (Instrumenter.HasMethodAdvice) instrumenter,
-                adviceShader,
-                adviceClasses,
-                helperPredicate::isHelperClass));
+                (Instrumenter.HasMethodAdvice) instrumenter, adviceShader, adviceClasses));
       }
     }
 
-    // Advice roots are inlined into the target method, not injected, so they aren't helpers.
+    // Inferred helpers = our classes referenced from the advice, minus the advice roots.
     Set<String> inferredHelpers = new LinkedHashSet<>();
     for (Reference reference : allReferences) {
       if (!adviceClasses.contains(reference.className)
@@ -153,21 +147,26 @@ private byte[] generateMuzzleClass(InstrumenterModule module) {
     }
 
     // Manual additions cover helpers the crawl can't see (reflection, SPI, advice-less injectors).
-    Set<String> helperClasses = new LinkedHashSet<>(inferredHelpers);
-    Collections.addAll(helperClasses, module.helperClassNames());
-    for (String helper : new ArrayList<>(helperClasses)) {
+    Set<String> manualHelpers = new LinkedHashSet<>(asList(module.helperClassNames()));
+    Set<String> seedHelpers = new LinkedHashSet<>(inferredHelpers);
+    seedHelpers.addAll(manualHelpers);
+    for (String helper : new ArrayList<>(seedHelpers)) {
       if (isOwnOutput(helper)) {
-        addNestedClasses(helper, helperClasses);
+        addNestedClasses(helper, seedHelpers);
       }
     }
 
     String[] orderedHelpers =
-        orderHelpers(helperClasses, Thread.currentThread().getContextClassLoader());
+        discoverAndOrderHelpers(
+            seedHelpers,
+            manualHelpers,
+            helperPredicate,
+            Thread.currentThread().getContextClassLoader());
 
-    writeInferenceReport(module, adviceClasses.isEmpty(), inferredHelpers, helperClasses);
+    writeInferenceReport(module, adviceClasses.isEmpty(), inferredHelpers, orderedHelpers);
 
     // Injected helpers are our own classes, so they must not be asserted as library references.
-    Set<String> ignoredClassNames = new HashSet<>(helperClasses);
+    Set<String> ignoredClassNames = new HashSet<>(asList(orderedHelpers));
     Collections.addAll(ignoredClassNames, module.muzzleIgnoredClassNames());
     List<Reference> references = new ArrayList<>();
     for (Reference reference : allReferences) {
@@ -225,19 +224,21 @@ private byte[] generateMuzzleClass(InstrumenterModule module) {
     mv.visitMaxs(0, 0);
     mv.visitEnd();
 
-    // Generated: public static String[] helperClassNames()
-    MethodVisitor hv =
-        cw.visitMethod(
-            Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC,
-            "helperClassNames",
-            "()[Ljava/lang/String;",
-            null,
-            null);
-    hv.visitCode();
-    writeStrings(hv, orderedHelpers);
-    hv.visitInsn(Opcodes.ARETURN);
-    hv.visitMaxs(0, 0);
-    hv.visitEnd();
+    // Generated: public static String[] helperClassNames() — omitted for helper-less modules.
+    if (orderedHelpers.length > 0) {
+      MethodVisitor hv =
+          cw.visitMethod(
+              Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC,
+              "helperClassNames",
+              "()[Ljava/lang/String;",
+              null,
+              null);
+      hv.visitCode();
+      writeStrings(hv, orderedHelpers);
+      hv.visitInsn(Opcodes.ARETURN);
+      hv.visitMaxs(0, 0);
+      hv.visitEnd();
+    }
 
     return cw.toByteArray();
   }
@@ -270,28 +271,35 @@ private void addNestedClasses(String className, Set<String> helperClasses) {
   }
 
   /**
-   * Load-orders helpers (dependencies first) via {@link HelperScanner}, keeping only our helpers
-   * (the scanner may pull in library classes) and retaining any helper it could not locate.
+   * Runs {@link HelperScanner} over the seed helpers to both discover their transitive helper
+   * dependencies and load-order the result (dependencies first). Keeps only our own helpers (plus
+   * manual additions), dropping library classes the scanner pulls in, and retains every seed even
+   * if it could not be located.
    */
-  private static String[] orderHelpers(Set<String> helpers, ClassLoader loader) {
-    if (helpers.isEmpty()) {
+  private static String[] discoverAndOrderHelpers(
+      Set<String> seedHelpers,
+      Set<String> manualHelpers,
+      HelperClassPredicate helperPredicate,
+      ClassLoader loader) {
+    if (seedHelpers.isEmpty()) {
       return new String[0];
     }
-    List<String> ordered = new ArrayList<>(helpers.size());
+    List<String> ordered = new ArrayList<>();
     try {
       for (String name :
           HelperScanner.withClassDependencies(
-              ClassFileLocator.ForClassLoader.of(loader), helpers.toArray(new String[0]))) {
-        if (helpers.contains(name) && !ordered.contains(name)) {
+              ClassFileLocator.ForClassLoader.of(loader), seedHelpers.toArray(new String[0]))) {
+        if ((helperPredicate.isHelperClass(name) || manualHelpers.contains(name))
+            && !ordered.contains(name)) {
           ordered.add(name);
         }
       }
     } catch (Throwable ignore) {
-      // best-effort ordering; unordered helpers are appended below
+      // best-effort ordering; unordered seeds are appended below
     }
-    for (String name : helpers) {
-      if (!ordered.contains(name)) {
-        ordered.add(name);
+    for (String seed : seedHelpers) {
+      if (!ordered.contains(seed)) {
+        ordered.add(seed);
       }
     }
     return ordered.toArray(new String[0]);
@@ -305,7 +313,7 @@ private void writeInferenceReport(
       InstrumenterModule module,
       boolean adviceLess,
       Set<String> inferredHelpers,
-      Set<String> allHelpers) {
+      String[] injectedHelpers) {
     try {
       Set<String> declared = new LinkedHashSet<>(asList(module.helperClassNames()));
       if (declared.isEmpty() && inferredHelpers.isEmpty()) {
@@ -323,7 +331,7 @@ private void writeInferenceReport(
       StringBuilder sb = new StringBuilder();
       sb.append("module: ").append(module.getClass().getName()).append('\n');
       sb.append("has-advice: ").append(!adviceLess).append('\n');
-      sb.append("injected-helpers: ").append(allHelpers.size()).append('\n');
+      sb.append("injected-helpers: ").append(injectedHelpers.length).append('\n');
       if (adviceLess && !declared.isEmpty()) {
         sb.append("crawl-cannot-cover: advice-less module; declared helpers are all manual-only\n");
       }

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

@@ -14,7 +14,6 @@
 import java.util.Map;
 import java.util.Queue;
 import java.util.Set;
-import java.util.function.Predicate;
 import net.bytebuddy.jar.asm.ClassReader;
 import net.bytebuddy.jar.asm.ClassVisitor;
 import net.bytebuddy.jar.asm.FieldVisitor;
@@ -39,10 +38,6 @@ public class ReferenceCreator extends ClassVisitor {
    */
   private static final String REFERENCE_CREATION_PACKAGE = "datadog.trace.instrumentation.";
 
-  /** Default recursion rule; callers that identify helpers by other means pass their own. */
-  private static final Predicate<String> DEFAULT_SHOULD_RECURSE =
-      name -> name.startsWith(REFERENCE_CREATION_PACKAGE);
-
   private static final int UNDEFINED_LINE = -1;
 
   /** Set containing name+descriptor signatures of Object methods. */
@@ -63,30 +58,9 @@ public class ReferenceCreator extends ClassVisitor {
    * @return Map of [referenceClassName -> Reference]
    * @throws IllegalStateException if class is not found or unable to be loaded.
    */
-  public static Map<String, Reference> createReferencesFrom(
-      final String entryPointClassName, final AdviceShader adviceShader, final ClassLoader loader)
-      throws IllegalStateException {
-    return createReferencesFrom(entryPointClassName, adviceShader, loader, DEFAULT_SHOULD_RECURSE);
-  }
-
-  /**
-   * Generate all references reachable from a given class, following references into classes
-   * accepted by {@code shouldRecurse}.
-   *
-   * @param entryPointClassName Starting point for generating references.
-   * @param adviceShader Optional shading to apply to the advice.
-   * @param loader Classloader used to read class bytes.
-   * @param shouldRecurse Decides whether a referenced class is one of ours to keep crawling into
-   *     (vs. a library/bootstrap leaf that is merely recorded).
-   * @return Map of [referenceClassName -> Reference]
-   * @throws IllegalStateException if class is not found or unable to be loaded.
-   */
   @SuppressForbidden
   public static Map<String, Reference> createReferencesFrom(
-      final String entryPointClassName,
-      final AdviceShader adviceShader,
-      final ClassLoader loader,
-      final Predicate<String> shouldRecurse)
+      final String entryPointClassName, final AdviceShader adviceShader, final ClassLoader loader)
       throws IllegalStateException {
     final Set<String> visitedSources = new HashSet<>();
     final Map<String, Reference> references = new LinkedHashMap<>();
@@ -113,8 +87,9 @@ public static Map<String, Reference> createReferencesFrom(
 
         final Map<String, Reference> instrumentationReferences = cv.getReferences();
         for (final Map.Entry<String, Reference> entry : instrumentationReferences.entrySet()) {
-          // Only keep crawling into classes that are ours
-          if (!visitedSources.contains(entry.getKey()) && shouldRecurse.test(entry.getKey())) {
+          // Don't generate references created outside of the datadog instrumentation package.
+          if (!visitedSources.contains(entry.getKey())
+              && entry.getKey().startsWith(REFERENCE_CREATION_PACKAGE)) {
             instrumentationQueue.add(entry.getKey());
           }
           Reference toMerge = references.get(entry.getKey());