diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java index 415d4886503..5a7e4b9c3df 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java @@ -130,12 +130,7 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati adviceShader = AdviceShader.with(module); - String[] helperClassNames = - InstrumenterModule.loadStaticMuzzleHelperClassNames( - Utils.getExtendedClassLoader(), module.getClass().getName()); - if (null == helperClassNames) { - helperClassNames = module.helperClassNames(); - } + String[] helperClassNames = module.helperClassNames(); if (module.injectHelperDependencies()) { helperClassNames = HelperScanner.withClassDependencies(helperClassNames); } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java index 833d573cfde..cd297fa65d8 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java @@ -113,23 +113,10 @@ public static ReferenceMatcher loadStaticMuzzleReferences( } /** - * @return the build-time inferred and manually-declared helper class names captured by {@code - * $Muzzle}, or {@code null} when none are available and fall back to {@link - * #helperClassNames()}. + * The helper classes to inject. Override this to declare them manually; otherwise {@code + * MuzzleGenerator} generates it at build time from the helpers inferred from the advice, so at + * runtime it returns every helper the module injects. */ - public 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) { - return null; - } - } - - /** Optional manual additions to the injected helper set. */ public String[] helperClassNames() { return NO_HELPERS; } 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 50cc27526de..d4fc8c508e3 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 @@ -77,14 +77,32 @@ 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); + File muzzleClass = new File(targetDir, moduleDefinition.getInternalName() + "$Muzzle.class"); try { muzzleClass.getParentFile().mkdirs(); - Files.write(muzzleClass.toPath(), generateMuzzleClass(module)); + Files.write(muzzleClass.toPath(), generateMuzzleClass(module, allReferences, orderedHelpers)); } catch (IOException e) { throw new RuntimeException(e); } - 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( @@ -123,23 +141,8 @@ private static Reference[] generateReferences( } /** This code is generated in a separate side-class. */ - private byte[] generateMuzzleClass(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)); - } - } - - String[] orderedHelpers = computeInjectedHelpers(module, allReferences, adviceClasses); + private byte[] generateMuzzleClass( + InstrumenterModule module, List allReferences, String[] orderedHelpers) { // Injected helpers are our own classes, so they don't need to be asserted as library // references. @@ -201,24 +204,48 @@ private byte[] generateMuzzleClass(InstrumenterModule module) { mv.visitMaxs(0, 0); mv.visitEnd(); - // Generate helperClassNames() with resolved helpers for the agent to read at load time; - // skip the method entirely when the module injects nothing. - 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(); + } + + /** + * 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. + */ + private static final class HelperClassNamesWriter extends ClassVisitor { + private static final String HELPER_METHOD = "helperClassNames"; + private static final String HELPER_DESCRIPTOR = "()[Ljava/lang/String;"; + + private final String[] helpers; + private boolean declared; + + HelperClassNamesWriter(ClassVisitor classVisitor, String[] helpers) { + super(Opcodes.ASM7, classVisitor); + this.helpers = helpers; } - return cw.toByteArray(); + @Override + public MethodVisitor visitMethod( + int access, String name, String descriptor, String signature, String[] exceptions) { + if (HELPER_METHOD.equals(name) && HELPER_DESCRIPTOR.equals(descriptor)) { + declared = true; + } + return super.visitMethod(access, name, descriptor, signature, exceptions); + } + + @Override + public void visitEnd() { + // Only generate when the module does not declare its own manually listed helpers. + if (!declared && helpers.length > 0) { + MethodVisitor mv = + super.visitMethod(Opcodes.ACC_PUBLIC, HELPER_METHOD, HELPER_DESCRIPTOR, null, null); + mv.visitCode(); + writeStrings(mv, helpers); + mv.visitInsn(Opcodes.ARETURN); + mv.visitMaxs(0, 0); + mv.visitEnd(); + } + super.visitEnd(); + } } /** Resolves the ordered set of helper classes to inject for a module. */ @@ -234,32 +261,37 @@ String[] computeInjectedHelpers( } } - // Add manually defined helpers. + // A module with a manually declared list uses it directly; otherwise the helpers + // inferred from its advice are used (with their nested classes, dependency-ordered, and any + // build-time-only muzzle providers dropped). Set manualHelpers = new LinkedHashSet<>(asList(module.helperClassNames())); - Set initialHelpers = new LinkedHashSet<>(inferredHelpers); - initialHelpers.addAll(manualHelpers); - for (String helper : new ArrayList<>(initialHelpers)) { - if (isOwnOutput(helper)) { - addNestedClasses(helper, initialHelpers); + String[] injectedHelpers; + if (!manualHelpers.isEmpty()) { + injectedHelpers = manualHelpers.toArray(new String[0]); + } else { + Set initialHelpers = new LinkedHashSet<>(inferredHelpers); + for (String helper : new ArrayList<>(initialHelpers)) { + if (isOwnOutput(helper)) { + addNestedClasses(helper, initialHelpers); + } } - } - - ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader(); - String[] orderedHelpers = - discoverAndOrderHelpers(initialHelpers, manualHelpers, helperPredicate, contextClassLoader); - - // Drop build-time-only muzzle providers. - ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(contextClassLoader); - List injectableHelpers = new ArrayList<>(orderedHelpers.length); - for (String helper : orderedHelpers) { - if (!isBuildTimeOnly(helper, locator)) { - injectableHelpers.add(helper); + ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader(); + String[] orderedHelpers = + discoverAndOrderHelpers( + initialHelpers, manualHelpers, helperPredicate, contextClassLoader); + // Drop build-time-only muzzle providers. + ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(contextClassLoader); + List injectableHelpers = new ArrayList<>(orderedHelpers.length); + for (String helper : orderedHelpers) { + if (!isBuildTimeOnly(helper, locator)) { + injectableHelpers.add(helper); + } } + injectedHelpers = injectableHelpers.toArray(new String[0]); } - orderedHelpers = injectableHelpers.toArray(new String[0]); - writeInferenceReport(module, adviceClasses.isEmpty(), inferredHelpers, orderedHelpers); - return orderedHelpers; + writeInferenceReport(module, adviceClasses.isEmpty(), inferredHelpers, injectedHelpers); + return injectedHelpers; } /** {@code true} if the class was compiled from this instrumentation subproject's own output. */ 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 2a979cb6c40..662f9df230f 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 @@ -42,11 +42,14 @@ static void apply() { } } - /** Module declaring a manual helper the advice crawl cannot see. */ - static class CombineModule extends TestInstrumentationClasses.BaseInst { + /** Module declaring a manual helper list. */ + static class ManualModule extends TestInstrumentationClasses.BaseInst { @Override public String[] helperClassNames() { return new String[] {ManualHelperFixture.class.getName()}; } } + + /** Module with no declared helper list, so its helpers come from auto-detection. */ + static class InferredModule extends TestInstrumentationClasses.BaseInst {} } 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/MuzzleGeneratorTest.java index e05d547a9e4..7032425bcbc 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/MuzzleGeneratorTest.java @@ -6,9 +6,10 @@ import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.muzzle.MuzzleGeneratorFixtures.BuildTimeProviderFixture; import datadog.trace.agent.tooling.muzzle.MuzzleGeneratorFixtures.CombineAdvice; -import datadog.trace.agent.tooling.muzzle.MuzzleGeneratorFixtures.CombineModule; import datadog.trace.agent.tooling.muzzle.MuzzleGeneratorFixtures.InferredHelperFixture; +import datadog.trace.agent.tooling.muzzle.MuzzleGeneratorFixtures.InferredModule; import datadog.trace.agent.tooling.muzzle.MuzzleGeneratorFixtures.ManualHelperFixture; +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; @@ -39,7 +40,30 @@ void isBuildTimeOnlyDetectsMuzzleReferenceProviders() { } @Test - void combinesInferredAndManualHelpersAndDropsMuzzleProviders() throws Exception { + void declaredHelperListIsUsedAsIsWithoutMergingInference() throws Exception { + List injected = injectedHelpers(new ManualModule()); + + // A module that declares helperClassNames() gets exactly that list - nothing inferred from the + // advice is merged in. + assertTrue(injected.contains(MANUAL), "manually declared helper should be injected"); + assertFalse(injected.contains(INFERRED), "inferred helper must not be merged into manual list"); + assertFalse(injected.contains(OWNER), "ownOutput helper must not be merged into manual list"); + } + + @Test + void inferredHelpersAreUsedWhenNoListDeclaredAndMuzzleProvidersDropped() throws Exception { + List injected = injectedHelpers(new InferredModule()); + + // If no declared list, the helpers inferred from the advice are injected minus the + // build-time-only muzzle provider and advice root. + assertTrue(injected.contains(INFERRED), "inferred helper should be injected"); + assertTrue(injected.contains(OWNER), "ownOutput helper should be injected"); + assertFalse(injected.contains(MUZZLE_HELPER), "build-time MuzzleHelper must not be injected"); + assertFalse(injected.contains(CombineAdvice.class.getName()), "advice root is not a helper"); + assertFalse(injected.contains(MANUAL), "helper the advice never references is not inferred"); + } + + 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); @@ -52,24 +76,9 @@ void combinesInferredAndManualHelpersAndDropsMuzzleProviders() throws Exception List references = new ArrayList<>(crawled.values()); Set adviceClasses = Collections.singleton(CombineAdvice.class.getName()); - List injected = - injectedHelpers(generator, new CombineModule(), references, adviceClasses); - - assertTrue(injected.contains(INFERRED), "inferred helper should be injected"); - assertTrue(injected.contains(MANUAL), "manually declared helper should be injected"); - assertTrue(injected.contains(OWNER), "ownOutput helper should be injected"); - assertFalse(injected.contains(MUZZLE_HELPER), "build-time MuzzleHelper must not be injected"); - assertFalse(injected.contains(CombineAdvice.class.getName()), "advice root is not a helper"); - } - - private static List injectedHelpers( - MuzzleGenerator generator, - InstrumenterModule module, - List references, - Set adviceClasses) { ClassLoader previous = Thread.currentThread().getContextClassLoader(); // computeInjectedHelpers resolves classes via the context class-loader. - Thread.currentThread().setContextClassLoader(MuzzleGeneratorTest.class.getClassLoader()); + Thread.currentThread().setContextClassLoader(loader); try { return Arrays.asList(generator.computeInjectedHelpers(module, references, adviceClasses)); } finally { diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletResponseInstrumentation.java b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletResponseInstrumentation.java index e708d76e79a..30736ff15e1 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletResponseInstrumentation.java @@ -42,6 +42,7 @@ public ElementMatcher hierarchyMatcher() { public String[] helperClassNames() { return new String[] { "datadog.trace.instrumentation.servlet.ServletRequestSetter", + "datadog.trace.instrumentation.servlet.http.HttpServletResponseDecorator", }; }