From 1c8c1b53809d6366bcb8d3263c220c0f3db1d8fa Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Mon, 13 Jan 2025 17:52:41 +0000 Subject: [PATCH 1/2] Delay creation of advice remapper+cache until needed (cherry picked from commit ede74e31902f08bca41aaf4793344928d7f95680) --- .../agent/tooling/ShadedAdviceLocator.java | 2 +- .../trace/agent/tooling/AdviceShader.java | 130 ++++++++++-------- .../trace/agent/tooling/HelperInjector.java | 8 +- .../muzzle/MuzzleVersionScanPlugin.java | 2 +- .../tooling/muzzle/ReferenceCreator.java | 2 +- 5 files changed, 80 insertions(+), 64 deletions(-) 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..1fea0336d7e 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 @@ -10,89 +10,105 @@ 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; - /** 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; + if (relocations != null) { + return new AdviceShader(relocations); + } + 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; - } - } + private AdviceShader(Map relocations) { + this.relocations = relocations; } /** 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 + final class AdviceMapper extends Remapper { + private final DDCache mappingCache = DDCaches.newFixedSizeCache(64); + + /** Flattened sequence of old-prefix, new-prefix relocations. */ + private final String[] prefixes; + + AdviceMapper() { + // 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; + } + } } - 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); + @Override + public String map(String internalName) { + 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..9d0e7551d45 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,12 @@ 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); } - 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(); From e431c3536c1c1e66157fadb8e3bf85c3abba6744 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Mon, 13 Jan 2025 20:35:56 +0000 Subject: [PATCH 2/2] Ensure shaded helpers have unique names when injected into class-loaders (cherry picked from commit e8e337c52d9712622d5b28100061fff11d3e602f) --- .../tooling/CombiningTransformerBuilder.java | 2 +- .../trace/agent/tooling/AdviceShader.java | 43 ++++++++++++++++++- .../trace/agent/tooling/HelperInjector.java | 1 + 3 files changed, 43 insertions(+), 3 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-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..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; @@ -12,18 +17,32 @@ /** 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 +61,29 @@ 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() { + // 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; @@ -68,6 +103,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); }