Muzzle method and access matching

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Utils.java

@@ -2,6 +2,7 @@ package datadog.trace.agent.tooling;
 
 import datadog.trace.bootstrap.DatadogClassLoader;
 import datadog.trace.bootstrap.DatadogClassLoader.BootstrapClassLoaderProxy;
+import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.net.URL;
 

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVisitor.java

@@ -263,10 +263,17 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
                 mv.visitInsn(Opcodes.DUP);
                 mv.visitLdcInsn(j);
 
-              mv.visitTypeInsn(Opcodes.NEW, "datadog/trace/agent/tooling/muzzle/Reference$Source");
+                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.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;
@@ -292,7 +299,8 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
 
               { // return type
                 mv.visitLdcInsn(method.getReturnType().getDescriptor());
-              mv.visitMethodInsn(Opcodes.INVOKESTATIC,
+                mv.visitMethodInsn(
+                    Opcodes.INVOKESTATIC,
                     Type.getInternalName(Type.class),
                     "getType",
                     Type.getMethodDescriptor(Type.class.getMethod("getType", String.class)),
@@ -307,7 +315,8 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
                 mv.visitLdcInsn(j);
 
                 mv.visitLdcInsn(parameterType.getDescriptor());
-              mv.visitMethodInsn(Opcodes.INVOKESTATIC,
+                mv.visitMethodInsn(
+                    Opcodes.INVOKESTATIC,
                     Type.getInternalName(Type.class),
                     "getType",
                     Type.getMethodDescriptor(Type.class.getMethod("getType", String.class)),
@@ -321,7 +330,14 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
                   Opcodes.INVOKEVIRTUAL,
                   "datadog/trace/agent/tooling/muzzle/Reference$Builder",
                   "withMethod",
-                Type.getMethodDescriptor(Reference.Builder.class.getMethod("withMethod", Reference.Source[].class, Reference.Flag[].class, String.class, Type.class, Type[].class)),
+                  Type.getMethodDescriptor(
+                      Reference.Builder.class.getMethod(
+                          "withMethod",
+                          Reference.Source[].class,
+                          Reference.Flag[].class,
+                          String.class,
+                          Type.class,
+                          Type[].class)),
                   false);
             }
             mv.visitMethodInsn(
@@ -356,11 +372,11 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
           mv.visitInsn(Opcodes.ARETURN);
           mv.visitLabel(finish);
 
-        mv.visitLocalVariable("this", "L" + instrumentationClassName + ";", null, start, finish, 0);
+          mv.visitLocalVariable(
+              "this", "L" + instrumentationClassName + ";", null, start, finish, 0);
           mv.visitMaxs(0, 0); // recomputed
           mv.visitEnd();
-        }
+        } catch (Exception e) {
-        catch(Exception e) {
           throw new RuntimeException(e);
         }
       }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/Reference.java

@@ -1,21 +1,15 @@
 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.Arrays;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
-
+import net.bytebuddy.jar.asm.Opcodes;
-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 {
@@ -88,7 +82,7 @@ public class Reference {
 
     return new Reference(
         merge(sources, anotherReference.sources),
-        merge(flags, anotherReference.flags),
+        mergeFlags(flags, anotherReference.flags),
         className,
         superName,
         merge(interfaces, anotherReference.interfaces),
@@ -110,39 +104,6 @@ public class Reference {
     return merged;
   }
 
-  /**
-   * Check this reference against a classloader's classpath.
-   *
-   * @param loader
-   * @return A list of mismatched sources. A list of size 0 means the reference matches the class.
-   */
-  public List<Mismatch> checkMatch(ClassLoader loader) {
-    if (loader == BOOTSTRAP_CLASSLOADER) {
-      throw new IllegalStateException("Cannot directly check against bootstrap classloader");
-    }
-    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) {
-    final String resourceName = Utils.getResourceName(className);
-    return loader.getResource(resourceName) != null
-        // we can also reach bootstrap classes
-        || Utils.getBootstrapProxy().getResource(resourceName) != null;
-  }
-
   public static class Source {
     private final String name;
     private final int line;
@@ -212,9 +173,26 @@ public class Reference {
       }
     }
 
-    /**
+    public static class MissingFlag extends Mismatch {
-     * Fallback mismatch in case an unexpected exception occurs during reference checking.
+      final Flag expectedFlag;
-     */
+      final String classMethodOrFieldDesc;
+      final int foundAccess;
+
+      public MissingFlag(
+          Source[] sources, String classMethodOrFieldDesc, Flag expectedFlag, int foundAccess) {
+        super(sources);
+        this.classMethodOrFieldDesc = classMethodOrFieldDesc;
+        this.expectedFlag = expectedFlag;
+        this.foundAccess = foundAccess;
+      }
+
+      @Override
+      String getMismatchDetails() {
+        return classMethodOrFieldDesc + " requires flag " + expectedFlag + " found " + foundAccess;
+      }
+    }
+
+    /** Fallback mismatch in case an unexpected exception occurs during reference checking. */
     public static class ReferenceCheckError extends Mismatch {
       private final Exception referenceCheckExcetpion;
 
@@ -252,16 +230,131 @@ public class Reference {
   }
 
   /** Expected flag (or lack of flag) on a class, method, or field reference. */
-  public static enum Flag {
+  public enum Flag {
-    PUBLIC,
+    PUBLIC {
-    PACKAGE_OR_HIGHER,
+      @Override
-    PROTECTED_OR_HIGHER,
+      public boolean supersedes(Flag anotherFlag) {
-    PRIVATE_OR_HIGHER,
+        switch (anotherFlag) {
-    NON_FINAL,
+          case PRIVATE_OR_HIGHER:
-    STATIC,
+          case PROTECTED_OR_HIGHER:
-    NON_STATIC,
+          case PACKAGE_OR_HIGHER:
-    INTERFACE,
+            return true;
-    NON_INTERFACE
+          default:
+            return false;
+        }
+      }
+
+      @Override
+      public boolean matches(int asmFlags) {
+        return (Opcodes.ACC_PUBLIC & asmFlags) != 0;
+      }
+    },
+    PACKAGE_OR_HIGHER {
+      @Override
+      public boolean supersedes(Flag anotherFlag) {
+        return anotherFlag == PRIVATE_OR_HIGHER;
+      }
+
+      @Override
+      public boolean matches(int asmFlags) {
+        return (Opcodes.ACC_PUBLIC & asmFlags) != 0
+            || ((Opcodes.ACC_PRIVATE & asmFlags) == 0 && (Opcodes.ACC_PROTECTED & asmFlags) == 0);
+      }
+    },
+    PROTECTED_OR_HIGHER {
+      @Override
+      public boolean supersedes(Flag anotherFlag) {
+        return anotherFlag == PRIVATE_OR_HIGHER;
+      }
+
+      @Override
+      public boolean matches(int asmFlags) {
+        return (Opcodes.ACC_PUBLIC & asmFlags) != 0 || (Opcodes.ACC_PROTECTED & asmFlags) != 0;
+      }
+    },
+    PRIVATE_OR_HIGHER {
+      @Override
+      public boolean matches(int asmFlags) {
+        // you can't out-private a private
+        return true;
+      }
+    },
+    NON_FINAL {
+      @Override
+      public boolean contradicts(Flag anotherFlag) {
+        return anotherFlag == FINAL;
+      }
+
+      @Override
+      public boolean matches(int asmFlags) {
+        return (Opcodes.ACC_FINAL & asmFlags) == 0;
+      }
+    },
+    FINAL {
+      @Override
+      public boolean contradicts(Flag anotherFlag) {
+        return anotherFlag == NON_FINAL;
+      }
+
+      @Override
+      public boolean matches(int asmFlags) {
+        return (Opcodes.ACC_FINAL & asmFlags) != 0;
+      }
+    },
+    STATIC {
+      @Override
+      public boolean contradicts(Flag anotherFlag) {
+        return anotherFlag == NON_STATIC;
+      }
+
+      @Override
+      public boolean matches(int asmFlags) {
+        return (Opcodes.ACC_STATIC & asmFlags) != 0;
+      }
+    },
+    NON_STATIC {
+      @Override
+      public boolean contradicts(Flag anotherFlag) {
+        return anotherFlag == STATIC;
+      }
+
+      @Override
+      public boolean matches(int asmFlags) {
+        return (Opcodes.ACC_STATIC & asmFlags) == 0;
+      }
+    },
+    INTERFACE {
+      @Override
+      public boolean contradicts(Flag anotherFlag) {
+        return anotherFlag == NON_INTERFACE;
+      }
+
+      @Override
+      public boolean matches(int asmFlags) {
+        return (Opcodes.ACC_INTERFACE & asmFlags) != 0;
+      }
+    },
+    NON_INTERFACE {
+      @Override
+      public boolean contradicts(Flag anotherFlag) {
+        return anotherFlag == INTERFACE;
+      }
+
+      @Override
+      public boolean matches(int asmFlags) {
+        return (Opcodes.ACC_INTERFACE & asmFlags) == 0;
+      }
+    };
+
+    public boolean contradicts(Flag anotherFlag) {
+      return false;
+    }
+
+    public boolean supersedes(Flag anotherFlag) {
+      return false;
+    }
+
+    public abstract boolean matches(int asmFlags);
   }
 
   public static class Method {
@@ -272,16 +365,22 @@ public class Reference {
     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());
+      this(
+          new Source[0],
+          new Flag[0],
+          name,
+          Type.getMethodType(descriptor).getReturnType(),
+          Type.getMethodType(descriptor).getArgumentTypes());
     }
 
     public Method(
-      Source[] sources,
+        Source[] sources, Flag[] flags, String name, Type returnType, Type[] parameterTypes) {
-      Flag[] flags,
+      this(
-      String name,
+          new HashSet<>(Arrays.asList(sources)),
-      Type returnType,
+          new HashSet<>(Arrays.asList(flags)),
-      Type[] parameterTypes) {
+          name,
-      this(new HashSet<>(Arrays.asList(sources)), new HashSet<>(Arrays.asList(flags)), name, returnType, Arrays.asList(parameterTypes));
+          returnType,
+          Arrays.asList(parameterTypes));
     }
 
     public Method(
@@ -297,19 +396,30 @@ public class Reference {
       this.parameterTypes = parameterTypes;
     }
 
-    public Set<Source> getSources() { return sources; }
+    public Set<Source> getSources() {
+      return sources;
+    }
 
-    public Set<Flag> getFlags() { return flags; }
+    public Set<Flag> getFlags() {
+      return flags;
+    }
 
-    public String getName() { return name; }
+    public String getName() {
+      return name;
+    }
 
-    public Type getReturnType() { return returnType; }
+    public Type getReturnType() {
+      return returnType;
+    }
 
-    public List<Type> getParameterTypes() { return parameterTypes; }
+    public List<Type> getParameterTypes() {
+      return parameterTypes;
+    }
 
     public Method merge(Method anotherMethod) {
       if (!this.equals(anotherMethod)) {
-        throw new IllegalStateException("Cannot merge incompatible methods " + this + " <> " + anotherMethod);
+        throw new IllegalStateException(
+            "Cannot merge incompatible methods " + this + " <> " + anotherMethod);
       }
 
       final Set<Source> mergedSources = new HashSet<>();
@@ -327,7 +437,8 @@ public class Reference {
     public String toString() {
       // <init>()V
       // toString()Ljava/lang/String;
-      return name + Type.getMethodType(returnType, parameterTypes.toArray(new Type[0])).getDescriptor();
+      return name
+          + Type.getMethodType(returnType, parameterTypes.toArray(new Type[0])).getDescriptor();
     }
 
     @Override
@@ -413,7 +524,7 @@ public class Reference {
     }
 
     public Builder withFlag(Flag flag) {
-      // TODO
+      flags.add(flag);
       return this;
     }
 
@@ -422,7 +533,12 @@ public class Reference {
       return this;
     }
 
-    public Builder withMethod(Source[] sources, Flag[] methodFlags, String methodName, Type returnType, Type... methodArgs) {
+    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) {
@@ -434,7 +550,14 @@ public class Reference {
     }
 
     public Reference build() {
-      return new Reference(sources, flags, className, superName, interfaces, new HashSet<>(fields), new HashSet<>(methods));
+      return new Reference(
+          sources,
+          flags,
+          className,
+          superName,
+          interfaces,
+          new HashSet<>(fields),
+          new HashSet<>(methods));
     }
   }
 }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java

@@ -1,28 +1,33 @@
 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;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
+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.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
+// Additional things we could check
 // - annotations on class
 // - outer class
 // - inner class
+// - cast opcodes in method bodies
 public class ReferenceCreator extends ClassVisitor {
-  public static Map<String, Reference> createReferencesFrom(String entryPointClassName, ClassLoader loader) {
+  public static Map<String, Reference> createReferencesFrom(
+      String entryPointClassName, ClassLoader loader) {
     return ReferenceCreator.createReferencesFrom(entryPointClassName, loader, true);
   }
 
@@ -31,9 +36,9 @@ public class ReferenceCreator extends ClassVisitor {
    *
    * @param entryPointClassName Starting point for generating references.
    * @param loader Classloader used to read class bytes.
+   * @param startFromMethodBodies if true only create refs from method bodies.
    * @return Map of [referenceClassName -> Reference]
    */
-  // TODO: document startFromMethodBodies and explain purpose
   private static Map<String, Reference> createReferencesFrom(
       String entryPointClassName, ClassLoader loader, boolean startFromMethodBodies) {
     final Set<String> visitedSources = new HashSet<>();
@@ -81,8 +86,44 @@ public class ReferenceCreator extends ClassVisitor {
     return references;
   }
 
+  private static String internalPackageName(String internalName) {
+    return internalName.replaceAll("/[^/]+$", "");
+  }
+
+  /**
+   * 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 Reference.Flag computeMinimumClassAccess(Type from, Type to) {
+    if (from.getInternalName().equalsIgnoreCase(to.getInternalName())) {
+      return Reference.Flag.PRIVATE_OR_HIGHER;
+    } else if (internalPackageName(from.getInternalName())
+        .equals(internalPackageName(to.getInternalName()))) {
+      return Reference.Flag.PACKAGE_OR_HIGHER;
+    } else {
+      return Reference.Flag.PUBLIC;
+    }
+  }
+
+  /**
+   * 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 Reference.Flag computeMinimumMethodAccess(Type from, Type to, Type methodType) {
+    if (from.getInternalName().equalsIgnoreCase(to.getInternalName())) {
+      return Reference.Flag.PRIVATE_OR_HIGHER;
+    } else {
+      // Additional references: check the type hierarchy of FROM to distinguish public from
+      // protected
+      return Reference.Flag.PROTECTED_OR_HIGHER;
+    }
+  }
+
   private Map<String, Reference> references = new HashMap<>();
   private String refSourceClassName;
+  private Type refSourceType;
   private boolean createFromMethodBodiesOnly;
 
   private ReferenceCreator(ClassVisitor classVisitor, boolean createFromMethodBodiesOnly) {
@@ -111,6 +152,7 @@ public class ReferenceCreator extends ClassVisitor {
       final String superName,
       final String[] interfaces) {
     refSourceClassName = Utils.getClassName(name);
+    refSourceType = Type.getType("L" + name + ";");
     // Additional references we could check
     // - supertype of class and visible from this package
     // - interfaces of class and visible from this package
@@ -118,7 +160,8 @@ public class ReferenceCreator extends ClassVisitor {
   }
 
   @Override
-  public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value) {
+  public FieldVisitor visitField(
+      int access, String name, String descriptor, String signature, Object value) {
     // Additional references we could check
     // - type of field + visible from this package
     return super.visitField(access, name, descriptor, signature, value);
@@ -163,15 +206,15 @@ public class ReferenceCreator extends ClassVisitor {
       //   * field-source visibility from this point (PRIVATE?)
 
       // owning class has a field
-      addReference(new Reference.Builder(owner)
+      addReference(
-        .withSource(refSourceClassName, currentLineNumber)
+          new Reference.Builder(owner).withSource(refSourceClassName, currentLineNumber).build());
-        .build());
       Type fieldType = Type.getType(descriptor);
       if (fieldType.getSort() == Type.ARRAY) {
         fieldType = fieldType.getElementType();
       }
       if (fieldType.getSort() == Type.OBJECT) {
-        addReference(new Reference.Builder(fieldType.getInternalName())
+        addReference(
+            new Reference.Builder(fieldType.getInternalName())
                 .withSource(refSourceClassName, currentLineNumber)
                 .build());
       }
@@ -188,24 +231,64 @@ public class ReferenceCreator extends ClassVisitor {
       // Additional references we could check
       // * DONE name of method owner's class
       //   * DONE is the owner an interface?
-      //   * owner's access from here (PRIVATE?)
+      //   * DONE owner's access from here (PRIVATE?)
-      //   * method on the owner class
+      //   * DONE method on the owner class
-      //   * is the method static? Is it visible from here?
+      //   * 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);
-      Source[] sources = new Reference.Source[]{new Reference.Source(refSourceClassName, currentLineNumber)};
+
+      { // ref for method return type
+        Type returnType = methodType.getReturnType();
+        if (returnType.getSort() == Type.ARRAY) {
+          returnType = returnType.getElementType();
+        }
+        if (returnType.getSort() == Type.OBJECT) {
+          addReference(
+              new Reference.Builder(returnType.getInternalName())
+                  .withSource(refSourceClassName, currentLineNumber)
+                  .withFlag(computeMinimumClassAccess(refSourceType, returnType))
+                  .build());
+        }
+      }
+      // refs for method param types
+      for (Type paramType : methodType.getArgumentTypes()) {
+        if (paramType.getSort() == Type.ARRAY) {
+          paramType = paramType.getElementType();
+        }
+        if (paramType.getSort() == Type.OBJECT) {
+          addReference(
+              new Reference.Builder(paramType.getInternalName())
+                  .withSource(refSourceClassName, currentLineNumber)
+                  .withFlag(computeMinimumClassAccess(refSourceType, paramType))
+                  .build());
+        }
+      }
 
       Type ownerType = Type.getType("L" + owner + ";");
       if (ownerType.getSort() == Type.ARRAY) {
         ownerType = ownerType.getElementType();
       }
-      addReference(new Reference.Builder(ownerType.getInternalName())
+
+      final List<Reference.Flag> methodFlags = new ArrayList<>();
+      methodFlags.add(
+          opcode == Opcodes.INVOKESTATIC ? Reference.Flag.STATIC : Reference.Flag.NON_STATIC);
+      methodFlags.add(computeMinimumMethodAccess(refSourceType, ownerType, methodType));
+
+      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())
+              .withFlag(computeMinimumClassAccess(refSourceType, ownerType))
+              .withMethod(
+                  new Reference.Source[] {
+                    new Reference.Source(refSourceClassName, currentLineNumber)
+                  },
+                  methodFlags.toArray(new Reference.Flag[0]),
+                  name,
+                  methodType.getReturnType(),
+                  methodType.getArgumentTypes())
               .build());
       super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
     }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceMatcher.java

@@ -1,16 +1,27 @@
 package datadog.trace.agent.tooling.muzzle;
 
+import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOADER;
 import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap;
 import static net.bytebuddy.dynamic.loading.ClassLoadingStrategy.BOOTSTRAP_LOADER;
 
 import datadog.trace.agent.tooling.Utils;
+import datadog.trace.agent.tooling.muzzle.Reference.Mismatch;
+import datadog.trace.agent.tooling.muzzle.Reference.Source;
 import datadog.trace.bootstrap.WeakMap;
+import java.io.InputStream;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.WeakHashMap;
+import java.util.concurrent.ConcurrentHashMap;
 import lombok.extern.slf4j.Slf4j;
+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 net.bytebuddy.jar.asm.Type;
 
 /** Matches a set of references against a classloader. */
 @Slf4j
@@ -54,7 +65,7 @@ public class ReferenceMatcher {
             // Don't reference-check helper classes.
             // They will be injected by the instrumentation's HelperInjector.
             if (!helperClassNames.contains(reference.getClassName())) {
-              mismatches.addAll(reference.checkMatch(loader));
+              mismatches.addAll(checkMatch(reference, loader));
             }
           }
           mismatchCache.put(loader, mismatches);
@@ -63,4 +74,245 @@ public class ReferenceMatcher {
     }
     return mismatches;
   }
+
+  /**
+   * Check a reference against a classloader's classpath.
+   *
+   * @param loader
+   * @return A list of mismatched sources. A list of size 0 means the reference matches the class.
+   */
+  private static List<Reference.Mismatch> checkMatch(Reference reference, ClassLoader loader) {
+    if (loader == BOOTSTRAP_CLASSLOADER) {
+      throw new IllegalStateException("Cannot directly check against bootstrap classloader");
+    }
+    if (!onClasspath(reference.getClassName(), loader)) {
+      return Collections.<Mismatch>singletonList(
+          new Mismatch.MissingClass(
+              reference.getSources().toArray(new Source[0]), reference.getClassName()));
+    }
+    final List<Mismatch> mismatches = new ArrayList<>(0);
+    try {
+      ReferenceMatcher.UnloadedType typeOnClasspath =
+          ReferenceMatcher.UnloadedType.of(reference.getClassName(), loader);
+      mismatches.addAll(typeOnClasspath.checkMatch(reference));
+      for (Reference.Method requiredMethod : reference.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 static boolean onClasspath(final String className, final ClassLoader loader) {
+    final String resourceName = Utils.getResourceName(className);
+    return loader.getResource(resourceName) != null
+        // we can also reach bootstrap classes
+        || Utils.getBootstrapProxy().getResource(resourceName) != null;
+  }
+
+  /**
+   * A representation of a jvm class created from a byte array without loading the class in
+   * question.
+   *
+   * <p>Used to compare an expected Reference with the actual runtime class without causing
+   * classloads.
+   */
+  public static class UnloadedType extends ClassVisitor {
+    private static final Map<ClassLoader, Map<String, UnloadedType>> typeCache =
+        Collections.synchronizedMap(new WeakHashMap<ClassLoader, Map<String, UnloadedType>>());
+
+    private String superName = null;
+    private String className = null;
+    private String[] interfaceNames = new String[0];
+    private UnloadedType unloadedSuper = null;
+    private final List<UnloadedType> unloadedInterfaces = new ArrayList<>();
+    private int flags;
+    private final List<Method> methods = new ArrayList<>();
+
+    public static UnloadedType of(String className, ClassLoader classLoader) throws Exception {
+      className = Utils.getInternalName(className);
+      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 int getFlags() {
+      return flags;
+    }
+
+    public List<Reference.Mismatch> checkMatch(Reference reference) {
+      final List<Reference.Mismatch> mismatches = new ArrayList<>(0);
+      for (Reference.Flag flag : reference.getFlags()) {
+        if (!flag.matches(getFlags())) {
+          final String desc = this.getClassName();
+          mismatches.add(
+              new Mismatch.MissingFlag(
+                  reference.getSources().toArray(new Source[0]), desc, flag, getFlags()));
+        }
+      }
+      return mismatches;
+    }
+
+    public List<Reference.Mismatch> checkMatch(Reference.Method method) {
+      final List<Reference.Mismatch> mismatches = new ArrayList<>(0);
+      // does the method exist?
+      Method unloadedMethod = findMethod(method, true);
+      if (unloadedMethod == null) {
+        mismatches.add(
+            new Reference.Mismatch.MissingMethod(
+                method.getSources().toArray(new Reference.Source[0]),
+                className,
+                method.toString()));
+      } else {
+        for (Reference.Flag flag : method.getFlags()) {
+          if (!flag.matches(unloadedMethod.getFlags())) {
+            final String desc = this.getClassName() + "#" + unloadedMethod.signature;
+            mismatches.add(
+                new Mismatch.MissingFlag(
+                    method.getSources().toArray(new Source[0]),
+                    desc,
+                    flag,
+                    unloadedMethod.getFlags()));
+          }
+        }
+      }
+      return mismatches;
+    }
+
+    private Method findMethod(Reference.Method method, boolean includePrivateMethods) {
+      Method unloadedMethod =
+          new Method(
+              0,
+              method.getName(),
+              Type.getMethodType(
+                      method.getReturnType(), method.getParameterTypes().toArray(new Type[0]))
+                  .getDescriptor());
+      for (Method meth : methods) {
+        if (meth.equals(unloadedMethod)) {
+          if (meth.is(Opcodes.ACC_PRIVATE)) {
+            return includePrivateMethods ? meth : null;
+          } else {
+            return meth;
+          }
+        }
+      }
+      if (null != unloadedSuper) {
+        final Method meth = unloadedSuper.findMethod(method, false);
+        if (null != meth) return meth;
+      }
+      for (UnloadedType unloadedInterface : unloadedInterfaces) {
+        final Method meth = unloadedInterface.findMethod(method, false);
+        if (null != meth) return meth;
+      }
+      return null;
+    }
+
+    public boolean hasField(Reference.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;
+      this.flags = access;
+      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 Method(access, name, descriptor));
+      return super.visitMethod(access, name, descriptor, signature, exceptions);
+    }
+
+    private static class Method {
+      private final int flags;
+      // name + descriptor
+      private final String signature;
+
+      public Method(int flags, String name, String desc) {
+        this.flags = flags;
+        this.signature = name + desc;
+      }
+
+      public boolean is(int flag) {
+        boolean result = (flags & flag) != 0;
+        return result;
+      }
+
+      public int getFlags() {
+        return flags;
+      }
+
+      @Override
+      public String toString() {
+        return new StringBuilder("Unloaded: ").append(signature).toString();
+      }
+
+      @Override
+      public boolean equals(Object o) {
+        if (o instanceof Method) {
+          return signature.toString().equals(((Method) o).signature);
+        }
+        return false;
+      }
+
+      @Override
+      public int hashCode() {
+        return signature.hashCode();
+      }
+    }
+  }
 }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/UnloadedType.java

@@ -1,132 +0,0 @@
-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);
-  }
-
-}

dd-java-agent/testing/src/test/groovy/muzzle/AdviceReferenceVisitorTest.groovy

@@ -1,39 +0,0 @@
-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.ReferenceCreator
-import datadog.trace.agent.tooling.muzzle.ReferenceMatcher
-
-class AdviceReferenceVisitorTest extends AgentTestRunner {
-
-  def "methods body references"() {
-    setup:
-    Map<String, Reference> references = ReferenceCreator.createReferencesFrom(AdviceClass.getName(), this.getClass().getClassLoader())
-
-    expect:
-    references.get('java.lang.Object') != null
-    references.get('muzzle.AdviceClass$A') != null
-    references.get('muzzle.AdviceClass$SomeInterface') != null
-    references.get('muzzle.AdviceClass$SomeImplementation') != null
-    references.keySet().size() == 4
-  }
-
-  def "match safe classpaths"() {
-    setup:
-    Reference[] refs = ReferenceCreator.createReferencesFrom(AdviceClass.getName(), this.getClass().getClassLoader()).values().toArray(new Reference[0])
-    ReferenceMatcher refMatcher = new ReferenceMatcher(refs)
-    ClassLoader safeClassloader = new URLClassLoader([TestUtils.createJarWithClasses(AdviceClass$A,
-      AdviceClass$SomeInterface,
-      AdviceClass$SomeImplementation)] as URL[],
-      (ClassLoader) null)
-    ClassLoader unsafeClassloader = new URLClassLoader([TestUtils.createJarWithClasses(AdviceClass$SomeInterface,
-      AdviceClass$SomeImplementation)] as URL[],
-      (ClassLoader) null)
-
-    expect:
-    refMatcher.getMismatchedReferenceSources(safeClassloader).size() == 0
-    refMatcher.getMismatchedReferenceSources(unsafeClassloader).size() == 1
-  }
-}

dd-java-agent/testing/src/test/groovy/muzzle/ReferenceCreatorTest.groovy

@@ -12,19 +12,47 @@ class ReferenceCreatorTest extends AgentTestRunner {
 
     expect:
     references.get('java.lang.Object') != null
+    references.get('java.lang.String') != null
     references.get('muzzle.TestClasses$MethodBodyAdvice$A') != null
     references.get('muzzle.TestClasses$MethodBodyAdvice$B') != null
     references.get('muzzle.TestClasses$MethodBodyAdvice$SomeInterface') != null
     references.get('muzzle.TestClasses$MethodBodyAdvice$SomeImplementation') != null
     references.keySet().size() == 6
 
+    // interface flags
+    references.get('muzzle.TestClasses$MethodBodyAdvice$B').getFlags().contains(Reference.Flag.NON_INTERFACE)
+    references.get('muzzle.TestClasses$MethodBodyAdvice$SomeInterface').getFlags().contains(Reference.Flag.INTERFACE)
+
+    // class access flags
+    references.get('java.lang.Object').getFlags().contains(Reference.Flag.PUBLIC)
+    references.get('muzzle.TestClasses$MethodBodyAdvice$B').getFlags().contains(Reference.Flag.PACKAGE_OR_HIGHER)
+
     Set<Reference.Method> bMethods = references.get('muzzle.TestClasses$MethodBodyAdvice$B').getMethods()
-    bMethods.contains(method("aMethod", "(Ljava/lang/String;)Ljava/lang/String;"))
+    findMethod(bMethods, "aMethod", "(Ljava/lang/String;)Ljava/lang/String;") != null
-    bMethods.contains(method("aMethodWithPrimitives", "(Z)V"))
+    findMethod(bMethods, "aMethodWithPrimitives", "(Z)V") != null
-    bMethods.contains(method("aMethodWithArrays", "([Ljava/lang/String;)[Ljava/lang/Object;"))
+    findMethod(bMethods, "aStaticMethod", "()V") != null
+    findMethod(bMethods, "aMethodWithArrays", "([Ljava/lang/String;)[Ljava/lang/Object;") != null
+
+    findMethod(bMethods, "aMethod", "(Ljava/lang/String;)Ljava/lang/String;").getFlags().contains(Reference.Flag.NON_STATIC)
+    findMethod(bMethods, "aStaticMethod", "()V").getFlags().contains(Reference.Flag.STATIC)
   }
 
-  private def method(String name, String descriptor) {
+  def "protected ref test"() {
-    return new Reference.Method(name, descriptor)
+    setup:
+    Map<String, Reference> references = ReferenceCreator.createReferencesFrom(MethodBodyAdvice.B2.getName(), this.getClass().getClassLoader())
+
+    expect:
+    Set<Reference.Method> bMethods = references.get('muzzle.TestClasses$MethodBodyAdvice$B').getMethods()
+    findMethod(bMethods, "protectedMethod", "()V") != null
+    findMethod(bMethods, "protectedMethod", "()V").getFlags().contains(Reference.Flag.PROTECTED_OR_HIGHER)
+  }
+
+  private static findMethod(Set<Reference.Method> methods, String methodName, String methodDesc) {
+    for (Reference.Method method : methods) {
+      if (method == new Reference.Method(methodName, methodDesc)) {
+        return method
+      }
+    }
+    return null
   }
 }

dd-java-agent/testing/src/test/groovy/muzzle/ReferenceMatcherTest.groovy

@@ -2,18 +2,17 @@ package muzzle
 
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.agent.test.TestUtils
+import datadog.trace.agent.tooling.Utils
 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 {
@@ -46,10 +45,30 @@ class ReferenceMatcherTest extends AgentTestRunner {
     MuzzleWeakReferenceTest.classLoaderRefIsGarbageCollected()
   }
 
+  def "match classes"() {
+    ReferenceMatcher.UnloadedType unloadedB = ReferenceMatcher.UnloadedType.of(MethodBodyAdvice.B.getName(), MethodBodyAdvice.B.getClassLoader())
+    Reference ref
+
+    when:
+    ref = new Reference.Builder(Utils.getInternalName(MethodBodyAdvice.B.getName()))
+      .withFlag(Flag.NON_INTERFACE)
+      .build()
+    then:
+    unloadedB.checkMatch(ref).size() == 0
+
+    when:
+    ref = new Reference.Builder(Utils.getInternalName(MethodBodyAdvice.B.getName()))
+      .withFlag(Flag.INTERFACE)
+      .build()
+    then:
+    unloadedB.checkMatch(ref).size() == 1
+  }
+
   def "match methods"() {
     setup:
-    UnloadedType unloadedB = UnloadedType.of(MethodBodyAdvice.B.getName(), MethodBodyAdvice.B.getClassLoader())
+    ReferenceMatcher.UnloadedType unloadedB = ReferenceMatcher.UnloadedType.of(MethodBodyAdvice.B.getName(), MethodBodyAdvice.B.getClassLoader())
-    UnloadedType unloadedInterface = UnloadedType.of(MethodBodyAdvice.AnotherInterface.getName(), MethodBodyAdvice.AnotherInterface.getClassLoader())
+    ReferenceMatcher.UnloadedType unloadedB2 = ReferenceMatcher.UnloadedType.of(MethodBodyAdvice.B2.getName(), MethodBodyAdvice.B2.getClassLoader())
+    ReferenceMatcher.UnloadedType unloadedInterface = ReferenceMatcher.UnloadedType.of(MethodBodyAdvice.AnotherInterface.getName(), MethodBodyAdvice.AnotherInterface.getClassLoader())
     Method methodRef
 
     // match method declared in the class
@@ -70,6 +89,24 @@ class ReferenceMatcherTest extends AgentTestRunner {
     then:
     unloadedInterface.checkMatch(methodRef).size() == 0
 
+    // match private method in the class
+    when:
+    methodRef = new Method("privateStuff", "()V")
+    then:
+    unloadedB.checkMatch(methodRef).size() == 0
+
+    // fail to match private method in superclass
+    when:
+    methodRef = new Method("privateStuff", "()V")
+    then:
+    unloadedB2.checkMatch(methodRef).size() == 1
+
+    // static method flag mismatch
+    when:
+    methodRef = new Method(new Source[0], [Flag.NON_STATIC] as Flag[], "aStaticMethod", Type.getType("V"))
+    then:
+    unloadedB2.checkMatch(methodRef).size() == 1
+
     // missing method mismatch
     when:
     methodRef = new Method(new Source[0], new Flag[0], "missingTestMethod", Type.VOID_TYPE, new Type[0])

dd-java-agent/testing/src/test/java/muzzle/MuzzleWeakReferenceTest.java

@@ -4,7 +4,6 @@ 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;
@@ -18,7 +17,12 @@ public class MuzzleWeakReferenceTest {
   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]);
+    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;

dd-java-agent/testing/src/test/java/muzzle/TestClasses.java

@@ -13,6 +13,7 @@ public class TestClasses {
       a.b.aMethod("foo");
       a.b.aMethodWithPrimitives(false);
       a.b.aMethodWithArrays(new String[0]);
+      B.aStaticMethod();
     }
 
     public static class A {
@@ -20,9 +21,28 @@ public class TestClasses {
     }
 
     public static class B {
-      public String aMethod(String s) { return s; }
+      public String aMethod(String s) {
+        return s;
+      }
+
       public void aMethodWithPrimitives(boolean b) {}
-      public Object[] aMethodWithArrays(String[] s) { return s; }
+
+      public Object[] aMethodWithArrays(String[] s) {
+        return s;
+      }
+
+      private void privateStuff() {}
+
+      protected void protectedMethod() {}
+
+      public static void aStaticMethod() {}
+    }
+
+    public static class B2 extends B {
+      public void stuff() {
+        B b = new B();
+        b.protectedMethod();
+      }
     }
 
     public interface SomeInterface {
@@ -42,5 +62,4 @@ public class TestClasses {
 
     public interface AnotherInterface extends SomeInterface {}
   }
-
 }