Muzzle method ref creation and matching
MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Utils.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Utils.java
@@ -76,6 +76,30 @@ public class Utils {
return resourceName.replaceAll("\\.class\\$", "").replace('/', '.');
}
+ /** com.foo.Bar -> com/foo/Bar */
+ public static String getInternalName(final String resourceName) {
+ return resourceName.replaceAll("\\.class\\$", "").replace('.', '/');
+ }
+
+ public static boolean isClassLoaded(final String className, final ClassLoader classLoader) {
+ Class<?> loadedClass = findLoadedClass(className, classLoader);
+ return loadedClass != null && loadedClass.getClassLoader() == classLoader;
+ }
+
+ public static Class<?> findLoadedClass(final String className, ClassLoader classLoader) {
+ if (classLoader == ClassLoaderMatcher.BOOTSTRAP_CLASSLOADER) {
+ classLoader = ClassLoader.getSystemClassLoader();
+ }
+ try {
+ findLoadedClassMethod.setAccessible(true);
+ return (Class<?>) findLoadedClassMethod.invoke(classLoader, className);
+ } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException e) {
+ throw new IllegalStateException(e);
+ } finally {
+ findLoadedClassMethod.setAccessible(false);
+ }
+ }
+
static boolean getConfigEnabled(final String name, final boolean fallback) {
final String property =
System.getProperty(

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVisitor.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVisitor.java
@@ -134,6 +134,7 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
* return this.instrumentationMuzzle;
* }
*/
+ try {
final MethodVisitor mv =
visitMethod(
Opcodes.ACC_PROTECTED + Opcodes.ACC_SYNCHRONIZED,
@@ -254,12 +255,27 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
false);
}
for (Reference.Method method : references[i].getMethods()) {
- mv.visitLdcInsn(method.getName());
+ mv.visitLdcInsn(method.getSources().size());
+ mv.visitTypeInsn(
+ Opcodes.ANEWARRAY, "datadog/trace/agent/tooling/muzzle/Reference$Source");
+ int j = 0;
+ for (Reference.Source source : method.getSources()) {
+ mv.visitInsn(Opcodes.DUP);
+ mv.visitLdcInsn(j);
+
+ mv.visitTypeInsn(Opcodes.NEW, "datadog/trace/agent/tooling/muzzle/Reference$Source");
+ mv.visitInsn(Opcodes.DUP);
+ mv.visitLdcInsn(source.getName());
+ mv.visitLdcInsn(source.getLine());mv.visitMethodInsn(Opcodes.INVOKESPECIAL, "datadog/trace/agent/tooling/muzzle/Reference$Source", "<init>", "(Ljava/lang/String;I)V", false);
+
+ mv.visitInsn(Opcodes.AASTORE);
+ ++j;
+ }
mv.visitLdcInsn(method.getFlags().size());
mv.visitTypeInsn(
Opcodes.ANEWARRAY, "datadog/trace/agent/tooling/muzzle/Reference$Flag");
- int j = 0;
+ j = 0;
for (Reference.Flag flag : method.getFlags()) {
mv.visitInsn(Opcodes.DUP);
mv.visitLdcInsn(j);
@@ -272,24 +288,40 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
++j;
}
- mv.visitLdcInsn(method.getReturnType());
+ mv.visitLdcInsn(method.getName());
- mv.visitLdcInsn(method.getParameterTypes().size());
- mv.visitTypeInsn(Opcodes.ANEWARRAY, "java/lang/String");
+ { // return type
+ mv.visitLdcInsn(method.getReturnType().getDescriptor());
+ mv.visitMethodInsn(Opcodes.INVOKESTATIC,
+ Type.getInternalName(Type.class),
+ "getType",
+ Type.getMethodDescriptor(Type.class.getMethod("getType", String.class)),
+ false);
+ }
- int k = 0;
- for (String parameterType : method.getParameterTypes()) {
+ mv.visitLdcInsn(method.getParameterTypes().size());
+ mv.visitTypeInsn(Opcodes.ANEWARRAY, Type.getInternalName(Type.class));
+ j = 0;
+ for (Type parameterType : method.getParameterTypes()) {
mv.visitInsn(Opcodes.DUP);
- mv.visitLdcInsn(k);
- mv.visitLdcInsn(parameterType);
+ mv.visitLdcInsn(j);
+
+ mv.visitLdcInsn(parameterType.getDescriptor());
+ mv.visitMethodInsn(Opcodes.INVOKESTATIC,
+ Type.getInternalName(Type.class),
+ "getType",
+ Type.getMethodDescriptor(Type.class.getMethod("getType", String.class)),
+ false);
+
mv.visitInsn(Opcodes.AASTORE);
+ j++;
}
mv.visitMethodInsn(
Opcodes.INVOKEVIRTUAL,
"datadog/trace/agent/tooling/muzzle/Reference$Builder",
"withMethod",
- "(Ljava/lang/String;[Ldatadog/trace/agent/tooling/muzzle/Reference$Flag;Ljava/lang/String;[Ljava/lang/String;)Ldatadog/trace/agent/tooling/muzzle/Reference$Builder;",
+ Type.getMethodDescriptor(Reference.Builder.class.getMethod("withMethod", Reference.Source[].class, Reference.Flag[].class, String.class, Type.class, Type[].class)),
false);
}
mv.visitMethodInsn(
@@ -328,6 +360,10 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
mv.visitMaxs(0, 0); // recomputed
mv.visitEnd();
}
+ catch(Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
super.visitField(
Opcodes.ACC_PRIVATE + Opcodes.ACC_VOLATILE,

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/Reference.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/Reference.java
package datadog.trace.agent.tooling.muzzle;
import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOADER;
+import static java.util.Collections.singletonList;
import datadog.trace.agent.tooling.Utils;
+
+import java.io.PrintWriter;
+import java.io.StringWriter;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
+import net.bytebuddy.description.type.TypeDefinition;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.jar.asm.Type;
+import net.bytebuddy.pool.TypePool;
+
/** An immutable reference to a jvm class. */
public class Reference {
private final Set<Source> sources;
@@ -111,13 +120,20 @@ public class Reference {
if (loader == BOOTSTRAP_CLASSLOADER) {
throw new IllegalStateException("Cannot directly check against bootstrap classloader");
}
- if (onClasspath(className, loader)) {
- return new ArrayList<>(0);
- } else {
- final List<Mismatch> mismatches = new ArrayList<>();
- mismatches.add(new Mismatch.MissingClass(sources.toArray(new Source[0]), className));
- return mismatches;
+ if (!onClasspath(className, loader)) {
+ return Collections.<Mismatch>singletonList(new Mismatch.MissingClass(sources.toArray(new Source[0]), className));
}
+ final List<Mismatch> mismatches = new ArrayList<>(0);
+ try {
+ UnloadedType typeOnClasspath = UnloadedType.of(className, loader);
+ for (Method requiredMethod : this.getMethods()) {
+ mismatches.addAll(typeOnClasspath.checkMatch(requiredMethod));
+ }
+ } catch (Exception e) {
+ // Shouldn't happen. Fail the reference check and add a mismatch for debug logging.
+ mismatches.add(new Mismatch.ReferenceCheckError(e));
+ }
+ return mismatches;
}
private boolean onClasspath(final String className, final ClassLoader loader) {
@@ -195,6 +211,44 @@ public class Reference {
return "Missing class " + className;
}
}
+
+ /**
+ * Fallback mismatch in case an unexpected exception occurs during reference checking.
+ */
+ public static class ReferenceCheckError extends Mismatch {
+ private final Exception referenceCheckExcetpion;
+
+ public ReferenceCheckError(Exception e) {
+ super(new Source[0]);
+ this.referenceCheckExcetpion = e;
+ }
+
+ @Override
+ String getMismatchDetails() {
+ final StringWriter sw = new StringWriter();
+ sw.write("Failed to generate reference check: ");
+ // add exception message and stack trace
+ final PrintWriter pw = new PrintWriter(sw);
+ referenceCheckExcetpion.printStackTrace(pw);
+ return sw.toString();
+ }
+ }
+
+ public static class MissingMethod extends Mismatch {
+ final String className;
+ final String method;
+
+ public MissingMethod(Source[] sources, String className, String method) {
+ super(sources);
+ this.className = className;
+ this.method = method;
+ }
+
+ @Override
+ String getMismatchDetails() {
+ return "Missing method " + className + "#" + method;
+ }
+ }
}
/** Expected flag (or lack of flag) on a class, method, or field reference. */
@@ -214,15 +268,28 @@ public class Reference {
private final Set<Source> sources;
private final Set<Flag> flags;
private final String name;
- private final String returnType;
- private final List<String> parameterTypes;
+ private final Type returnType;
+ private final List<Type> parameterTypes;
+
+ public Method(String name, String descriptor) {
+ this(new Source[0], new Flag[0], name, Type.getMethodType(descriptor).getReturnType(), Type.getMethodType(descriptor).getArgumentTypes());
+ }
+
+ public Method(
+ Source[] sources,
+ Flag[] flags,
+ String name,
+ Type returnType,
+ Type[] parameterTypes) {
+ this(new HashSet<>(Arrays.asList(sources)), new HashSet<>(Arrays.asList(flags)), name, returnType, Arrays.asList(parameterTypes));
+ }
public Method(
Set<Source> sources,
Set<Flag> flags,
String name,
- String returnType,
- List<String> parameterTypes) {
+ Type returnType,
+ List<Type> parameterTypes) {
this.sources = sources;
this.flags = flags;
this.name = name;
@@ -230,50 +297,50 @@ public class Reference {
this.parameterTypes = parameterTypes;
}
- public String getName() {
- return name;
- }
+ public Set<Source> getSources() { return sources; }
- public Set<Flag> getFlags() {
- return flags;
- }
+ public Set<Flag> getFlags() { return flags; }
+
+ public String getName() { return name; }
+
+ public Type getReturnType() { return returnType; }
- public String getReturnType() {
- return returnType;
+ public List<Type> getParameterTypes() { return parameterTypes; }
+
+ public Method merge(Method anotherMethod) {
+ if (!this.equals(anotherMethod)) {
+ throw new IllegalStateException("Cannot merge incompatible methods " + this + " <> " + anotherMethod);
}
- public List<String> getParameterTypes() {
- return parameterTypes;
+ final Set<Source> mergedSources = new HashSet<>();
+ mergedSources.addAll(sources);
+ mergedSources.addAll(anotherMethod.sources);
+
+ final Set<Flag> mergedFlags = new HashSet<>();
+ mergedFlags.addAll(flags);
+ mergedFlags.addAll(anotherMethod.flags);
+
+ return new Method(mergedSources, mergedFlags, name, returnType, parameterTypes);
}
- public Method merge(Method anotherMethod) {
- // TODO
- return this;
+ @Override
+ public String toString() {
+ // <init>()V
+ // toString()Ljava/lang/String;
+ return name + Type.getMethodType(returnType, parameterTypes.toArray(new Type[0])).getDescriptor();
}
@Override
public boolean equals(Object o) {
if (o instanceof Method) {
- Method other = (Method) o;
- if ((!name.equals(other.name))
- || (!returnType.equals(other.returnType))
- || parameterTypes.size() != other.parameterTypes.size()) {
- return false;
- }
- for (int i = 0; i < parameterTypes.size(); ++i) {
- if (!parameterTypes.get(i).equals(other.parameterTypes.get(i))) {
- return false;
- }
- }
- return true;
+ return toString().equals(o.toString());
}
return false;
}
@Override
public int hashCode() {
- // will cause collisions for overloaded method refs but performance hit should be negligable
- return name.hashCode();
+ return toString().hashCode();
}
}
@@ -323,8 +390,8 @@ public class Reference {
private final String className;
private String superName = null;
private final Set<String> interfaces = new HashSet<>();
- private final Set<Field> fields = new HashSet<>();
- private final Set<Method> methods = new HashSet<>();
+ private final List<Field> fields = new ArrayList<>();
+ private final List<Method> methods = new ArrayList<>();
public Builder(final String className) {
this.className = className;
@@ -355,14 +422,19 @@ public class Reference {
return this;
}
- public Builder withMethod(
- String methodName, Flag[] methodFlags, String returnType, String[] methodArgs) {
- // TODO
+ public Builder withMethod(Source[] sources, Flag[] methodFlags, String methodName, Type returnType, Type... methodArgs) {
+ final Method method = new Method(sources, methodFlags, methodName, returnType, methodArgs);
+ int existingIndex = methods.indexOf(method);
+ if (existingIndex == -1) {
+ methods.add(method);
+ } else {
+ methods.set(existingIndex, method.merge(methods.get(existingIndex)));
+ }
return this;
}
public Reference build() {
- return new Reference(sources, flags, className, superName, interfaces, fields, methods);
+ return new Reference(sources, flags, className, superName, interfaces, new HashSet<>(fields), new HashSet<>(methods));
}
}
}

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java
package datadog.trace.agent.tooling.muzzle;
import datadog.trace.agent.tooling.Utils;
+import datadog.trace.agent.tooling.muzzle.Reference.Source;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayDeque;
@@ -21,6 +22,10 @@ import net.bytebuddy.jar.asm.Opcodes;
// - outer class
// - inner class
public class ReferenceCreator extends ClassVisitor {
+ public static Map<String, Reference> createReferencesFrom(String entryPointClassName, ClassLoader loader) {
+ return ReferenceCreator.createReferencesFrom(entryPointClassName, loader, true);
+ }
+
/**
* Generate all references reachable from a given class.
*
@@ -28,8 +33,9 @@ public class ReferenceCreator extends ClassVisitor {
* @param loader Classloader used to read class bytes.
* @return Map of [referenceClassName -> Reference]
*/
- public static Map<String, Reference> createReferencesFrom(
- String entryPointClassName, ClassLoader loader) {
+ // TODO: document startFromMethodBodies and explain purpose
+ private static Map<String, Reference> createReferencesFrom(
+ String entryPointClassName, ClassLoader loader, boolean startFromMethodBodies) {
final Set<String> visitedSources = new HashSet<>();
final Map<String, Reference> references = new HashMap<>();
@@ -42,7 +48,9 @@ public class ReferenceCreator extends ClassVisitor {
try {
final InputStream in = loader.getResourceAsStream(Utils.getResourceName(className));
try {
- final ReferenceCreator cv = new ReferenceCreator(null);
+ final ReferenceCreator cv = new ReferenceCreator(null, startFromMethodBodies);
+ // only start from method bodies on the first pass
+ startFromMethodBodies = false;
final ClassReader reader = new ClassReader(in);
reader.accept(cv, ClassReader.SKIP_FRAMES);
@@ -75,9 +83,11 @@ public class ReferenceCreator extends ClassVisitor {
private Map<String, Reference> references = new HashMap<>();
private String refSourceClassName;
+ private boolean createFromMethodBodiesOnly;
- private ReferenceCreator(ClassVisitor classVisitor) {
+ private ReferenceCreator(ClassVisitor classVisitor, boolean createFromMethodBodiesOnly) {
super(Opcodes.ASM6, classVisitor);
+ this.createFromMethodBodiesOnly = createFromMethodBodiesOnly;
}
public Map<String, Reference> getReferences() {
@@ -156,9 +166,15 @@ public class ReferenceCreator extends ClassVisitor {
addReference(new Reference.Builder(owner)
.withSource(refSourceClassName, currentLineNumber)
.build());
- addReference(new Reference.Builder(Type.getType(descriptor).getInternalName())
+ Type fieldType = Type.getType(descriptor);
+ if (fieldType.getSort() == Type.ARRAY) {
+ fieldType = fieldType.getElementType();
+ }
+ if (fieldType.getSort() == Type.OBJECT) {
+ addReference(new Reference.Builder(fieldType.getInternalName())
.withSource(refSourceClassName, currentLineNumber)
.build());
+ }
super.visitFieldInsn(opcode, owner, name, descriptor);
}
@@ -171,7 +187,7 @@ public class ReferenceCreator extends ClassVisitor {
final boolean isInterface) {
// Additional references we could check
// * DONE name of method owner's class
- // * is the owner an interface?
+ // * DONE is the owner an interface?
// * owner's access from here (PRIVATE?)
// * method on the owner class
// * is the method static? Is it visible from here?
@@ -179,8 +195,17 @@ public class ReferenceCreator extends ClassVisitor {
// * params classes
// * return type
- addReference(new Reference.Builder(owner)
+ final Type methodType = Type.getMethodType(descriptor);
+ Source[] sources = new Reference.Source[]{new Reference.Source(refSourceClassName, currentLineNumber)};
+
+ Type ownerType = Type.getType("L"+owner+";");
+ if (ownerType.getSort() == Type.ARRAY) {
+ ownerType = ownerType.getElementType();
+ }
+ addReference(new Reference.Builder(ownerType.getInternalName())
.withSource(refSourceClassName, currentLineNumber)
+ .withFlag(isInterface ? Reference.Flag.INTERFACE : Reference.Flag.NON_INTERFACE)
+ .withMethod(sources, new Reference.Flag[0], name, methodType.getReturnType(), methodType.getArgumentTypes())
.build());
super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
}

ADD
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/UnloadedType.java
+package datadog.trace.agent.tooling.muzzle;
+
+import datadog.trace.agent.tooling.Utils;
+import net.bytebuddy.jar.asm.ClassReader;
+import net.bytebuddy.jar.asm.ClassVisitor;
+import net.bytebuddy.jar.asm.MethodVisitor;
+import net.bytebuddy.jar.asm.Opcodes;
+import datadog.trace.agent.tooling.muzzle.Reference.Method;
+import datadog.trace.agent.tooling.muzzle.Reference.Mismatch;
+import datadog.trace.agent.tooling.muzzle.Reference.Source;
+import datadog.trace.agent.tooling.muzzle.Reference.Field;
+
+import java.io.InputStream;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.WeakHashMap;
+import java.util.concurrent.ConcurrentHashMap;
+
+public class UnloadedType extends ClassVisitor {
+ private static final Map<ClassLoader, Map<String, UnloadedType>> typeCache = Collections.synchronizedMap(new WeakHashMap<ClassLoader, Map<String, UnloadedType>>());
+
+ private volatile String superName = null;
+ private volatile String className = null;
+ private volatile String[] interfaceNames = new String[0];
+ private volatile UnloadedType unloadedSuper = null;
+ private final List<UnloadedType> unloadedInterfaces = new ArrayList<>();
+ private final List<Method> methods = new ArrayList<>();
+
+ public static UnloadedType of(String className, ClassLoader classLoader) throws Exception {
+ className = Utils.getInternalName(className);
+ // TODO: triple-check cache logic and weak-refs
+ Map<String, UnloadedType> classLoaderCache = typeCache.get(classLoader);
+ if (classLoaderCache == null) {
+ synchronized (classLoader) {
+ classLoaderCache = typeCache.get(classLoader);
+ if (classLoaderCache == null) {
+ classLoaderCache = new ConcurrentHashMap<>();
+ typeCache.put(classLoader, classLoaderCache);
+ }
+ }
+ }
+ UnloadedType unloadedType = classLoaderCache.get(className);
+ if (unloadedType == null) {
+ final InputStream in = classLoader.getResourceAsStream(Utils.getResourceName(className));
+ unloadedType = new UnloadedType(null);
+ final ClassReader reader = new ClassReader(in);
+ reader.accept(unloadedType, ClassReader.SKIP_CODE);
+ if (unloadedType.superName != null) {
+ unloadedType.unloadedSuper = UnloadedType.of(unloadedType.superName, classLoader);
+ }
+ for (String interfaceName : unloadedType.interfaceNames) {
+ unloadedType.unloadedInterfaces.add(UnloadedType.of(interfaceName, classLoader));
+ }
+ classLoaderCache.put(className, unloadedType);
+ }
+ return unloadedType;
+ }
+
+ private UnloadedType(ClassVisitor cv) {
+ super(Opcodes.ASM6, cv);
+ }
+
+ public String getClassName() {
+ return className;
+ }
+
+ public String getSuperName() {
+ return superName;
+ }
+
+ public List<Mismatch> checkMatch(Method method) {
+ final List<Mismatch> mismatches = new ArrayList<>(0);
+ // does the method exist?
+ if (!hasMethod(method)) {
+ mismatches.add(new Mismatch.MissingMethod(method.getSources().toArray(new Source[0]), className, method.toString()));
+ } else {
+ // TODO: are the expected method flags present (static, public, etc)
+ }
+ return mismatches;
+ }
+
+ private boolean hasMethod(Method method) {
+ if (methods.contains(method)) {
+ return true;
+ }
+ // FIXME: private methods on the super type are not reachable!
+ if (null != unloadedSuper && unloadedSuper.hasMethod(method)) {
+ return true;
+ }
+ for (UnloadedType unloadedInterface : unloadedInterfaces) {
+ if (unloadedInterface.hasMethod(method)) { return true; }
+ }
+ return false;
+ }
+
+ public boolean hasField(Field field) {
+ // TODO does the field exist?
+ // TODO are the expected field flags present (static, public, etc)
+ throw new RuntimeException("TODO");
+ }
+
+ @Override
+ public void visit(
+ final int version,
+ final int access,
+ final String name,
+ final String signature,
+ final String superName,
+ final String[] interfaces) {
+ className = Utils.getClassName(name);
+ if (null != superName) this.superName = Utils.getClassName(superName);
+ if (null != interfaces) this.interfaceNames = interfaces;
+ super.visit(version, access, name, signature, superName, interfaces);
+ }
+
+
+ @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
+ methods.add(new Reference.Method(name, descriptor));
+ return super.visitMethod(access, name, descriptor, signature, exceptions);
+ }
+
+}

MODIFY
dd-java-agent/testing/src/test/groovy/muzzle/ReferenceCreatorTest.groovy
dd-java-agent/testing/src/test/groovy/muzzle/ReferenceCreatorTest.groovy
@@ -16,7 +16,15 @@ class ReferenceCreatorTest extends AgentTestRunner {
references.get('muzzle.TestClasses$MethodBodyAdvice$B') != null
references.get('muzzle.TestClasses$MethodBodyAdvice$SomeInterface') != null
references.get('muzzle.TestClasses$MethodBodyAdvice$SomeImplementation') != null
- references.keySet().size() == 5
+ references.keySet().size() == 6
+
+ Set<Reference.Method> bMethods = references.get('muzzle.TestClasses$MethodBodyAdvice$B').getMethods()
+ bMethods.contains(method("aMethod", "(Ljava/lang/String;)Ljava/lang/String;"))
+ bMethods.contains(method("aMethodWithPrimitives", "(Z)V"))
+ bMethods.contains(method("aMethodWithArrays", "([Ljava/lang/String;)[Ljava/lang/Object;"))
}
+ private def method(String name, String descriptor) {
+ return new Reference.Method(name, descriptor)
+ }
}

MODIFY
dd-java-agent/testing/src/test/groovy/muzzle/ReferenceMatcherTest.groovy
dd-java-agent/testing/src/test/groovy/muzzle/ReferenceMatcherTest.groovy
@@ -3,10 +3,17 @@ package muzzle
import datadog.trace.agent.test.AgentTestRunner
import datadog.trace.agent.test.TestUtils
import datadog.trace.agent.tooling.muzzle.Reference
+import datadog.trace.agent.tooling.muzzle.Reference.Method
+import datadog.trace.agent.tooling.muzzle.Reference.Source
+import datadog.trace.agent.tooling.muzzle.Reference.Flag
import datadog.trace.agent.tooling.muzzle.ReferenceCreator
import datadog.trace.agent.tooling.muzzle.ReferenceMatcher
+import datadog.trace.agent.tooling.muzzle.UnloadedType
+import net.bytebuddy.jar.asm.Type
import spock.lang.Shared
+import java.lang.ref.WeakReference
+
import static muzzle.TestClasses.*
class ReferenceMatcherTest extends AgentTestRunner {
@@ -33,4 +40,40 @@ class ReferenceMatcherTest extends AgentTestRunner {
refMatcher.getMismatchedReferenceSources(safeClasspath).size() == 0
refMatcher.getMismatchedReferenceSources(unsafeClasspath).size() == 1
}
+
+ def "matching does not hold a strong reference to classloaders"() {
+ expect:
+ MuzzleWeakReferenceTest.classLoaderRefIsGarbageCollected()
+ }
+
+ def "match methods"() {
+ setup:
+ UnloadedType unloadedB = UnloadedType.of(MethodBodyAdvice.B.getName(), MethodBodyAdvice.B.getClassLoader())
+ UnloadedType unloadedInterface = UnloadedType.of(MethodBodyAdvice.AnotherInterface.getName(), MethodBodyAdvice.AnotherInterface.getClassLoader())
+ Method methodRef
+
+ // match method declared in the class
+ when:
+ methodRef = new Method("aMethod", "(Ljava/lang/String;)Ljava/lang/String;")
+ then:
+ unloadedB.checkMatch(methodRef).size() == 0
+
+ // match method declared in the supertype
+ when:
+ methodRef = new Method("hashCode", "()I")
+ then:
+ unloadedB.checkMatch(methodRef).size() == 0
+
+ // match method declared in interface
+ when:
+ methodRef = new Method("someMethod", "()V")
+ then:
+ unloadedInterface.checkMatch(methodRef).size() == 0
+
+ // missing method mismatch
+ when:
+ methodRef = new Method(new Source[0], new Flag[0], "missingTestMethod", Type.VOID_TYPE, new Type[0])
+ then:
+ unloadedB.checkMatch(methodRef).size() == 1
+ }
}

ADD
dd-java-agent/testing/src/test/java/muzzle/MuzzleWeakReferenceTest.java
+package muzzle;
+
+import datadog.trace.agent.test.TestUtils;
+import datadog.trace.agent.tooling.muzzle.Reference;
+import datadog.trace.agent.tooling.muzzle.ReferenceCreator;
+import datadog.trace.agent.tooling.muzzle.ReferenceMatcher;
+
+import java.lang.ref.WeakReference;
+import java.net.URL;
+import java.net.URLClassLoader;
+
+public class MuzzleWeakReferenceTest {
+ /*
+ * Spock holds strong references to all local variables. For weak reference testing we must create our strong references away from Spock in this java class.
+ *
+ * Even returning a WeakReference<ClassLoader> is enough for spock to create a strong ref.
+ */
+ public static boolean classLoaderRefIsGarbageCollected() {
+ ClassLoader loader = new URLClassLoader(new URL[0], null);
+ WeakReference<ClassLoader> clRef = new WeakReference<>(loader);
+ Reference[] refs = ReferenceCreator.createReferencesFrom(TestClasses.MethodBodyAdvice.class.getName(), MuzzleWeakReferenceTest.class.getClassLoader()).values().toArray(new Reference[0]);
+ ReferenceMatcher refMatcher = new ReferenceMatcher(refs);
+ refMatcher.getMismatchedReferenceSources(loader);
+ loader = null;
+ TestUtils.awaitGC();
+ return clRef.get() == null;
+ }
+}

MODIFY
dd-java-agent/testing/src/test/java/muzzle/TestClasses.java
dd-java-agent/testing/src/test/java/muzzle/TestClasses.java
@@ -10,14 +10,20 @@ public class TestClasses {
A a = new A();
SomeInterface inter = new SomeImplementation();
inter.someMethod();
- a.b.toString();
+ a.b.aMethod("foo");
+ a.b.aMethodWithPrimitives(false);
+ a.b.aMethodWithArrays(new String[0]);
}
public static class A {
public B b = new B();
}
- public static class B {}
+ public static class B {
+ public String aMethod(String s) { return s; }
+ public void aMethodWithPrimitives(boolean b) {}
+ public Object[] aMethodWithArrays(String[] s) { return s; }
+ }
public interface SomeInterface {
void someMethod();
@@ -33,6 +39,8 @@ public class TestClasses {
public static int staticField = 0;
public final int finalField = 0;
}
+
+ public interface AnotherInterface extends SomeInterface { }
}
}

Muzzle method ref creation and matching