diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/HelperResolver.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/HelperResolver.java new file mode 100644 index 00000000000..da8fe96111a --- /dev/null +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/HelperResolver.java @@ -0,0 +1,218 @@ +package datadog.trace.agent.tooling.muzzle; + +import static java.util.Arrays.asList; + +import datadog.trace.agent.tooling.AdviceShader; +import datadog.trace.agent.tooling.HelperScanner; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import java.io.File; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import net.bytebuddy.dynamic.ClassFileLocator; + +/** + * Resolves the helper classes an {@link InstrumenterModule} injects: a module with a manually + * declared {@code helperClassNames()} list uses it directly, otherwise the helpers inferred from + * its advice are used (dependency-ordered, build-time-only classes dropped). + * + *

Used by {@link MuzzleGenerator}, which resolves once per module and emits both outputs from + * that single crawl: the {@code $Muzzle} side-class (excluding these helpers from the asserted + * references) and the module's own {@code helperClassNames()} (holding the resolved list). + */ +final class HelperResolver { + private static final String MUZZLE_REFERENCE_API = "datadog/trace/agent/tooling/muzzle/Reference"; + + private final File sourceDir; + + HelperResolver(File sourceDir) { + this.sourceDir = sourceDir; + } + + /** The crawled advice references and the resolved helper set for a module. */ + static final class Result { + final List references; + final Set adviceClasses; + final String[] injectedHelpers; + + Result(List references, Set adviceClasses, String[] injectedHelpers) { + this.references = references; + this.adviceClasses = adviceClasses; + this.injectedHelpers = injectedHelpers; + } + } + + Result resolve(InstrumenterModule module) { + AdviceShader adviceShader = AdviceShader.with(module.adviceShading()); + + // Collect the muzzle references from every advice the module defines. + Set adviceClasses = new HashSet<>(); + List allReferences = new ArrayList<>(); + for (Instrumenter instrumenter : module.typeInstrumentations()) { + if (instrumenter instanceof Instrumenter.HasMethodAdvice) { + Collections.addAll( + allReferences, + generateReferences( + (Instrumenter.HasMethodAdvice) instrumenter, adviceShader, adviceClasses)); + } + } + return new Result( + allReferences, adviceClasses, computeInjectedHelpers(module, allReferences, adviceClasses)); + } + + private static Reference[] generateReferences( + Instrumenter.HasMethodAdvice instrumenter, + AdviceShader adviceShader, + Set allAdviceClasses) { + // track sources we've generated references from to avoid recursion + final Set referenceSources = new HashSet<>(); + final Map references = new LinkedHashMap<>(); + final Set adviceClasses = new HashSet<>(); + instrumenter.methodAdvice( + (matcher, adviceClass, additionalClasses) -> { + adviceClasses.add(adviceClass); + if (additionalClasses != null) { + adviceClasses.addAll(asList(additionalClasses)); + } + }); + // remember the advice roots so callers can exclude them from the injected helper set + allAdviceClasses.addAll(adviceClasses); + ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader(); + for (String adviceClass : adviceClasses) { + if (referenceSources.add(adviceClass)) { + for (Map.Entry entry : + ReferenceCreator.createReferencesFrom(adviceClass, adviceShader, contextClassLoader) + .entrySet()) { + Reference toMerge = references.get(entry.getKey()); + if (null == toMerge) { + references.put(entry.getKey(), entry.getValue()); + } else { + references.put(entry.getKey(), toMerge.merge(entry.getValue())); + } + } + } + } + return references.values().toArray(new Reference[0]); + } + + /** Resolves the ordered set of helper classes to inject for a module. */ + String[] computeInjectedHelpers( + InstrumenterModule module, List allReferences, Set adviceClasses) { + // A module that declares its own helper list uses it directly. + Set manualHelpers = new LinkedHashSet<>(asList(module.helperClassNames())); + if (!manualHelpers.isEmpty()) { + return manualHelpers.toArray(new String[0]); + } + + // Otherwise infer them: our classes referenced from the advice (minus the advice roots), plus + // their nested classes, dependency-ordered, with build-time-only muzzle providers dropped. + HelperClassPredicate helperPredicate = new HelperClassPredicate(this::isOwnOutput); + Set helpers = new LinkedHashSet<>(); + for (Reference reference : allReferences) { + if (!adviceClasses.contains(reference.className) + && helperPredicate.isHelperClass(reference.className)) { + helpers.add(reference.className); + } + } + for (String helper : new ArrayList<>(helpers)) { + if (isOwnOutput(helper)) { + addNestedClasses(helper, helpers); + } + } + ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader(); + String[] orderedHelpers = discoverAndOrderHelpers(helpers, helperPredicate, contextClassLoader); + ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(contextClassLoader); + List injectable = new ArrayList<>(orderedHelpers.length); + for (String helper : orderedHelpers) { + if (!isBuildTimeOnly(helper, locator)) { + injectable.add(helper); + } + } + return injectable.toArray(new String[0]); + } + + /** {@code true} if the class was compiled from this instrumentation subproject's own output. */ + private boolean isOwnOutput(String className) { + return new File(sourceDir, className.replace('.', '/') + ".class").isFile(); + } + + /** Adds the nested classes ({@code Foo$Bar}, {@code Foo$1}, ...) of an ownOutput helper. */ + private void addNestedClasses(String className, Set helperClasses) { + File classFile = new File(sourceDir, className.replace('.', '/') + ".class"); + File dir = classFile.getParentFile(); + if (dir == null || !dir.isDirectory()) { + return; + } + int lastDot = className.lastIndexOf('.'); + String pkg = lastDot < 0 ? "" : className.substring(0, lastDot + 1); + String prefix = (lastDot < 0 ? className : className.substring(lastDot + 1)) + "$"; + File[] siblings = dir.listFiles(); + if (siblings == null) { + return; + } + for (File sibling : siblings) { + String fileName = sibling.getName(); + if (fileName.startsWith(prefix) && fileName.endsWith(".class")) { + helperClasses.add(pkg + fileName.substring(0, fileName.length() - ".class".length())); + } + } + } + + /** + * {@code true} if the class uses the muzzle {@link Reference} API (as a {@link ReferenceProvider} + * or via {@code compileReferences}) - this means the class is build-time only and should not be + * injected. + */ + static boolean isBuildTimeOnly(String className, ClassFileLocator locator) { + try { + ClassFileLocator.Resolution resolution = locator.locate(className); + if (!resolution.isResolved()) { + return false; + } + // The muzzle type appears as a constant-pool entry when the class references it. + return new String(resolution.resolve(), StandardCharsets.ISO_8859_1) + .contains(MUZZLE_REFERENCE_API); + } catch (IOException e) { + return false; + } + } + + /** + * Expands the given helpers with any helper classes they depend on and returns them in + * dependency-first load order (required by {@link datadog.trace.agent.tooling.HelperInjector}) + * via {@link HelperScanner}. Library classes the scanner pulls in are dropped, but helpers that + * could not be located are kept (appended, unordered). + */ + private static String[] discoverAndOrderHelpers( + Set initialHelpers, HelperClassPredicate helperPredicate, ClassLoader loader) { + if (initialHelpers.isEmpty()) { + return new String[0]; + } + List ordered = new ArrayList<>(); + try { + for (String name : + HelperScanner.withClassDependencies( + ClassFileLocator.ForClassLoader.of(loader), initialHelpers.toArray(new String[0]))) { + if (helperPredicate.isHelperClass(name) && !ordered.contains(name)) { + ordered.add(name); + } + } + } catch (Throwable ignore) { + // best-effort ordering; unlocatable helpers are appended below + } + for (String helper : initialHelpers) { + if (!ordered.contains(helper)) { + ordered.add(helper); + } + } + return ordered.toArray(new String[0]); + } +} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java index 4319dbbee79..6c7bacdca73 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java @@ -2,28 +2,20 @@ import static java.util.Arrays.asList; -import datadog.trace.agent.tooling.AdviceShader; -import datadog.trace.agent.tooling.HelperScanner; -import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import java.io.File; import java.io.IOException; -import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.util.ArrayList; import java.util.Collections; import java.util.HashSet; -import java.util.LinkedHashMap; -import java.util.LinkedHashSet; import java.util.List; -import java.util.Map; import java.util.Set; import net.bytebuddy.asm.AsmVisitorWrapper; import net.bytebuddy.description.field.FieldDescription; import net.bytebuddy.description.field.FieldList; import net.bytebuddy.description.method.MethodList; import net.bytebuddy.description.type.TypeDescription; -import net.bytebuddy.dynamic.ClassFileLocator; import net.bytebuddy.implementation.Implementation; import net.bytebuddy.jar.asm.ClassVisitor; import net.bytebuddy.jar.asm.ClassWriter; @@ -77,67 +69,20 @@ public ClassVisitor wrap( throw new RuntimeException(e); } - AdviceShader adviceShader = AdviceShader.with(module.adviceShading()); - - // Collect the muzzle references from every advice the module defines. - Set adviceClasses = new HashSet<>(); - List allReferences = new ArrayList<>(); - for (Instrumenter instrumenter : module.typeInstrumentations()) { - if (instrumenter instanceof Instrumenter.HasMethodAdvice) { - Collections.addAll( - allReferences, - generateReferences( - (Instrumenter.HasMethodAdvice) instrumenter, adviceShader, adviceClasses)); - } - } - - String[] orderedHelpers = computeInjectedHelpers(module, allReferences, adviceClasses); + // Resolve helpers once, then emit: the $Muzzle side-class with references and the module's + // helperClassNames() with the resolved helper list. + HelperResolver.Result resolved = new HelperResolver(sourceDir).resolve(module); File muzzleClass = new File(targetDir, moduleDefinition.getInternalName() + "$Muzzle.class"); try { muzzleClass.getParentFile().mkdirs(); - Files.write(muzzleClass.toPath(), generateMuzzleClass(module, allReferences, orderedHelpers)); + Files.write( + muzzleClass.toPath(), + generateMuzzleClass(module, resolved.references, resolved.injectedHelpers)); } catch (IOException e) { throw new RuntimeException(e); } - - // Write the resolved helpers into the module's helperClassNames() so agent reads them directly. - return new HelperClassNamesWriter(classVisitor, orderedHelpers); - } - - private static Reference[] generateReferences( - Instrumenter.HasMethodAdvice instrumenter, - AdviceShader adviceShader, - Set allAdviceClasses) { - // track sources we've generated references from to avoid recursion - final Set referenceSources = new HashSet<>(); - final Map references = new LinkedHashMap<>(); - final Set adviceClasses = new HashSet<>(); - instrumenter.methodAdvice( - (matcher, adviceClass, additionalClasses) -> { - adviceClasses.add(adviceClass); - if (additionalClasses != null) { - adviceClasses.addAll(asList(additionalClasses)); - } - }); - // remember the advice roots so callers can exclude them from the injected helper set - allAdviceClasses.addAll(adviceClasses); - ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader(); - for (String adviceClass : adviceClasses) { - if (referenceSources.add(adviceClass)) { - for (Map.Entry entry : - ReferenceCreator.createReferencesFrom(adviceClass, adviceShader, contextClassLoader) - .entrySet()) { - Reference toMerge = references.get(entry.getKey()); - if (null == toMerge) { - references.put(entry.getKey(), entry.getValue()); - } else { - references.put(entry.getKey(), toMerge.merge(entry.getValue())); - } - } - } - } - return references.values().toArray(new Reference[0]); + return new HelperClassNamesWriter(classVisitor, resolved.injectedHelpers); } /** This code is generated in a separate side-class. */ @@ -209,7 +154,7 @@ private byte[] generateMuzzleClass( /** * Adds a {@code helperClassNames()} returning the build-time-resolved helper list to modules that - * don't declare one; a module that declares its own keeps it. + * don't declare one. */ private static final class HelperClassNamesWriter extends ClassVisitor { private static final String HELPER_METHOD = "helperClassNames"; @@ -248,120 +193,6 @@ public void visitEnd() { } } - /** Resolves the ordered set of helper classes to inject for a module. */ - String[] computeInjectedHelpers( - InstrumenterModule module, List allReferences, Set adviceClasses) { - // A module that declares its own helper list uses it directly. - Set manualHelpers = new LinkedHashSet<>(asList(module.helperClassNames())); - if (!manualHelpers.isEmpty()) { - return manualHelpers.toArray(new String[0]); - } - - // Otherwise infer them - HelperClassPredicate helperPredicate = new HelperClassPredicate(this::isOwnOutput); - Set helpers = new LinkedHashSet<>(); - for (Reference reference : allReferences) { - if (!adviceClasses.contains(reference.className) - && helperPredicate.isHelperClass(reference.className)) { - helpers.add(reference.className); - } - } - for (String helper : new ArrayList<>(helpers)) { - if (isOwnOutput(helper)) { - addNestedClasses(helper, helpers); - } - } - ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader(); - String[] orderedHelpers = discoverAndOrderHelpers(helpers, helperPredicate, contextClassLoader); - ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(contextClassLoader); - List injectableHelpers = new ArrayList<>(orderedHelpers.length); - for (String helper : orderedHelpers) { - if (!isBuildTimeOnly(helper, locator)) { - injectableHelpers.add(helper); - } - } - return injectableHelpers.toArray(new String[0]); - } - - /** {@code true} if the class was compiled from this instrumentation subproject's own output. */ - private boolean isOwnOutput(String className) { - return new File(sourceDir, className.replace('.', '/') + ".class").isFile(); - } - - /** Adds the nested classes ({@code Foo$Bar}, {@code Foo$1}, ...) of an ownOutput helper. */ - private void addNestedClasses(String className, Set helperClasses) { - File classFile = new File(sourceDir, className.replace('.', '/') + ".class"); - File dir = classFile.getParentFile(); - if (dir == null || !dir.isDirectory()) { - return; - } - int lastDot = className.lastIndexOf('.'); - String pkg = lastDot < 0 ? "" : className.substring(0, lastDot + 1); - String prefix = (lastDot < 0 ? className : className.substring(lastDot + 1)) + "$"; - File[] siblings = dir.listFiles(); - if (siblings == null) { - return; - } - for (File sibling : siblings) { - String fileName = sibling.getName(); - if (fileName.startsWith(prefix) && fileName.endsWith(".class")) { - helperClasses.add(pkg + fileName.substring(0, fileName.length() - ".class".length())); - } - } - } - - private static final String MUZZLE_REFERENCE_API = "datadog/trace/agent/tooling/muzzle/Reference"; - - /** - * {@code true} if the class uses the muzzle {@link Reference} API (as a {@link ReferenceProvider} - * or via {@code compileReferences}). This method is used to avoid injecting build-time-only - * classes. - */ - static boolean isBuildTimeOnly(String className, ClassFileLocator locator) { - try { - ClassFileLocator.Resolution resolution = locator.locate(className); - if (!resolution.isResolved()) { - return false; - } - // The muzzle type appears as a constant-pool entry when the class references it. - return new String(resolution.resolve(), StandardCharsets.ISO_8859_1) - .contains(MUZZLE_REFERENCE_API); - } catch (IOException e) { - return false; - } - } - - /** - * Expands the given helpers with any helper classes they depend on and returns them in - * dependency-first load order (required by {@link datadog.trace.agent.tooling.HelperInjector}) - * via {@link HelperScanner}. Library classes the scanner pulls in are dropped, but helpers that - * could not be located are kept (appended, unordered). - */ - private static String[] discoverAndOrderHelpers( - Set initialHelpers, HelperClassPredicate helperPredicate, ClassLoader loader) { - if (initialHelpers.isEmpty()) { - return new String[0]; - } - List ordered = new ArrayList<>(); - try { - for (String name : - HelperScanner.withClassDependencies( - ClassFileLocator.ForClassLoader.of(loader), initialHelpers.toArray(new String[0]))) { - if (helperPredicate.isHelperClass(name) && !ordered.contains(name)) { - ordered.add(name); - } - } - } catch (Throwable ignore) { - // best-effort ordering; unlocatable helpers are appended below - } - for (String helper : initialHelpers) { - if (!ordered.contains(helper)) { - ordered.add(helper); - } - } - return ordered.toArray(new String[0]); - } - private static void writeReference(MethodVisitor mv, Reference reference) { if (reference instanceof OrReference) { mv.visitTypeInsn(Opcodes.NEW, "datadog/trace/agent/tooling/muzzle/OrReference"); diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/MuzzleGeneratorTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/HelperResolverTest.java similarity index 84% rename from dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/MuzzleGeneratorTest.java rename to dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/HelperResolverTest.java index 7032425bcbc..be2b2aac296 100644 --- a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/MuzzleGeneratorTest.java +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/HelperResolverTest.java @@ -12,7 +12,6 @@ import datadog.trace.agent.tooling.muzzle.MuzzleGeneratorFixtures.ManualModule; import datadog.trace.agent.tooling.muzzle.MuzzleGeneratorFixtures.OwnerWithMuzzleFixture; import java.io.File; -import java.nio.file.Files; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -22,7 +21,7 @@ import net.bytebuddy.dynamic.ClassFileLocator; import org.junit.jupiter.api.Test; -class MuzzleGeneratorTest { +class HelperResolverTest { private static final String INFERRED = InferredHelperFixture.class.getName(); private static final String MANUAL = ManualHelperFixture.class.getName(); @@ -33,10 +32,10 @@ class MuzzleGeneratorTest { void isBuildTimeOnlyDetectsMuzzleReferenceProviders() { ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(getClass().getClassLoader()); // classes that use the muzzle Reference API are build-time only and must not be injected - assertTrue(MuzzleGenerator.isBuildTimeOnly(BuildTimeProviderFixture.class.getName(), locator)); - assertTrue(MuzzleGenerator.isBuildTimeOnly(MUZZLE_HELPER, locator)); + assertTrue(HelperResolver.isBuildTimeOnly(BuildTimeProviderFixture.class.getName(), locator)); + assertTrue(HelperResolver.isBuildTimeOnly(MUZZLE_HELPER, locator)); // ordinary helper is injectable - assertFalse(MuzzleGenerator.isBuildTimeOnly(INFERRED, locator)); + assertFalse(HelperResolver.isBuildTimeOnly(INFERRED, locator)); } @Test @@ -64,13 +63,12 @@ void inferredHelpersAreUsedWhenNoListDeclaredAndMuzzleProvidersDropped() throws } private static List injectedHelpers(InstrumenterModule module) throws Exception { - // Point the generator's "ownOutput" at this module's compiled test classes so the fixtures + // Point the resolver's "ownOutput" at this module's compiled test classes so the fixtures // count as this subproject's own helpers. File sourceDir = classesRootOf(MuzzleGeneratorFixtures.class); - File targetDir = Files.createTempDirectory("muzzle-generator-test").toFile(); - MuzzleGenerator generator = new MuzzleGenerator(sourceDir, targetDir); + HelperResolver resolver = new HelperResolver(sourceDir); - ClassLoader loader = MuzzleGeneratorTest.class.getClassLoader(); + ClassLoader loader = HelperResolverTest.class.getClassLoader(); Map crawled = ReferenceCreator.createReferencesFrom(CombineAdvice.class.getName(), loader); List references = new ArrayList<>(crawled.values()); @@ -80,7 +78,7 @@ private static List injectedHelpers(InstrumenterModule module) throws Ex // computeInjectedHelpers resolves classes via the context class-loader. Thread.currentThread().setContextClassLoader(loader); try { - return Arrays.asList(generator.computeInjectedHelpers(module, references, adviceClasses)); + return Arrays.asList(resolver.computeInjectedHelpers(module, references, adviceClasses)); } finally { Thread.currentThread().setContextClassLoader(previous); } diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/MuzzleGeneratorFixtures.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/MuzzleGeneratorFixtures.java index 662f9df230f..42e3229ec4a 100644 --- a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/MuzzleGeneratorFixtures.java +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/MuzzleGeneratorFixtures.java @@ -3,7 +3,7 @@ import java.util.Collections; import net.bytebuddy.pool.TypePool; -/** Fixtures for {@link MuzzleGeneratorTest}. */ +/** Fixtures for {@link HelperResolverTest}. */ final class MuzzleGeneratorFixtures { private MuzzleGeneratorFixtures() {}