Merge pull request #383 from DataDog/ark/muzzle-hookup

Enable Muzzle validation for Default Instrumentation Match phase

Author: realark

.circleci/config.yml

@@ -187,8 +187,8 @@ jobs:
             - dd-trace-java-version-scan
 
       - run:
-          name: Verify Version Scan
-          command: ./gradlew verifyVersionScan --parallel --stacktrace --no-daemon --max-workers=6
+          name: Verify Version Scan and Muzzle
+          command: ./gradlew verifyVersionScan muzzle --parallel --stacktrace --no-daemon --max-workers=6
 
       - save_cache:
           key: dd-trace-java-version-scan-{{ checksum "dd-trace-java.gradle" }}

buildSrc/src/main/groovy/MuzzleExtension.groovy

@@ -0,0 +1,4 @@
+class MuzzleExtension {
+  String group
+  String module
+}

buildSrc/src/main/groovy/MuzzlePlugin.groovy

@@ -0,0 +1,58 @@
+import org.gradle.api.Plugin
+import org.gradle.api.Project
+
+import java.lang.reflect.Method
+
+/**
+ * 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) {
+    def bootstrapProject = project.rootProject.getChildProjects().get('dd-java-agent').getChildProjects().get('agent-bootstrap')
+    def toolingProject = project.rootProject.getChildProjects().get('dd-java-agent').getChildProjects().get('agent-tooling')
+    project.extensions.create("muzzle", MuzzleExtension)
+    def muzzle = project.task('muzzle') {
+      group = 'Muzzle'
+      description = "Run instrumentation muzzle on compile time dependencies"
+      doLast {
+        List<URL> userUrls = new ArrayList<>()
+        project.getLogger().info("Creating user classpath for: " + project.getName())
+        for (File f : project.configurations.compileOnly.getFiles()) {
+          project.getLogger().info( '--' + f)
+          userUrls.add(f.toURI().toURL())
+        }
+        for (File f : bootstrapProject.sourceSets.main.runtimeClasspath.getFiles()) {
+          project.getLogger().info( '--' + f)
+          userUrls.add(f.toURI().toURL())
+        }
+        final ClassLoader userCL = new URLClassLoader(userUrls.toArray(new URL[0]), (ClassLoader) null)
+
+        project.getLogger().info("Creating dd classpath for: " + project.getName())
+        Set<URL> ddUrls = new HashSet<>()
+        for (File f : toolingProject.sourceSets.main.runtimeClasspath.getFiles()) {
+          project.getLogger().info( '--' + f)
+          ddUrls.add(f.toURI().toURL())
+        }
+        for(File f : project.sourceSets.main.runtimeClasspath.getFiles()) {
+          project.getLogger().info( '--' + f)
+          ddUrls.add(f.toURI().toURL())
+        }
+
+        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.MuzzleVersionScanPlugin')
+          .getMethod('assertInstrumentationNotMuzzled', ClassLoader.class)
+        assertionMethod.invoke(null, userCL)
+      }
+    }
+    // project.tasks.muzzle.dependsOn(bootstrapProject.tasks.shadowJar)
+    project.tasks.muzzle.dependsOn(bootstrapProject.tasks.compileJava)
+    project.tasks.muzzle.dependsOn(toolingProject.tasks.compileJava)
+    project.afterEvaluate {
+      project.tasks.muzzle.dependsOn(project.tasks.compileJava)
+    }
+  }
+}

buildSrc/src/main/resources/META-INF/gradle-plugins/muzzle.properties

@@ -0,0 +1 @@
+implementation-class=MuzzlePlugin

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 == []
   }
-  */
 
 }

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

@@ -8,8 +8,6 @@
 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,19 +89,11 @@ public void onError(
         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());
-      }
+      log.debug(
+          "Failed to handle {} for transformation on classloader {}: {}",
+          typeName,
+          classLoader,
+          throwable.getMessage());
     }
 
     @Override

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

@@ -3,14 +3,19 @@
 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 lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.utility.JavaModule;
 
 /**
  * Built-in bytebuddy-based instrumentation for the datadog javaagent.
@@ -42,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.
@@ -74,6 +81,35 @@ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
       AgentBuilder.Identified.Extendable advice =
           agentBuilder
               .type(typeMatcher(), classLoaderMatcher())
+              .and(
+                  new AgentBuilder.RawMatcher() {
+                    @Override
+                    public boolean matches(
+                        TypeDescription typeDescription,
+                        ClassLoader classLoader,
+                        JavaModule module,
+                        Class<?> classBeingRedefined,
+                        ProtectionDomain protectionDomain) {
+                      // Optimization: calling getInstrumentationMuzzle() inside this method prevents unnecessary loading of muzzle references during agentBuilder setup.
+                      final ReferenceMatcher muzzle = getInstrumentationMuzzle();
+                      if (null != muzzle) {
+                        List<Reference.Mismatch> mismatches =
+                            muzzle.getMismatchedReferenceSources(classLoader);
+                        if (mismatches.size() > 0) {
+                          log.debug(
+                              "Instrumentation muzzled: {} -- {} on {}",
+                              instrumentationPrimaryName,
+                              this.getClass().getName(),
+                              classLoader);
+                        }
+                        for (Reference.Mismatch mismatch : mismatches) {
+                          log.debug("-- {}", mismatch);
+                        }
+                        return mismatches.size() == 0;
+                      }
+                      return true;
+                    }
+                  })
               .transform(DDTransformers.defaultTransformers());
       final String[] helperClassNames = helperClassNames();
       if (helperClassNames.length > 0) {
@@ -85,6 +121,15 @@ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
       return advice.asDecorator();
     }
 
+    /**
+     * This method is implemented dynamically by compile-time bytecode transformations.
+     *
+     * <p>{@see datadog.trace.agent.tooling.muzzle.MuzzleGradlePlugin}
+     */
+    protected ReferenceMatcher getInstrumentationMuzzle() {
+      return null;
+    }
+
     @Override
     public String[] helperClassNames() {
       return new String[0];
@@ -105,11 +150,13 @@ protected boolean defaultEnabled() {
       return getConfigEnabled("dd.integrations.enabled", true);
     }
 
-    protected static String getPropOrEnv(final String name) {
+    // TODO: move common config helpers to Utils
+
+    public static String getPropOrEnv(final String name) {
       return System.getProperty(name, System.getenv(propToEnvName(name)));
     }
 
-    private static String propToEnvName(final String name) {
+    public static String propToEnvName(final String name) {
       return name.toUpperCase().replace(".", "_");
     }
   }

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

@@ -6,51 +6,32 @@
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.DynamicType.Builder;
 
+/** Bytebuddy gradle plugin which creates muzzle-references at compile time. */
 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
-  //     - 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
-
-  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;
+    TypeDefinition instrumenter = null == target ? null : target.getSuperClass();
     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) {
     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) {

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

@@ -0,0 +1,86 @@
+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() {}
+}