Support instrumentation-based class injection via dd-instrument-java (#9901)

Author: mcculls

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java

@@ -21,6 +21,7 @@
 import datadog.environment.JavaVirtualMachine;
 import datadog.environment.OperatingSystem;
 import datadog.environment.SystemProperties;
+import datadog.instrument.classinject.ClassInjector;
 import datadog.trace.api.Config;
 import datadog.trace.api.Platform;
 import datadog.trace.api.StatsDClientManager;
@@ -208,6 +209,15 @@ public static void start(
     StaticEventLogger.begin("Agent");
     StaticEventLogger.begin("Agent.start");
 
+    try {
+      ClassInjector.enableClassInjection(inst);
+    } catch (Throwable e) {
+      log.debug("Instrumentation-based class injection is not available", e);
+      setSystemPropertyDefault(
+          propertyNameToSystemPropertyName(TraceInstrumentationConfig.UNSAFE_CLASS_INJECTION),
+          "true");
+    }
+
     createAgentClassloader(agentJarURL);
 
     if (Platform.isNativeImageBuilder()) {

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

@@ -3,13 +3,15 @@
 import static datadog.trace.bootstrap.AgentClassLoading.INJECTING_HELPERS;
 import static java.util.Arrays.asList;
 
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.instrumentation.api.EagerHelper;
 import datadog.trace.util.JDK9ModuleAccess;
 import java.io.IOException;
 import java.lang.ref.WeakReference;
 import java.lang.reflect.AnnotatedElement;
 import java.security.CodeSource;
 import java.security.ProtectionDomain;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
@@ -22,7 +24,6 @@
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.ClassFileLocator;
 import net.bytebuddy.dynamic.DynamicType;
-import net.bytebuddy.dynamic.loading.ClassInjector;
 import net.bytebuddy.utility.JavaModule;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -34,6 +35,9 @@ public class HelperInjector implements Instrumenter.TransformingAdvice {
   private static final ClassFileLocator classFileLocator =
       ClassFileLocator.ForClassLoader.of(Utils.getExtendedClassLoader());
 
+  private static final boolean unsafeClassInjection =
+      InstrumenterConfig.get().isUnsafeClassInjection();
+
   private final boolean useAgentCodeSource;
   private final AdviceShader adviceShader;
   private final String requestingName;
@@ -132,7 +136,7 @@ public DynamicType.Builder<?> transform(
           }
 
           final Map<String, byte[]> classnameToBytes = getHelperMap();
-          final Map<String, Class<?>> classes = injectClassLoader(classLoader, classnameToBytes);
+          final Collection<Class<?>> classes = injectClassLoader(classLoader, classnameToBytes);
 
           // all datadog helper classes are in the unnamed module
           // and there's exactly one unnamed module per classloader
@@ -141,7 +145,7 @@ public DynamicType.Builder<?> transform(
           }
 
           // forcibly initialize any eager helpers
-          for (Class<?> clazz : classes.values()) {
+          for (Class<?> clazz : classes) {
             if (EagerHelper.class.isAssignableFrom(clazz)) {
               try {
                 clazz.getMethod("init").invoke(null);
@@ -173,16 +177,30 @@ public DynamicType.Builder<?> transform(
     return builder;
   }
 
-  private Map<String, Class<?>> injectClassLoader(
+  private Collection<Class<?>> injectClassLoader(
       final ClassLoader classLoader, final Map<String, byte[]> classnameToBytes) {
     INJECTING_HELPERS.begin();
     try {
       if (useAgentCodeSource) {
         ProtectionDomain protectionDomain = createProtectionDomain(classLoader);
-        return new ClassInjector.UsingReflection(classLoader, protectionDomain)
-            .injectRaw(classnameToBytes);
+        if (unsafeClassInjection) {
+          return new net.bytebuddy.dynamic.loading.ClassInjector.UsingReflection(
+                  classLoader, protectionDomain)
+              .injectRaw(classnameToBytes)
+              .values();
+        } else {
+          return datadog.instrument.classinject.ClassInjector.injectClasses(
+              classnameToBytes, protectionDomain);
+        }
       } else {
-        return new ClassInjector.UsingReflection(classLoader).injectRaw(classnameToBytes);
+        if (unsafeClassInjection) {
+          return new net.bytebuddy.dynamic.loading.ClassInjector.UsingReflection(classLoader)
+              .injectRaw(classnameToBytes)
+              .values();
+        } else {
+          return datadog.instrument.classinject.ClassInjector.injectClasses(
+              classnameToBytes, classLoader);
+        }
       }
     } finally {
       INJECTING_HELPERS.end();

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/HelperInjectionTest.groovy

@@ -1,8 +1,10 @@
 package datadog.trace.agent.test
 
+import datadog.instrument.classinject.ClassInjector
 import datadog.trace.agent.tooling.HelperInjector
 import datadog.trace.agent.tooling.Utils
 import datadog.trace.test.util.DDSpecification
+import net.bytebuddy.agent.ByteBuddyAgent
 
 import java.lang.ref.WeakReference
 import java.util.concurrent.atomic.AtomicReference
@@ -15,6 +17,7 @@ class HelperInjectionTest extends DDSpecification {
 
   def "helpers injected to non-delegating classloader"() {
     setup:
+    ClassInjector.enableClassInjection(ByteBuddyAgent.getInstrumentation())
     HelperInjector injector = new HelperInjector(false, "test", HELPER_CLASS_NAME)
     AtomicReference<URLClassLoader> emptyLoader = new AtomicReference<>(new URLClassLoader(new URL[0], (ClassLoader) null))
 

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteTransformerTest.groovy

@@ -1,10 +1,12 @@
 package datadog.trace.agent.tooling.csi
 
+import datadog.instrument.classinject.ClassInjector
 import datadog.trace.agent.tooling.bytebuddy.csi.Advices
 import datadog.trace.agent.tooling.bytebuddy.csi.CallSiteTransformer
 import datadog.trace.agent.tooling.csi.CallSiteAdvice.MethodHandler
 import datadog.trace.api.function.TriFunction
 import groovy.transform.CompileDynamic
+import net.bytebuddy.agent.ByteBuddyAgent
 import net.bytebuddy.description.type.TypeDescription
 import net.bytebuddy.dynamic.DynamicType
 import net.bytebuddy.jar.asm.Opcodes
@@ -206,6 +208,7 @@ class CallSiteTransformerTest extends BaseCallSiteTest {
   @SuppressWarnings(['GroovyAccessibility', 'GroovyAssignabilityCheck'])
   void 'test call site transformer with helpers'() {
     setup:
+    ClassInjector.enableClassInjection(ByteBuddyAgent.getInstrumentation())
     final source = StringConcatExample
     final helper = InstrumentationHelper
     final customClassLoader = new ClassLoader() { }

dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/InstrumentationSpecification.groovy

@@ -16,6 +16,7 @@ import com.datadog.debugger.sink.ProbeStatusSink
 import com.google.common.collect.Sets
 import datadog.communication.ddagent.DDAgentFeaturesDiscovery
 import datadog.communication.monitor.Monitoring
+import datadog.instrument.classinject.ClassInjector
 import datadog.trace.agent.test.asserts.ListWriterAssert
 import datadog.trace.agent.test.datastreams.MockFeaturesDiscovery
 import datadog.trace.agent.test.datastreams.RecordingDatastreamsPayloadWriter
@@ -401,6 +402,8 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
       return trackingSpan
     }
 
+    ClassInjector.enableClassInjection(INSTRUMENTATION)
+
     // if a test enables the instrumentation it verifies,
     // the cache needs to be recomputed taking into account that instrumentation's matchers
     ClassLoaderMatchers.resetState()

dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java

@@ -160,6 +160,9 @@ public final class TraceInstrumentationConfig {
   public static final String RESOLVER_USE_URL_CACHES = "resolver.use.url.caches";
   public static final String RESOLVER_RESET_INTERVAL = "resolver.reset.interval";
   public static final String RESOLVER_NAMES_ARE_UNIQUE = "resolver.names.are.unique";
+
+  public static final String UNSAFE_CLASS_INJECTION = "unsafe.class.injection";
+
   public static final String CASSANDRA_KEYSPACE_STATEMENT_EXTRACTION_ENABLED =
       "trace.cassandra.keyspace.statement.extraction.enabled";
   public static final String COUCHBASE_INTERNAL_SPANS_ENABLED =

internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java

@@ -73,6 +73,7 @@
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_PEKKO_SCHEDULER_ENABLED;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_THREAD_POOL_EXECUTORS_EXCLUDE;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_ENABLED;
+import static datadog.trace.api.config.TraceInstrumentationConfig.UNSAFE_CLASS_INJECTION;
 import static datadog.trace.api.config.UsmConfig.USM_ENABLED;
 import static datadog.trace.util.CollectionUtils.tryMakeImmutableList;
 import static datadog.trace.util.CollectionUtils.tryMakeImmutableSet;
@@ -162,6 +163,8 @@ public class InstrumenterConfig {
   private final Boolean resolverUseUrlCaches;
   private final int resolverResetInterval;
 
+  private final boolean unsafeClassInjection;
+
   private final boolean runtimeContextFieldInjection;
   private final boolean serialVersionUIDFieldInjection;
 
@@ -280,6 +283,8 @@ private InstrumenterConfig() {
             ? 0
             : configProvider.getInteger(RESOLVER_RESET_INTERVAL, DEFAULT_RESOLVER_RESET_INTERVAL);
 
+    unsafeClassInjection = configProvider.getBoolean(UNSAFE_CLASS_INJECTION, false);
+
     runtimeContextFieldInjection =
         configProvider.getBoolean(
             RUNTIME_CONTEXT_FIELD_INJECTION, DEFAULT_RUNTIME_CONTEXT_FIELD_INJECTION);
@@ -505,6 +510,10 @@ public String getResolverCacheDir() {
     return resolverCacheDir;
   }
 
+  public boolean isUnsafeClassInjection() {
+    return unsafeClassInjection;
+  }
+
   public String getInstrumentationConfigId() {
     return instrumentationConfigId;
   }

metadata/supported-configurations.json

@@ -1340,6 +1340,7 @@
     "DD_TRIAGE_REPORT_DIR": ["A"],
     "DD_TRIAGE_REPORT_TRIGGER": ["A"],
     "DD_UNDERTOW_CONTINUATION": ["A"],
+    "DD_UNSAFE_CLASS_INJECTION": ["A"],
     "DD_USM_ENABLED": ["A"],
     "DD_VERSION": ["A"],
     "DD_WRITER_BAGGAGE_INJECT": ["A"],