Enable instrumentation to scan for classes for version matching

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

@@ -1,14 +1,18 @@
 package com.datadoghq.trace.agent;
 
 import com.datadoghq.trace.resolver.FactoryUtils;
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.core.type.TypeReference;
 import java.io.File;
 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;
 
 /**
@@ -19,20 +23,23 @@ import lombok.extern.slf4j.Slf4j;
 public class InstrumentationChecker {
 
   private static final String CONFIG_FILE = "dd-trace-supported-framework";
-  private final Map<String, List<Map<String, String>>> rules;
+  private final Map<String, List<ArtifactSupport>> rules;
   private final Map<String, String> frameworks;
 
   private static InstrumentationChecker INSTANCE;
+  private final ClassLoader classLoader;
 
   /* For testing purpose */
   InstrumentationChecker(
-      final Map<String, List<Map<String, String>>> rules, final Map<String, String> frameworks) {
+      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() {
+  private InstrumentationChecker(final ClassLoader classLoader) {
+    this.classLoader = classLoader;
     rules =
         FactoryUtils.loadConfigFromResource(
             CONFIG_FILE, new TypeReference<Map<String, List<Map<String, String>>>>() {});
@@ -43,11 +50,12 @@ public class InstrumentationChecker {
    * Return a list of unsupported rules regarding loading deps
    *
    * @return the list of unsupported rules
+   * @param classLoader
    */
-  public static synchronized List<String> getUnsupportedRules() {
+  public static synchronized List<String> getUnsupportedRules(final ClassLoader classLoader) {
 
     if (INSTANCE == null) {
-      INSTANCE = new InstrumentationChecker();
+      INSTANCE = new InstrumentationChecker(classLoader);
     }
 
     return INSTANCE.doGetUnsupportedRules();
@@ -60,16 +68,24 @@ public class InstrumentationChecker {
 
       // Check rules
       boolean supported = false;
-      for (final Map<String, String> check : rules.get(rule)) {
+      for (final ArtifactSupport check : rules.get(rule)) {
-        if (frameworks.containsKey(check.get("artifact"))) {
+        if (frameworks.containsKey(check.artifact)) {
-          final boolean matched =
+          // If no classes to scan, fall back on version regex.
-              Pattern.matches(
+          boolean matched =
-                  check.get("supported_version"), frameworks.get(check.get("artifact")));
+              check.identifyingPresentClasses.isEmpty() && check.identifyingMissingClasses.isEmpty()
+                  ? Pattern.matches(check.supportedVersion, frameworks.get(check.artifact))
+                  : true;
+          for (final String identifyingClass : check.identifyingPresentClasses) {
+            matched &= isClassPresent(identifyingClass);
+          }
+          for (final String identifyingClass : check.identifyingMissingClasses) {
+            matched &= !isClassPresent(identifyingClass);
+          }
           if (!matched) {
             log.debug(
                 "Library conflict: supported_version={}, actual_version={}",
-                check.get("supported_version"),
+                check.supportedVersion,
-                frameworks.get(check.get("artifact")));
+                frameworks.get(check.artifact));
             supported = false;
             break;
           }
@@ -87,6 +103,15 @@ public class InstrumentationChecker {
     return unsupportedRules;
   }
 
+  private boolean isClassPresent(final String identifyingPresentClass) {
+    try {
+      return identifyingPresentClass != null
+          && Class.forName(identifyingPresentClass, false, classLoader) != null;
+    } catch (final ClassNotFoundException e) {
+      return false;
+    }
+  }
+
   private static Map<String, String> scanLoadedLibraries() {
 
     final Map<String, String> frameworks = new HashMap<>();
@@ -143,4 +168,19 @@ public class InstrumentationChecker {
       return null;
     }
   }
+
+  @Data
+  @JsonIgnoreProperties("check")
+  static class ArtifactSupport {
+    private String artifact;
+
+    @JsonProperty("supported_version")
+    private String supportedVersion;
+
+    @JsonProperty("identifying_present_classes")
+    private List<String> identifyingPresentClasses = Collections.emptyList();
+
+    @JsonProperty("identifying_missing_classes")
+    private List<String> identifyingMissingClasses = Collections.emptyList();
+  }
 }

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

@@ -11,7 +11,12 @@ import java.lang.reflect.Method;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URL;
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Enumeration;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import javassist.ClassPool;
@@ -58,7 +63,8 @@ public class TraceAnnotationsManager {
     final List<String> loadedScripts = loadRules(ClassLoader.getSystemClassLoader());
 
     //Check if some rules have to be uninstalled
-    final List<String> uninstallScripts = InstrumentationChecker.getUnsupportedRules();
+    final List<String> uninstallScripts =
+        InstrumentationChecker.getUnsupportedRules(ClassLoader.getSystemClassLoader());
     if (agentTracerConfig != null) {
       final List<String> disabledInstrumentations = agentTracerConfig.getDisabledInstrumentations();
       if (disabledInstrumentations != null && !disabledInstrumentations.isEmpty()) {

dd-java-agent/src/main/resources/dd-trace-supported-framework.yaml

@@ -1,55 +1,56 @@
 ### This file define all supported libraries
 ### Syntax:
 ###   <rulename>:                         # the rule name defined in the otarules.btm
-###     - check:                    # a list of tests, if one not matched, thus the rule is removed
+###     - artifact: <artifactId>          # the artifact name to be tested
-###       articfact: <artifactId>   # the artifact name to be tested
 ###       supported_version:              # a regex expression to express the version required by the rule
-###     - check: ...
+###       identifying_present_classes:    # a list of classes that distinctly identify the range of supported libraries
+###       - some.key.FrameworkClass       #    only if missing, supported_version regex is used.
+###     - ...
 
 opentracing-apache-httpclient:
-  - check:
+  - artifact: httpclient
-    artifact: httpclient
     supported_version: 4\.[3|4|5]\..*
-  - check:
+
-    artifact: commons-httpclient
+  - artifact: commons-httpclient
     supported_version: none
 
 opentracing-aws-sdk:
-  - check:
+  - artifact: aws-java-sdk
-    artifact: aws-java-sdk
     supported_version: 1\.11\..*
 
 opentracing-cassandra-driver:
-  - check:
+  - artifact: cassandra-driver-core
-    artifact: cassandra-driver-core
     supported_version: 3\.2.*
 
 opentracing-web-servlet-filter:
-  - check:
+  - artifact: jetty-server
-    artifact: jetty-server
     supported_version: (8\.|9\.).*
-  - check:
+
-    artifact: tomcat_catalina
+  - artifact: tomcat_catalina
     supported_version: (8\.|9\.).*
-  - check:
+
-    artifact: tomcat-embed-core
+  - artifact: tomcat-embed-core
     supported_version: (8\.|9\.).*
 
 opentracing-okhttp3:
-  - check:
+  - artifact: okhttp
-    artifact: okhttp
     supported_version: 3\..*
 
 # For rules opentracing-jms-2_{consumer,producer}
 opentracing-jms-2:
-  - check:
+  - artifact: javax.jms-api
-    artifact: javax.jms-api
     supported_version: 2\..*
 
 opentracing-mongo-driver:
-  - check:
+  - artifact: mongo-java-driver
-    artifact: mongo-java-driver
     supported_version: 3\..*
+    identifying_present_classes:
+      - com.mongodb.operation.AsyncReadOperation.class
+      - com.mongodb.client.model.MapReduceAction.class
+
   - check:
     artifact: mongodb-driver-async
     supported_version: 3\..*
+    identifying_present_classes:
+      - com.mongodb.operation.AsyncReadOperation.class
+      - com.mongodb.client.model.MapReduceAction.class

dd-java-agent/src/test/groovy/com/datadoghq/trace/agent/InstrumentationCheckerTest.groovy

@@ -6,7 +6,7 @@ import spock.lang.Specification
 
 class InstrumentationCheckerTest extends Specification {
   Map<String, List<Map<String, String>>> rules =
-    FactoryUtils.loadConfigFromResource("supported-version-test", new TypeReference<Map<String, List<Map<String, String>>>>() {
+    FactoryUtils.loadConfigFromResource("supported-version-test", new TypeReference<Map<String, List<InstrumentationChecker.ArtifactSupport>>>() {
     });
   Map<String, String> frameworks = [
     "artifact-1": "1.2.3.1232",
@@ -18,10 +18,16 @@ class InstrumentationCheckerTest extends Specification {
 
   def "test rules"() {
     setup:
-    def rules = InstrumentationChecker.getUnsupportedRules();
+    def rules = InstrumentationChecker.getUnsupportedRules(java.lang.ClassLoader.getSystemClassLoader());
 
     expect:
     rules.size() == 3
     rules.sort() == ["unsupportedRuleOne", "unsupportedRuleThree", "unsupportedRuleTwo"]
   }
+
+  static class DemoClass1 {}
+
+  static class DemoClass2 {}
+
+  static class DemoClass3 {}
 }

dd-java-agent/src/test/resources/supported-version-test.yaml

@@ -1,26 +1,41 @@
 unsupportedRuleOne:
-  - check:
+  - artifact: artifact-1
-    artifact: artifact-1
     supported_version: 1\.2\.3\..*
-  - check:
+
-    artifact: artifact-2
+  - artifact: artifact-2
     supported_version: none
+    identifying_present_classes:
+    -
+
 unsupportedRuleTwo:
-    - check:
+  - artifact: artifact-1
-      artifact: artifact-1
     supported_version: 2\.3\.
+    identifying_present_classes:
+    - com.datadoghq.trace.agent.InstrumentationCheckerTest$MissingClass
+
 supportedRuleOne:
-  - check:
+  - artifact: artifact-3
-    artifact: artifact-3
     supported_version: 5\..*
+    identifying_present_classes:
+      - com.datadoghq.trace.agent.InstrumentationCheckerTest$DemoClass1
+      - com.datadoghq.trace.agent.InstrumentationCheckerTest$DemoClass2
+      - com.datadoghq.trace.agent.InstrumentationCheckerTest$DemoClass3
+
 supportedRuleTwo:
-  - check:
+  - artifact: artifact-1
-    artifact: artifact-1
     supported_version: 1\.2\.3\..*
-  - check:
+    identifying_present_classes:
-    artifact: artifact-2
+      - com.datadoghq.trace.agent.InstrumentationCheckerTest$DemoClass1
+      - com.datadoghq.trace.agent.InstrumentationCheckerTest$DemoClass2
+
+  - artifact: artifact-2
     supported_version: 4\..*
+    identifying_present_classes:
+      - com.datadoghq.trace.agent.InstrumentationCheckerTest$DemoClass3
+
 unsupportedRuleThree:
-  - check:
+  - artifact: foo
-    artifact: foo
     supported_version: 1
+    identifying_present_classes:
+      - com.datadoghq.trace.agent.InstrumentationCheckerTest$DemoClass1
+      - com.datadoghq.trace.agent.InstrumentationCheckerTest$MissingClass