From 51b4d946d97c7392507d5018040a4a4fb2b8c1b5 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Tue, 14 Jan 2025 09:51:16 +0000 Subject: [PATCH] Ensure shaded helpers have unique names when injected into class-loaders (#8193) * Delay creation of advice remapper+cache until needed (cherry picked from commit ede74e31902f08bca41aaf4793344928d7f95680) * Ensure shaded helpers have unique names when injected into class-loaders (cherry picked from commit e8e337c52d9712622d5b28100061fff11d3e602f) --- .../tooling/CombiningTransformerBuilder.java | 2 +- .../agent/tooling/ShadedAdviceLocator.java | 2 +- .../trace/agent/tooling/AdviceShader.java | 167 ++++++++++++------ .../trace/agent/tooling/HelperInjector.java | 9 +- .../muzzle/MuzzleVersionScanPlugin.java | 2 +- .../tooling/muzzle/ReferenceCreator.java | 2 +- 6 files changed, 120 insertions(+), 64 deletions(-) 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-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 index 3386291959d..484fa9fa55e 100644 --- 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 @@ -17,7 +17,7 @@ public ShadedAdviceLocator(ClassLoader adviceLoader, AdviceShader adviceShader) public Resolution locate(String className) throws IOException { final Resolution resolution = adviceLocator.locate(className); if (resolution.isResolved()) { - return new Resolution.Explicit(adviceShader.shade(resolution.resolve())); + return new Resolution.Explicit(adviceShader.shadeClass(resolution.resolve())); } else { return resolution; } 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 ac588c78d8b..0ad3fb1da34 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; @@ -10,89 +15,139 @@ 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); +public final class AdviceShader { + private final Map relocations; + private final List helperNames; - /** Flattened sequence of old-prefix, new-prefix relocations. */ - private final String[] prefixes; + private volatile Remapper remapper; - public static AdviceShader with(Map relocations) { - return relocations != null ? new AdviceShader(relocations) : null; + /** + * 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; } - 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; - } + /** 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, emptyList()); } + return null; + } + + private AdviceShader(Map relocations, List helperNames) { + this.relocations = relocations; + this.helperNames = helperNames; } /** Applies shading before calling the given {@link ClassVisitor}. */ - public ClassVisitor shade(ClassVisitor cv) { - return new ClassRemapper(cv, this); + public ClassVisitor shadeClass(ClassVisitor cv) { + if (null == remapper) { + remapper = new AdviceMapper(); + } + return new ClassRemapper(cv, remapper); } /** Returns the result of shading the given bytecode. */ - public byte[] shade(byte[] bytecode) { + public byte[] shadeClass(byte[] bytecode) { ClassReader cr = new ClassReader(bytecode); ClassWriter cw = new ClassWriter(null, 0); - cr.accept(shade(cw), 0); + cr.accept(shadeClass(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 + /** 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 cache.computeIfAbsent(internalName, this::shade); + return dottedName; } - @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); + 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() { + // record the unique names that we've given to injected helpers + 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; + 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; + } + } + } + + @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/")) { + return internalName; // never shade these references + } + return mappingCache.computeIfAbsent(internalName, this::shadeInternalName); + } + + @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 shadeInternalName(text); + } } else { - return shade(text); + return super.mapValue(value); } - } 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()); + private String shadeInternalName(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; } - 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()); + 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; } - 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 d0dd6e6d1ab..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 @@ -89,12 +89,13 @@ public HelperInjector( private Map getHelperMap() throws IOException { if (dynamicTypeMap.isEmpty()) { final Map classnameToBytes = new LinkedHashMap<>(); - for (final String helperClassName : helperClassNames) { - byte[] classBytes = classFileLocator.locate(helperClassName).resolve(); + for (String helperName : helperClassNames) { + byte[] classBytes = classFileLocator.locate(helperName).resolve(); if (adviceShader != null) { - classBytes = adviceShader.shade(classBytes); + classBytes = adviceShader.shadeClass(classBytes); + helperName = adviceShader.uniqueHelper(helperName); } - classnameToBytes.put(helperClassName, classBytes); + classnameToBytes.put(helperName, classBytes); } return classnameToBytes; 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 6f6c3016dc7..5e695deb7e2 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 @@ -132,7 +132,7 @@ private static Map createHelperMap(final InstrumenterModule modu ClassFileLocator.ForClassLoader.of(module.getClass().getClassLoader()); byte[] classBytes = locator.locate(helperName).resolve(); if (null != adviceShader) { - classBytes = adviceShader.shade(classBytes); + classBytes = adviceShader.shadeClass(classBytes); } helperMap.put(helperName, classBytes); } 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 2bc44533088..e2052c546c3 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 @@ -70,7 +70,7 @@ public static Map createReferencesFrom( if (null == adviceShader) { reader.accept(cv, ClassReader.SKIP_FRAMES); } else { - reader.accept(adviceShader.shade(cv), ClassReader.SKIP_FRAMES); + reader.accept(adviceShader.shadeClass(cv), ClassReader.SKIP_FRAMES); } final Map instrumentationReferences = cv.getReferences();