diff --git a/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java b/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java index 91aadc39b58..c8c55036e3e 100644 --- a/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java +++ b/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java @@ -71,6 +71,7 @@ public final class CombiningTransformerBuilder private ElementMatcher classLoaderMatcher; private Map contextStore; private AgentBuilder.Transformer contextRequestRewriter; + private AdviceShader adviceShader; private HelperTransformer helperTransformer; private Advice.PostProcessor.Factory postProcessor; private MuzzleCheck muzzle; @@ -118,6 +119,8 @@ 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); @@ -125,7 +128,10 @@ 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(); @@ -238,11 +244,17 @@ public void applyAdvice(ElementMatcher 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); } } diff --git a/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/ShadedAdviceLocator.java b/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/ShadedAdviceLocator.java new file mode 100644 index 00000000000..3386291959d --- /dev/null +++ b/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(); + } +} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AdviceShader.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AdviceShader.java new file mode 100644 index 00000000000..ac588c78d8b --- /dev/null +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AdviceShader.java @@ -0,0 +1,98 @@ +package datadog.trace.agent.tooling; + +import datadog.trace.api.cache.DDCache; +import datadog.trace.api.cache.DDCaches; +import java.util.Map; +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 relocations to all references. */ +public final class AdviceShader extends Remapper { + private final DDCache cache = DDCaches.newFixedSizeCache(64); + + /** Flattened sequence of old-prefix, new-prefix relocations. */ + private final String[] prefixes; + + public static AdviceShader with(Map relocations) { + return relocations != null ? new AdviceShader(relocations) : null; + } + + AdviceShader(Map relocations) { + // convert relocations to a flattened sequence: old-prefix, new-prefix, etc. + this.prefixes = new String[relocations.size() * 2]; + int i = 0; + for (Map.Entry e : relocations.entrySet()) { + String oldPrefix = e.getKey(); + String newPrefix = e.getValue(); + if (oldPrefix.indexOf('.') > 0) { + // accept dotted prefixes, but store them in their internal form + this.prefixes[i++] = oldPrefix.replace('.', '/'); + this.prefixes[i++] = newPrefix.replace('.', '/'); + } else { + this.prefixes[i++] = oldPrefix; + this.prefixes[i++] = newPrefix; + } + } + } + + /** 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 internalName) { + if (internalName.startsWith("java/") + || internalName.startsWith("datadog/") + || internalName.startsWith("net/bytebuddy/")) { + return internalName; // never shade these references + } + return cache.computeIfAbsent(internalName, this::shade); + } + + @Override + public Object mapValue(Object value) { + if (value instanceof String) { + String text = (String) value; + if (text.isEmpty()) { + return text; + } else if (text.indexOf('.') > 0) { + return shadeDottedName(text); + } else { + return shade(text); + } + } else { + return super.mapValue(value); + } + } + + private String shade(String internalName) { + for (int i = 0; i < prefixes.length; i += 2) { + if (internalName.startsWith(prefixes[i])) { + return prefixes[i + 1] + internalName.substring(prefixes[i].length()); + } + } + return internalName; + } + + private String shadeDottedName(String name) { + String internalName = name.replace('.', '/'); + for (int i = 0; i < prefixes.length; i += 2) { + if (internalName.startsWith(prefixes[i])) { + return prefixes[i + 1].replace('/', '.') + name.substring(prefixes[i].length()); + } + } + return name; + } +} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java index dafc828bc2a..d0dd6e6d1ab 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java +++ b/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 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 helperMap) { this.useAgentCodeSource = useAgentCodeSource; this.requestingName = requestingName; + this.adviceShader = null; helperClassNames = helperMap.keySet(); dynamicTypeMap.putAll(helperMap); @@ -78,9 +89,11 @@ public HelperInjector( private Map getHelperMap() throws IOException { if (dynamicTypeMap.isEmpty()) { final Map 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); } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java index 9a6d0d0b35f..0f4fa35dd92 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java @@ -156,6 +156,11 @@ public ElementMatcher methodIgnoreMatcher() { return isSynthetic(); } + /** Override this to apply shading to method advice and injected helpers. */ + public Map adviceShading() { + return null; + } + /** Override this to post-process the operand stack of any transformed methods. */ public Advice.PostProcessor.Factory postProcessor() { return null; diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java index 430644d028a..a2d5a035b1c 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java +++ b/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 referenceSources = new HashSet<>(); final Map 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 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 ignoredClassNames = new HashSet<>(Arrays.asList(module.muzzleIgnoredClassNames())); + AdviceShader adviceShader = AdviceShader.with(module.adviceShading()); List 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); diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPlugin.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPlugin.java index c4dc304ffcc..6f6c3016dc7 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPlugin.java +++ b/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 createHelperMap(final InstrumenterModule modu String[] helperClasses = module.helperClassNames(); final Map helperMap = new LinkedHashMap<>(helperClasses.length); Set 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 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; diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java index 65559894a08..2bc44533088 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java +++ b/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 createReferencesFrom( - final String entryPointClassName, final ClassLoader loader) throws IllegalStateException { + final String entryPointClassName, final AdviceShader adviceShader, final ClassLoader loader) + throws IllegalStateException { final Set visitedSources = new HashSet<>(); final Map references = new LinkedHashMap<>(); @@ -64,7 +67,11 @@ public static Map 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 instrumentationReferences = cv.getReferences(); for (final Map.Entry entry : instrumentationReferences.entrySet()) { @@ -88,6 +95,11 @@ public static Map createReferencesFrom( return references; } + public static Map createReferencesFrom( + final String entryPointClassName, final ClassLoader loader) { + return createReferencesFrom(entryPointClassName, null, loader); + } + private static boolean samePackage(String from, String to) { int fromLength = from.lastIndexOf('/'); int toLength = to.lastIndexOf('/');