Clean up and document
MODIFY
buildSrc/src/main/groovy/MuzzlePlugin.groovy
buildSrc/src/main/groovy/MuzzlePlugin.groovy
@@ -3,6 +3,11 @@ import org.gradle.api.Project
import java.lang.reflect.Method
+/**
+ * POC muzzle task plugin which runs muzzle validation against an instrumentation's compile-time dependencies.
+ *
+ * <p/>TODO: merge this with version scan
+ */
class MuzzlePlugin implements Plugin<Project> {
@Override
void apply(Project project) {
@@ -38,8 +43,8 @@ class MuzzlePlugin implements Plugin<Project> {
final ClassLoader agentCL = new URLClassLoader(ddUrls.toArray(new URL[0]), (ClassLoader) null)
// find all instrumenters, get muzzle, and assert
- Method assertionMethod = agentCL.loadClass('datadog.trace.agent.tooling.muzzle.MuzzleGradlePlugin')
- .getMethod('assertAllInstrumentationIsMuzzled', ClassLoader.class)
+ Method assertionMethod = agentCL.loadClass('datadog.trace.agent.tooling.muzzle.MuzzleVersionScanPlugin')
+ .getMethod('assertInstrumentationNotMuzzled', ClassLoader.class)
assertionMethod.invoke(null, userCL)
}
}

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java
@@ -123,9 +123,8 @@ public interface Instrumenter {
/**
* This method is implemented dynamically by compile-time bytecode transformations.
*
- * <p>TODO bytecode magic and documentation
+ * <p>{@see datadog.trace.agent.tooling.muzzle.MuzzleGradlePlugin}
*/
- // TODO: Make final
protected ReferenceMatcher getInstrumentationMuzzle() {
return null;
}

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGradlePlugin.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGradlePlugin.java
package datadog.trace.agent.tooling.muzzle;
-import datadog.trace.agent.tooling.HelperInjector;
import datadog.trace.agent.tooling.Instrumenter;
-import java.lang.reflect.Method;
-import java.util.List;
-import java.util.ServiceLoader;
+import net.bytebuddy.asm.AsmVisitorWrapper;
import net.bytebuddy.build.Plugin;
+import net.bytebuddy.description.field.FieldDescription;
+import net.bytebuddy.description.field.FieldList;
+import net.bytebuddy.description.method.MethodList;
import net.bytebuddy.description.type.TypeDefinition;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.DynamicType.Builder;
+import net.bytebuddy.implementation.Implementation;
+import net.bytebuddy.jar.asm.ClassVisitor;
+import net.bytebuddy.jar.asm.MethodVisitor;
+import net.bytebuddy.jar.asm.Opcodes;
+import net.bytebuddy.pool.TypePool;
+/** Bytebuddy gradle plugin which creates muzzle-references at compile time. */
public class MuzzleGradlePlugin implements Plugin {
- // TODO:
- // - Additional references to check
- // - Fields
- // - methods
- // - visit annotations
- // - visit parameter types
- // - visit method instructions
- // - method invoke type
- // - access flags (including implicit package-private)
- // - supertypes
- // - Misc
- // - Also match interfaces which extend Instrumenter
- // - Expose config instead of hardcoding datadog namespace (or reconfigure classpath)
- // - Run muzzle in matching phase (may require a rewrite of the instrumentation api)
- // - Documentation
- // - Fix TraceConfigInstrumentation
- // - assert no muzzle field defined in instrumentation
- // - make getMuzzle final in default and remove in gradle plugin
- // - pull muzzle field + method names into static constants
-
private static final TypeDescription DefaultInstrumenterTypeDesc =
new TypeDescription.ForLoadedType(Instrumenter.Default.class);
@@ -38,7 +24,7 @@ public class MuzzleGradlePlugin implements Plugin {
public boolean matches(final TypeDescription target) {
// AutoService annotation is not retained at runtime. Check for Instrumenter.Default supertype
boolean isInstrumenter = false;
- TypeDefinition instrumenter = target;
+ TypeDefinition instrumenter = null == target ? null : target.getSuperClass();
while (instrumenter != null) {
if (instrumenter.equals(DefaultInstrumenterTypeDesc)) {
isInstrumenter = true;
@@ -51,15 +37,10 @@ public class MuzzleGradlePlugin implements Plugin {
@Override
public Builder<?> apply(Builder<?> builder, TypeDescription typeDescription) {
- if (typeDescription.equals(DefaultInstrumenterTypeDesc)) {
- // FIXME
- System.out.println("~~~~FIXME: remove final modifier on Default: " + typeDescription);
- return builder;
- } else {
return builder.visit(new MuzzleVisitor());
}
- }
+ /** Compile-time Optimization used by gradle buildscripts. */
public static class NoOp implements Plugin {
@Override
public boolean matches(final TypeDescription target) {
@@ -72,53 +53,56 @@ public class MuzzleGradlePlugin implements Plugin {
}
}
- public static void assertAllInstrumentationIsMuzzled(ClassLoader cl) throws Exception {
- for (Instrumenter instrumenter :
- ServiceLoader.load(Instrumenter.class, MuzzleGradlePlugin.class.getClassLoader())) {
- if (instrumenter.getClass().getName().endsWith("TraceConfigInstrumentation")) {
- // TraceConfigInstrumentation doesn't do muzzle checks
- // check on TracerClassInstrumentation instead
- instrumenter =
- (Instrumenter)
- MuzzleGradlePlugin.class
- .getClassLoader()
- .loadClass(instrumenter.getClass().getName() + "$TracerClassInstrumentation")
- .getDeclaredConstructor()
- .newInstance();
- }
- Method m = null;
- try {
- m = instrumenter.getClass().getDeclaredMethod("getInstrumentationMuzzle");
- m.setAccessible(true);
- ReferenceMatcher muzzle = (ReferenceMatcher) m.invoke(instrumenter);
- List<Reference.Mismatch> mismatches = muzzle.getMismatchedReferenceSources(cl);
- if (mismatches.size() > 0) {
- System.err.println(
- "FAILED MUZZLE VALIDATION: " + instrumenter.getClass().getName() + " mismatches:");
- for (Reference.Mismatch mismatch : mismatches) {
- System.err.println("-- " + mismatch);
- }
- throw new RuntimeException("Instrumentation failed Muzzle validation");
- }
- } finally {
- if (null != m) {
- m.setAccessible(false);
- }
- }
- try {
- // Ratpack injects the scope manager as a helper.
- // This is likely a bug, but we'll grandfather it out of the helper checks for now.
- if (!instrumenter.getClass().getName().contains("Ratpack")) {
- // verify helper injector works
- final String[] helperClassNames = instrumenter.helperClassNames();
- if (helperClassNames.length > 0) {
- new HelperInjector(helperClassNames).transform(null, null, cl, null);
- }
- }
- } catch (Exception e) {
- System.err.println(
- "FAILED HELPER INJECTION. Are Helpers being injected in the correct order?");
- throw e;
+ private static class RemoveFinalFlagVisitor implements AsmVisitorWrapper {
+ final String methodName;
+
+ public RemoveFinalFlagVisitor(String methodName) {
+ this.methodName = methodName;
+ }
+
+ @Override
+ public int mergeWriter(int flags) {
+ return flags;
+ }
+
+ @Override
+ public int mergeReader(int flags) {
+ return flags;
+ }
+
+ @Override
+ public ClassVisitor wrap(
+ TypeDescription instrumentedType,
+ ClassVisitor classVisitor,
+ Implementation.Context implementationContext,
+ TypePool typePool,
+ FieldList<FieldDescription.InDefinedShape> fields,
+ MethodList<?> methods,
+ int writerFlags,
+ int readerFlags) {
+ return new Visitor(classVisitor);
+ }
+
+ private class Visitor extends ClassVisitor {
+ public Visitor(ClassVisitor cv) {
+ super(Opcodes.ASM6, cv);
+ }
+
+ @Override
+ public MethodVisitor visitMethod(
+ final int access,
+ final String name,
+ final String descriptor,
+ final String signature,
+ final String[] exceptions) {
+ MethodVisitor methodVisitor =
+ super.visitMethod(access, name, descriptor, signature, exceptions);
+ if (name.equals(methodName) && (access & Opcodes.ACC_FINAL) != 0) {
+ return super.visitMethod(
+ access ^ Opcodes.ACC_FINAL, name, descriptor, signature, exceptions);
+ } else {
+ return super.visitMethod(access, name, descriptor, signature, exceptions);
+ }
}
}
}

ADD
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPlugin.java
+package datadog.trace.agent.tooling.muzzle;
+
+import datadog.trace.agent.tooling.HelperInjector;
+import datadog.trace.agent.tooling.Instrumenter;
+import java.lang.reflect.Method;
+import java.util.List;
+import java.util.ServiceLoader;
+
+/**
+ * Entry point for muzzle version scan gradle plugin.
+ *
+ * <p>For each instrumenter on the classpath, run muzzle validation and throw an exception if any
+ * mismatches are detected.
+ *
+ * <p>Additionally, after a successful muzzle validation run each instrumenter's helper injector.
+ */
+public class MuzzleVersionScanPlugin {
+ public static void assertInstrumentationNotMuzzled(ClassLoader cl) throws Exception {
+ // muzzle validate all instrumenters
+ for (Instrumenter instrumenter :
+ ServiceLoader.load(Instrumenter.class, MuzzleGradlePlugin.class.getClassLoader())) {
+ if (instrumenter.getClass().getName().endsWith("TraceConfigInstrumentation")) {
+ // TraceConfigInstrumentation doesn't do muzzle checks
+ // check on TracerClassInstrumentation instead
+ instrumenter =
+ (Instrumenter)
+ MuzzleGradlePlugin.class
+ .getClassLoader()
+ .loadClass(instrumenter.getClass().getName() + "$TracerClassInstrumentation")
+ .getDeclaredConstructor()
+ .newInstance();
+ }
+ Method m = null;
+ try {
+ m = instrumenter.getClass().getDeclaredMethod("getInstrumentationMuzzle");
+ m.setAccessible(true);
+ ReferenceMatcher muzzle = (ReferenceMatcher) m.invoke(instrumenter);
+ List<Reference.Mismatch> mismatches = muzzle.getMismatchedReferenceSources(cl);
+ if (mismatches.size() > 0) {
+ System.err.println(
+ "FAILED MUZZLE VALIDATION: " + instrumenter.getClass().getName() + " mismatches:");
+ for (Reference.Mismatch mismatch : mismatches) {
+ System.err.println("-- " + mismatch);
+ }
+ throw new RuntimeException("Instrumentation failed Muzzle validation");
+ }
+ } finally {
+ if (null != m) {
+ m.setAccessible(false);
+ }
+ }
+ }
+ // run helper injector on all instrumenters
+ for (Instrumenter instrumenter :
+ ServiceLoader.load(Instrumenter.class, MuzzleGradlePlugin.class.getClassLoader())) {
+ if (instrumenter.getClass().getName().endsWith("TraceConfigInstrumentation")) {
+ // TraceConfigInstrumentation doesn't do muzzle checks
+ // check on TracerClassInstrumentation instead
+ instrumenter =
+ (Instrumenter)
+ MuzzleGradlePlugin.class
+ .getClassLoader()
+ .loadClass(instrumenter.getClass().getName() + "$TracerClassInstrumentation")
+ .getDeclaredConstructor()
+ .newInstance();
+ }
+ try {
+ // Ratpack injects the scope manager as a helper.
+ // This is likely a bug, but we'll grandfather it out of the helper checks for now.
+ if (!instrumenter.getClass().getName().contains("Ratpack")) {
+ // verify helper injector works
+ final String[] helperClassNames = instrumenter.helperClassNames();
+ if (helperClassNames.length > 0) {
+ new HelperInjector(helperClassNames).transform(null, null, cl, null);
+ }
+ }
+ } catch (Exception e) {
+ System.err.println(
+ "FAILED HELPER INJECTION. Are Helpers being injected in the correct order?");
+ throw e;
+ }
+ }
+ }
+
+ private MuzzleVersionScanPlugin() {}
+}

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
@@ -15,11 +15,7 @@ import net.bytebuddy.implementation.Implementation;
import net.bytebuddy.jar.asm.*;
import net.bytebuddy.pool.TypePool;
-/**
- * TODO: update doc Visit a class and add: 1) a private instrumenationMuzzle field and getter AND 2)
- * logic to the end of the instrumentation transformer to assert classpath is safe to apply
- * instrumentation to.
- */
+/** Visit a class and add: a private instrumenationMuzzle field and getter */
public class MuzzleVisitor implements AsmVisitorWrapper {
@Override
public int mergeWriter(int flags) {
@@ -104,7 +100,7 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
if (!referenceSources.contains(adviceClass)) {
referenceSources.add(adviceClass);
for (Map.Entry<String, Reference> entry :
- AdviceReferenceVisitor.createReferencesFrom(
+ ReferenceCreator.createReferencesFrom(
adviceClass, ReferenceMatcher.class.getClassLoader())
.entrySet()) {
if (references.containsKey(entry.getKey())) {

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
@@ -5,7 +5,7 @@ import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOAD
import datadog.trace.agent.tooling.Utils;
import java.util.*;
-/** An immutable reference to a single class file. */
+/** An immutable reference to a jvm class. */
public class Reference {
private final Set<Source> sources;
private final String className;

RENAME
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/AdviceReferenceVisitor.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java
@@ -7,11 +7,63 @@ import java.util.*;
import net.bytebuddy.jar.asm.*;
/** Visit a class and collect all references made by the visited class. */
-public class AdviceReferenceVisitor extends ClassVisitor {
+public class ReferenceCreator extends ClassVisitor {
+ /**
+ * Generate all references reachable from a given class.
+ *
+ * @param entryPointClassName Starting point for generating references.
+ * @param loader Classloader used to read class bytes.
+ * @return Map of [referenceClassName -> Reference]
+ */
+ public static Map<String, Reference> createReferencesFrom(
+ String entryPointClassName, ClassLoader loader) {
+ final Set<String> visitedSources = new HashSet<>();
+ final Map<String, Reference> references = new HashMap<>();
+
+ final Queue<String> instrumentationQueue = new ArrayDeque<>();
+ instrumentationQueue.add(entryPointClassName);
+
+ while (!instrumentationQueue.isEmpty()) {
+ final String className = instrumentationQueue.remove();
+ visitedSources.add(className);
+ try {
+ final InputStream in = loader.getResourceAsStream(Utils.getResourceName(className));
+ try {
+ final ReferenceCreator cv = new ReferenceCreator(null);
+ final ClassReader reader = new ClassReader(in);
+ reader.accept(cv, ClassReader.SKIP_FRAMES);
+
+ Map<String, Reference> instrumentationReferences = cv.getReferences();
+ for (Map.Entry<String, Reference> entry : instrumentationReferences.entrySet()) {
+ // Don't generate references created outside of the datadog instrumentation package.
+ if (!visitedSources.contains(entry.getKey())
+ && entry.getKey().startsWith("datadog.trace.instrumentation.")) {
+ instrumentationQueue.add(entry.getKey());
+ }
+ if (references.containsKey(entry.getKey())) {
+ references.put(
+ entry.getKey(), references.get(entry.getKey()).merge(entry.getValue()));
+ } else {
+ references.put(entry.getKey(), entry.getValue());
+ }
+ }
+
+ } finally {
+ if (in != null) {
+ in.close();
+ }
+ }
+ } catch (IOException ioe) {
+ throw new IllegalStateException(ioe);
+ }
+ }
+ return references;
+ }
+
private Map<String, Reference> references = new HashMap<>();
private String refSourceClassName;
- public AdviceReferenceVisitor(ClassVisitor classVisitor) {
+ private ReferenceCreator(ClassVisitor classVisitor) {
super(Opcodes.ASM6, classVisitor);
}
@@ -51,7 +103,6 @@ public class AdviceReferenceVisitor extends ClassVisitor {
}
private class AdviceReferenceMethodVisitor extends MethodVisitor {
- private boolean isAdviceMethod = false;
private int currentLineNumber = -1;
public AdviceReferenceMethodVisitor(MethodVisitor methodVisitor) {
@@ -75,56 +126,4 @@ public class AdviceReferenceVisitor extends ClassVisitor {
new Reference.Builder(owner).withSource(refSourceClassName, currentLineNumber).build());
}
}
-
- /**
- * Generate all references reachable from a given class.
- *
- * @param entryPointClassName Starting point for generating references.
- * @param loader Classloader used to read class bytes.
- * @return Map of [referenceClassName -> Reference]
- */
- public static Map<String, Reference> createReferencesFrom(
- String entryPointClassName, ClassLoader loader) {
- final Set<String> visitedSources = new HashSet<>();
- final Map<String, Reference> references = new HashMap<>();
-
- final Queue<String> instrumentationQueue = new ArrayDeque<>();
- instrumentationQueue.add(entryPointClassName);
-
- while (!instrumentationQueue.isEmpty()) {
- final String className = instrumentationQueue.remove();
- visitedSources.add(className);
- try {
- final InputStream in = loader.getResourceAsStream(Utils.getResourceName(className));
- try {
- final AdviceReferenceVisitor cv = new AdviceReferenceVisitor(null);
- final ClassReader reader = new ClassReader(in);
- reader.accept(cv, ClassReader.SKIP_FRAMES);
-
- Map<String, Reference> instrumentationReferences = cv.getReferences();
- for (Map.Entry<String, Reference> entry : instrumentationReferences.entrySet()) {
- // TODO: expose config instead of hardcoding datadog instrumentation namespace
- if (!visitedSources.contains(entry.getKey())
- && entry.getKey().startsWith("datadog.trace.instrumentation.")) {
- instrumentationQueue.add(entry.getKey());
- }
- if (references.containsKey(entry.getKey())) {
- references.put(
- entry.getKey(), references.get(entry.getKey()).merge(entry.getValue()));
- } else {
- references.put(entry.getKey(), entry.getValue());
- }
- }
-
- } finally {
- if (in != null) {
- in.close();
- }
- }
- } catch (IOException ioe) {
- throw new IllegalStateException(ioe);
- }
- }
- return references;
- }
}

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceMatcher.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceMatcher.java
@@ -3,19 +3,12 @@ package datadog.trace.agent.tooling.muzzle;
import static net.bytebuddy.dynamic.loading.ClassLoadingStrategy.BOOTSTRAP_LOADER;
import datadog.trace.agent.tooling.Utils;
-import java.security.ProtectionDomain;
import java.util.*;
import lombok.extern.slf4j.Slf4j;
-import net.bytebuddy.agent.builder.AgentBuilder;
-import net.bytebuddy.description.type.TypeDescription;
-import net.bytebuddy.utility.JavaModule;
-/**
- * A bytebuddy matcher that matches if expected references (classes, fields, methods, visibility)
- * are present on the classpath.
- */
+/** Matches a set of references against a classloader. */
@Slf4j
-public class ReferenceMatcher implements AgentBuilder.RawMatcher {
+public class ReferenceMatcher {
private final Map<ClassLoader, List<Reference.Mismatch>> mismatchCache =
Collections.synchronizedMap(new WeakHashMap<ClassLoader, List<Reference.Mismatch>>());
private final Reference[] references;
@@ -30,16 +23,6 @@ public class ReferenceMatcher implements AgentBuilder.RawMatcher {
this.helperClassNames = new HashSet<>(Arrays.asList(helperClassNames));
}
- @Override
- public boolean matches(
- TypeDescription typeDescription,
- ClassLoader classLoader,
- JavaModule module,
- Class<?> classBeingRedefined,
- ProtectionDomain protectionDomain) {
- return matches(classLoader);
- }
-
/**
* @param loader Classloader to validate against (or null for bootstrap)
* @return true if all references match the classpath of loader

MODIFY
dd-java-agent/testing/src/test/groovy/muzzle/AdviceReferenceVisitorTest.groovy
dd-java-agent/testing/src/test/groovy/muzzle/AdviceReferenceVisitorTest.groovy
@@ -2,7 +2,7 @@ package muzzle
import datadog.trace.agent.test.AgentTestRunner
import datadog.trace.agent.test.TestUtils
-import datadog.trace.agent.tooling.muzzle.AdviceReferenceVisitor
+import datadog.trace.agent.tooling.muzzle.ReferenceCreator
import datadog.trace.agent.tooling.muzzle.ReferenceMatcher
import datadog.trace.agent.tooling.muzzle.Reference
@@ -10,7 +10,7 @@ class AdviceReferenceVisitorTest extends AgentTestRunner {
def "methods body references"() {
setup:
- Map<String, Reference> references = AdviceReferenceVisitor.createReferencesFrom(AdviceClass.getName(), this.getClass().getClassLoader())
+ Map<String, Reference> references = ReferenceCreator.createReferencesFrom(AdviceClass.getName(), this.getClass().getClassLoader())
expect:
references.get('java.lang.Object') != null
@@ -22,7 +22,7 @@ class AdviceReferenceVisitorTest extends AgentTestRunner {
def "match safe classpaths"() {
setup:
- Reference[] refs = AdviceReferenceVisitor.createReferencesFrom(AdviceClass.getName(), this.getClass().getClassLoader()).values().toArray(new Reference[0])
+ 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,

Clean up and document