Generate at compile-time and lazy-load reference matchers
MODIFY
dd-java-agent-ittests/src/test/groovy/datadog/trace/agent/integration/muzzle/MuzzleBytecodeTransformTest.groovy
dd-java-agent-ittests/src/test/groovy/datadog/trace/agent/integration/muzzle/MuzzleBytecodeTransformTest.groovy
package datadog.trace.agent.integration.muzzle
import datadog.trace.agent.test.IntegrationTestUtils
+import java.lang.reflect.Field
+import java.lang.reflect.Method
import spock.lang.Specification
@@ -9,15 +11,38 @@ class MuzzleBytecodeTransformTest extends Specification {
def "muzzle fields added to all instrumentation"() {
setup:
List<Class> unMuzzledClasses = []
+ List<Class> nonLazyFields = []
+ List<Class> unInitFields = []
for (final Object instrumenter : ServiceLoader.load(IntegrationTestUtils.getAgentClassLoader().loadClass("datadog.trace.agent.tooling.Instrumenter"), IntegrationTestUtils.getAgentClassLoader())) {
+ Field f
+ Method m
try {
- instrumenter.getClass().getDeclaredField("instrumentationMuzzle")
- } catch(NoSuchFieldException nsfe) {
+ f = instrumenter.getClass().getDeclaredField("instrumentationMuzzle")
+ f.setAccessible(true)
+ if (f.get(instrumenter) != null) {
+ nonLazyFields.add(instrumenter.getClass())
+ }
+ m = instrumenter.getClass().getDeclaredMethod("getInstrumentationMuzzle")
+ m.setAccessible(true)
+ m.invoke(instrumenter)
+ if (f.get(instrumenter) == null) {
+ unInitFields.add(instrumenter.getClass())
+ }
+ } catch(NoSuchFieldException | NoSuchMethodException e) {
unMuzzledClasses.add(instrumenter.getClass())
+ } finally {
+ if (null != f) {
+ f.setAccessible(false)
+ }
+ if (null != m) {
+ m.setAccessible(false)
+ }
}
}
expect:
unMuzzledClasses == []
+ nonLazyFields == []
+ unInitFields == []
}
}

MODIFY
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/AdviceReferenceVisitor.java
@@ -80,9 +80,16 @@ public class AdviceReferenceVisitor 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<String>();
+ final Set<String> visitedSources = new HashSet<>();
final Map<String, Reference> references = new HashMap<>();
final Queue<String> instrumentationQueue = new ArrayDeque<>();
@@ -92,10 +99,7 @@ public class AdviceReferenceVisitor extends ClassVisitor {
final String className = instrumentationQueue.remove();
visitedSources.add(className);
try {
- final InputStream in =
- ReferenceMatcher.class
- .getClassLoader()
- .getResourceAsStream(Utils.getResourceName(className));
+ final InputStream in = loader.getResourceAsStream(Utils.getResourceName(className));
try {
final AdviceReferenceVisitor cv = new AdviceReferenceVisitor(null);
final ClassReader reader = new ClassReader(in);
@@ -117,8 +121,10 @@ public class AdviceReferenceVisitor extends ClassVisitor {
}
} finally {
+ if (in != null) {
in.close();
}
+ }
} catch (IOException ioe) {
throw new IllegalStateException(ioe);
}

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
@@ -8,17 +8,10 @@ import net.bytebuddy.dynamic.DynamicType.Builder;
public class MuzzleGradlePlugin implements Plugin {
// TODO:
- // - PR1
- // - Make errorprone happy
- // - Unit tests
- // - Review design and ease of adding future features
- // - better checking
- // - versionScan
// - Optimizations
// - Cache safe and unsafe classloaders
- // - Do reference checking at compile time
+ // - Do reference generation at compile time
// - lazy-load reference muzzle field
- // - Also match interfaces which extend Instrumenter
// - Additional references to check
// - Fields
// - methods
@@ -29,11 +22,9 @@ public class MuzzleGradlePlugin implements Plugin {
// - access flags (including implicit package-private)
// - supertypes
// - Misc
+ // - Also match interfaces which extend Instrumenter
// - Expose config instead of hardcoding datadog namespace (or reconfigure classpath)
- // - Throw checked exception for better logging of mismatches
// - Run muzzle in matching phase (may require a rewrite of the instrumentation api)
- // - Remove ReplaceIsSafeVisitor
- // - IntegrationTest: assert all instrumentation has matcher field
// - Documentation
private static final TypeDescription InstrumenterTypeDesc =
@@ -58,4 +49,16 @@ public class MuzzleGradlePlugin implements Plugin {
public Builder<?> apply(Builder<?> builder, TypeDescription typeDescription) {
return builder.visit(new MuzzleVisitor());
}
+
+ public static class NoOp implements Plugin {
+ @Override
+ public boolean matches(final TypeDescription target) {
+ return false;
+ }
+
+ @Override
+ public Builder<?> apply(Builder<?> builder, TypeDescription typeDescription) {
+ return builder;
+ }
+ }
}

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
package datadog.trace.agent.tooling.muzzle;
import java.util.Collection;
+import java.util.HashMap;
import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
import net.bytebuddy.asm.AsmVisitorWrapper;
import net.bytebuddy.description.field.FieldDescription;
import net.bytebuddy.description.field.FieldList;
@@ -12,8 +15,8 @@ import net.bytebuddy.jar.asm.*;
import net.bytebuddy.pool.TypePool;
/**
- * Add a instrumentationMuzzle field and for each decorator add a final transformer to assert the
- * classpath is safe instrument.
+ * 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.
*/
public class MuzzleVisitor implements AsmVisitorWrapper {
@Override
@@ -40,12 +43,13 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
}
public static class InsertSafetyMatcher extends ClassVisitor {
+ private String instrumentationClassName;
+ private Set<String> adviceClassNames = new HashSet<>();
+
public InsertSafetyMatcher(ClassVisitor classVisitor) {
super(Opcodes.ASM6, classVisitor);
}
- public String instrumentationClassName;
-
@Override
public void visit(
final int version,
@@ -73,10 +77,152 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
return new InsertMuzzleTransformer(methodVisitor);
}
+ public Reference[] generateReferences() {
+ // track sources we've generated references from to avoid recursion
+ final Set<String> referenceSources = new HashSet<>();
+ final Map<String, Reference> references = new HashMap<>();
+
+ for (String adviceClass : adviceClassNames) {
+ if (!referenceSources.contains(adviceClass)) {
+ referenceSources.add(adviceClass);
+ for (Map.Entry<String, Reference> entry :
+ AdviceReferenceVisitor.createReferencesFrom(
+ adviceClass, ReferenceMatcher.class.getClassLoader())
+ .entrySet()) {
+ if (references.containsKey(entry.getKey())) {
+ references.put(
+ entry.getKey(), references.get(entry.getKey()).merge(entry.getValue()));
+ } else {
+ references.put(entry.getKey(), entry.getValue());
+ }
+ }
+ }
+ }
+ return references.values().toArray(new Reference[0]);
+ }
+
@Override
public void visitEnd() {
+ { // generate getInstrumentationMuzzle method
+ /*
+ * private synchronized ReferenceMatcher getInstrumentationMuzzle() {
+ * if (null == this.instrumentationMuzzle) {
+ * this.instrumentationMuzzle = new ReferenceMatcher(new Reference[]{
+ * //reference builders
+ * });
+ * }
+ * return this.instrumentationMuzzle;
+ * }
+ */
+ final MethodVisitor mv =
+ visitMethod(
+ Opcodes.ACC_PRIVATE + Opcodes.ACC_SYNCHRONIZED,
+ "getInstrumentationMuzzle",
+ "()Ldatadog/trace/agent/tooling/muzzle/ReferenceMatcher;",
+ null,
+ null);
+
+ mv.visitCode();
+ final Label start = new Label();
+ final Label ret = new Label();
+ final Label finish = new Label();
+
+ mv.visitLabel(start);
+ mv.visitInsn(Opcodes.ACONST_NULL);
+ mv.visitVarInsn(Opcodes.ALOAD, 0);
+ mv.visitFieldInsn(
+ Opcodes.GETFIELD,
+ instrumentationClassName,
+ "instrumentationMuzzle",
+ "Ldatadog/trace/agent/tooling/muzzle/ReferenceMatcher;");
+ mv.visitJumpInsn(Opcodes.IF_ACMPNE, ret);
+
+ final Reference[] references = generateReferences();
+ mv.visitVarInsn(Opcodes.ALOAD, 0);
+ mv.visitTypeInsn(Opcodes.NEW, "datadog/trace/agent/tooling/muzzle/ReferenceMatcher");
+ mv.visitInsn(Opcodes.DUP);
+ mv.visitIntInsn(Opcodes.BIPUSH, references.length);
+ mv.visitTypeInsn(Opcodes.ANEWARRAY, "datadog/trace/agent/tooling/muzzle/Reference");
+
+ for (int i = 0; i < references.length; ++i) {
+ mv.visitInsn(Opcodes.DUP);
+ mv.visitIntInsn(Opcodes.BIPUSH, i);
+ mv.visitTypeInsn(Opcodes.NEW, "datadog/trace/agent/tooling/muzzle/Reference$Builder");
+ mv.visitInsn(Opcodes.DUP);
+ mv.visitLdcInsn(references[i].getClassName());
+ mv.visitMethodInsn(
+ Opcodes.INVOKESPECIAL,
+ "datadog/trace/agent/tooling/muzzle/Reference$Builder",
+ "<init>",
+ "(Ljava/lang/String;)V",
+ false);
+ if (null != references[i].getSuperName()) {
+ mv.visitLdcInsn(references[i].getSuperName());
+ mv.visitMethodInsn(
+ Opcodes.INVOKEVIRTUAL,
+ "datadog/trace/agent/tooling/muzzle/Reference$Builder",
+ "withSuperName",
+ "(Ljava/lang/String;)Ldatadog/trace/agent/tooling/muzzle/Reference$Builder;",
+ false);
+ }
+ for (String interfaceName : references[i].getInterfaces()) {
+ mv.visitLdcInsn(interfaceName);
+ mv.visitMethodInsn(
+ Opcodes.INVOKEVIRTUAL,
+ "datadog/trace/agent/tooling/muzzle/Reference$Builder",
+ "withInterface",
+ "(Ljava/lang/String;)Ldatadog/trace/agent/tooling/muzzle/Reference$Builder;",
+ false);
+ }
+ for (Reference.Source source : references[i].getSources()) {
+ mv.visitLdcInsn(source.getName());
+ mv.visitIntInsn(Opcodes.BIPUSH, source.getLine());
+ mv.visitMethodInsn(
+ Opcodes.INVOKEVIRTUAL,
+ "datadog/trace/agent/tooling/muzzle/Reference$Builder",
+ "withSource",
+ "(Ljava/lang/String;I)Ldatadog/trace/agent/tooling/muzzle/Reference$Builder;",
+ false);
+ }
+ mv.visitMethodInsn(
+ Opcodes.INVOKEVIRTUAL,
+ "datadog/trace/agent/tooling/muzzle/Reference$Builder",
+ "build",
+ "()Ldatadog/trace/agent/tooling/muzzle/Reference;",
+ false);
+ mv.visitInsn(Opcodes.AASTORE);
+ }
+
+ mv.visitMethodInsn(
+ Opcodes.INVOKESPECIAL,
+ "datadog/trace/agent/tooling/muzzle/ReferenceMatcher",
+ "<init>",
+ "([Ldatadog/trace/agent/tooling/muzzle/Reference;)V",
+ false);
+ mv.visitFieldInsn(
+ Opcodes.PUTFIELD,
+ instrumentationClassName,
+ "instrumentationMuzzle",
+ "Ldatadog/trace/agent/tooling/muzzle/ReferenceMatcher;");
+
+ mv.visitLabel(ret);
+ mv.visitFrame(Opcodes.F_SAME, 1, null, 0, null);
+ mv.visitVarInsn(Opcodes.ALOAD, 0);
+ mv.visitFieldInsn(
+ Opcodes.GETFIELD,
+ instrumentationClassName,
+ "instrumentationMuzzle",
+ "Ldatadog/trace/agent/tooling/muzzle/ReferenceMatcher;");
+ mv.visitInsn(Opcodes.ARETURN);
+ mv.visitLabel(finish);
+
+ mv.visitLocalVariable("this", "L" + instrumentationClassName + ";", null, start, finish, 0);
+ mv.visitMaxs(0, 0); // recomputed
+ mv.visitEnd();
+ }
+
super.visitField(
- Opcodes.ACC_PUBLIC,
+ Opcodes.ACC_PRIVATE + Opcodes.ACC_VOLATILE,
"instrumentationMuzzle",
Type.getDescriptor(ReferenceMatcher.class),
null,
@@ -129,15 +275,17 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
} else if (name.equals("advice")) {
if (STATE == PREVIOUS_INSTRUCTION_GET_CLASS_NAME) {
adviceClassNames.add(lastClassLDC);
+ InsertSafetyMatcher.this.adviceClassNames.add(lastClassLDC);
}
// add last LDC/ToString to adivce list
} else if (name.equals("asDecorator")) {
this.visitVarInsn(Opcodes.ALOAD, 0);
- this.visitFieldInsn(
- Opcodes.GETFIELD,
+ mv.visitMethodInsn(
+ Opcodes.INVOKESPECIAL,
instrumentationClassName,
- "instrumentationMuzzle",
- Type.getDescriptor(ReferenceMatcher.class));
+ "getInstrumentationMuzzle",
+ "()Ldatadog/trace/agent/tooling/muzzle/ReferenceMatcher;",
+ false);
mv.visitIntInsn(Opcodes.BIPUSH, adviceClassNames.size());
mv.visitTypeInsn(Opcodes.ANEWARRAY, "java/lang/String");
int i = 0;
@@ -192,17 +340,10 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
public void visitInsn(final int opcode) {
if (opcode == Opcodes.RETURN) {
super.visitVarInsn(Opcodes.ALOAD, 0);
- mv.visitTypeInsn(Opcodes.NEW, "datadog/trace/agent/tooling/muzzle/ReferenceMatcher");
- mv.visitInsn(Opcodes.DUP);
- mv.visitMethodInsn(
- Opcodes.INVOKESPECIAL,
- "datadog/trace/agent/tooling/muzzle/ReferenceMatcher",
- "<init>",
- "()V",
- false);
+ super.visitInsn(Opcodes.ACONST_NULL);
super.visitFieldInsn(
Opcodes.PUTFIELD,
- instrumentationClassName.replace('.', '/'),
+ instrumentationClassName,
"instrumentationMuzzle",
Type.getDescriptor(ReferenceMatcher.class));
}

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
@@ -26,6 +26,18 @@ public class Reference {
return className;
}
+ public String getSuperName() {
+ return superName;
+ }
+
+ public String[] getInterfaces() {
+ return interfaceNames;
+ }
+
+ public Source[] getSources() {
+ return sources;
+ }
+
/**
* TODO: doc
*
@@ -138,4 +150,35 @@ public class Reference {
}
}
}
+
+ public static class Builder {
+ private final String className;
+ private String superName = null;
+ private Set<String> interfaces = new HashSet<>();
+ private Set<Source> sources = new HashSet<>();
+
+ public Builder(final String className) {
+ this.className = className;
+ }
+
+ public Builder withSuperName(String superName) {
+ this.superName = superName;
+ return this;
+ }
+
+ public Builder withInterface(String interfaceName) {
+ interfaces.add(interfaceName);
+ return this;
+ }
+
+ public Builder withSource(String sourceName, int line) {
+ sources.add(new Source(sourceName, line));
+ return this;
+ }
+
+ public Reference build() {
+ return new Reference(
+ sources.toArray(new Source[0]), className, superName, interfaces.toArray(new String[0]));
+ }
+ }
}

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
@@ -5,6 +5,8 @@ import static net.bytebuddy.dynamic.loading.ClassLoadingStrategy.BOOTSTRAP_LOADE
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.agent.builder.AgentBuilder.Transformer;
import net.bytebuddy.description.type.TypeDescription;
@@ -15,16 +17,14 @@ import net.bytebuddy.utility.JavaModule;
* A bytebuddy matcher that matches if expected references (classes, fields, methods, visibility)
* are present on the classpath.
*/
+@Slf4j
public class ReferenceMatcher implements AgentBuilder.RawMatcher {
// TODO: Cache safe and unsafe classloaders
- // list of unique references (by class name)
- Map<String, Reference> references = new HashMap<>();
- Set<String> referenceSources = new HashSet<>();
+ private final Reference[] references;
- // take a list of references
- public ReferenceMatcher() {
- // TODO: pass in references
+ public ReferenceMatcher(Reference... references) {
+ this.references = references;
}
@Override
@@ -37,49 +37,33 @@ public class ReferenceMatcher implements AgentBuilder.RawMatcher {
return matches(classLoader);
}
+ /**
+ * @param loader Classloader to validate against (or null for bootstrap)
+ * @return true if all references match the classpath of loader
+ */
public boolean matches(ClassLoader loader) {
return getMismatchedReferenceSources(loader).size() == 0;
}
+ /**
+ * @param loader Classloader to validate against (or null for bootstrap)
+ * @return A list of all mismatches between this ReferenceMatcher and loader's classpath.
+ */
public List<Reference.Mismatch> getMismatchedReferenceSources(ClassLoader loader) {
if (loader == BOOTSTRAP_LOADER) {
loader = Utils.getBootstrapProxy();
}
final List<Reference.Mismatch> mismatchedReferences = new ArrayList<>(0);
- for (Reference reference : references.values()) {
+ for (Reference reference : references) {
mismatchedReferences.addAll(reference.checkMatch(loader));
}
- // TODO: log mismatches
- /*
- if (mismatchedReferences.size() > 0) {
- System.out.println(mismatchedReferences.size() + " mismatches on classloader " + loader);
- for (Reference.Mismatch mismatch : mismatchedReferences) {
- // TODO: log more info about why mismatch occurred. Missing method, missing field, signature mismatch.
- System.out.println("--" + mismatch.toString());
- }
- }
- */
return mismatchedReferences;
}
+ /**
+ * Create a bytebuddy matcher which throws a MismatchException when there are mismatches with the classloader under transformation.
+ */
public Transformer assertSafeTransformation(String... adviceClassNames) {
- // load or check cache for advice class names
- for (String adviceClass : adviceClassNames) {
- if (!referenceSources.contains(adviceClass)) {
- referenceSources.add(adviceClass);
- for (Map.Entry<String, Reference> entry :
- AdviceReferenceVisitor.createReferencesFrom(
- adviceClass, ReferenceMatcher.class.getClassLoader())
- .entrySet()) {
- if (references.containsKey(entry.getKey())) {
- references.put(entry.getKey(), references.get(entry.getKey()).merge(entry.getValue()));
- } else {
- references.put(entry.getKey(), entry.getValue());
- }
- }
- }
- }
-
return new Transformer() {
@Override
public DynamicType.Builder<?> transform(
@@ -91,7 +75,6 @@ public class ReferenceMatcher implements AgentBuilder.RawMatcher {
if (mismatches.size() == 0) {
return builder;
} else {
- // TODO: make mismatch exception type and add more descriptive logging at listener level.
throw new MismatchException(classLoader, mismatches);
}
}

MODIFY
dd-java-agent/instrumentation/instrumentation.gradle
dd-java-agent/instrumentation/instrumentation.gradle
@@ -14,11 +14,22 @@ Project instr_project = project
subprojects { subProj ->
if (subProj.getParent() == instr_project) {
apply plugin: "net.bytebuddy.byte-buddy"
+
subProj.byteBuddy {
transformation {
- tasks = ["compileJava"]
- plugin = "datadog.trace.agent.tooling.muzzle.MuzzleGradlePlugin"
- classPath = project(':dd-java-agent:agent-tooling').configurations.instrumentationMuzzle
+ // Applying NoOp optimizes build by applying bytebuddy plugin to only compileJava task
+ tasks = ['compileJava']
+ plugin = 'datadog.trace.agent.tooling.muzzle.MuzzleGradlePlugin$NoOp'
+ }
+ }
+
+ subProj.afterEvaluate {
+ subProj.byteBuddy {
+ transformation {
+ tasks = ['compileJava']
+ plugin = 'datadog.trace.agent.tooling.muzzle.MuzzleGradlePlugin'
+ classPath = project(':dd-java-agent:agent-tooling').configurations.instrumentationMuzzle + subProj.configurations.compile + subProj.sourceSets.main.output
+ }
}
}

MODIFY
dd-java-agent/testing/src/test/groovy/muzzle/AdviceReferenceVisitorTest.groovy
dd-java-agent/testing/src/test/groovy/muzzle/AdviceReferenceVisitorTest.groovy
@@ -4,6 +4,7 @@ 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.ReferenceMatcher
+import datadog.trace.agent.tooling.muzzle.Reference
class AdviceReferenceVisitorTest extends AgentTestRunner {
@@ -21,8 +22,9 @@ class AdviceReferenceVisitorTest extends AgentTestRunner {
def "match safe classpaths"() {
setup:
- ReferenceMatcher refMatcher = new ReferenceMatcher()
- refMatcher.assertSafeTransformation(AdviceClass.getName())
+ Reference[] refs = AdviceReferenceVisitor.createReferencesFrom(AdviceClass.getName(), this.getClass().getClassLoader()).values().toArray(new Reference[0])
+ ReferenceMatcher refMatcher = new ReferenceMatcher(refs)
+ refMatcher.assertSafeTransformation()
ClassLoader safeClassloader = new URLClassLoader([TestUtils.createJarWithClasses(AdviceClass$A,
AdviceClass$SomeInterface,
AdviceClass$SomeImplementation)] as URL[],

Generate at compile-time and lazy-load reference matchers