Merge pull request #781 from DataDog/mar-kolya/fix-helper-injection

Fix helper injection in context provider

Author: mar-kolya

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

@@ -63,7 +63,6 @@ public final AgentBuilder instrument(final AgentBuilder parentAgentBuilder) {
         return parentAgentBuilder;
       }
 
-      final MuzzleMatcher muzzleMatcher = new MuzzleMatcher();
       AgentBuilder.Identified.Extendable agentBuilder =
           parentAgentBuilder
               .type(
@@ -74,13 +73,13 @@ public final AgentBuilder instrument(final AgentBuilder parentAgentBuilder) {
                       classLoaderMatcher(),
                       "Instrumentation class loader matcher unexpected exception: "
                           + getClass().getName()))
-              .and(muzzleMatcher)
+              .and(new MuzzleMatcher())
               .and(new PostMatchHook())
               .transform(DDTransformers.defaultTransformers());
       agentBuilder = injectHelperClasses(agentBuilder);
       agentBuilder = contextProvider.instrumentationTransformer(agentBuilder);
       agentBuilder = applyInstrumentationTransformers(agentBuilder);
-      agentBuilder = contextProvider.additionalInstrumentation(agentBuilder, muzzleMatcher);
+      agentBuilder = contextProvider.additionalInstrumentation(agentBuilder);
       return agentBuilder;
     }
 

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedProvider.java

@@ -102,16 +102,23 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
   /** fields-accessor-interface-name -> fields-accessor-interface-dynamic-type */
   private final Map<String, DynamicType.Unloaded<?>> fieldAccessorInterfaces;
 
+  private final AgentBuilder.Transformer fieldAccessorInterfacesInjector;
+
   /** context-store-type-name -> context-store-type-name-dynamic-type */
   private final Map<String, DynamicType.Unloaded<?>> contextStoreImplementations;
 
+  private final AgentBuilder.Transformer contextStoreImplementationsInjector;
+
   private final boolean fieldInjectionEnabled;
 
   public FieldBackedProvider(final Instrumenter.Default instrumenter) {
     this.instrumenter = instrumenter;
     byteBuddy = new ByteBuddy();
     fieldAccessorInterfaces = generateFieldAccessorInterfaces();
+    fieldAccessorInterfacesInjector = bootstrapHelperInjector(fieldAccessorInterfaces.values());
     contextStoreImplementations = generateContextStoreImplementationClasses();
+    contextStoreImplementationsInjector =
+        bootstrapHelperInjector(contextStoreImplementations.values());
     fieldInjectionEnabled = Config.get().isRuntimeContextFieldInjection();
   }
 
@@ -125,46 +132,11 @@ public AgentBuilder.Identified.Extendable instrumentationTransformer(
        */
       builder =
           builder.transform(getTransformerForASMVisitor(getContextStoreReadsRewritingVisitor()));
-
-      /**
-       * We inject into bootstrap classloader because field accessor interfaces are needed by
-       * context store implementations. Unfortunately this forces us to remove stored type checking
-       * because actual classes may not be available at this point.
-       */
-      builder = builder.transform(bootstrapHelperInjector(fieldAccessorInterfaces.values()));
-
-      /**
-       * We inject context store implementation into bootstrap classloader because same
-       * implementation may be used by different instrumentations and it has to use same static map
-       * in case of fallback to map-backed storage.
-       */
-      builder = builder.transform(bootstrapHelperInjector(contextStoreImplementations.values()));
+      builder = injectHelpersIntoBootstrapClassloader(builder);
     }
     return builder;
   }
 
-  /** Get transformer that forces helper injection onto bootstrap classloader. */
-  private AgentBuilder.Transformer bootstrapHelperInjector(
-      final Collection<DynamicType.Unloaded<?>> helpers) {
-    return new AgentBuilder.Transformer() {
-      final HelperInjector injector = HelperInjector.forDynamicTypes(helpers);
-
-      @Override
-      public DynamicType.Builder<?> transform(
-          final DynamicType.Builder<?> builder,
-          final TypeDescription typeDescription,
-          final ClassLoader classLoader,
-          final JavaModule module) {
-        return injector.transform(
-            builder,
-            typeDescription,
-            // context store implementation classes will always go to the bootstrap
-            BOOTSTRAP_CLASSLOADER,
-            module);
-      }
-    };
-  }
-
   private AsmVisitorWrapper getContextStoreReadsRewritingVisitor() {
     return new AsmVisitorWrapper() {
       @Override
@@ -328,9 +300,49 @@ public void visitLdcInsn(final Object value) {
     };
   }
 
+  private AgentBuilder.Identified.Extendable injectHelpersIntoBootstrapClassloader(
+      AgentBuilder.Identified.Extendable builder) {
+    /**
+     * We inject into bootstrap classloader because field accessor interfaces are needed by context
+     * store implementations. Unfortunately this forces us to remove stored type checking because
+     * actual classes may not be available at this point.
+     */
+    builder = builder.transform(fieldAccessorInterfacesInjector);
+
+    /**
+     * We inject context store implementation into bootstrap classloader because same implementation
+     * may be used by different instrumentations and it has to use same static map in case of
+     * fallback to map-backed storage.
+     */
+    builder = builder.transform(contextStoreImplementationsInjector);
+    return builder;
+  }
+
+  /** Get transformer that forces helper injection onto bootstrap classloader. */
+  private AgentBuilder.Transformer bootstrapHelperInjector(
+      final Collection<DynamicType.Unloaded<?>> helpers) {
+    return new AgentBuilder.Transformer() {
+      final HelperInjector injector = HelperInjector.forDynamicTypes(helpers);
+
+      @Override
+      public DynamicType.Builder<?> transform(
+          final DynamicType.Builder<?> builder,
+          final TypeDescription typeDescription,
+          final ClassLoader classLoader,
+          final JavaModule module) {
+        return injector.transform(
+            builder,
+            typeDescription,
+            // context store implementation classes will always go to the bootstrap
+            BOOTSTRAP_CLASSLOADER,
+            module);
+      }
+    };
+  }
+
   @Override
   public AgentBuilder.Identified.Extendable additionalInstrumentation(
-      AgentBuilder.Identified.Extendable builder, final AgentBuilder.RawMatcher muzzleMatcher) {
+      AgentBuilder.Identified.Extendable builder) {
 
     if (fieldInjectionEnabled) {
       for (final Map.Entry<String, String> entry : instrumenter.contextStore().entrySet()) {
@@ -344,16 +356,19 @@ public AgentBuilder.Identified.Extendable additionalInstrumentation(
                     safeHasSuperType(named(entry.getKey())).and(not(isInterface())),
                     instrumenter.classLoaderMatcher())
                 .and(safeToInjectFieldsMatcher())
-                /**
-                 * By adding the muzzleMatcher here, we are adding risk that the rules for injecting
-                 * the classes into the classloader and the rules for adding the field to the class
-                 * might be different. However the consequences are much greater if the class is not
-                 * injected but the field is added, since that results in a NoClassDef error.
-                 */
-                .and(muzzleMatcher)
-                .transform(
-                    getTransformerForASMVisitor(
-                        getFieldInjectionVisitor(entry.getKey(), entry.getValue())));
+                .transform(AgentBuilder.Transformer.NoOp.INSTANCE);
+
+        /**
+         * We inject helpers here as well as when instrumentation is applied to ensure that helpers
+         * are present even if instrumented classes are not loaded, but classes with state fields
+         * added are loaded (e.g. sun.net.www.protocol.https.HttpsURLConnectionImpl).
+         */
+        builder = injectHelpersIntoBootstrapClassloader(builder);
+
+        builder =
+            builder.transform(
+                getTransformerForASMVisitor(
+                    getFieldInjectionVisitor(entry.getKey(), entry.getValue())));
       }
     }
     return builder;

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContextProvider.java

@@ -13,5 +13,5 @@ AgentBuilder.Identified.Extendable instrumentationTransformer(
 
   /** Hook to define additional instrumentation. Run at instrumentation advice is hooked up. */
   AgentBuilder.Identified.Extendable additionalInstrumentation(
-      AgentBuilder.Identified.Extendable builder, final AgentBuilder.RawMatcher muzzleMatcher);
+      AgentBuilder.Identified.Extendable builder);
 }

dd-java-agent/instrumentation/http-url-connection/src/test/groovy/HttpUrlConnectionTest.groovy

@@ -5,7 +5,9 @@ import io.opentracing.tag.Tags
 import io.opentracing.util.GlobalTracer
 import org.springframework.web.client.RestTemplate
 import spock.lang.AutoCleanup
+import spock.lang.Requires
 import spock.lang.Shared
+import sun.net.www.protocol.https.HttpsURLConnectionImpl
 
 import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
@@ -472,4 +474,14 @@ class HttpUrlConnectionTest extends AgentTestRunner {
     where:
     renameService << [false, true]
   }
+
+  // This test makes no sense on IBM JVM because there is no HttpsURLConnectionImpl class there
+  @Requires({ !System.getProperty("java.vm.name").contains("IBM J9 VM") })
+  def "Make sure we can load HttpsURLConnectionImpl"() {
+    when:
+    def instance = new HttpsURLConnectionImpl(null, null, null)
+
+    then:
+    instance != null
+  }
 }