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 c8c55036e3e..21a0b4c6e47 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 @@ -119,7 +119,7 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati new FieldBackedContextRequestRewriter(contextStore, module.name())) : null; - adviceShader = AdviceShader.with(module.adviceShading()); + adviceShader = AdviceShader.with(module); String[] helperClassNames = module.helperClassNames(); if (module.injectHelperDependencies()) { 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 index 1fea0336d7e..3a02c7b78b4 100644 --- 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 @@ -1,7 +1,12 @@ package datadog.trace.agent.tooling; +import static java.util.Arrays.asList; +import static java.util.Collections.emptyList; + import datadog.trace.api.cache.DDCache; import datadog.trace.api.cache.DDCaches; +import java.util.HashMap; +import java.util.List; import java.util.Map; import net.bytebuddy.jar.asm.ClassReader; import net.bytebuddy.jar.asm.ClassVisitor; @@ -12,18 +17,34 @@ /** Shades advice bytecode by applying relocations to all references. */ public final class AdviceShader { private final Map relocations; + private final List helperNames; private volatile Remapper remapper; + /** + * Used when installing {@link InstrumenterModule}s. Ensures any injected helpers have unique + * names so the original and relocated modules can inject helpers into the same class-loader. + */ + public static AdviceShader with(InstrumenterModule module) { + if (module.adviceShading() != null) { + return new AdviceShader(module.adviceShading(), asList(module.helperClassNames())); + } + return null; + } + + /** + * Used to generate and check muzzle references. Only applies relocations declared in modules. + */ public static AdviceShader with(Map relocations) { if (relocations != null) { - return new AdviceShader(relocations); + return new AdviceShader(relocations, emptyList()); } return null; } - private AdviceShader(Map relocations) { + private AdviceShader(Map relocations, List helperNames) { this.relocations = relocations; + this.helperNames = helperNames; } /** Applies shading before calling the given {@link ClassVisitor}. */ @@ -42,13 +63,28 @@ public byte[] shadeClass(byte[] bytecode) { return cw.toByteArray(); } + /** Generates a unique shaded name for the given helper. */ + public String uniqueHelper(String dottedName) { + int packageEnd = dottedName.lastIndexOf('.'); + if (packageEnd > 0) { + return dottedName.substring(0, packageEnd + 1) + "shaded" + dottedName.substring(packageEnd); + } + return dottedName; + } + final class AdviceMapper extends Remapper { private final DDCache mappingCache = DDCaches.newFixedSizeCache(64); /** Flattened sequence of old-prefix, new-prefix relocations. */ private final String[] prefixes; + private final Map helperMapping; + AdviceMapper() { + this.helperMapping = new HashMap<>(helperNames.size() + 1, 1f); + for (String h : helperNames) { + this.helperMapping.put(h.replace('.', '/'), uniqueHelper(h).replace('.', '/')); + } // convert relocations to a flattened sequence: old-prefix, new-prefix, etc. this.prefixes = new String[relocations.size() * 2]; int i = 0; @@ -68,6 +104,10 @@ final class AdviceMapper extends Remapper { @Override public String map(String internalName) { + String uniqueName = helperMapping.get(internalName); + if (uniqueName != null) { + return uniqueName; + } if (internalName.startsWith("java/") || internalName.startsWith("datadog/") || internalName.startsWith("net/bytebuddy/")) { 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 9d0e7551d45..3c543731625 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 @@ -93,6 +93,7 @@ private Map getHelperMap() throws IOException { byte[] classBytes = classFileLocator.locate(helperName).resolve(); if (adviceShader != null) { classBytes = adviceShader.shadeClass(classBytes); + helperName = adviceShader.uniqueHelper(helperName); } classnameToBytes.put(helperName, classBytes); }