Add support for instrumentation based class injection via dd-instrument-java

Author: mcculls

dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java

@@ -6,6 +6,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.isDefaultFinalizer;
 
 import datadog.environment.SystemProperties;
+import datadog.instrument.classinject.ClassInjector;
 import datadog.trace.agent.tooling.bytebuddy.SharedTypePools;
 import datadog.trace.agent.tooling.bytebuddy.iast.TaintableRedefinitionStrategyListener;
 import datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers;
@@ -108,6 +109,10 @@ public static ClassFileTransformer installBytebuddyAgent(
       final AgentBuilder.Listener... listeners) {
     Utils.setInstrumentation(inst);
 
+    if (!InstrumenterConfig.get().isUnsafeClassInjection()) {
+      ClassInjector.enableClassInjection(inst);
+    }
+
     TypePoolFacade.registerAsSupplier();
 
     if (InstrumenterConfig.get().isResolverMemoizingEnabled()) {

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;
@@ -34,6 +36,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 +137,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 +146,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 +178,29 @@ 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 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 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-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java

@@ -159,6 +159,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;
   }