Hook up muzzle to Instrumenter.Default's matcher
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
@@ -8,13 +8,17 @@ import spock.lang.Specification
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())) {
+ for (Object instrumenter : ServiceLoader.load(IntegrationTestUtils.getAgentClassLoader().loadClass("datadog.trace.agent.tooling.Instrumenter"), IntegrationTestUtils.getAgentClassLoader())) {
+ if (instrumenter.getClass().getName().endsWith("TraceConfigInstrumentation")) {
+ // TraceConfigInstrumentation doesn't do muzzle checks
+ // check on TracerClassInstrumentation instead
+ instrumenter = IntegrationTestUtils.getAgentClassLoader().loadClass(instrumenter.getClass().getName() + '$TracerClassInstrumentation').newInstance()
+ }
Field f
Method m
try {
@@ -45,6 +49,5 @@ class MuzzleBytecodeTransformTest extends Specification {
nonLazyFields == []
unInitFields == []
}
- */
}

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java
@@ -8,8 +8,6 @@ import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
import static net.bytebuddy.matcher.ElementMatchers.named;
import static net.bytebuddy.matcher.ElementMatchers.not;
-import datadog.trace.agent.tooling.muzzle.Reference.Mismatch;
-import datadog.trace.agent.tooling.muzzle.ReferenceMatcher.MismatchException;
import java.lang.instrument.Instrumentation;
import java.util.ServiceLoader;
import lombok.extern.slf4j.Slf4j;
@@ -91,20 +89,12 @@ public class AgentInstaller {
final JavaModule module,
final boolean loaded,
final Throwable throwable) {
- if (throwable instanceof MismatchException) {
- final MismatchException mismatchException = (MismatchException) throwable;
- log.debug("{}", mismatchException.getMessage());
- for (final Mismatch mismatch : mismatchException.getMismatches()) {
- log.debug("--{}", mismatch);
- }
- } else {
log.debug(
"Failed to handle {} for transformation on classloader {}: {}",
typeName,
classLoader,
throwable.getMessage());
}
- }
@Override
public void onTransformation(

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
@@ -3,13 +3,14 @@ package datadog.trace.agent.tooling;
import static datadog.trace.agent.tooling.Utils.getConfigEnabled;
import static net.bytebuddy.matcher.ElementMatchers.any;
+import datadog.trace.agent.tooling.muzzle.Reference;
+import datadog.trace.agent.tooling.muzzle.ReferenceMatcher;
import java.security.ProtectionDomain;
import java.util.Arrays;
import java.util.HashSet;
+import java.util.List;
import java.util.Map;
import java.util.Set;
-
-import datadog.trace.agent.tooling.muzzle.ReferenceMatcher;
import lombok.extern.slf4j.Slf4j;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.description.type.TypeDescription;
@@ -46,11 +47,13 @@ public interface Instrumenter {
@Slf4j
abstract class Default implements Instrumenter {
private final Set<String> instrumentationNames;
+ private final String instrumentationPrimaryName;
protected final boolean enabled;
public Default(final String instrumentationName, final String... additionalNames) {
this.instrumentationNames = new HashSet<>(Arrays.asList(additionalNames));
instrumentationNames.add(instrumentationName);
+ instrumentationPrimaryName = instrumentationName;
// If default is enabled, we want to enable individually,
// if default is disabled, we want to disable individually.
@@ -78,11 +81,28 @@ public interface Instrumenter {
AgentBuilder.Identified.Extendable advice =
agentBuilder
.type(typeMatcher(), classLoaderMatcher())
- .and(new AgentBuilder.RawMatcher() {
+ .and(
+ new AgentBuilder.RawMatcher() {
@Override
- public boolean matches(TypeDescription typeDescription, ClassLoader classLoader, JavaModule module, Class<?> classBeingRedefined, ProtectionDomain protectionDomain) {
+ public boolean matches(
+ TypeDescription typeDescription,
+ ClassLoader classLoader,
+ JavaModule module,
+ Class<?> classBeingRedefined,
+ ProtectionDomain protectionDomain) {
// Optimization: calling getMuzzleReferenceMatcher() inside this method prevents unnecessary loading of muzzle references during agentBuilder setup.
- return getInstrumentationMuzzle().matches(classLoader);
+ final ReferenceMatcher muzzle = getInstrumentationMuzzle();
+ if (null != muzzle) {
+ List<Reference.Mismatch> mismatches = muzzle.getMismatchedReferenceSources(classLoader);
+ if (mismatches.size() > 0) {
+ log.debug("Instrumentation muzzled: {} on {}", instrumentationPrimaryName, classLoader);
+ }
+ for (Reference.Mismatch mismatch : mismatches) {
+ log.debug("-- {}", mismatch);
+ }
+ return mismatches.size() == 0;
+ }
+ return true;
}
})
.transform(DDTransformers.defaultTransformers());
@@ -99,7 +119,7 @@ public interface Instrumenter {
/**
* This method is implemented dynamically by compile-time bytecode transformations.
*
- * TODO bytecode magic and documentation
+ * <p>TODO bytecode magic and documentation
*/
// TODO: Make final
protected ReferenceMatcher getInstrumentationMuzzle() {

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,10 +8,6 @@ import net.bytebuddy.dynamic.DynamicType.Builder;
public class MuzzleGradlePlugin implements Plugin {
// TODO:
- // - Optimizations
- // - Cache safe and unsafe classloaders
- // - Do reference generation at compile time
- // - lazy-load reference muzzle field
// - Additional references to check
// - Fields
// - methods
@@ -26,30 +22,39 @@ public class MuzzleGradlePlugin implements Plugin {
// - 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 InstrumenterTypeDesc =
- new TypeDescription.ForLoadedType(Instrumenter.class);
+ private static final TypeDescription DefaultInstrumenterTypeDesc =
+ new TypeDescription.ForLoadedType(Instrumenter.Default.class);
@Override
public boolean matches(final TypeDescription target) {
- // AutoService annotation is not retained at runtime. Check for instrumenter supertype
+ // AutoService annotation is not retained at runtime. Check for Instrumenter.Default supertype
boolean isInstrumenter = false;
TypeDefinition instrumenter = target;
while (instrumenter != null) {
- if (instrumenter.getInterfaces().contains(InstrumenterTypeDesc)) {
+ if (instrumenter.equals(DefaultInstrumenterTypeDesc)) {
isInstrumenter = true;
break;
}
instrumenter = instrumenter.getSuperClass();
}
- // return isInstrumenter;
- return false;
+ return isInstrumenter;
}
@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());
}
+ }
public static class NoOp implements Plugin {
@Override

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 datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.Utils;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
@@ -15,8 +16,9 @@ import net.bytebuddy.jar.asm.*;
import net.bytebuddy.pool.TypePool;
/**
- * 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.
+ * 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.
*/
public class MuzzleVisitor implements AsmVisitorWrapper {
@Override
@@ -44,7 +46,7 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
public static class InsertSafetyMatcher extends ClassVisitor {
private String instrumentationClassName;
- private Set<String> adviceClassNames = new HashSet<>();
+ private Instrumenter.Default instrumenter;
public InsertSafetyMatcher(ClassVisitor classVisitor) {
super(Opcodes.ASM6, classVisitor);
@@ -59,6 +61,17 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
final String superName,
final String[] interfaces) {
this.instrumentationClassName = name;
+ try {
+ instrumenter =
+ (Instrumenter.Default)
+ MuzzleVisitor.class
+ .getClassLoader()
+ .loadClass(Utils.getClassName(instrumentationClassName))
+ .getDeclaredConstructor()
+ .newInstance();
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
super.visit(version, access, name, signature, superName, interfaces);
}
@@ -74,13 +87,18 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
if ("<init>".equals(name)) {
methodVisitor = new InitializeFieldVisitor(methodVisitor);
}
- return new InsertMuzzleTransformer(methodVisitor);
+ return 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<>();
+ final Set<String> adviceClassNames = new HashSet<>();
+
+ for (String adviceClassName : instrumenter.transformers().values()) {
+ adviceClassNames.add(adviceClassName);
+ }
for (String adviceClass : adviceClassNames) {
if (!referenceSources.contains(adviceClass)) {
@@ -105,9 +123,10 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
public void visitEnd() {
{ // generate getInstrumentationMuzzle method
/*
- * private synchronized ReferenceMatcher getInstrumentationMuzzle() {
+ * protected synchronized ReferenceMatcher getInstrumentationMuzzle() {
* if (null == this.instrumentationMuzzle) {
- * this.instrumentationMuzzle = new ReferenceMatcher(new Reference[]{
+ * this.instrumentationMuzzle = new ReferenceMatcher(this.helperClassNames(),
+ * new Reference[]{
* //reference builders
* });
* }
@@ -116,7 +135,7 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
*/
final MethodVisitor mv =
visitMethod(
- Opcodes.ACC_PRIVATE + Opcodes.ACC_SYNCHRONIZED,
+ Opcodes.ACC_PROTECTED + Opcodes.ACC_SYNCHRONIZED,
"getInstrumentationMuzzle",
"()Ldatadog/trace/agent/tooling/muzzle/ReferenceMatcher;",
null,
@@ -137,10 +156,21 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
"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.visitVarInsn(Opcodes.ALOAD, 0);
+ mv.visitMethodInsn(
+ Opcodes.INVOKEVIRTUAL,
+ instrumentationClassName,
+ "helperClassNames",
+ "()[Ljava/lang/String;",
+ false);
+
+ final Reference[] references = generateReferences();
mv.visitIntInsn(Opcodes.BIPUSH, references.length);
mv.visitTypeInsn(Opcodes.ANEWARRAY, "datadog/trace/agent/tooling/muzzle/Reference");
@@ -275,7 +305,7 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
Opcodes.INVOKESPECIAL,
"datadog/trace/agent/tooling/muzzle/ReferenceMatcher",
"<init>",
- "([Ldatadog/trace/agent/tooling/muzzle/Reference;)V",
+ "([Ljava/lang/String;[Ldatadog/trace/agent/tooling/muzzle/Reference;)V",
false);
mv.visitFieldInsn(
Opcodes.PUTFIELD,
@@ -308,106 +338,6 @@ public class MuzzleVisitor implements AsmVisitorWrapper {
super.visitEnd();
}
- /**
- * Changes this:<br>
- * &nbsp.transform(DDAdvice.create().advice(named("fooMethod", FooAdvice.class.getname())))
- * &nbsp.asDecorator(); Into this:<br>
- * &nbsp.transform(DDAdvice.create().advice(named("fooMethod", FooAdvice.class.getname())))
- * &nbsp.transform(this.instrumentationMuzzle.assertSafeTransformation("foo.package.FooAdvice"));
- * &nbsp.asDecorator(); className)
- */
- public class InsertMuzzleTransformer extends MethodVisitor {
- // it would be nice to manage the state with an enum, but that requires this class to be non-static
- private final int INIT = 0;
- // SomeClass
- private final int PREVIOUS_INSTRUCTION_LDC = 1;
- // SomeClass.getName()
- private final int PREVIOUS_INSTRUCTION_GET_CLASS_NAME = 2;
-
- private String lastClassLDC = null;
-
- private Collection<String> adviceClassNames = new HashSet<>();
- private int STATE = INIT;
-
- public InsertMuzzleTransformer(MethodVisitor methodVisitor) {
- super(Opcodes.ASM6, methodVisitor);
- }
-
- public void reset() {
- STATE = INIT;
- lastClassLDC = null;
- adviceClassNames.clear();
- }
-
- @Override
- public void visitMethodInsn(
- final int opcode,
- final String owner,
- final String name,
- final String descriptor,
- final boolean isInterface) {
- if (name.equals("getName")) {
- if (STATE == PREVIOUS_INSTRUCTION_LDC) {
- STATE = PREVIOUS_INSTRUCTION_GET_CLASS_NAME;
- }
- } 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);
- mv.visitMethodInsn(
- Opcodes.INVOKESPECIAL,
- instrumentationClassName,
- "getInstrumentationMuzzle",
- "()Ldatadog/trace/agent/tooling/muzzle/ReferenceMatcher;",
- false);
- mv.visitIntInsn(Opcodes.BIPUSH, adviceClassNames.size());
- mv.visitTypeInsn(Opcodes.ANEWARRAY, "java/lang/String");
- int i = 0;
- for (String adviceClassName : adviceClassNames) {
- mv.visitInsn(Opcodes.DUP);
- mv.visitIntInsn(Opcodes.BIPUSH, i);
- mv.visitLdcInsn(adviceClassName);
- mv.visitInsn(Opcodes.AASTORE);
- ++i;
- }
- mv.visitMethodInsn(
- Opcodes.INVOKEVIRTUAL,
- "datadog/trace/agent/tooling/muzzle/ReferenceMatcher",
- "assertSafeTransformation",
- "([Ljava/lang/String;)Lnet/bytebuddy/agent/builder/AgentBuilder$Transformer;",
- false);
- mv.visitMethodInsn(
- Opcodes.INVOKEINTERFACE,
- "net/bytebuddy/agent/builder/AgentBuilder$Identified$Narrowable",
- "transform",
- "(Lnet/bytebuddy/agent/builder/AgentBuilder$Transformer;)Lnet/bytebuddy/agent/builder/AgentBuilder$Identified$Extendable;",
- true);
- reset();
- } else {
- STATE = INIT;
- lastClassLDC = null;
- }
- super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
- }
-
- @Override
- public void visitLdcInsn(final Object value) {
- if (value instanceof Type) {
- Type type = (Type) value;
- if (type.getSort() == Type.OBJECT) {
- lastClassLDC = type.getClassName();
- STATE = PREVIOUS_INSTRUCTION_LDC;
- type.getClassName();
- }
- }
- super.visitLdcInsn(value);
- }
- }
-
/** Append a field initializer to the end of a method. */
public class InitializeFieldVisitor extends MethodVisitor {
public InitializeFieldVisitor(MethodVisitor methodVisitor) {

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
@@ -7,9 +7,7 @@ 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;
-import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.utility.JavaModule;
/**
@@ -21,9 +19,15 @@ public class ReferenceMatcher implements AgentBuilder.RawMatcher {
private final Map<ClassLoader, List<Reference.Mismatch>> mismatchCache =
Collections.synchronizedMap(new WeakHashMap<ClassLoader, List<Reference.Mismatch>>());
private final Reference[] references;
+ private final Set<String> helperClassNames;
public ReferenceMatcher(Reference... references) {
+ this(new String[0], references);
+ }
+
+ public ReferenceMatcher(String[] helperClassNames, Reference[] references) {
this.references = references;
+ this.helperClassNames = new HashSet<>(Arrays.asList(helperClassNames));
}
@Override
@@ -52,54 +56,23 @@ public class ReferenceMatcher implements AgentBuilder.RawMatcher {
if (loader == BOOTSTRAP_LOADER) {
loader = Utils.getBootstrapProxy();
}
- final List<Reference.Mismatch> mismatchedReferences = new ArrayList<>(0);
- for (Reference reference : references) {
- mismatchedReferences.addAll(reference.checkMatch(loader));
- }
- return mismatchedReferences;
- }
-
- /**
- * Create a bytebuddy matcher which throws a MismatchException when there are mismatches with the
- * classloader under transformation.
- */
- public Transformer assertSafeTransformation(String... adviceClassNames) {
- return new Transformer() {
- @Override
- public DynamicType.Builder<?> transform(
- DynamicType.Builder<?> builder,
- TypeDescription typeDescription,
- ClassLoader classLoader,
- JavaModule module) {
- if (classLoader == BOOTSTRAP_LOADER) {
- classLoader = Utils.getBootstrapProxy();
- }
- List<Reference.Mismatch> mismatches =
- mismatchCache.get(getMismatchedReferenceSources(classLoader));
+ List<Reference.Mismatch> mismatches = mismatchCache.get(loader);
if (null == mismatches) {
- // okay if entered by multiple callers during initialization
- mismatches = getMismatchedReferenceSources(classLoader);
- mismatchCache.put(classLoader, mismatches);
- }
- if (mismatches.size() == 0) {
- return builder;
- } else {
- throw new MismatchException(classLoader, mismatches);
+ synchronized (loader) {
+ mismatches = mismatchCache.get(loader);
+ if (null == mismatches) {
+ mismatches = new ArrayList<>(0);
+ for (Reference reference : references) {
+ // 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));
}
}
- };
+ mismatchCache.put(loader, mismatches);
}
-
- public static class MismatchException extends RuntimeException {
- private final List<Reference.Mismatch> mismatches;
-
- public MismatchException(ClassLoader classLoader, List<Reference.Mismatch> mismatches) {
- super(mismatches.size() + " mismatches on classloader: " + classLoader);
- this.mismatches = mismatches;
}
-
- public List<Reference.Mismatch> getMismatches() {
- return this.mismatches;
}
+ return mismatches;
}
}

MODIFY
dd-java-agent/instrumentation/akka-http-10.0/src/lagomTest/java/EchoService.java
dd-java-agent/instrumentation/akka-http-10.0/src/lagomTest/java/EchoService.java
@@ -10,6 +10,7 @@ public interface EchoService extends Service {
ServiceCall<Source<String, NotUsed>, Source<String, NotUsed>> error();
+ @Override
default Descriptor descriptor() {
return named("echo").withCalls(namedCall("echo", this::echo), namedCall("error", this::error));
}

MODIFY
dd-java-agent/instrumentation/trace-annotation/src/main/java/datadog/trace/instrumentation/trace_annotation/TraceConfigInstrumentation.java
dd-java-agent/instrumentation/trace-annotation/src/main/java/datadog/trace/instrumentation/trace_annotation/TraceConfigInstrumentation.java
@@ -104,6 +104,11 @@ public class TraceConfigInstrumentation implements Instrumenter {
private final String className;
private final Set<String> methodNames;
+ /** No-arg constructor only used by muzzle and tests. */
+ public TracerClassInstrumentation() {
+ this("noop", Collections.singleton("noop"));
+ }
+
public TracerClassInstrumentation(String className, Set<String> methodNames) {
super("trace", "trace-config");
this.className = className;

MODIFY
dd-java-agent/testing/src/test/groovy/muzzle/AdviceReferenceVisitorTest.groovy
dd-java-agent/testing/src/test/groovy/muzzle/AdviceReferenceVisitorTest.groovy
@@ -24,7 +24,6 @@ class AdviceReferenceVisitorTest extends AgentTestRunner {
setup:
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[],

Hook up muzzle to Instrumenter.Default's matcher