Inject helper classes into application’s classloader

This is important for things like Spring Boot’s executable jar classloader which keeps all the jars individually rather than on the system classloader.

Author: tylerbenson

dd-java-agent-ittests/src/test/java/com/datadoghq/trace/agent/test/integration/ApacheHTTPClientTest.java

@@ -9,6 +9,10 @@ public class ApacheHTTPClientTest {
 
   @Test
   public void test() throws Exception {
+    // Since the HttpClientBuilder initializer doesn't work, invoke manually.
+    Class.forName("com.datadoghq.trace.agent.InstrumentationRulesManager")
+        .getMethod("registerClassLoad")
+        .invoke(null);
 
     final HttpClientBuilder builder = HttpClientBuilder.create();
     assertThat(builder.getClass().getSimpleName()).isEqualTo("TracingHttpClientBuilder");

dd-java-agent/dd-java-agent.gradle

@@ -13,14 +13,12 @@ whitelistedInstructionClasses += whitelistedBranchClasses += [
   'com.datadoghq.trace.agent.*',
   'com.datadoghq.trace.agent.integration.*',
   'io.opentracing.contrib.mongo.TracingCommandListenerFactory',
+  'io.opentracing.contrib.*',
 ]
 
 dependencies {
   compile project(':dd-trace')
   compile project(':dd-trace-annotations')
-  compile(project(path: ':dd-java-agent:integrations:helpers', configuration: "shadow")) {
-    transitive = false
-  }
 
   compile group: 'org.jboss.byteman', name: 'byteman', version: '4.0.0-BETA5'
 
@@ -35,12 +33,24 @@ dependencies {
   testCompile group: 'org.mongodb', name: 'mongo-java-driver', version: '3.4.2'
   testCompile group: 'io.opentracing', name: 'opentracing-mock', version: '0.30.0'
 
+  testCompile(project(path: ':dd-java-agent:integrations:helpers', configuration: "shadow")) {
+    transitive = false
+  }
+
   // Not bundled in with the agent.  Usage requires being on the app's classpath (eg. Spring Boot's executable jar)
   compileOnly group: 'io.opentracing.contrib', name: 'opentracing-jdbc', version: '0.0.3'
 }
 
 project(':dd-java-agent:integrations:helpers').afterEvaluate { helperProject ->
-  project.compileJava.dependsOn helperProject.tasks.shadowJar
+  project.processResources {
+    from(helperProject.tasks.shadowJar)
+    rename {
+      it.startsWith("helpers") && it.endsWith(".jar") ?
+        "helpers.jar.zip" :
+        it
+    }
+  }
+  project.processResources.dependsOn helperProject.tasks.shadowJar
 }
 
 jar {

dd-java-agent/src/main/java/com/datadoghq/trace/agent/AgentRulesManager.java

@@ -25,7 +25,7 @@ public class AgentRulesManager {
   private static final String ddTraceVersion = DDTraceInfo.VERSION;
   private static final String ddTraceAnnotationsVersion = DDTraceAnnotationsInfo.VERSION;
 
-  private static final String SPRING_BOOT_RULE = "spring-boot-rule.btm";
+  private static final String INITIALIZER_RULES = "initializer-rules.btm";
 
   protected static volatile AgentRulesManager INSTANCE;
 
@@ -61,7 +61,7 @@ public static void initialize(final Retransformer trans) {
 
     INSTANCE = manager;
 
-    manager.loadRules(SPRING_BOOT_RULE, ClassLoader.getSystemClassLoader());
+    manager.loadRules(INITIALIZER_RULES, ClassLoader.getSystemClassLoader());
     manager.traceAnnotationsManager.initialize();
   }
 

dd-java-agent/src/main/java/com/datadoghq/trace/agent/ClassLoaderIntegrationInjector.java

@@ -0,0 +1,75 @@
+package com.datadoghq.trace.agent;
+
+import com.google.common.collect.Maps;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.util.Map;
+import java.util.zip.ZipEntry;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+public class ClassLoaderIntegrationInjector {
+  private final Map<ZipEntry, byte[]> entries;
+  private final Map<ClassLoader, Method> invocationPoints = Maps.newConcurrentMap();
+
+  public ClassLoaderIntegrationInjector(Map<ZipEntry, byte[]> entries) {
+    this.entries = entries;
+  }
+
+  public void inject(ClassLoader cl) {
+    try {
+      Method inovcationPoint = getInovcationPoint(cl);
+      Map<ZipEntry, byte[]> toInject = Maps.newHashMap(entries);
+      Map<ZipEntry, byte[]> injectedEntries = Maps.newHashMap();
+      boolean successfulyAdded = true;
+      while (!toInject.isEmpty() && successfulyAdded) {
+        log.debug("Attempting to inject {} entries into {}", toInject.size(), cl);
+        successfulyAdded = false;
+        for (Map.Entry<ZipEntry, byte[]> entry : toInject.entrySet()) {
+          String name = entry.getKey().getName();
+          if (!name.endsWith(".class")) {
+            continue;
+          }
+          byte[] bytes = entry.getValue();
+          try {
+            inovcationPoint.invoke(cl, bytes, 0, bytes.length);
+            injectedEntries.put(entry.getKey(), entry.getValue());
+            successfulyAdded = true;
+          } catch (InvocationTargetException e) {
+            log.debug("Error calling 'defineClass' method on {} for entry {}", cl, entry);
+          }
+        }
+        toInject.keySet().removeAll(injectedEntries.keySet());
+      }
+
+    } catch (NoSuchMethodException e) {
+      log.error("Error getting 'defineClass' method from {}", cl);
+    } catch (IllegalAccessException e) {
+      log.error("Error accessing 'defineClass' method on {}", cl);
+    }
+  }
+
+  private Method getInovcationPoint(ClassLoader cl) throws NoSuchMethodException {
+    if (invocationPoints.containsKey(invocationPoints)) {
+      return invocationPoints.get(invocationPoints);
+    }
+    Class<?> clazz = cl.getClass();
+    NoSuchMethodException firstException = null;
+    while (clazz != null) {
+      try {
+        // defineClass is protected so we may need to check up the class hierarchy.
+        Method invocationPoint =
+            clazz.getDeclaredMethod("defineClass", byte[].class, int.class, int.class);
+        invocationPoint.setAccessible(true);
+        invocationPoints.put(cl, invocationPoint);
+        return invocationPoint;
+      } catch (NoSuchMethodException e) {
+        if (firstException == null) {
+          firstException = e;
+        }
+        clazz = clazz.getSuperclass();
+      }
+    }
+    throw firstException;
+  }
+}

dd-java-agent/src/main/java/com/datadoghq/trace/agent/InstrumentationChecker.java

@@ -3,15 +3,11 @@
 import com.datadoghq.trace.resolver.FactoryUtils;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.core.type.TypeReference;
-import java.io.File;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
 import lombok.Data;
 import lombok.extern.slf4j.Slf4j;
 
@@ -23,105 +19,23 @@
 public class InstrumentationChecker {
 
   private static final String CONFIG_FILE = "dd-trace-supported-framework";
-  private static InstrumentationChecker INSTANCE;
 
   private final Map<String, List<ArtifactSupport>> rules;
-  private final Map<String, String> frameworks;
-
-  private final ClassLoader classLoader;
 
   /* For testing purpose */
   InstrumentationChecker(
       final Map<String, List<ArtifactSupport>> rules, final Map<String, String> frameworks) {
     this.rules = rules;
-    this.frameworks = frameworks;
-    this.classLoader = ClassLoader.getSystemClassLoader();
-    INSTANCE = this;
   }
 
-  private InstrumentationChecker(final ClassLoader classLoader) {
-    this.classLoader = classLoader;
+  public InstrumentationChecker() {
     rules =
         FactoryUtils.loadConfigFromResource(
             CONFIG_FILE, new TypeReference<Map<String, List<ArtifactSupport>>>() {});
-    frameworks = scanLoadedLibraries();
-  }
-
-  /**
-   * Return a list of unsupported rules regarding loading deps
-   *
-   * @return the list of unsupported rules
-   * @param classLoader
-   */
-  public static synchronized List<String> getUnsupportedRules(final ClassLoader classLoader) {
-
-    if (INSTANCE == null) {
-      INSTANCE = new InstrumentationChecker(classLoader);
-    }
-
-    return INSTANCE.doGetUnsupportedRules();
-  }
-
-  private static Map<String, String> scanLoadedLibraries() {
-
-    final Map<String, String> frameworks = new HashMap<>();
-
-    // Scan classpath provided jars
-    final List<File> jars = getJarFiles(System.getProperty("java.class.path"));
-    for (final File file : jars) {
-
-      final String jarName = file.getName();
-      final String version = extractJarVersion(jarName);
-
-      if (version != null) {
-
-        // Extract artifactId
-        final String artifactId = file.getName().substring(0, jarName.indexOf(version) - 1);
-
-        // Store it
-        frameworks.put(artifactId, version);
-      }
-    }
-    log.debug("{} libraries found in the class-path", frameworks.size());
-
-    return frameworks;
-  }
-
-  private static List<File> getJarFiles(final String paths) {
-    final List<File> filesList = new ArrayList<>();
-    for (final String path : paths.split(File.pathSeparator)) {
-      final File file = new File(path);
-      if (file.isDirectory()) {
-        recurse(filesList, file);
-      } else {
-        if (file.getName().endsWith(".jar")) {
-          log.trace("{} found in the classpath", file.getName());
-          filesList.add(file);
-        }
-      }
-    }
-    return filesList;
-  }
-
-  private static void recurse(final List<File> filesList, final File f) {
-    final File[] list = f.listFiles();
-    for (final File file : list) {
-      getJarFiles(file.getPath());
-    }
   }
 
-  private static String extractJarVersion(final String jarName) {
-
-    final Pattern versionPattern = Pattern.compile("-(\\d+\\..+)\\.jar");
-    final Matcher matcher = versionPattern.matcher(jarName);
-    if (matcher.find()) {
-      return matcher.group(1);
-    } else {
-      return null;
-    }
-  }
-
-  private List<String> doGetUnsupportedRules() {
+  public List<String> getUnsupportedRules(ClassLoader classLoader) {
+    log.debug("Checking rule compatibility on classloader {}", classLoader);
 
     final List<String> unsupportedRules = new ArrayList<>();
     for (final String rule : rules.keySet()) {
@@ -131,58 +45,54 @@ private List<String> doGetUnsupportedRules() {
       for (final ArtifactSupport check : rules.get(rule)) {
         log.debug("Checking rule {}", check);
 
-        boolean matched = true;
-        for (final Map.Entry<String, String> identifier :
-            check.identifyingPresentClasses.entrySet()) {
-          final boolean classPresent = isClassPresent(identifier.getKey());
-          if (!classPresent) {
-            log.debug("Instrumentation {} not applied due to missing class {}.", rule, identifier);
-          } else {
-            String identifyingMethod = identifier.getValue();
-            if (identifyingMethod != null && !identifyingMethod.isEmpty()) {
-              Class clazz = getClassIfPresent(identifier.getKey(), classLoader);
-              // already confirmed above the class is there.
-              Method[] declaredMethods = clazz.getDeclaredMethods();
-              boolean methodFound = false;
-              for (Method m : declaredMethods) {
-                if (m.getName().equals(identifyingMethod)) {
-                  methodFound = true;
-                  break;
+        boolean matched =
+            (check.identifyingPresentClasses != null
+                    && !check.identifyingPresentClasses.entrySet().isEmpty())
+                || (check.identifyingMissingClasses != null
+                    && !check.identifyingMissingClasses.isEmpty());
+        if (check.identifyingPresentClasses != null) {
+          for (final Map.Entry<String, String> identifier :
+              check.identifyingPresentClasses.entrySet()) {
+            final boolean classPresent = isClassPresent(identifier.getKey(), classLoader);
+            if (!classPresent) {
+              log.debug(
+                  "Instrumentation {} not applied due to missing class {}.", rule, identifier);
+            } else {
+              String identifyingMethod = identifier.getValue();
+              if (identifyingMethod != null && !identifyingMethod.isEmpty()) {
+                Class clazz = getClassIfPresent(identifier.getKey(), classLoader);
+                // already confirmed above the class is there.
+                Method[] declaredMethods = clazz.getDeclaredMethods();
+                boolean methodFound = false;
+                for (Method m : declaredMethods) {
+                  if (m.getName().equals(identifyingMethod)) {
+                    methodFound = true;
+                    break;
+                  }
+                }
+                if (!methodFound) {
+                  log.debug(
+                      "Instrumentation {} not applied due to missing method {}.{}",
+                      rule,
+                      identifier.getKey(),
+                      identifyingMethod);
+                  matched = false;
                 }
-              }
-              if (!methodFound) {
-                log.debug(
-                    "Instrumentation {} not applied due to missing method {}.{}",
-                    rule,
-                    identifier.getKey(),
-                    identifyingMethod);
-                matched = false;
               }
             }
+            matched &= classPresent;
           }
-          matched &= classPresent;
-        }
-        for (final String identifyingClass : check.identifyingMissingClasses) {
-          final boolean classMissing = !isClassPresent(identifyingClass);
-          if (!classMissing) {
-            log.debug(
-                "Instrumentation {} not applied due to present class {}.", rule, identifyingClass);
-          }
-          matched &= classMissing;
         }
-
-        final boolean useVersionMatching =
-            frameworks.containsKey(check.artifact)
-                && check.identifyingMissingClasses.isEmpty()
-                && check.identifyingPresentClasses.isEmpty();
-        if (useVersionMatching) {
-          // If no classes to scan, fall back on version regex.
-          matched = Pattern.matches(check.supportedVersion, frameworks.get(check.artifact));
-          if (!matched) {
-            log.debug(
-                "Library conflict: supported_version={}, actual_version={}",
-                check.supportedVersion,
-                frameworks.get(check.artifact));
+        if (check.identifyingMissingClasses != null) {
+          for (final String identifyingClass : check.identifyingMissingClasses) {
+            final boolean classMissing = !isClassPresent(identifyingClass, classLoader);
+            if (!classMissing) {
+              log.debug(
+                  "Instrumentation {} not applied due to present class {}.",
+                  rule,
+                  identifyingClass);
+            }
+            matched &= classMissing;
           }
         }
 
@@ -201,10 +111,6 @@ private List<String> doGetUnsupportedRules() {
     return unsupportedRules;
   }
 
-  private boolean isClassPresent(final String identifyingPresentClass) {
-    return isClassPresent(identifyingPresentClass, classLoader);
-  }
-
   static boolean isClassPresent(final String identifyingPresentClass, ClassLoader classLoader) {
     return getClassIfPresent(identifyingPresentClass, classLoader) != null;
   }