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..3dd87f46251 --- /dev/null +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/HelperResolver.java @@ -0,0 +1,246 @@ +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.net.URISyntaxException; +import java.nio.charset.StandardCharsets; +import java.security.CodeSource; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Comparator; +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 HelperResolver() {} + + /** 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; + } + } + + /** Crawls the module's advice once and resolves both its references and injected helpers. */ + static Result resolve(InstrumenterModule module) { + File sourceRoot = sourceRootFor(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, sourceRoot)); + } + + 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. */ + static String[] computeInjectedHelpers( + InstrumenterModule module, + List allReferences, + Set adviceClasses, + File sourceRoot) { + // A module that declares its own helper list uses it directly. + String[] declaredHelpers = module.helperClassNames(); + if (declaredHelpers.length > 0) { + return declaredHelpers; + } + + // Otherwise infer them + HelperClassPredicate helperPredicate = + new HelperClassPredicate(name -> isOwnOutput(sourceRoot, name)); + 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(sourceRoot, helper)) { + addNestedClasses(sourceRoot, 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]); + } + + /** + * The subproject's compiled-output root, taken from the loaded module's code source (the + * classpath entry it was loaded from, i.e. the raw-classes folder). + */ + private static File sourceRootFor(InstrumenterModule module) { + CodeSource codeSource = module.getClass().getProtectionDomain().getCodeSource(); + if (codeSource == null || codeSource.getLocation() == null) { + throw new IllegalStateException( + "Cannot locate compiled output for " + module.getClass().getName()); + } + try { + return new File(codeSource.getLocation().toURI()); + } catch (URISyntaxException e) { + throw new IllegalStateException( + "Cannot resolve compiled output for " + codeSource.getLocation(), e); + } + } + + /** {@code true} if the class was compiled from this instrumentation subproject's own output. */ + private static boolean isOwnOutput(File sourceRoot, String className) { + return new File(sourceRoot, className.replace('.', '/') + ".class").isFile(); + } + + /** Adds the nested classes ({@code Foo$Bar}, {@code Foo$1}, ...) of an ownOutput helper. */ + private static void addNestedClasses( + File sourceRoot, String className, Set helperClasses) { + File classFile = new File(sourceRoot, 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; + } + // listFiles() order is filesystem-dependent; sort for reproducible helper ordering. + Arrays.sort(siblings, Comparator.comparing(File::getName)); + 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 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]); + } +} 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 59b00800a3c..616a16084f8 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,32 +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.net.URISyntaxException; -import java.nio.charset.StandardCharsets; import java.nio.file.Files; -import java.security.CodeSource; import java.util.ArrayList; -import java.util.Arrays; import java.util.Collections; -import java.util.Comparator; 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; @@ -78,29 +66,15 @@ public ClassVisitor wrap( throw new RuntimeException(e); } - File sourceRoot = sourceRootFor(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)); - } - } - - String[] orderedHelpers = - computeInjectedHelpers(module, allReferences, adviceClasses, sourceRoot); + // Resolve helpers from a single advice crawl, then emit both outputs below. + HelperResolver.Result resolved = HelperResolver.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); } @@ -110,42 +84,7 @@ public ClassVisitor wrap( return classVisitor; } // 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. */ @@ -256,144 +195,6 @@ public void visitEnd() { } } - /** Resolves the ordered set of helper classes to inject for a module. */ - String[] computeInjectedHelpers( - InstrumenterModule module, - List allReferences, - Set adviceClasses, - File sourceRoot) { - // A module that declares its own helper list uses it directly. - String[] declaredHelpers = module.helperClassNames(); - if (declaredHelpers.length > 0) { - return declaredHelpers; - } - - // Otherwise infer them - HelperClassPredicate helperPredicate = - new HelperClassPredicate(name -> isOwnOutput(sourceRoot, name)); - 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(sourceRoot, helper)) { - addNestedClasses(sourceRoot, 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]); - } - - /** - * The subproject's compiled-output root, taken from the loaded module's code source (the - * classpath entry it was loaded from, i.e. the raw-classes folder). - */ - private static File sourceRootFor(InstrumenterModule module) { - CodeSource codeSource = module.getClass().getProtectionDomain().getCodeSource(); - if (codeSource == null || codeSource.getLocation() == null) { - throw new IllegalStateException( - "Cannot locate compiled output for " + module.getClass().getName()); - } - try { - return new File(codeSource.getLocation().toURI()); - } catch (URISyntaxException e) { - throw new IllegalStateException( - "Cannot resolve compiled output for " + codeSource.getLocation(), e); - } - } - - /** {@code true} if the class was compiled from this instrumentation subproject's own output. */ - private static boolean isOwnOutput(File sourceRoot, String className) { - return new File(sourceRoot, className.replace('.', '/') + ".class").isFile(); - } - - /** Adds the nested classes ({@code Foo$Bar}, {@code Foo$1}, ...) of an ownOutput helper. */ - private static void addNestedClasses( - File sourceRoot, String className, Set helperClasses) { - File classFile = new File(sourceRoot, 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; - } - Arrays.sort(siblings, Comparator.comparing(File::getName)); - 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 82% 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 ed2bb3dd8a0..83d1d901671 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,11 @@ 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 - // count as this subproject's own helpers. - File sourceDir = classesRootOf(MuzzleGeneratorFixtures.class); - File targetDir = Files.createTempDirectory("muzzle-generator-test").toFile(); - MuzzleGenerator generator = new MuzzleGenerator(targetDir); + // Point "ownOutput" at this module's compiled test classes so the fixtures count as this + // subproject's own helpers. + File sourceRoot = classesRootOf(MuzzleGeneratorFixtures.class); - ClassLoader loader = MuzzleGeneratorTest.class.getClassLoader(); + ClassLoader loader = HelperResolverTest.class.getClassLoader(); Map crawled = ReferenceCreator.createReferencesFrom(CombineAdvice.class.getName(), loader); List references = new ArrayList<>(crawled.values()); @@ -81,7 +78,7 @@ private static List injectedHelpers(InstrumenterModule module) throws Ex Thread.currentThread().setContextClassLoader(loader); try { return Arrays.asList( - generator.computeInjectedHelpers(module, references, adviceClasses, sourceDir)); + HelperResolver.computeInjectedHelpers(module, references, adviceClasses, sourceRoot)); } 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() {}