Support instrumentation of repackaged libraries

Author: mcculls

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

@@ -71,6 +71,7 @@ public final class CombiningTransformerBuilder
   private ElementMatcher<ClassLoader> classLoaderMatcher;
   private Map<String, String> contextStore;
   private AgentBuilder.Transformer contextRequestRewriter;
+  private AdviceShader adviceShader;
   private HelperTransformer helperTransformer;
   private Advice.PostProcessor.Factory postProcessor;
   private MuzzleCheck muzzle;
@@ -118,14 +119,19 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati
                 new FieldBackedContextRequestRewriter(contextStore, module.name()))
             : null;
 
+    adviceShader = AdviceShader.with(module.adviceShading());
+
     String[] helperClassNames = module.helperClassNames();
     if (module.injectHelperDependencies()) {
       helperClassNames = HelperScanner.withClassDependencies(helperClassNames);
     }
     helperTransformer =
         helperClassNames.length > 0
             ? new HelperTransformer(
-                module.useAgentCodeSource(), module.getClass().getSimpleName(), helperClassNames)
+                module.useAgentCodeSource(),
+                adviceShader,
+                module.getClass().getSimpleName(),
+                helperClassNames)
             : null;
 
     postProcessor = module.postProcessor();
@@ -238,11 +244,17 @@ public void applyAdvice(ElementMatcher<? super MethodDescription> matcher, Strin
     if (postProcessor != null) {
       customMapping = customMapping.with(postProcessor);
     }
-    advice.add(
+    AgentBuilder.Transformer.ForAdvice forAdvice =
         new AgentBuilder.Transformer.ForAdvice(customMapping)
-            .include(Utils.getBootstrapProxy(), Utils.getExtendedClassLoader())
             .withExceptionHandler(ExceptionHandlers.defaultExceptionHandler())
-            .advice(not(ignoredMethods).and(matcher), adviceClass));
+            .include(Utils.getBootstrapProxy());
+    ClassLoader adviceLoader = Utils.getExtendedClassLoader();
+    if (adviceShader != null) {
+      forAdvice = forAdvice.include(new ShadedAdviceLocator(adviceLoader, adviceShader));
+    } else {
+      forAdvice = forAdvice.include(adviceLoader);
+    }
+    advice.add(forAdvice.advice(not(ignoredMethods).and(matcher), adviceClass));
   }
 
   public ClassFileTransformer installOn(Instrumentation instrumentation) {
@@ -342,8 +354,11 @@ public DynamicType.Builder<?> transform(
 
   static final class HelperTransformer extends HelperInjector implements AgentBuilder.Transformer {
     HelperTransformer(
-        boolean useAgentCodeSource, String requestingName, String... helperClassNames) {
-      super(useAgentCodeSource, requestingName, helperClassNames);
+        boolean useAgentCodeSource,
+        AdviceShader adviceShader,
+        String requestingName,
+        String... helperClassNames) {
+      super(useAgentCodeSource, adviceShader, requestingName, helperClassNames);
     }
   }
 

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

@@ -0,0 +1,30 @@
+package datadog.trace.agent.tooling;
+
+import java.io.IOException;
+import net.bytebuddy.dynamic.ClassFileLocator;
+
+/** Locates and shades class-file resources from the advice class-loader. */
+public final class ShadedAdviceLocator implements ClassFileLocator {
+  private final ClassFileLocator adviceLocator;
+  private final AdviceShader adviceShader;
+
+  public ShadedAdviceLocator(ClassLoader adviceLoader, AdviceShader adviceShader) {
+    this.adviceLocator = ClassFileLocator.ForClassLoader.of(adviceLoader);
+    this.adviceShader = adviceShader;
+  }
+
+  @Override
+  public Resolution locate(String className) throws IOException {
+    final Resolution resolution = adviceLocator.locate(className);
+    if (resolution.isResolved()) {
+      return new Resolution.Explicit(adviceShader.shade(resolution.resolve()));
+    } else {
+      return resolution;
+    }
+  }
+
+  @Override
+  public void close() throws IOException {
+    adviceLocator.close();
+  }
+}

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

@@ -0,0 +1,42 @@
+package datadog.trace.agent.tooling;
+
+import datadog.trace.api.cache.DDCache;
+import datadog.trace.api.cache.DDCaches;
+import java.util.function.Function;
+import net.bytebuddy.jar.asm.ClassReader;
+import net.bytebuddy.jar.asm.ClassVisitor;
+import net.bytebuddy.jar.asm.ClassWriter;
+import net.bytebuddy.jar.asm.commons.ClassRemapper;
+import net.bytebuddy.jar.asm.commons.Remapper;
+
+/** Shades advice bytecode by applying a shading function to all references. */
+public final class AdviceShader extends Remapper {
+  private final DDCache<String, String> cache = DDCaches.newFixedSizeCache(64);
+  private final Function<String, String> shading;
+
+  public static AdviceShader with(Function<String, String> shading) {
+    return shading != null ? new AdviceShader(shading) : null;
+  }
+
+  AdviceShader(Function<String, String> shading) {
+    this.shading = shading;
+  }
+
+  /** Applies shading before calling the given {@link ClassVisitor}. */
+  public ClassVisitor shade(ClassVisitor cv) {
+    return new ClassRemapper(cv, this);
+  }
+
+  /** Returns the result of shading the given bytecode. */
+  public byte[] shade(byte[] bytecode) {
+    ClassReader cr = new ClassReader(bytecode);
+    ClassWriter cw = new ClassWriter(null, 0);
+    cr.accept(shade(cw), 0);
+    return cw.toByteArray();
+  }
+
+  @Override
+  public String map(String name) {
+    return cache.computeIfAbsent(name, shading);
+  }
+}

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

@@ -33,6 +33,7 @@ public class HelperInjector implements Instrumenter.TransformingAdvice {
       ClassFileLocator.ForClassLoader.of(Utils.getExtendedClassLoader());
 
   private final boolean useAgentCodeSource;
+  private final AdviceShader adviceShader;
   private final String requestingName;
 
   private final Set<String> helperClassNames;
@@ -58,8 +59,17 @@ public HelperInjector(
       final boolean useAgentCodeSource,
       final String requestingName,
       final String... helperClassNames) {
+    this(useAgentCodeSource, null, requestingName, helperClassNames);
+  }
+
+  public HelperInjector(
+      final boolean useAgentCodeSource,
+      final AdviceShader adviceShader,
+      final String requestingName,
+      final String... helperClassNames) {
     this.useAgentCodeSource = useAgentCodeSource;
     this.requestingName = requestingName;
+    this.adviceShader = adviceShader;
 
     this.helperClassNames = new LinkedHashSet<>(Arrays.asList(helperClassNames));
   }
@@ -70,6 +80,7 @@ public HelperInjector(
       final Map<String, byte[]> helperMap) {
     this.useAgentCodeSource = useAgentCodeSource;
     this.requestingName = requestingName;
+    this.adviceShader = null;
 
     helperClassNames = helperMap.keySet();
     dynamicTypeMap.putAll(helperMap);
@@ -78,9 +89,11 @@ public HelperInjector(
   private Map<String, byte[]> getHelperMap() throws IOException {
     if (dynamicTypeMap.isEmpty()) {
       final Map<String, byte[]> classnameToBytes = new LinkedHashMap<>();
-
       for (final String helperClassName : helperClassNames) {
-        final byte[] classBytes = classFileLocator.locate(helperClassName).resolve();
+        byte[] classBytes = classFileLocator.locate(helperClassName).resolve();
+        if (adviceShader != null) {
+          classBytes = adviceShader.shade(classBytes);
+        }
         classnameToBytes.put(helperClassName, classBytes);
       }
 

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

@@ -21,6 +21,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -156,6 +157,11 @@ public ElementMatcher<? super MethodDescription> methodIgnoreMatcher() {
     return isSynthetic();
   }
 
+  /** Override this to apply shading to method advice and injected helpers. */
+  public Function<String, String> adviceShading() {
+    return null;
+  }
+
   /** Override this to post-process the operand stack of any transformed methods. */
   public Advice.PostProcessor.Factory postProcessor() {
     return null;

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java

@@ -1,5 +1,6 @@
 package datadog.trace.agent.tooling.muzzle;
 
+import datadog.trace.agent.tooling.AdviceShader;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import java.io.File;
@@ -78,7 +79,8 @@ public ClassVisitor wrap(
     return classVisitor;
   }
 
-  private static Reference[] generateReferences(Instrumenter.HasMethodAdvice instrumenter) {
+  private static Reference[] generateReferences(
+      Instrumenter.HasMethodAdvice instrumenter, AdviceShader adviceShader) {
     // track sources we've generated references from to avoid recursion
     final Set<String> referenceSources = new HashSet<>();
     final Map<String, Reference> references = new LinkedHashMap<>();
@@ -88,7 +90,8 @@ private static Reference[] generateReferences(Instrumenter.HasMethodAdvice instr
     for (String adviceClass : adviceClasses) {
       if (referenceSources.add(adviceClass)) {
         for (Map.Entry<String, Reference> entry :
-            ReferenceCreator.createReferencesFrom(adviceClass, contextClassLoader).entrySet()) {
+            ReferenceCreator.createReferencesFrom(adviceClass, adviceShader, contextClassLoader)
+                .entrySet()) {
           Reference toMerge = references.get(entry.getKey());
           if (null == toMerge) {
             references.put(entry.getKey(), entry.getValue());
@@ -105,12 +108,13 @@ private static Reference[] generateReferences(Instrumenter.HasMethodAdvice instr
   private static byte[] generateMuzzleClass(InstrumenterModule module) {
 
     Set<String> ignoredClassNames = new HashSet<>(Arrays.asList(module.muzzleIgnoredClassNames()));
+    AdviceShader adviceShader = AdviceShader.with(module.adviceShading());
 
     List<Reference> references = new ArrayList<>();
     for (Instrumenter instrumenter : module.typeInstrumentations()) {
       if (instrumenter instanceof Instrumenter.HasMethodAdvice) {
         for (Reference reference :
-            generateReferences((Instrumenter.HasMethodAdvice) instrumenter)) {
+            generateReferences((Instrumenter.HasMethodAdvice) instrumenter, adviceShader)) {
           // ignore helper classes, they will be injected by the instrumentation's HelperInjector.
           if (!ignoredClassNames.contains(reference.className)) {
             references.add(reference);

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPlugin.java

@@ -1,5 +1,6 @@
 package datadog.trace.agent.tooling.muzzle;
 
+import datadog.trace.agent.tooling.AdviceShader;
 import datadog.trace.agent.tooling.HelperInjector;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.agent.tooling.bytebuddy.SharedTypePools;
@@ -109,6 +110,7 @@ private static Map<String, byte[]> createHelperMap(final InstrumenterModule modu
     String[] helperClasses = module.helperClassNames();
     final Map<String, byte[]> helperMap = new LinkedHashMap<>(helperClasses.length);
     Set<String> helperClassNames = new HashSet<>(Arrays.asList(helperClasses));
+    AdviceShader adviceShader = AdviceShader.with(module.adviceShading());
     for (final String helperName : helperClasses) {
       int nestedClassIndex = helperName.lastIndexOf('$');
       if (nestedClassIndex > 0) {
@@ -128,7 +130,10 @@ private static Map<String, byte[]> createHelperMap(final InstrumenterModule modu
       }
       final ClassFileLocator locator =
           ClassFileLocator.ForClassLoader.of(module.getClass().getClassLoader());
-      final byte[] classBytes = locator.locate(helperName).resolve();
+      byte[] classBytes = locator.locate(helperName).resolve();
+      if (null != adviceShader) {
+        classBytes = adviceShader.shade(classBytes);
+      }
       helperMap.put(helperName, classBytes);
     }
     return helperMap;

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java

@@ -3,6 +3,7 @@
 import static datadog.trace.util.Strings.getClassName;
 import static datadog.trace.util.Strings.getResourceName;
 
+import datadog.trace.agent.tooling.AdviceShader;
 import datadog.trace.bootstrap.Constants;
 import java.io.InputStream;
 import java.util.ArrayDeque;
@@ -41,12 +42,14 @@ public class ReferenceCreator extends ClassVisitor {
    * Generate all references reachable from a given class.
    *
    * @param entryPointClassName Starting point for generating references.
+   * @param adviceShader Optional shading to apply to the advice.
    * @param loader Classloader used to read class bytes.
    * @return Map of [referenceClassName -> Reference]
    * @throws IllegalStateException if class is not found or unable to be loaded.
    */
   public static Map<String, Reference> createReferencesFrom(
-      final String entryPointClassName, final ClassLoader loader) throws IllegalStateException {
+      final String entryPointClassName, final AdviceShader adviceShader, final ClassLoader loader)
+      throws IllegalStateException {
     final Set<String> visitedSources = new HashSet<>();
     final Map<String, Reference> references = new LinkedHashMap<>();
 
@@ -64,7 +67,11 @@ public static Map<String, Reference> createReferencesFrom(
         }
         final ReferenceCreator cv = new ReferenceCreator(null);
         final ClassReader reader = new ClassReader(in);
-        reader.accept(cv, ClassReader.SKIP_FRAMES);
+        if (null == adviceShader) {
+          reader.accept(cv, ClassReader.SKIP_FRAMES);
+        } else {
+          reader.accept(adviceShader.shade(cv), ClassReader.SKIP_FRAMES);
+        }
 
         final Map<String, Reference> instrumentationReferences = cv.getReferences();
         for (final Map.Entry<String, Reference> entry : instrumentationReferences.entrySet()) {