Re-use AWS SDK module to instrument EMR's shaded copy

Author: mcculls

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

@@ -73,6 +73,7 @@ public final class CombiningTransformerBuilder
   private AgentBuilder.Transformer contextRequestRewriter;
   private HelperTransformer helperTransformer;
   private Advice.PostProcessor.Factory postProcessor;
+  private AdviceShader adviceShader;
   private MuzzleCheck muzzle;
 
   // temporary buffer for collecting advice; reset for each instrumenter
@@ -125,10 +126,14 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati
     helperTransformer =
         helperClassNames.length > 0
             ? new HelperTransformer(
-                module.useAgentCodeSource(), module.getClass().getSimpleName(), helperClassNames)
+                module.useAgentCodeSource(),
+                adviceShader,
+                module.getClass().getSimpleName(),
+                helperClassNames)
             : null;
 
     postProcessor = module.postProcessor();
+    adviceShader = AdviceShader.with(module.adviceShading());
 
     muzzle = new MuzzleCheck(module, instrumentationId);
   }
@@ -238,11 +243,14 @@ 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));
+            .withExceptionHandler(ExceptionHandlers.defaultExceptionHandler());
+    if (adviceShader != null) {
+      forAdvice = forAdvice.with(ShadedClassFiles.strategy(adviceShader));
+    }
+    advice.add(forAdvice.advice(not(ignoredMethods).and(matcher), adviceClass));
   }
 
   public ClassFileTransformer installOn(Instrumentation instrumentation) {
@@ -342,8 +350,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/ShadedClassFiles.java

@@ -0,0 +1,39 @@
+package datadog.trace.agent.tooling;
+
+import static net.bytebuddy.agent.builder.AgentBuilder.LocationStrategy.ForClassLoader.STRONG;
+
+import java.io.IOException;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.dynamic.ClassFileLocator;
+
+/** Locates and shades class-file resources. */
+public final class ShadedClassFiles implements ClassFileLocator {
+  private final ClassFileLocator delegate;
+  private final AdviceShader adviceShader;
+
+  /** Strategy that uses the given shader to shade any located class-file resources. */
+  public static AgentBuilder.LocationStrategy strategy(AdviceShader adviceShader) {
+    return (classLoader, javaModule) ->
+        new ShadedClassFiles(STRONG.classFileLocator(classLoader, javaModule), adviceShader);
+  }
+
+  ShadedClassFiles(ClassFileLocator delegate, AdviceShader adviceShader) {
+    this.delegate = delegate;
+    this.adviceShader = adviceShader;
+  }
+
+  @Override
+  public Resolution locate(String className) throws IOException {
+    final Resolution resolution = delegate.locate(className);
+    if (resolution.isResolved()) {
+      return new Resolution.Explicit(adviceShader.shade(resolution.resolve()));
+    } else {
+      return resolution;
+    }
+  }
+
+  @Override
+  public void close() throws IOException {
+    delegate.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(cr, 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;
@@ -161,6 +162,11 @@ public Advice.PostProcessor.Factory postProcessor() {
     return null;
   }
 
+  /** Override this to apply shading to method advice and injected helpers. */
+  public Function<String, String> adviceShading() {
+    return null;
+  }
+
   /**
    * Context stores to define for this instrumentation. Are added to matching class loaders.
    *

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()) {

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AWSHttpClientInstrumentation.java

@@ -22,10 +22,15 @@
  */
 public class AWSHttpClientInstrumentation
     implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+  private final String targetPackage;
+
+  public AWSHttpClientInstrumentation(String targetPackage) {
+    this.targetPackage = targetPackage;
+  }
 
   @Override
   public String instrumentedType() {
-    return "com.amazonaws.http.AmazonHttpClient";
+    return targetPackage + ".http.AmazonHttpClient";
   }
 
   @Override

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkModule.java

@@ -9,10 +9,16 @@
 
 /** Groups the instrumentations for AWS SDK 1.11.0+. */
 @AutoService(InstrumenterModule.class)
-public final class AwsSdkModule extends InstrumenterModule.Tracing {
+public class AwsSdkModule extends InstrumenterModule.Tracing {
+  private final String targetPackage;
 
   public AwsSdkModule() {
+    this("com.amazonaws");
+  }
+
+  public AwsSdkModule(String targetPackage) {
     super("aws-sdk");
+    this.targetPackage = targetPackage;
   }
 
   @Override
@@ -30,18 +36,18 @@ public String[] helperClassNames() {
   @Override
   public Map<String, String> contextStore() {
     Map<String, String> map = new java.util.HashMap<>();
-    map.put("com.amazonaws.services.sqs.model.ReceiveMessageResult", "java.lang.String");
+    map.put(targetPackage + ".services.sqs.model.ReceiveMessageResult", "java.lang.String");
     map.put(
-        "com.amazonaws.AmazonWebServiceRequest",
+        targetPackage + ".AmazonWebServiceRequest",
         "datadog.trace.bootstrap.instrumentation.api.AgentSpan");
     return map;
   }
 
   @Override
   public List<Instrumenter> typeInstrumentations() {
     return Arrays.asList(
-        new AWSHttpClientInstrumentation(),
-        new RequestExecutorInstrumentation(),
-        new HandlerChainFactoryInstrumentation());
+        new AWSHttpClientInstrumentation(targetPackage),
+        new RequestExecutorInstrumentation(targetPackage),
+        new HandlerChainFactoryInstrumentation(targetPackage));
   }
 }