From 4b574ea4e245d6b3fa986f35834f440b952890f9 Mon Sep 17 00:00:00 2001 From: Sarah Chen Date: Tue, 1 Sep 2026 16:57:35 -0400 Subject: [PATCH 1/3] Implement generic AdviceProcessor --- dd-java-agent/agent-tooling/build.gradle | 1 + .../trace/agent/tooling/AdviceShader.java | 21 +- .../agent/tooling/InstrumenterModule.java | 2 +- .../agent/tooling/advice/AdviceProcessor.java | 9 + .../advice/AdviceProcessorContext.java | 53 ++ .../tooling/advice/AdviceScanResult.java | 195 ++++++ .../agent/tooling/advice/AdviceScanner.java | 461 +++++++++++++ .../advice/AdviceScanningGradlePlugin.java | 73 ++ .../muzzle/MuzzleGenerationProcessor.java | 43 ++ .../muzzle/MuzzleGenerationResult.java | 23 + .../agent/tooling/muzzle/MuzzleGenerator.java | 133 +--- .../tooling/muzzle/MuzzleGradlePlugin.java | 45 -- .../tooling/muzzle/ReferenceCreator.java | 626 ++++++------------ .../muzzle/MuzzleVersionScanPluginTest.groovy | 6 +- .../muzzle/ReferenceCreatorTest.groovy | 16 +- .../muzzle/ReferenceMatcherTest.groovy | 2 +- .../advice/AdviceProcessorPipelineTest.java | 74 +++ .../tooling/advice/AdviceScannerTest.java | 90 +++ .../advice/AdviceScanningFixtures.java | 76 +++ .../muzzle/MuzzleWeakReferenceTest.java | 4 +- .../ReferenceCreatorConversionTest.java | 66 ++ .../muzzle/ReferenceCreatorTestSupport.java | 39 ++ .../muzzle/TestInstrumentationClasses.java | 3 +- dd-java-agent/instrumentation/build.gradle | 2 +- 24 files changed, 1458 insertions(+), 605 deletions(-) create mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceProcessor.java create mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceProcessorContext.java create mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanResult.java create mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanner.java create mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanningGradlePlugin.java create mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationProcessor.java create mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationResult.java delete mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGradlePlugin.java create mode 100644 dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceProcessorPipelineTest.java create mode 100644 dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScannerTest.java create mode 100644 dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScanningFixtures.java create mode 100644 dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorConversionTest.java create mode 100644 dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorTestSupport.java diff --git a/dd-java-agent/agent-tooling/build.gradle b/dd-java-agent/agent-tooling/build.gradle index 06272944f25..965dcc3d5d9 100644 --- a/dd-java-agent/agent-tooling/build.gradle +++ b/dd-java-agent/agent-tooling/build.gradle @@ -50,6 +50,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' diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AdviceShader.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AdviceShader.java index 5443e9af7af..07b4426767d 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AdviceShader.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AdviceShader.java @@ -47,10 +47,29 @@ private AdviceShader(Map relocations, List helperNames) /** Applies shading before calling the given {@link ClassVisitor}. */ public ClassVisitor shadeClass(ClassVisitor cv) { + return new ClassRemapper(cv, remapper()); + } + + /** Applies advice relocation rules to a binary class name. */ + public String shadeClassName(String className) { + return remapper().mapType(className.replace('.', '/')).replace('/', '.'); + } + + /** Applies advice relocation rules to a field/type descriptor. */ + public String shadeTypeDescriptor(String descriptor) { + return remapper().mapDesc(descriptor); + } + + /** Applies advice relocation rules to a method descriptor. */ + public String shadeMethodDescriptor(String descriptor) { + return remapper().mapMethodDesc(descriptor); + } + + private Remapper remapper() { if (null == remapper) { remapper = new AdviceMapper(); } - return new ClassRemapper(cv, remapper); + return remapper; } /** Returns the result of shading the given bytecode. */ 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 d2abbc265e5..68fbe0affba 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 @@ -104,7 +104,7 @@ public static ReferenceMatcher loadStaticMuzzleReferences( String muzzleClass = instrumentationClass + "$Muzzle"; try { // Muzzle class contains static references captured at build-time - // see datadog.trace.agent.tooling.muzzle.MuzzleGenerator + // see datadog.trace.agent.tooling.muzzle.MuzzleGenerationProcessor return (ReferenceMatcher) classLoader.loadClass(muzzleClass).getMethod("create").invoke(null); } catch (Throwable e) { log.warn("Failed to load - muzzle.class={}", muzzleClass, e); diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceProcessor.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceProcessor.java new file mode 100644 index 00000000000..9e48c8b239a --- /dev/null +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceProcessor.java @@ -0,0 +1,9 @@ +package datadog.trace.agent.tooling.advice; + +/** Ordered build-time consumer of a reusable {@link AdviceScanResult}. */ +public interface AdviceProcessor { + /** Type used to publish this processor's output to later processors. */ + Class resultType(); + + T process(AdviceScanResult scanResult, AdviceProcessorContext context); +} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceProcessorContext.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceProcessorContext.java new file mode 100644 index 00000000000..b2d4ec125e8 --- /dev/null +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceProcessorContext.java @@ -0,0 +1,53 @@ +package datadog.trace.agent.tooling.advice; + +import datadog.trace.agent.tooling.InstrumenterModule; +import java.io.File; +import java.util.HashMap; +import java.util.Map; +import net.bytebuddy.dynamic.DynamicType; + +/** Mutable pipeline context shared by ordered advice processors for one module. */ +public final class AdviceProcessorContext { + private final InstrumenterModule module; + private final File targetDirectory; + private final Map, Object> results = new HashMap<>(); + private DynamicType.Builder builder; + + AdviceProcessorContext( + InstrumenterModule module, File targetDirectory, DynamicType.Builder builder) { + this.module = module; + this.targetDirectory = targetDirectory; + this.builder = builder; + } + + public InstrumenterModule getModule() { + return module; + } + + public File getTargetDirectory() { + return targetDirectory; + } + + public DynamicType.Builder getBuilder() { + return builder; + } + + public void setBuilder(DynamicType.Builder builder) { + this.builder = builder; + } + + public T getResult(Class resultType) { + Object result = results.get(resultType); + if (result == null) { + throw new IllegalStateException("No advice processor result of type " + resultType.getName()); + } + return resultType.cast(result); + } + + void putResult(Class resultType, T result) { + if (results.put(resultType, result) != null) { + throw new IllegalStateException( + "Duplicate advice processor result type " + resultType.getName()); + } + } +} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanResult.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanResult.java new file mode 100644 index 00000000000..5b2179c0434 --- /dev/null +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanResult.java @@ -0,0 +1,195 @@ +package datadog.trace.agent.tooling.advice; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** Immutable, neutral result of scanning a module's method advice. */ +public final class AdviceScanResult { + public enum UsageKind { + TYPE, + FIELD, + METHOD, + HANDLE, + INVOKEDYNAMIC + } + + /** Source location of a bytecode use. */ + public static final class SourceLocation { + private final String className; + private final int line; + + SourceLocation(String className, int line) { + this.className = className; + this.line = line; + } + + public String getClassName() { + return className; + } + + public int getLine() { + return line; + } + } + + /** A method-handle target without an ASM dependency in the public model. */ + public static final class HandleUse { + private final int tag; + private final String owner; + private final String name; + private final String descriptor; + private final boolean interfaceOwner; + + HandleUse(int tag, String owner, String name, String descriptor, boolean interfaceOwner) { + this.tag = tag; + this.owner = owner; + this.name = name; + this.descriptor = descriptor; + this.interfaceOwner = interfaceOwner; + } + + public int getTag() { + return tag; + } + + public String getOwner() { + return owner; + } + + public String getName() { + return name; + } + + public String getDescriptor() { + return descriptor; + } + + public boolean isInterfaceOwner() { + return interfaceOwner; + } + } + + /** One ordered bytecode use made by a scanned class. */ + public static final class Usage { + private final UsageKind kind; + private final SourceLocation source; + private final int opcode; + private final String owner; + private final String name; + private final String descriptor; + private final boolean interfaceOwner; + private final boolean implementedInterface; + private final List handles; + + Usage( + UsageKind kind, + SourceLocation source, + int opcode, + String owner, + String name, + String descriptor, + boolean interfaceOwner, + boolean implementedInterface, + List handles) { + this.kind = kind; + this.source = source; + this.opcode = opcode; + this.owner = owner; + this.name = name; + this.descriptor = descriptor; + this.interfaceOwner = interfaceOwner; + this.implementedInterface = implementedInterface; + this.handles = immutableCopy(handles); + } + + public UsageKind getKind() { + return kind; + } + + public SourceLocation getSource() { + return source; + } + + public int getOpcode() { + return opcode; + } + + /** Binary name of the used type or member owner. */ + public String getOwner() { + return owner; + } + + public String getName() { + return name; + } + + public String getDescriptor() { + return descriptor; + } + + public boolean isInterfaceOwner() { + return interfaceOwner; + } + + public boolean isImplementedInterface() { + return implementedInterface; + } + + public List getHandles() { + return handles; + } + } + + /** One discovered class. External leaves have no scanned bytecode. */ + public static final class ClassInfo { + private final String className; + private final boolean scanned; + private final List usages; + + ClassInfo(String className, boolean scanned, List usages) { + this.className = className; + this.scanned = scanned; + this.usages = immutableCopy(usages); + } + + public String getClassName() { + return className; + } + + public boolean isScanned() { + return scanned; + } + + public List getUsages() { + return usages; + } + } + + private final List adviceRoots; + private final Map classes; + + AdviceScanResult(Collection adviceRoots, Map classes) { + this.adviceRoots = immutableCopy(adviceRoots); + this.classes = Collections.unmodifiableMap(new LinkedHashMap<>(classes)); + } + + public List getAdviceRoots() { + return adviceRoots; + } + + public Map getClasses() { + return classes; + } + + public ClassInfo getClassInfo(String className) { + return classes.get(className); + } + + private static List immutableCopy(Collection values) { + return Collections.unmodifiableList(new ArrayList<>(values)); + } +} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanner.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanner.java new file mode 100644 index 00000000000..1e0701b226c --- /dev/null +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanner.java @@ -0,0 +1,461 @@ +package datadog.trace.agent.tooling.advice; + +import static java.util.Collections.addAll; +import static java.util.Collections.emptyList; +import static java.util.Collections.singletonList; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.advice.AdviceScanResult.ClassInfo; +import datadog.trace.agent.tooling.advice.AdviceScanResult.HandleUse; +import datadog.trace.agent.tooling.advice.AdviceScanResult.SourceLocation; +import datadog.trace.agent.tooling.advice.AdviceScanResult.Usage; +import datadog.trace.agent.tooling.advice.AdviceScanResult.UsageKind; +import de.thetaphi.forbiddenapis.SuppressForbidden; +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.net.URISyntaxException; +import java.security.CodeSource; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import net.bytebuddy.jar.asm.ClassReader; +import net.bytebuddy.jar.asm.ClassVisitor; +import net.bytebuddy.jar.asm.Handle; +import net.bytebuddy.jar.asm.Label; +import net.bytebuddy.jar.asm.MethodVisitor; +import net.bytebuddy.jar.asm.Opcodes; +import net.bytebuddy.jar.asm.Type; +import net.bytebuddy.utility.StreamDrainer; + +/** Scans all advice and reachable module bytecode for one instrumenter module. */ +public final class AdviceScanner { + private static final int UNDEFINED_LINE = -1; + + private final InstrumenterModule module; + private final File ownOutput; + private final ClassLoader loader; + private final LinkedHashSet adviceRoots = new LinkedHashSet<>(); + private final LinkedHashMap classes = new LinkedHashMap<>(); + private final Deque scanQueue = new ArrayDeque<>(); + private final Set queued = new LinkedHashSet<>(); + + private AdviceScanner(InstrumenterModule module, File ownOutput, ClassLoader loader) { + this.module = module; + this.ownOutput = ownOutput; + this.loader = loader; + } + + public static AdviceScanResult scan(InstrumenterModule module) { + return scan(module, sourceRootFor(module), Thread.currentThread().getContextClassLoader()); + } + + static AdviceScanResult scan(InstrumenterModule module, File ownOutput, ClassLoader loader) { + return new AdviceScanner(module, ownOutput, loader).scan(); + } + + private AdviceScanResult scan() { + collectAdviceRoots(); + for (String adviceRoot : adviceRoots) { + MutableClassInfo info = discover(adviceRoot, adviceRoot); + enqueue(info, true); + } + + String className; + while ((className = scanQueue.pollFirst()) != null) { + MutableClassInfo info = classes.get(className); + if (info.scanned) { + continue; + } + byte[] bytecode = readClass(info); + try { + new ClassReader(bytecode).accept(new ScanningVisitor(info), ClassReader.SKIP_FRAMES); + info.scanned = true; + } catch (Throwable error) { + throw scanFailure(className, owningAdvice(className), "cannot parse class bytecode", error); + } + } + + Map immutableClasses = new LinkedHashMap<>(); + for (Map.Entry entry : classes.entrySet()) { + immutableClasses.put(entry.getKey(), entry.getValue().freeze()); + } + return new AdviceScanResult(adviceRoots, immutableClasses); + } + + private void collectAdviceRoots() { + List instrumenters = module.typeInstrumentations(); + for (Instrumenter instrumenter : instrumenters) { + if (instrumenter instanceof Instrumenter.HasMethodAdvice) { + ((Instrumenter.HasMethodAdvice) instrumenter) + .methodAdvice( + (matcher, adviceClass, additionalClasses) -> { + adviceRoots.add(adviceClass); + if (additionalClasses != null) { + addAll(adviceRoots, additionalClasses); + } + }); + } + } + } + + private MutableClassInfo discover(String className, String adviceClass) { + if (className == null) { + return null; + } + MutableClassInfo existing = classes.get(className); + if (existing != null) { + if (adviceClass != null && existing.adviceClass == null) { + existing.adviceClass = adviceClass; + } + return existing; + } + MutableClassInfo created = new MutableClassInfo(className, isOwnOutput(className), adviceClass); + classes.put(className, created); + return created; + } + + private void enqueue(MutableClassInfo info, boolean adviceRoot) { + if (info != null && !info.scanned && (adviceRoot || info.owned) && queued.add(info.className)) { + scanQueue.addLast(info.className); + } + } + + private void addDependency(MutableClassInfo from, String className) { + if (className == null || className.equals(from.className)) { + return; + } + MutableClassInfo target = discover(className, from.adviceClass); + enqueue(target, false); + } + + private void addTypeDependency(MutableClassInfo from, Type type) { + if (type == null) { + return; + } + while (type.getSort() == Type.ARRAY) { + type = type.getElementType(); + } + if (type.getSort() == Type.METHOD) { + for (Type argument : type.getArgumentTypes()) { + addTypeDependency(from, argument); + } + addTypeDependency(from, type.getReturnType()); + } else if (type.getSort() == Type.OBJECT) { + addDependency(from, type.getClassName()); + } + } + + private void addHandleDependencies(MutableClassInfo from, Handle handle) { + addDependency(from, binaryName(handle.getOwner())); + } + + private byte[] readClass(MutableClassInfo info) { + try { + if (info.owned) { + File classFile = classFile(info.className); + if (!classFile.isFile()) { + throw scanFailure( + info.className, owningAdvice(info.className), "owned class file is missing", null); + } + return java.nio.file.Files.readAllBytes(classFile.toPath()); + } + String resource = info.className.replace('.', '/') + ".class"; + try (InputStream input = loader.getResourceAsStream(resource)) { + if (input == null) { + throw scanFailure( + info.className, owningAdvice(info.className), "advice class is missing", null); + } + return StreamDrainer.DEFAULT.drain(input); + } + } catch (IOException error) { + throw scanFailure( + info.className, owningAdvice(info.className), "cannot read class bytecode", error); + } + } + + private boolean isOwnOutput(String className) { + return classFile(className).isFile(); + } + + private File classFile(String className) { + return new File(ownOutput, className.replace('.', File.separatorChar) + ".class"); + } + + private String owningAdvice(String className) { + MutableClassInfo info = classes.get(className); + return info == null || info.adviceClass == null ? "" : info.adviceClass; + } + + private IllegalStateException scanFailure( + String className, String adviceClass, String detail, Throwable cause) { + String message = + "Advice scan failed for module " + + module.getClass().getName() + + ", advice " + + adviceClass + + ", class " + + className + + ": " + + detail; + return cause == null + ? new IllegalStateException(message) + : new IllegalStateException(message, cause); + } + + @SuppressForbidden + 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 " + module.getClass().getName()); + } + try { + return new File(codeSource.getLocation().toURI()); + } catch (URISyntaxException error) { + throw new IllegalStateException( + "Cannot resolve compiled output for module " + module.getClass().getName(), error); + } + } + + private static String binaryName(String internalName) { + return internalName.replace('/', '.'); + } + + private static Type underlyingType(Type type) { + while (type.getSort() == Type.ARRAY) { + type = type.getElementType(); + } + return type; + } + + private final class ScanningVisitor extends ClassVisitor { + private final MutableClassInfo info; + + private ScanningVisitor(MutableClassInfo info) { + super(Opcodes.ASM7); + this.info = info; + } + + @Override + public void visit( + int version, + int access, + String name, + String signature, + String superName, + String[] interfaces) { + if (interfaces != null) { + for (String interfaceName : interfaces) { + String binaryInterface = binaryName(interfaceName); + addDependency(info, binaryInterface); + info.usages.add( + usage( + UsageKind.TYPE, + UNDEFINED_LINE, + -1, + binaryInterface, + null, + null, + true, + true, + emptyList())); + } + } + } + + @Override + public MethodVisitor visitMethod( + int access, String name, String descriptor, String signature, String[] exceptions) { + return new ScanningMethodVisitor(info); + } + + private Usage usage( + UsageKind kind, + int line, + int opcode, + String owner, + String name, + String descriptor, + boolean interfaceOwner, + boolean implementedInterface, + List handles) { + return new Usage( + kind, + new SourceLocation(info.className, line), + opcode, + owner, + name, + descriptor, + interfaceOwner, + implementedInterface, + handles); + } + } + + private final class ScanningMethodVisitor extends MethodVisitor { + private final MutableClassInfo info; + private int line = UNDEFINED_LINE; + + private ScanningMethodVisitor(MutableClassInfo info) { + super(Opcodes.ASM7); + this.info = info; + } + + @Override + public void visitLineNumber(int line, Label start) { + this.line = line; + } + + @Override + public void visitFieldInsn(int opcode, String owner, String name, String descriptor) { + String binaryOwner = binaryName(owner); + addDependency(info, binaryOwner); + addTypeDependency(info, Type.getType(descriptor)); + info.usages.add( + usage(UsageKind.FIELD, opcode, binaryOwner, name, descriptor, false, emptyList())); + } + + @Override + public void visitMethodInsn( + int opcode, String owner, String name, String descriptor, boolean isInterface) { + String binaryOwner = binaryName(owner); + addDependency(info, binaryOwner); + addTypeDependency(info, Type.getMethodType(descriptor)); + info.usages.add( + usage(UsageKind.METHOD, opcode, binaryOwner, name, descriptor, isInterface, emptyList())); + } + + @Override + public void visitTypeInsn(int opcode, String typeName) { + Type type = underlyingType(Type.getObjectType(typeName)); + if (type.getSort() == Type.OBJECT) { + String binaryType = type.getClassName(); + addDependency(info, binaryType); + info.usages.add(usage(UsageKind.TYPE, opcode, binaryType, null, null, false, emptyList())); + } + } + + @Override + public void visitMultiANewArrayInsn(String descriptor, int dimensions) { + Type type = underlyingType(Type.getType(descriptor)); + if (type.getSort() == Type.OBJECT) { + String binaryType = type.getClassName(); + addDependency(info, binaryType); + info.usages.add( + usage( + UsageKind.TYPE, + Opcodes.MULTIANEWARRAY, + binaryType, + null, + descriptor, + false, + emptyList())); + } + } + + @Override + public void visitInvokeDynamicInsn( + String name, String descriptor, Handle bootstrapMethodHandle, Object... arguments) { + List handles = new ArrayList<>(); + addHandle(info, handles, bootstrapMethodHandle); + for (Object argument : arguments) { + if (argument instanceof Handle) { + addHandle(info, handles, (Handle) argument); + } + } + info.usages.add( + usage( + UsageKind.INVOKEDYNAMIC, + Opcodes.INVOKEDYNAMIC, + null, + name, + descriptor, + false, + handles)); + } + + @Override + public void visitLdcInsn(Object value) { + if (value instanceof Type) { + Type type = underlyingType((Type) value); + if (type.getSort() == Type.OBJECT) { + String binaryType = type.getClassName(); + addDependency(info, binaryType); + info.usages.add( + usage(UsageKind.TYPE, Opcodes.LDC, binaryType, null, null, false, emptyList())); + } + } else if (value instanceof Handle) { + Handle handle = (Handle) value; + addHandleDependencies(info, handle); + info.usages.add( + usage( + UsageKind.HANDLE, + Opcodes.LDC, + binaryName(handle.getOwner()), + handle.getName(), + handle.getDesc(), + handle.isInterface(), + singletonList(toHandleUse(handle)))); + } + } + + private void addHandle(MutableClassInfo info, List handles, Handle handle) { + addHandleDependencies(info, handle); + handles.add(toHandleUse(handle)); + } + + private Usage usage( + UsageKind kind, + int opcode, + String owner, + String name, + String descriptor, + boolean interfaceOwner, + List handles) { + return new Usage( + kind, + new SourceLocation(info.className, line), + opcode, + owner, + name, + descriptor, + interfaceOwner, + false, + handles); + } + } + + private static HandleUse toHandleUse(Handle handle) { + return new HandleUse( + handle.getTag(), + binaryName(handle.getOwner()), + handle.getName(), + handle.getDesc(), + handle.isInterface()); + } + + private static final class MutableClassInfo { + private final String className; + private final boolean owned; + private String adviceClass; + private boolean scanned; + private final List usages = new ArrayList<>(); + + private MutableClassInfo(String className, boolean owned, String adviceClass) { + this.className = className; + this.owned = owned; + this.adviceClass = adviceClass; + } + + private ClassInfo freeze() { + return new ClassInfo(className, scanned, usages); + } + } +} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanningGradlePlugin.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanningGradlePlugin.java new file mode 100644 index 00000000000..100296b6896 --- /dev/null +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanningGradlePlugin.java @@ -0,0 +1,73 @@ +package datadog.trace.agent.tooling.advice; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.concreteClass; +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static java.util.Collections.singletonList; + +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.bytebuddy.SharedTypePools; +import datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers; +import datadog.trace.agent.tooling.muzzle.MuzzleGenerationProcessor; +import java.io.File; +import java.io.IOException; +import java.util.List; +import net.bytebuddy.build.Plugin; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.dynamic.ClassFileLocator; +import net.bytebuddy.dynamic.DynamicType; + +/** Build-time plugin that scans each module once and runs ordered scan consumers. */ +public class AdviceScanningGradlePlugin extends Plugin.ForElementMatcher { + static { + SharedTypePools.registerIfAbsent(SharedTypePools.simpleCache()); + HierarchyMatchers.registerIfAbsent(HierarchyMatchers.simpleChecks()); + } + + private final File targetDirectory; + private final List> processors; + + public AdviceScanningGradlePlugin(File targetDirectory) { + this(targetDirectory, singletonList(new MuzzleGenerationProcessor())); + } + + AdviceScanningGradlePlugin(File targetDirectory, List> processors) { + super(concreteClass().and(extendsClass(named(InstrumenterModule.class.getName())))); + this.targetDirectory = targetDirectory; + this.processors = processors; + } + + @Override + public DynamicType.Builder apply( + DynamicType.Builder builder, + TypeDescription typeDescription, + ClassFileLocator classFileLocator) { + ClassLoader loader = Thread.currentThread().getContextClassLoader(); + InstrumenterModule module; + try { + // The module instance is shared by scanning and every processor. + module = + (InstrumenterModule) + loader.loadClass(typeDescription.getName()).getConstructor().newInstance(); + } catch (ReflectiveOperationException error) { + throw new IllegalStateException( + "Cannot instantiate instrumenter module " + typeDescription.getName(), error); + } + + AdviceScanResult scanResult = AdviceScanner.scan(module); + AdviceProcessorContext context = new AdviceProcessorContext(module, targetDirectory, builder); + for (AdviceProcessor processor : processors) { + runProcessor(processor, scanResult, context); + } + return context.getBuilder(); + } + + private static void runProcessor( + AdviceProcessor processor, AdviceScanResult scanResult, AdviceProcessorContext context) { + T result = processor.process(scanResult, context); + context.putResult(processor.resultType(), result); + } + + @Override + public void close() throws IOException {} +} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationProcessor.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationProcessor.java new file mode 100644 index 00000000000..cf6d93baa56 --- /dev/null +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationProcessor.java @@ -0,0 +1,43 @@ +package datadog.trace.agent.tooling.muzzle; + +import static java.util.Arrays.asList; +import static java.util.Collections.addAll; + +import datadog.trace.agent.tooling.AdviceShader; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.advice.AdviceProcessor; +import datadog.trace.agent.tooling.advice.AdviceProcessorContext; +import datadog.trace.agent.tooling.advice.AdviceScanResult; +import java.io.File; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +/** + * Converts a neutral advice scan to references and emits the existing {@code $Muzzle} side class. + */ +public final class MuzzleGenerationProcessor implements AdviceProcessor { + @Override + public Class resultType() { + return MuzzleGenerationResult.class; + } + + @Override + public MuzzleGenerationResult process( + AdviceScanResult scanResult, AdviceProcessorContext context) { + InstrumenterModule module = context.getModule(); + + Set ignoredClasses = new HashSet<>(asList(module.muzzleIgnoredClassNames())); + AdviceShader shader = AdviceShader.with(module.adviceShading()); + List references = + ReferenceCreator.createReferences(scanResult, scanResult.getAdviceRoots(), shader); + references.removeIf(reference -> ignoredClasses.contains(reference.className)); + Reference[] additionalReferences = module.additionalMuzzleReferences(); + if (additionalReferences != null) { + addAll(references, additionalReferences); + } + + File muzzleClass = MuzzleGenerator.generate(context.getTargetDirectory(), module, references); + return new MuzzleGenerationResult(references, muzzleClass); + } +} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationResult.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationResult.java new file mode 100644 index 00000000000..8fdca7c3f5c --- /dev/null +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationResult.java @@ -0,0 +1,23 @@ +package datadog.trace.agent.tooling.muzzle; + +import java.io.File; +import java.util.List; + +/** Typed muzzle processor output available to later advice processors. */ +public final class MuzzleGenerationResult { + private final Reference[] references; + private final File generatedClass; + + MuzzleGenerationResult(List references, File generatedClass) { + this.references = references.toArray(new Reference[0]); + this.generatedClass = generatedClass; + } + + public Reference[] getReferences() { + return references.clone(); + } + + public File getGeneratedClass() { + return generatedClass; + } +} 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 69421f6d16a..9c0d288b1d8 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 @@ -1,139 +1,34 @@ package datadog.trace.agent.tooling.muzzle; -import static java.util.Arrays.asList; - -import datadog.trace.agent.tooling.AdviceShader; -import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import java.io.File; import java.io.IOException; import java.nio.file.Files; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashSet; -import java.util.LinkedHashMap; 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.implementation.Implementation; -import net.bytebuddy.jar.asm.ClassVisitor; import net.bytebuddy.jar.asm.ClassWriter; import net.bytebuddy.jar.asm.MethodVisitor; import net.bytebuddy.jar.asm.Opcodes; import net.bytebuddy.jar.asm.Type; -import net.bytebuddy.pool.TypePool; - -/** Generates a 'Muzzle' side-class for each {@link InstrumenterModule}. */ -public class MuzzleGenerator implements AsmVisitorWrapper { - private final File targetDir; - - public MuzzleGenerator(File targetDir) { - this.targetDir = targetDir; - } - - @Override - public int mergeWriter(int flags) { - return flags | ClassWriter.COMPUTE_MAXS; - } - @Override - public int mergeReader(int flags) { - return flags; - } - - @Override - public ClassVisitor wrap( - final TypeDescription moduleDefinition, - final ClassVisitor classVisitor, - final Implementation.Context implementationContext, - final TypePool typePool, - final FieldList fields, - final MethodList methods, - final int writerFlags, - final int readerFlags) { - - InstrumenterModule module; - try { - module = - (InstrumenterModule) - Thread.currentThread() - .getContextClassLoader() - .loadClass(moduleDefinition.getName()) - .getConstructor() - .newInstance(); - } catch (ReflectiveOperationException e) { - throw new RuntimeException(e); - } +/** Generates a {@code $Muzzle} side class from resolved references. */ +final class MuzzleGenerator { + private MuzzleGenerator() {} - File muzzleClass = new File(targetDir, moduleDefinition.getInternalName() + "$Muzzle.class"); + static File generate( + File targetDirectory, InstrumenterModule module, List references) { + File muzzleClass = + new File(targetDirectory, Type.getInternalName(module.getClass()) + "$Muzzle.class"); try { muzzleClass.getParentFile().mkdirs(); - Files.write(muzzleClass.toPath(), generateMuzzleClass(module)); - } catch (IOException e) { - throw new RuntimeException(e); + Files.write(muzzleClass.toPath(), generateMuzzleClass(module, references)); + } catch (IOException error) { + throw new IllegalStateException( + "Cannot write muzzle class for " + module.getClass().getName(), error); } - return classVisitor; + return muzzleClass; } - private static Reference[] generateReferences( - Instrumenter.HasMethodAdvice instrumenter, AdviceShader adviceShader) { - // 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)); - } - }); - 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]); - } - - /** This code is generated in a separate side-class. */ - private static byte[] generateMuzzleClass(InstrumenterModule module) { - - Set ignoredClassNames = new HashSet<>(asList(module.muzzleIgnoredClassNames())); - AdviceShader adviceShader = AdviceShader.with(module.adviceShading()); - - List references = new ArrayList<>(); - for (Instrumenter instrumenter : module.typeInstrumentations()) { - if (instrumenter instanceof Instrumenter.HasMethodAdvice) { - for (Reference reference : - generateReferences((Instrumenter.HasMethodAdvice) instrumenter, adviceShader)) { - // ignore helper classes, they will be injected by the instrumentation's HelperInjector. - if (!ignoredClassNames.contains(reference.className)) { - references.add(reference); - } - } - } - } - Reference[] additionalReferences = module.additionalMuzzleReferences(); - if (null != additionalReferences) { - Collections.addAll(references, additionalReferences); - } - + private static byte[] generateMuzzleClass(InstrumenterModule module, List references) { ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES); cw.visit( Opcodes.V1_8, @@ -150,7 +45,6 @@ private static byte[] generateMuzzleClass(InstrumenterModule module) { "()Ldatadog/trace/agent/tooling/muzzle/ReferenceMatcher;", null, null); - mv.visitCode(); mv.visitTypeInsn(Opcodes.NEW, "datadog/trace/agent/tooling/muzzle/ReferenceMatcher"); @@ -173,7 +67,6 @@ private static byte[] generateMuzzleClass(InstrumenterModule module) { "", "([Ldatadog/trace/agent/tooling/muzzle/Reference;)V", false); - mv.visitInsn(Opcodes.ARETURN); mv.visitMaxs(0, 0); diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGradlePlugin.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGradlePlugin.java deleted file mode 100644 index 5a22ac484be..00000000000 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGradlePlugin.java +++ /dev/null @@ -1,45 +0,0 @@ -package datadog.trace.agent.tooling.muzzle; - -import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.concreteClass; -import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass; -import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; - -import datadog.trace.agent.tooling.InstrumenterModule; -import datadog.trace.agent.tooling.bytebuddy.SharedTypePools; -import datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers; -import java.io.File; -import java.io.IOException; -import net.bytebuddy.build.Plugin; -import net.bytebuddy.description.type.TypeDescription; -import net.bytebuddy.dynamic.ClassFileLocator; -import net.bytebuddy.dynamic.DynamicType; - -/** - * Byte-Buddy gradle plugin which creates muzzle-references at compile time. - * - * @see datadog.gradle.plugin.instrument.BuildTimeInstrumentationPlugin - */ -public class MuzzleGradlePlugin extends Plugin.ForElementMatcher { - static { - SharedTypePools.registerIfAbsent(SharedTypePools.simpleCache()); - HierarchyMatchers.registerIfAbsent(HierarchyMatchers.simpleChecks()); - } - - private final File targetDir; - - public MuzzleGradlePlugin(File targetDir) { - super(concreteClass().and(extendsClass(named(InstrumenterModule.class.getName())))); - this.targetDir = targetDir; - } - - @Override - public DynamicType.Builder apply( - final DynamicType.Builder builder, - final TypeDescription typeDescription, - final ClassFileLocator classFileLocator) { - return builder.visit(new MuzzleGenerator(targetDir)); - } - - @Override - public void close() throws IOException {} -} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java index 5fdc4d3f827..572683ea8bc 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java @@ -1,496 +1,272 @@ package datadog.trace.agent.tooling.muzzle; -import static datadog.trace.util.Strings.getClassName; -import static datadog.trace.util.Strings.getResourceName; - import datadog.trace.agent.tooling.AdviceShader; +import datadog.trace.agent.tooling.advice.AdviceScanResult; +import datadog.trace.agent.tooling.advice.AdviceScanResult.ClassInfo; +import datadog.trace.agent.tooling.advice.AdviceScanResult.HandleUse; +import datadog.trace.agent.tooling.advice.AdviceScanResult.SourceLocation; +import datadog.trace.agent.tooling.advice.AdviceScanResult.Usage; import datadog.trace.bootstrap.Constants; -import de.thetaphi.forbiddenapis.SuppressForbidden; -import java.io.InputStream; import java.lang.reflect.Method; import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collection; import java.util.HashSet; import java.util.LinkedHashMap; +import java.util.List; import java.util.Map; import java.util.Queue; import java.util.Set; -import net.bytebuddy.jar.asm.ClassReader; -import net.bytebuddy.jar.asm.ClassVisitor; -import net.bytebuddy.jar.asm.FieldVisitor; -import net.bytebuddy.jar.asm.Handle; -import net.bytebuddy.jar.asm.Label; -import net.bytebuddy.jar.asm.MethodVisitor; import net.bytebuddy.jar.asm.Opcodes; import net.bytebuddy.jar.asm.Type; -/** Visit a class and collect all references made by the visited class. */ -// Additional things we could check -// - annotations on class -// - outer class -// - inner class -// - cast opcodes in method bodies -public class ReferenceCreator extends ClassVisitor { - /** - * Classes in this namespace will be scanned and used to create references. - * - *

For now we're hardcoding this to the instrumentation package so we only create references - * from the method advice and helper classes. - */ +/** Converts neutral advice uses into the existing member-level muzzle reference model. */ +public final class ReferenceCreator { private static final String REFERENCE_CREATION_PACKAGE = "datadog.trace.instrumentation."; - private static final int UNDEFINED_LINE = -1; - - /** Set containing name+descriptor signatures of Object methods. */ private static final Set OBJECT_METHODS = new HashSet<>(); static { - for (Method m : Object.class.getMethods()) { - OBJECT_METHODS.add(methodSig(m.getName(), Type.getMethodDescriptor(m))); + for (Method method : Object.class.getMethods()) { + OBJECT_METHODS.add(methodSig(method.getName(), Type.getMethodDescriptor(method))); } } - /** - * Generate all references reachable from a given class. - * - * @param entryPointClassName Starting point for generating references. - * @param adviceShader Optional shading to apply to the advice. - * @param loader Classloader used to read class bytes. - * @return Map of [referenceClassName -> Reference] - * @throws IllegalStateException if class is not found or unable to be loaded. - */ - @SuppressForbidden - public static Map createReferencesFrom( - final String entryPointClassName, final AdviceShader adviceShader, final ClassLoader loader) - throws IllegalStateException { - final Set visitedSources = new HashSet<>(); - final Map references = new LinkedHashMap<>(); + private final AdviceScanResult scanResult; + private final AdviceShader shader; + private final Map references = new LinkedHashMap<>(); + private final Queue sources = new ArrayDeque<>(); + private final Set visitedSources = new HashSet<>(); - final Queue instrumentationQueue = new ArrayDeque<>(); - instrumentationQueue.add(entryPointClassName); + private ReferenceCreator(AdviceScanResult scanResult, AdviceShader shader) { + this.scanResult = scanResult; + this.shader = shader; + } - while (!instrumentationQueue.isEmpty()) { - final String className = instrumentationQueue.remove(); - visitedSources.add(className); - String resourceName = getResourceName(className); - try (InputStream in = loader.getResourceAsStream(resourceName)) { - if (null == in) { - System.err.println(resourceName + " not found, skipping"); - continue; - } - final ReferenceCreator cv = new ReferenceCreator(null); - final ClassReader reader = new ClassReader(in); - if (null == adviceShader) { - reader.accept(cv, ClassReader.SKIP_FRAMES); - } else { - reader.accept(adviceShader.shadeClass(cv), ClassReader.SKIP_FRAMES); - } + static List createReferences( + AdviceScanResult scanResult, Collection sourceClasses, AdviceShader shader) { + ReferenceCreator creator = new ReferenceCreator(scanResult, shader); + creator.sources.addAll(sourceClasses); + creator.createReferences(); + return new ArrayList<>(creator.references.values()); + } - final Map instrumentationReferences = cv.getReferences(); - for (final Map.Entry entry : instrumentationReferences.entrySet()) { - // 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()); - if (null == toMerge) { - references.put(entry.getKey(), entry.getValue()); - } else { - references.put(entry.getKey(), toMerge.merge(entry.getValue())); - } + private void createReferences() { + String sourceClass; + while ((sourceClass = sources.poll()) != null) { + if (!visitedSources.add(sourceClass)) { + continue; + } + ClassInfo info = scanResult.getClassInfo(sourceClass); + if (info != null && info.isScanned()) { + for (Usage usage : info.getUsages()) { + add(usage); } - - } catch (final Throwable t) { - throw new IllegalStateException("Error reading class " + className, t); } } - return references; } - public static Map createReferencesFrom( - final String entryPointClassName, final ClassLoader loader) { - return createReferencesFrom(entryPointClassName, null, loader); + private void add(Usage usage) { + switch (usage.getKind()) { + case TYPE: + // ReferenceCreator did not visit MULTIANEWARRAY instructions. + if (usage.getOpcode() != Opcodes.MULTIANEWARRAY) { + addTypeReference(usage.getOwner(), usage.getSource(), usage.isImplementedInterface()); + } + break; + case FIELD: + addFieldReference(usage); + break; + case METHOD: + addMethodReference(usage); + break; + case HANDLE: + // ReferenceCreator did not treat method handles loaded with ldc as muzzle references. + break; + case INVOKEDYNAMIC: + addInvokeDynamicReferences(usage.getHandles(), usage.getSource()); + break; + default: + throw new IllegalStateException("Unhandled advice usage " + usage.getKind()); + } } - private static boolean samePackage(String from, String to) { - int fromLength = from.lastIndexOf('/'); - int toLength = to.lastIndexOf('/'); - return fromLength == toLength && from.regionMatches(0, to, 0, fromLength + 1); + private void addFieldReference(Usage usage) { + String shadedOwner = shade(usage.getOwner()); + if (ignoreReference(shadedOwner)) { + return; + } + String source = shade(usage.getSource().getClassName()); + String sourceInternal = internalName(source); + String ownerInternal = internalName(shadedOwner); + int fieldFlags = computeMinimumFieldAccess(sourceInternal, ownerInternal); + fieldFlags |= + usage.getOpcode() == Opcodes.GETSTATIC || usage.getOpcode() == Opcodes.PUTSTATIC + ? Reference.EXPECTS_STATIC + : Reference.EXPECTS_NON_STATIC; + merge( + usage.getOwner(), + new Reference.Builder(ownerInternal) + .withSource(source, usage.getSource().getLine()) + .withFlag(computeMinimumClassAccess(sourceInternal, ownerInternal)) + .withField( + new String[] {source + ":" + usage.getSource().getLine()}, + fieldFlags, + usage.getName(), + shadeTypeDescriptor(usage.getDescriptor())) + .build()); + addDescriptorType(Type.getType(usage.getDescriptor()), usage.getSource()); } - /** - * Compute the minimum required access for FROM class to access the TO class. - * - * @return A reference flag with the required level of access. - */ - private static int computeMinimumClassAccess(final String from, final String to) { - if (from.equalsIgnoreCase(to)) { - return 0; // same access; nothing to assert - } else if (samePackage(from, to)) { - return Reference.EXPECTS_NON_PRIVATE; - } else { - return Reference.EXPECTS_PUBLIC; + private void addMethodReference(Usage usage) { + String shadedOwner = shade(usage.getOwner()); + String descriptor = shadeMethodDescriptor(usage.getDescriptor()); + if (ignoreReference(shadedOwner) || ignoreObjectMethod(usage.getName(), descriptor)) { + return; } - } - /** - * Compute the minimum required access for FROM class to access a field on the TO class. - * - * @return A reference flag with the required level of access. - */ - private static int computeMinimumFieldAccess(final String from, final String to) { - if (from.equalsIgnoreCase(to)) { - return 0; // same access; nothing to assert - } else if (samePackage(from, to)) { - return Reference.EXPECTS_NON_PRIVATE; - } else { - // Additional references: check the type hierarchy of FROM to distinguish public from - // protected - return Reference.EXPECTS_PUBLIC_OR_PROTECTED; + Type originalMethodType = Type.getMethodType(usage.getDescriptor()); + addDescriptorType(originalMethodType.getReturnType(), usage.getSource()); + for (Type argument : originalMethodType.getArgumentTypes()) { + addDescriptorType(argument, usage.getSource()); } + + String source = shade(usage.getSource().getClassName()); + String sourceInternal = internalName(source); + String ownerInternal = internalName(shadedOwner); + int methodFlags = computeMinimumMethodAccess(sourceInternal, ownerInternal); + methodFlags |= + usage.getOpcode() == Opcodes.INVOKESTATIC + ? Reference.EXPECTS_STATIC + : Reference.EXPECTS_NON_STATIC; + Type methodType = Type.getMethodType(descriptor); + merge( + usage.getOwner(), + new Reference.Builder(ownerInternal) + .withSource(source, usage.getSource().getLine()) + .withFlag( + usage.isInterfaceOwner() + ? Reference.EXPECTS_INTERFACE + : Reference.EXPECTS_NON_INTERFACE) + .withFlag(computeMinimumClassAccess(sourceInternal, ownerInternal)) + .withMethod( + new String[] {source + ":" + usage.getSource().getLine()}, + methodFlags, + usage.getName(), + methodType.getReturnType(), + methodType.getArgumentTypes()) + .build()); } - /** - * Compute the minimum required access for FROM class to access METHODTYPE on the TO class. - * - * @return A reference flag with the required level of access. - */ - private static int computeMinimumMethodAccess(final String from, final String to) { - if (from.equalsIgnoreCase(to)) { - return 0; // same access; nothing to assert - } else { - // Additional references: check the type hierarchy of FROM to distinguish public from - // protected - return Reference.EXPECTS_PUBLIC_OR_PROTECTED; + private void addInvokeDynamicReferences(List handles, SourceLocation source) { + for (HandleUse handle : handles) { + String className = handle.getOwner(); + String shadedClass = shade(className); + if (shadedClass.startsWith("java.")) { + continue; + } + String shadedSource = shade(source.getClassName()); + String sourceInternal = internalName(shadedSource); + String classInternal = internalName(shadedClass); + merge( + className, + new Reference.Builder(classInternal) + .withSource(shadedSource, source.getLine()) + .withFlag(computeMinimumClassAccess(sourceInternal, classInternal)) + .build()); } } - /** - * @return If TYPE is an array, return the underlying type. If TYPE is not an array simply return - * the type. - */ - private static Type underlyingType(Type type) { + private void addDescriptorType(Type type, SourceLocation source) { while (type.getSort() == Type.ARRAY) { type = type.getElementType(); } - return type; - } - - private final Map references = new LinkedHashMap<>(); - private String refSourceClassName; - private String refSourceTypeInternalName; - - private ReferenceCreator(final ClassVisitor classVisitor) { - super(Opcodes.ASM7, classVisitor); - } - - public Map getReferences() { - return references; + if (type.getSort() == Type.OBJECT) { + addTypeReference(type.getClassName(), source, false); + } } - private void addReference(final Reference ref) { - if (!ref.className.startsWith("java.")) { - Reference reference = references.get(ref.className); - if (null == reference) { - references.put(ref.className, ref); - } else { - references.put(ref.className, reference.merge(ref)); - } + private void addTypeReference( + String className, SourceLocation source, boolean implementedInterface) { + String shadedClass = shade(className); + if (ignoreReference(shadedClass)) { + return; } + String shadedSource = shade(source.getClassName()); + String sourceInternal = internalName(shadedSource); + String classInternal = internalName(shadedClass); + int flags = + implementedInterface + ? Reference.EXPECTS_PUBLIC + : computeMinimumClassAccess(sourceInternal, classInternal); + merge( + className, + new Reference.Builder(classInternal) + .withSource(shadedSource, source.getLine()) + .withFlag(flags) + .build()); } - @Override - public void visit( - final int version, - final int access, - final String name, - final String signature, - final String superName, - final String[] interfaces) { - refSourceClassName = getClassName(name); - Type refSourceType = Type.getType("L" + name + ";"); - refSourceTypeInternalName = refSourceType.getInternalName(); - - // Add references to each of the interfaces. - for (String iface : interfaces) { - if (!ignoreReference(iface)) { - addReference( - new Reference.Builder(iface) - .withSource( - refSourceClassName, - UNDEFINED_LINE) // We don't have a specific line number to use. - .withFlag(Reference.EXPECTS_PUBLIC) - .build()); - } + private void merge(String unshadedClassName, Reference reference) { + Reference previous = references.get(reference.className); + references.put(reference.className, previous == null ? reference : previous.merge(reference)); + if (isReferenceSource(reference.className)) { + sources.add(unshadedClassName); } - // the super type is handled by the method visitor to the constructor. - super.visit(version, access, name, signature, superName, interfaces); } - @Override - public FieldVisitor visitField( - final int access, - final String name, - final String descriptor, - final String signature, - final Object value) { - // Additional references we could check - // - annotations on field - - // intentionally not creating refs to fields here. - // Will create refs in method instructions to include line numbers. - return super.visitField(access, name, descriptor, signature, value); + private String shade(String className) { + return shader == null ? className : shader.shadeClassName(className); } - @Override - public MethodVisitor visitMethod( - final int access, - final String name, - final String descriptor, - final String signature, - final String[] exceptions) { - // Additional references we could check - // - Classes in signature (return type, params) and visible from this package - return new AdviceReferenceMethodVisitor( - super.visitMethod(access, name, descriptor, signature, exceptions)); + private String shadeTypeDescriptor(String descriptor) { + return shader == null ? descriptor : shader.shadeTypeDescriptor(descriptor); } - private class AdviceReferenceMethodVisitor extends MethodVisitor { - private int currentLineNumber = UNDEFINED_LINE; - - public AdviceReferenceMethodVisitor(final MethodVisitor methodVisitor) { - super(Opcodes.ASM7, methodVisitor); - } - - @Override - public void visitLineNumber(final int line, final Label start) { - currentLineNumber = line; - super.visitLineNumber(line, start); - } - - @Override - public void visitFieldInsn( - final int opcode, final String owner, final String name, final String descriptor) { - if (ignoreReference(owner)) { - return; - } - - // Additional references we could check - // * DONE owner class - // * DONE owner class has a field (name) - // * DONE field is static or non-static - // * DONE field's visibility from this point (NON_PRIVATE?) - // * DONE owner class's visibility from this point (NON_PRIVATE?) - // - // * DONE field-source class (descriptor) - // * DONE field-source visibility from this point (PRIVATE?) - - final Type ownerType = - owner.startsWith("[") - ? underlyingType(Type.getType(owner)) - : Type.getType("L" + owner + ";"); - final Type fieldType = Type.getType(descriptor); - - String ownerTypeInternalName = ownerType.getInternalName(); - - int fieldFlags = 0; - fieldFlags |= computeMinimumFieldAccess(refSourceTypeInternalName, ownerTypeInternalName); - fieldFlags |= - opcode == Opcodes.GETSTATIC || opcode == Opcodes.PUTSTATIC - ? Reference.EXPECTS_STATIC - : Reference.EXPECTS_NON_STATIC; - - addReference( - new Reference.Builder(ownerTypeInternalName) - .withSource(refSourceClassName, currentLineNumber) - .withFlag(computeMinimumClassAccess(refSourceTypeInternalName, ownerTypeInternalName)) - .withField( - new String[] {refSourceClassName + ":" + currentLineNumber}, - fieldFlags, - name, - fieldType) - .build()); - - final Type underlyingFieldType = underlyingType(fieldType); - String underlyingFieldTypeInternalName = underlyingFieldType.getInternalName(); - if (underlyingFieldType.getSort() == Type.OBJECT - && !ignoreReference(underlyingFieldTypeInternalName)) { - addReference( - new Reference.Builder(underlyingFieldTypeInternalName) - .withSource(refSourceClassName, currentLineNumber) - .withFlag( - computeMinimumClassAccess( - refSourceTypeInternalName, underlyingFieldTypeInternalName)) - .build()); - } - super.visitFieldInsn(opcode, owner, name, descriptor); - } - - @Override - public void visitMethodInsn( - final int opcode, - final String owner, - final String name, - final String descriptor, - final boolean isInterface) { - if (ignoreReference(owner) || ignoreObjectMethod(name, descriptor)) { - return; - } - - // Additional references we could check - // * DONE name of method owner's class - // * DONE is the owner an interface? - // * DONE owner's access from here (PRIVATE?) - // * DONE method on the owner class - // * DONE is the method static? Is it visible from here? - // * Class names from the method descriptor - // * params classes - // * return type - final Type methodType = Type.getMethodType(descriptor); - - { // ref for method return type - final Type returnType = underlyingType(methodType.getReturnType()); - String returnTypeInternalName = returnType.getInternalName(); - if (returnType.getSort() == Type.OBJECT && !ignoreReference(returnTypeInternalName)) { - addReference( - new Reference.Builder(returnTypeInternalName) - .withSource(refSourceClassName, currentLineNumber) - .withFlag( - computeMinimumClassAccess(refSourceTypeInternalName, returnTypeInternalName)) - .build()); - } - } - // refs for method param types - for (Type paramType : methodType.getArgumentTypes()) { - paramType = underlyingType(paramType); - String paramTypeInternalName = paramType.getInternalName(); - if (paramType.getSort() == Type.OBJECT && !ignoreReference(paramTypeInternalName)) { - addReference( - new Reference.Builder(paramTypeInternalName) - .withSource(refSourceClassName, currentLineNumber) - .withFlag( - computeMinimumClassAccess(refSourceTypeInternalName, paramTypeInternalName)) - .build()); - } - } - - final Type ownerType = - owner.startsWith("[") - ? underlyingType(Type.getType(owner)) - : Type.getType("L" + owner + ";"); - String ownerTypeInternalName = ownerType.getInternalName(); - - int methodFlags = 0; - methodFlags |= - opcode == Opcodes.INVOKESTATIC ? Reference.EXPECTS_STATIC : Reference.EXPECTS_NON_STATIC; - methodFlags |= computeMinimumMethodAccess(refSourceTypeInternalName, ownerTypeInternalName); + private String shadeMethodDescriptor(String descriptor) { + return shader == null ? descriptor : shader.shadeMethodDescriptor(descriptor); + } - addReference( - new Reference.Builder(ownerTypeInternalName) - .withSource(refSourceClassName, currentLineNumber) - .withFlag(isInterface ? Reference.EXPECTS_INTERFACE : Reference.EXPECTS_NON_INTERFACE) - .withFlag(computeMinimumClassAccess(refSourceTypeInternalName, ownerTypeInternalName)) - .withMethod( - new String[] {refSourceClassName + ":" + currentLineNumber}, - methodFlags, - name, - methodType.getReturnType(), - methodType.getArgumentTypes()) - .build()); - super.visitMethodInsn(opcode, owner, name, descriptor, isInterface); + private static int computeMinimumClassAccess(String from, String to) { + if (from.equalsIgnoreCase(to)) { + return 0; + } else if (samePackage(from, to)) { + return Reference.EXPECTS_NON_PRIVATE; + } else { + return Reference.EXPECTS_PUBLIC; } + } - @Override - public void visitTypeInsn(final int opcode, final String stype) { - if (ignoreReference(stype)) { - return; - } - Type type = underlyingType(Type.getObjectType(stype)); - - if (ignoreReference(type.getInternalName())) { - return; - } - - addReference( - new Reference.Builder(type.getInternalName()) - .withSource(refSourceClassName, currentLineNumber) - .withFlag( - computeMinimumClassAccess(refSourceTypeInternalName, type.getInternalName())) - .build()); - super.visitTypeInsn(opcode, stype); + private static int computeMinimumFieldAccess(String from, String to) { + if (from.equalsIgnoreCase(to)) { + return 0; + } else if (samePackage(from, to)) { + return Reference.EXPECTS_NON_PRIVATE; + } else { + return Reference.EXPECTS_PUBLIC_OR_PROTECTED; } + } - @Override - public void visitInvokeDynamicInsn( - String name, - String descriptor, - Handle bootstrapMethodHandle, - Object... bootstrapMethodArguments) { - // This part might be unnecessary... - addReference( - new Reference.Builder(bootstrapMethodHandle.getOwner()) - .withSource(refSourceClassName, currentLineNumber) - .withFlag( - computeMinimumClassAccess( - refSourceTypeInternalName, - Type.getObjectType(bootstrapMethodHandle.getOwner()).getInternalName())) - .build()); - for (Object arg : bootstrapMethodArguments) { - if (arg instanceof Handle) { - Handle handle = (Handle) arg; - addReference( - new Reference.Builder(handle.getOwner()) - .withSource(refSourceClassName, currentLineNumber) - .withFlag( - computeMinimumClassAccess( - refSourceTypeInternalName, - Type.getObjectType(handle.getOwner()).getInternalName())) - .build()); - } - } - super.visitInvokeDynamicInsn( - name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments); - } + private static int computeMinimumMethodAccess(String from, String to) { + return from.equalsIgnoreCase(to) ? 0 : Reference.EXPECTS_PUBLIC_OR_PROTECTED; + } - @Override - public void visitLdcInsn(final Object value) { - if (value instanceof Type) { - final Type type = underlyingType((Type) value); - String typeInternalName = type.getInternalName(); - if (type.getSort() == Type.OBJECT && !ignoreReference(typeInternalName)) { - addReference( - new Reference.Builder(typeInternalName) - .withSource(refSourceClassName, currentLineNumber) - .withFlag(computeMinimumClassAccess(refSourceTypeInternalName, typeInternalName)) - .build()); - } - } - super.visitLdcInsn(value); - } + private static boolean samePackage(String from, String to) { + int fromLength = from.lastIndexOf('/'); + int toLength = to.lastIndexOf('/'); + return fromLength == toLength && from.regionMatches(0, to, 0, fromLength + 1); } - /** - * {@code true} if we know this internal class is always available and doesn't require checking. - * - *

Optimization to avoid storing and checking muzzle references that will never fail. - */ private static boolean ignoreReference(String name) { String dottedName = name.replace('/', '.'); - // drop any array prefix, so we can check the actual component type if (dottedName.startsWith("[")) { int componentMarker = dottedName.lastIndexOf("[L"); if (componentMarker < 0) { - return true; // ignore primitive array references - } else { - dottedName = dottedName.substring(componentMarker + 2); + return true; } + dottedName = dottedName.substring(componentMarker + 2); } - // ignore references to core JDK types (see existing check in addReference) - if (dottedName.startsWith("java.")) { - return true; - } - // ignore SLF4J references which will be changed to datadog.slf4j in final jar - if (dottedName.startsWith("org.slf4j.")) { + if (dottedName.startsWith("java.") || dottedName.startsWith("org.slf4j.")) { return true; } for (String prefix : Constants.BOOTSTRAP_PACKAGE_PREFIXES) { @@ -505,7 +281,15 @@ private static boolean ignoreObjectMethod(String methodName, String methodDescri return OBJECT_METHODS.contains(methodSig(methodName, methodDescriptor)); } + private static boolean isReferenceSource(String className) { + return className.startsWith(REFERENCE_CREATION_PACKAGE); + } + private static String methodSig(String methodName, String methodDescriptor) { return methodName + methodDescriptor; } + + private static String internalName(String className) { + return className.replace('.', '/'); + } } diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPluginTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPluginTest.groovy index 48003a016e3..9283cbc6ce3 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPluginTest.groovy +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPluginTest.groovy @@ -28,7 +28,8 @@ class MuzzleVersionScanPluginTest extends DDSpecification { def "test assertInstrumentationMuzzled advice"() { setup: def instrumentationLoader = new ServiceEnabledClassLoader(InstrumenterModule, - Instrumenter.HasMethodAdvice, ElementMatcher, ReferenceMatcher, Reference, ReferenceCreator) + Instrumenter.HasMethodAdvice, ElementMatcher, ReferenceMatcher, Reference, ReferenceCreator, + ReferenceCreatorTestSupport) instrumentationLoader.addClass(TestInstrumentationClasses) instrumentationLoader.addClass(BaseInst) instCP.each { instrumentationLoader.addClass(it) } @@ -54,7 +55,8 @@ class MuzzleVersionScanPluginTest extends DDSpecification { def "verify advice match failure"() { setup: def instrumentationLoader = new ServiceEnabledClassLoader(InstrumenterModule, - Instrumenter.HasMethodAdvice, ElementMatcher, ReferenceMatcher, Reference, ReferenceCreator) + Instrumenter.HasMethodAdvice, ElementMatcher, ReferenceMatcher, Reference, ReferenceCreator, + ReferenceCreatorTestSupport) instrumentationLoader.addClass(TestInstrumentationClasses) instrumentationLoader.addClass(BaseInst) instCP.each { instrumentationLoader.addClass(it) } diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/ReferenceCreatorTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/ReferenceCreatorTest.groovy index 1b821256f3e..633df93eea4 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/ReferenceCreatorTest.groovy +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/ReferenceCreatorTest.groovy @@ -15,7 +15,7 @@ import static datadog.trace.agent.tooling.muzzle.Reference.EXPECTS_STATIC class ReferenceCreatorTest extends DDSpecification { def "method body creates references"() { setup: - Map references = ReferenceCreator.createReferencesFrom(MethodBodyAdvice.name, this.class.classLoader) + Map references = ReferenceCreatorTestSupport.referencesFrom(MethodBodyAdvice) expect: references.get('datadog.trace.agent.tooling.muzzle.TestAdviceClasses$MethodBodyAdvice$A') != null @@ -63,7 +63,7 @@ class ReferenceCreatorTest extends DDSpecification { // declares the method explicitly or not." setup: - Map references = ReferenceCreator.createReferencesFrom(CompiledWithInvokeinterfaceForObjectMethods.name, this.class.classLoader) + Map references = ReferenceCreatorTestSupport.referencesFrom(CompiledWithInvokeinterfaceForObjectMethods) expect: references.get(CompiledWithInvokeinterfaceForObjectMethods.DatadogInterface.name) == null @@ -71,7 +71,7 @@ class ReferenceCreatorTest extends DDSpecification { def "protected ref test"() { setup: - Map references = ReferenceCreator.createReferencesFrom(MethodBodyAdvice.B2.name, this.class.classLoader) + Map references = ReferenceCreatorTestSupport.referencesFrom(MethodBodyAdvice.B2) expect: Set bMethods = references.get('datadog.trace.agent.tooling.muzzle.TestAdviceClasses$MethodBodyAdvice$B').methods @@ -81,7 +81,7 @@ class ReferenceCreatorTest extends DDSpecification { def "ldc creates references"() { setup: - Map references = ReferenceCreator.createReferencesFrom(LdcAdvice.name, this.class.classLoader) + Map references = ReferenceCreatorTestSupport.referencesFrom(LdcAdvice) expect: references.get('datadog.trace.agent.tooling.muzzle.TestAdviceClasses$MethodBodyAdvice$A') != null @@ -89,7 +89,7 @@ class ReferenceCreatorTest extends DDSpecification { def "interface impl creates references"() { setup: - Map references = ReferenceCreator.createReferencesFrom(MethodBodyAdvice.SomeImplementation.name, this.class.classLoader) + Map references = ReferenceCreatorTestSupport.referencesFrom(MethodBodyAdvice.SomeImplementation) expect: references.get('datadog.trace.agent.tooling.muzzle.TestAdviceClasses$MethodBodyAdvice$SomeInterface') != null @@ -98,7 +98,7 @@ class ReferenceCreatorTest extends DDSpecification { def "child class creates references"() { setup: - Map references = ReferenceCreator.createReferencesFrom(MethodBodyAdvice.A2.name, this.class.classLoader) + Map references = ReferenceCreatorTestSupport.referencesFrom(MethodBodyAdvice.A2) expect: references.get('datadog.trace.agent.tooling.muzzle.TestAdviceClasses$MethodBodyAdvice$A') != null @@ -107,7 +107,7 @@ class ReferenceCreatorTest extends DDSpecification { def "instanceof creates references"() { setup: - Map references = ReferenceCreator.createReferencesFrom(InstanceofAdvice.name, this.class.classLoader) + Map references = ReferenceCreatorTestSupport.referencesFrom(InstanceofAdvice) expect: references.get('datadog.trace.agent.tooling.muzzle.TestAdviceClasses$MethodBodyAdvice$A') != null @@ -115,7 +115,7 @@ class ReferenceCreatorTest extends DDSpecification { def "invokedynamic creates references"() { setup: - Map references = ReferenceCreator.createReferencesFrom(TestAdviceClasses.InDyAdvice.name, this.class.classLoader) + Map references = ReferenceCreatorTestSupport.referencesFrom(TestAdviceClasses.InDyAdvice) expect: references.get('datadog.trace.agent.tooling.muzzle.TestAdviceClasses$MethodBodyAdvice$HasMethod') != null diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/ReferenceMatcherTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/ReferenceMatcherTest.groovy index cdfcd656e51..fd835d6dd54 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/ReferenceMatcherTest.groovy +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/ReferenceMatcherTest.groovy @@ -53,7 +53,7 @@ class ReferenceMatcherTest extends DDSpecification { def "match safe classpaths"() { setup: - Reference[] refs = ReferenceCreator.createReferencesFrom(MethodBodyAdvice.getName(), testClasspath).values().toArray(new Reference[0]) + Reference[] refs = ReferenceCreatorTestSupport.referencesFrom(MethodBodyAdvice).values().toArray(new Reference[0]) ReferenceMatcher refMatcher = new ReferenceMatcher(refs) expect: diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceProcessorPipelineTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceProcessorPipelineTest.java new file mode 100644 index 00000000000..29ea7a8419a --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceProcessorPipelineTest.java @@ -0,0 +1,74 @@ +package datadog.trace.agent.tooling.advice; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.agent.tooling.advice.AdviceScanningFixtures.PipelineModule; +import datadog.trace.agent.tooling.advice.AdviceScanningFixtures.ScanModule; +import datadog.trace.agent.tooling.muzzle.MuzzleGenerationProcessor; +import datadog.trace.agent.tooling.muzzle.MuzzleGenerationResult; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.stream.Stream; +import net.bytebuddy.ByteBuddy; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.dynamic.ClassFileLocator; +import net.bytebuddy.dynamic.DynamicType; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class AdviceProcessorPipelineTest { + @Test + void oneScanFeedsMuzzleAndThirdProcessor(@TempDir Path temp) throws Exception { + PipelineModule.instances = 0; + ScanModule.adviceRegistrations = 0; + CapturingProcessor third = new CapturingProcessor(); + AdviceScanningGradlePlugin plugin = + new AdviceScanningGradlePlugin( + temp.toFile(), Arrays.asList(new MuzzleGenerationProcessor(), third)); + TypeDescription description = new TypeDescription.ForLoadedType(PipelineModule.class); + DynamicType.Builder transformed = + plugin.apply( + new ByteBuddy().redefine(PipelineModule.class), + description, + ClassFileLocator.ForClassLoader.of(getClass().getClassLoader())); + + byte[] generatedModule = transformed.make().getBytes(); + + assertTrue(generatedModule.length > 0); + assertEquals(1, PipelineModule.instances); + assertEquals(1, ScanModule.adviceRegistrations); + assertNotNull(third.scanResult); + assertTrue(Files.isRegularFile(third.muzzle.getGeneratedClass().toPath())); + assertTrue( + Stream.of(third.muzzle.getReferences()) + .anyMatch(reference -> reference.className.equals("extra.AddedReference"))); + assertTrue( + Stream.of(third.muzzle.getReferences()) + .noneMatch( + reference -> reference.className.equals("net.bytebuddy.jar.asm.ClassReader"))); + assertTrue( + Stream.of(third.muzzle.getReferences()) + .noneMatch( + reference -> reference.className.equals("net.bytebuddy.jar.asm.ClassWriter"))); + } + + private static final class CapturingProcessor implements AdviceProcessor { + private AdviceScanResult scanResult; + private MuzzleGenerationResult muzzle; + + @Override + public Class resultType() { + return String.class; + } + + @Override + public String process(AdviceScanResult scanResult, AdviceProcessorContext context) { + this.scanResult = scanResult; + this.muzzle = context.getResult(MuzzleGenerationResult.class); + return "captured"; + } + } +} diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScannerTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScannerTest.java new file mode 100644 index 00000000000..5e8d4d7bd8d --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScannerTest.java @@ -0,0 +1,90 @@ +package datadog.trace.agent.tooling.advice; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.agent.tooling.advice.AdviceScanResult.ClassInfo; +import datadog.trace.agent.tooling.advice.AdviceScanResult.Usage; +import datadog.trace.agent.tooling.advice.AdviceScanResult.UsageKind; +import datadog.trace.agent.tooling.advice.AdviceScanningFixtures.AdditionalAdvice; +import datadog.trace.agent.tooling.advice.AdviceScanningFixtures.AdviceRoot; +import datadog.trace.agent.tooling.advice.AdviceScanningFixtures.Dependency; +import datadog.trace.agent.tooling.advice.AdviceScanningFixtures.ScanModule; +import java.io.File; +import java.nio.file.Path; +import java.util.Arrays; +import net.bytebuddy.jar.asm.ClassReader; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class AdviceScannerTest { + @Test + void scansNeutralUsesAndAdditionalClasses() throws Exception { + AdviceScanResult result = scan(new ScanModule()); + + assertEquals( + Arrays.asList(AdviceRoot.class.getName(), AdditionalAdvice.class.getName()), + result.getAdviceRoots()); + ClassInfo root = result.getClassInfo(AdviceRoot.class.getName()); + assertTrue(root.isScanned()); + assertTrue(hasUsage(root, UsageKind.FIELD, "field")); + assertTrue(hasUsage(root, UsageKind.METHOD, "")); + assertTrue(hasUsage(root, UsageKind.METHOD, "method")); + assertTrue(hasUsage(root, UsageKind.TYPE, null)); + + Usage invokeDynamic = firstUsage(root, UsageKind.INVOKEDYNAMIC); + assertNotNull(invokeDynamic); + assertFalse(invokeDynamic.getHandles().isEmpty()); + assertTrue( + root.getUsages().stream() + .anyMatch( + use -> + use.getKind() == UsageKind.TYPE + && use.getOwner().equals(Dependency.class.getName()))); + + assertFalse(result.getClassInfo(ClassReader.class.getName()).isScanned()); + assertFalse(result.getClassInfo(String.class.getName()).isScanned()); + assertTrue(result.getClassInfo(Dependency.class.getName()).isScanned()); + assertTrue(result.getClassInfo(AdditionalAdvice.class.getName()).isScanned()); + } + + @Test + void producesDeterministicResults() throws Exception { + AdviceScanResult first = scan(new ScanModule()); + AdviceScanResult second = scan(new ScanModule()); + + assertEquals(first.getClasses().keySet(), second.getClasses().keySet()); + assertEquals(first.getAdviceRoots(), second.getAdviceRoots()); + } + + @Test + void scansResolvableAdviceRootsOutsideCurrentOutput(@TempDir Path temp) { + AdviceScanResult result = + AdviceScanner.scan(new ScanModule(), temp.toFile(), getClass().getClassLoader()); + + ClassInfo root = result.getClassInfo(AdviceRoot.class.getName()); + assertTrue(root.isScanned()); + assertTrue(hasUsage(root, UsageKind.METHOD, "method")); + assertFalse(result.getClassInfo(Dependency.class.getName()).isScanned()); + } + + private static AdviceScanResult scan(ScanModule module) throws Exception { + return AdviceScanner.scan(module, classesRoot(), AdviceScannerTest.class.getClassLoader()); + } + + private static File classesRoot() throws Exception { + return new File( + AdviceScannerTest.class.getProtectionDomain().getCodeSource().getLocation().toURI()); + } + + private static boolean hasUsage(ClassInfo info, UsageKind kind, String name) { + return info.getUsages().stream() + .anyMatch(use -> use.getKind() == kind && (name == null || name.equals(use.getName()))); + } + + private static Usage firstUsage(ClassInfo info, UsageKind kind) { + return info.getUsages().stream().filter(use -> use.getKind() == kind).findFirst().orElse(null); + } +} diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScanningFixtures.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScanningFixtures.java new file mode 100644 index 00000000000..acc8a209d43 --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScanningFixtures.java @@ -0,0 +1,76 @@ +package datadog.trace.agent.tooling.advice; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.muzzle.Reference; +import java.util.ArrayList; +import java.util.List; +import java.util.function.Supplier; +import net.bytebuddy.jar.asm.ClassReader; + +final class AdviceScanningFixtures { + private AdviceScanningFixtures() {} + + static final class Dependency { + static String field; + + Dependency() {} + + String method(String value) { + return value; + } + } + + static class AdviceRoot { + static String apply(String value) { + Dependency.field = value; + Dependency dependency = new Dependency(); + Dependency[] array = new Dependency[1]; + Class type = Dependency.class; + Supplier constructor = Dependency::new; + List library = new ArrayList<>(); + Class externalLibrary = ClassReader.class; + return dependency.method( + array.length + type.getName() + constructor.get() + library + externalLibrary); + } + } + + static class AdditionalAdvice { + static void apply() { + new Dependency(); + } + } + + public static class ScanModule extends InstrumenterModule + implements Instrumenter.HasMethodAdvice { + static int adviceRegistrations; + + public ScanModule() { + super("advice-scan-test"); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + adviceRegistrations++; + transformer.applyAdvices(null, AdviceRoot.class.getName(), AdditionalAdvice.class.getName()); + } + } + + public static final class PipelineModule extends ScanModule { + static int instances; + + public PipelineModule() { + instances++; + } + + @Override + public String[] muzzleIgnoredClassNames() { + return new String[] {ClassReader.class.getName()}; + } + + @Override + public Reference[] additionalMuzzleReferences() { + return new Reference[] {new Reference.Builder("extra/AddedReference").build()}; + } + } +} diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/MuzzleWeakReferenceTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/MuzzleWeakReferenceTest.java index 4f7400df19e..e5cc688e96e 100644 --- a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/MuzzleWeakReferenceTest.java +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/MuzzleWeakReferenceTest.java @@ -15,9 +15,7 @@ public static boolean classLoaderRefIsGarbageCollected() throws InterruptedExcep ClassLoader loader = new URLClassLoader(new URL[0], null); final WeakReference clRef = new WeakReference<>(loader); final Reference[] refs = - ReferenceCreator.createReferencesFrom( - TestAdviceClasses.MethodBodyAdvice.class.getName(), - MuzzleWeakReferenceTest.class.getClassLoader()) + ReferenceCreatorTestSupport.referencesFrom(TestAdviceClasses.MethodBodyAdvice.class) .values() .toArray(new Reference[0]); final ReferenceMatcher refMatcher = new ReferenceMatcher(refs); diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorConversionTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorConversionTest.java new file mode 100644 index 00000000000..15d3e23100b --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorConversionTest.java @@ -0,0 +1,66 @@ +package datadog.trace.agent.tooling.muzzle; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.agent.tooling.AdviceShader; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.advice.AdviceScanResult; +import datadog.trace.agent.tooling.advice.AdviceScanner; +import java.util.Collection; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInfo; + +class ReferenceCreatorConversionTest { + @Test + void appliesAdviceShadingOnlyDuringConversion() { + ShadingModule module = new ShadingModule(); + AdviceScanResult scan = AdviceScanner.scan(module); + AdviceShader shader = AdviceShader.with(module.adviceShading()); + + List converted = + ReferenceCreator.createReferences(scan, scan.getAdviceRoots(), shader); + Map references = byName(converted); + + assertTrue(references.containsKey("relocated.library.TestInfo")); + assertFalse(references.containsKey(TestInfo.class.getName())); + assertNotNull(scan.getClassInfo(TestInfo.class.getName())); + } + + private static Map byName(Collection references) { + Map result = new LinkedHashMap<>(); + for (Reference reference : references) { + result.put(reference.className, reference); + } + return result; + } + + public static final class ShadingModule extends InstrumenterModule + implements Instrumenter.HasMethodAdvice { + public ShadingModule() { + super("muzzle-shading"); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice(null, ShadingAdvice.class.getName()); + } + + @Override + public Map adviceShading() { + return Collections.singletonMap("org.junit.jupiter.api", "relocated.library"); + } + } + + static final class ShadingAdvice { + static String apply(TestInfo testInfo) { + return testInfo.getDisplayName(); + } + } +} diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorTestSupport.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorTestSupport.java new file mode 100644 index 00000000000..8ddc01f8baf --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorTestSupport.java @@ -0,0 +1,39 @@ +package datadog.trace.agent.tooling.muzzle; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.advice.AdviceScanResult; +import datadog.trace.agent.tooling.advice.AdviceScanner; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +public final class ReferenceCreatorTestSupport { + private ReferenceCreatorTestSupport() {} + + public static Map referencesFrom(Class adviceClass) { + AdviceScanResult scanResult = AdviceScanner.scan(new AdviceModule(adviceClass.getName())); + List references = + ReferenceCreator.createReferences(scanResult, scanResult.getAdviceRoots(), null); + Map referencesByName = new LinkedHashMap<>(); + for (Reference reference : references) { + referencesByName.put(reference.className, reference); + } + return referencesByName; + } + + private static final class AdviceModule extends InstrumenterModule + implements Instrumenter.HasMethodAdvice { + private final String adviceClass; + + private AdviceModule(String adviceClass) { + super("jdbc"); + this.adviceClass = adviceClass; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice(null, adviceClass); + } + } +} diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/TestInstrumentationClasses.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/TestInstrumentationClasses.java index e5b5a28a3a6..a841619685d 100644 --- a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/TestInstrumentationClasses.java +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/TestInstrumentationClasses.java @@ -9,8 +9,7 @@ public abstract class TestInstrumentationClasses { static { Map references = - ReferenceCreator.createReferencesFrom( - SomeAdvice.class.getName(), SomeAdvice.class.getClassLoader()); + ReferenceCreatorTestSupport.referencesFrom(SomeAdvice.class); SOME_ADVICE_REFS = references.values().toArray(new Reference[0]); } diff --git a/dd-java-agent/instrumentation/build.gradle b/dd-java-agent/instrumentation/build.gradle index 924b627e8b6..144e4c367c7 100644 --- a/dd-java-agent/instrumentation/build.gradle +++ b/dd-java-agent/instrumentation/build.gradle @@ -22,7 +22,7 @@ subprojects { Project subProj -> subProj.pluginManager.withPlugin("dd-trace-java.build-time-instrumentation") { subProj.extensions.configure(BuildTimeInstrumentationExtension) { it.plugins.addAll( - 'datadog.trace.agent.tooling.muzzle.MuzzleGradlePlugin', + 'datadog.trace.agent.tooling.advice.AdviceScanningGradlePlugin', 'datadog.trace.agent.tooling.bytebuddy.NewTaskForGradlePlugin', 'datadog.trace.agent.tooling.bytebuddy.reqctx.RewriteRequestContextAdvicePlugin', ) From 25979b2446f61e758f2ed30f510639b5a7a89dfa Mon Sep 17 00:00:00 2001 From: Sarah Chen Date: Wed, 2 Sep 2026 10:35:52 -0400 Subject: [PATCH 2/3] Fix muzzle reference generation for shared instrumentation helpers --- .../tooling/advice/AdviceScanResult.java | 10 ++++++ .../agent/tooling/advice/AdviceScanner.java | 5 ++- .../tooling/muzzle/ReferenceCreator.java | 36 +++++++++++++------ .../tooling/advice/AdviceScannerTest.java | 6 ++++ .../advice/AdviceScanningFixtures.java | 8 ++++- .../ReferenceCreatorConversionTest.java | 4 ++- .../testing/ExternalHelper.java | 11 ++++++ 7 files changed, 66 insertions(+), 14 deletions(-) create mode 100644 dd-java-agent/agent-tooling/src/test/java/datadog/trace/instrumentation/testing/ExternalHelper.java diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanResult.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanResult.java index 5b2179c0434..8a9df288257 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanResult.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanResult.java @@ -9,6 +9,8 @@ /** Immutable, neutral result of scanning a module's method advice. */ public final class AdviceScanResult { + private static final String INSTRUMENTATION_PACKAGE = "datadog.trace.instrumentation."; + public enum UsageKind { TYPE, FIELD, @@ -164,6 +166,10 @@ public boolean isScanned() { return scanned; } + public boolean isInstrumentationClass() { + return AdviceScanResult.isInstrumentationClass(className); + } + public List getUsages() { return usages; } @@ -189,6 +195,10 @@ public ClassInfo getClassInfo(String className) { return classes.get(className); } + static boolean isInstrumentationClass(String className) { + return className.startsWith(INSTRUMENTATION_PACKAGE); + } + private static List immutableCopy(Collection values) { return Collections.unmodifiableList(new ArrayList<>(values)); } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanner.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanner.java index 1e0701b226c..1db102c5c53 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanner.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanner.java @@ -122,7 +122,10 @@ private MutableClassInfo discover(String className, String adviceClass) { } private void enqueue(MutableClassInfo info, boolean adviceRoot) { - if (info != null && !info.scanned && (adviceRoot || info.owned) && queued.add(info.className)) { + if (info != null + && !info.scanned + && (adviceRoot || info.owned || AdviceScanResult.isInstrumentationClass(info.className)) + && queued.add(info.className)) { scanQueue.addLast(info.className); } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java index 572683ea8bc..af07d041910 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java @@ -22,8 +22,6 @@ /** Converts neutral advice uses into the existing member-level muzzle reference model. */ public final class ReferenceCreator { - private static final String REFERENCE_CREATION_PACKAGE = "datadog.trace.instrumentation."; - private static final Set OBJECT_METHODS = new HashSet<>(); static { @@ -92,7 +90,11 @@ private void add(Usage usage) { } private void addFieldReference(Usage usage) { - String shadedOwner = shade(usage.getOwner()); + String owner = underlyingClassName(usage.getOwner()); + if (owner == null) { + return; + } + String shadedOwner = shade(owner); if (ignoreReference(shadedOwner)) { return; } @@ -105,7 +107,7 @@ private void addFieldReference(Usage usage) { ? Reference.EXPECTS_STATIC : Reference.EXPECTS_NON_STATIC; merge( - usage.getOwner(), + owner, new Reference.Builder(ownerInternal) .withSource(source, usage.getSource().getLine()) .withFlag(computeMinimumClassAccess(sourceInternal, ownerInternal)) @@ -119,7 +121,11 @@ private void addFieldReference(Usage usage) { } private void addMethodReference(Usage usage) { - String shadedOwner = shade(usage.getOwner()); + String owner = underlyingClassName(usage.getOwner()); + if (owner == null) { + return; + } + String shadedOwner = shade(owner); String descriptor = shadeMethodDescriptor(usage.getDescriptor()); if (ignoreReference(shadedOwner) || ignoreObjectMethod(usage.getName(), descriptor)) { return; @@ -141,7 +147,7 @@ private void addMethodReference(Usage usage) { : Reference.EXPECTS_NON_STATIC; Type methodType = Type.getMethodType(descriptor); merge( - usage.getOwner(), + owner, new Reference.Builder(ownerInternal) .withSource(source, usage.getSource().getLine()) .withFlag( @@ -210,7 +216,8 @@ private void addTypeReference( private void merge(String unshadedClassName, Reference reference) { Reference previous = references.get(reference.className); references.put(reference.className, previous == null ? reference : previous.merge(reference)); - if (isReferenceSource(reference.className)) { + ClassInfo info = scanResult.getClassInfo(unshadedClassName); + if (info != null && info.isScanned() && info.isInstrumentationClass()) { sources.add(unshadedClassName); } } @@ -281,10 +288,6 @@ private static boolean ignoreObjectMethod(String methodName, String methodDescri return OBJECT_METHODS.contains(methodSig(methodName, methodDescriptor)); } - private static boolean isReferenceSource(String className) { - return className.startsWith(REFERENCE_CREATION_PACKAGE); - } - private static String methodSig(String methodName, String methodDescriptor) { return methodName + methodDescriptor; } @@ -292,4 +295,15 @@ private static String methodSig(String methodName, String methodDescriptor) { private static String internalName(String className) { return className.replace('.', '/'); } + + private static String underlyingClassName(String className) { + if (!className.startsWith("[")) { + return className; + } + Type type = Type.getType(internalName(className)); + while (type.getSort() == Type.ARRAY) { + type = type.getElementType(); + } + return type.getSort() == Type.OBJECT ? type.getClassName() : null; + } } diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScannerTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScannerTest.java index 5e8d4d7bd8d..0116cb4ab2c 100644 --- a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScannerTest.java +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScannerTest.java @@ -12,10 +12,12 @@ import datadog.trace.agent.tooling.advice.AdviceScanningFixtures.AdviceRoot; import datadog.trace.agent.tooling.advice.AdviceScanningFixtures.Dependency; import datadog.trace.agent.tooling.advice.AdviceScanningFixtures.ScanModule; +import datadog.trace.instrumentation.testing.ExternalHelper; import java.io.File; import java.nio.file.Path; import java.util.Arrays; import net.bytebuddy.jar.asm.ClassReader; +import net.bytebuddy.jar.asm.Type; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; @@ -68,6 +70,10 @@ void scansResolvableAdviceRootsOutsideCurrentOutput(@TempDir Path temp) { assertTrue(root.isScanned()); assertTrue(hasUsage(root, UsageKind.METHOD, "method")); assertFalse(result.getClassInfo(Dependency.class.getName()).isScanned()); + ClassInfo helper = result.getClassInfo(ExternalHelper.class.getName()); + assertTrue(helper.isScanned()); + assertTrue(hasUsage(helper, UsageKind.METHOD, "getClassName")); + assertFalse(result.getClassInfo(Type.class.getName()).isScanned()); } private static AdviceScanResult scan(ScanModule module) throws Exception { diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScanningFixtures.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScanningFixtures.java index acc8a209d43..7f4c5909e4a 100644 --- a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScanningFixtures.java +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScanningFixtures.java @@ -3,6 +3,7 @@ import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.muzzle.Reference; +import datadog.trace.instrumentation.testing.ExternalHelper; import java.util.ArrayList; import java.util.List; import java.util.function.Supplier; @@ -31,7 +32,12 @@ static String apply(String value) { List library = new ArrayList<>(); Class externalLibrary = ClassReader.class; return dependency.method( - array.length + type.getName() + constructor.get() + library + externalLibrary); + array.length + + type.getName() + + constructor.get() + + library + + externalLibrary + + ExternalHelper.typeName()); } } diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorConversionTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorConversionTest.java index 15d3e23100b..ef743ca6c5f 100644 --- a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorConversionTest.java +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorConversionTest.java @@ -30,6 +30,7 @@ void appliesAdviceShadingOnlyDuringConversion() { assertTrue(references.containsKey("relocated.library.TestInfo")); assertFalse(references.containsKey(TestInfo.class.getName())); + assertFalse(references.keySet().stream().anyMatch(name -> name.startsWith("["))); assertNotNull(scan.getClassInfo(TestInfo.class.getName())); } @@ -60,7 +61,8 @@ public Map adviceShading() { static final class ShadingAdvice { static String apply(TestInfo testInfo) { - return testInfo.getDisplayName(); + TestInfo[] testInfos = {testInfo}; + return testInfos.clone()[0].getDisplayName(); } } } diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/instrumentation/testing/ExternalHelper.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/instrumentation/testing/ExternalHelper.java new file mode 100644 index 00000000000..76778dd53cf --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/instrumentation/testing/ExternalHelper.java @@ -0,0 +1,11 @@ +package datadog.trace.instrumentation.testing; + +import net.bytebuddy.jar.asm.Type; + +public final class ExternalHelper { + private ExternalHelper() {} + + public static String typeName() { + return Type.getType(Object.class).getClassName(); + } +} From 2c0e0f6bf6b9755a6b32c9bc1bda617dbfd8c891 Mon Sep 17 00:00:00 2001 From: Sarah Chen Date: Wed, 2 Sep 2026 16:01:44 -0400 Subject: [PATCH 3/3] Clean up --- .../advice/AdviceProcessorContext.java | 14 +---- .../agent/tooling/advice/AdviceScanner.java | 58 +++++++++---------- .../advice/AdviceScanningGradlePlugin.java | 4 +- .../muzzle/MuzzleGenerationProcessor.java | 17 +++--- .../muzzle/MuzzleGenerationResult.java | 23 -------- .../agent/tooling/muzzle/MuzzleGenerator.java | 3 +- .../tooling/muzzle/ReferenceCreator.java | 6 +- .../advice/AdviceProcessorPipelineTest.java | 29 +++++----- .../tooling/advice/AdviceScannerTest.java | 20 ++----- .../ReferenceCreatorConversionTest.java | 15 +---- .../muzzle/ReferenceCreatorTestSupport.java | 8 ++- 11 files changed, 64 insertions(+), 133 deletions(-) delete mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationResult.java diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceProcessorContext.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceProcessorContext.java index b2d4ec125e8..66cffc6abd2 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceProcessorContext.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceProcessorContext.java @@ -4,20 +4,16 @@ import java.io.File; import java.util.HashMap; import java.util.Map; -import net.bytebuddy.dynamic.DynamicType; /** Mutable pipeline context shared by ordered advice processors for one module. */ public final class AdviceProcessorContext { private final InstrumenterModule module; private final File targetDirectory; private final Map, Object> results = new HashMap<>(); - private DynamicType.Builder builder; - AdviceProcessorContext( - InstrumenterModule module, File targetDirectory, DynamicType.Builder builder) { + AdviceProcessorContext(InstrumenterModule module, File targetDirectory) { this.module = module; this.targetDirectory = targetDirectory; - this.builder = builder; } public InstrumenterModule getModule() { @@ -28,14 +24,6 @@ public File getTargetDirectory() { return targetDirectory; } - public DynamicType.Builder getBuilder() { - return builder; - } - - public void setBuilder(DynamicType.Builder builder) { - this.builder = builder; - } - public T getResult(Class resultType) { Object result = results.get(resultType); if (result == null) { diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanner.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanner.java index 1db102c5c53..81fa88a0a4c 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanner.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanner.java @@ -238,6 +238,29 @@ private static Type underlyingType(Type type) { return type; } + private static Usage usage( + MutableClassInfo info, + UsageKind kind, + int line, + int opcode, + String owner, + String name, + String descriptor, + boolean interfaceOwner, + boolean implementedInterface, + List handles) { + return new Usage( + kind, + new SourceLocation(info.className, line), + opcode, + owner, + name, + descriptor, + interfaceOwner, + implementedInterface, + handles); + } + private final class ScanningVisitor extends ClassVisitor { private final MutableClassInfo info; @@ -260,6 +283,7 @@ public void visit( addDependency(info, binaryInterface); info.usages.add( usage( + info, UsageKind.TYPE, UNDEFINED_LINE, -1, @@ -278,28 +302,6 @@ public MethodVisitor visitMethod( int access, String name, String descriptor, String signature, String[] exceptions) { return new ScanningMethodVisitor(info); } - - private Usage usage( - UsageKind kind, - int line, - int opcode, - String owner, - String name, - String descriptor, - boolean interfaceOwner, - boolean implementedInterface, - List handles) { - return new Usage( - kind, - new SourceLocation(info.className, line), - opcode, - owner, - name, - descriptor, - interfaceOwner, - implementedInterface, - handles); - } } private final class ScanningMethodVisitor extends MethodVisitor { @@ -422,16 +424,8 @@ private Usage usage( String descriptor, boolean interfaceOwner, List handles) { - return new Usage( - kind, - new SourceLocation(info.className, line), - opcode, - owner, - name, - descriptor, - interfaceOwner, - false, - handles); + return AdviceScanner.usage( + info, kind, line, opcode, owner, name, descriptor, interfaceOwner, false, handles); } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanningGradlePlugin.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanningGradlePlugin.java index 100296b6896..48495b9a229 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanningGradlePlugin.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/advice/AdviceScanningGradlePlugin.java @@ -55,11 +55,11 @@ public DynamicType.Builder apply( } AdviceScanResult scanResult = AdviceScanner.scan(module); - AdviceProcessorContext context = new AdviceProcessorContext(module, targetDirectory, builder); + AdviceProcessorContext context = new AdviceProcessorContext(module, targetDirectory); for (AdviceProcessor processor : processors) { runProcessor(processor, scanResult, context); } - return context.getBuilder(); + return builder; } private static void runProcessor( diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationProcessor.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationProcessor.java index cf6d93baa56..e336dfe1475 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationProcessor.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationProcessor.java @@ -8,7 +8,6 @@ import datadog.trace.agent.tooling.advice.AdviceProcessor; import datadog.trace.agent.tooling.advice.AdviceProcessorContext; import datadog.trace.agent.tooling.advice.AdviceScanResult; -import java.io.File; import java.util.HashSet; import java.util.List; import java.util.Set; @@ -16,28 +15,26 @@ /** * Converts a neutral advice scan to references and emits the existing {@code $Muzzle} side class. */ -public final class MuzzleGenerationProcessor implements AdviceProcessor { +public final class MuzzleGenerationProcessor implements AdviceProcessor { @Override - public Class resultType() { - return MuzzleGenerationResult.class; + public Class resultType() { + return Reference[].class; } @Override - public MuzzleGenerationResult process( - AdviceScanResult scanResult, AdviceProcessorContext context) { + public Reference[] process(AdviceScanResult scanResult, AdviceProcessorContext context) { InstrumenterModule module = context.getModule(); Set ignoredClasses = new HashSet<>(asList(module.muzzleIgnoredClassNames())); AdviceShader shader = AdviceShader.with(module.adviceShading()); - List references = - ReferenceCreator.createReferences(scanResult, scanResult.getAdviceRoots(), shader); + List references = ReferenceCreator.createReferences(scanResult, shader); references.removeIf(reference -> ignoredClasses.contains(reference.className)); Reference[] additionalReferences = module.additionalMuzzleReferences(); if (additionalReferences != null) { addAll(references, additionalReferences); } - File muzzleClass = MuzzleGenerator.generate(context.getTargetDirectory(), module, references); - return new MuzzleGenerationResult(references, muzzleClass); + MuzzleGenerator.generate(context.getTargetDirectory(), module, references); + return references.toArray(new Reference[0]); } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationResult.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationResult.java deleted file mode 100644 index 8fdca7c3f5c..00000000000 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerationResult.java +++ /dev/null @@ -1,23 +0,0 @@ -package datadog.trace.agent.tooling.muzzle; - -import java.io.File; -import java.util.List; - -/** Typed muzzle processor output available to later advice processors. */ -public final class MuzzleGenerationResult { - private final Reference[] references; - private final File generatedClass; - - MuzzleGenerationResult(List references, File generatedClass) { - this.references = references.toArray(new Reference[0]); - this.generatedClass = generatedClass; - } - - public Reference[] getReferences() { - return references.clone(); - } - - public File getGeneratedClass() { - return generatedClass; - } -} 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 9c0d288b1d8..70fa5a7d1dc 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 @@ -14,7 +14,7 @@ final class MuzzleGenerator { private MuzzleGenerator() {} - static File generate( + static void generate( File targetDirectory, InstrumenterModule module, List references) { File muzzleClass = new File(targetDirectory, Type.getInternalName(module.getClass()) + "$Muzzle.class"); @@ -25,7 +25,6 @@ static File generate( throw new IllegalStateException( "Cannot write muzzle class for " + module.getClass().getName(), error); } - return muzzleClass; } private static byte[] generateMuzzleClass(InstrumenterModule module, List references) { diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java index af07d041910..7b82e0f7ca1 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java @@ -10,7 +10,6 @@ import java.lang.reflect.Method; import java.util.ArrayDeque; import java.util.ArrayList; -import java.util.Collection; import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; @@ -41,10 +40,9 @@ private ReferenceCreator(AdviceScanResult scanResult, AdviceShader shader) { this.shader = shader; } - static List createReferences( - AdviceScanResult scanResult, Collection sourceClasses, AdviceShader shader) { + static List createReferences(AdviceScanResult scanResult, AdviceShader shader) { ReferenceCreator creator = new ReferenceCreator(scanResult, shader); - creator.sources.addAll(sourceClasses); + creator.sources.addAll(scanResult.getAdviceRoots()); creator.createReferences(); return new ArrayList<>(creator.references.values()); } diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceProcessorPipelineTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceProcessorPipelineTest.java index 29ea7a8419a..9b1f1f1b1e4 100644 --- a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceProcessorPipelineTest.java +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceProcessorPipelineTest.java @@ -7,7 +7,7 @@ import datadog.trace.agent.tooling.advice.AdviceScanningFixtures.PipelineModule; import datadog.trace.agent.tooling.advice.AdviceScanningFixtures.ScanModule; import datadog.trace.agent.tooling.muzzle.MuzzleGenerationProcessor; -import datadog.trace.agent.tooling.muzzle.MuzzleGenerationResult; +import datadog.trace.agent.tooling.muzzle.Reference; import java.nio.file.Files; import java.nio.file.Path; import java.util.Arrays; @@ -15,7 +15,6 @@ import net.bytebuddy.ByteBuddy; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.dynamic.ClassFileLocator; -import net.bytebuddy.dynamic.DynamicType; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; @@ -29,35 +28,33 @@ void oneScanFeedsMuzzleAndThirdProcessor(@TempDir Path temp) throws Exception { new AdviceScanningGradlePlugin( temp.toFile(), Arrays.asList(new MuzzleGenerationProcessor(), third)); TypeDescription description = new TypeDescription.ForLoadedType(PipelineModule.class); - DynamicType.Builder transformed = - plugin.apply( - new ByteBuddy().redefine(PipelineModule.class), - description, - ClassFileLocator.ForClassLoader.of(getClass().getClassLoader())); + plugin.apply( + new ByteBuddy().redefine(PipelineModule.class), + description, + ClassFileLocator.ForClassLoader.of(getClass().getClassLoader())); + Path muzzleClass = + temp.resolve(PipelineModule.class.getName().replace('.', '/') + "$Muzzle.class"); - byte[] generatedModule = transformed.make().getBytes(); - - assertTrue(generatedModule.length > 0); assertEquals(1, PipelineModule.instances); assertEquals(1, ScanModule.adviceRegistrations); assertNotNull(third.scanResult); - assertTrue(Files.isRegularFile(third.muzzle.getGeneratedClass().toPath())); + assertTrue(Files.isRegularFile(muzzleClass)); assertTrue( - Stream.of(third.muzzle.getReferences()) + Stream.of(third.muzzle) .anyMatch(reference -> reference.className.equals("extra.AddedReference"))); assertTrue( - Stream.of(third.muzzle.getReferences()) + Stream.of(third.muzzle) .noneMatch( reference -> reference.className.equals("net.bytebuddy.jar.asm.ClassReader"))); assertTrue( - Stream.of(third.muzzle.getReferences()) + Stream.of(third.muzzle) .noneMatch( reference -> reference.className.equals("net.bytebuddy.jar.asm.ClassWriter"))); } private static final class CapturingProcessor implements AdviceProcessor { private AdviceScanResult scanResult; - private MuzzleGenerationResult muzzle; + private Reference[] muzzle; @Override public Class resultType() { @@ -67,7 +64,7 @@ public Class resultType() { @Override public String process(AdviceScanResult scanResult, AdviceProcessorContext context) { this.scanResult = scanResult; - this.muzzle = context.getResult(MuzzleGenerationResult.class); + this.muzzle = context.getResult(Reference[].class); return "captured"; } } diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScannerTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScannerTest.java index 0116cb4ab2c..887273b9c54 100644 --- a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScannerTest.java +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/advice/AdviceScannerTest.java @@ -13,7 +13,6 @@ import datadog.trace.agent.tooling.advice.AdviceScanningFixtures.Dependency; import datadog.trace.agent.tooling.advice.AdviceScanningFixtures.ScanModule; import datadog.trace.instrumentation.testing.ExternalHelper; -import java.io.File; import java.nio.file.Path; import java.util.Arrays; import net.bytebuddy.jar.asm.ClassReader; @@ -23,8 +22,8 @@ class AdviceScannerTest { @Test - void scansNeutralUsesAndAdditionalClasses() throws Exception { - AdviceScanResult result = scan(new ScanModule()); + void scansNeutralUsesAndAdditionalClasses() { + AdviceScanResult result = AdviceScanner.scan(new ScanModule()); assertEquals( Arrays.asList(AdviceRoot.class.getName(), AdditionalAdvice.class.getName()), @@ -53,9 +52,9 @@ void scansNeutralUsesAndAdditionalClasses() throws Exception { } @Test - void producesDeterministicResults() throws Exception { - AdviceScanResult first = scan(new ScanModule()); - AdviceScanResult second = scan(new ScanModule()); + void producesDeterministicResults() { + AdviceScanResult first = AdviceScanner.scan(new ScanModule()); + AdviceScanResult second = AdviceScanner.scan(new ScanModule()); assertEquals(first.getClasses().keySet(), second.getClasses().keySet()); assertEquals(first.getAdviceRoots(), second.getAdviceRoots()); @@ -76,15 +75,6 @@ void scansResolvableAdviceRootsOutsideCurrentOutput(@TempDir Path temp) { assertFalse(result.getClassInfo(Type.class.getName()).isScanned()); } - private static AdviceScanResult scan(ScanModule module) throws Exception { - return AdviceScanner.scan(module, classesRoot(), AdviceScannerTest.class.getClassLoader()); - } - - private static File classesRoot() throws Exception { - return new File( - AdviceScannerTest.class.getProtectionDomain().getCodeSource().getLocation().toURI()); - } - private static boolean hasUsage(ClassInfo info, UsageKind kind, String name) { return info.getUsages().stream() .anyMatch(use -> use.getKind() == kind && (name == null || name.equals(use.getName()))); diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorConversionTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorConversionTest.java index ef743ca6c5f..06b8d14ecac 100644 --- a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorConversionTest.java +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorConversionTest.java @@ -9,9 +9,7 @@ import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.advice.AdviceScanResult; import datadog.trace.agent.tooling.advice.AdviceScanner; -import java.util.Collection; import java.util.Collections; -import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import org.junit.jupiter.api.Test; @@ -24,9 +22,8 @@ void appliesAdviceShadingOnlyDuringConversion() { AdviceScanResult scan = AdviceScanner.scan(module); AdviceShader shader = AdviceShader.with(module.adviceShading()); - List converted = - ReferenceCreator.createReferences(scan, scan.getAdviceRoots(), shader); - Map references = byName(converted); + List converted = ReferenceCreator.createReferences(scan, shader); + Map references = ReferenceCreatorTestSupport.byName(converted); assertTrue(references.containsKey("relocated.library.TestInfo")); assertFalse(references.containsKey(TestInfo.class.getName())); @@ -34,14 +31,6 @@ void appliesAdviceShadingOnlyDuringConversion() { assertNotNull(scan.getClassInfo(TestInfo.class.getName())); } - private static Map byName(Collection references) { - Map result = new LinkedHashMap<>(); - for (Reference reference : references) { - result.put(reference.className, reference); - } - return result; - } - public static final class ShadingModule extends InstrumenterModule implements Instrumenter.HasMethodAdvice { public ShadingModule() { diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorTestSupport.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorTestSupport.java index 8ddc01f8baf..874a5cec23e 100644 --- a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorTestSupport.java +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/ReferenceCreatorTestSupport.java @@ -4,8 +4,8 @@ import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.advice.AdviceScanResult; import datadog.trace.agent.tooling.advice.AdviceScanner; +import java.util.Collection; import java.util.LinkedHashMap; -import java.util.List; import java.util.Map; public final class ReferenceCreatorTestSupport { @@ -13,8 +13,10 @@ private ReferenceCreatorTestSupport() {} public static Map referencesFrom(Class adviceClass) { AdviceScanResult scanResult = AdviceScanner.scan(new AdviceModule(adviceClass.getName())); - List references = - ReferenceCreator.createReferences(scanResult, scanResult.getAdviceRoots(), null); + return byName(ReferenceCreator.createReferences(scanResult, null)); + } + + static Map byName(Collection references) { Map referencesByName = new LinkedHashMap<>(); for (Reference reference : references) { referencesByName.put(reference.className, reference);