Fix helper injection in context provider

Currently helpers will not be injected if instrumentation doesn't
apply. Unfortunately it is possible for some classes to have context
fields added even if instrumentation is not allied (i.g. instrumented
classes are not used). Fix this by always injecting all helpers if we
inject context fields.

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
+  }
 }